fix(shared): retain SQLite connection ownership through execution

packages/shared/src/nodeSqliteClient.test.ts

@@ -873,3 +873,134 @@ it.effect("does not warn when the client stops its checkpoint worker", () => {
     assert.lengthOf(warnings, 0);
   }).pipe(Effect.provide(logs));
 });
+it.effect("finishes masked database cleanup after its caller is interrupted", () =>
+  Effect.gen(function* () {
+    const sql = yield* SqlClient.SqlClient;
+    yield* sql`CREATE TABLE cleanup_writes(value INTEGER PRIMARY KEY)`;
+    const entered = yield* Deferred.make<void>();
+    const work = yield* Effect.forkChild(
+      Deferred.succeed(entered, undefined).pipe(
+        Effect.andThen(Effect.never),
+        Effect.onInterrupt(() =>
+          Effect.gen(function* () {
+            yield* sql`INSERT INTO cleanup_writes VALUES (1)`;
+            yield* sql.withTransaction(sql`INSERT INTO cleanup_writes VALUES (2)`);
+          }),
+        ),
+      ),
+      { startImmediately: true },
+    );
+    yield* Deferred.await(entered);
+    yield* Fiber.interrupt(work);
+    assert.deepEqual(yield* sql`SELECT value FROM cleanup_writes ORDER BY value`.values, [
+      [1],
+      [2],
+    ]);
+  }).pipe(Effect.provide(SqliteClient.layer({ filename: ":memory:" }))),
+);
+
+const eventLoopTurn = Effect.promise(() => new Promise<void>((resolve) => setImmediate(resolve)));
+
+it.effect("owns the reader until a standalone statement has finished", () =>
+  withReaderClient((sql) =>
+    Effect.gen(function* () {
+      yield* countEntries(sql).pipe(SqliteClient.readOnly);
+      const worker = startedWorkers.at(-1)!;
+      const send = worker.postMessage.bind(worker);
+      let pending: Parameters<typeof worker.postMessage> | undefined;
+      let interleaved = false;
+      let announce = () => {};
+      const ready = new Promise<void>((resolve) => {
+        announce = resolve;
+      });
+      const spy = vi.spyOn(worker, "postMessage").mockImplementation((...args) => {
+        if (pending === undefined) {
+          pending = args;
+          announce();
+        } else {
+          interleaved = true;
+          send(...args);
+        }
+      });
+      const release = () => {
+        spy.mockRestore();
+        if (pending !== undefined) send(...pending);
+        pending = undefined;
+      };
+      yield* Effect.gen(function* () {
+        const read = yield* Effect.forkChild(countEntries(sql).pipe(SqliteClient.readOnly), {
+          startImmediately: true,
+        });
+        yield* Effect.promise(() => ready);
+        const entered = yield* Deferred.make<void>();
+        const transaction = yield* Effect.forkChild(
+          sql.withTransaction(Deferred.succeed(entered, undefined)).pipe(SqliteClient.readOnly),
+          { startImmediately: true },
+        );
+        // Let the competing acquisition run while the first request remains held.
+        yield* eventLoopTurn;
+        assert.isFalse(interleaved, "the transaction must wait for the standalone read");
+        assert.isFalse(yield* Deferred.isDone(entered));
+        release();
+        assert.equal(yield* Fiber.join(read), 1);
+        yield* Fiber.join(transaction);
+        assert.equal(yield* countEntries(sql).pipe(SqliteClient.readOnly), 1);
+      }).pipe(Effect.ensuring(Effect.sync(release)));
+    }),
+  ),
+);
+
+it.effect.each([
+  { reservation: false, masked: false },
+  { reservation: false, masked: true },
+  { reservation: true, masked: false },
+  { reservation: true, masked: true },
+])(
+  "preserves caller cancellation while waiting (reservation=$reservation, masked=$masked)",
+  ({ reservation, masked }) =>
+    Effect.gen(function* () {
+      const sql = yield* SqlClient.SqlClient;
+      yield* sql`CREATE TABLE entries(value INTEGER)`;
+      const held = yield* Deferred.make<void>();
+      const release = yield* Deferred.make<void>();
+      const holder = yield* Effect.forkChild(
+        sql.withTransaction(
+          Deferred.succeed(held, undefined).pipe(Effect.andThen(Deferred.await(release))),
+        ),
+        { startImmediately: true },
+      );
+      yield* Deferred.await(held);
+      const entered = yield* Deferred.make<void>();
+      const body = sql`INSERT INTO entries VALUES (1)`.pipe(
+        Effect.andThen(Deferred.succeed(entered, undefined)),
+      );
+      const acquire = reservation
+        ? Effect.scoped(
+            Effect.gen(function* () {
+              const connection = yield* sql.reserve;
+              yield* connection.executeUnprepared("INSERT INTO entries VALUES (1)", [], undefined);
+              yield* Deferred.succeed(entered, undefined);
+            }),
+          )
+        : sql.withTransaction(body);
+      const waiter = yield* Effect.forkChild(masked ? Effect.uninterruptible(acquire) : acquire, {
+        startImmediately: true,
+      });
+      const cancelled = yield* Effect.forkChild(Fiber.interrupt(waiter), {
+        startImmediately: true,
+      });
+      yield* Effect.gen(function* () {
+        yield* eventLoopTurn;
+        assert.equal(cancelled.pollUnsafe() !== undefined, !masked);
+        assert.isFalse(yield* Deferred.isDone(entered));
+      }).pipe(Effect.ensuring(Deferred.succeed(release, undefined)));
+      yield* Fiber.join(holder);
+      yield* Fiber.join(cancelled);
+      assert.equal(yield* Deferred.isDone(entered), masked);
+      yield* sql.withTransaction(sql`INSERT INTO entries VALUES (2)`);
+      assert.deepEqual(
+        yield* sql`SELECT value FROM entries ORDER BY value`.values,
+        masked ? [[1], [2]] : [[2]],
+      );
+    }).pipe(Effect.provide(SqliteClient.layer({ filename: ":memory:" }))),
+);

