user.rs

  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 request = client
326                .ok_or_else(|| anyhow!("can't upgrade client reference"))?
327                .request(request);
328            request.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            Ok(())
341        })
342    }
343
344    #[cfg(any(test, feature = "test-support"))]
345    pub fn clear_contacts(&mut self) {
346        self.contacts.clear();
347        self.incoming_contact_requests.clear();
348        self.outgoing_contact_requests.clear();
349    }
350
351    pub fn get_users(
352        &mut self,
353        mut user_ids: Vec<u64>,
354        cx: &mut ModelContext<Self>,
355    ) -> Task<Result<Vec<Arc<User>>>> {
356        user_ids.retain(|id| !self.users.contains_key(id));
357        self.load_users(proto::GetUsers { user_ids }, cx)
358    }
359
360    pub fn fuzzy_search_users(
361        &mut self,
362        query: String,
363        cx: &mut ModelContext<Self>,
364    ) -> Task<Result<Vec<Arc<User>>>> {
365        self.load_users(proto::FuzzySearchUsers { query }, cx)
366    }
367
368    pub fn fetch_user(
369        &mut self,
370        user_id: u64,
371        cx: &mut ModelContext<Self>,
372    ) -> Task<Result<Arc<User>>> {
373        if let Some(user) = self.users.get(&user_id).cloned() {
374            return cx.foreground().spawn(async move { Ok(user) });
375        }
376
377        let load_users = self.get_users(vec![user_id], cx);
378        cx.spawn(|this, mut cx| async move {
379            load_users.await?;
380            this.update(&mut cx, |this, _| {
381                this.users
382                    .get(&user_id)
383                    .cloned()
384                    .ok_or_else(|| anyhow!("server responded with no users"))
385            })
386        })
387    }
388
389    pub fn current_user(&self) -> Option<Arc<User>> {
390        self.current_user.borrow().clone()
391    }
392
393    pub fn watch_current_user(&self) -> watch::Receiver<Option<Arc<User>>> {
394        self.current_user.clone()
395    }
396
397    fn load_users(
398        &mut self,
399        request: impl RequestMessage<Response = UsersResponse>,
400        cx: &mut ModelContext<Self>,
401    ) -> Task<Result<Vec<Arc<User>>>> {
402        let client = self.client.clone();
403        let http = self.http.clone();
404        cx.spawn_weak(|this, mut cx| async move {
405            if let Some(rpc) = client.upgrade() {
406                let response = rpc.request(request).await.context("error loading users")?;
407                let users = future::join_all(
408                    response
409                        .users
410                        .into_iter()
411                        .map(|user| User::new(user, http.as_ref())),
412                )
413                .await;
414
415                if let Some(this) = this.upgrade(&cx) {
416                    this.update(&mut cx, |this, _| {
417                        for user in &users {
418                            this.users.insert(user.id, user.clone());
419                        }
420                    });
421                }
422                Ok(users)
423            } else {
424                Ok(Vec::new())
425            }
426        })
427    }
428}
429
430impl User {
431    async fn new(message: proto::User, http: &dyn HttpClient) -> Arc<Self> {
432        Arc::new(User {
433            id: message.id,
434            github_login: message.github_login,
435            avatar: fetch_avatar(http, &message.avatar_url).warn_on_err().await,
436        })
437    }
438}
439
440impl Contact {
441    async fn from_proto(
442        contact: proto::Contact,
443        user_store: &ModelHandle<UserStore>,
444        cx: &mut AsyncAppContext,
445    ) -> Result<Self> {
446        let user = user_store
447            .update(cx, |user_store, cx| {
448                user_store.fetch_user(contact.user_id, cx)
449            })
450            .await?;
451        let mut projects = Vec::new();
452        for project in contact.projects {
453            let mut guests = Vec::new();
454            for participant_id in project.guests {
455                guests.push(
456                    user_store
457                        .update(cx, |user_store, cx| {
458                            user_store.fetch_user(participant_id, cx)
459                        })
460                        .await?,
461                );
462            }
463            projects.push(ProjectMetadata {
464                id: project.id,
465                worktree_root_names: project.worktree_root_names.clone(),
466                is_shared: project.is_shared,
467                guests,
468            });
469        }
470        Ok(Self { user, projects })
471    }
472}
473
474async fn fetch_avatar(http: &dyn HttpClient, url: &str) -> Result<Arc<ImageData>> {
475    let mut response = http
476        .get(url, Default::default(), true)
477        .await
478        .map_err(|e| anyhow!("failed to send user avatar request: {}", e))?;
479
480    if !response.status().is_success() {
481        return Err(anyhow!("avatar request failed {:?}", response.status()));
482    }
483
484    let mut body = Vec::new();
485    response
486        .body_mut()
487        .read_to_end(&mut body)
488        .await
489        .map_err(|e| anyhow!("failed to read user avatar response body: {}", e))?;
490    let format = image::guess_format(&body)?;
491    let image = image::load_from_memory_with_format(&body, format)?.into_bgra8();
492    Ok(ImageData::new(image))
493}