feat: Adding duckdb instance shutdown during inactivity (#32)

* Adding instance manager

* Adding shutdown DB test

* Adding missed changes

Author: itsTalwar

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

@@ -3,11 +3,15 @@ import log from 'loglevel';
 import React, { 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 IndexedDBMProvider = ({ children }: { children: JSX.Element }) => {
   const fileManagerRef = React.useRef<IndexedDBFileManager | null>(null);
   const [dbm, setdbm] = useState<DBM | null>(null);
+  const instanceManagerRef = React.useRef<InstanceManager>(
+    new InstanceManager()
+  );
 
   const dbState = useAsyncDuckDB();
 
@@ -16,7 +20,7 @@ export const IndexedDBMProvider = ({ children }: { children: JSX.Element }) => {
       return;
     }
     fileManagerRef.current = new IndexedDBFileManager({
-      db: dbState,
+      instanceManager: instanceManagerRef.current,
       fetchTableFileBuffers: async (table) => {
         return [];
       },
@@ -25,12 +29,15 @@ export const IndexedDBMProvider = ({ children }: { children: JSX.Element }) => {
     fileManagerRef.current.initializeDB();
 
     const dbm = new DBM({
-      db: dbState,
+      instanceManager: instanceManagerRef.current,
       fileManager: fileManagerRef.current,
       onEvent: (event) => {
         console.info(event);
       },
       logger: log,
+      options: {
+        shutdownInactiveTime: 1000,
+      },
     });
 
     setdbm(dbm);

benchmarking/benchmarking-app/src/app/dbm-context/instance-manager.ts

@@ -0,0 +1,42 @@
+import { InstanceManagerType } from '@devrev/meerkat-dbm';
+import * as duckdb from '@duckdb/duckdb-wasm';
+import { AsyncDuckDB, LogEntryVariant } from '@duckdb/duckdb-wasm';
+const JSDELIVR_BUNDLES = duckdb.getJsDelivrBundles();
+
+export class InstanceManager implements InstanceManagerType {
+  private db: AsyncDuckDB | null = null;
+
+  private async initDB() {
+    const bundle = await duckdb.selectBundle(JSDELIVR_BUNDLES);
+
+    const worker_url = URL.createObjectURL(
+      new Blob([`importScripts("${bundle.mainWorker!}");`], {
+        type: 'text/javascript',
+      })
+    );
+
+    // Instantiate the asynchronus version of DuckDB-wasm
+    const worker = new Worker(worker_url);
+    const logger = {
+      log: (msg: LogEntryVariant) => console.log(msg),
+    };
+    const db = new duckdb.AsyncDuckDB(logger, worker);
+    await db.instantiate(bundle.mainModule, bundle.pthreadWorker);
+    URL.revokeObjectURL(worker_url);
+    return db;
+  }
+
+  async getDB() {
+    if (!this.db) {
+      console.info('Creating new DB');
+      this.db = await this.initDB();
+    }
+    return this.db;
+  }
+
+  async terminateDB() {
+    console.info('terminateDB');
+    await this.db?.terminate();
+    this.db = null;
+  }
+}

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

@@ -3,11 +3,15 @@ import log from 'loglevel';
 import React, { 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 MemoryDBMProvider = ({ children }: { children: JSX.Element }) => {
   const fileManagerRef = React.useRef<FileManagerType | null>(null);
   const [dbm, setdbm] = useState<DBM | null>(null);
+  const instanceManagerRef = React.useRef<InstanceManager>(
+    new InstanceManager()
+  );
 
   const dbState = useAsyncDuckDB();
 
@@ -16,14 +20,14 @@ export const MemoryDBMProvider = ({ children }: { children: JSX.Element }) => {
       return;
     }
     fileManagerRef.current = new MemoryDBFileManager({
-      db: dbState,
+      instanceManager: instanceManagerRef.current,
       fetchTableFileBuffers: async (table) => {
         return [];
       },
     });
     log.setLevel('DEBUG');
     const dbm = new DBM({
-      db: dbState,
+      instanceManager: instanceManagerRef.current,
       fileManager: fileManagerRef.current,
       logger: log,
       onEvent: (event) => {

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

@@ -3,11 +3,15 @@ import log from 'loglevel';
 import React, { 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 RawDBMProvider = ({ children }: { children: JSX.Element }) => {
   const fileManagerRef = React.useRef<FileManagerType | null>(null);
   const [dbm, setdbm] = useState<DBM | null>(null);
+  const instanceManagerRef = React.useRef<InstanceManager>(
+    new InstanceManager()
+  );
 
   const dbState = useAsyncDuckDB();
 
@@ -16,13 +20,13 @@ export const RawDBMProvider = ({ children }: { children: JSX.Element }) => {
       return;
     }
     fileManagerRef.current = new MemoryDBFileManager({
-      db: dbState,
+      instanceManager: instanceManagerRef.current,
       fetchTableFileBuffers: async (table) => {
         return [];
       },
     });
     const dbm = new DBM({
-      db: dbState,
+      instanceManager: instanceManagerRef.current,
       fileManager: fileManagerRef.current,
       logger: log,
       onEvent: (event) => {

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

@@ -5,6 +5,7 @@ import {
   FileManagerType,
 } from '../file-manager/file-manager-type';
 import { DBM, DBMConstructorOptions } from './dbm';
+import { InstanceManagerType } from './instance-manager';
 
 export class MockFileManager implements FileManagerType {
   private fileBufferStore: Record<string, FileBufferStore> = {};
@@ -64,21 +65,34 @@ const mockDB = {
   },
 };
 
+export class InstanceManager implements InstanceManagerType {
+  async getDB() {
+    return mockDB as AsyncDuckDB;
+  }
+
+  async terminateDB() {
+    // do nothing
+  }
+}
+
 describe('DBM', () => {
   let db: AsyncDuckDB;
   let fileManager: FileManagerType;
   let dbm: DBM;
+  let instanceManager;
 
   beforeAll(async () => {
     // eslint-disable-next-line @typescript-eslint/ban-ts-comment
     //@ts-ignore
     db = mockDB;
     fileManager = new MockFileManager();
+    instanceManager = new InstanceManager();
   });
 
   beforeEach(() => {
+    const instanceManager = new InstanceManager();
     const options: DBMConstructorOptions = {
-      db,
+      instanceManager: instanceManager,
       fileManager,
       logger: log,
       onEvent: (event) => {
@@ -149,4 +163,113 @@ describe('DBM', () => {
       expect(result3).toEqual(['SELECT * FROM table3']);
     });
   });
+
+  describe('shutdown the db test', () => {
+    it('should shutdown the db if there are no queries in the queue', async () => {
+      const instanceManager = new InstanceManager();
+      // If instanceManager.terminateDB is a method
+      jest.spyOn(instanceManager, 'terminateDB');
+
+      // If instanceManager.terminateDB is a function
+      instanceManager.terminateDB = jest.fn();
+      const options: DBMConstructorOptions = {
+        instanceManager: instanceManager,
+        fileManager,
+        logger: log,
+        onEvent: (event) => {
+          console.log(event);
+        },
+        options: {
+          shutdownInactiveTime: 100,
+        },
+      };
+      const dbm = new DBM(options);
+
+      /**
+       * Execute a query
+       */
+      const promise1 = dbm.queryWithTableNames('SELECT * FROM table1', [
+        'table1',
+      ]);
+
+      /**
+       * Execute another query
+       */
+      const promise2 = dbm.queryWithTableNames('SELECT * FROM table2', [
+        'table1',
+      ]);
+
+      /**
+       * Wait for the queries to complete
+       */
+      await Promise.all([promise1, promise2]);
+
+      await new Promise((resolve) => setTimeout(resolve, 10));
+      /**
+       * Expect instanceManager.terminateDB to not be called
+       */
+      expect(instanceManager.terminateDB).not.toBeCalled();
+
+      /**
+       * wait for 200ms
+       */
+      await new Promise((resolve) => setTimeout(resolve, 200));
+      /**
+       * Expect instanceManager.terminateDB to be called
+       */
+      expect(instanceManager.terminateDB).toBeCalled();
+    });
+
+    it('should not shutdown the db if option is not set', async () => {
+      const instanceManager = new InstanceManager();
+      // If instanceManager.terminateDB is a method
+      jest.spyOn(instanceManager, 'terminateDB');
+
+      // If instanceManager.terminateDB is a function
+      instanceManager.terminateDB = jest.fn();
+      const options: DBMConstructorOptions = {
+        instanceManager: instanceManager,
+        fileManager,
+        logger: log,
+        onEvent: (event) => {
+          console.log(event);
+        },
+      };
+      const dbm = new DBM(options);
+
+      /**
+       * Execute a query
+       */
+      const promise1 = dbm.queryWithTableNames('SELECT * FROM table1', [
+        'table1',
+      ]);
+
+      /**
+       * Execute another query
+       */
+      const promise2 = dbm.queryWithTableNames('SELECT * FROM table2', [
+        'table1',
+      ]);
+
+      /**
+       * Wait for the queries to complete
+       */
+      await Promise.all([promise1, promise2]);
+
+      await new Promise((resolve) => setTimeout(resolve, 10));
+      /**
+       * Expect instanceManager.terminateDB to not be called
+       */
+      expect(instanceManager.terminateDB).not.toBeCalled();
+
+      /**
+       * wait for 200ms
+       */
+      await new Promise((resolve) => setTimeout(resolve, 200));
+      /**
+       * Expect instanceManager.terminateDB to be called
+       */
+      expect(instanceManager.terminateDB).not.toBeCalled();
+    });
+  });
 });