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