@powersync/lib-service-postgres 0.5.2 → 0.5.3

package/src/db/connection/ConnectionSlot.ts

@@ -1,12 +1,22 @@
 import * as framework from '@powersync/lib-services-framework';
 import * as pgwire from '@powersync/service-jpgwire';
+import * as timers from 'node:timers/promises';
+
+export type NotificationEvent =
+  | { type: 'notification'; notification: pgwire.PgNotification }
+  | { type: 'channels-registered' }
+  | { type: 'connection-error'; error: unknown };
 
 export interface NotificationListener {
-  notification?: (payload: pgwire.PgNotification) => void;
+  /**
+   * Reports notifications and notification-connection lifecycle events.
+   */
+  notificationEvent?: (event: NotificationEvent) => void;
 }
 
 export interface ConnectionSlotListener extends NotificationListener {
   connectionAvailable?: () => void;
+  /** Reports that this slot exhausted its connection attempts. */
   connectionError?: (exception: any) => void;
   connectionCreated?: (connection: pgwire.PgConnection) => Promise<void>;
 }
@@ -23,22 +33,23 @@ export type ConnectionSlotOptions = {
 };
 
 export const MAX_CONNECTION_ATTEMPTS = 5;
+const CONNECTION_RETRY_DELAY_MS = 100;
 
 export class ConnectionSlot extends framework.BaseObserver<ConnectionSlotListener> {
-  isAvailable: boolean;
   isPoking: boolean;
 
   closed: boolean;
 
   protected connection: pgwire.PgConnection | null;
   protected connectingPromise: Promise<pgwire.PgConnection> | null;
+  protected activeLease: symbol | null;
 
   constructor(protected options: ConnectionSlotOptions) {
     super();
-    this.isAvailable = false;
     this.connection = null;
     this.isPoking = false;
     this.connectingPromise = null;
+    this.activeLease = null;
     this.closed = false;
   }
 
@@ -46,21 +57,8 @@ export class ConnectionSlot extends framework.BaseObserver<ConnectionSlotListene
     return !!this.connection;
   }
 
