Move app state creating Broadcasts

.gitignore

@@ -14,4 +14,4 @@ backend/databases
 
 *.log
 
-plugin/coverage
+*.sqlx

backend/config-e2e.yml

@@ -1,10 +1,13 @@
 database:
   databases_directory_path: databases
   max_connections: 12
+
 server:
   host: 0.0.0.0
   port: 3000
   max_body_size_mb: 512
+  max_clients_per_vault: 256
+
 users:
   user_tokens:
   - name: admin

backend/sync_server/src/app_state.rs

@@ -0,0 +1,34 @@
+pub mod broadcasts;
+pub mod database;
+
+use std::ffi::OsString;
+
+use anyhow::Result;
+use broadcasts::Broadcasts;
+use database::Database;
+
+use crate::{config::Config, consts::DEFAULT_CONFIG_PATH};
+
+#[derive(Clone, Debug)]
+pub struct AppState {
+    pub config: Config,
+    pub database: Database,
+    pub broadcasts: Broadcasts,
+}
+
+impl AppState {
+    pub async fn try_new(config_path: Option<OsString>) -> Result<Self> {
+        let config_path = config_path.unwrap_or_else(|| OsString::from(DEFAULT_CONFIG_PATH));
+        let path = std::path::PathBuf::from(config_path);
+
+        let config = Config::read_or_create(&path).await?;
+        let database = Database::try_new(&config.database).await?;
+        let broadcasts = Broadcasts::new(&config.server);
+
+        Ok(Self {
+            config,
+            database,
+            broadcasts,
+        })
+    }
+}

backend/sync_server/src/app_state/broadcasts.rs

@@ -0,0 +1,59 @@
+use std::{collections::HashMap, sync::Arc};
+
+use anyhow::Context;
+use tokio::sync::{Mutex, broadcast};
+
+use super::database::models::{DocumentVersionWithoutContent, VaultId};
+use crate::{
+    config::server_config::ServerConfig,
+    errors::{SyncServerError, server_error},
+};
+
+#[derive(Debug, Clone)]
+pub struct Broadcasts {
+    max_clients_per_vault: usize,
+    tx: Arc<Mutex<HashMap<VaultId, broadcast::Sender<DocumentVersionWithoutContent>>>>,
+}
+
+impl Broadcasts {
+    pub fn new(server_config: &ServerConfig) -> Self {
+        Self {
+            max_clients_per_vault: server_config.max_clients_per_vault,
+            tx: Arc::new(Mutex::new(HashMap::new())),
+        }
+    }
+
+    pub async fn get_receiver(
+        &self,
+        vault: VaultId,
+    ) -> broadcast::Receiver<DocumentVersionWithoutContent> {
+        let tx = self.get_or_create(vault).await;
+
+        tx.subscribe()
+    }
+
+    pub async fn send(
+        &self,
+        vault: VaultId,
+        document: DocumentVersionWithoutContent,
+    ) -> Result<(), SyncServerError> {
+        let tx = self.get_or_create(vault).await;
+
+        tx.send(document)
+            .context("Cannot broadcast update message to websocket listeners")
+            .map_err(server_error)?;
+
+        Ok(())
+    }
+
+    async fn get_or_create(
+        &self,
+        vault: VaultId,
+    ) -> broadcast::Sender<DocumentVersionWithoutContent> {
+        let mut tx = self.tx.lock().await;
+
+        tx.entry(vault)
+            .or_insert_with(|| broadcast::channel(self.max_clients_per_vault).0.clone())
+            .clone()
+    }
+}

backend/sync_server/src/app_state/database.rs

