use super::{http::HttpClient, proto, Client, Status, TypedEnvelope}; use anyhow::{anyhow, Result}; use futures::{future, AsyncReadExt, Future}; use gpui::{AsyncAppContext, Entity, ImageData, ModelContext, ModelHandle, Task}; use postage::{prelude::Stream, sink::Sink, watch}; use rpc::proto::{RequestMessage, UsersResponse}; use std::{ collections::{HashMap, HashSet}, sync::{Arc, Weak}, }; use util::TryFutureExt as _; #[derive(Debug)] pub struct User { pub id: u64, pub github_login: String, pub avatar: Option>, } #[derive(Debug)] pub struct Contact { pub user: Arc, pub projects: Vec, } #[derive(Debug)] pub struct ProjectMetadata { pub id: u64, pub is_shared: bool, pub worktree_root_names: Vec, pub guests: Vec>, } pub struct UserStore { users: HashMap>, update_contacts_tx: watch::Sender>, current_user: watch::Receiver>>, contacts: Vec>, incoming_contact_requests: Vec>, outgoing_contact_requests: Vec>, client: Weak, http: Arc, _maintain_contacts: Task<()>, _maintain_current_user: Task<()>, } pub enum Event {} impl Entity for UserStore { type Event = Event; } impl UserStore { pub fn new( client: Arc, http: Arc, cx: &mut ModelContext, ) -> Self { let (mut current_user_tx, current_user_rx) = watch::channel(); let (update_contacts_tx, mut update_contacts_rx) = watch::channel::>(); let rpc_subscription = client.add_message_handler(cx.handle(), Self::handle_update_contacts); Self { users: Default::default(), current_user: current_user_rx, contacts: Default::default(), incoming_contact_requests: Default::default(), outgoing_contact_requests: Default::default(), client: Arc::downgrade(&client), update_contacts_tx, http, _maintain_contacts: cx.spawn_weak(|this, mut cx| async move { let _subscription = rpc_subscription; while let Some(message) = update_contacts_rx.recv().await { if let Some((message, this)) = message.zip(this.upgrade(&cx)) { this.update(&mut cx, |this, cx| this.update_contacts(message, cx)) .log_err() .await; } } }), _maintain_current_user: cx.spawn_weak(|this, mut cx| async move { let mut status = client.status(); while let Some(status) = status.recv().await { match status { Status::Connected { .. } => { if let Some((this, user_id)) = this.upgrade(&cx).zip(client.user_id()) { let user = this .update(&mut cx, |this, cx| this.fetch_user(user_id, cx)) .log_err() .await; current_user_tx.send(user).await.ok(); } } Status::SignedOut => { current_user_tx.send(None).await.ok(); } _ => {} } } }), } } async fn handle_update_contacts( this: ModelHandle, msg: TypedEnvelope, _: Arc, mut cx: AsyncAppContext, ) -> Result<()> { this.update(&mut cx, |this, _| { *this.update_contacts_tx.borrow_mut() = Some(msg.payload); }); Ok(()) } fn update_contacts( &mut self, message: proto::UpdateContacts, cx: &mut ModelContext, ) -> Task> { let mut user_ids = HashSet::new(); for contact in &message.contacts { user_ids.insert(contact.user_id); user_ids.extend(contact.projects.iter().flat_map(|w| &w.guests).copied()); } user_ids.extend(message.incoming_requests.iter().map(|req| req.user_id)); user_ids.extend(message.outgoing_requests.iter()); let load_users = self.get_users(user_ids.into_iter().collect(), cx); cx.spawn(|this, mut cx| async move { load_users.await?; // Users are fetched in parallel above and cached in call to get_users // No need to paralellize here let mut updated_contacts = Vec::new(); for contact in message.contacts { updated_contacts.push(Arc::new( Contact::from_proto(contact, &this, &mut cx).await?, )); } let mut incoming_requests = Vec::new(); for request in message.incoming_requests { incoming_requests.push( this.update(&mut cx, |this, cx| this.fetch_user(request.user_id, cx)) .await?, ); } let mut outgoing_requests = Vec::new(); for requested_user_id in message.outgoing_requests { outgoing_requests.push( this.update(&mut cx, |this, cx| this.fetch_user(requested_user_id, cx)) .await?, ); } let removed_contacts = HashSet::::from_iter(message.remove_contacts.iter().copied()); let removed_incoming_requests = HashSet::::from_iter(message.remove_incoming_requests.iter().copied()); let removed_outgoing_requests = HashSet::::from_iter(message.remove_outgoing_requests.iter().copied()); this.update(&mut cx, |this, cx| { this.contacts .retain(|contact| !removed_contacts.contains(&contact.user.id)); this.contacts.extend(updated_contacts); this.contacts .sort_by(|a, b| a.user.github_login.cmp(&b.user.github_login)); cx.notify(); this.incoming_contact_requests .retain(|user| !removed_incoming_requests.contains(&user.id)); this.incoming_contact_requests.extend(incoming_requests); this.incoming_contact_requests .sort_by(|a, b| a.github_login.cmp(&b.github_login)); this.outgoing_contact_requests .retain(|user| !removed_outgoing_requests.contains(&user.id)); this.outgoing_contact_requests.extend(outgoing_requests); this.outgoing_contact_requests .sort_by(|a, b| a.github_login.cmp(&b.github_login)); }); Ok(()) }) } pub fn contacts(&self) -> &[Arc] { &self.contacts } pub fn has_contact(&self, user: &Arc) -> bool { self.contacts .binary_search_by_key(&&user.github_login, |contact| &contact.user.github_login) .is_ok() } pub fn request_contact(&self, to_user_id: u64) -> impl Future> { let client = self.client.upgrade(); async move { client .ok_or_else(|| anyhow!("not logged in"))? .request(proto::RequestContact { to_user_id }) .await?; Ok(()) } } pub fn respond_to_contact_request( &self, from_user_id: u64, accept: bool, ) -> impl Future> { let client = self.client.upgrade(); async move { client .ok_or_else(|| anyhow!("not logged in"))? .request(proto::RespondToContactRequest { requesting_user_id: from_user_id, response: if accept { proto::ContactRequestResponse::Accept } else { proto::ContactRequestResponse::Reject } as i32, }) .await?; Ok(()) } } pub fn get_users( &mut self, mut user_ids: Vec, cx: &mut ModelContext, ) -> Task>>> { user_ids.retain(|id| !self.users.contains_key(id)); self.load_users(proto::GetUsers { user_ids }, cx) } pub fn fuzzy_search_users( &mut self, query: String, cx: &mut ModelContext, ) -> Task>>> { self.load_users(proto::FuzzySearchUsers { query }, cx) } pub fn fetch_user( &mut self, user_id: u64, cx: &mut ModelContext, ) -> Task>> { if let Some(user) = self.users.get(&user_id).cloned() { return cx.foreground().spawn(async move { Ok(user) }); } let load_users = self.get_users(vec![user_id], cx); cx.spawn(|this, mut cx| async move { load_users.await?; this.update(&mut cx, |this, _| { this.users .get(&user_id) .cloned() .ok_or_else(|| anyhow!("server responded with no users")) }) }) } pub fn current_user(&self) -> Option> { self.current_user.borrow().clone() } pub fn watch_current_user(&self) -> watch::Receiver>> { self.current_user.clone() } fn load_users( &mut self, request: impl RequestMessage, cx: &mut ModelContext, ) -> Task>>> { let client = self.client.clone(); let http = self.http.clone(); cx.spawn_weak(|this, mut cx| async move { if let Some(rpc) = client.upgrade() { let response = rpc.request(request).await?; let users = future::join_all( response .users .into_iter() .map(|user| User::new(user, http.as_ref())), ) .await; if let Some(this) = this.upgrade(&cx) { this.update(&mut cx, |this, _| { for user in &users { this.users.insert(user.id, user.clone()); } }); } Ok(users) } else { Ok(Vec::new()) } }) } } impl User { async fn new(message: proto::User, http: &dyn HttpClient) -> Arc { Arc::new(User { id: message.id, github_login: message.github_login, avatar: fetch_avatar(http, &message.avatar_url).warn_on_err().await, }) } } impl Contact { async fn from_proto( contact: proto::Contact, user_store: &ModelHandle, cx: &mut AsyncAppContext, ) -> Result { let user = user_store .update(cx, |user_store, cx| { user_store.fetch_user(contact.user_id, cx) }) .await?; let mut projects = Vec::new(); for project in contact.projects { let mut guests = Vec::new(); for participant_id in project.guests { guests.push( user_store .update(cx, |user_store, cx| { user_store.fetch_user(participant_id, cx) }) .await?, ); } projects.push(ProjectMetadata { id: project.id, worktree_root_names: project.worktree_root_names.clone(), is_shared: project.is_shared, guests, }); } Ok(Self { user, projects }) } } async fn fetch_avatar(http: &dyn HttpClient, url: &str) -> Result> { let mut response = http .get(url, Default::default(), true) .await .map_err(|e| anyhow!("failed to send user avatar request: {}", e))?; if !response.status().is_success() { return Err(anyhow!("avatar request failed {:?}", response.status())); } let mut body = Vec::new(); response .body_mut() .read_to_end(&mut body) .await .map_err(|e| anyhow!("failed to read user avatar response body: {}", e))?; let format = image::guess_format(&body)?; let image = image::load_from_memory_with_format(&body, format)?.into_bgra8(); Ok(ImageData::new(image)) }