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