From db0e24086174f2d2b4aee364f36c862b35838002 Mon Sep 17 00:00:00 2001 From: Cristina Poncela Cubeiro <140309543+cristinaponcela@users.noreply.github.com> Date: Mon, 17 Aug 2026 09:28:04 +0200 Subject: [PATCH] feat(sqlite): lease lifecycle and renewal Checked what is usually done, and seems to be common pattern to use half the max time for interval (e.g. Kubernetes leader election docs). --- .../sqlite-node/src/sqlite/repo.ts | 55 ++++++++++++++----- .../sqlite-node/src/sqlite/session.ts | 43 +++++++++++---- .../src/sqlite/session/writer-lease.ts | 17 ++++++ .../sqlite-node/src/sqlite/storage.ts | 4 ++ .../sqlite-node/test/repo.test.ts | 8 ++- 5 files changed, 100 insertions(+), 27 deletions(-) diff --git a/packages/session-backends/sqlite-node/src/sqlite/repo.ts b/packages/session-backends/sqlite-node/src/sqlite/repo.ts index 5cd3df0df..a156b4045 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/repo.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/repo.ts @@ -12,7 +12,7 @@ import { readSingleSessionRow, type SqliteSessionMetadata, } from "./session/session-row.ts"; -import { claimWriterLease, releaseWriterLease } from "./session/writer-lease.ts"; +import { claimWriterLease, releaseWriterLease, renewWriterLease, type WriterLeaseRow } from "./session/writer-lease.ts"; import { SqliteOpenSession } from "./session.ts"; import { SqliteStorage } from "./storage.ts"; import type { SqliteDatabase, SqliteDatabaseFactory } from "./types.ts"; @@ -21,6 +21,7 @@ export const SQLITE_STORAGE_VERSION = 1; export const SQLITE_SESSION_EXTENSION = ".sqlite"; const DEFAULT_WRITER_LEASE_MS = 30_000; +const WRITER_LEASE_RENEW_INTERVAL_MS = DEFAULT_WRITER_LEASE_MS / 2; const FIRST_AVAILABLE_COMMIT_SEQ = 3; export type SqliteSessionCreateOptions = SessionCreateOptions; @@ -69,6 +70,7 @@ export class SqliteSessionRepo { this.reserveId(id); const path = sessionPath(this.directory, id); let db: SqliteDatabase | undefined; + let lease: WriterLeaseRow | undefined; let reservedFile = false; let initialized = false; let session: SqliteOpenSession | undefined; @@ -87,22 +89,25 @@ export class SqliteSessionRepo { ...(options.parentSessionId === undefined ? {} : { parentSessionId: options.parentSessionId }), path, }; - activeDb.transaction(() => { + lease = activeDb.transaction(() => { if (readSessionRowCount(activeDb) !== 0) throw new Error(`SQLite session already exists at ${path}`); insertSessionRow(activeDb, metadata, SQLITE_STORAGE_VERSION, FIRST_AVAILABLE_COMMIT_SEQ); insertInitialMainLaneRegisters(activeDb); - // TODO: Keep this lease until SqliteOpenSession.close() and release only the matching fenced claim. - const lease = claimWriterLease(activeDb, uuidv7(this.now()), this.now(), DEFAULT_WRITER_LEASE_MS); - releaseWriterLease(activeDb, lease.owner_id, lease.fence); + return claimWriterLease(activeDb, uuidv7(this.now()), this.now(), DEFAULT_WRITER_LEASE_MS); }); initialized = true; - session = this.openStorageBackedSession(metadata, activeDb); + session = this.openStorageBackedSession(metadata, activeDb, lease); return session; } catch (error) { if (reservedFile && !initialized) await removeSessionFiles(path, { force: true }); throw error; } finally { if (session === undefined) { + const failedDb = db; + const failedLease = lease; + if (failedDb !== undefined && failedLease !== undefined) { + failedDb.transaction(() => releaseWriterLease(failedDb, failedLease.owner_id, failedLease.fence)); + } db?.close(); this.pendingIds.delete(id); } @@ -112,6 +117,7 @@ export class SqliteSessionRepo { async open(metadata: SqliteSessionMetadata): Promise { this.reserveId(metadata.id); let db: SqliteDatabase | undefined; + let lease: WriterLeaseRow | undefined; let session: SqliteOpenSession | undefined; try { await access(metadata.path); @@ -129,15 +135,19 @@ export class SqliteSessionRepo { if (stored.id !== metadata.id) { throw new Error(`SQLite session path ${metadata.path} contains ${stored.id}, not ${metadata.id}`); } - // TODO: Keep this lease until SqliteOpenSession.close() and release only the matching fenced claim. - const lease = claimWriterLease(activeDb, uuidv7(this.now()), this.now(), DEFAULT_WRITER_LEASE_MS); - releaseWriterLease(activeDb, lease.owner_id, lease.fence); + lease = claimWriterLease(activeDb, uuidv7(this.now()), this.now(), DEFAULT_WRITER_LEASE_MS); return stored; }); - session = this.openStorageBackedSession(stored, activeDb); + if (lease === undefined) throw new Error("Failed to claim SQLite writer lease"); + session = this.openStorageBackedSession(stored, activeDb, lease); return session; } finally { if (session === undefined) { + const failedDb = db; + const failedLease = lease; + if (failedDb !== undefined && failedLease !== undefined) { + failedDb.transaction(() => releaseWriterLease(failedDb, failedLease.owner_id, failedLease.fence)); + } db?.close(); this.pendingIds.delete(metadata.id); } @@ -171,12 +181,27 @@ export class SqliteSessionRepo { await rm(metadata.path, { force: true }); } - private openStorageBackedSession(metadata: SqliteSessionMetadata, db: SqliteDatabase): SqliteOpenSession { - const storage = new SqliteStorage(db, { now: this.now }); + private openStorageBackedSession( + metadata: SqliteSessionMetadata, + db: SqliteDatabase, + initialLease: WriterLeaseRow, + ): SqliteOpenSession { + let lease = initialLease; + const renew = () => { + lease = db.transaction(() => + renewWriterLease(db, lease.owner_id, lease.fence, this.now(), DEFAULT_WRITER_LEASE_MS), + ); + }; + const storage = new SqliteStorage(db, { now: this.now, beforeCommit: renew }); const session = new StorageBackedSession(metadata, storage); - return new SqliteOpenSession(session, () => { - db.close(); - this.pendingIds.delete(metadata.id); + return new SqliteOpenSession(session, { + renewWriterLease: renew, + releaseWriterLease: () => db.transaction(() => releaseWriterLease(db, lease.owner_id, lease.fence)), + renewIntervalMs: WRITER_LEASE_RENEW_INTERVAL_MS, + onClose: () => { + db.close(); + this.pendingIds.delete(metadata.id); + }, }); } diff --git a/packages/session-backends/sqlite-node/src/sqlite/session.ts b/packages/session-backends/sqlite-node/src/sqlite/session.ts index e349422c9..42c993a33 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/session.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/session.ts @@ -14,27 +14,44 @@ import type { } from "@earendil-works/pi-agent-core"; import type { SqliteSessionMetadata } from "./session/session-row.ts"; -/** - * SQLite-specific open-session lifecycle wrapper. - * - * TODO: Keep the writer lease claimed for the lifetime of this open session, - * renew it while open, and release only the matching fenced lease on close. - */ +export interface SqliteOpenSessionOptions { + onClose: () => void; + renewWriterLease: () => void; + releaseWriterLease: () => void; + renewIntervalMs: number; +} + +/** SQLite-specific open-session lifecycle wrapper. */ export class SqliteOpenSession implements Session { readonly metadata: SqliteSessionMetadata; readonly idGenerator: Session["idGenerator"]; private readonly session: Session; private readonly onClose: () => void; + private readonly renewWriterLease: () => void; + private readonly releaseWriterLease: () => void; + private readonly renewalTimer: ReturnType; private readonly admitted = new Set>(); + private leaseError: unknown; private readonly closedError = new Error("Session is closed"); private state: "open" | "closing" | "closed" = "open"; private closePromise: Promise | undefined; - constructor(session: Session, onClose: () => void) { + constructor(session: Session, options: SqliteOpenSessionOptions) { this.session = session; this.metadata = session.metadata; this.idGenerator = session.idGenerator; - this.onClose = onClose; + this.onClose = options.onClose; + this.renewWriterLease = options.renewWriterLease; + this.releaseWriterLease = options.releaseWriterLease; + this.renewalTimer = setInterval(() => { + try { + this.renewWriterLease(); + } catch (error) { + this.leaseError = error; + clearInterval(this.renewalTimer); + } + }, options.renewIntervalMs); + this.renewalTimer.unref?.(); } mutate(lane: string, mutation: (mutator: SessionMutator) => T | Promise): Promise { @@ -153,14 +170,20 @@ export class SqliteOpenSession implements Session { this.closePromise = Promise.allSettled([...this.admitted]) .then(() => this.session.close()) .finally(() => { - this.state = "closed"; - this.onClose(); + clearInterval(this.renewalTimer); + try { + this.releaseWriterLease(); + } finally { + this.state = "closed"; + this.onClose(); + } }); return this.closePromise; } private admit(operation: () => Promise): Promise { if (this.state !== "open") return Promise.reject(this.closedError); + if (this.leaseError !== undefined) return Promise.reject(this.leaseError); let result: Promise; try { result = operation(); diff --git a/packages/session-backends/sqlite-node/src/sqlite/session/writer-lease.ts b/packages/session-backends/sqlite-node/src/sqlite/session/writer-lease.ts index b25134d9a..45913a3d6 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/session/writer-lease.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/session/writer-lease.ts @@ -26,6 +26,23 @@ export function claimWriterLease(db: SqliteDatabase, ownerId: string, now: numbe return { rowid: current.rowid, owner_id: ownerId, fence, expires_at_ms: expiresAt }; } +export function renewWriterLease( + db: SqliteDatabase, + ownerId: string, + fence: number, + now: number, + leaseMs: number, +): WriterLeaseRow { + const expiresAt = now + leaseMs; + const result = sql`UPDATE writer_lease SET expires_at_ms = ${expiresAt} + WHERE owner_id = ${ownerId} AND fence = ${fence}`.run(db); + if (result.changes !== 1) throw new Error("SQLite writer lease was lost"); + const row = sql`SELECT rowid, owner_id, fence, expires_at_ms FROM writer_lease + WHERE owner_id = ${ownerId} AND fence = ${fence}`.get(db); + if (row === undefined) throw new Error("SQLite writer lease was lost"); + return row; +} + export function releaseWriterLease(db: SqliteDatabase, ownerId: string, fence: number): void { sql`DELETE FROM writer_lease WHERE owner_id = ${ownerId} AND fence = ${fence}`.run(db); } diff --git a/packages/session-backends/sqlite-node/src/sqlite/storage.ts b/packages/session-backends/sqlite-node/src/sqlite/storage.ts index 83345a57a..949014560 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/storage.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/storage.ts @@ -23,11 +23,13 @@ import type { SqliteDatabase } from "./types.ts"; export interface SqliteStorageOptions { now?: () => number; + beforeCommit?: () => void; } export class SqliteStorage implements Storage { private readonly db: SqliteDatabase; private readonly now: () => number; + private readonly beforeCommit: () => void; private readonly entryWriter: EntryRowWriter; private readonly usageWriter: UsageLedgerRowWriter; private commitQueue: Promise = Promise.resolve(); @@ -37,6 +39,7 @@ export class SqliteStorage implements Storage { constructor(db: SqliteDatabase, options: SqliteStorageOptions = {}) { this.db = db; this.now = options.now ?? Date.now; + this.beforeCommit = options.beforeCommit ?? (() => undefined); this.entryWriter = new EntryRowWriter(db); this.usageWriter = new UsageLedgerRowWriter(db); } @@ -104,6 +107,7 @@ export class SqliteStorage implements Storage { } private applyCommit(transaction: Transaction): CommitResult { + this.beforeCommit(); return this.db.transaction(() => { const firstSeq = readNextSeq(this.db); const prepared = prepareStorageCommit(transaction, firstSeq, this.now()); diff --git a/packages/session-backends/sqlite-node/test/repo.test.ts b/packages/session-backends/sqlite-node/test/repo.test.ts index 1408df784..a40da7f37 100644 --- a/packages/session-backends/sqlite-node/test/repo.test.ts +++ b/packages/session-backends/sqlite-node/test/repo.test.ts @@ -87,7 +87,7 @@ describe("SqliteSessionRepo", () => { }); }); - it("lists sessions without taking the writer lease", async () => { + it("lists sessions without changing the writer lease", async () => { await withTempDir(async (directory) => { const repo = new SqliteSessionRepo({ directory, @@ -96,10 +96,14 @@ describe("SqliteSessionRepo", () => { }); const session = await repo.create({ id: "session" }); const { metadata } = session; + let leaseBeforeList: unknown[] = []; + await withDb(metadata.path, (db) => { + leaseBeforeList = sql`SELECT owner_id, fence FROM writer_lease`.all(db); + }); await expect(repo.list()).resolves.toMatchObject([{ id: "session", path: metadata.path }]); await withDb(metadata.path, (db) => { - expect(sql`SELECT owner_id, fence FROM writer_lease`.all(db)).toEqual([]); + expect(sql`SELECT owner_id, fence FROM writer_lease`.all(db)).toEqual(leaseBeforeList); }); await session.close(); });