Support larger read pools too

Author: simolus3

packages/web/src/db/adapters/AsyncWebAdapter.ts

@@ -3,7 +3,10 @@ import {
   DBAdapterDefaultMixin,
   DBAdapterListener,
   DBLockOptions,
-  LockContext
+  LockContext,
+  Mutex,
+  Semaphore,
+  UnlockFn
 } from '@powersync/common';
 import { SharedConnectionWorker, WebDBAdapter, WebDBAdapterConfiguration } from './WebDBAdapter.js';
 import { DatabaseClient } from './wa-sqlite/DatabaseClient.js';
@@ -14,77 +17,52 @@ type PendingListener = { listener: Partial<DBAdapterListener>; closeAfterRegiste
  * A connection pool implementation delegating to another pool opened asynchronnously.
  */
 class AsyncConnectionPool implements ConnectionPool {
-  protected readonly inner: Promise<PoolConnection>;
-
+  protected readonly state: Promise<PoolState>;
   protected resolvedWriter?: DatabaseClient;
-  private activeOnWriter = 0;
-  private activeOnReader = 0;
 
   private readonly pendingListeners = new Set<PendingListener>();
 
   constructor(
     inner: Promise<PoolConnection>,
     readonly name: string
   ) {
-    this.inner = inner.then((client) => {
+    this.state = inner.then((client) => {
       for (const pending of this.pendingListeners) {
         pending.closeAfterRegisteredOnResolvedPool = client.writer.registerListener(pending.listener);
       }
       this.pendingListeners.clear();
 
       this.resolvedWriter = client.writer;
-      return client;
+      if (client.additionalReaders.length) {
+        return readWritePoolState(client.writer, client.additionalReaders);
+      }
+
+      return singleConnectionPoolState(client.writer);
     });
   }
 
   async init() {
-    await this.inner;
+    await this.state;
   }
 
   async close() {
-    const inner = await this.inner;
-
-    await inner.writer.close();
-    await inner.additionalReader?.close();
+    const state = await this.state;
+    await state.close();
   }
 
   async readLock<T>(fn: (tx: LockContext) => Promise<T>, options?: DBLockOptions): Promise<T> {
-    const inner = await this.inner;
-
-    // This is a crude load balancing scheme between the writer and an additional read connection (if available).
-    // Ideally, we should support abortable requests (which would allow us to request a lock from both and just use
-    // whatever completes first). For now, this at least gives us some concurrency. We can improve this in the future.
-    if (inner.additionalReader && this.activeOnReader <= this.activeOnWriter) {
-      try {
-        this.activeOnReader++;
-        return await inner.additionalReader.readLock(fn, options);
-      } finally {
-        this.activeOnReader--;
-      }
-    }
-
-    try {
-      this.activeOnWriter++;
-      return await inner.writer.readLock(fn, options);
-    } finally {
-      this.activeOnWriter--;
-    }
+    const state = await this.state;
+    return state.withConnection(true, fn, options);
   }
 
   async writeLock<T>(fn: (tx: LockContext) => Promise<T>, options?: DBLockOptions): Promise<T> {
-    const inner = await this.inner;
-    try {
-      this.activeOnWriter++;
-      return await inner.writer.writeLock(fn, options);
-    } finally {
-      this.activeOnWriter--;
-    }
+    const state = await this.state;
+    return state.withConnection(false, fn, options);
   }
 
   async refreshSchema(): Promise<void> {
-    const inner = await this.inner;
-    await inner.writer.refreshSchema();
-    await inner.additionalReader?.refreshSchema();
+    const state = await this.state;
+    await state.refreshSchema();
   }
 
   registerListener(listener: Partial<DBAdapterListener>): () => void {
@@ -107,13 +85,116 @@ class AsyncConnectionPool implements ConnectionPool {
 
 export interface PoolConnection {
   writer: DatabaseClient;
-  additionalReader?: DatabaseClient;
+  additionalReaders: DatabaseClient[];
+}
+
+interface PoolState {
+  writer: DatabaseClient;
+  withConnection<T>(allowReadOnly: boolean, fn: (tx: LockContext) => Promise<T>, options?: DBLockOptions): Promise<T>;
+  close(): Promise<void>;
+  refreshSchema(): Promise<void>;
+}
+
+function singleConnectionPoolState(connection: DatabaseClient): PoolState {
+  return {
+    writer: connection,
+    withConnection: (allowReadOnly, fn, options) => {
+      if (allowReadOnly) {
+        return connection.readLock(fn, options);
+      } else {
+        return connection.writeLock(fn, options);
+      }
+    },
+    close: () => connection.close(),
+    refreshSchema: () => connection.refreshSchema()
+  };
+}
+
+function readWritePoolState(writer: DatabaseClient, readers: DatabaseClient[]): PoolState {
+  // DatabaseClients have locks internally, so these aren't necessary for correctness. However, our mutex and semaphore
+  // implementations are very cheap to cancel, which we use to dispatch reads to the first available connection (by
+  // simply requesting all of them and sticking with the first connection we get).
+  const writerMutex = new Mutex();
+  const readerSemaphore = new Semaphore(readers);
+
+  return {
+    writer,
+    async withConnection(allowReadOnly, fn, options) {
+      const abortController = new AbortController();
+      const abortSignal = abortController.signal;
+
+      let timeout: any = null;
+      let release: UnlockFn | undefined;
+      if (options?.timeoutMs) {
+        timeout = setTimeout(() => abortController.abort, options.timeoutMs);
+      }
+
+      try {
+        if (allowReadOnly) {
+          let connection: DatabaseClient;
+
+          // Even if we have a pool of read connections, it's typically very small and we assume that most queries are
+          // reads. So, we want to request any connection from the read pool and the dedicated write connection (which
+          // can also serve reads). We race for the first connection we can obtain this way, and then abort the other
+          // request.
+          [connection, release] = await new Promise<[DatabaseClient, UnlockFn]>((resolve, reject) => {
+            let didComplete = false;
+            function complete() {
+              didComplete = true;
+              abortController.abort();
+            }
+
+            function completeSuccess(connection: DatabaseClient, returnFn: UnlockFn) {
+              if (didComplete) {
+                // We're not going to use this connection, so return it immediately.
+                returnFn();
+              } else {
+                complete();
+                resolve([connection, returnFn]);
+              }
+            }
+
+            function completeError(error: unknown) {
+              // We either have a working connection already, or we've rejected the promise. Either way, we don't need
+              // to do either thing again.
+              if (didComplete) return;
+
+              complete();
+              reject(error);
+            }
+
+            writerMutex.acquire(abortSignal).then((unlock) => completeSuccess(writer, unlock), completeError);
+            readerSemaphore
+              .requestOne(abortSignal)
+              .then(({ item, release }) => completeSuccess(item, release), completeError);
+          });
+
+          return await connection.readLock(fn);
+        } else {
+          return await writerMutex.runExclusive(() => writer.writeLock(fn), abortSignal);
+        }
+      } finally {
+        if (timeout != null) {
+          clearTimeout(timeout);
+        }
+        release?.();
+      }
+    },
+    async close() {
+      await writer.close();
+      await Promise.all(readers.map((r) => r.close()));
+    },
+    async refreshSchema() {
+      await writer.refreshSchema();
+      await Promise.all(readers.map((r) => r.refreshSchema()));
+    }
+  };
 }
 
 export class AsyncDbAdapter extends DBAdapterDefaultMixin(AsyncConnectionPool) implements WebDBAdapter {
   async shareConnection(): Promise<SharedConnectionWorker> {
-    const inner = await this.inner;
-    return inner.writer.shareConnection();
+    const state = await this.state;
+    return state.writer.shareConnection();
   }
 
   getConfiguration(): WebDBAdapterConfiguration {

packages/web/src/db/adapters/wa-sqlite/WASQLiteOpenFactory.ts

@@ -19,6 +19,16 @@ import { AsyncDbAdapter, PoolConnection } from '../AsyncWebAdapter.js';
 
 export interface WASQLiteOpenFactoryOptions extends WebSQLOpenFactoryOptions {
   vfs?: WASQLiteVFS;
+  /**
+   * If the {@link vfs} supports it, an additional amount of read-only connections to open. Using additional read
+   * connections can speed up queries by dispatching them to multiple workers running them concurrently.
+   *
+   * {@link WASQLiteVFS.OPFSWriteAheadVFS} is the only VFS with support for multiple connections, so this option is
+   * ignored for other VFS implementations.
+   *
+   * Defaults to 1.
+   */
+  additionalReaders?: number;
 }
 
 export interface ResolvedWASQLiteOpenFactoryOptions extends ResolvedWebSQLOpenOptions {
@@ -113,7 +123,7 @@ export class WASQLiteOpenFactory implements SQLOpenFactory {
     });
 
     let client: DatabaseClient;
-    let additionalReader: DatabaseClient | undefined;
+    let additionalReaders: DatabaseClient[] = [];
     let requiresPersistentTriggers = vfsRequiresDedicatedWorkers(vfs);
 
     if (useWebWorker) {
@@ -167,9 +177,12 @@ export class WASQLiteOpenFactory implements SQLOpenFactory {
 
       if (vfs == WASQLiteVFS.OPFSWriteAheadVFS) {
         // This VFS supports concurrent reads, so we can open additional workers to host read-only connections for
-        // concurrent reads / writes. To avoid excessive resource usage, we currently add one additional reader per
-        // tab. In the future, we might revisit this to use a growable pool of readers.
-        additionalReader = await openDatabaseWorker(resolveOptions(true));
+        // concurrent reads / writes.
+        const additionalReadersCount = this.options.additionalReaders ?? 1;
+        for (let i = 0; i < additionalReadersCount; i++) {
+          const reader = await openDatabaseWorker(resolveOptions(true));
+          additionalReaders.push(reader);
+        }
       }
     } else {
       // Don't use a web worker. Instead, open the MultiDatabaseServer a worker would use locally.
@@ -189,7 +202,7 @@ export class WASQLiteOpenFactory implements SQLOpenFactory {
 
     return {
       writer: client,
-      additionalReader
+      additionalReaders
     };
   }
 }

packages/web/tests/src/db/write_ahead_log_opfs.test.ts

@@ -7,7 +7,8 @@ test('supports concurrent reads', async () => {
   const db = generateTestDb({
     database: new WASQLiteOpenFactory({
       dbFilename: 'basic-opfs.sqlite',
-      vfs: WASQLiteVFS.OPFSWriteAheadVFS
+      vfs: WASQLiteVFS.OPFSWriteAheadVFS,
+      additionalReaders: 1
     }),
     schema: TEST_SCHEMA
   });
@@ -20,4 +21,12 @@ test('supports concurrent reads', async () => {
   });
 
   expect(await db.getAll('SELECT * FROM customers')).toHaveLength(1);
+
+  // Despite only using one additional read connection, we should be able to support two concurrent readers by using
+  // the write connection too.
+  await db.readLock(async (ctx1) => {
+    await db.readLock(async (ctx2) => {
+      await Promise.all([ctx1.execute('SELECT 1'), ctx2.execute('SELECT 2')]);
+    });
+  });
 });