packages/shared/src/nodeSqliteClient.ts

@@ -14,14 +14,15 @@ import * as Option from "effect/Option";
 import * as Predicate from "effect/Predicate";
 import * as References from "effect/References";
 import * as Schema from "effect/Schema";
+import * as Scheduler from "effect/Scheduler";
 import * as Scope from "effect/Scope";
 import * as Semaphore from "effect/Semaphore";
 import * as Context from "effect/Context";
 import * as Stream from "effect/Stream";
 import * as Tracer from "effect/Tracer";
 import * as Reactivity from "effect/reactivity/Reactivity";
 import * as Client from "effect/sql/SqlClient";
-import type { Connection } from "effect/sql/SqlConnection";
+import type { Borrower, Connection } from "effect/sql/SqlConnection";
 import { SqlError, classifySqliteError } from "effect/sql/SqlError";
 import * as Statement from "effect/sql/Statement";
 
@@ -72,10 +73,9 @@ const annotateAcquired = (fiber: Fiber.Fiber<unknown, unknown>, acquired: Acquir
   span.attribute("db.lock_wait_ms", acquired.lockWaitMs);
 };
 
-const TransactionAcquire = Context.Reference<{ acquired?: Acquired } | undefined>(
-  "@t3tools/shared/nodeSqliteClient/TransactionAcquire",
-  { defaultValue: () => undefined },
-);
+const TransactionAcquire = Context.Reference<
+  { acquired?: Acquired; restore: typeof Effect.interruptible | undefined } | undefined
+>("@t3tools/shared/nodeSqliteClient/TransactionAcquire", { defaultValue: () => undefined });
 
 export class UnsupportedNodeSqliteVersionError extends Schema.TaggedError<UnsupportedNodeSqliteVersionError>()(
   "UnsupportedNodeSqliteVersionError",
@@ -773,43 +773,48 @@ const make = Effect.fn("makeWithDatabase")(function* (
     lockWaitMs: Number((clock.monotonicTimeNanosUnsafe() - start) / 1000n) / 1000,
   });
 
-  const acquirer = Effect.flatMap(route, (target) =>
-    Effect.withFiber((fiber) => {
-      if (
-        !fiber.getRef(References.TracerEnabled) ||
-        !fiber.getRef(Statement.SpanPropagationEnabled)
-      ) {
-        return target.lock.withPermits(1)(Effect.succeed(target.connection));
-      }
-      const clock = fiber.getRef(Clock.Clock);
-      const start = clock.monotonicTimeNanosUnsafe();
-      return target.lock.withPermits(1)(
-        Effect.sync(() => {
-          annotateAcquired(fiber, acquired(target, clock, start));
-          return target.connection;
-        }),
-      );
-    }),
-  );
+  const borrower: Borrower = (use) =>
+    Effect.flatMap(route, (target) =>
+      Effect.withFiber((fiber) => {
+        const effect = Effect.suspend(() => use(target.connection));
+        if (
+          !fiber.getRef(References.TracerEnabled) ||
+          !fiber.getRef(Statement.SpanPropagationEnabled)
+        ) {
+          return target.lock.withPermits(1)(effect);
+        }
+        const clock = fiber.getRef(Clock.Clock);
+        const start = clock.monotonicTimeNanosUnsafe();
+        return target.lock.withPermits(1)(
+          Effect.andThen(
+            Effect.sync(() => annotateAcquired(fiber, acquired(target, clock, start))),
+            effect,
+          ),
+        );
+      }),
+    );
   const transactionAcquirer = Effect.flatMap(route, (target) =>
     Effect.uninterruptibleMask((restore) => {
       const fiber = Fiber.getCurrent()!;
       const scope = Context.getUnsafe(fiber.context, Scope.Scope);
       const clock = fiber.getRef(Clock.Clock);
       const start = clock.monotonicTimeNanosUnsafe();
+      const pending = fiber.getRef(TransactionAcquire);
+      const acquire = Effect.tap((pending?.restore ?? restore)(target.lock.take(1)), () => {
+        if (pending !== undefined) pending.acquired = acquired(target, clock, start);
+        return Scope.addFinalizer(scope, target.lock.release(1));
+      });
       return Effect.as(
-        Effect.tap(restore(target.lock.take(1)), () => {
-          const pending = fiber.getRef(TransactionAcquire);
-          if (pending !== undefined) pending.acquired = acquired(target, clock, start);
-          return Scope.addFinalizer(scope, target.lock.release(1));
-        }),
+        // A scheduler yield after taking the permit must not precede its release finalizer.
+        Effect.provideService(acquire, Scheduler.PreventSchedulerYield, true),
         target.connection,
       );
     }),
   );
 
   const client = yield* Client.make({
-    acquirer,
+    acquirer: transactionAcquirer,
+    borrower,
     compiler,
     transactionAcquirer,
     // A deferred BEGIN only takes the write lock at the first write. If another
@@ -825,13 +830,19 @@ const make = Effect.fn("makeWithDatabase")(function* (
   const withTransaction = client.withTransaction;
   return Object.assign(client, {
     withTransaction: <A, E, R>(effect: Effect.Effect<A, E, R>) =>
-      Effect.suspend(() => {
-        const pending: { acquired?: Acquired } = {};
+      Effect.uninterruptibleMask((restore) => {
+        // SqlClient masks setup; acquisition must preserve the caller's interruptibility.
+        const pending: { acquired?: Acquired; restore: typeof Effect.interruptible | undefined } = {
+          restore,
+        };
         const body = Effect.withFiber<A, E, R>((fiber) => {
           annotateAcquired(fiber, pending.acquired);
+          pending.restore = undefined;
           return effect;
         });
-        return withTransaction(body).pipe(Effect.provideService(TransactionAcquire, pending));
+        return restore(
+          withTransaction(body).pipe(Effect.provideService(TransactionAcquire, pending)),
+        );
       }),
   });
 });