-  protected async connect() {
-    this.connectingPromise = pgwire.connectPgWire(this.options.config, {
-      type: 'standard',
-      applicationName: this.options.applicationName
-    });
-    const connection = await this.connectingPromise;
-    this.connectingPromise = null;
-    await this.iterateAsyncListeners(async (l) => l.connectionCreated?.(connection));
-
-    /**
-     * Configure the Postgres connection to listen to notifications.
-     * Subscribing to notifications, even without a registered listener, should not add much overhead.
-     */
-    await this.configureConnectionNotifications(connection);
-    return connection;
+  get isAvailable() {
+    return this.connection != null && this.activeLease == null && !this.closed;
   }
 
   async [Symbol.asyncDispose]() {
@@ -70,84 +68,166 @@ export class ConnectionSlot extends framework.BaseObserver<ConnectionSlotListene
     super.clearListeners();
   }
 
-  protected async configureConnectionNotifications(connection: pgwire.PgConnection) {
-    connection.onnotification = this.handleNotification;
+  /**
+   * Ensure this slot has a connection and signal when it can be leased.
+   */
+  async poke() {
+    if (this.closed || this.connection || this.isPoking) {
+      return;
+    }
 
-    for (const channelName of this.options.notificationChannels ?? []) {
-      await connection.query({
-        statement: `LISTEN ${channelName}`
-      });
+    this.isPoking = true;
+    try {
+      for (let retryCounter = 0; retryCounter <= MAX_CONNECTION_ATTEMPTS; retryCounter++) {
+        try {
+          const connection = await this.connect();
+
+          if (this.closed) {
+            connection.destroy();
+            return;
+          }
+
+          this.connection = connection;
+
+          // Register this only after the slot owns the connection. If
+          // `whenDestroyed` has already settled, its callback runs in a later
+          // microtask and must see this exact connection assigned. Registering
+          // it before assignment could ignore the destruction and leave the
+          // slot holding an already-destroyed connection.
+          connection.whenDestroyed
+            .catch((error) => framework.logger.debug('Postgres connection destroyed with an error', error))
+            .then(() => this.handleConnectionDestroyed(connection));
+
+          if (this.activeLease == null) {
+            this.notifyConnectionAvailable();
+          }
+          break;
+        } catch (ex) {
+          if (retryCounter >= MAX_CONNECTION_ATTEMPTS) {
+            this.iterateListeners((cb) => {
+              cb.connectionError?.(ex);
+              cb.notificationEvent?.({ type: 'connection-error', error: ex });
+            });
+          } else {
+            await timers.setTimeout(CONNECTION_RETRY_DELAY_MS);
+            if (this.closed) {
+              return;
+            }
+          }
+        }
+      }
+    } finally {
+      this.isPoking = false;
     }
   }
 
-  protected handleNotification = (payload: pgwire.PgNotification) => {
-    if (!this.options.notificationChannels?.includes(payload.channel)) {
-      return;
+  protected async connect() {
+    // Only allow a single connect to run at-a-time
+    if (this.connectingPromise) {
+      return this.connectingPromise;
     }
-    this.iterateListeners((l) => l.notification?.(payload));
-  };
 
-  protected hasNotificationListener() {
-    return !!Object.values(this.listeners).find((l) => !!l.notification);
+    const connectInternal = async () => {
+      let connection: pgwire.PgConnection | null = null;
+      try {
+        const connected = await pgwire.connectPgWire(this.options.config, {
+          type: 'standard',
+          applicationName: this.options.applicationName
+        });
+        connection = connected;
+
+        await this.iterateAsyncListeners(async (l) => l.connectionCreated?.(connected));
+
+        /**
+         * Configure the Postgres connection to listen to notifications.
+         * Subscribing to notifications, even without a registered listener, should not add much overhead.
+         */
+        await this.configureConnectionNotifications(connected);
+
+        return connected;
+      } catch (error) {
+        connection?.destroy();
+        throw error;
+      }
+    };
+
+    this.connectingPromise = connectInternal();
+    try {
+      return await this.connectingPromise;
+    } finally {
+      this.connectingPromise = null;
+    }
   }
 
-  /**
-   * Test the connection if it can be reached.
-   */
-  async poke() {
-    if (this.isPoking || (this.isConnected && this.isAvailable == false) || this.closed) {
+  protected handleConnectionDestroyed(connection: pgwire.PgConnection) {
+    // Guard that the closed connection is actually the one in use by the slot.
+    if (this.connection != connection) {
       return;
     }
-    this.isPoking = true;
-    for (let retryCounter = 0; retryCounter <= MAX_CONNECTION_ATTEMPTS; retryCounter++) {
-      try {
-        const connection = this.connection ?? (await this.connect());
 
-        await connection.query({
-          statement: 'SELECT 1'
-        });
+    // Clear the slot reference, marking the slot as unavailable
+    this.connection = null;
 
-        if (!this.connection) {
-          this.connection = connection;
-          this.setAvailable();
-        } else if (this.isAvailable) {
-          this.iterateListeners((cb) => cb.connectionAvailable?.());
-        }
+    // Notification slots must restore their LISTEN subscriptions immediately.
+    // Other slots reconnect lazily when the next lease is requested.
+    if (!this.closed && this.options.notificationChannels?.length) {
+      this.poke().catch((error) => framework.logger.error('Failed to restore Postgres notification connection', error));
+    }
+  }
 
-        // Connection is alive and healthy
-        break;
-      } catch (ex) {
-        // Should be valid for all cases
-        this.isAvailable = false;
-        if (this.connection) {
-          this.connection.onnotification = () => {};
-          this.connection.destroy();
-          this.connection = null;
-        }
-        if (retryCounter >= MAX_CONNECTION_ATTEMPTS) {
-          this.iterateListeners((cb) => cb.connectionError?.(ex));
-        }
-      }
+  protected async configureConnectionNotifications(connection: pgwire.PgConnection) {
+    connection.onnotification = this.handleNotification;
+
+    const notificationChannels = this.options.notificationChannels ?? [];
+    for (const channelName of notificationChannels) {
+      await connection.query({
+        statement: `LISTEN ${channelName}`
+      });
+    }
+
+    if (notificationChannels.length > 0) {
+      this.iterateListeners((l) => l.notificationEvent?.({ type: 'channels-registered' }));
     }
-    this.isPoking = false;
   }
 
-  protected setAvailable() {
-    this.isAvailable = true;
+  protected handleNotification = (payload: pgwire.PgNotification) => {
+    if (!this.options.notificationChannels?.includes(payload.channel)) {
+      return;
+    }
+    this.iterateListeners((l) => l.notificationEvent?.({ type: 'notification', notification: payload }));
+  };
+
+  protected notifyConnectionAvailable() {
     this.iterateListeners((l) => l.connectionAvailable?.());
   }
 
   lock(): ConnectionLease | null {
-    if (!this.isAvailable || !this.connection || this.closed) {
+    if (!this.isAvailable || !this.connection || this.activeLease != null || this.closed) {
       return null;
     }
 
-    this.isAvailable = false;
+    // Create a unique symbol to identify this lease
+    const lease = Symbol();
+    this.activeLease = lease;
 
     return {
       connection: this.connection,
       release: () => {
-        this.setAvailable();
+        // Only release if this lease is the current active lease
+        if (this.activeLease != lease) {
+          return;
+        }
+        this.activeLease = null;
+        if (this.closed) {
+          return;
+        }
+        if (this.connection) {
+          this.notifyConnectionAvailable();
+        } else {
+          // The leased connection was destroyed. Reconnect now in case a
+          // request was queued while this slot still held the lease.
+          this.poke().catch((error) => framework.logger.error('Failed to restore Postgres connection', error));
+        }
       }
     };
   }

package/src/db/connection/DatabaseClient.ts

@@ -42,6 +42,8 @@ export class DatabaseClient extends AbstractPostgresConnection<DatabaseClientLis
 
   protected initialized: Promise<void>;
   protected queue: PromiseWithResolvers<ConnectionLease>[];
+  /** Latest exhausted connection attempt for each slot in the current attempt round. */
+  protected failedConnectionSlots: Map<ConnectionSlot, unknown>;
 
   constructor(protected options: DatabaseClientOptions) {
     super();
@@ -50,6 +52,7 @@ export class DatabaseClient extends AbstractPostgresConnection<DatabaseClientLis
       maxSize: options.config.max_pool_size,
       applicationName: options.applicationName
     });
+    this.failedConnectionSlots = new Map();
     this.connections = Array.from({ length: TRANSACTION_CONNECTION_COUNT }, (v, index) => {
       // Only listen to notifications on a single (the first) connection
       const notificationChannels = index == 0 ? options.notificationChannels : [];
@@ -59,8 +62,8 @@ export class DatabaseClient extends AbstractPostgresConnection<DatabaseClientLis
         applicationName: options.applicationName
       });
       slot.registerListener({
-        connectionAvailable: () => this.processConnectionQueue(),
-        connectionError: (ex) => this.handleConnectionError(ex),
+        connectionAvailable: () => this.handleConnectionAvailable(slot),
+        connectionError: (ex) => this.handleConnectionError(slot, ex),
         connectionCreated: (connection) => this.iterateAsyncListeners(async (l) => l.connectionCreated?.(connection))
       });
       return slot;
@@ -85,15 +88,13 @@ export class DatabaseClient extends AbstractPostgresConnection<DatabaseClientLis
 
   registerListener(listener: Partial<DatabaseClientListener>): () => void {
     let disposeNotification: (() => void) | null = null;
-    if ('notification' in listener) {
+    if ('notificationEvent' in listener) {
       // Pass this on to the first connection slot
       // It will only actively listen on the connection once a listener has been registered
-      disposeNotification = this.connections[0].registerListener({
-        notification: listener.notification
-      });
+      disposeNotification = this.connections[0].registerListener({ notificationEvent: listener.notificationEvent });
       this.pokeSlots();
 
-      delete listener['notification'];
+      delete listener.notificationEvent;
     }
 
     const superDispose = super.registerListener(listener);
@@ -202,14 +203,23 @@ export class DatabaseClient extends AbstractPostgresConnection<DatabaseClientLis
     const deferred = Promise.withResolvers<ConnectionLease>();
     this.queue.push(deferred);
 
+    // Try already-connected slots first. poke() only creates missing
+    // connections and does not emit another availability event for an existing
+    // connection, so skipping this could leave the request queued indefinitely.
+    this.processConnectionQueue();
     this.pokeSlots();
 
     return deferred.promise;
   }
 
   protected pokeSlots() {
-    // Poke the slots to check if they are alive
+    // Ensure the slots have connections and notify any queued requests.
     for (const slot of this.connections) {
+      if (!slot.isConnected && !slot.isPoking) {
+        // This slot is starting a fresh attempt, so a failure from an earlier
+        // attempt must not count against the current request.
+        this.failedConnectionSlots.delete(slot);
+      }
       // No need to await this. Errors are reported asynchronously
       slot.poke();
     }
@@ -244,12 +254,26 @@ export class DatabaseClient extends AbstractPostgresConnection<DatabaseClientLis
     }
   }
 
+  protected handleConnectionAvailable(slot: ConnectionSlot) {
+    // This slot recovered, so its earlier failure must no longer contribute to
+    // deciding whether the shared lease queue is unreachable.
+    this.failedConnectionSlots.delete(slot);
+    this.processConnectionQueue();
+  }
+
   /**
-   * Reports connection errors which might occur from bad configuration or
-   * a server which is no longer available.
-   * This fails all pending requests.
+   * Lease requests wait in one shared queue and are not assigned to a slot
+   * until that slot becomes available. A failure from one slot therefore
+   * cannot fail a particular request: another slot may still serve it.
    */
-  protected handleConnectionError(exception: any) {
+  protected handleConnectionError(slot: ConnectionSlot, exception: any) {
+    this.failedConnectionSlots.set(slot, exception);
+    if (this.failedConnectionSlots.size < this.connections.length) {
+      return;
+    }
+
+    // Every slot has exhausted its attempts, so no slot can currently make
+    // progress on the shared queue.
     for (const q of this.queue) {
       q.reject(exception);
     }