refactor: dbm and file manager (#124)

* refactor: dbm

* fix: test

* remove: log

* fix: test

Author: vpbs2

.gitignore

@@ -43,3 +43,7 @@ Thumbs.db
 
 # NX
 .nx
+
+
+# Parquet data
+/electron-app/src/data

benchmarking/src/app/benchmarking-tests/dbm-benchmarking.spec.ts

@@ -20,7 +20,9 @@ describe('Benchmarking DBMs', () => {
 
     browser = await puppeteer.launch({
       headless: 'new',
+      args: ['--no-sandbox', '--disable-setuid-sandbox'],
     });
+
     page = await browser.newPage();
 
     //Wait for the server to start by visiting the page

benchmarking/src/app/dbm-context/parallel-memory-dbm-context.tsx

@@ -15,7 +15,7 @@ export const ParallelMemoryDBMProvider = ({
 }: {
   children: JSX.Element;
 }) => {
-  const [dbm, setdbm] = useState<DBMParallel<SharedArrayBuffer> | null>(null);
+  const [dbm, setdbm] = useState<DBMParallel | null>(null);
   const instanceManagerRef = useRef<InstanceManager>(new InstanceManager());
   const fileManagerRef = useRef<ParallelMemoryFileManager>(
     new ParallelMemoryFileManager({

benchmarking/src/app/hooks/dbm-context.tsx

@@ -1,17 +1,17 @@
 import { DBM, DBMParallel, FileManagerType } from '@devrev/meerkat-dbm';
 import React from 'react';
 
-export type DBMContextType<T> = {
-  dbm: DBM | DBMParallel<T>;
+export type DBMContextType = {
+  dbm: DBM | DBMParallel;
   fileManager: FileManagerType;
 };
 
-export const DBMContext = React.createContext<DBMContextType<any> | undefined>(
+export const DBMContext = React.createContext<DBMContextType | undefined>(
   undefined
 );
 
-export const useDBM = <T,>() => {
-  const context = React.useContext(DBMContext) as DBMContextType<T> | undefined;
+export const useDBM = () => {
+  const context = React.useContext(DBMContext) as DBMContextType | undefined;
   if (context === undefined) {
     throw new Error('useDBM must be used within a DBMProvider');
   }

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

@@ -37,8 +37,8 @@ iFrameRunnerManager.iFrameManagers.set('1', runnerMock as any);
 iFrameRunnerManager.iFrameManagers.set('2', runnerMock as any);
 
 describe('DBMParallel', () => {
-  let dbmParallel: DBMParallel<SharedArrayBuffer>;
-  let fileManager: FileManagerType<SharedArrayBuffer>;
+  let dbmParallel: DBMParallel;
+  let fileManager: FileManagerType;
   let instanceManager: InstanceManager;
 
   beforeAll(async () => {

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

@@ -9,11 +9,11 @@ import { DBM } from '../dbm';
 import { DBMConstructorOptions, TableConfig } from '../types';
 import { InstanceManager } from './mock';
 
-export class MockFileManager<T> implements FileManagerType<T> {
-  private fileBufferStore: Record<string, FileBufferStore<T>> = {};
+export class MockFileManager implements FileManagerType {
+  private fileBufferStore: Record<string, FileBufferStore> = {};
   private tables: Record<string, Table> = {};
 
-  async bulkRegisterFileBuffer(props: FileBufferStore<T>[]): Promise<void> {
+  async bulkRegisterFileBuffer(props: FileBufferStore[]): Promise<void> {
     for (const prop of props) {
       this.fileBufferStore[prop.fileName] = prop;
       this.tables[prop.tableName] = this.tables[prop.tableName] || {
@@ -23,7 +23,7 @@ export class MockFileManager<T> implements FileManagerType<T> {
     }
   }
 
-  async registerFileBuffer(prop: FileBufferStore<T>): Promise<void> {
+  async registerFileBuffer(prop: FileBufferStore): Promise<void> {
     this.fileBufferStore[prop.fileName] = prop;
     this.tables[prop.tableName] = this.tables[prop.tableName] || { files: [] };
     this.tables[prop.tableName].files.push(prop);
@@ -40,18 +40,10 @@ export class MockFileManager<T> implements FileManagerType<T> {
 
     this.registerFileBuffer({
       ...fileData,
-      buffer: [] as T,
+      buffer: new Uint8Array(),
     });
   }
 
-  async getFileBuffer(name: string): Promise<T> {
-    const fileBuffer = this.fileBufferStore[name];
-    if (!fileBuffer) {
-      throw new Error(`File buffer for ${name} not found`);
-    }
-    return fileBuffer.buffer;
-  }
-
   async mountFileBufferByTables(tables: TableConfig[]): Promise<void> {
     const tableNames = tables.map((table) => table.name);
     for (const tableName of tableNames) {
@@ -462,69 +454,6 @@ describe('DBM', () => {
     });
   });
 
-  describe('table locks', () => {
-    it('should lock the table and release it', async () => {
-      const tableName = 'exampleTable';
-
-      // Request the lock for the table and then release it
-      await dbm.lockTables([tableName]);
-
-      expect(dbm.isTableLocked(tableName)).toBe(true);
-
-      await dbm.unlockTables([tableName]);
-
-      expect(dbm.isTableLocked(tableName)).toBe(false);
-
-      // Again request the lock for the table
-      await dbm.lockTables([tableName]);
-
-      await dbm.unlockTables([tableName]);
-    });
-
-    it('two consumers requesting lock for the same table', async () => {
-      const tableName = 'exampleTable';
-
-      // Set up promises for the two consumers
-      const consumer1Promise = dbm.lockTables([tableName]);
-      const consumer2Promise = dbm.lockTables([tableName]);
-
-      // Wait for the first consumer to get the lock
-      await expect(consumer1Promise).resolves.toBeUndefined();
-
-      expect(dbm.isTableLocked(tableName)).toBe(true);
-
-      const timeout1 = new Promise((resolve) => {
-        setTimeout(resolve, 1000, 'TIMEOUT');
-      });
-
-      // 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'
-      );
-
-      // Release the lock for the first consumer
-      await dbm.unlockTables([tableName]);
-
-      // Check if the table is still locked as the consumer2 will get the lock
-      expect(dbm.isTableLocked(tableName)).toBe(true);
-
-      const timeout2 = new Promise((resolve) => {
-        setTimeout(resolve, 1000, 'TIMEOUT');
-      });
-
-      // This time the consumer2 will get the lock
-      await expect(
-        Promise.race([consumer2Promise, timeout2])
-      ).resolves.toBeUndefined();
-
-      // Release the lock
-      await dbm.unlockTables([tableName]);
-
-      expect(dbm.isTableLocked(tableName)).toBe(false);
-    });
-  });
-
   describe('create connection callback', () => {
     it('should call the create connection callback', async () => {
       await dbm.query('SELECT 1');

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

@@ -0,0 +1,70 @@
+import { TableLockManager } from '../table-lock-manager';
+
+describe('Table Lock Manager', () => {
+  let tableLockManager: TableLockManager;
+
+  beforeEach(() => {
+    tableLockManager = new TableLockManager();
+  });
+
+  it('should lock the table and release it', async () => {
+    const tableName = 'exampleTable';
+
+    // Request the lock for the table and then release it
+    await tableLockManager.lockTables([tableName]);
+
+    expect(tableLockManager.isTableLocked(tableName)).toBe(true);
+
+    await tableLockManager.unlockTables([tableName]);
+
+    expect(tableLockManager.isTableLocked(tableName)).toBe(false);
+
+    // Again request the lock for the table
+    await tableLockManager.lockTables([tableName]);
+
+    await tableLockManager.unlockTables([tableName]);
+  });
+
+  it('two consumers requesting lock for the same table', async () => {
+    const tableName = 'exampleTable';
+
+    // Set up promises for the two consumers
+    const consumer1Promise = tableLockManager.lockTables([tableName]);
+    const consumer2Promise = tableLockManager.lockTables([tableName]);
+
+    // Wait for the first consumer to get the lock
+    await expect(consumer1Promise).resolves.toBeUndefined();
+
+    expect(tableLockManager.isTableLocked(tableName)).toBe(true);
+
+    const timeout1 = new Promise((resolve) => {
+      setTimeout(resolve, 1000, 'TIMEOUT');
+    });
+
+    // 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'
+    );
+
+    // Release the lock for the first consumer
+    await tableLockManager.unlockTables([tableName]);
+
+    // Check if the table is still locked as the consumer2 will get the lock
+    expect(tableLockManager.isTableLocked(tableName)).toBe(true);
+
+    const timeout2 = new Promise((resolve) => {
+      setTimeout(resolve, 1000, 'TIMEOUT');
+    });
+
+    // This time the consumer2 will get the lock
+    await expect(
+      Promise.race([consumer2Promise, timeout2])
+    ).resolves.toBeUndefined();
+
+    // Release the lock
+    await tableLockManager.unlockTables([tableName]);
+
+    expect(tableLockManager.isTableLocked(tableName)).toBe(false);
+  });
+});

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

@@ -4,12 +4,8 @@ import {
   BROWSER_RUNNER_TYPE,
   BrowserRunnerExecQueryMessageResponse,
 } from '../../window-communication/runner-types';
-import {
-  DBMConstructorOptions,
-  QueryOptions,
-  TableConfig,
-  TableLock,
-} from '../types';
+import { TableLockManager } from '../table-lock-manager';
+import { DBMConstructorOptions, QueryOptions, TableConfig } from '../types';
 import { IFrameRunnerManager } from './runner-manager';
 
 //Round Robin for multiple runners like 10
@@ -21,10 +17,9 @@ const roundRobin = (counter: number, maxValue: number): number => {
   return counter + 1;
 };
 
-export class DBMParallel<BufferType = Uint8Array> {
-  private fileManager: FileManagerType<BufferType>;
+export class DBMParallel extends TableLockManager {
+  private fileManager: FileManagerType;
   private logger: DBMLogger;
-  private tableLockRegistry: Record<string, TableLock> = {};
 
   private onEvent?: (event: DBMEvent) => void;
   private options: DBMConstructorOptions['options'];
@@ -42,9 +37,11 @@ export class DBMParallel<BufferType = Uint8Array> {
     instanceManager,
     onDuckDBShutdown,
     iFrameRunnerManager,
-  }: DBMConstructorOptions<BufferType> & {
+  }: DBMConstructorOptions & {
     iFrameRunnerManager: IFrameRunnerManager;
   }) {
+    super();
+
     this.fileManager = fileManager;
     this.logger = logger;
     this.onEvent = onEvent;
@@ -91,66 +88,6 @@ export class DBMParallel<BufferType = Uint8Array> {
     }, this.options.shutdownInactiveTime);
   }
 
-  async lockTables(tableNames: string[]): Promise<void> {
-    const promises = [];
-
-    for (const tableName of tableNames) {
-      const tableLock = this.tableLockRegistry[tableName];
-
-      // If the table lock doesn't exist, create a new lock
-      if (!tableLock) {
-        this.tableLockRegistry[tableName] = {
-          isLocked: true,
-          promiseQueue: [],
-        };
-        continue;
-      }
-
-      // If the table is already locked, add the promise to the queue
-      if (tableLock.isLocked) {
-        const promise = new Promise<void>((resolve, reject) => {
-          tableLock.promiseQueue.push({ reject, resolve });
-        });
-        promises.push(promise);
-      }
-
-      // Set the table as locked
-      tableLock.isLocked = true;
-    }
-
-    // Wait for all promises to resolve (locks to be acquired)
-    await Promise.all(promises);
-  }
-
-  async unlockTables(tableNames: string[]): Promise<void> {
-    for (const tableName of tableNames) {
-      const tableLock = this.tableLockRegistry[tableName];
-
-      // If the table lock doesn't exist, create a new lock
-      if (!tableLock) {
-        this.tableLockRegistry[tableName] = {
-          isLocked: false,
-          promiseQueue: [],
-        };
-      }
-
-      const nextPromiseInQueue = tableLock?.promiseQueue?.shift();
-
-      // If there is a promise in the queue, resolve it and keep the table as locked
-      if (nextPromiseInQueue) {
-        tableLock.isLocked = true;
-        nextPromiseInQueue.resolve();
-      } else {
-        // If there are no promises in the queue, set the table as unlocked
-        tableLock.isLocked = false;
-      }
-    }
-  }
-
-  isTableLocked(tableName: string): boolean {
-    return this.tableLockRegistry[tableName]?.isLocked ?? false;
-  }
-
   public async queryWithTables({
     query,
     tables,

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

@@ -12,9 +12,7 @@ import { IFrameManager } from './iframe-manager';
 export interface IFrameRunnerManagerConstructor {
   runnerURL: string;
   origin: string;
-  fetchTableFileBuffers: (
-    tables: TableConfig[]
-  ) => Promise<FileBufferStore<SharedArrayBuffer>[]>;
+  fetchTableFileBuffers: (tables: TableConfig[]) => Promise<FileBufferStore[]>;
   fetchPreQuery: (
     runnerId: string,
     tableWiseFiles: TableWiseFiles[]
@@ -47,7 +45,7 @@ export class IFrameRunnerManager {
 
   private fetchTableFileBuffers: (
     tables: TableConfig[]
-  ) => Promise<FileBufferStore<SharedArrayBuffer>[]>;
+  ) => Promise<FileBufferStore[]>;
   private fetchPreQuery: (
     runnerId: string,
     tableWiseFiles: TableWiseFiles[]