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