Managing file in duckdb for better swapping (#42)

* WIP

* Adding max allowed memory in file

* Fix test

* Adding test

* Upgrade version

* Update comments

Author: itsTalwar

meerkat-dbm/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@devrev/meerkat-dbm",
-  "version": "0.0.134",
+  "version": "0.0.135",
   "dependencies": {
     "tslib": "^2.3.0",
     "@duckdb/duckdb-wasm": "^1.28.0",

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

@@ -100,6 +100,10 @@ export class MockFileManager implements FileManagerType {
 
     return data;
   }
+
+  onDBShutdownHandler = jest.fn(async () => {
+    // do nothing
+  });
 }
 
 const mockDB = {
@@ -307,6 +311,7 @@ describe('DBM', () => {
       /**
        * Expect instanceManager.terminateDB to be called
        */
+      expect(fileManager.onDBShutdownHandler).toBeCalled();
       expect(instanceManager.terminateDB).toBeCalled();
     });
 

meerkat-dbm/src/dbm/dbm.ts

@@ -89,6 +89,7 @@ export class DBM {
     if (this.onDuckdbShutdown) {
       this.onDuckdbShutdown();
     }
+    await this.fileManager.onDBShutdownHandler();
     await this.instanceManager.terminateDB();
   }
 
@@ -186,26 +187,6 @@ export class DBM {
       metadata: options?.metadata,
     });
 
-    /**
-     * Unload all the files from the database, so that the files can be removed from memory
-     */
-    const startUnmountTime = Date.now();
-    await this.fileManager.unmountFileBufferByTableNames(tableNames);
-    const endUnmountTime = Date.now();
-
-    this.logger.debug(
-      'Time spent in unmounting files:',
-      endUnmountTime - startUnmountTime,
-      'ms',
-      query
-    );
-
-    this._emitEvent({
-      event_name: 'unmount_file_buffer_duration',
-      duration: endUnmountTime - startUnmountTime,
-      metadata: options?.metadata,
-    });
-
     return result;
   }
 

meerkat-dbm/src/file-manager/file-manager-type.ts

@@ -14,7 +14,6 @@ export interface FileManagerType {
   registerFileBuffer: (props: FileBufferStore) => Promise<void>;
   getFileBuffer: (name: string) => Promise<Uint8Array | undefined>;
   mountFileBufferByTableNames: (tableName: string[]) => Promise<void>;
-  unmountFileBufferByTableNames: (tableName: string[]) => Promise<void>;
   getFilesByTableName(tableName: string): Promise<FileData[]>;
   dropFilesByTableName(tableName: string, fileNames: string[]): Promise<void>;
   getFilesNameForTables(tableNames: string[]): Promise<
@@ -23,11 +22,18 @@ export interface FileManagerType {
       files: string[];
     }[]
   >;
+  onDBShutdownHandler: () => Promise<void>;
 }
 
 export interface FileManagerConstructorOptions {
   fetchTableFileBuffers: (tableName: string) => Promise<FileBufferStore[]>;
   instanceManager: InstanceManagerType;
+  options?: {
+    /**
+     * Maximum size of the file in DB in bytes
+     */
+    maxFileSize?: number;
+  };
 }
 
 export const FILE_TYPES = {

meerkat-dbm/src/file-manager/indexed-db/file-registerer.spec.ts

@@ -0,0 +1,82 @@
+import { InstanceManagerType } from '../../dbm/instance-manager';
+import { FileRegisterer } from './file-registerer';
+
+const mockDB = {
+  registerFileBuffer: jest.fn(async (fileName: string, buffer: Uint8Array) => {
+    return new Promise((resolve) => {
+      setTimeout(() => {
+        resolve([fileName]);
+      }, 200);
+    });
+  }),
+  registerEmptyFileBuffer: jest.fn(async (fileName: string) => {
+    return new Promise((resolve) => {
+      setTimeout(() => {
+        resolve([fileName]);
+      }, 200);
+    });
+  }),
+};
+
+describe('FileRegisterer', () => {
+  let db: typeof mockDB;
+  let instanceManager: InstanceManagerType;
+  let fileRegisterer: FileRegisterer;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    db = mockDB;
+    instanceManager = {
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      //@ts-ignore
+      getDB: async () => {
+        return db;
+      },
+      terminateDB: async () => {
+        return;
+      },
+    };
+    fileRegisterer = new FileRegisterer({ instanceManager });
+  });
+
+  it('should register file buffer', async () => {
+    const fileName = 'test.txt';
+    const buffer = new Uint8Array();
+    await fileRegisterer.registerFileBuffer(fileName, buffer);
+    expect(db.registerFileBuffer).toHaveBeenCalledWith(fileName, buffer);
+    expect(fileRegisterer.isFileRegisteredInDB(fileName)).toBeTruthy();
+  });
+
+  it('should register empty file buffer', async () => {
+    const fileName = 'test.txt';
+    await fileRegisterer.registerEmptyFileBuffer(fileName);
+    expect(db.registerEmptyFileBuffer).toHaveBeenCalledWith(fileName);
+    expect(fileRegisterer.isFileRegisteredInDB(fileName)).toBeFalsy();
+  });
+
+  it('should check if file is registered in DB', () => {
+    const fileName = 'test.txt';
+    expect(fileRegisterer.isFileRegisteredInDB(fileName)).toBeFalsy();
+  });
+
+  it('should flush file cache', () => {
+    fileRegisterer.flushFileCache();
+    expect(fileRegisterer.getAllFilesInDB()).toEqual([]);
+  });
+
+  it('should get total byte length', () => {
+    const fileName = 'test.txt';
+    const buffer = new Uint8Array(5);
+    fileRegisterer.registerFileBuffer(fileName, buffer);
+    expect(fileRegisterer.totalByteLength()).toEqual(5);
+  });
+
+  it('should get all files in DB', () => {
+    const fileName1 = 'test1.txt';
+    const fileName2 = 'test2.txt';
+    const buffer = new Uint8Array();
+    fileRegisterer.registerFileBuffer(fileName1, buffer);
+    fileRegisterer.registerFileBuffer(fileName2, buffer);
+    expect(fileRegisterer.getAllFilesInDB()).toEqual([fileName1, fileName2]);
+  });
+});

