1use crate::db::{self, ChannelId, UserId};
2use anyhow::{anyhow, Result};
3use collections::{hash_map::Entry, BTreeMap, HashMap, HashSet};
4use rpc::{proto, ConnectionId, Receipt};
5use serde::Serialize;
6use std::{collections::hash_map, mem, path::PathBuf};
7use tracing::instrument;
8
9#[derive(Default, Serialize)]
10pub struct Store {
11 connections: HashMap<ConnectionId, ConnectionState>,
12 connections_by_user_id: HashMap<UserId, HashSet<ConnectionId>>,
13 projects: HashMap<u64, Project>,
14 #[serde(skip)]
15 channels: HashMap<ChannelId, Channel>,
16 next_project_id: u64,
17}
18
19#[derive(Serialize)]
20struct ConnectionState {
21 user_id: UserId,
22 projects: HashSet<u64>,
23 requested_projects: HashSet<u64>,
24 channels: HashSet<ChannelId>,
25}
26
27#[derive(Serialize)]
28pub struct Project {
29 pub host_connection_id: ConnectionId,
30 pub host_user_id: UserId,
31 pub guests: HashMap<ConnectionId, (ReplicaId, UserId)>,
32 #[serde(skip)]
33 pub join_requests: HashMap<UserId, Vec<Receipt<proto::JoinProject>>>,
34 pub active_replica_ids: HashSet<ReplicaId>,
35 pub worktrees: BTreeMap<u64, Worktree>,
36 pub language_servers: Vec<proto::LanguageServer>,
37}
38
39#[derive(Default, Serialize)]
40pub struct Worktree {
41 pub root_name: String,
42 pub visible: bool,
43 #[serde(skip)]
44 pub entries: HashMap<u64, proto::Entry>,
45 #[serde(skip)]
46 pub diagnostic_summaries: BTreeMap<PathBuf, proto::DiagnosticSummary>,
47 pub scan_id: u64,
48}
49
50#[derive(Default)]
51pub struct Channel {
52 pub connection_ids: HashSet<ConnectionId>,
53}
54
55pub type ReplicaId = u16;
56
57#[derive(Default)]
58pub struct RemovedConnectionState {
59 pub user_id: UserId,
60 pub hosted_projects: HashMap<u64, Project>,
61 pub guest_project_ids: HashSet<u64>,
62 pub contact_ids: HashSet<UserId>,
63}
64
65pub struct LeftProject {
66 pub host_user_id: UserId,
67 pub host_connection_id: ConnectionId,
68 pub connection_ids: Vec<ConnectionId>,
69 pub remove_collaborator: bool,
70 pub cancel_request: Option<UserId>,
71 pub unshare: bool,
72}
73
74#[derive(Copy, Clone)]
75pub struct Metrics {
76 pub connections: usize,
77 pub registered_projects: usize,
78 pub shared_projects: usize,
79}
80
81impl Store {
82 pub fn metrics(&self) -> Metrics {
83 let connections = self.connections.len();
84 let mut registered_projects = 0;
85 let mut shared_projects = 0;
86 for project in self.projects.values() {
87 registered_projects += 1;
88 if !project.guests.is_empty() {
89 shared_projects += 1;
90 }
91 }
92
93 Metrics {
94 connections,
95 registered_projects,
96 shared_projects,
97 }
98 }
99
100 #[instrument(skip(self))]
101 pub fn add_connection(&mut self, connection_id: ConnectionId, user_id: UserId) {
102 self.connections.insert(
103 connection_id,
104 ConnectionState {
105 user_id,
106 projects: Default::default(),
107 requested_projects: Default::default(),
108 channels: Default::default(),
109 },
110 );
111 self.connections_by_user_id
112 .entry(user_id)
113 .or_default()
114 .insert(connection_id);
115 }
116
117 #[instrument(skip(self))]
118 pub fn remove_connection(
119 &mut self,
120 connection_id: ConnectionId,
121 ) -> Result<RemovedConnectionState> {
122 let connection = self
123 .connections
124 .get_mut(&connection_id)
125 .ok_or_else(|| anyhow!("no such connection"))?;
126
127 let user_id = connection.user_id;
128 let connection_projects = mem::take(&mut connection.projects);
129 let connection_channels = mem::take(&mut connection.channels);
130
131 let mut result = RemovedConnectionState::default();
132 result.user_id = user_id;
133
134 // Leave all channels.
135 for channel_id in connection_channels {
136 self.leave_channel(connection_id, channel_id);
137 }
138
139 // Unregister and leave all projects.
140 for project_id in connection_projects {
141 if let Ok(project) = self.unregister_project(project_id, connection_id) {
142 result.hosted_projects.insert(project_id, project);
143 } else if self.leave_project(connection_id, project_id).is_ok() {
144 result.guest_project_ids.insert(project_id);
145 }
146 }
147
148 let user_connections = self.connections_by_user_id.get_mut(&user_id).unwrap();
149 user_connections.remove(&connection_id);
150 if user_connections.is_empty() {
151 self.connections_by_user_id.remove(&user_id);
152 }
153
154 self.connections.remove(&connection_id).unwrap();
155
156 Ok(result)
157 }
158
159 #[cfg(test)]
160 pub fn channel(&self, id: ChannelId) -> Option<&Channel> {
161 self.channels.get(&id)
162 }
163
164 pub fn join_channel(&mut self, connection_id: ConnectionId, channel_id: ChannelId) {
165 if let Some(connection) = self.connections.get_mut(&connection_id) {
166 connection.channels.insert(channel_id);
167 self.channels
168 .entry(channel_id)
169 .or_default()
170 .connection_ids
171 .insert(connection_id);
172 }
173 }
174
175 pub fn leave_channel(&mut self, connection_id: ConnectionId, channel_id: ChannelId) {
176 if let Some(connection) = self.connections.get_mut(&connection_id) {
177 connection.channels.remove(&channel_id);
178 if let hash_map::Entry::Occupied(mut entry) = self.channels.entry(channel_id) {
179 entry.get_mut().connection_ids.remove(&connection_id);
180 if entry.get_mut().connection_ids.is_empty() {
181 entry.remove();
182 }
183 }
184 }
185 }
186
187 pub fn user_id_for_connection(&self, connection_id: ConnectionId) -> Result<UserId> {
188 Ok(self
189 .connections
190 .get(&connection_id)
191 .ok_or_else(|| anyhow!("unknown connection"))?
192 .user_id)
193 }
194
195 pub fn connection_ids_for_user<'a>(
196 &'a self,
197 user_id: UserId,
198 ) -> impl 'a + Iterator<Item = ConnectionId> {
199 self.connections_by_user_id
200 .get(&user_id)
201 .into_iter()
202 .flatten()
203 .copied()
204 }
205
206 pub fn is_user_online(&self, user_id: UserId) -> bool {
207 !self
208 .connections_by_user_id
209 .get(&user_id)
210 .unwrap_or(&Default::default())
211 .is_empty()
212 }
213
214 pub fn build_initial_contacts_update(
215 &self,
216 contacts: Vec<db::Contact>,
217 ) -> proto::UpdateContacts {
218 let mut update = proto::UpdateContacts::default();
219
220 for contact in contacts {
221 match contact {
222 db::Contact::Accepted {
223 user_id,
224 should_notify,
225 } => {
226 update
227 .contacts
228 .push(self.contact_for_user(user_id, should_notify));
229 }
230 db::Contact::Outgoing { user_id } => {
231 update.outgoing_requests.push(user_id.to_proto())
232 }
233 db::Contact::Incoming {
234 user_id,
235 should_notify,
236 } => update
237 .incoming_requests
238 .push(proto::IncomingContactRequest {
239 requester_id: user_id.to_proto(),
240 should_notify,
241 }),
242 }
243 }
244
245 update
246 }
247
248 pub fn contact_for_user(&self, user_id: UserId, should_notify: bool) -> proto::Contact {
249 proto::Contact {
250 user_id: user_id.to_proto(),
251 projects: self.project_metadata_for_user(user_id),
252 online: self.is_user_online(user_id),
253 should_notify,
254 }
255 }
256
257 pub fn project_metadata_for_user(&self, user_id: UserId) -> Vec<proto::ProjectMetadata> {
258 let connection_ids = self.connections_by_user_id.get(&user_id);
259 let project_ids = connection_ids.iter().flat_map(|connection_ids| {
260 connection_ids
261 .iter()
262 .filter_map(|connection_id| self.connections.get(connection_id))
263 .flat_map(|connection| connection.projects.iter().copied())
264 });
265
266 let mut metadata = Vec::new();
267 for project_id in project_ids {
268 if let Some(project) = self.projects.get(&project_id) {
269 if project.host_user_id == user_id {
270 metadata.push(proto::ProjectMetadata {
271 id: project_id,
272 worktree_root_names: project
273 .worktrees
274 .values()
275 .map(|worktree| worktree.root_name.clone())
276 .collect(),
277 guests: project
278 .guests
279 .values()
280 .map(|(_, user_id)| user_id.to_proto())
281 .collect(),
282 });
283 }
284 }
285 }
286
287 metadata
288 }
289
290 pub fn register_project(
291 &mut self,
292 host_connection_id: ConnectionId,
293 host_user_id: UserId,
294 ) -> u64 {
295 let project_id = self.next_project_id;
296 self.projects.insert(
297 project_id,
298 Project {
299 host_connection_id,
300 host_user_id,
301 guests: Default::default(),
302 join_requests: Default::default(),
303 active_replica_ids: Default::default(),
304 worktrees: Default::default(),
305 language_servers: Default::default(),
306 },
307 );
308 if let Some(connection) = self.connections.get_mut(&host_connection_id) {
309 connection.projects.insert(project_id);
310 }
311 self.next_project_id += 1;
312 project_id
313 }
314
315 pub fn update_project(
316 &mut self,
317 project_id: u64,
318 worktrees: &[proto::WorktreeMetadata],
319 connection_id: ConnectionId,
320 ) -> Result<()> {
321 let project = self
322 .projects
323 .get_mut(&project_id)
324 .ok_or_else(|| anyhow!("no such project"))?;
325 if project.host_connection_id == connection_id {
326 let mut old_worktrees = mem::take(&mut project.worktrees);
327 for worktree in worktrees {
328 if let Some(old_worktree) = old_worktrees.remove(&worktree.id) {
329 project.worktrees.insert(worktree.id, old_worktree);
330 } else {
331 project.worktrees.insert(
332 worktree.id,
333 Worktree {
334 root_name: worktree.root_name.clone(),
335 visible: worktree.visible,
336 ..Default::default()
337 },
338 );
339 }
340 }
341 Ok(())
342 } else {
343 Err(anyhow!("no such project"))?
344 }
345 }
346
347 pub fn unregister_project(
348 &mut self,
349 project_id: u64,
350 connection_id: ConnectionId,
351 ) -> Result<Project> {
352 match self.projects.entry(project_id) {
353 hash_map::Entry::Occupied(e) => {
354 if e.get().host_connection_id == connection_id {
355 let project = e.remove();
356
357 if let Some(host_connection) = self.connections.get_mut(&connection_id) {
358 host_connection.projects.remove(&project_id);
359 }
360
361 for guest_connection in project.guests.keys() {
362 if let Some(connection) = self.connections.get_mut(&guest_connection) {
363 connection.projects.remove(&project_id);
364 }
365 }
366
367 for requester_user_id in project.join_requests.keys() {
368 if let Some(requester_connection_ids) =
369 self.connections_by_user_id.get_mut(&requester_user_id)
370 {
371 for requester_connection_id in requester_connection_ids.iter() {
372 if let Some(requester_connection) =
373 self.connections.get_mut(requester_connection_id)
374 {
375 requester_connection.requested_projects.remove(&project_id);
376 }
377 }
378 }
379 }
380
381 Ok(project)
382 } else {
383 Err(anyhow!("no such project"))?
384 }
385 }
386 hash_map::Entry::Vacant(_) => Err(anyhow!("no such project"))?,
387 }
388 }
389
390 pub fn update_diagnostic_summary(
391 &mut self,
392 project_id: u64,
393 worktree_id: u64,
394 connection_id: ConnectionId,
395 summary: proto::DiagnosticSummary,
396 ) -> Result<Vec<ConnectionId>> {
397 let project = self
398 .projects
399 .get_mut(&project_id)
400 .ok_or_else(|| anyhow!("no such project"))?;
401 if project.host_connection_id == connection_id {
402 let worktree = project
403 .worktrees
404 .get_mut(&worktree_id)
405 .ok_or_else(|| anyhow!("no such worktree"))?;
406 worktree
407 .diagnostic_summaries
408 .insert(summary.path.clone().into(), summary);
409 return Ok(project.connection_ids());
410 }
411
412 Err(anyhow!("no such worktree"))?
413 }
414
415 pub fn start_language_server(
416 &mut self,
417 project_id: u64,
418 connection_id: ConnectionId,
419 language_server: proto::LanguageServer,
420 ) -> Result<Vec<ConnectionId>> {
421 let project = self
422 .projects
423 .get_mut(&project_id)
424 .ok_or_else(|| anyhow!("no such project"))?;
425 if project.host_connection_id == connection_id {
426 project.language_servers.push(language_server);
427 return Ok(project.connection_ids());
428 }
429
430 Err(anyhow!("no such project"))?
431 }
432
433 pub fn request_join_project(
434 &mut self,
435 requester_id: UserId,
436 project_id: u64,
437 receipt: Receipt<proto::JoinProject>,
438 ) -> Result<()> {
439 let connection = self
440 .connections
441 .get_mut(&receipt.sender_id)
442 .ok_or_else(|| anyhow!("no such connection"))?;
443 let project = self
444 .projects
445 .get_mut(&project_id)
446 .ok_or_else(|| anyhow!("no such project"))?;
447 connection.requested_projects.insert(project_id);
448 project
449 .join_requests
450 .entry(requester_id)
451 .or_default()
452 .push(receipt);
453 Ok(())
454 }
455
456 pub fn deny_join_project_request(
457 &mut self,
458 responder_connection_id: ConnectionId,
459 requester_id: UserId,
460 project_id: u64,
461 ) -> Option<Vec<Receipt<proto::JoinProject>>> {
462 let project = self.projects.get_mut(&project_id)?;
463 if responder_connection_id != project.host_connection_id {
464 return None;
465 }
466
467 let receipts = project.join_requests.remove(&requester_id)?;
468 for receipt in &receipts {
469 let requester_connection = self.connections.get_mut(&receipt.sender_id)?;
470 requester_connection.requested_projects.remove(&project_id);
471 }
472 Some(receipts)
473 }
474
475 pub fn accept_join_project_request(
476 &mut self,
477 responder_connection_id: ConnectionId,
478 requester_id: UserId,
479 project_id: u64,
480 ) -> Option<(Vec<(Receipt<proto::JoinProject>, ReplicaId)>, &Project)> {
481 let project = self.projects.get_mut(&project_id)?;
482 if responder_connection_id != project.host_connection_id {
483 return None;
484 }
485
486 let receipts = project.join_requests.remove(&requester_id)?;
487 let mut receipts_with_replica_ids = Vec::new();
488 for receipt in receipts {
489 let requester_connection = self.connections.get_mut(&receipt.sender_id)?;
490 requester_connection.requested_projects.remove(&project_id);
491 requester_connection.projects.insert(project_id);
492 let mut replica_id = 1;
493 while project.active_replica_ids.contains(&replica_id) {
494 replica_id += 1;
495 }
496 project.active_replica_ids.insert(replica_id);
497 project
498 .guests
499 .insert(receipt.sender_id, (replica_id, requester_id));
500 receipts_with_replica_ids.push((receipt, replica_id));
501 }
502
503 Some((receipts_with_replica_ids, project))
504 }
505
506 pub fn leave_project(
507 &mut self,
508 connection_id: ConnectionId,
509 project_id: u64,
510 ) -> Result<LeftProject> {
511 let user_id = self.user_id_for_connection(connection_id)?;
512 let project = self
513 .projects
514 .get_mut(&project_id)
515 .ok_or_else(|| anyhow!("no such project"))?;
516
517 // If the connection leaving the project is a collaborator, remove it.
518 let remove_collaborator =
519 if let Some((replica_id, _)) = project.guests.remove(&connection_id) {
520 project.active_replica_ids.remove(&replica_id);
521 true
522 } else {
523 false
524 };
525
526 // If the connection leaving the project has a pending request, remove it.
527 // If that user has no other pending requests on other connections, indicate that the request should be cancelled.
528 let mut cancel_request = None;
529 if let Entry::Occupied(mut entry) = project.join_requests.entry(user_id) {
530 entry
531 .get_mut()
532 .retain(|receipt| receipt.sender_id != connection_id);
533 if entry.get().is_empty() {
534 entry.remove();
535 cancel_request = Some(user_id);
536 }
537 }
538
539 if let Some(connection) = self.connections.get_mut(&connection_id) {
540 connection.projects.remove(&project_id);
541 }
542
543 let connection_ids = project.connection_ids();
544 let unshare = connection_ids.len() <= 1 && project.join_requests.is_empty();
545 if unshare {
546 project.language_servers.clear();
547 for worktree in project.worktrees.values_mut() {
548 worktree.diagnostic_summaries.clear();
549 worktree.entries.clear();
550 }
551 }
552
553 Ok(LeftProject {
554 host_connection_id: project.host_connection_id,
555 host_user_id: project.host_user_id,
556 connection_ids,
557 cancel_request,
558 unshare,
559 remove_collaborator,
560 })
561 }
562
563 pub fn update_worktree(
564 &mut self,
565 connection_id: ConnectionId,
566 project_id: u64,
567 worktree_id: u64,
568 worktree_root_name: &str,
569 removed_entries: &[u64],
570 updated_entries: &[proto::Entry],
571 scan_id: u64,
572 ) -> Result<(Vec<ConnectionId>, bool)> {
573 let project = self.write_project(project_id, connection_id)?;
574 let mut worktree = project.worktrees.entry(worktree_id).or_default();
575 let metadata_changed = worktree_root_name != worktree.root_name;
576 worktree.root_name = worktree_root_name.to_string();
577 for entry_id in removed_entries {
578 worktree.entries.remove(&entry_id);
579 }
580 for entry in updated_entries {
581 worktree.entries.insert(entry.id, entry.clone());
582 }
583 worktree.scan_id = scan_id;
584 let connection_ids = project.connection_ids();
585 Ok((connection_ids, metadata_changed))
586 }
587
588 pub fn project_connection_ids(
589 &self,
590 project_id: u64,
591 acting_connection_id: ConnectionId,
592 ) -> Result<Vec<ConnectionId>> {
593 Ok(self
594 .read_project(project_id, acting_connection_id)?
595 .connection_ids())
596 }
597
598 pub fn channel_connection_ids(&self, channel_id: ChannelId) -> Result<Vec<ConnectionId>> {
599 Ok(self
600 .channels
601 .get(&channel_id)
602 .ok_or_else(|| anyhow!("no such channel"))?
603 .connection_ids())
604 }
605
606 pub fn project(&self, project_id: u64) -> Result<&Project> {
607 self.projects
608 .get(&project_id)
609 .ok_or_else(|| anyhow!("no such project"))
610 }
611
612 pub fn read_project(&self, project_id: u64, connection_id: ConnectionId) -> Result<&Project> {
613 let project = self
614 .projects
615 .get(&project_id)
616 .ok_or_else(|| anyhow!("no such project"))?;
617 if project.host_connection_id == connection_id
618 || project.guests.contains_key(&connection_id)
619 {
620 Ok(project)
621 } else {
622 Err(anyhow!("no such project"))?
623 }
624 }
625
626 fn write_project(
627 &mut self,
628 project_id: u64,
629 connection_id: ConnectionId,
630 ) -> Result<&mut Project> {
631 let project = self
632 .projects
633 .get_mut(&project_id)
634 .ok_or_else(|| anyhow!("no such project"))?;
635 if project.host_connection_id == connection_id
636 || project.guests.contains_key(&connection_id)
637 {
638 Ok(project)
639 } else {
640 Err(anyhow!("no such project"))?
641 }
642 }
643
644 #[cfg(test)]
645 pub fn check_invariants(&self) {
646 for (connection_id, connection) in &self.connections {
647 for project_id in &connection.projects {
648 let project = &self.projects.get(&project_id).unwrap();
649 if project.host_connection_id != *connection_id {
650 assert!(project.guests.contains_key(connection_id));
651 }
652
653 for (worktree_id, worktree) in project.worktrees.iter() {
654 let mut paths = HashMap::default();
655 for entry in worktree.entries.values() {
656 let prev_entry = paths.insert(&entry.path, entry);
657 assert_eq!(
658 prev_entry,
659 None,
660 "worktree {:?}, duplicate path for entries {:?} and {:?}",
661 worktree_id,
662 prev_entry.unwrap(),
663 entry
664 );
665 }
666 }
667 }
668 for channel_id in &connection.channels {
669 let channel = self.channels.get(channel_id).unwrap();
670 assert!(channel.connection_ids.contains(connection_id));
671 }
672 assert!(self
673 .connections_by_user_id
674 .get(&connection.user_id)
675 .unwrap()
676 .contains(connection_id));
677 }
678
679 for (user_id, connection_ids) in &self.connections_by_user_id {
680 for connection_id in connection_ids {
681 assert_eq!(
682 self.connections.get(connection_id).unwrap().user_id,
683 *user_id
684 );
685 }
686 }
687
688 for (project_id, project) in &self.projects {
689 let host_connection = self.connections.get(&project.host_connection_id).unwrap();
690 assert!(host_connection.projects.contains(project_id));
691
692 for guest_connection_id in project.guests.keys() {
693 let guest_connection = self.connections.get(guest_connection_id).unwrap();
694 assert!(guest_connection.projects.contains(project_id));
695 }
696 assert_eq!(project.active_replica_ids.len(), project.guests.len(),);
697 assert_eq!(
698 project.active_replica_ids,
699 project
700 .guests
701 .values()
702 .map(|(replica_id, _)| *replica_id)
703 .collect::<HashSet<_>>(),
704 );
705 }
706
707 for (channel_id, channel) in &self.channels {
708 for connection_id in &channel.connection_ids {
709 let connection = self.connections.get(connection_id).unwrap();
710 assert!(connection.channels.contains(channel_id));
711 }
712 }
713 }
714}
715
716impl Project {
717 pub fn guest_connection_ids(&self) -> Vec<ConnectionId> {
718 self.guests.keys().copied().collect()
719 }
720
721 pub fn connection_ids(&self) -> Vec<ConnectionId> {
722 self.guests
723 .keys()
724 .copied()
725 .chain(Some(self.host_connection_id))
726 .collect()
727 }
728}
729
730impl Channel {
731 fn connection_ids(&self) -> Vec<ConnectionId> {
732 self.connection_ids.iter().copied().collect()
733 }
734}