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).

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<SqliteOpenSession> {
 		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);
+			},
 		});
 	}
 

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<SqliteSessionMetadata> {
 	readonly metadata: SqliteSessionMetadata;
 	readonly idGenerator: Session<SqliteSessionMetadata>["idGenerator"];
 	private readonly session: Session<SqliteSessionMetadata>;
 	private readonly onClose: () => void;
+	private readonly renewWriterLease: () => void;
+	private readonly releaseWriterLease: () => void;
+	private readonly renewalTimer: ReturnType<typeof setInterval>;
 	private readonly admitted = new Set<Promise<unknown>>();
+	private leaseError: unknown;
 	private readonly closedError = new Error("Session is closed");
 	private state: "open" | "closing" | "closed" = "open";
 	private closePromise: Promise<void> | undefined;
 
-	constructor(session: Session<SqliteSessionMetadata>, onClose: () => void) {
+	constructor(session: Session<SqliteSessionMetadata>, 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<T>(lane: string, mutation: (mutator: SessionMutator) => T | Promise<T>): Promise<T> {
@@ -153,14 +170,20 @@ export class SqliteOpenSession implements Session<SqliteSessionMetadata> {
 		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<T>(operation: () => Promise<T>): Promise<T> {
 		if (this.state !== "open") return Promise.reject(this.closedError);
+		if (this.leaseError !== undefined) return Promise.reject(this.leaseError);
 		let result: Promise<T>;
 		try {
 			result = operation();

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<WriterLeaseRow>(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);
 }

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<void> = 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());