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 fn perform_contact_request<T: RequestMessage>(
360 &mut self,
361 user_id: u64,
362 request: T,
363 cx: &mut ModelContext<Self>,
364 ) -> Task<Result<()>> {
365 let client = self.client.upgrade();
366 *self.pending_contact_requests.entry(user_id).or_insert(0) += 1;
367 cx.notify();
368
369 cx.spawn(|this, mut cx| async move {
370 let response = client
371 .ok_or_else(|| anyhow!("can't upgrade client reference"))?
372 .request(request)
373 .await;
374 this.update(&mut cx, |this, cx| {
375 if let Entry::Occupied(mut request_count) =
376 this.pending_contact_requests.entry(user_id)
377 {
378 *request_count.get_mut() -= 1;
379 if *request_count.get() == 0 {
380 request_count.remove();
381 }
382 }
383 cx.notify();
384 });
385 response?;
386 Ok(())
387 })
388 }
389
390 pub fn clear_contacts(&mut self) -> impl Future<Output = ()> {
391 let (tx, mut rx) = postage::barrier::channel();
392 self.update_contacts_tx
393 .unbounded_send(UpdateContacts::Clear(tx))
394 .unwrap();
395 async move {
396 rx.recv().await;
397 }
398 }
399
400 pub fn get_users(
401 &mut self,
402 mut user_ids: Vec<u64>,
403 cx: &mut ModelContext<Self>,
404 ) -> Task<Result<()>> {
405 user_ids.retain(|id| !self.users.contains_key(id));
406 if user_ids.is_empty() {
407 Task::ready(Ok(()))
408 } else {
409 let load = self.load_users(proto::GetUsers { user_ids }, cx);
410 cx.foreground().spawn(async move {
411 load.await?;
412 Ok(())
413 })
414 }
415 }
416
417 pub fn fuzzy_search_users(
418 &mut self,
419 query: String,
420 cx: &mut ModelContext<Self>,
421 ) -> Task<Result<Vec<Arc<User>>>> {
422 self.load_users(proto::FuzzySearchUsers { query }, cx)
423 }
424
425 pub fn fetch_user(
426 &mut self,
427 user_id: u64,
428 cx: &mut ModelContext<Self>,
429 ) -> Task<Result<Arc<User>>> {
430 if let Some(user) = self.users.get(&user_id).cloned() {
431 return cx.foreground().spawn(async move { Ok(user) });
432 }
433
434 let load_users = self.get_users(vec![user_id], cx);
435 cx.spawn(|this, mut cx| async move {
436 load_users.await?;
437 this.update(&mut cx, |this, _| {
438 this.users
439 .get(&user_id)
440 .cloned()
441 .ok_or_else(|| anyhow!("server responded with no users"))
442 })
443 })
444 }
445
446 pub fn current_user(&self) -> Option<Arc<User>> {
447 self.current_user.borrow().clone()
448 }
449
450 pub fn watch_current_user(&self) -> watch::Receiver<Option<Arc<User>>> {
451 self.current_user.clone()
452 }
453
454 fn load_users(
455 &mut self,
456 request: impl RequestMessage<Response = UsersResponse>,
457 cx: &mut ModelContext<Self>,
458 ) -> Task<Result<Vec<Arc<User>>>> {
459 let client = self.client.clone();
460 let http = self.http.clone();
461 cx.spawn_weak(|this, mut cx| async move {
462 if let Some(rpc) = client.upgrade() {
463 let response = rpc.request(request).await.context("error loading users")?;
464 let users = future::join_all(
465 response
466 .users
467 .into_iter()
468 .map(|user| User::new(user, http.as_ref())),
469 )
470 .await;
471
472 if let Some(this) = this.upgrade(&cx) {
473 this.update(&mut cx, |this, _| {
474 for user in &users {
475 this.users.insert(user.id, user.clone());
476 }
477 });
478 }
479 Ok(users)
480 } else {
481 Ok(Vec::new())
482 }
483 })
484 }
485}
486
487impl User {
488 async fn new(message: proto::User, http: &dyn HttpClient) -> Arc<Self> {
489 Arc::new(User {
490 id: message.id,
491 github_login: message.github_login,
492 avatar: fetch_avatar(http, &message.avatar_url).warn_on_err().await,
493 })
494 }
495}
496
497impl Contact {
498 async fn from_proto(
499 contact: proto::Contact,
500 user_store: &ModelHandle<UserStore>,
501 cx: &mut AsyncAppContext,
502 ) -> Result<Self> {
503 let user = user_store
504 .update(cx, |user_store, cx| {
505 user_store.fetch_user(contact.user_id, cx)
506 })
507 .await?;
508 let mut projects = Vec::new();
509 for project in contact.projects {
510 let mut guests = Vec::new();
511 for participant_id in project.guests {
512 guests.push(
513 user_store
514 .update(cx, |user_store, cx| {
515 user_store.fetch_user(participant_id, cx)
516 })
517 .await?,
518 );
519 }
520 projects.push(ProjectMetadata {
521 id: project.id,
522 worktree_root_names: project.worktree_root_names.clone(),
523 is_shared: project.is_shared,
524 guests,
525 });
526 }
527 Ok(Self {
528 user,
529 online: contact.online,
530 projects,
531 })
532 }
533
534 pub fn non_empty_projects(&self) -> impl Iterator<Item = &ProjectMetadata> {
535 self.projects
536 .iter()
537 .filter(|project| !project.worktree_root_names.is_empty())
538 }
539}
540
541async fn fetch_avatar(http: &dyn HttpClient, url: &str) -> Result<Arc<ImageData>> {
542 let mut response = http
543 .get(url, Default::default(), true)
544 .await
545 .map_err(|e| anyhow!("failed to send user avatar request: {}", e))?;
546
547 if !response.status().is_success() {
548 return Err(anyhow!("avatar request failed {:?}", response.status()));
549 }
550
551 let mut body = Vec::new();
552 response
553 .body_mut()
554 .read_to_end(&mut body)
555 .await
556 .map_err(|e| anyhow!("failed to read user avatar response body: {}", e))?;
557 let format = image::guess_format(&body)?;
558 let image = image::load_from_memory_with_format(&body, format)?.into_bgra8();
559 Ok(ImageData::new(image))
560}