@beignet/core/src/locks/index.ts

Core framework primitives for Beignet

/**
 * @beignet/core/locks
 *
 * Provider-neutral distributed lock and lease primitives for Beignet apps.
 */

import {
  createProvider,
  createProviderInstrumentation,
} from "../providers/index.js";

type MaybePromise<T> = T | Promise<T>;

/**
 * Metadata attached to lock acquisition instrumentation and diagnostics.
 */
export type LeaseMetadata = Record<
  string,
  string | number | boolean | null | undefined
>;

/**
 * Options for acquiring a lease.
 */
export type LeaseAcquireOptions = {
  /**
   * Lease time-to-live in milliseconds. The provider must treat ownership as
   * expired after this duration unless the lease is renewed.
   */
  ttlMs: number;
  /**
   * How long to wait for the lease before returning a non-acquired result.
   *
   * @default 0
   */
  waitMs?: number;
  /**
   * Delay between acquisition attempts while waiting.
   *
   * @default 50
   */
  retryDelayMs?: number;
  /**
   * Optional owner token supplied by the caller. Providers generate one when
   * omitted. Release and renew operations must only succeed for the current
   * owner token.
   */
  ownerToken?: string;
  /**
   * Optional diagnostics metadata. Providers should not use this for lock
   * correctness.
   */
  metadata?: LeaseMetadata;
};

/**
 * Options for renewing a lease.
 */
export type LeaseRenewOptions = {
  /**
   * New lease time-to-live in milliseconds. Defaults to the original ttlMs
   * used to acquire the lease.
   */
  ttlMs?: number;
};

/**
 * State required to restore a lease handle in another runtime invocation.
 */
export type LeaseRestoreOptions = {
  /**
   * TTL to use when `renew()` is called without an explicit override.
   */
  ttlMs: number;
  /**
   * Persisted expiry when the caller has it. Omitted values stay unknown.
   */
  expiresAt?: Date;
  /**
   * Persisted fencing token when the caller has it. Omitted values stay
   * unknown rather than being fabricated.
   */
  fencingToken?: string | number;
};

/**
 * Active lease handle returned by a lock provider.
 */
export type LeaseHandle = {
  key: string;
  ownerToken: string;
  expiresAt?: Date;
  /**
   * Optional monotonically increasing token. Apps can store this with writes to
   * reject stale lease owners when the underlying resource supports fencing.
   */
  fencingToken?: string | number;
  renew(options?: LeaseRenewOptions): Promise<boolean>;
  release(): Promise<boolean>;
};

/**
 * Lease acquisition result.
 */
export type LeaseAcquireResult =
  | { acquired: true; lease: LeaseHandle }
  | { acquired: false; reason: "unavailable" | "timeout" };

/**
 * App-facing lock port.
 */
export type LocksPort = {
  acquire(
    key: string,
    options: LeaseAcquireOptions,
  ): Promise<LeaseAcquireResult>;
  withLease<T>(
    key: string,
    options: LeaseAcquireOptions,
    fn: (ctx: { lease: LeaseHandle }) => MaybePromise<T>,
  ): Promise<T | undefined>;
  restore(
    key: string,
    ownerToken: string,
    options: LeaseRestoreOptions,
  ): LeaseHandle;
  forceRelease(key: string): Promise<boolean>;
};

/**
 * Captured memory lease state exposed for tests.
 */
export type MemoryLeaseRecord = {
  key: string;
  ownerToken: string;
  expiresAt: Date;
  ttlMs: number;
  fencingToken: number;
  metadata?: LeaseMetadata;
};

/**
 * In-memory locks port exposed for assertions in tests.
 */
export type MemoryLocksPort = LocksPort & {
  leases: Map<string, MemoryLeaseRecord>;
  reset(): void;
};

/**
 * Options for the in-memory locks adapter.
 */
export type CreateMemoryLocksOptions = {
  now?: () => Date;
  createOwnerToken?: () => string;
  sleep?: (ms: number) => Promise<void>;
};

/**
 * Options for the in-memory locks provider.
 */
export type MemoryLocksProviderOptions = CreateMemoryLocksOptions & {
  name?: string;
};

/**
 * Ports contributed by the memory locks provider.
 */
export interface MemoryLocksProviderPorts {
  locks: LocksPort;
}

/**
 * Error thrown when a lease option is invalid.
 */
export class LeaseOptionsError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "LeaseOptionsError";
  }
}

