fix(unstable): kv watch should stop when db is closed (#21665)

Fixes #21634.

2 files changed, 30 insertions, 5 deletions

diff --git a/cli/tests/unit/kv_test.ts b/cli/tests/unit/kv_test.ts
index 28c913f21..4963882e1 100644
--- a/cli/tests/unit/kv_test.ts
+++ b/cli/tests/unit/kv_test.ts
@@ -2248,3 +2248,23 @@ dbTest("set with key versionstamp suffix", async (db) => {
     "expected string, number, bigint, ArrayBufferView, boolean",
   );
 });
+
+Deno.test({
+  name: "watch should stop when db closed",
+  async fn() {
+    const db = await Deno.openKv(":memory:");
+
+    const watch = db.watch([["a"]]);
+    const completion = (async () => {
+      for await (const _item of watch) {
+        // pass
+      }
+    })();
+
+    setTimeout(() => {
+      db.close();
+    }, 100);
+
+    await completion;
+  },
+});
diff --git a/ext/kv/lib.rs b/ext/kv/lib.rs
index e78526856..032a16863 100644
--- a/ext/kv/lib.rs
+++ b/ext/kv/lib.rs
@@ -444,16 +444,21 @@ async fn op_kv_watch_next(
   let cancel_handle = resource.cancel_handle.clone();
   let stream = RcRef::map(resource, |r| &r.stream)
     .borrow_mut()
-    .or_cancel(db_cancel_handle)
-    .or_cancel(cancel_handle)
+    .or_cancel(db_cancel_handle.clone())
+    .or_cancel(cancel_handle.clone())
     .await;
   let Ok(Ok(mut stream)) = stream else {
     return Ok(None);
   };
 
-  // doesn't need a cancel handle because the stream ends when the database
-  // connection is closed
-  let Some(res) = stream.next().await else {
+  // We hold a strong reference to `resource`, so we can't rely on the stream
+  // being dropped when the db connection is closed
+  let Ok(Ok(Some(res))) = stream
+    .next()
+    .or_cancel(db_cancel_handle)
+    .or_cancel(cancel_handle)
+    .await
+  else {
     return Ok(None);
   };