From 7da9e7c45d4068f895c08fe8e2d4cdee472bd796 Mon Sep 17 00:00:00 2001
From: Jarred Sumner <jarred@jarredsumner.com>
Date: Mon, 5 Aug 2024 22:07:49 -0700
Subject: [PATCH] Add test to #13082 and use `WTF_MAKE_FAST_ALLOCATED` (#13105)

Co-authored-by: FuPeiJiang <42662615+FuPeiJiang@users.noreply.github.com>
---
 src/bun.js/bindings/sqlite/JSSQLStatement.cpp | 40 ++++++++++++++++---
 test/js/bun/sqlite/sqlite.test.js             | 25 ++++++++++++
 2 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/src/bun.js/bindings/sqlite/JSSQLStatement.cpp b/src/bun.js/bindings/sqlite/JSSQLStatement.cpp
index 9c1214b5fe..21a8950f75 100644
--- a/src/bun.js/bindings/sqlite/JSSQLStatement.cpp
+++ b/src/bun.js/bindings/sqlite/JSSQLStatement.cpp
@@ -199,17 +199,38 @@ static inline JSC::JSValue jsBigIntFromSQLite(JSC::JSGlobalObject* globalObject,
         return {};                                                                                                 \
     }
 
+DECLARE_ALLOCATOR_WITH_HEAP_IDENTIFIER(VersionSqlite3);
+
 class VersionSqlite3 {
+    WTF_MAKE_FAST_ALLOCATED_WITH_HEAP_IDENTIFIER(VersionSqlite3);
+
 public:
     explicit VersionSqlite3(sqlite3* db)
         : db(db)
         , version(0)
+        , reference_count(1)
     {
     }
     sqlite3* db;
     std::atomic<uint64_t> version;
+    size_t reference_count;
+
+    void release()
+    {
+        ASSERT(reference_count > 0);
+        --reference_count;
+        if (reference_count == 0) {
+            if (!db) {
+                return;
+            }
+            sqlite3_close_v2(db);
+            db = nullptr;
+        }
+    };
 };
 
+DEFINE_ALLOCATOR_WITH_HEAP_IDENTIFIER(VersionSqlite3);
+
 class SQLiteSingleton {
 public:
     Vector<VersionSqlite3*> databases;
@@ -400,6 +421,9 @@ public:
     {
         Structure* structure = globalObject->JSSQLStatementStructure();
         JSSQLStatement* ptr = new (NotNull, JSC::allocateCell<JSSQLStatement>(globalObject->vm())) JSSQLStatement(structure, *globalObject, stmt, version_db, memorySizeChange);
+        if (version_db) {
+            ++version_db->reference_count;
+        }
         ptr->finishCreation(globalObject->vm());
         return ptr;
     }
@@ -1616,12 +1640,7 @@ JSC_DEFINE_HOST_FUNCTION(jsSQLStatementOpenStatementFunction, (JSC::JSGlobalObje
     databases().append(new VersionSqlite3(db));
     if (finalizationTarget.isObject()) {
         vm.heap.addFinalizer(finalizationTarget.getObject(), [index](JSC::JSCell* ptr) -> void {
-            auto* db = databases()[index];
-            if (!db->db) {
-                return;
-            }
-            sqlite3_close_v2(db->db);
-            databases()[index]->db = nullptr;
+            databases()[index]->release();
         });
     }
     RELEASE_AND_RETURN(scope, JSValue::encode(jsNumber(index)));
@@ -2225,6 +2244,11 @@ JSC_DEFINE_HOST_FUNCTION(jsSQLStatementExecuteStatementFunctionRun, (JSC::JSGlob
         DO_REBIND(arg0);
     }
 
+    if (UNLIKELY(!castedThis->version_db->db)) {
+        throwException(lexicalGlobalObject, scope, createError(lexicalGlobalObject, "Database has closed"_s));
+        return JSValue::encode(JSC::jsUndefined());
+    }
+
     int total_changes_before = sqlite3_total_changes(castedThis->version_db->db);
 
     int status = sqlite3_step(stmt);
@@ -2393,6 +2417,10 @@ JSSQLStatement::~JSSQLStatement()
         columnNames->releaseData();
         this->columnNames = nullptr;
     }
+
+    if (this->version_db) {
+        this->version_db->release();
+    }
 }
 
 void JSSQLStatement::analyzeHeap(JSCell* cell, HeapAnalyzer& analyzer)
diff --git a/test/js/bun/sqlite/sqlite.test.js b/test/js/bun/sqlite/sqlite.test.js
index e6d939607a..fdac0b6af0 100644
--- a/test/js/bun/sqlite/sqlite.test.js
+++ b/test/js/bun/sqlite/sqlite.test.js
@@ -1261,3 +1261,28 @@ it("reports changes in Statement#run", () => {
   expect(db.prepare(sql).run().changes).toBe(2);
   expect(db.query(sql).run().changes).toBe(2);
 });
+
+it("#13082", async () => {
+  async function run() {
+    const stmt = (() => {
+      const db = new Database(":memory:");
+      let stmt = db.prepare("select 1");
+      db.close();
+      return stmt;
+    })();
+    Bun.gc(true);
+    await Bun.sleep(100);
+    Bun.gc(true);
+    stmt.all();
+    stmt.get();
+    stmt.run();
+  }
+
+  const count = 100;
+  const runs = new Array(count);
+  for (let i = 0; i < count; i++) {
+    runs[i] = run();
+  }
+
+  await Promise.allSettled(runs);
+});