@@ -0,0 +1,357 @@
+use core::time::Duration;
+use std::{collections::HashMap, sync::Arc};
+
+use anyhow::{Context as _, Result};
+use models::{
+    DocumentId, DocumentVersionWithoutContent, StoredDocumentVersion, VaultId, VaultUpdateId,
+};
+use sqlx::{sqlite::SqliteConnectOptions, types::chrono::Utc};
+pub mod models;
+use sqlx::{Pool, Sqlite, sqlite::SqlitePoolOptions};
+use tokio::sync::Mutex;
+use uuid::fmt::Hyphenated;
+
+use crate::config::database_config::DatabaseConfig;
+
+#[derive(Clone, Debug)]
+pub struct Database {
+    config: DatabaseConfig,
+    connection_pools: Arc<Mutex<HashMap<VaultId, Pool<Sqlite>>>>,
+}
+
+pub type Transaction<'a> = sqlx::Transaction<'a, Sqlite>;
+
+impl Database {
+    pub async fn try_new(config: &DatabaseConfig) -> Result<Self> {
+        tokio::fs::create_dir_all(&config.databases_directory_path)
+            .await
+            .with_context(|| {
+                format!(
+                    "Failed to create databases directory: {}",
+                    config.databases_directory_path.to_string_lossy()
+                )
+            })?;
+
+        let mut connection_pools = std::collections::HashMap::new();
+
+        let mut entries = tokio::fs::read_dir(&config.databases_directory_path).await?;
+        while let Some(entry) = entries.next_entry().await? {
+            if !entry.file_name().to_string_lossy().ends_with(".sqlite") {
+                continue;
+            }
+
+            let vault: VaultId = entry
+                .file_name()
+                .to_string_lossy()
+                .trim_end_matches(".sqlite")
+                .to_owned();
+
+            connection_pools.insert(
+                vault.clone(),
+                Self::create_vault_database(config, &vault).await?,
+            );
+        }
+
+        Ok(Self {
+            config: config.clone(),
+            connection_pools: Arc::new(Mutex::new(connection_pools)),
+        })
+    }
+
+    async fn create_vault_database(
+        config: &DatabaseConfig,
+        vault: &VaultId,
+    ) -> Result<Pool<Sqlite>> {
+        let file_name = config
+            .databases_directory_path
+            .join(format!("{vault}.sqlite"));
+
+        let connection_options = SqliteConnectOptions::new()
+            .filename(file_name.clone())
+            .create_if_missing(true)
+            .busy_timeout(Duration::from_secs(3600))
+            .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal);
+
+        let pool = SqlitePoolOptions::new()
+            .max_connections(config.max_connections)
+            .test_before_acquire(true)
+            .connect_with(connection_options)
+            .await
+            .with_context(|| format!("Cannot open database at {}", file_name.display()))?;
+
+        Self::run_migrations(&pool).await?;
+
+        Ok(pool)
+    }
+
+    async fn run_migrations(pool: &Pool<Sqlite>) -> Result<()> {
+        sqlx::migrate!("src/app_state/database/migrations")
+            .run(pool)
+            .await
+            .context("Cannot check for pending migrations")
+    }
+
+    async fn get_connection_pool(&self, vault: &VaultId) -> Result<Pool<Sqlite>> {
+        let mut pools = self.connection_pools.lock().await;
+        if !pools.contains_key(vault) {
+            let pool = Self::create_vault_database(&self.config, vault).await?;
+            pools.insert(vault.clone(), pool);
+        }
+
+        let pool = pools
+            .get(vault)
+            .expect("Pool was just inserted or already exists");
+
+        Ok(pool.clone())
+    }
+
+    /// Attempting to write from this transaction might result in a
+    /// database locked error. Use this transaction for read-only operations.
+    pub async fn create_readonly_transaction(
+        &self,
+        vault: &VaultId,
+    ) -> Result<Transaction<'static>> {
+        self.get_connection_pool(vault)
+            .await?
+            .begin()
+            .await
+            .context("Cannot create transaction")
+    }
+
+    pub async fn create_write_transaction(&self, vault: &VaultId) -> Result<Transaction<'static>> {
+        let mut transaction = self.create_readonly_transaction(vault).await?;
+
+        // sqlx doesn't support immediate transactions for sqlite: https://github.com/launchbadge/sqlx/issues/481
+        sqlx::query!("END; BEGIN IMMEDIATE;")
+            .execute(&mut *transaction)
+            .await?;
+
+        Ok(transaction)
+    }
+
+    /// Return the latest state of all documents in the vault
+    pub async fn get_latest_documents(
+        &self,
+        vault: &VaultId,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<Vec<DocumentVersionWithoutContent>> {
+        let query = sqlx::query_as!(
+            DocumentVersionWithoutContent,
+            r#"
+            select 
+                vault_update_id,
+                document_id as "document_id: Hyphenated", 
+                relative_path,
+                updated_date as "updated_date: chrono::DateTime<Utc>",
+                is_deleted
+            from latest_document_versions
+            order by vault_update_id desc
+            "#,
+        );
+
+        if let Some(transaction) = transaction {
+            query.fetch_all(&mut **transaction).await
+        } else {
+            query
+                .fetch_all(&self.get_connection_pool(vault).await?)
+                .await
+        }
+        .context("Cannot fetch latest documents")
+    }
+
+    /// Return the latest state of all documents (including deleted) in the
+    /// vault which have changed since the given update id
+    pub async fn get_latest_documents_since(
+        &self,
+        vault: &VaultId,
+        vault_update_id: VaultUpdateId,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<Vec<DocumentVersionWithoutContent>> {
+        let query = sqlx::query_as!(
+            DocumentVersionWithoutContent,
+            r#"
+            select
+                vault_update_id,
+                document_id as "document_id: Hyphenated",
+                relative_path,
+                updated_date as "updated_date: chrono::DateTime<Utc>",
+                is_deleted
+            from latest_document_versions
+            where vault_update_id > ?
+            order by vault_update_id desc
+            "#,
+            vault_update_id
+        );
+
+        if let Some(transaction) = transaction {
+            query.fetch_all(&mut **transaction).await
+        } else {
+            query
+                .fetch_all(&self.get_connection_pool(vault).await?)
+                .await
+        }
+        .with_context(|| {
+            format!("Cannot fetch latest documents since vault_update_id {vault_update_id}")
+        })
+    }
+
+    pub async fn get_max_update_id_in_vault(
+        &self,
+        vault: &VaultId,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<i64> {
+        let query = sqlx::query!(
+            r#"
+            select coalesce(max(vault_update_id), 0) as max_vault_update_id
+            from documents
+            "#,
+        );
+
+        if let Some(transaction) = transaction {
+            query.fetch_one(&mut **transaction).await
+        } else {
+            query
+                .fetch_one(&self.get_connection_pool(vault).await?)
+                .await
+        }
+        .map(|row| row.max_vault_update_id)
+        .context("Cannot fetch max update id in vault")
+    }
+
+    pub async fn get_latest_document_by_path(
+        &self,
+        vault: &VaultId,
+        relative_path: &str,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<Option<StoredDocumentVersion>> {
+        let query = sqlx::query_as!(
+            StoredDocumentVersion,
+            r#"
+            select 
+                vault_update_id,
+                document_id as "document_id: Hyphenated", 
+                relative_path,
+                updated_date as "updated_date: chrono::DateTime<Utc>",
+                content,
+                is_deleted
+            from latest_document_versions
+            where relative_path = ?
+            order by vault_update_id desc  -- `latest_document_versions` only contains a single latest version of each document, however,
+                                           -- multiple documents can have the same `relative_path`, if they have been deleted. That's
+                                           -- why we only care about the latest version of the document with the given relative path.
+            limit 1
+            "#,
+            relative_path
+        );
+
+        if let Some(transaction) = transaction {
+            query.fetch_optional(&mut **transaction).await
+        } else {
+            query
+                .fetch_optional(&self.get_connection_pool(vault).await?)
+                .await
+        }
+        .context("Cannot fetch latest document version")
+    }
+
+    pub async fn get_latest_document(
+        &self,
+        vault: &VaultId,
+        document_id: &DocumentId,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<Option<StoredDocumentVersion>> {
+        let document_id = document_id.as_hyphenated();
+        let query = sqlx::query_as!(
+            StoredDocumentVersion,
+            r#"
+            select 
+                vault_update_id,
+                document_id as "document_id: Hyphenated", 
+                relative_path,
+                updated_date as "updated_date: chrono::DateTime<Utc>",
+                content,
+                is_deleted
+            from latest_document_versions
+            where document_id = ?
+            "#,
+            document_id
+        );
+
+        if let Some(transaction) = transaction {
+            query.fetch_optional(&mut **transaction).await
+        } else {
+            query
+                .fetch_optional(&self.get_connection_pool(vault).await?)
+                .await
+        }
+        .context("Cannot fetch latest document version")
+    }
+
+    pub async fn get_document_version(
+        &self,
+        vault: &VaultId,
+        vault_update_id: VaultUpdateId,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<Option<StoredDocumentVersion>> {
+        let query = sqlx::query_as!(
+            StoredDocumentVersion,
+            r#"
+            select 
+                vault_update_id,
+                document_id as "document_id: Hyphenated", 
+                relative_path,
+                updated_date as "updated_date: chrono::DateTime<Utc>",
+                content,
+                is_deleted
+            from documents
+            where vault_update_id = ?"#,
+            vault_update_id
+        );
+
+        if let Some(transaction) = transaction {
+            query.fetch_optional(&mut **transaction).await
+        } else {
+            query
+                .fetch_optional(&self.get_connection_pool(vault).await?)
+                .await
+        }
+        .context("Cannot fetch document version")
+    }
+
+    pub async fn insert_document_version(
+        &self,
+        vault: &VaultId,
+        version: &StoredDocumentVersion,
+        transaction: Option<&mut Transaction<'_>>,
+    ) -> Result<()> {
+        let document_id = version.document_id.as_hyphenated();
+        let query = sqlx::query!(
+            r#"
+            insert into documents (
+                vault_update_id,
+                document_id, 
+                relative_path,
+                updated_date,
+                content,
+                is_deleted
+            )
+            values (?, ?, ?, ?, ?, ?)
+            "#,
+            version.vault_update_id,
+            document_id,
+            version.relative_path,
+            version.updated_date,
+            version.content,
+            version.is_deleted
+        );
+
+        if let Some(transaction) = transaction {
+            query.execute(&mut **transaction).await
+        } else {
+            query.execute(&self.get_connection_pool(vault).await?).await
+        }
+        .context("Cannot insert document version")?;
+
+        Ok(())
+    }
+}

backend/sync_server/src/app_state/database/migrations/20241207143519_bootstrap.sql

@@ -0,0 +1,21 @@
+CREATE TABLE IF NOT EXISTS documents (
+    vault_update_id INTEGER NOT NULL PRIMARY KEY,
+    document_id TEXT NOT NULL,
+    relative_path TEXT NOT NULL,
+    updated_date TIMESTAMP NOT NULL,
+    content BLOB NOT NULL,
+    is_deleted BOOLEAN NOT NULL
+);
+
+CREATE VIEW IF NOT EXISTS latest_document_versions AS
+SELECT d.*
+FROM documents d
+INNER JOIN (
+    SELECT MAX(vault_update_id) AS max_version_id
+    FROM documents
+    GROUP BY document_id
+) max_versions
+ON d.vault_update_id = max_versions.max_version_id;
+
+CREATE INDEX IF NOT EXISTS idx_documents_vault_id_relative_path
+ON documents (relative_path);

backend/sync_server/src/app_state/database/models.rs

@@ -0,0 +1,68 @@
+use chrono::{DateTime, Utc};
+use schemars::JsonSchema;
+use serde::Serialize;
+use sync_lib::bytes_to_base64;
+
+pub type VaultId = String;
+pub type VaultUpdateId = i64;
+pub type DocumentId = uuid::Uuid;
+
+#[derive(Debug, Clone)]
+pub struct StoredDocumentVersion {
+    pub vault_update_id: VaultUpdateId,
+    pub document_id: DocumentId,
+    pub relative_path: String,
+    pub updated_date: DateTime<Utc>,
+    pub content: Vec<u8>,
+    pub is_deleted: bool,
+}
+
+impl PartialEq<Self> for StoredDocumentVersion {
+    fn eq(&self, other: &Self) -> bool { self.vault_update_id == other.vault_update_id }
+}
+
+#[derive(Debug, Clone, Serialize, JsonSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct DocumentVersionWithoutContent {
+    pub vault_update_id: VaultUpdateId,
+    pub document_id: DocumentId,
+    pub relative_path: String,
+    pub updated_date: DateTime<Utc>,
+    pub is_deleted: bool,
+}
+
+impl From<StoredDocumentVersion> for DocumentVersionWithoutContent {
+    fn from(value: StoredDocumentVersion) -> Self {
+        Self {
+            vault_update_id: value.vault_update_id,
+            document_id: value.document_id,
+            relative_path: value.relative_path,
+            updated_date: value.updated_date,
+            is_deleted: value.is_deleted,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, JsonSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct DocumentVersion {
+    pub vault_update_id: VaultUpdateId,
+    pub document_id: DocumentId,
+    pub relative_path: String,
+    pub updated_date: DateTime<Utc>,
+    pub content_base64: String,
+    pub is_deleted: bool,
+}
+
+impl From<StoredDocumentVersion> for DocumentVersion {
+    fn from(value: StoredDocumentVersion) -> Self {
+        Self {
+            vault_update_id: value.vault_update_id,
+            document_id: value.document_id,
+            relative_path: value.relative_path,
+            updated_date: value.updated_date,
+            content_base64: bytes_to_base64(&value.content),
+            is_deleted: value.is_deleted,
+        }
+    }
+}

backend/sync_server/src/config/server_config.rs

@@ -1,7 +1,10 @@
 use log::debug;
 use serde::{Deserialize, Serialize};
 
-use crate::consts::{DEFAULT_HOST, DEFAULT_MAX_BODY_SIZE_MB, DEFAULT_PORT};
+use crate::consts::{
+    DEFAULT_HOST, DEFAULT_MAX_BODY_SIZE_MB, DEFAULT_MAX_CLIENTS_PER_VAULT, DEFAULT_PORT,
+};
+
 #[derive(Debug, Deserialize, Serialize, Clone)]
 pub struct ServerConfig {
     #[serde(default = "default_host")]
@@ -12,6 +15,9 @@ pub struct ServerConfig {
 
     #[serde(default = "default_max_body_size_mb")]
     pub max_body_size_mb: usize,
+
+    #[serde(default = "default_max_clients_per_vault")]
+    pub max_clients_per_vault: usize,
 }
 
 fn default_host() -> String {
@@ -29,12 +35,18 @@ fn default_max_body_size_mb() -> usize {
     DEFAULT_MAX_BODY_SIZE_MB
 }
 
+fn default_max_clients_per_vault() -> usize {
+    debug!("Using default max clients per vault: {DEFAULT_MAX_CLIENTS_PER_VAULT}");
+    DEFAULT_MAX_CLIENTS_PER_VAULT
+}
+
 impl Default for ServerConfig {
     fn default() -> Self {
         Self {
             host: default_host(),
             port: default_port(),
             max_body_size_mb: default_max_body_size_mb(),
+            max_clients_per_vault: default_max_clients_per_vault(),
         }
     }
 }

backend/sync_server/src/consts.rs

@@ -4,3 +4,4 @@ pub const DEFAULT_HOST: &str = "127.0.0.1";
 pub const DEFAULT_PORT: u16 = 3000;
 pub const DEFAULT_MAX_CONNECTIONS: u32 = 12;
 pub const DEFAULT_MAX_BODY_SIZE_MB: usize = 4096;
+pub const DEFAULT_MAX_CLIENTS_PER_VAULT: usize = 256;