Merge branch 'master' into centril/clear-table-and-fix-drop-table

Author: Shubham8287

crates/core/src/db/relational_db.rs

@@ -1621,13 +1621,22 @@ pub async fn local_durability(commitlog_dir: CommitLogDir) -> io::Result<(LocalD
 /// [DurabilityProvider]: crate::host::host_controller::DurabilityProvider
 pub async fn snapshot_watching_commitlog_compressor(
     mut snapshot_rx: watch::Receiver<u64>,
+    mut clog_tx: Option<tokio::sync::mpsc::Sender<u64>>,
+    mut snap_tx: Option<tokio::sync::mpsc::Sender<u64>>,
     durability: LocalDurability,
 ) {
     let mut prev_snapshot_offset = *snapshot_rx.borrow_and_update();
     while snapshot_rx.changed().await.is_ok() {
         let snapshot_offset = *snapshot_rx.borrow_and_update();
         let durability = durability.clone();
-        let res = asyncify(move || {
+
+        if let Some(snap_tx) = &mut snap_tx {
+            if let Err(err) = snap_tx.try_send(snapshot_offset) {
+                tracing::warn!("failed to send offset {snapshot_offset} after snapshot creation: {err}");
+            }
+        }
+
+        let res: io::Result<_> = asyncify(move || {
             let segment_offsets = durability.existing_segment_offsets()?;
             let start_idx = segment_offsets
                 .binary_search(&prev_snapshot_offset)
@@ -1641,15 +1650,27 @@ pub async fn snapshot_watching_commitlog_compressor(
             // in this case, segment_offsets[end_idx] is the segment that contains the snapshot,
             // which we don't want to compress, so an exclusive range is correct.
             let segment_offsets = &segment_offsets[..end_idx];
-            durability.compress_segments(segment_offsets)
+            durability.compress_segments(segment_offsets)?;
+            let n = segment_offsets.len();
+            let last_compressed_segment = if n > 0 { Some(segment_offsets[n - 1]) } else { None };
+            Ok(last_compressed_segment)
         })
         .await;
 
-        if let Err(e) = res {
-            tracing::warn!("failed to compress segments: {e}");
-            continue;
-        }
+        let last_compressed_segment = match res {
+            Ok(opt_offset) => opt_offset,
+            Err(err) => {
+                tracing::warn!("failed to compress segments: {err}");
+                continue;
+            }
+        };
         prev_snapshot_offset = snapshot_offset;
+
+        if let Some((clog_tx, last_compressed_segment)) = clog_tx.as_mut().zip(last_compressed_segment) {
+            if let Err(err) = clog_tx.try_send(last_compressed_segment) {
+                tracing::warn!("failed to send offset {last_compressed_segment} after compression: {err}");
+            }
+        }
     }
 }
 

crates/core/src/host/host_controller.rs

@@ -16,7 +16,7 @@ use crate::subscription::module_subscription_manager::{spawn_send_worker, Subscr
 use crate::util::asyncify;
 use crate::util::jobs::{JobCore, JobCores};
 use crate::worker_metrics::WORKER_METRICS;
-use anyhow::{anyhow, ensure, Context};
+use anyhow::{anyhow, Context};
 use async_trait::async_trait;
 use durability::{Durability, EmptyHistory};
 use log::{info, trace, warn};
@@ -448,117 +448,6 @@ impl HostController {
         host.migrate_plan(host_type, program, style).await
     }
 
-    /// Start the host `replica_id` and conditionally update it.
-    ///
-    /// If the host was not initialized before, it is initialized with the
-    /// program [`Database::initial_program`], which is loaded from the
-    /// controller's [`ProgramStorage`].
-    ///
-    /// If it was already initialized and its stored program hash matches
-    /// [`Database::initial_program`], no further action is taken.
-    ///
-    /// Otherwise, if `expected_hash` is `Some` and does **not** match the
-    /// stored hash, an error is returned.
-    ///
-    /// Otherwise, the host is updated to [`Database::initial_program`], loading
-    /// the program data from the controller's [`ProgramStorage`].
-    ///
-    /// > Note that this ascribes different semantics to [`Database::initial_program`]
-    /// > than elsewhere, where the [`Database`] value is provided by the control
-    /// > database. The method is mainly useful for bootstrapping the control
-    /// > database itself.
-    pub async fn init_maybe_update_module_host(
-        &self,
-        database: Database,
-        replica_id: u64,
-        expected_hash: Option<Hash>,
-    ) -> anyhow::Result<watch::Receiver<ModuleHost>> {
-        trace!("custom bootstrap {}/{}", database.database_identity, replica_id);
-
-        let db_addr = database.database_identity;
-        let host_type = database.host_type;
-        let program_hash = database.initial_program;
-
-        let mut guard = self.acquire_write_lock(replica_id).await;
-
-        // `HostController::clone` is fast,
-        // as all of its fields are either `Copy` or wrapped in `Arc`.
-        let this = self.clone();
-
-        // `try_init_host` is not cancel safe, as it will spawn other async tasks
-        // which hold a filesystem lock past when `try_init_host` returns or is cancelled.
-        // This means that, if `try_init_host` is cancelled, subsequent calls will fail.
-        //
-        // The rest of this future is also not cancel safe, as it will `Option::take` out of the guard
-        // at the start of the block and then store back into it at the end.
-        //
-        // This is problematic because Axum will cancel its handler tasks if the client disconnects,
-        // and this method is called from Axum handlers, e.g. for the publish route.
-        // `tokio::spawn` a task to update the contents of `guard`,
-        // so that it will run to completion even if the caller goes away.
-        //
-        // Note that `tokio::spawn` only cancels its tasks when the runtime shuts down,
-        // at which point we won't be calling `try_init_host` again anyways.
-        let module = tokio::spawn(async move {
-            let mut host = match guard.take() {
-                Some(host) => host,
-                None => this.try_init_host(database, replica_id).await?,
-            };
-            let module = host.module.subscribe();
-
-            // The program is now either:
-            //
-            // - the desired one from [Database], in which case we do nothing
-            // - `Some` expected hash, in which case we update to the desired one
-            // - `None` expected hash, in which case we also update
-            let stored_hash = stored_program_hash(host.db())?
-                .with_context(|| format!("[{db_addr}] database improperly initialized"))?;
-            if stored_hash == program_hash {
-                info!("[{db_addr}] database up-to-date with {program_hash}");
-                *guard = Some(host);
-            } else {
-                if let Some(expected_hash) = expected_hash {
-                    ensure!(
-                        expected_hash == stored_hash,
-                        "[{}] expected program {} found {}",
-                        db_addr,
-                        expected_hash,
-                        stored_hash
-                    );
-                }
-                info!("[{db_addr}] updating database from `{stored_hash}` to `{program_hash}`");
-                let program = load_program(&this.program_storage, program_hash).await?;
-                let update_result = host
-                    .update_module(
-                        this.runtimes.clone(),
-                        host_type,
-                        program,
-                        MigrationPolicy::Compatible,
-                        this.energy_monitor.clone(),
-                        this.unregister_fn(replica_id),
-                        this.db_cores.take(),
-                    )
-                    .await?;
-                match update_result {
-                    UpdateDatabaseResult::NoUpdateNeeded | UpdateDatabaseResult::UpdatePerformed => {
-                        *guard = Some(host);
-                    }
-                    UpdateDatabaseResult::AutoMigrateError(e) => {
-                        return Err(anyhow::anyhow!(e));
-                    }
-                    UpdateDatabaseResult::ErrorExecutingMigration(e) => {
-                        return Err(e);
-                    }
-                }
-            }
-
-            Ok::<_, anyhow::Error>(module)
-        })
-        .await??;
-
-        Ok(module)
-    }
-
     /// Release all resources of the [`ModuleHost`] identified by `replica_id`,
     /// and deregister it from the controller.
     #[tracing::instrument(level = "trace", skip_all)]
@@ -1115,10 +1004,6 @@ impl Host {
 
         Ok(res)
     }
-
-    fn db(&self) -> &RelationalDB {
-        &self.replica_ctx.relational_db
-    }
 }
 
 impl Drop for Host {

crates/core/src/host/module_host.rs

@@ -32,6 +32,7 @@ use spacetimedb_data_structures::map::{HashCollectionExt as _, IntMap};
 use spacetimedb_datastore::execution_context::{ExecutionContext, ReducerContext, Workload, WorkloadType};
 use spacetimedb_datastore::locking_tx_datastore::MutTxId;
 use spacetimedb_datastore::traits::{IsolationLevel, Program, TxData};
+use spacetimedb_durability::DurableOffset;
 use spacetimedb_execution::pipelined::PipelinedProject;
 use spacetimedb_lib::db::raw_def::v9::Lifecycle;
 use spacetimedb_lib::identity::{AuthCtx, RequestId};
@@ -1233,6 +1234,10 @@ impl ModuleHost {
         &self.replica_ctx().database
     }
 
+    pub fn durable_tx_offset(&self) -> Option<DurableOffset> {
+        self.replica_ctx().relational_db.durable_tx_offset()
+    }
+
     pub(crate) fn replica_ctx(&self) -> &ReplicaContext {
         self.module.replica_ctx()
     }

crates/standalone/src/lib.rs

@@ -127,6 +127,8 @@ impl DurabilityProvider for StandaloneDurabilityProvider {
             |snapshot_rx| {
                 tokio::spawn(relational_db::snapshot_watching_commitlog_compressor(
                     snapshot_rx,
+                    None,
+                    None,
                     durability,
                 ));
             }

sdks/typescript/packages/sdk/src/client_cache.ts

@@ -5,7 +5,7 @@ export class ClientCache {
   /**
    * The tables in the database.
    */
-  tables: Map<string, TableCache>;
+  tables: Map<string, TableCache<any>>;
 
   constructor() {
     this.tables = new Map();
@@ -16,7 +16,9 @@ export class ClientCache {
    * @param name The name of the table.
    * @returns The table
    */
-  getTable(name: string): TableCache {
+  getTable<RowType extends Record<string, any>>(
+    name: string
+  ): TableCache<RowType> {
     const table = this.tables.get(name);
 
     // ! This should not happen as the table should be available but an exception is thrown just in case.
@@ -30,10 +32,10 @@ export class ClientCache {
     return table;
   }
 
-  getOrCreateTable<RowType>(
+  getOrCreateTable<RowType extends Record<string, any>>(
     tableTypeInfo: TableRuntimeTypeInfo
   ): TableCache<RowType> {
-    let table: TableCache;
+    let table: TableCache<RowType>;
     if (!this.tables.has(tableTypeInfo.tableName)) {
       table = new TableCache<RowType>(tableTypeInfo);
       this.tables.set(tableTypeInfo.tableName, table);

sdks/typescript/packages/sdk/src/db_connection_impl.ts

@@ -483,7 +483,9 @@ export class DbConnectionImpl<
         const parsedTableUpdates = await parseDatabaseUpdate(
           message.value.update
         );
-        const subscribeAppliedMessage: SubscribeAppliedMessage = {
+        const subscribeAppliedMessage: SubscribeAppliedMessage<
+          Record<string, any>
+        > = {
           tag: 'SubscribeApplied',
           queryId: message.value.queryId.id,
           tableUpdates: parsedTableUpdates,
@@ -495,7 +497,9 @@ export class DbConnectionImpl<
         const parsedTableUpdates = await parseDatabaseUpdate(
           message.value.update
         );
-        const unsubscribeAppliedMessage: UnsubscribeAppliedMessage = {
+        const unsubscribeAppliedMessage: UnsubscribeAppliedMessage<
+          Record<string, any>
+        > = {
           tag: 'UnsubscribeApplied',
           queryId: message.value.queryId.id,
           tableUpdates: parsedTableUpdates,

sdks/typescript/packages/sdk/src/message_types.ts

@@ -4,14 +4,14 @@ import { Identity } from 'spacetimedb';
 import type { TableUpdate } from './table_cache.ts';
 import { Timestamp } from 'spacetimedb';
 
-export type InitialSubscriptionMessage = {
+export type InitialSubscriptionMessage<RowType extends Record<string, any>> = {
   tag: 'InitialSubscription';
-  tableUpdates: TableUpdate[];
+  tableUpdates: TableUpdate<RowType>[];
 };
 
-export type TransactionUpdateMessage = {
+export type TransactionUpdateMessage<RowType extends Record<string, any>> = {
   tag: 'TransactionUpdate';
-  tableUpdates: TableUpdate[];
+  tableUpdates: TableUpdate<RowType>[];
   identity: Identity;
   connectionId: ConnectionId | null;
   reducerInfo?: {
@@ -24,10 +24,11 @@ export type TransactionUpdateMessage = {
   energyConsumed: bigint;
 };
 
-export type TransactionUpdateLightMessage = {
-  tag: 'TransactionUpdateLight';
-  tableUpdates: TableUpdate[];
-};
+export type TransactionUpdateLightMessage<RowType extends Record<string, any>> =
+  {
+    tag: 'TransactionUpdateLight';
+    tableUpdates: TableUpdate<RowType>[];
+  };
 
 export type IdentityTokenMessage = {
   tag: 'IdentityToken';
@@ -36,16 +37,16 @@ export type IdentityTokenMessage = {
   connectionId: ConnectionId;
 };
 
-export type SubscribeAppliedMessage = {
+export type SubscribeAppliedMessage<RowType extends Record<string, any>> = {
   tag: 'SubscribeApplied';
   queryId: number;
-  tableUpdates: TableUpdate[];
+  tableUpdates: TableUpdate<RowType>[];
 };
 
-export type UnsubscribeAppliedMessage = {
+export type UnsubscribeAppliedMessage<RowType extends Record<string, any>> = {
   tag: 'UnsubscribeApplied';
   queryId: number;
-  tableUpdates: TableUpdate[];
+  tableUpdates: TableUpdate<RowType>[];
 };
 
 export type SubscriptionError = {
@@ -54,11 +55,12 @@ export type SubscriptionError = {
   error: string;
 };
 
-export type Message =
-  | InitialSubscriptionMessage
-  | TransactionUpdateMessage
-  | TransactionUpdateLightMessage
-  | IdentityTokenMessage
-  | SubscribeAppliedMessage
-  | UnsubscribeAppliedMessage
-  | SubscriptionError;
+export type Message<RowType extends Record<string, any> = Record<string, any>> =
+
+    | InitialSubscriptionMessage<RowType>
+    | TransactionUpdateMessage<RowType>
+    | TransactionUpdateLightMessage<RowType>
+    | IdentityTokenMessage
+    | SubscribeAppliedMessage<RowType>
+    | UnsubscribeAppliedMessage<RowType>
+    | SubscriptionError;