From 9a3f59aede9a6356efc13c8e8b31f6592b7aa64b Mon Sep 17 00:00:00 2001
From: Vincent Weevers <mail@vincentweevers.nl>
Date: Sun, 19 Sep 2021 01:26:19 +0200
Subject: [PATCH] Prevent GC of db during `clear()` and other operations

---
 .github/workflows/test.yml |  2 ++
 binding.cc                 | 19 ++++++++++++---
 test/clear-gc-test.js      | 47 ++++++++++++++++++++++++++++++++++++++
 test/gc.js                 |  1 +
 4 files changed, 66 insertions(+), 3 deletions(-)
 create mode 100644 test/clear-gc-test.js

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 444f20f2..15227ae9 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -39,3 +39,5 @@ jobs:
         uses: GabrielBB/xvfb-action@v1
         with:
           run: npm run test-electron
+      - name: Test GC
+        run: npm run test-gc
diff --git a/binding.cc b/binding.cc
index cc6e57cb..94a4c1d3 100644
--- a/binding.cc
+++ b/binding.cc
@@ -361,9 +361,14 @@ struct Database {
       filterPolicy_(leveldb::NewBloomFilterPolicy(10)),
       currentIteratorId_(0),
       pendingCloseWorker_(NULL),
+      ref_(NULL),
       priorityWork_(0) {}
 
   ~Database () {
+    if (ref_ != NULL) {
+      napi_delete_reference(env_, ref_);
+    }
+
     if (db_ != NULL) {
       delete db_;
       db_ = NULL;
@@ -444,11 +449,13 @@ struct Database {
   }
 
   void IncrementPriorityWork () {
-    ++priorityWork_;
+    napi_reference_ref(env_, ref_, &priorityWork_);
   }
 
   void DecrementPriorityWork () {
-    if (--priorityWork_ == 0 && pendingCloseWorker_ != NULL) {
+    napi_reference_unref(env_, ref_, &priorityWork_);
+
+    if (priorityWork_ == 0 && pendingCloseWorker_ != NULL) {
       pendingCloseWorker_->Queue();
       pendingCloseWorker_ = NULL;
     }
@@ -465,6 +472,7 @@ struct Database {
   uint32_t currentIteratorId_;
   BaseWorker *pendingCloseWorker_;
   std::map< uint32_t, Iterator * > iterators_;
+  napi_ref ref_;
 
 private:
   uint32_t priorityWork_;
@@ -828,11 +836,16 @@ NAPI_METHOD(db_init) {
   NAPI_STATUS_THROWS(napi_create_external(env, database,
                                           FinalizeDatabase,
                                           NULL, &result));
+
+  // Reference counter to prevent GC of database while priority workers are active
+  NAPI_STATUS_THROWS(napi_create_reference(env, result, 0, &database->ref_));
+
   return result;
 }
 
 /**
  * Worker class for opening a database.
+ * TODO: shouldn't this be a PriorityWorker?
  */
 struct OpenWorker final : public BaseWorker {
   OpenWorker (napi_env env,
@@ -1133,7 +1146,6 @@ struct ClearWorker final : public PriorityWorker {
   }
 
   ~ClearWorker () {
-    // TODO: write GC tests
     delete baseIterator_;
     delete writeOptions_;
   }
@@ -1476,6 +1488,7 @@ struct EndWorker final : public BaseWorker {
   }
 
   void HandleOKCallback () override {
+    // TODO: if we don't use EndWorker, do we still delete the reference?
     napi_delete_reference(env_, iterator_->Detach());
     BaseWorker::HandleOKCallback();
   }
diff --git a/test/clear-gc-test.js b/test/clear-gc-test.js
new file mode 100644
index 00000000..2794878a
--- /dev/null
+++ b/test/clear-gc-test.js
@@ -0,0 +1,47 @@
+'use strict'
+
+const test = require('tape')
+const testCommon = require('./common')
+const sourceData = []
+
+for (let i = 0; i < 1e3; i++) {
+  sourceData.push({
+    type: 'put',
+    key: i.toString(),
+    value: Math.random().toString()
+  })
+}
+
+test('db without ref does not get GCed while clear() is in progress', function (t) {
+  t.plan(4)
+
+  let db = testCommon.factory()
+
+  db.open(function (err) {
+    t.ifError(err, 'no open error')
+
+    // Insert test data
+    db.batch(sourceData.slice(), function (err) {
+      t.ifError(err, 'no batch error')
+
+      // Start async work
+      db.clear(function () {
+        t.pass('got callback')
+
+        // Give GC another chance to run, to rule out other issues.
+        setImmediate(function () {
+          if (global.gc) global.gc()
+          t.pass()
+        })
+      })
+
+      // Remove reference. The db should not get garbage collected
+      // until after the clear() callback, thanks to a napi_ref.
+      db = null
+
+      // Useful for manual testing with "node --expose-gc".
+      // The pending tap assertion may also allow GC to kick in.
+      if (global.gc) global.gc()
+    })
+  })
+})
diff --git a/test/gc.js b/test/gc.js
index a5655f91..956cc2ad 100644
--- a/test/gc.js
+++ b/test/gc.js
@@ -10,3 +10,4 @@ if (!global.gc) {
 require('./cleanup-hanging-iterators-test')
 require('./iterator-gc-test')
 require('./chained-batch-gc-test')
+require('./clear-gc-test')