meerkat-dbm/src/file-manager/indexed-db/file-registerer.ts

@@ -0,0 +1,78 @@
+import { AsyncDuckDB } from '@duckdb/duckdb-wasm';
+import { InstanceManagerType } from '../../dbm/instance-manager';
+
+interface FileRegistererConstructorOptions {
+  instanceManager: InstanceManagerType;
+}
+
+export interface FileRegistererType {
+  registerFileBuffer: AsyncDuckDB['registerFileBuffer'];
+  registerEmptyFileBuffer: AsyncDuckDB['registerEmptyFileBuffer'];
+  isFileRegisteredInDB: (fileName: string) => boolean;
+  flushFileCache: () => void;
+  totalByteLength: () => number;
+  getAllFilesInDB: () => string[];
+}
+
+export class FileRegisterer implements FileRegistererType {
+  instanceManager: InstanceManagerType;
+
+  private registeredFilesSet = new Map<
+    string,
+    {
+      byteLength: number;
+    }
+  >();
+
+  constructor({ instanceManager }: FileRegistererConstructorOptions) {
+    this.instanceManager = instanceManager;
+  }
+
+  registerFileBuffer: AsyncDuckDB['registerFileBuffer'] = async (
+    fileName,
+    buffer
+  ) => {
+    if (this.registeredFilesSet.has(fileName)) {
+      return;
+    }
+    this.registeredFilesSet.set(fileName, {
+      byteLength: buffer.byteLength,
+    });
+    return this.instanceManager
+      .getDB()
+      .then((db) => db.registerFileBuffer(fileName, buffer));
+  };
+
+  registerEmptyFileBuffer: AsyncDuckDB['registerEmptyFileBuffer'] = async (
+    fileName
+  ) => {
+    await this.instanceManager
+      .getDB()
+      .then((db) => db.registerEmptyFileBuffer(fileName));
+
+    this.registeredFilesSet.delete(fileName);
+  };
+
+  isFileRegisteredInDB(fileName: string): boolean {
+    return this.registeredFilesSet.has(fileName);
+  }
+
+  /**
+   * This function on purpose does not clear the files from the DB
+   * As this is called when the DB is shutdown and we just want to tell our file registerer to clear the cache
+   */
+  flushFileCache(): void {
+    this.registeredFilesSet.clear();
+  }
+
+  totalByteLength(): number {
+    return Array.from(this.registeredFilesSet.values()).reduce(
+      (total, { byteLength }) => total + byteLength,
+      0
+    );
+  }
+
+  getAllFilesInDB(): string[] {
+    return Array.from(this.registeredFilesSet.keys());
+  }
+}