feat: parallel indexed dbm (#100)

vpbs2 · web-flow · commit 003667ced8c3 · 2024-09-05T02:42:03.000+05:30
* add: parallel indexedb runner

* update: handler

* add: working changes

* chore: cleanup

* update: runner indexedb

* update: parallel indexeddb

* update: runner

* update: tests

* update: indexedb runner

* update: test route

* fix: benchmarking test

* increase: timeout

* update: test
diff --git a/benchmarking/public/runner/indexeddb-runner.html b/benchmarking/public/runner/indexeddb-runner.html
diff --git a/benchmarking/public/runner/memory-runner.html b/benchmarking/public/runner/memory-runner.html
diff --git a/benchmarking/src/app/app.tsx b/benchmarking/src/app/app.tsx
@@ -1,6 +1,7 @@
 import { Route, BrowserRouter as Router, Routes } from 'react-router-dom';
 import { IndexedDBMProvider } from './dbm-context/indexed-dbm-context';
 import { MemoryDBMProvider } from './dbm-context/memory-dbm-context';
+import { ParallelIndexedDBMProvider } from './dbm-context/parallel-indexed-dbm-context';
 import { ParallelMemoryDBMProvider } from './dbm-context/parallel-memory-dbm-context';
 import { RawDBMProvider } from './dbm-context/raw-dbm-context';
 import { FileLoader } from './file-loader/file-loader';
@@ -63,6 +64,19 @@ export function App() {
             </div>
           }
         />
+        <Route
+          path="/parallel-indexed-dbm"
+          element={
+            <div>
+              <h1>Parallel Indexed DuckDB</h1>
+              <ParallelIndexedDBMProvider>
+                <FileLoader bufferType="uint8Array">
+                  <QueryBenchmarking />
+                </FileLoader>
+              </ParallelIndexedDBMProvider>
+            </div>
+          }
+        />
       </Routes>
     </Router>
   );
