add: multi read table lock manager (#153)

Author: vpbs2

meerkat-browser-runner/src/app/app.tsx

@@ -6,7 +6,7 @@ import {
   FileManagerType,
   getMainAppName,
   getRunnerAppName,
-  RunnerMemoryDBFileManager,
+  RunnerIndexedDBFileManager,
   WindowCommunication,
 } from '@devrev/meerkat-dbm';
 
@@ -72,7 +72,7 @@ export function App() {
   }
 
   if (!fileManagerRef.current) {
-    fileManagerRef.current = new RunnerMemoryDBFileManager({
+    fileManagerRef.current = new RunnerIndexedDBFileManager({
       instanceManager: instanceManagerRef.current,
       fetchTableFileBuffers: async () => [],
       logger: log,
@@ -82,7 +82,6 @@ export function App() {
           payload: event,
         });
       },
-      communication: communicationRef.current,
     });
   }
 

meerkat-dbm/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@devrev/meerkat-dbm",
-  "version": "0.1.32",
+  "version": "0.1.33",
   "dependencies": {
     "tslib": "^2.3.0",
     "@duckdb/duckdb-wasm": "1.28.1-dev258.0",

meerkat-dbm/src/dbm/__test__/dbm-parallel.spec.ts

@@ -156,4 +156,36 @@ describe('DBMParallel', () => {
 
     jest.useRealTimers();
   });
+
+  it('should wait for a write lock to be released before executing a query', async () => {
+    const tableName = 'lockedTable';
+    let queryExecuted = false;
+
+    // Mock sendRequest to track execution
+    runnerMock.communication.sendRequest.mockImplementation(() => {
+      queryExecuted = true;
+      return Promise.resolve({
+        message: { isError: false, data: [{ data: 1 }] },
+      });
+    });
+
+    // Acquire a write lock
+    await dbmParallel.lockTables([tableName], 'write');
+
+    // Execute a query that requires a read lock on the same table
+    const queryPromise = dbmParallel.queryWithTables({
+      query: `SELECT * FROM ${tableName}`,
+      tables: [{ name: tableName }],
+    });
+
+    // The query should not have been executed yet because of the write lock
+    expect(queryExecuted).toBe(false);
+
+    // Release the write lock
+    dbmParallel.unlockTables([tableName], 'write');
+
+    // Now the query should be able to execute
+    await queryPromise;
+    expect(queryExecuted).toBe(true);
+  });
 });

meerkat-dbm/src/dbm/__test__/table-manager.spec.ts

@@ -1,70 +1,165 @@
 import { TableLockManager } from '../table-lock-manager';
 
