Get collab2 green

crates/collab2/src/db.rs

@@ -0,0 +1,672 @@
+#[cfg(test)]
+pub mod tests;
+
+#[cfg(test)]
+pub use tests::TestDb;
+
+mod ids;
+mod queries;
+mod tables;
+
+use crate::{executor::Executor, Error, Result};
+use anyhow::anyhow;
+use collections::{BTreeMap, HashMap, HashSet};
+use dashmap::DashMap;
+use futures::StreamExt;
+use rand::{prelude::StdRng, Rng, SeedableRng};
+use rpc::{
+    proto::{self},
+    ConnectionId,
+};
+use sea_orm::{
+    entity::prelude::*,
+    sea_query::{Alias, Expr, OnConflict},
+    ActiveValue, Condition, ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbErr,
+    FromQueryResult, IntoActiveModel, IsolationLevel, JoinType, QueryOrder, QuerySelect, Statement,
+    TransactionTrait,
+};
+use serde::{Deserialize, Serialize};
+use sqlx::{
+    migrate::{Migrate, Migration, MigrationSource},
+    Connection,
+};
+use std::{
+    fmt::Write as _,
+    future::Future,
+    marker::PhantomData,
+    ops::{Deref, DerefMut},
+    path::Path,
+    rc::Rc,
+    sync::Arc,
+    time::Duration,
+};
+use tables::*;
+use tokio::sync::{Mutex, OwnedMutexGuard};
+
+pub use ids::*;
+pub use sea_orm::ConnectOptions;
+pub use tables::user::Model as User;
+
+pub struct Database {
+    options: ConnectOptions,
+    pool: DatabaseConnection,
+    rooms: DashMap<RoomId, Arc<Mutex<()>>>,
+    rng: Mutex<StdRng>,
+    executor: Executor,
+    notification_kinds_by_id: HashMap<NotificationKindId, &'static str>,
+    notification_kinds_by_name: HashMap<String, NotificationKindId>,
+    #[cfg(test)]
+    runtime: Option<tokio::runtime::Runtime>,
+}
+
+// The `Database` type has so many methods that its impl blocks are split into
+// separate files in the `queries` folder.
+impl Database {
+    pub async fn new(options: ConnectOptions, executor: Executor) -> Result<Self> {
+        sqlx::any::install_default_drivers();
+        Ok(Self {
+            options: options.clone(),
+            pool: sea_orm::Database::connect(options).await?,
+            rooms: DashMap::with_capacity(16384),
+            rng: Mutex::new(StdRng::seed_from_u64(0)),
+            notification_kinds_by_id: HashMap::default(),
+            notification_kinds_by_name: HashMap::default(),
+            executor,
+            #[cfg(test)]
+            runtime: None,
+        })
+    }
+
+    #[cfg(test)]
+    pub fn reset(&self) {
+        self.rooms.clear();
+    }
+
+    pub async fn migrate(
+        &self,
+        migrations_path: &Path,
+        ignore_checksum_mismatch: bool,
+    ) -> anyhow::Result<Vec<(Migration, Duration)>> {
+        let migrations = MigrationSource::resolve(migrations_path)
+            .await
+            .map_err(|err| anyhow!("failed to load migrations: {err:?}"))?;
+
+        let mut connection = sqlx::AnyConnection::connect(self.options.get_url()).await?;
+
+        connection.ensure_migrations_table().await?;
+        let applied_migrations: HashMap<_, _> = connection
+            .list_applied_migrations()
+            .await?
+            .into_iter()
+            .map(|m| (m.version, m))
+            .collect();
+
+        let mut new_migrations = Vec::new();
+        for migration in migrations {
+            match applied_migrations.get(&migration.version) {
+                Some(applied_migration) => {
+                    if migration.checksum != applied_migration.checksum && !ignore_checksum_mismatch
+                    {
+                        Err(anyhow!(
+                            "checksum mismatch for applied migration {}",
+                            migration.description
+                        ))?;
+                    }
+                }
+                None => {
+                    let elapsed = connection.apply(&migration).await?;
+                    new_migrations.push((migration, elapsed));
+                }
+            }
+        }
+
+        Ok(new_migrations)
+    }
+
+    pub async fn initialize_static_data(&mut self) -> Result<()> {
+        self.initialize_notification_kinds().await?;
+        Ok(())
+    }
+
+    pub async fn transaction<F, Fut, T>(&self, f: F) -> Result<T>
+    where
+        F: Send + Fn(TransactionHandle) -> Fut,
+        Fut: Send + Future<Output = Result<T>>,
+    {
+        let body = async {
+            let mut i = 0;
+            loop {
+                let (tx, result) = self.with_transaction(&f).await?;
+                match result {
+                    Ok(result) => match tx.commit().await.map_err(Into::into) {
+                        Ok(()) => return Ok(result),
+                        Err(error) => {
+                            if !self.retry_on_serialization_error(&error, i).await {
+                                return Err(error);
+                            }
+                        }
+                    },
+                    Err(error) => {
+                        tx.rollback().await?;
+                        if !self.retry_on_serialization_error(&error, i).await {
+                            return Err(error);
+                        }
+                    }
+                }
+                i += 1;
+            }
+        };
+
+        self.run(body).await
+    }
+
+    async fn optional_room_transaction<F, Fut, T>(&self, f: F) -> Result<Option<RoomGuard<T>>>
+    where
+        F: Send + Fn(TransactionHandle) -> Fut,
+        Fut: Send + Future<Output = Result<Option<(RoomId, T)>>>,
+    {
+        let body = async {
+            let mut i = 0;
+            loop {
+                let (tx, result) = self.with_transaction(&f).await?;
+                match result {
+                    Ok(Some((room_id, data))) => {
+                        let lock = self.rooms.entry(room_id).or_default().clone();
+                        let _guard = lock.lock_owned().await;
+                        match tx.commit().await.map_err(Into::into) {
+                            Ok(()) => {
+                                return Ok(Some(RoomGuard {
+                                    data,
+                                    _guard,
+                                    _not_send: PhantomData,
+                                }));
+                            }
+                            Err(error) => {
+                                if !self.retry_on_serialization_error(&error, i).await {
+                                    return Err(error);
+                                }
+                            }
+                        }
+                    }
+                    Ok(None) => match tx.commit().await.map_err(Into::into) {
+                        Ok(()) => return Ok(None),
+                        Err(error) => {
+                            if !self.retry_on_serialization_error(&error, i).await {
+                                return Err(error);
+                            }
+                        }
+                    },
+                    Err(error) => {
+                        tx.rollback().await?;
+                        if !self.retry_on_serialization_error(&error, i).await {
+                            return Err(error);
+                        }
+                    }
+                }
+                i += 1;
+            }
+        };
+
+        self.run(body).await
+    }
+
+    async fn room_transaction<F, Fut, T>(&self, room_id: RoomId, f: F) -> Result<RoomGuard<T>>
+    where
+        F: Send + Fn(TransactionHandle) -> Fut,
+        Fut: Send + Future<Output = Result<T>>,
+    {
+        let body = async {
+            let mut i = 0;
+            loop {
+                let lock = self.rooms.entry(room_id).or_default().clone();
+                let _guard = lock.lock_owned().await;
+                let (tx, result) = self.with_transaction(&f).await?;
+                match result {
+                    Ok(data) => match tx.commit().await.map_err(Into::into) {
+                        Ok(()) => {
+                            return Ok(RoomGuard {
+                                data,
+                                _guard,
+                                _not_send: PhantomData,
+                            });
+                        }
+                        Err(error) => {
+                            if !self.retry_on_serialization_error(&error, i).await {
+                                return Err(error);
+                            }
+                        }
+                    },
+                    Err(error) => {
+                        tx.rollback().await?;
+                        if !self.retry_on_serialization_error(&error, i).await {
+                            return Err(error);
+                        }
+                    }
+                }
+                i += 1;
+            }
+        };
+
+        self.run(body).await
+    }
+
+    async fn with_transaction<F, Fut, T>(&self, f: &F) -> Result<(DatabaseTransaction, Result<T>)>
+    where
+        F: Send + Fn(TransactionHandle) -> Fut,
+        Fut: Send + Future<Output = Result<T>>,
+    {
+        let tx = self
+            .pool
+            .begin_with_config(Some(IsolationLevel::Serializable), None)
+            .await?;
+
+        let mut tx = Arc::new(Some(tx));
+        let result = f(TransactionHandle(tx.clone())).await;
+        let Some(tx) = Arc::get_mut(&mut tx).and_then(|tx| tx.take()) else {
+            return Err(anyhow!(
+                "couldn't complete transaction because it's still in use"
+            ))?;
+        };
+
+        Ok((tx, result))
+    }
+
+    async fn run<F, T>(&self, future: F) -> Result<T>
+    where
+        F: Future<Output = Result<T>>,
+    {
+        #[cfg(test)]
+        {
+            if let Executor::Deterministic(executor) = &self.executor {
+                executor.simulate_random_delay().await;
+            }
+
+            self.runtime.as_ref().unwrap().block_on(future)
+        }
+
+        #[cfg(not(test))]
+        {
+            future.await
+        }
+    }
+
+    async fn retry_on_serialization_error(&self, error: &Error, prev_attempt_count: u32) -> bool {
+        // If the error is due to a failure to serialize concurrent transactions, then retry
+        // this transaction after a delay. With each subsequent retry, double the delay duration.
+        // Also vary the delay randomly in order to ensure different database connections retry
+        // at different times.
+        if is_serialization_error(error) {
+            let base_delay = 4_u64 << prev_attempt_count.min(16);
+            let randomized_delay = base_delay as f32 * self.rng.lock().await.gen_range(0.5..=2.0);
+            log::info!(
+                "retrying transaction after serialization error. delay: {} ms.",
+                randomized_delay
+            );
+            self.executor
+                .sleep(Duration::from_millis(randomized_delay as u64))
+                .await;
+            true
+        } else {
+            false
+        }
+    }
+}
+
+fn is_serialization_error(error: &Error) -> bool {
+    const SERIALIZATION_FAILURE_CODE: &'static str = "40001";
+    match error {
+        Error::Database(
+            DbErr::Exec(sea_orm::RuntimeErr::SqlxError(error))
+            | DbErr::Query(sea_orm::RuntimeErr::SqlxError(error)),
+        ) if error
+            .as_database_error()
+            .and_then(|error| error.code())
+            .as_deref()
+            == Some(SERIALIZATION_FAILURE_CODE) =>
+        {
+            true
+        }
+        _ => false,
+    }
+}
+
+pub struct TransactionHandle(Arc<Option<DatabaseTransaction>>);
+
+impl Deref for TransactionHandle {
+    type Target = DatabaseTransaction;
+
+    fn deref(&self) -> &Self::Target {
+        self.0.as_ref().as_ref().unwrap()
+    }
+}
+
+pub struct RoomGuard<T> {
+    data: T,
+    _guard: OwnedMutexGuard<()>,
+    _not_send: PhantomData<Rc<()>>,
+}
+
+impl<T> Deref for RoomGuard<T> {
+    type Target = T;
+
+    fn deref(&self) -> &T {
+        &self.data
+    }
+}
+
+impl<T> DerefMut for RoomGuard<T> {
+    fn deref_mut(&mut self) -> &mut T {
+        &mut self.data
+    }
+}
+
+impl<T> RoomGuard<T> {
+    pub fn into_inner(self) -> T {
+        self.data
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum Contact {
+    Accepted { user_id: UserId, busy: bool },
+    Outgoing { user_id: UserId },
+    Incoming { user_id: UserId },
+}
+
+impl Contact {
+    pub fn user_id(&self) -> UserId {
+        match self {
+            Contact::Accepted { user_id, .. } => *user_id,
+            Contact::Outgoing { user_id } => *user_id,
+            Contact::Incoming { user_id, .. } => *user_id,
+        }
+    }
+}
+
+pub type NotificationBatch = Vec<(UserId, proto::Notification)>;
+
+pub struct CreatedChannelMessage {
+    pub message_id: MessageId,
+    pub participant_connection_ids: Vec<ConnectionId>,
+    pub channel_members: Vec<UserId>,
+    pub notifications: NotificationBatch,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, FromQueryResult, Serialize, Deserialize)]
+pub struct Invite {
+    pub email_address: String,
+    pub email_confirmation_code: String,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct NewSignup {
+    pub email_address: String,
+    pub platform_mac: bool,
+    pub platform_windows: bool,
+    pub platform_linux: bool,
+    pub editor_features: Vec<String>,
+    pub programming_languages: Vec<String>,
+    pub device_id: Option<String>,
+    pub added_to_mailing_list: bool,
+    pub created_at: Option<DateTime>,
+}
+
+#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, FromQueryResult)]
+pub struct WaitlistSummary {
+    pub count: i64,
+    pub linux_count: i64,
+    pub mac_count: i64,
+    pub windows_count: i64,
+    pub unknown_count: i64,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct NewUserParams {
+    pub github_login: String,
+    pub github_user_id: i32,
+}
+
+#[derive(Debug)]
+pub struct NewUserResult {
+    pub user_id: UserId,
+    pub metrics_id: String,
+    pub inviting_user_id: Option<UserId>,
+    pub signup_device_id: Option<String>,
+}
+
+#[derive(Debug)]
+pub struct MoveChannelResult {
+    pub participants_to_update: HashMap<UserId, ChannelsForUser>,
+    pub participants_to_remove: HashSet<UserId>,
+    pub moved_channels: HashSet<ChannelId>,
+}
+
+#[derive(Debug)]
+pub struct RenameChannelResult {
+    pub channel: Channel,
+    pub participants_to_update: HashMap<UserId, Channel>,
+}
+
+#[derive(Debug)]
+pub struct CreateChannelResult {
+    pub channel: Channel,
+    pub participants_to_update: Vec<(UserId, ChannelsForUser)>,
+}
+
+#[derive(Debug)]
+pub struct SetChannelVisibilityResult {
+    pub participants_to_update: HashMap<UserId, ChannelsForUser>,
+    pub participants_to_remove: HashSet<UserId>,
+    pub channels_to_remove: Vec<ChannelId>,
+}
+
+#[derive(Debug)]
+pub struct MembershipUpdated {
+    pub channel_id: ChannelId,
+    pub new_channels: ChannelsForUser,
+    pub removed_channels: Vec<ChannelId>,
+}
+
+#[derive(Debug)]
+pub enum SetMemberRoleResult {
+    InviteUpdated(Channel),
+    MembershipUpdated(MembershipUpdated),
+}
+
+#[derive(Debug)]
+pub struct InviteMemberResult {
+    pub channel: Channel,
+    pub notifications: NotificationBatch,
+}
+
+#[derive(Debug)]
+pub struct RespondToChannelInvite {
+    pub membership_update: Option<MembershipUpdated>,
+    pub notifications: NotificationBatch,
+}
+
+#[derive(Debug)]
+pub struct RemoveChannelMemberResult {
+    pub membership_update: MembershipUpdated,
+    pub notification_id: Option<NotificationId>,
+}
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct Channel {
+    pub id: ChannelId,
+    pub name: String,
+    pub visibility: ChannelVisibility,
+    pub role: ChannelRole,
+    pub parent_path: Vec<ChannelId>,
+}
+
+impl Channel {
+    fn from_model(value: channel::Model, role: ChannelRole) -> Self {
+        Channel {
+            id: value.id,
+            visibility: value.visibility,
+            name: value.clone().name,
+            role,
+            parent_path: value.ancestors().collect(),
+        }
+    }
+
+    pub fn to_proto(&self) -> proto::Channel {
+        proto::Channel {
+            id: self.id.to_proto(),
+            name: self.name.clone(),
+            visibility: self.visibility.into(),
+            role: self.role.into(),
+            parent_path: self.parent_path.iter().map(|c| c.to_proto()).collect(),
+        }
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ChannelMember {
+    pub role: ChannelRole,
+    pub user_id: UserId,
+    pub kind: proto::channel_member::Kind,
+}
+
+impl ChannelMember {
+    pub fn to_proto(&self) -> proto::ChannelMember {
+        proto::ChannelMember {
+            role: self.role.into(),
+            user_id: self.user_id.to_proto(),
+            kind: self.kind.into(),
+        }
+    }
+}
+
+#[derive(Debug, PartialEq)]
+pub struct ChannelsForUser {
+    pub channels: Vec<Channel>,
+    pub channel_participants: HashMap<ChannelId, Vec<UserId>>,
+    pub unseen_buffer_changes: Vec<proto::UnseenChannelBufferChange>,
+    pub channel_messages: Vec<proto::UnseenChannelMessage>,
+}
+
+#[derive(Debug)]
+pub struct RejoinedChannelBuffer {
+    pub buffer: proto::RejoinedChannelBuffer,
+    pub old_connection_id: ConnectionId,
+}
+
+#[derive(Clone)]
+pub struct JoinRoom {
+    pub room: proto::Room,
+    pub channel_id: Option<ChannelId>,
+    pub channel_members: Vec<UserId>,
+}
+
+pub struct RejoinedRoom {
+    pub room: proto::Room,
+    pub rejoined_projects: Vec<RejoinedProject>,
+    pub reshared_projects: Vec<ResharedProject>,
+    pub channel_id: Option<ChannelId>,
+    pub channel_members: Vec<UserId>,
+}
+
+pub struct ResharedProject {
+    pub id: ProjectId,
+    pub old_connection_id: ConnectionId,
+    pub collaborators: Vec<ProjectCollaborator>,
+    pub worktrees: Vec<proto::WorktreeMetadata>,
+}
+
+pub struct RejoinedProject {
+    pub id: ProjectId,
+    pub old_connection_id: ConnectionId,
+    pub collaborators: Vec<ProjectCollaborator>,
+    pub worktrees: Vec<RejoinedWorktree>,
+    pub language_servers: Vec<proto::LanguageServer>,
+}
+
+#[derive(Debug)]
+pub struct RejoinedWorktree {
+    pub id: u64,
+    pub abs_path: String,
+    pub root_name: String,
+    pub visible: bool,
+    pub updated_entries: Vec<proto::Entry>,
+    pub removed_entries: Vec<u64>,
+    pub updated_repositories: Vec<proto::RepositoryEntry>,
+    pub removed_repositories: Vec<u64>,
+    pub diagnostic_summaries: Vec<proto::DiagnosticSummary>,
+    pub settings_files: Vec<WorktreeSettingsFile>,
+    pub scan_id: u64,
+    pub completed_scan_id: u64,
+}
+
+pub struct LeftRoom {
+    pub room: proto::Room,
+    pub channel_id: Option<ChannelId>,
+    pub channel_members: Vec<UserId>,
+    pub left_projects: HashMap<ProjectId, LeftProject>,
+    pub canceled_calls_to_user_ids: Vec<UserId>,
+    pub deleted: bool,
+}
+
+pub struct RefreshedRoom {
+    pub room: proto::Room,
+    pub channel_id: Option<ChannelId>,
+    pub channel_members: Vec<UserId>,
+    pub stale_participant_user_ids: Vec<UserId>,
+    pub canceled_calls_to_user_ids: Vec<UserId>,
+}
+
+pub struct RefreshedChannelBuffer {
+    pub connection_ids: Vec<ConnectionId>,
+    pub collaborators: Vec<proto::Collaborator>,
+}
+
+pub struct Project {
+    pub collaborators: Vec<ProjectCollaborator>,
+    pub worktrees: BTreeMap<u64, Worktree>,
+    pub language_servers: Vec<proto::LanguageServer>,
+}
+
+pub struct ProjectCollaborator {
+    pub connection_id: ConnectionId,
+    pub user_id: UserId,
+    pub replica_id: ReplicaId,
+    pub is_host: bool,
+}
+
+impl ProjectCollaborator {
+    pub fn to_proto(&self) -> proto::Collaborator {
+        proto::Collaborator {
+            peer_id: Some(self.connection_id.into()),
+            replica_id: self.replica_id.0 as u32,
+            user_id: self.user_id.to_proto(),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct LeftProject {
+    pub id: ProjectId,
+    pub host_user_id: UserId,
+    pub host_connection_id: ConnectionId,
+    pub connection_ids: Vec<ConnectionId>,
+}
+
+pub struct Worktree {
+    pub id: u64,
+    pub abs_path: String,
+    pub root_name: String,
+    pub visible: bool,
+    pub entries: Vec<proto::Entry>,
+    pub repository_entries: BTreeMap<u64, proto::RepositoryEntry>,
+    pub diagnostic_summaries: Vec<proto::DiagnosticSummary>,
+    pub settings_files: Vec<WorktreeSettingsFile>,
+    pub scan_id: u64,
+    pub completed_scan_id: u64,
+}
+
+#[derive(Debug)]
+pub struct WorktreeSettingsFile {
+    pub path: String,
+    pub content: String,
+}