diff --git a/benchmarking/src/app/dbm-context/parallel-indexed-dbm-context.tsx b/benchmarking/src/app/dbm-context/parallel-indexed-dbm-context.tsx
@@ -0,0 +1,85 @@
+import {
+  DBMParallel,
+  IFrameRunnerManager,
+  ParallelIndexedDBFileManager,
+} from '@devrev/meerkat-dbm';
+import log from 'loglevel';
+import { useRef, useState } from 'react';
+import { DBMContext } from '../hooks/dbm-context';
+import { useClassicEffect } from '../hooks/use-classic-effect';
+import { InstanceManager } from './instance-manager';
+import { useAsyncDuckDB } from './use-async-duckdb';
+
+export const ParallelIndexedDBMProvider = ({
+  children,
+}: {
+  children: JSX.Element;
+}) => {
+  const [dbm, setdbm] = useState<DBMParallel | null>(null);
+  const instanceManagerRef = useRef<InstanceManager>(new InstanceManager());
+  const fileManagerRef = useRef<ParallelIndexedDBFileManager>(
+    new ParallelIndexedDBFileManager({
+      instanceManager: instanceManagerRef.current,
+      fetchTableFileBuffers: async (table) => {
+        return [];
+      },
+      logger: log,
+      onEvent: (event) => {
+        console.info(event);
+      },
+    })
+  );
+
+  const dbState = useAsyncDuckDB();
+
+  useClassicEffect(() => {
+    if (!dbState) {
+      return;
+    }
+    const iframeManager = new IFrameRunnerManager({
+      runnerURL: 'http://localhost:4200/runner/indexeddb-runner.html',
+      origin: 'http://localhost:4200',
+      totalRunners: 4,
+      fetchTableFileBuffers: async (table) => {
+        return [];
+      },
+      fetchPreQuery: () => {
+        return [];
+      },
+      onEvent: (event) => {
+        console.info(event);
+      },
+      logger: log,
+    });
+
+    const dbm = new DBMParallel({
+      instanceManager: instanceManagerRef.current,
+      fileManager: fileManagerRef.current,
+      onEvent: (event) => {
+        console.info(event);
+      },
+      logger: log,
+      options: {
+        shutdownInactiveTime: 1000,
+      },
+      iFrameRunnerManager: iframeManager,
+    });
+
+    setdbm(dbm);
+  }, [dbState]);
+
+  if (!dbm || !fileManagerRef.current) {
+    return <div>Loading...</div>;
+  }
+
+  return (
+    <DBMContext.Provider
+      value={{
+        dbm,
+        fileManager: fileManagerRef.current as any,
+      }}
+    >
+      {children}
+    </DBMContext.Provider>
+  );
+};
diff --git a/benchmarking/src/app/dbm-context/parallel-memory-dbm-context.tsx b/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 | null>(null);
+  const [dbm, setdbm] = useState<DBMParallel<SharedArrayBuffer> | null>(null);
   const instanceManagerRef = useRef<InstanceManager>(new InstanceManager());
   const fileManagerRef = useRef<ParallelMemoryFileManager>(
     new ParallelMemoryFileManager({
@@ -37,7 +37,7 @@ export const ParallelMemoryDBMProvider = ({
       return;
     }
     const iframeManager = new IFrameRunnerManager({
-      runnerURL: 'http://localhost:4200/runner/index.html',
+      runnerURL: 'http://localhost:4200/runner/memory-runner.html',
       origin: 'http://localhost:4200',
       totalRunners: 4,
       fetchTableFileBuffers: async (table) => {
diff --git a/benchmarking/src/app/hooks/dbm-context.tsx b/benchmarking/src/app/hooks/dbm-context.tsx
@@ -1,13 +1,17 @@
 import { DBM, DBMParallel, FileManagerType } from '@devrev/meerkat-dbm';
 import React from 'react';
 
-export const DBMContext = React.createContext<{
-  dbm: DBM | DBMParallel;
+export type DBMContextType<T> = {
+  dbm: DBM | DBMParallel<T>;
   fileManager: FileManagerType;
-}>(null as any);
+};
+
+export const DBMContext = React.createContext<DBMContextType<any> | undefined>(
+  undefined
+);
 
-export const useDBM = () => {
-  const context = React.useContext(DBMContext);
+export const useDBM = <T,>() => {
+  const context = React.useContext(DBMContext) as DBMContextType<T> | undefined;
   if (context === undefined) {
     throw new Error('useDBM must be used within a DBMProvider');
   }
diff --git a/meerkat-dbm/package.json b/meerkat-dbm/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@devrev/meerkat-dbm",
-  "version": "0.1.1",
+  "version": "0.1.2",
   "dependencies": {
     "tslib": "^2.3.0",
     "@devrev/duckdb-wasm": "1.14.3",
diff --git a/meerkat-dbm/src/dbm/__test__/dbm-parallel.spec.ts b/meerkat-dbm/src/dbm/__test__/dbm-parallel.spec.ts
@@ -37,7 +37,7 @@ iFrameRunnerManager.iFrameManagers.set('1', runnerMock as any);
 iFrameRunnerManager.iFrameManagers.set('2', runnerMock as any);
 
 describe('DBMParallel', () => {
-  let dbmParallel: DBMParallel;
+  let dbmParallel: DBMParallel<SharedArrayBuffer>;
   let fileManager: FileManagerType<SharedArrayBuffer>;
   let instanceManager: InstanceManager;
 
diff --git a/meerkat-dbm/src/dbm/dbm-parallel/dbm-parallel.ts b/meerkat-dbm/src/dbm/dbm-parallel/dbm-parallel.ts
@@ -4,7 +4,12 @@ import {
   BROWSER_RUNNER_TYPE,
   BrowserRunnerExecQueryMessageResponse,
 } from '../../window-communication/runner-types';
-import { DBMConstructorOptions, QueryOptions, TableConfig } from '../types';
+import {
+  DBMConstructorOptions,
+  QueryOptions,
+  TableConfig,
+  TableLock,
+} from '../types';
 import { IFrameRunnerManager } from './runner-manager';
 
 //Round Robin for multiple runners like 10
@@ -16,9 +21,11 @@ const roundRobin = (counter: number, maxValue: number): number => {
   return counter + 1;
 };
 
-export class DBMParallel {
-  private fileManager: FileManagerType<SharedArrayBuffer>;
+export class DBMParallel<BufferType = Uint8Array> {
+  private fileManager: FileManagerType<BufferType>;
   private logger: DBMLogger;
+  private tableLockRegistry: Record<string, TableLock> = {};
+
   private onEvent?: (event: DBMEvent) => void;
   private options: DBMConstructorOptions['options'];
   private onDuckDBShutdown?: () => void;
@@ -35,7 +42,7 @@ export class DBMParallel {
     instanceManager,
     onDuckDBShutdown,
     iFrameRunnerManager,
-  }: DBMConstructorOptions<SharedArrayBuffer> & {
+  }: DBMConstructorOptions<BufferType> & {
     iFrameRunnerManager: IFrameRunnerManager;
   }) {
     this.fileManager = fileManager;
@@ -84,6 +91,66 @@ export class DBMParallel {
     }, 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,
diff --git a/meerkat-dbm/src/file-manager/index.ts b/meerkat-dbm/src/file-manager/index.ts
@@ -1,5 +1,7 @@
 export * from './file-manager-type';
 export * from './indexed-db/indexed-db-file-manager';
+export * from './indexed-db/parallel-indexed-db-file-manager';
+export * from './indexed-db/runner-indexed-db-file-manager';
 export * from './memory/memory-file-manager';
 export * from './memory/parallel-memory-file-manager';
 export * from './memory/runner-memory-file-manager';
diff --git a/meerkat-dbm/src/file-manager/indexed-db/parallel-indexed-db-file-manager.ts b/meerkat-dbm/src/file-manager/indexed-db/parallel-indexed-db-file-manager.ts
@@ -0,0 +1,19 @@
+import { TableConfig } from '../../dbm/types';
+import {
+  FileManagerConstructorOptions,
+  FileManagerType,
+} from '../file-manager-type';
+import { IndexedDBFileManager } from './indexed-db-file-manager';
+
+export class ParallelIndexedDBFileManager
+  extends IndexedDBFileManager
+  implements FileManagerType<Uint8Array>
+{
+  constructor(options: FileManagerConstructorOptions) {
+    super(options);
+  }
+
+  override async mountFileBufferByTables(tables: TableConfig[]): Promise<void> {
+    // no-op
+  }
+}
diff --git a/meerkat-dbm/src/file-manager/indexed-db/runner-indexed-db-file-manager.ts b/meerkat-dbm/src/file-manager/indexed-db/runner-indexed-db-file-manager.ts
@@ -0,0 +1,50 @@
+import {
+  FileBufferStore,
+  FileJsonStore,
+  FileManagerConstructorOptions,
+  FileManagerType,
+} from '../file-manager-type';
+import { IndexedDBFileManager } from './indexed-db-file-manager';
+
+export class RunnerIndexedDBFileManager
+  extends IndexedDBFileManager
+  implements FileManagerType<Uint8Array>
+{
+  constructor(options: FileManagerConstructorOptions) {
+    super(options);
+  }
+
+  override async bulkRegisterFileBuffer(
+    fileBuffers: FileBufferStore[]
+  ): Promise<void> {
+    // no-op
+  }
+
+  override async registerFileBuffer(
+    fileBuffer: FileBufferStore
+  ): Promise<void> {
+    // no-op
+  }
+
+  override async bulkRegisterJSON(jsonData: FileJsonStore[]): Promise<void> {
+    // no-op
+  }
+
+  override async registerJSON(jsonData: FileJsonStore): Promise<void> {
+    // no-op
+  }
+
+  override async setTableMetadata(
+    tableName: string,
+    metadata: object
+  ): Promise<void> {
+    // no-op
+  }
+
+  override async dropFilesByTableName(
+    tableName: string,
+    fileNames: string[]
+  ): Promise<void> {
+    // no-op
+  }
+}
diff --git a/meerkat-dbm/src/utils/merge-file-buffer-store-into-table.ts b/meerkat-dbm/src/utils/merge-file-buffer-store-into-table.ts
@@ -20,8 +20,10 @@ import { Table } from '../types';
  *
  */
 
-export const mergeFileBufferStoreIntoTable = (
-  fileBufferStore: FileBufferStore[],
+export const mergeFileBufferStoreIntoTable = <
+  T extends Uint8Array | SharedArrayBuffer
+>(
+  fileBufferStore: FileBufferStore<T>[],
   currentTableState: Table[]
 ): Map<string, Table> => {
   const tableMap = new Map<string, Table>(

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`{`
`2`	`2`	`"name": "@devrev/meerkat-dbm",`
`3`		`- "version": "0.1.1",`
	`3`	`+ "version": "0.1.2",`
`4`	`4`	`"dependencies": {`
`5`	`5`	`"tslib": "^2.3.0",`
`6`	`6`	`"@devrev/duckdb-wasm": "1.14.3",`