1use super::{http::HttpClient, proto, Client, Status, TypedEnvelope};
2use anyhow::{anyhow, Context, Result};
3use futures::{future, AsyncReadExt};
4use gpui::{AsyncAppContext, Entity, ImageData, ModelContext, ModelHandle, Task};
5use postage::{prelude::Stream, sink::Sink, watch};
6use rpc::proto::{RequestMessage, UsersResponse};
7use std::{
8 collections::{hash_map::Entry, HashMap, HashSet},
9 sync::{Arc, Weak},
10};
11use util::TryFutureExt as _;
12
13#[derive(Debug)]
14pub struct User {
15 pub id: u64,
16 pub github_login: String,
17 pub avatar: Option<Arc<ImageData>>,
18}
19
20#[derive(Debug)]
21pub struct Contact {
22 pub user: Arc<User>,
23 pub projects: Vec<ProjectMetadata>,
24}
25
26#[derive(Debug)]
27pub struct ProjectMetadata {
28 pub id: u64,
29 pub is_shared: bool,
30 pub worktree_root_names: Vec<String>,
31 pub guests: Vec<Arc<User>>,
32}
33
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum ContactRequestStatus {
36 None,
37 Pending,
38 RequestSent,
39 RequestReceived,
40 RequestAccepted,
41}
42
43pub struct UserStore {
44 users: HashMap<u64, Arc<User>>,
45 update_contacts_tx: watch::Sender<Option<proto::UpdateContacts>>,
46 current_user: watch::Receiver<Option<Arc<User>>>,
47 contacts: Vec<Arc<Contact>>,
48 incoming_contact_requests: Vec<Arc<User>>,
49 outgoing_contact_requests: Vec<Arc<User>>,
50 pending_contact_requests: HashMap<u64, usize>,
51 client: Weak<Client>,
52 http: Arc<dyn HttpClient>,
53 _maintain_contacts: Task<()>,
54 _maintain_current_user: Task<()>,
55}
56
57pub enum Event {}
58
59impl Entity for UserStore {
60 type Event = Event;
61}
62
63impl UserStore {
64 pub fn new(
65 client: Arc<Client>,
66 http: Arc<dyn HttpClient>,
67 cx: &mut ModelContext<Self>,
68 ) -> Self {
69 let (mut current_user_tx, current_user_rx) = watch::channel();
70 let (update_contacts_tx, mut update_contacts_rx) =
71 watch::channel::<Option<proto::UpdateContacts>>();
72 let rpc_subscription =
73 client.add_message_handler(cx.handle(), Self::handle_update_contacts);
74 Self {
75 users: Default::default(),
76 current_user: current_user_rx,
77 contacts: Default::default(),
78 incoming_contact_requests: Default::default(),
79 outgoing_contact_requests: Default::default(),
80 client: Arc::downgrade(&client),
81 update_contacts_tx,
82 http,
83 _maintain_contacts: cx.spawn_weak(|this, mut cx| async move {
84 let _subscription = rpc_subscription;
85 while let Some(message) = update_contacts_rx.recv().await {
86 if let Some((message, this)) = message.zip(this.upgrade(&cx)) {
87 this.update(&mut cx, |this, cx| this.update_contacts(message, cx))
88 .log_err()
89 .await;
90 }
91 }
92 }),
93 _maintain_current_user: cx.spawn_weak(|this, mut cx| async move {
94 let mut status = client.status();
95 while let Some(status) = status.recv().await {
96 match status {
97 Status::Connected { .. } => {
98 if let Some((this, user_id)) = this.upgrade(&cx).zip(client.user_id()) {
99 let user = this
100 .update(&mut cx, |this, cx| this.fetch_user(user_id, cx))
101 .log_err()
102 .await;
103 current_user_tx.send(user).await.ok();
104 }
105 }
106 Status::SignedOut => {
107 current_user_tx.send(None).await.ok();
108 }
109 _ => {}
110 }
111 }
112 }),
113 pending_contact_requests: Default::default(),
114 }
115 }
116
117 async fn handle_update_contacts(
118 this: ModelHandle<Self>,
119 msg: TypedEnvelope<proto::UpdateContacts>,
120 _: Arc<Client>,
121 mut cx: AsyncAppContext,
122 ) -> Result<()> {
123 this.update(&mut cx, |this, _| {
124 *this.update_contacts_tx.borrow_mut() = Some(msg.payload);
125 });
126 Ok(())
127 }
128
129 fn update_contacts(
130 &mut self,
131 message: proto::UpdateContacts,
132 cx: &mut ModelContext<Self>,
133 ) -> Task<Result<()>> {
134 log::info!("update contacts on client {:?}", message);
135 let mut user_ids = HashSet::new();
136 for contact in &message.contacts {
137 user_ids.insert(contact.user_id);
138 user_ids.extend(contact.projects.iter().flat_map(|w| &w.guests).copied());
139 }
140 user_ids.extend(message.incoming_requests.iter().map(|req| req.requester_id));
141 user_ids.extend(message.outgoing_requests.iter());
142
143 let load_users = self.get_users(user_ids.into_iter().collect(), cx);
144 cx.spawn(|this, mut cx| async move {
145 load_users.await?;
146
147 // Users are fetched in parallel above and cached in call to get_users
148 // No need to paralellize here
149 let mut updated_contacts = Vec::new();
150 for contact in message.contacts {
151 updated_contacts.push(Arc::new(
152 Contact::from_proto(contact, &this, &mut cx).await?,
153 ));
154 }
155
156 let mut incoming_requests = Vec::new();
157 for request in message.incoming_requests {
158 incoming_requests.push(
159 this.update(&mut cx, |this, cx| {
160 this.fetch_user(request.requester_id, cx)
161 })
162 .await?,
163 );
164 }
165
166 let mut outgoing_requests = Vec::new();
167 for requested_user_id in message.outgoing_requests {
168 outgoing_requests.push(
169 this.update(&mut cx, |this, cx| this.fetch_user(requested_user_id, cx))
170 .await?,
171 );
172 }
173
174 let removed_contacts =
175 HashSet::<u64>::from_iter(message.remove_contacts.iter().copied());
176 let removed_incoming_requests =
177 HashSet::<u64>::from_iter(message.remove_incoming_requests.iter().copied());
178 let removed_outgoing_requests =
179 HashSet::<u64>::from_iter(message.remove_outgoing_requests.iter().copied());
180
181 this.update(&mut cx, |this, cx| {
182 // Remove contacts
183 this.contacts
184 .retain(|contact| !removed_contacts.contains(&contact.user.id));
185 // Update existing contacts and insert new ones
186 for updated_contact in updated_contacts {
187 match this
188 .contacts
189 .binary_search_by_key(&&updated_contact.user.github_login, |contact| {
190 &contact.user.github_login
191 }) {
192 Ok(ix) => this.contacts[ix] = updated_contact,
193 Err(ix) => this.contacts.insert(ix, updated_contact),
194 }
195 }
196
197 // Remove incoming contact requests
198 this.incoming_contact_requests
199 .retain(|user| !removed_incoming_requests.contains(&user.id));
200 // Update existing incoming requests and insert new ones
201 for request in incoming_requests {
202 match this
203 .incoming_contact_requests
204 .binary_search_by_key(&&request.github_login, |contact| {
205 &contact.github_login
206 }) {
207 Ok(ix) => this.incoming_contact_requests[ix] = request,
208 Err(ix) => this.incoming_contact_requests.insert(ix, request),
209 }
210 }
211
212 // Remove outgoing contact requests
213 this.outgoing_contact_requests
214 .retain(|user| !removed_outgoing_requests.contains(&user.id));
215 // Update existing incoming requests and insert new ones
216 for request in outgoing_requests {
217 match this
218 .outgoing_contact_requests
219 .binary_search_by_key(&&request.github_login, |contact| {
220 &contact.github_login
221 }) {
222 Ok(ix) => this.outgoing_contact_requests[ix] = request,
223 Err(ix) => this.outgoing_contact_requests.insert(ix, request),
224 }
225 }
226
227 cx.notify();
228 });
229
230 Ok(())
231 })
232 }
233
234 pub fn contacts(&self) -> &[Arc<Contact>] {
235 &self.contacts
236 }
237
238 pub fn has_contact(&self, user: &Arc<User>) -> bool {
239 self.contacts
240 .binary_search_by_key(&&user.github_login, |contact| &contact.user.github_login)
241 .is_ok()
242 }
243
244 pub fn incoming_contact_requests(&self) -> &[Arc<User>] {
245 &self.incoming_contact_requests
246 }
247
248 pub fn outgoing_contact_requests(&self) -> &[Arc<User>] {
249 &self.outgoing_contact_requests
250 }
251
252 pub fn contact_request_status(&self, user: &User) -> ContactRequestStatus {
253 if self.pending_contact_requests.contains_key(&user.id) {
254 ContactRequestStatus::Pending
255 } else if self
256 .contacts
257 .binary_search_by_key(&&user.id, |contact| &contact.user.id)
258 .is_ok()
259 {
260 ContactRequestStatus::RequestAccepted
261 } else if self
262 .outgoing_contact_requests
263 .binary_search_by_key(&&user.id, |user| &user.id)
264 .is_ok()
265 {
266 ContactRequestStatus::RequestSent
267 } else if self
268 .incoming_contact_requests
269 .binary_search_by_key(&&user.id, |user| &user.id)
270 .is_ok()
271 {
272 ContactRequestStatus::RequestReceived
273 } else {
274 ContactRequestStatus::None
275 }
276 }
277
278 pub fn request_contact(
279 &mut self,
280 responder_id: u64,
281 cx: &mut ModelContext<Self>,
282 ) -> Task<Result<()>> {
283 self.perform_contact_request(responder_id, proto::RequestContact { responder_id }, cx)
284 }
285
286 pub fn remove_contact(
287 &mut self,
288 user_id: u64,
289 cx: &mut ModelContext<Self>,
290 ) -> Task<Result<()>> {
291 self.perform_contact_request(user_id, proto::RemoveContact { user_id }, cx)
292 }
293
294 pub fn respond_to_contact_request(
295 &mut self,
296 requester_id: u64,
297 accept: bool,
298 cx: &mut ModelContext<Self>,
299 ) -> Task<Result<()>> {
300 self.perform_contact_request(
301 requester_id,
302 proto::RespondToContactRequest {
303 requester_id,
304 response: if accept {
305 proto::ContactRequestResponse::Accept
306 } else {
307 proto::ContactRequestResponse::Reject
308 } as i32,
309 },
310 cx,
311 )
312 }
313
314 fn perform_contact_request<T: RequestMessage>(
315 &mut self,
316 user_id: u64,
317 request: T,
318 cx: &mut ModelContext<Self>,
319 ) -> Task<Result<()>> {
320 let client = self.client.upgrade();
321 *self.pending_contact_requests.entry(user_id).or_insert(0) += 1;
322 cx.notify();
323
324 cx.spawn(|this, mut cx| async move {
325 let response = client
326 .ok_or_else(|| anyhow!("can't upgrade client reference"))?
327 .request(request)
328 .await;
329 this.update(&mut cx, |this, cx| {
330 if let Entry::Occupied(mut request_count) =
331 this.pending_contact_requests.entry(user_id)
332 {
333 *request_count.get_mut() -= 1;
334 if *request_count.get() == 0 {
335 request_count.remove();
336 }
337 }
338 cx.notify();
339 });
340 response?;
341 Ok(())
342 })
343 }
344
345 #[cfg(any(test, feature = "test-support"))]
346 pub fn clear_contacts(&mut self) {
347 self.contacts.clear();
348 self.incoming_contact_requests.clear();
349 self.outgoing_contact_requests.clear();
350 }
351
352 pub fn get_users(
353 &mut self,
354 mut user_ids: Vec<u64>,
355 cx: &mut ModelContext<Self>,
356 ) -> Task<Result<Vec<Arc<User>>>> {
357 user_ids.retain(|id| !self.users.contains_key(id));
358 self.load_users(proto::GetUsers { user_ids }, cx)
359 }
360
361 pub fn fuzzy_search_users(
362 &mut self,
363 query: String,
364 cx: &mut ModelContext<Self>,
365 ) -> Task<Result<Vec<Arc<User>>>> {
366 self.load_users(proto::FuzzySearchUsers { query }, cx)
367 }
368
369 pub fn fetch_user(
370 &mut self,
371 user_id: u64,
372 cx: &mut ModelContext<Self>,
373 ) -> Task<Result<Arc<User>>> {
374 if let Some(user) = self.users.get(&user_id).cloned() {
375 return cx.foreground().spawn(async move { Ok(user) });
376 }
377
378 let load_users = self.get_users(vec![user_id], cx);
379 cx.spawn(|this, mut cx| async move {
380 load_users.await?;
381 this.update(&mut cx, |this, _| {
382 this.users
383 .get(&user_id)
384 .cloned()
385 .ok_or_else(|| anyhow!("server responded with no users"))
386 })
387 })
388 }
389
390 pub fn current_user(&self) -> Option<Arc<User>> {
391 self.current_user.borrow().clone()
392 }
393
394 pub fn watch_current_user(&self) -> watch::Receiver<Option<Arc<User>>> {
395 self.current_user.clone()
396 }
397
398 fn load_users(
399 &mut self,
400 request: impl RequestMessage<Response = UsersResponse>,
401 cx: &mut ModelContext<Self>,
402 ) -> Task<Result<Vec<Arc<User>>>> {
403 let client = self.client.clone();
404 let http = self.http.clone();
405 cx.spawn_weak(|this, mut cx| async move {
406 if let Some(rpc) = client.upgrade() {
407 let response = rpc.request(request).await.context("error loading users")?;
408 let users = future::join_all(
409 response
410 .users
411 .into_iter()
412 .map(|user| User::new(user, http.as_ref())),
413 )
414 .await;
415
416 if let Some(this) = this.upgrade(&cx) {
417 this.update(&mut cx, |this, _| {
418 for user in &users {
419 this.users.insert(user.id, user.clone());
420 }
421 });
422 }
423 Ok(users)
424 } else {
425 Ok(Vec::new())
426 }
427 })
428 }
429}
430
431impl User {
432 async fn new(message: proto::User, http: &dyn HttpClient) -> Arc<Self> {
433 Arc::new(User {
434 id: message.id,
435 github_login: message.github_login,
436 avatar: fetch_avatar(http, &message.avatar_url).warn_on_err().await,
437 })
438 }
439}
440
441impl Contact {
442 async fn from_proto(
443 contact: proto::Contact,
444 user_store: &ModelHandle<UserStore>,
445 cx: &mut AsyncAppContext,
446 ) -> Result<Self> {
447 let user = user_store
448 .update(cx, |user_store, cx| {
449 user_store.fetch_user(contact.user_id, cx)
450 })
451 .await?;
452 let mut projects = Vec::new();
453 for project in contact.projects {
454 let mut guests = Vec::new();
455 for participant_id in project.guests {
456 guests.push(
457 user_store
458 .update(cx, |user_store, cx| {
459 user_store.fetch_user(participant_id, cx)
460 })
461 .await?,
462 );
463 }
464 projects.push(ProjectMetadata {
465 id: project.id,
466 worktree_root_names: project.worktree_root_names.clone(),
467 is_shared: project.is_shared,
468 guests,
469 });
470 }
471 Ok(Self { user, projects })
472 }
473}
474
475async fn fetch_avatar(http: &dyn HttpClient, url: &str) -> Result<Arc<ImageData>> {
476 let mut response = http
477 .get(url, Default::default(), true)
478 .await
479 .map_err(|e| anyhow!("failed to send user avatar request: {}", e))?;
480
481 if !response.status().is_success() {
482 return Err(anyhow!("avatar request failed {:?}", response.status()));
483 }
484
485 let mut body = Vec::new();
486 response
487 .body_mut()
488 .read_to_end(&mut body)
489 .await
490 .map_err(|e| anyhow!("failed to read user avatar response body: {}", e))?;
491 let format = image::guess_format(&body)?;
492 let image = image::load_from_memory_with_format(&body, format)?.into_bgra8();
493 Ok(ImageData::new(image))
494}