/**
 * Create an in-memory locks port for tests and single-process development.
 */
export function createMemoryLocks(
  options: CreateMemoryLocksOptions = {},
): MemoryLocksPort {
  const leases = new Map<string, MemoryLeaseRecord>();
  const now = options.now ?? (() => new Date());
  const sleep = options.sleep ?? defaultSleep;
  const createOwnerToken = options.createOwnerToken ?? createRandomToken;
  let nextFencingToken = 1;

  const port: MemoryLocksPort = {
    leases,
    async acquire(key, acquireOptions) {
      validateAcquireOptions(key, acquireOptions);
      const waitMs = acquireOptions.waitMs ?? 0;
      const retryDelayMs = acquireOptions.retryDelayMs ?? 50;
      const deadline = now().getTime() + waitMs;
      const ownerToken = acquireOptions.ownerToken ?? createOwnerToken();

      while (true) {
        pruneExpired(key, leases, now);
        if (!leases.has(key)) {
          const record: MemoryLeaseRecord = {
            key,
            ownerToken,
            expiresAt: new Date(now().getTime() + acquireOptions.ttlMs),
            ttlMs: acquireOptions.ttlMs,
            fencingToken: nextFencingToken++,
            metadata: acquireOptions.metadata,
          };
          leases.set(key, record);
          return {
            acquired: true,
            lease: createMemoryLease(port, record, now),
          };
        }

        if (waitMs <= 0) {
          return { acquired: false, reason: "unavailable" };
        }

        const remainingMs = deadline - now().getTime();
        if (remainingMs <= 0) {
          pruneExpired(key, leases, now);
          if (!leases.has(key)) continue;
          return { acquired: false, reason: "timeout" };
        }

        await sleep(Math.min(retryDelayMs, remainingMs));
      }
    },
    async withLease(key, acquireOptions, fn) {
      const result = await port.acquire(key, acquireOptions);
      if (!result.acquired) return undefined;

      try {
        return await fn({ lease: result.lease });
      } finally {
        await result.lease.release();
      }
    },
    restore(key, ownerToken, restoreOptions) {
      if (!key) throw new LeaseOptionsError("Lease key is required.");
      if (!ownerToken) {
        throw new LeaseOptionsError("Lease owner token is required.");
      }
      validateRestoreOptions(restoreOptions);

      return createMemoryLease(
        port,
        {
          key,
          ownerToken,
          ttlMs: restoreOptions.ttlMs,
          expiresAt: restoreOptions.expiresAt,
          fencingToken: restoreOptions.fencingToken,
        },
        now,
      );
    },
    async forceRelease(key) {
      pruneExpired(key, leases, now);
      return leases.delete(key);
    },
    reset() {
      leases.clear();
      nextFencingToken = 1;
    },
  };

  return port;
}

/**
 * Create a provider that contributes an in-memory locks port.
 */
export function createMemoryLocksProvider(
  options: MemoryLocksProviderOptions = {},
) {
  const { name = "memory-locks", ...locksOptions } = options;

  return createProvider({
    name,
    metadata: {
      ports: ["locks"],
      watchers: ["locks"],
    },
    setup({ ports }) {
      const instrumentation = createProviderInstrumentation(ports, {
        providerName: name,
        watcher: "locks",
      });
      const locks = createMemoryLocks(locksOptions);

      return {
        ports: {
          locks: instrumentLocks(locks, instrumentation),
        } satisfies MemoryLocksProviderPorts,
      };
    },
  });
}

/**
 * Convenience helper for any lock port.
 */
export function acquireLease(
  locks: LocksPort,
  key: string,
  options: LeaseAcquireOptions,
): Promise<LeaseAcquireResult> {
  return locks.acquire(key, options);
}

/**
 * Convenience helper for any lock port.
 */
export function withLease<T>(
  locks: LocksPort,
  key: string,
  options: LeaseAcquireOptions,
  fn: (ctx: { lease: LeaseHandle }) => MaybePromise<T>,
): Promise<T | undefined> {
  return locks.withLease(key, options, fn);
}

