Merge branch 'master' into jsdt/store-client-creds

Author: jsdt

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)]
@@ -1121,10 +1010,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

@@ -35,6 +35,7 @@ use spacetimedb_datastore::execution_context::{ExecutionContext, ReducerContext,
 use spacetimedb_datastore::locking_tx_datastore::MutTxId;
 use spacetimedb_datastore::system_tables::{ST_CLIENT_ID, ST_CONNECTION_CREDENTIALS_ID};
 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};
@@ -1263,6 +1264,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()
     }

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;

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

@@ -8,18 +8,21 @@ import {
 import { stdbLogger } from './logger.ts';
 import type { ComparablePrimitive } from 'spacetimedb';
 
-export type Operation = {
+export type Operation<
+  RowType extends Record<string, any> = Record<string, any>,
+> = {
   type: 'insert' | 'delete';
   // For tables with a primary key, this is the primary key value, as a primitive or string.
   // Otherwise, it is an encoding of the full row.
   rowId: ComparablePrimitive;
-  // TODO: Refine this type to at least reflect that it is a product.
-  row: any;
+  row: RowType;
 };
 
-export type TableUpdate = {
+export type TableUpdate<
+  RowType extends Record<string, any> = Record<string, any>,
+> = {
   tableName: string;
-  operations: Operation[];
+  operations: Operation<RowType>[];
 };
 
 export type PendingCallback = {
@@ -30,7 +33,9 @@ export type PendingCallback = {
 /**
  * Builder to generate calls to query a `table` in the database
  */
-export class TableCache<RowType = any> {
+export class TableCache<
+  RowType extends Record<string, any> = Record<string, any>,
+> {
   private rows: Map<ComparablePrimitive, [RowType, number]>;
   private tableTypeInfo: TableRuntimeTypeInfo;
   private emitter: EventEmitter<'insert' | 'delete' | 'update'>;
@@ -57,18 +62,24 @@ export class TableCache<RowType = any> {
   /**
    * @returns The values of the rows in the table
    */
-  iter(): any[] {
+  iter(): RowType[] {
     return Array.from(this.rows.values()).map(([row]) => row);
   }
 
   applyOperations = (
-    operations: Operation[],
+    operations: Operation<RowType>[],
     ctx: EventContextInterface
   ): PendingCallback[] => {
     const pendingCallbacks: PendingCallback[] = [];
     if (this.tableTypeInfo.primaryKeyInfo !== undefined) {
-      const insertMap = new Map<ComparablePrimitive, [Operation, number]>();
-      const deleteMap = new Map<ComparablePrimitive, [Operation, number]>();
+      const insertMap = new Map<
+        ComparablePrimitive,
+        [Operation<RowType>, number]
+      >();
+      const deleteMap = new Map<
+        ComparablePrimitive,
+        [Operation<RowType>, number]
+      >();
       for (const op of operations) {
         if (op.type === 'insert') {
           const [_, prevCount] = insertMap.get(op.rowId) || [op, 0];
@@ -177,7 +188,7 @@ export class TableCache<RowType = any> {
 
   insert = (
     ctx: EventContextInterface,
-    operation: Operation,
+    operation: Operation<RowType>,
     count: number = 1
   ): PendingCallback | undefined => {
     const [_, previousCount] = this.rows.get(operation.rowId) || [
@@ -200,7 +211,7 @@ export class TableCache<RowType = any> {
 
   delete = (
     ctx: EventContextInterface,
-    operation: Operation,
+    operation: Operation<RowType>,
     count: number = 1
   ): PendingCallback | undefined => {
     const [_, previousCount] = this.rows.get(operation.rowId) || [

smoketests/__init__.py

@@ -216,7 +216,7 @@ def log_records(self, n):
         logs = self.spacetime("logs", "--format=json", "-n", str(n), "--", self.database_identity)
         return list(map(json.loads, logs.splitlines()))
 
-    def publish_module(self, domain=None, *, clear=True, capture_stderr=True):
+    def publish_module(self, domain=None, *, clear=True, capture_stderr=True, num_replicas=None):
         print("publishing module", self.publish_module)
         publish_output = self.spacetime(
             "publish",
@@ -227,10 +227,11 @@ def publish_module(self, domain=None, *, clear=True, capture_stderr=True):
             # because the server address is `node` which doesn't look like `localhost` or `127.0.0.1`
             # and so the publish step prompts for confirmation.
             "--yes",
+            *["--num-replicas", f"{num_replicas}"] if num_replicas is not None else [],
             capture_stderr=capture_stderr,
         )
         self.resolved_identity = re.search(r"identity: ([0-9a-fA-F]+)", publish_output)[1]
-        self.database_identity = domain if domain is not None else self.resolved_identity
+        self.database_identity = self.resolved_identity
 
     @classmethod
     def reset_config(cls):