1use crate::db::{self, ProjectId, UserId};
2use anyhow::{anyhow, Result};
3use collections::{btree_map, BTreeMap, BTreeSet, HashMap, HashSet};
4use rpc::{proto, ConnectionId};
5use serde::Serialize;
6use std::path::PathBuf;
7use tracing::instrument;
8
9pub type RoomId = u64;
10
11#[derive(Default, Serialize)]
12pub struct Store {
13 connections: BTreeMap<ConnectionId, ConnectionState>,
14 connected_users: BTreeMap<UserId, ConnectedUser>,
15 next_room_id: RoomId,
16 rooms: BTreeMap<RoomId, proto::Room>,
17 projects: BTreeMap<ProjectId, Project>,
18}
19
20#[derive(Default, Serialize)]
21struct ConnectedUser {
22 connection_ids: HashSet<ConnectionId>,
23 active_call: Option<Call>,
24}
25
26#[derive(Serialize)]
27struct ConnectionState {
28 user_id: UserId,
29 admin: bool,
30 projects: BTreeSet<ProjectId>,
31}
32
33#[derive(Copy, Clone, Eq, PartialEq, Serialize)]
34pub struct Call {
35 pub calling_user_id: UserId,
36 pub room_id: RoomId,
37 pub connection_id: Option<ConnectionId>,
38 pub initial_project_id: Option<ProjectId>,
39}
40
41#[derive(Serialize)]
42pub struct Project {
43 pub id: ProjectId,
44 pub room_id: RoomId,
45 pub host_connection_id: ConnectionId,
46 pub host: Collaborator,
47 pub guests: HashMap<ConnectionId, Collaborator>,
48 pub active_replica_ids: HashSet<ReplicaId>,
49 pub worktrees: BTreeMap<u64, Worktree>,
50 pub language_servers: Vec<proto::LanguageServer>,
51}
52
53#[derive(Serialize)]
54pub struct Collaborator {
55 pub replica_id: ReplicaId,
56 pub user_id: UserId,
57 pub admin: bool,
58}
59
60#[derive(Default, Serialize)]
61pub struct Worktree {
62 pub abs_path: PathBuf,
63 pub root_name: String,
64 pub visible: bool,
65 #[serde(skip)]
66 pub entries: BTreeMap<u64, proto::Entry>,
67 #[serde(skip)]
68 pub diagnostic_summaries: BTreeMap<PathBuf, proto::DiagnosticSummary>,
69 pub scan_id: u64,
70 pub is_complete: bool,
71}
72
73pub type ReplicaId = u16;
74
75pub struct LeftProject {
76 pub id: ProjectId,
77 pub host_user_id: UserId,
78 pub host_connection_id: ConnectionId,
79 pub connection_ids: Vec<ConnectionId>,
80 pub remove_collaborator: bool,
81}
82
83#[derive(Copy, Clone)]
84pub struct Metrics {
85 pub connections: usize,
86 pub shared_projects: usize,
87}
88
89impl Store {
90 pub fn metrics(&self) -> Metrics {
91 let connections = self.connections.values().filter(|c| !c.admin).count();
92 let mut shared_projects = 0;
93 for project in self.projects.values() {
94 if let Some(connection) = self.connections.get(&project.host_connection_id) {
95 if !connection.admin {
96 shared_projects += 1;
97 }
98 }
99 }
100
101 Metrics {
102 connections,
103 shared_projects,
104 }
105 }
106
107 #[instrument(skip(self))]
108 pub fn add_connection(&mut self, connection_id: ConnectionId, user_id: UserId, admin: bool) {
109 self.connections.insert(
110 connection_id,
111 ConnectionState {
112 user_id,
113 admin,
114 projects: Default::default(),
115 },
116 );
117 let connected_user = self.connected_users.entry(user_id).or_default();
118 connected_user.connection_ids.insert(connection_id);
119 }
120
121 #[instrument(skip(self))]
122 pub fn remove_connection(&mut self, connection_id: ConnectionId) -> Result<()> {
123 let connection = self
124 .connections
125 .get_mut(&connection_id)
126 .ok_or_else(|| anyhow!("no such connection"))?;
127
128 let user_id = connection.user_id;
129 let connected_user = self.connected_users.get_mut(&user_id).unwrap();
130 connected_user.connection_ids.remove(&connection_id);
131 if connected_user.connection_ids.is_empty() {
132 self.connected_users.remove(&user_id);
133 }
134 self.connections.remove(&connection_id).unwrap();
135 Ok(())
136 }
137
138 pub fn connection_ids_for_user(
139 &self,
140 user_id: UserId,
141 ) -> impl Iterator<Item = ConnectionId> + '_ {
142 self.connected_users
143 .get(&user_id)
144 .into_iter()
145 .map(|state| &state.connection_ids)
146 .flatten()
147 .copied()
148 }
149
150 pub fn is_user_online(&self, user_id: UserId) -> bool {
151 !self
152 .connected_users
153 .get(&user_id)
154 .unwrap_or(&Default::default())
155 .connection_ids
156 .is_empty()
157 }
158
159 pub fn build_initial_contacts_update(
160 &self,
161 contacts: Vec<db::Contact>,
162 ) -> proto::UpdateContacts {
163 let mut update = proto::UpdateContacts::default();
164
165 for contact in contacts {
166 match contact {
167 db::Contact::Accepted {
168 user_id,
169 should_notify,
170 busy,
171 } => {
172 update
173 .contacts
174 .push(self.contact_for_user(user_id, should_notify, busy));
175 }
176 db::Contact::Outgoing { user_id } => {
177 update.outgoing_requests.push(user_id.to_proto())
178 }
179 db::Contact::Incoming {
180 user_id,
181 should_notify,
182 } => update
183 .incoming_requests
184 .push(proto::IncomingContactRequest {
185 requester_id: user_id.to_proto(),
186 should_notify,
187 }),
188 }
189 }
190
191 update
192 }
193
194 pub fn contact_for_user(
195 &self,
196 user_id: UserId,
197 should_notify: bool,
198 busy: bool,
199 ) -> proto::Contact {
200 proto::Contact {
201 user_id: user_id.to_proto(),
202 online: self.is_user_online(user_id),
203 busy,
204 should_notify,
205 }
206 }
207
208 pub fn rooms(&self) -> &BTreeMap<RoomId, proto::Room> {
209 &self.rooms
210 }
211
212 pub fn unshare_project(
213 &mut self,
214 project_id: ProjectId,
215 connection_id: ConnectionId,
216 ) -> Result<(&proto::Room, Project)> {
217 match self.projects.entry(project_id) {
218 btree_map::Entry::Occupied(e) => {
219 if e.get().host_connection_id == connection_id {
220 let project = e.remove();
221
222 if let Some(host_connection) = self.connections.get_mut(&connection_id) {
223 host_connection.projects.remove(&project_id);
224 }
225
226 for guest_connection in project.guests.keys() {
227 if let Some(connection) = self.connections.get_mut(guest_connection) {
228 connection.projects.remove(&project_id);
229 }
230 }
231
232 let room = self
233 .rooms
234 .get_mut(&project.room_id)
235 .ok_or_else(|| anyhow!("no such room"))?;
236 let participant = room
237 .participants
238 .iter_mut()
239 .find(|participant| participant.peer_id == connection_id.0)
240 .ok_or_else(|| anyhow!("no such room"))?;
241 participant
242 .projects
243 .retain(|project| project.id != project_id.to_proto());
244
245 Ok((room, project))
246 } else {
247 Err(anyhow!("no such project"))?
248 }
249 }
250 btree_map::Entry::Vacant(_) => Err(anyhow!("no such project"))?,
251 }
252 }
253
254 pub fn leave_project(
255 &mut self,
256 project_id: ProjectId,
257 connection_id: ConnectionId,
258 ) -> Result<LeftProject> {
259 let project = self
260 .projects
261 .get_mut(&project_id)
262 .ok_or_else(|| anyhow!("no such project"))?;
263
264 // If the connection leaving the project is a collaborator, remove it.
265 let remove_collaborator = if let Some(guest) = project.guests.remove(&connection_id) {
266 project.active_replica_ids.remove(&guest.replica_id);
267 true
268 } else {
269 false
270 };
271
272 if let Some(connection) = self.connections.get_mut(&connection_id) {
273 connection.projects.remove(&project_id);
274 }
275
276 Ok(LeftProject {
277 id: project.id,
278 host_connection_id: project.host_connection_id,
279 host_user_id: project.host.user_id,
280 connection_ids: project.connection_ids(),
281 remove_collaborator,
282 })
283 }
284
285 #[cfg(test)]
286 pub fn check_invariants(&self) {
287 for (connection_id, connection) in &self.connections {
288 for project_id in &connection.projects {
289 let project = &self.projects.get(project_id).unwrap();
290 if project.host_connection_id != *connection_id {
291 assert!(project.guests.contains_key(connection_id));
292 }
293
294 for (worktree_id, worktree) in project.worktrees.iter() {
295 let mut paths = HashMap::default();
296 for entry in worktree.entries.values() {
297 let prev_entry = paths.insert(&entry.path, entry);
298 assert_eq!(
299 prev_entry,
300 None,
301 "worktree {:?}, duplicate path for entries {:?} and {:?}",
302 worktree_id,
303 prev_entry.unwrap(),
304 entry
305 );
306 }
307 }
308 }
309
310 assert!(self
311 .connected_users
312 .get(&connection.user_id)
313 .unwrap()
314 .connection_ids
315 .contains(connection_id));
316 }
317
318 for (user_id, state) in &self.connected_users {
319 for connection_id in &state.connection_ids {
320 assert_eq!(
321 self.connections.get(connection_id).unwrap().user_id,
322 *user_id
323 );
324 }
325
326 if let Some(active_call) = state.active_call.as_ref() {
327 if let Some(active_call_connection_id) = active_call.connection_id {
328 assert!(
329 state.connection_ids.contains(&active_call_connection_id),
330 "call is active on a dead connection"
331 );
332 assert!(
333 state.connection_ids.contains(&active_call_connection_id),
334 "call is active on a dead connection"
335 );
336 }
337 }
338 }
339
340 for (room_id, room) in &self.rooms {
341 // for pending_user_id in &room.pending_participant_user_ids {
342 // assert!(
343 // self.connected_users
344 // .contains_key(&UserId::from_proto(*pending_user_id)),
345 // "call is active on a user that has disconnected"
346 // );
347 // }
348
349 for participant in &room.participants {
350 assert!(
351 self.connections
352 .contains_key(&ConnectionId(participant.peer_id)),
353 "room {} contains participant {:?} that has disconnected",
354 room_id,
355 participant
356 );
357
358 for participant_project in &participant.projects {
359 let project = &self.projects[&ProjectId::from_proto(participant_project.id)];
360 assert_eq!(
361 project.room_id, *room_id,
362 "project was shared on a different room"
363 );
364 }
365 }
366
367 // assert!(
368 // !room.pending_participant_user_ids.is_empty() || !room.participants.is_empty(),
369 // "room can't be empty"
370 // );
371 }
372
373 for (project_id, project) in &self.projects {
374 let host_connection = self.connections.get(&project.host_connection_id).unwrap();
375 assert!(host_connection.projects.contains(project_id));
376
377 for guest_connection_id in project.guests.keys() {
378 let guest_connection = self.connections.get(guest_connection_id).unwrap();
379 assert!(guest_connection.projects.contains(project_id));
380 }
381 assert_eq!(project.active_replica_ids.len(), project.guests.len());
382 assert_eq!(
383 project.active_replica_ids,
384 project
385 .guests
386 .values()
387 .map(|guest| guest.replica_id)
388 .collect::<HashSet<_>>(),
389 );
390
391 let room = &self.rooms[&project.room_id];
392 let room_participant = room
393 .participants
394 .iter()
395 .find(|participant| participant.peer_id == project.host_connection_id.0)
396 .unwrap();
397 assert!(
398 room_participant
399 .projects
400 .iter()
401 .any(|project| project.id == project_id.to_proto()),
402 "project was not shared in room"
403 );
404 }
405 }
406}
407
408impl Project {
409 pub fn guest_connection_ids(&self) -> Vec<ConnectionId> {
410 self.guests.keys().copied().collect()
411 }
412
413 pub fn connection_ids(&self) -> Vec<ConnectionId> {
414 self.guests
415 .keys()
416 .copied()
417 .chain(Some(self.host_connection_id))
418 .collect()
419 }
420}