function createMemoryLease(
  port: MemoryLocksPort,
  record: {
    key: string;
    ownerToken: string;
    ttlMs: number;
    expiresAt?: Date;
    fencingToken?: string | number;
  },
  now: () => Date,
): LeaseHandle {
  const lease: LeaseHandle = {
    key: record.key,
    ownerToken: record.ownerToken,
    ...(record.expiresAt ? { expiresAt: record.expiresAt } : {}),
    ...(record.fencingToken !== undefined
      ? { fencingToken: record.fencingToken }
      : {}),
    async renew(options) {
      const active = port.leases.get(record.key);
      if (!active || active.ownerToken !== record.ownerToken) {
        return false;
      }
      if (active.expiresAt.getTime() <= now().getTime()) {
        port.leases.delete(record.key);
        return false;
      }

      const ttlMs = options?.ttlMs ?? record.ttlMs;
      if (!Number.isFinite(ttlMs) || ttlMs <= 0) {
        throw new LeaseOptionsError("Lease ttlMs must be a positive number.");
      }

      active.ttlMs = ttlMs;
      active.expiresAt = new Date(now().getTime() + ttlMs);
      record.ttlMs = ttlMs;
      lease.expiresAt = active.expiresAt;
      return true;
    },
    async release() {
      const active = port.leases.get(record.key);
      if (!active || active.ownerToken !== record.ownerToken) {
        return false;
      }
      if (active.expiresAt.getTime() <= now().getTime()) {
        port.leases.delete(record.key);
        return false;
      }

      return port.leases.delete(record.key);
    },
  };

  return lease;
}

function validateRestoreOptions(options: LeaseRestoreOptions): void {
  if (!options || !Number.isFinite(options.ttlMs) || options.ttlMs <= 0) {
    throw new LeaseOptionsError(
      "Lease restore ttlMs must be a positive number.",
    );
  }
  if (
    options.expiresAt !== undefined &&
    (!(options.expiresAt instanceof Date) ||
      !Number.isFinite(options.expiresAt.getTime()))
  ) {
    throw new LeaseOptionsError(
      "Lease restore expiresAt must be a valid Date.",
    );
  }
}

function instrumentLocks(
  locks: LocksPort,
  instrumentation: ReturnType<typeof createProviderInstrumentation>,
): LocksPort {
  const acquire: LocksPort["acquire"] = async (key, options) => {
    const startedAt = Date.now();
    const result = await locks.acquire(key, options);
    instrumentation.custom({
      name: "locks.acquire",
      label: "Lock acquire",
      summary: result.acquired ? "Lease acquired" : "Lease not acquired",
      details: {
        key,
        acquired: result.acquired,
        reason: result.acquired ? undefined : result.reason,
        ttlMs: options.ttlMs,
        waitMs: options.waitMs ?? 0,
        durationMs: Date.now() - startedAt,
      },
    });
    return result;
  };

  return {
    acquire,
    async withLease(key, options, fn) {
      const result = await acquire(key, options);
      if (!result.acquired) return undefined;

      try {
        return await fn({ lease: result.lease });
      } finally {
        await result.lease.release();
      }
    },
    restore(key, ownerToken, options) {
      return locks.restore(key, ownerToken, options);
    },
    async forceRelease(key) {
      const released = await locks.forceRelease(key);
      instrumentation.custom({
        name: "locks.forceRelease",
        label: "Lock force release",
        summary: released ? "Lease force released" : "Lease not found",
        details: { key, released },
      });
      return released;
    },
  };
}

function validateAcquireOptions(key: string, options: LeaseAcquireOptions) {
  if (!key) throw new LeaseOptionsError("Lease key is required.");
  if (!Number.isFinite(options.ttlMs) || options.ttlMs <= 0) {
    throw new LeaseOptionsError("Lease ttlMs must be a positive number.");
  }
  if (options.waitMs !== undefined && options.waitMs < 0) {
    throw new LeaseOptionsError("Lease waitMs must be zero or greater.");
  }
  if (options.retryDelayMs !== undefined && options.retryDelayMs <= 0) {
    throw new LeaseOptionsError("Lease retryDelayMs must be positive.");
  }
}

function pruneExpired(
  key: string,
  leases: Map<string, MemoryLeaseRecord>,
  now: () => Date,
) {
  const active = leases.get(key);
  if (active && active.expiresAt.getTime() <= now().getTime()) {
    leases.delete(key);
  }
}

function createRandomToken(): string {
  if (typeof crypto !== "undefined" && "randomUUID" in crypto) {
    return crypto.randomUUID();
  }

  return `${Date.now().toString(36)}-${Math.random().toString(36).slice(2)}`;
}

function defaultSleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}