@jbagatta/johnny-locke/dist/jetstream/jetstream-distributed-lock.js

A robust, strongly-consistent distributed locking library that provides atomic operations across multiple processes

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.JetstreamDistributedLock = void 0;
const nats_1 = require("nats");
const util_1 = require("../util");
class JetstreamDistributedLock {
    constructor(kv, config) {
        this.kv = kv;
        this.config = config;
        this.state = new Map();
        this.active = false;
        this.initialized = false;
    }
    static async connect(natsClient, config) {
        (0, util_1.validateLockConfiguration)(config);
        const kv = await natsClient.jetstream().views.kv(config.namespace, { bindOnly: true });
        console.log(`JetstreamDistributedLock connected to namespace ${config.namespace}: ${JSON.stringify(await kv.status())}`);
        const distributor = new JetstreamDistributedLock(kv, config);
        await new Promise(distributor.initialize.bind(distributor));
        return distributor;
    }
    static async create(natsClient, config) {
        (0, util_1.validateLockConfiguration)(config);
        const kv = await natsClient.jetstream().views.kv(config.namespace, {
            history: 1,
            ttl: config.objectExpiryMs ? (0, nats_1.nanos)(config.objectExpiryMs) : undefined,
            storage: nats_1.StorageType.Memory,
            replicas: 1
        });
        console.log(`JetstreamDistributedLock created namespace ${config.namespace}: ${JSON.stringify(await kv.status())}`);
        const distributor = new JetstreamDistributedLock(kv, config);
        await new Promise(distributor.initialize.bind(distributor));
        return distributor;
    }
    async initialize(resolve, reject) {
        try {
            if (this.initialized) {
                throw new Error('JetstreamDistributedLock already initialized');
            }
            const watch = await this.kv.watch({
                key: `${this.config.namespace}.>`,
                initializedFn: () => {
                    this.active = true;
                    this.initialized = true;
                    resolve();
                }
            });
            this.watch = watch;
            (async () => {
                for await (const entry of watch) {
                    if (this.initialized && !this.active) {
                        throw new Error('JetstreamDistributedLock closed');
                    }
                    this.processEntry(entry);
                }
            }).bind(this)().catch(reject);
        }
        catch (error) {
            reject(error);
        }
    }
    processEntry(entry) {
        try {
            const lockState = entry.operation === 'PUT'
                ? JSON.parse(entry.string())
                : { status: 'expired', ttlMs: 0, createdTime: entry.created.getTime() };
            this.updateLocalState(entry.key, lockState, entry.revision, entry.created.getTime());
        }
        catch (error) {
            console.error(`Could not process entry ${entry.key}, revision: ${entry.revision}, error: ${error}`);
        }
    }
    close() {
        this.watch?.stop();
        this.active = false;
        this.state.clear();
    }
    async withLock(key, timeoutMs, callback, lockDuration) {
        const lock = await this.acquireLock(key, timeoutMs, lockDuration);
        try {
            const updatedState = await callback(lock.value);
            const result = lock.update(updatedState);
            const updated = await this.releaseLock(key, result);
            if (!updated) {
                throw new util_1.TimeoutError(key);
            }
            return result;
        }
        catch (error) {
            await this.releaseLock(key, lock);
            throw error;
        }
    }
    async acquireLock(key, timeoutMs, lockDuration) {
        this.checkActive();
        const namespacedKey = this.toNamespacedKey(key);
        const duration = (0, util_1.computeLockDuration)(this.config.defaultLockDurationMs, lockDuration);
        let now = Date.now();
        const deadline = now + timeoutMs;
        while (now < deadline) {
            try {
                const unlockTimeout = deadline - now;
                const lockState = await new Promise((resolve, reject) => this.resolveOnUnlock.bind(this)(namespacedKey, unlockTimeout, resolve, reject));
                const lock = await this.createOrUpdateLock(namespacedKey, duration, lockState);
                const lockObj = lock.value ? JSON.parse(lock.value) : null;
                return new util_1.WritableObject(lockObj, lock.lockId);
            }
            catch (err) {
                // suppress timeouts and lock acquire failures, retry 
                now = Date.now();
            }
        }
        throw new util_1.TimeoutError(namespacedKey);
    }
    async tryAcquireLock(key, lockDuration) {
        this.checkActive();
        const namespacedKey = this.toNamespacedKey(key);
        try {
            await this.updateKey(namespacedKey);
            const lockState = this.state.get(namespacedKey);
            if (lockState && this.isLockActive(lockState.state)) {
                return { acquired: false, value: undefined };
            }
            const duration = (0, util_1.computeLockDuration)(this.config.defaultLockDurationMs, lockDuration);
            const lock = await this.createOrUpdateLock(namespacedKey, duration, lockState?.state);
            const lockObj = lock.value ? JSON.parse(lock.value) : null;
            return { acquired: true, value: new util_1.WritableObject(lockObj, lock.lockId) };
        }
        catch {
            return { acquired: false, value: undefined };
        }
    }
    async releaseLock(key, lockObj) {
        this.checkActive();
        const namespacedKey = this.toNamespacedKey(key);
        const lockState = this.state.get(namespacedKey);
        if (lockState && this.isLockActive(lockState.state) && lockState.state.lockId === lockObj.lockId) {
            const newLock = {
                status: 'unlocked',
                lockId: lockObj.lockId,
                value: lockObj.value ? JSON.stringify(lockObj.value) : undefined,
                createdTime: lockState.state.createdTime,
                ttlMs: lockState.state.ttlMs
            };
            try {
                const revision = await this.kv.update(namespacedKey, JSON.stringify(newLock), lockState.state.revision);
                this.updateLocalState(namespacedKey, newLock, revision, newLock.createdTime);
                return true;
            }
            catch (error) {
                console.error(`Could not release lock ${namespacedKey}, error: ${error}`);
                return false;
            }
        }
        return false;
    }
    async wait(key, timeoutMs) {
        this.checkActive();
        const namespacedKey = this.toNamespacedKey(key);
        const lockState = await new Promise((resolve, reject) => this.resolveOnUnlock.bind(this)(namespacedKey, timeoutMs, resolve, reject));
        const value = lockState?.value ? JSON.parse(lockState.value) : null;
        return { value };
    }
    async delete(key) {
        this.checkActive();
        const lock = await this.acquireLock(key, this.config.defaultLockDurationMs);
        return await this.releaseLock(key, lock.update(null));
    }
    async createOrUpdateLock(namespacedKey, duration, lockState) {
        const newLock = {
            status: 'locked',
            lockId: crypto.randomUUID(),
            value: lockState?.value,
            ttlMs: duration,
            createdTime: Date.now(),
        };
        const revision = lockState?.revision
            ? await this.kv.update(namespacedKey, JSON.stringify(newLock), lockState.revision)
            : await this.kv.create(namespacedKey, JSON.stringify(newLock));
        return this.updateLocalState(namespacedKey, newLock, revision, newLock.createdTime);
    }
    async resolveOnUnlock(namespacedKey, timeoutMs, resolve, reject) {
        await this.updateKey(namespacedKey);
        const initialState = this.state.get(namespacedKey);
        if (!initialState || !this.isLockActive(initialState.state)) {
            return resolve(initialState?.state);
        }
        const watchId = crypto.randomUUID();
        const timeout = setTimeout(() => {
            const state = this.state.get(namespacedKey);
            state?.watchers.delete(watchId);
            !state || !this.isLockActive(state.state)
                ? resolve(state?.state)
                : reject(new util_1.TimeoutError(namespacedKey));
        }, timeoutMs);
        const callback = async (entry) => {
            if (!this.isLockActive(entry)) {
                const state = this.state.get(namespacedKey);
                state?.watchers.delete(watchId);
                clearTimeout(timeout);
                resolve(entry);
            }
        };
        initialState.watchers.set(watchId, callback);
    }
    isLockActive(lockState) {
        const drift = lockState.timestamp - lockState.createdTime;
        const expiryWithDrift = Date.now() + drift;
        return lockState.status === 'locked' && (lockState.createdTime + lockState.ttlMs >= expiryWithDrift);
    }
    updateLocalState(namespacedKey, lockMessage, newRevision, time) {
        const oldValue = this.state.get(namespacedKey);
        if (oldValue?.state?.revision && oldValue.state.revision >= newRevision) {
            return oldValue.state;
        }
        const newState = {
            state: {
                ...lockMessage,
                revision: newRevision,
                timestamp: time
            },
            watchers: oldValue?.watchers ?? new Map()
        };
        this.state.set(namespacedKey, newState);
        newState.watchers.forEach(watcher => watcher(newState.state).catch(console.error));
        return newState.state;
    }
    async updateKey(namespacedKey) {
        const latest = await this.kv.get(namespacedKey);
        if (latest !== null) {
            this.processEntry(latest);
        }
    }
    checkActive() {
        if (!this.active) {
            throw new Error('JetstreamDistributedLock closed');
        }
    }
    toNamespacedKey(key) {
        return `${this.config.namespace}.${key}`;
    }
}
exports.JetstreamDistributedLock = JetstreamDistributedLock;
//# sourceMappingURL=jetstream-distributed-lock.js.map