fix(client): Better throttle snapshot stop strategy

.changeset/red-insects-provide.md

@@ -0,0 +1,5 @@
+---
+"electric-sql": patch
+---
+
+Better throttle snapshot stop strategy

clients/typescript/src/drivers/node-postgres/database.ts

@@ -58,10 +58,12 @@ export async function createEmbeddedPostgres(
   const db = pg.getPgClient()
   await db.connect()
 
+  let stopPromise: Promise<void>
+
   // We use the database directory as the name
   // because it uniquely identifies the DB
   return {
     db,
-    stop: () => pg.stop(),
+    stop: () => (stopPromise ??= pg.stop()),
   }
 }

clients/typescript/src/satellite/mock.ts

@@ -208,6 +208,8 @@ export class MockSatelliteClient
   transactionsCb?: TransactionCallback
   additionalDataCb?: AdditionalDataCallback
 
+  outboundStartedCallback?: OutboundStartedCallback
+
   relationData: Record<string, DataRecord[]> = {}
 
   deliverFirst = false
@@ -423,15 +425,15 @@ export class MockSatelliteClient
   }
 
   unsubscribeToTransactions(): void {
-    throw new Error('Method not implemented.')
+    this.transactionsCb = undefined
   }
 
   subscribeToAdditionalData(callback: AdditionalDataCallback): void {
     this.additionalDataCb = callback
   }
 
   unsubscribeToAdditionalData(_cb: AdditionalDataCallback): void {
-    throw new Error('Method not implemented.')
+    this.additionalDataCb = undefined
   }
 
   enqueueTransaction(transaction: DataTransaction): void {
@@ -448,10 +450,13 @@ export class MockSatelliteClient
 
   subscribeToOutboundStarted(callback: OutboundStartedCallback): void {
     this.on('outbound_started', callback)
+    this.outboundStartedCallback = callback
   }
 
   unsubscribeToOutboundStarted(): void {
-    throw new Error('Method not implemented.')
+    if (!this.outboundStartedCallback) return
+    this.removeListener('outbound_started', this.outboundStartedCallback)
+    this.outboundStartedCallback = undefined
   }
 
   sendErrorAfterTimeout(subscriptionId: string, timeout: number): void {

clients/typescript/src/satellite/process.ts

@@ -39,6 +39,7 @@ import {
   Uuid,
   DbRecord as DataRecord,
   ReplicatedRowTransformer,
+  ServerTransaction,
 } from '../util/types'
 import { SatelliteOpts } from './config'
 import { Client, Satellite } from './index'
@@ -155,6 +156,8 @@ export class SatelliteProcess implements Satellite {
   private _connectRetryHandler: ConnectRetryHandler
   private initializing?: Waiter
 
+  private _removeClientListeners?: () => void
+
   constructor(
     dbName: DbName,
     adapter: DatabaseAdapter,
@@ -191,8 +194,6 @@ export class SatelliteProcess implements Satellite {
     this.shapeRequestIdGenerator = this.subscriptionIdGenerator
 
     this._connectRetryHandler = connectRetryHandler
-
-    this.setClientListeners()
   }
 
   /**
@@ -212,6 +213,8 @@ export class SatelliteProcess implements Satellite {
       await this.logDatabaseVersion()
     }
 
+    this.setClientListeners()
+
     await this.migrator.up()
 
     const isVerified = await this._verifyTableStructure()
@@ -251,6 +254,7 @@ export class SatelliteProcess implements Satellite {
       this.notifier.subscribeToPotentialDataChanges(this._throttledSnapshot)
 
     // Start polling to request a snapshot every `pollingInterval` ms.
+    clearInterval(this._pollingInterval)
     this._pollingInterval = setInterval(
       this._throttledSnapshot,
       this.opts.pollingInterval
@@ -314,17 +318,54 @@ export class SatelliteProcess implements Satellite {
     await this.adapter.runInTransaction(...stmtsWithTriggers)
   }
 
+  // Adds all the necessary listeners to the satellite client
+  // They can be cleared up by calling the function `_removeClientListeners`
   setClientListeners(): void {
-    this.client.subscribeToError(this._handleClientError.bind(this))
-    this.client.subscribeToRelations(this._updateRelations.bind(this))
-    this.client.subscribeToTransactions(this._applyTransaction.bind(this))
-    this.client.subscribeToAdditionalData(this._applyAdditionalData.bind(this))
-    this.client.subscribeToOutboundStarted(this._throttledSnapshot.bind(this))
+    // Remove any existing listeners
+    if (this._removeClientListeners) {
+      this._removeClientListeners?.()
+      this._removeClientListeners = undefined
+    }
 
-    this.client.subscribeToSubscriptionEvents(
-      this._handleSubscriptionData.bind(this),
+    const clientErrorCallback = this._handleClientError.bind(this)
+    this.client.subscribeToError(clientErrorCallback)
+
+    const clientRelationsCallback = this._handleClientRelations.bind(this)
+    this.client.subscribeToRelations(clientRelationsCallback)
+
+    const clientTransactionsCallback = this._handleClientTransactions.bind(this)
+    this.client.subscribeToTransactions(clientTransactionsCallback)
+
+    const clientAdditionalDataCallback =
+      this._handleClientAdditionalData.bind(this)
+    this.client.subscribeToAdditionalData(clientAdditionalDataCallback)
+
+    const clientOutboundStartedCallback =
+      this._handleClientOutboundStarted.bind(this)
+    this.client.subscribeToOutboundStarted(clientOutboundStartedCallback)
+
+    const clientSubscriptionDataCallback =
+      this._handleSubscriptionData.bind(this)
+    const clientSubscriptionErrorCallback =
       this._handleSubscriptionError.bind(this)
+    this.client.subscribeToSubscriptionEvents(
+      clientSubscriptionDataCallback,
+      clientSubscriptionErrorCallback
     )
+
+    // Keep a way to remove the client listeners
+    this._removeClientListeners = () => {
+      this.client.unsubscribeToError(clientErrorCallback)
+      this.client.unsubscribeToRelations(clientRelationsCallback)
+      this.client.unsubscribeToTransactions(clientTransactionsCallback)
+      this.client.unsubscribeToAdditionalData(clientAdditionalDataCallback)
+      this.client.unsubscribeToOutboundStarted(clientOutboundStartedCallback)
+
+      this.client.unsubscribeToSubscriptionEvents(
+        clientSubscriptionDataCallback,
+        clientSubscriptionErrorCallback
+      )
+    }
   }
 
   // Unsubscribe from data changes and stop polling
@@ -333,19 +374,9 @@ export class SatelliteProcess implements Satellite {
   }
 
   private async _stop(shutdown?: boolean): Promise<void> {
-    // Stop snapshotting and polling for changes.
-    this._throttledSnapshot.cancel()
-
-    // Ensure that no snapshot is left running in the background
-    // by acquiring the mutex and releasing it immediately.
-    const releaseMutex = await this.snapshotMutex.acquire()
-    releaseMutex()
-
-    if (this._pollingInterval !== undefined) {
-      clearInterval(this._pollingInterval)
-
-      this._pollingInterval = undefined
-    }
+    // Stop snapshot polling
+    clearInterval(this._pollingInterval)
+    this._pollingInterval = undefined
 
     // Unsubscribe all listeners and remove them
     const unsubscribers = [
@@ -362,13 +393,30 @@ export class SatelliteProcess implements Satellite {
       }
     })
 
+    this._removeClientListeners?.()
+    this._removeClientListeners = undefined
+
+    // Cancel the snapshot throttle
+    this._throttledSnapshot.cancel()
+
+    // Make sure no snapshot is running after we stop the process, otherwise we might be trying to
+    // interact with a closed database connection
+    await this._waitForActiveSnapshots()
+
     this.disconnect()
 
     if (shutdown) {
       this.client.shutdown()
     }
   }
 
+  // Ensure that no snapshot is left running in the background
+  // by acquiring the mutex and releasing it immediately.
+  async _waitForActiveSnapshots(): Promise<void> {
+    const releaseMutex = await this.snapshotMutex.acquire()
+    releaseMutex()
+  }
+
   async subscribe(shapeDefinitions: Shape[]): Promise<ShapeSubscription> {
     // Await for client to be ready before doing anything else
     await this.initializing?.waitOn()
@@ -674,6 +722,22 @@ export class SatelliteProcess implements Satellite {
     }
   }
 
+  _handleClientRelations(relation: Relation): void {
+    this._updateRelations(relation)
+  }
+
+  async _handleClientTransactions(tx: ServerTransaction) {
+    await this._applyTransaction(tx)
+  }
+
+  async _handleClientAdditionalData(data: AdditionalData) {
+    await this._applyAdditionalData(data)
+  }
+
+  async _handleClientOutboundStarted() {
+    await this._throttledSnapshot()
+  }
+
   // handles async client errors: can be a socket error or a server error message
   _handleClientError(satelliteError: SatelliteError) {
     if (this.initializing && !this.initializing.finished()) {

clients/typescript/test/satellite/common.ts

@@ -311,9 +311,14 @@ export const makeContext = async (
   await mkdir('.tmp', { recursive: true })
   const dbName = `.tmp/test-${randomValue()}.db`
   const db = new SqliteDatabase(dbName)
+  const stop = async () => {
+    if (!db.open) return
+    db.close()
+  }
   const adapter = new SqliteDatabaseAdapter(db)
   const migrator = new SqliteBundleMigrator(adapter, sqliteMigrations)
   makeContextInternal(t, dbName, adapter, migrator, namespace, options)
+  t.context.stop = stop
 }
 
 export const makePgContext = async (
@@ -337,7 +342,10 @@ export const makePgliteContext = async (
 ) => {
   const dbName = `test-${randomValue()}`
   const db = new PGlite()
-  const stop = () => db.close()
+  const stop = async () => {
+    if (db.closed) return
+    await db.close()
+  }
   const adapter = new PgliteDatabaseAdapter(db)
   const migrator = new PgBundleMigrator(adapter, pgMigrations)
   makeContextInternal(t, dbName, adapter, migrator, namespace, options)
@@ -380,8 +388,12 @@ export const mockElectricClient = async (
   return electric
 }
 
-export const clean = async (t: ExecutionContext<{ dbName: string }>) => {
-  const { dbName } = t.context
+export const cleanAndStopDb = async (
+  t: ExecutionContext<{ dbName: string; stop?: () => Promise<void> }>
+) => {
+  const { dbName, stop } = t.context
+
+  await stop?.()
 
   await removeFile(dbName, { force: true })
   await removeFile(`${dbName}-journal`, { force: true })
@@ -396,8 +408,7 @@ export const cleanAndStopSatellite = async (
 ) => {
   const { satellite } = t.context
   await satellite.stop()
-  await clean(t)
-  await t.context.stop?.()
+  await cleanAndStopDb(t)
 }
 
 export async function migrateDb(