From b69b1e3ec8318d1ef83d9f1edea3a862442c5561 Mon Sep 17 00:00:00 2001
From: gsilvestrin <gsilvestrin@users.noreply.github.com>
Date: Fri, 4 Aug 2023 20:18:23 -0700
Subject: [PATCH] fix(node) Unit tests hangs and don't exit (#396)

---
 node/package.json          |  2 +-
 node/src/index.ts          | 16 +++++++++++++++-
 node/src/remote/index.ts   |  4 ++++
 rust/ffi/node/src/lib.rs   |  1 +
 rust/ffi/node/src/table.rs | 21 ++++++++++++++++++++-
 5 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/node/package.json b/node/package.json
index 6a382db2..0769de35 100644
--- a/node/package.json
+++ b/node/package.json
@@ -8,7 +8,7 @@
     "tsc": "tsc -b",
     "build": "cargo-cp-artifact --artifact cdylib vectordb-node index.node -- cargo build --message-format=json",
     "build-release": "npm run build -- --release",
-    "test": "npm run tsc && mocha -recursive dist/test",
+    "test": "npm run tsc && mocha -recursive dist/test --exit",
     "lint": "eslint native.js src --ext .js,.ts",
     "clean": "rm -rf node_modules *.node dist/",
     "pack-build": "neon pack-build",
diff --git a/node/src/index.ts b/node/src/index.ts
index b05729d9..2dd1ed16 100644
--- a/node/src/index.ts
+++ b/node/src/index.ts
@@ -23,7 +23,7 @@ import { Query } from './query'
 import { isEmbeddingFunction } from './embedding/embedding_function'
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
-const { databaseNew, databaseTableNames, databaseOpenTable, databaseDropTable, tableCreate, tableAdd, tableCreateVectorIndex, tableCountRows, tableDelete } = require('../native.js')
+const { databaseNew, databaseTableNames, databaseOpenTable, databaseDropTable, tableCreate, tableAdd, tableCreateVectorIndex, tableCountRows, tableDelete, tableClose } = require('../native.js')
 
 export { Query }
 export type { EmbeddingFunction }
@@ -215,6 +215,12 @@ export interface Table<T = number[]> {
    * ```
    */
   delete: (filter: string) => Promise<void>
+
+  /**
+   * Immediately closes the connection to this Table. After close is called,
+   * all operations on this Table will fail.
+   */
+  close: () => Promise<void>
 }
 
 /**
@@ -402,6 +408,14 @@ export class LocalTable<T = number[]> implements Table<T> {
   async delete (filter: string): Promise<void> {
     return tableDelete.call(this._tbl, filter)
   }
+
+  /**
+   * Immediately closes the connection to this Table. After close is called,
+   * all operations on this Table will fail.
+   */
+  async close (): Promise<void> {
+    return tableClose.call(this._tbl)
+  }
 }
 
 /// Config to build IVF_PQ index.
diff --git a/node/src/remote/index.ts b/node/src/remote/index.ts
index 4078a04d..0eb9f24d 100644
--- a/node/src/remote/index.ts
+++ b/node/src/remote/index.ts
@@ -165,4 +165,8 @@ export class RemoteTable<T = number[]> implements Table<T> {
   async delete (filter: string): Promise<void> {
     throw new Error('Not implemented')
   }
+
+  async close (): Promise<void> {
+    throw new Error('Not implemented')
+  }
 }
diff --git a/rust/ffi/node/src/lib.rs b/rust/ffi/node/src/lib.rs
index d428e944..2560f75b 100644
--- a/rust/ffi/node/src/lib.rs
+++ b/rust/ffi/node/src/lib.rs
@@ -226,6 +226,7 @@ fn main(mut cx: ModuleContext) -> NeonResult<()> {
     cx.export_function("tableAdd", JsTable::js_add)?;
     cx.export_function("tableCountRows", JsTable::js_count_rows)?;
     cx.export_function("tableDelete", JsTable::js_delete)?;
+    cx.export_function("tableClose", JsTable::js_close)?;
     cx.export_function(
         "tableCreateVectorIndex",
         index::vector::table_create_vector_index,
diff --git a/rust/ffi/node/src/table.rs b/rust/ffi/node/src/table.rs
index e956e2db..e041e18c 100644
--- a/rust/ffi/node/src/table.rs
+++ b/rust/ffi/node/src/table.rs
@@ -38,6 +38,8 @@ impl Finalize for JsTable {}
 pub(crate) enum JsTableMessage {
     // Promise to resolve and callback to be executed
     Callback(Deferred, TableCallback),
+    // Forces to shutdown the thread
+    Close,
 }
 
 impl JsTable {
@@ -56,7 +58,8 @@ impl JsTable {
                 match message {
                     JsTableMessage::Callback(deferred, f) => {
                         f(&mut table, &channel, deferred);
-                    }
+                    },
+                    JsTableMessage::Close => break
                 }
             }
         });
@@ -64,6 +67,14 @@ impl JsTable {
         Ok(Self { tx })
     }
 
+    // It is not necessary to call `close` since the database will be closed when the wrapping
+    // `JsBox` is garbage collected. However, calling `close` allows the process to exit
+    // immediately instead of waiting on garbage collection. This is useful in tests.
+    pub(crate) fn close(&self) -> Result<()> {
+        self.tx.send(JsTableMessage::Close)
+            .map_err(Error::from)
+    }
+
     pub(crate) fn send(
         &self,
         deferred: Deferred,
@@ -215,4 +226,12 @@ impl JsTable {
             .or_throw(&mut cx)?;
         Ok(promise)
     }
+
+    pub(crate) fn js_close(mut cx: FunctionContext) -> JsResult<JsUndefined> {
+        cx.this()
+            .downcast_or_throw::<JsBox<JsTable>, _>(&mut cx)?
+            .close()
+            .or_throw(&mut cx)?;
+        Ok(cx.undefined())
+    }
 }