-describe('Table Lock Manager', () => {
+const delay = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
+
+describe('TableLockManager', () => {
   let tableLockManager: TableLockManager;
+  const tableName = 'testTable';
 
   beforeEach(() => {
     tableLockManager = new TableLockManager();
   });
 
-  it('should lock the table and release it', async () => {
-    const tableName = 'exampleTable';
+  describe('write locks', () => {
+    it('should acquire and release a write lock', async () => {
+      await tableLockManager.lockTables([tableName], 'write');
+      expect(tableLockManager.isTableLocked(tableName)).toBe(true);
+
+      tableLockManager.unlockTables([tableName], 'write');
+      expect(tableLockManager.isTableLocked(tableName)).toBe(false);
+    });
+
+    it('a second writer should wait for the first writer to release the lock', async () => {
+      let writer1Finished = false;
+      const writer1 = async () => {
+        await tableLockManager.lockTables([tableName], 'write');
+
+        await delay(50); // Simulate work
+
+        expect(tableLockManager.isTableLocked(tableName)).toBe(true);
+        tableLockManager.unlockTables([tableName], 'write');
+        writer1Finished = true;
+      };
 
-    // Request the lock for the table and then release it
-    await tableLockManager.lockTables([tableName]);
+      let writer2Finished = false;
+      const writer2 = async () => {
+        await tableLockManager.lockTables([tableName], 'write');
 
-    expect(tableLockManager.isTableLocked(tableName)).toBe(true);
+        // When writer2 gets the lock, writer1 should be finished.
+        expect(writer1Finished).toBe(true);
 
-    await tableLockManager.unlockTables([tableName]);
+        expect(tableLockManager.isTableLocked(tableName)).toBe(true);
+        tableLockManager.unlockTables([tableName], 'write');
+        writer2Finished = true;
+      };
 
-    expect(tableLockManager.isTableLocked(tableName)).toBe(false);
+      await Promise.all([writer1(), writer2()]);
 
-    // Again request the lock for the table
-    await tableLockManager.lockTables([tableName]);
+      expect(writer1Finished).toBe(true);
+      expect(writer2Finished).toBe(true);
 
-    await tableLockManager.unlockTables([tableName]);
+      expect(tableLockManager.isTableLocked(tableName)).toBe(false);
+    });
   });
 
-  it('two consumers requesting lock for the same table', async () => {
-    const tableName = 'exampleTable';
+  describe('read locks', () => {
+    it('should allow multiple readers to acquire a lock simultaneously', async () => {
+      const readerPromises = [
+        tableLockManager.lockTables([tableName], 'read'),
+        tableLockManager.lockTables([tableName], 'read'),
+        tableLockManager.lockTables([tableName], 'read'),
+      ];
+
+      await Promise.all(readerPromises);
+      expect(tableLockManager.isTableLocked(tableName)).toBe(true);
+
+      tableLockManager.unlockTables([tableName], 'read');
+      tableLockManager.unlockTables([tableName], 'read');
+      tableLockManager.unlockTables([tableName], 'read');
+      expect(tableLockManager.isTableLocked(tableName)).toBe(false);
+    });
+  });
+
+  describe('mixed read/write locks', () => {
+    it('should not allow a writer if readers have the lock', async () => {
+      await tableLockManager.lockTables([tableName], 'read');
+
+      let writerAcquiredLock = false;
+      const writerPromise = tableLockManager
+        .lockTables([tableName], 'write')
+        .then(() => {
+          writerAcquiredLock = true;
+        });
 
-    // Set up promises for the two consumers
-    const consumer1Promise = tableLockManager.lockTables([tableName]);
-    const consumer2Promise = tableLockManager.lockTables([tableName]);
+      await delay(10); // Give writer time to wait
+      expect(writerAcquiredLock).toBe(false);
+      expect(tableLockManager.isTableLocked(tableName)).toBe(true);
 
-    // Wait for the first consumer to get the lock
-    await expect(consumer1Promise).resolves.toBeUndefined();
+      // Reader releases the lock
+      tableLockManager.unlockTables([tableName], 'read');
 
-    expect(tableLockManager.isTableLocked(tableName)).toBe(true);
+      await writerPromise;
+      expect(writerAcquiredLock).toBe(true);
+      expect(tableLockManager.isTableLocked(tableName)).toBe(true);
 
-    const timeout1 = new Promise((resolve) => {
-      setTimeout(resolve, 1000, 'TIMEOUT');
+      tableLockManager.unlockTables([tableName], 'write');
+      expect(tableLockManager.isTableLocked(tableName)).toBe(false);
     });
 
-    // Promise.race will wait for either the promises be resolved
-    // consumer2 will not be able to get the lock as it is already locked by consumer1
-    await expect(Promise.race([consumer2Promise, timeout1])).resolves.toBe(
-      'TIMEOUT'
-    );
+    it('should not allow a reader if a writer has the lock', async () => {
+      await tableLockManager.lockTables([tableName], 'write');
 
-    // Release the lock for the first consumer
-    await tableLockManager.unlockTables([tableName]);
+      let readerAcquiredLock = false;
+      const readerPromise = tableLockManager
+        .lockTables([tableName], 'read')
+        .then(() => {
+          readerAcquiredLock = true;
+        });
 
-    // Check if the table is still locked as the consumer2 will get the lock
-    expect(tableLockManager.isTableLocked(tableName)).toBe(true);
+      await delay(10);
+      expect(readerAcquiredLock).toBe(false);
+      expect(tableLockManager.isTableLocked(tableName)).toBe(true);
 
-    const timeout2 = new Promise((resolve) => {
-      setTimeout(resolve, 1000, 'TIMEOUT');
+      // Writer releases lock
+      tableLockManager.unlockTables([tableName], 'write');
+
+      await readerPromise;
+      expect(readerAcquiredLock).toBe(true);
+
+      tableLockManager.unlockTables([tableName], 'read');
+      expect(tableLockManager.isTableLocked(tableName)).toBe(false);
     });
 
-    // This time the consumer2 will get the lock
-    await expect(
-      Promise.race([consumer2Promise, timeout2])
-    ).resolves.toBeUndefined();
+    it('should prioritize waiting readers over a new writer', async () => {
+      // Writer1 acquires the lock
+      await tableLockManager.lockTables([tableName], 'write');
+
+      // Readers start waiting
+      const readerPromises = [
+        tableLockManager.lockTables([tableName], 'read'),
+        tableLockManager.lockTables([tableName], 'read'),
+      ];
+
+      // Writer2 starts waiting
+      const writer2Promise = tableLockManager.lockTables([tableName], 'write');
+
+      // Release writer1's lock
+      tableLockManager.unlockTables([tableName], 'write');
 
-    // Release the lock
-    await tableLockManager.unlockTables([tableName]);
+      // The waiting readers should get the lock next
+      await Promise.all(readerPromises);
+      expect(tableLockManager.isTableLocked(tableName)).toBe(true);
 
-    expect(tableLockManager.isTableLocked(tableName)).toBe(false);
+      let writer2AcquiredLock = false;
+      writer2Promise.then(() => {
+        writer2AcquiredLock = true;
+      });
+
+      await delay(10);
+      // Writer2 should still be waiting
+      expect(writer2AcquiredLock).toBe(false);
+
+      // Readers release locks
+      tableLockManager.unlockTables([tableName], 'read');
+      tableLockManager.unlockTables([tableName], 'read');
+
+      // Now writer2 should get the lock
+      await writer2Promise;
+      expect(writer2AcquiredLock).toBe(true);
+      expect(tableLockManager.isTableLocked(tableName)).toBe(true);
+
+      tableLockManager.unlockTables([tableName], 'write');
+      expect(tableLockManager.isTableLocked(tableName)).toBe(false);
+    });
   });
 });

meerkat-dbm/src/dbm/dbm-parallel/dbm-parallel.ts

@@ -128,6 +128,14 @@ export class DBMParallel extends TableLockManager {
         throw new Error('No runner found');
       }
 
+      /**
+       * Lock the tables
+       */
+      await this.lockTables(
+        tables.map((table) => table.name),
+        'read'
+      );
+
       const response =
         await runner.communication.sendRequest<BrowserRunnerExecQueryMessageResponse>(
           {
@@ -143,6 +151,14 @@ export class DBMParallel extends TableLockManager {
       const end = performance.now();
       this.logger.info(`Time to execute by DBM Parallel`, end - start);
 
+      /**
+       * Unlock the tables
+       */
+      this.unlockTables(
+        tables.map((table) => table.name),
+        'read'
+      );
+
       /**
        * The implementation is based on postMessage API, so we don't have the ability to throw an error from the runner
        * We have to check the response and throw an error if isError is true