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