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