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