1#[cfg(test)]
2pub mod tests;
3
4#[cfg(test)]
5pub use tests::TestDb;
6
7mod ids;
8mod queries;
9mod tables;
10
11use crate::{executor::Executor, Error, Result};
12use anyhow::anyhow;
13use collections::{BTreeMap, HashMap, HashSet};
14use dashmap::DashMap;
15use futures::StreamExt;
16use queries::channels::ChannelGraph;
17use rand::{prelude::StdRng, Rng, SeedableRng};
18use rpc::{
19 proto::{self},
20 ConnectionId,
21};
22use sea_orm::{
23 entity::prelude::*,
24 sea_query::{Alias, Expr, OnConflict},
25 ActiveValue, Condition, ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbErr,
26 FromQueryResult, IntoActiveModel, IsolationLevel, JoinType, QueryOrder, QuerySelect, Statement,
27 TransactionTrait,
28};
29use serde::{Deserialize, Serialize};
30use sqlx::{
31 migrate::{Migrate, Migration, MigrationSource},
32 Connection,
33};
34use std::{
35 fmt::Write as _,
36 future::Future,
37 marker::PhantomData,
38 ops::{Deref, DerefMut},
39 path::Path,
40 rc::Rc,
41 sync::Arc,
42 time::Duration,
43};
44use tables::*;
45use tokio::sync::{Mutex, OwnedMutexGuard};
46
47pub use ids::*;
48pub use sea_orm::ConnectOptions;
49pub use tables::user::Model as User;
50
51pub struct Database {
52 options: ConnectOptions,
53 pool: DatabaseConnection,
54 rooms: DashMap<RoomId, Arc<Mutex<()>>>,
55 rng: Mutex<StdRng>,
56 executor: Executor,
57 notification_kinds_by_id: HashMap<NotificationKindId, &'static str>,
58 notification_kinds_by_name: HashMap<String, NotificationKindId>,
59 #[cfg(test)]
60 runtime: Option<tokio::runtime::Runtime>,
61}
62
63// The `Database` type has so many methods that its impl blocks are split into
64// separate files in the `queries` folder.
65impl Database {
66 pub async fn new(options: ConnectOptions, executor: Executor) -> Result<Self> {
67 sqlx::any::install_default_drivers();
68 Ok(Self {
69 options: options.clone(),
70 pool: sea_orm::Database::connect(options).await?,
71 rooms: DashMap::with_capacity(16384),
72 rng: Mutex::new(StdRng::seed_from_u64(0)),
73 notification_kinds_by_id: HashMap::default(),
74 notification_kinds_by_name: HashMap::default(),
75 executor,
76 #[cfg(test)]
77 runtime: None,
78 })
79 }
80
81 #[cfg(test)]
82 pub fn reset(&self) {
83 self.rooms.clear();
84 }
85
86 pub async fn migrate(
87 &self,
88 migrations_path: &Path,
89 ignore_checksum_mismatch: bool,
90 ) -> anyhow::Result<Vec<(Migration, Duration)>> {
91 let migrations = MigrationSource::resolve(migrations_path)
92 .await
93 .map_err(|err| anyhow!("failed to load migrations: {err:?}"))?;
94
95 let mut connection = sqlx::AnyConnection::connect(self.options.get_url()).await?;
96
97 connection.ensure_migrations_table().await?;
98 let applied_migrations: HashMap<_, _> = connection
99 .list_applied_migrations()
100 .await?
101 .into_iter()
102 .map(|m| (m.version, m))
103 .collect();
104
105 let mut new_migrations = Vec::new();
106 for migration in migrations {
107 match applied_migrations.get(&migration.version) {
108 Some(applied_migration) => {
109 if migration.checksum != applied_migration.checksum && !ignore_checksum_mismatch
110 {
111 Err(anyhow!(
112 "checksum mismatch for applied migration {}",
113 migration.description
114 ))?;
115 }
116 }
117 None => {
118 let elapsed = connection.apply(&migration).await?;
119 new_migrations.push((migration, elapsed));
120 }
121 }
122 }
123
124 Ok(new_migrations)
125 }
126
127 pub async fn initialize_static_data(&mut self) -> Result<()> {
128 self.initialize_notification_kinds().await?;
129 Ok(())
130 }
131
132 pub async fn transaction<F, Fut, T>(&self, f: F) -> Result<T>
133 where
134 F: Send + Fn(TransactionHandle) -> Fut,
135 Fut: Send + Future<Output = Result<T>>,
136 {
137 let body = async {
138 let mut i = 0;
139 loop {
140 let (tx, result) = self.with_transaction(&f).await?;
141 match result {
142 Ok(result) => match tx.commit().await.map_err(Into::into) {
143 Ok(()) => return Ok(result),
144 Err(error) => {
145 if !self.retry_on_serialization_error(&error, i).await {
146 return Err(error);
147 }
148 }
149 },
150 Err(error) => {
151 tx.rollback().await?;
152 if !self.retry_on_serialization_error(&error, i).await {
153 return Err(error);
154 }
155 }
156 }
157 i += 1;
158 }
159 };
160
161 self.run(body).await
162 }
163
164 async fn optional_room_transaction<F, Fut, T>(&self, f: F) -> Result<Option<RoomGuard<T>>>
165 where
166 F: Send + Fn(TransactionHandle) -> Fut,
167 Fut: Send + Future<Output = Result<Option<(RoomId, T)>>>,
168 {
169 let body = async {
170 let mut i = 0;
171 loop {
172 let (tx, result) = self.with_transaction(&f).await?;
173 match result {
174 Ok(Some((room_id, data))) => {
175 let lock = self.rooms.entry(room_id).or_default().clone();
176 let _guard = lock.lock_owned().await;
177 match tx.commit().await.map_err(Into::into) {
178 Ok(()) => {
179 return Ok(Some(RoomGuard {
180 data,
181 _guard,
182 _not_send: PhantomData,
183 }));
184 }
185 Err(error) => {
186 if !self.retry_on_serialization_error(&error, i).await {
187 return Err(error);
188 }
189 }
190 }
191 }
192 Ok(None) => match tx.commit().await.map_err(Into::into) {
193 Ok(()) => return Ok(None),
194 Err(error) => {
195 if !self.retry_on_serialization_error(&error, i).await {
196 return Err(error);
197 }
198 }
199 },
200 Err(error) => {
201 tx.rollback().await?;
202 if !self.retry_on_serialization_error(&error, i).await {
203 return Err(error);
204 }
205 }
206 }
207 i += 1;
208 }
209 };
210
211 self.run(body).await
212 }
213
214 async fn room_transaction<F, Fut, T>(&self, room_id: RoomId, f: F) -> Result<RoomGuard<T>>
215 where
216 F: Send + Fn(TransactionHandle) -> Fut,
217 Fut: Send + Future<Output = Result<T>>,
218 {
219 let body = async {
220 let mut i = 0;
221 loop {
222 let lock = self.rooms.entry(room_id).or_default().clone();
223 let _guard = lock.lock_owned().await;
224 let (tx, result) = self.with_transaction(&f).await?;
225 match result {
226 Ok(data) => match tx.commit().await.map_err(Into::into) {
227 Ok(()) => {
228 return Ok(RoomGuard {
229 data,
230 _guard,
231 _not_send: PhantomData,
232 });
233 }
234 Err(error) => {
235 if !self.retry_on_serialization_error(&error, i).await {
236 return Err(error);
237 }
238 }
239 },
240 Err(error) => {
241 tx.rollback().await?;
242 if !self.retry_on_serialization_error(&error, i).await {
243 return Err(error);
244 }
245 }
246 }
247 i += 1;
248 }
249 };
250
251 self.run(body).await
252 }
253
254 async fn with_transaction<F, Fut, T>(&self, f: &F) -> Result<(DatabaseTransaction, Result<T>)>
255 where
256 F: Send + Fn(TransactionHandle) -> Fut,
257 Fut: Send + Future<Output = Result<T>>,
258 {
259 let tx = self
260 .pool
261 .begin_with_config(Some(IsolationLevel::Serializable), None)
262 .await?;
263
264 let mut tx = Arc::new(Some(tx));
265 let result = f(TransactionHandle(tx.clone())).await;
266 let Some(tx) = Arc::get_mut(&mut tx).and_then(|tx| tx.take()) else {
267 return Err(anyhow!(
268 "couldn't complete transaction because it's still in use"
269 ))?;
270 };
271
272 Ok((tx, result))
273 }
274
275 async fn run<F, T>(&self, future: F) -> Result<T>
276 where
277 F: Future<Output = Result<T>>,
278 {
279 #[cfg(test)]
280 {
281 if let Executor::Deterministic(executor) = &self.executor {
282 executor.simulate_random_delay().await;
283 }
284
285 self.runtime.as_ref().unwrap().block_on(future)
286 }
287
288 #[cfg(not(test))]
289 {
290 future.await
291 }
292 }
293
294 async fn retry_on_serialization_error(&self, error: &Error, prev_attempt_count: u32) -> bool {
295 // If the error is due to a failure to serialize concurrent transactions, then retry
296 // this transaction after a delay. With each subsequent retry, double the delay duration.
297 // Also vary the delay randomly in order to ensure different database connections retry
298 // at different times.
299 if is_serialization_error(error) {
300 let base_delay = 4_u64 << prev_attempt_count.min(16);
301 let randomized_delay = base_delay as f32 * self.rng.lock().await.gen_range(0.5..=2.0);
302 log::info!(
303 "retrying transaction after serialization error. delay: {} ms.",
304 randomized_delay
305 );
306 self.executor
307 .sleep(Duration::from_millis(randomized_delay as u64))
308 .await;
309 true
310 } else {
311 false
312 }
313 }
314}
315
316fn is_serialization_error(error: &Error) -> bool {
317 const SERIALIZATION_FAILURE_CODE: &'static str = "40001";
318 match error {
319 Error::Database(
320 DbErr::Exec(sea_orm::RuntimeErr::SqlxError(error))
321 | DbErr::Query(sea_orm::RuntimeErr::SqlxError(error)),
322 ) if error
323 .as_database_error()
324 .and_then(|error| error.code())
325 .as_deref()
326 == Some(SERIALIZATION_FAILURE_CODE) =>
327 {
328 true
329 }
330 _ => false,
331 }
332}
333
334pub struct TransactionHandle(Arc<Option<DatabaseTransaction>>);
335
336impl Deref for TransactionHandle {
337 type Target = DatabaseTransaction;
338
339 fn deref(&self) -> &Self::Target {
340 self.0.as_ref().as_ref().unwrap()
341 }
342}
343
344pub struct RoomGuard<T> {
345 data: T,
346 _guard: OwnedMutexGuard<()>,
347 _not_send: PhantomData<Rc<()>>,
348}
349
350impl<T> Deref for RoomGuard<T> {
351 type Target = T;
352
353 fn deref(&self) -> &T {
354 &self.data
355 }
356}
357
358impl<T> DerefMut for RoomGuard<T> {
359 fn deref_mut(&mut self) -> &mut T {
360 &mut self.data
361 }
362}
363
364impl<T> RoomGuard<T> {
365 pub fn into_inner(self) -> T {
366 self.data
367 }
368}
369
370#[derive(Clone, Debug, PartialEq, Eq)]
371pub enum Contact {
372 Accepted { user_id: UserId, busy: bool },
373 Outgoing { user_id: UserId },
374 Incoming { user_id: UserId },
375}
376
377impl Contact {
378 pub fn user_id(&self) -> UserId {
379 match self {
380 Contact::Accepted { user_id, .. } => *user_id,
381 Contact::Outgoing { user_id } => *user_id,
382 Contact::Incoming { user_id, .. } => *user_id,
383 }
384 }
385}
386
387pub type NotificationBatch = Vec<(UserId, proto::Notification)>;
388
389pub struct CreatedChannelMessage {
390 pub message_id: MessageId,
391 pub participant_connection_ids: Vec<ConnectionId>,
392 pub channel_members: Vec<UserId>,
393 pub notifications: NotificationBatch,
394}
395
396#[derive(Clone, Debug, PartialEq, Eq, FromQueryResult, Serialize, Deserialize)]
397pub struct Invite {
398 pub email_address: String,
399 pub email_confirmation_code: String,
400}
401
402#[derive(Clone, Debug, Deserialize)]
403pub struct NewSignup {
404 pub email_address: String,
405 pub platform_mac: bool,
406 pub platform_windows: bool,
407 pub platform_linux: bool,
408 pub editor_features: Vec<String>,
409 pub programming_languages: Vec<String>,
410 pub device_id: Option<String>,
411 pub added_to_mailing_list: bool,
412 pub created_at: Option<DateTime>,
413}
414
415#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, FromQueryResult)]
416pub struct WaitlistSummary {
417 pub count: i64,
418 pub linux_count: i64,
419 pub mac_count: i64,
420 pub windows_count: i64,
421 pub unknown_count: i64,
422}
423
424#[derive(Debug, Serialize, Deserialize)]
425pub struct NewUserParams {
426 pub github_login: String,
427 pub github_user_id: i32,
428}
429
430#[derive(Debug)]
431pub struct NewUserResult {
432 pub user_id: UserId,
433 pub metrics_id: String,
434 pub inviting_user_id: Option<UserId>,
435 pub signup_device_id: Option<String>,
436}
437
438#[derive(Debug)]
439pub struct MoveChannelResult {
440 pub participants_to_update: HashMap<UserId, ChannelsForUser>,
441 pub participants_to_remove: HashSet<UserId>,
442 pub moved_channels: HashSet<ChannelId>,
443}
444
445#[derive(Debug)]
446pub struct RenameChannelResult {
447 pub channel: Channel,
448 pub participants_to_update: HashMap<UserId, Channel>,
449}
450
451#[derive(Debug)]
452pub struct CreateChannelResult {
453 pub channel: Channel,
454 pub participants_to_update: Vec<(UserId, ChannelsForUser)>,
455}
456
457#[derive(Debug)]
458pub struct SetChannelVisibilityResult {
459 pub participants_to_update: HashMap<UserId, ChannelsForUser>,
460 pub participants_to_remove: HashSet<UserId>,
461}
462
463#[derive(Debug)]
464pub struct MembershipUpdated {
465 pub channel_id: ChannelId,
466 pub new_channels: ChannelsForUser,
467 pub removed_channels: Vec<ChannelId>,
468}
469
470#[derive(Debug)]
471pub enum SetMemberRoleResult {
472 InviteUpdated(Channel),
473 MembershipUpdated(MembershipUpdated),
474}
475
476#[derive(Debug)]
477pub struct InviteMemberResult {
478 pub channel: Channel,
479 pub notifications: NotificationBatch,
480}
481
482#[derive(Debug)]
483pub struct RespondToChannelInvite {
484 pub membership_update: Option<MembershipUpdated>,
485 pub notifications: NotificationBatch,
486}
487
488#[derive(Debug)]
489pub struct RemoveChannelMemberResult {
490 pub membership_update: MembershipUpdated,
491 pub notification_id: Option<NotificationId>,
492}
493
494#[derive(FromQueryResult, Debug, PartialEq, Eq, Hash)]
495pub struct Channel {
496 pub id: ChannelId,
497 pub name: String,
498 pub visibility: ChannelVisibility,
499 pub role: ChannelRole,
500}
501
502impl Channel {
503 pub fn to_proto(&self) -> proto::Channel {
504 proto::Channel {
505 id: self.id.to_proto(),
506 name: self.name.clone(),
507 visibility: self.visibility.into(),
508 role: self.role.into(),
509 }
510 }
511}
512
513#[derive(Debug, PartialEq, Eq, Hash)]
514pub struct ChannelMember {
515 pub role: ChannelRole,
516 pub user_id: UserId,
517 pub kind: proto::channel_member::Kind,
518}
519
520impl ChannelMember {
521 pub fn to_proto(&self) -> proto::ChannelMember {
522 proto::ChannelMember {
523 role: self.role.into(),
524 user_id: self.user_id.to_proto(),
525 kind: self.kind.into(),
526 }
527 }
528}
529
530#[derive(Debug, PartialEq)]
531pub struct ChannelsForUser {
532 pub channels: ChannelGraph,
533 pub channel_participants: HashMap<ChannelId, Vec<UserId>>,
534 pub unseen_buffer_changes: Vec<proto::UnseenChannelBufferChange>,
535 pub channel_messages: Vec<proto::UnseenChannelMessage>,
536}
537
538#[derive(Debug)]
539pub struct RejoinedChannelBuffer {
540 pub buffer: proto::RejoinedChannelBuffer,
541 pub old_connection_id: ConnectionId,
542}
543
544#[derive(Clone)]
545pub struct JoinRoom {
546 pub room: proto::Room,
547 pub channel_id: Option<ChannelId>,
548 pub channel_members: Vec<UserId>,
549}
550
551pub struct RejoinedRoom {
552 pub room: proto::Room,
553 pub rejoined_projects: Vec<RejoinedProject>,
554 pub reshared_projects: Vec<ResharedProject>,
555 pub channel_id: Option<ChannelId>,
556 pub channel_members: Vec<UserId>,
557}
558
559pub struct ResharedProject {
560 pub id: ProjectId,
561 pub old_connection_id: ConnectionId,
562 pub collaborators: Vec<ProjectCollaborator>,
563 pub worktrees: Vec<proto::WorktreeMetadata>,
564}
565
566pub struct RejoinedProject {
567 pub id: ProjectId,
568 pub old_connection_id: ConnectionId,
569 pub collaborators: Vec<ProjectCollaborator>,
570 pub worktrees: Vec<RejoinedWorktree>,
571 pub language_servers: Vec<proto::LanguageServer>,
572}
573
574#[derive(Debug)]
575pub struct RejoinedWorktree {
576 pub id: u64,
577 pub abs_path: String,
578 pub root_name: String,
579 pub visible: bool,
580 pub updated_entries: Vec<proto::Entry>,
581 pub removed_entries: Vec<u64>,
582 pub updated_repositories: Vec<proto::RepositoryEntry>,
583 pub removed_repositories: Vec<u64>,
584 pub diagnostic_summaries: Vec<proto::DiagnosticSummary>,
585 pub settings_files: Vec<WorktreeSettingsFile>,
586 pub scan_id: u64,
587 pub completed_scan_id: u64,
588}
589
590pub struct LeftRoom {
591 pub room: proto::Room,
592 pub channel_id: Option<ChannelId>,
593 pub channel_members: Vec<UserId>,
594 pub left_projects: HashMap<ProjectId, LeftProject>,
595 pub canceled_calls_to_user_ids: Vec<UserId>,
596 pub deleted: bool,
597}
598
599pub struct RefreshedRoom {
600 pub room: proto::Room,
601 pub channel_id: Option<ChannelId>,
602 pub channel_members: Vec<UserId>,
603 pub stale_participant_user_ids: Vec<UserId>,
604 pub canceled_calls_to_user_ids: Vec<UserId>,
605}
606
607pub struct RefreshedChannelBuffer {
608 pub connection_ids: Vec<ConnectionId>,
609 pub collaborators: Vec<proto::Collaborator>,
610}
611
612pub struct Project {
613 pub collaborators: Vec<ProjectCollaborator>,
614 pub worktrees: BTreeMap<u64, Worktree>,
615 pub language_servers: Vec<proto::LanguageServer>,
616}
617
618pub struct ProjectCollaborator {
619 pub connection_id: ConnectionId,
620 pub user_id: UserId,
621 pub replica_id: ReplicaId,
622 pub is_host: bool,
623}
624
625impl ProjectCollaborator {
626 pub fn to_proto(&self) -> proto::Collaborator {
627 proto::Collaborator {
628 peer_id: Some(self.connection_id.into()),
629 replica_id: self.replica_id.0 as u32,
630 user_id: self.user_id.to_proto(),
631 }
632 }
633}
634
635#[derive(Debug)]
636pub struct LeftProject {
637 pub id: ProjectId,
638 pub host_user_id: UserId,
639 pub host_connection_id: ConnectionId,
640 pub connection_ids: Vec<ConnectionId>,
641}
642
643pub struct Worktree {
644 pub id: u64,
645 pub abs_path: String,
646 pub root_name: String,
647 pub visible: bool,
648 pub entries: Vec<proto::Entry>,
649 pub repository_entries: BTreeMap<u64, proto::RepositoryEntry>,
650 pub diagnostic_summaries: Vec<proto::DiagnosticSummary>,
651 pub settings_files: Vec<WorktreeSettingsFile>,
652 pub scan_id: u64,
653 pub completed_scan_id: u64,
654}
655
656#[derive(Debug)]
657pub struct WorktreeSettingsFile {
658 pub path: String,
659 pub content: String,
660}