@platless/axom-cli 0.3.13 → 0.5.9

package/dist/workstation/keepers/kit/server.js

@@ -0,0 +1,518 @@
+// The keeper-side half of the kit: listen on the 0600 socket, authenticate with the connection token, answer the generic
+// `keeper.*` methods, dispatch to registered method handlers and push stream records under credit.
+// Keeper-side module: node builtins and kit siblings only (the import guard enforces it).
+import { randomBytes, randomUUID, timingSafeEqual } from 'node:crypto';
+import { promises as fs, renameSync, unlinkSync } from 'node:fs';
+import net from 'node:net';
+import { sampleActivity } from './activity.js';
+import { encodeFrame, LineDecoder, parseServerBound } from './ipc.js';
+import { ErrorCode, frameLimitsOf, isProtocolWindow, KeeperMethodError, MAX_CONNECTIONS, MAX_EVT_RECORDS, PRE_HELLO_MAX_BYTES, } from './protocol.js';
+const GENERIC = new Set(['keeper.hello', 'keeper.ping', 'keeper.info', 'keeper.activity', 'keeper.drain', 'keeper.stop', 'keeper.subscribe', 'keeper.unsubscribe']);
+const MAX_CREDIT = 1_000_000_000;
+const MAX_CREDIT_BYTES = Number.MAX_SAFE_INTEGER;
+const wireError = (code, message) => ({ code, message });
+class Subscription {
+    connection;
+    stream;
+    source;
+    credit;
+    /** Bytes the subscriber may still be sent; `Infinity` on a subscription without a byte window. */
+    byteCredit;
+    byteWindowed;
+    next;
+    pumping = false;
+    dirty = false;
+    closed = false;
+    started = false;
+    unwatch;
+    constructor(connection, stream, source, from, credit, creditBytes) {
+        this.connection = connection;
+        this.stream = stream;
+        this.source = source;
+        this.next = from;
+        this.credit = credit;
+        this.byteWindowed = creditBytes !== undefined;
+        this.byteCredit = creditBytes ?? Number.POSITIVE_INFINITY;
+        this.unwatch = source.watch(() => this.kick());
+    }
+    /** Begins pushing. Called after the subscribe reply is written so no `evt` can overtake its `res`. */
+    start() { this.started = true; this.kick(); }
+    grant(n, bytes) {
+        this.credit = Math.min(MAX_CREDIT, this.credit + n);
+        if (this.byteWindowed && bytes !== undefined)
+            this.byteCredit = Math.min(MAX_CREDIT_BYTES, this.byteCredit + bytes);
+        this.kick();
+    }
+    close() { if (this.closed)
+        return; this.closed = true; this.unwatch(); }
+    kick() {
+        this.dirty = true;
+        if (this.started && !this.pumping && !this.closed)
+            void this.pump();
+    }
+    async pump() {
+        this.pumping = true;
+        try {
+            while (this.dirty && !this.closed) {
+                this.dirty = false;
+                while (!this.closed && this.credit > 0 && this.byteCredit > 0) {
+                    if (!await this.connection.writable())
+                        return;
+                    const limit = Math.min(this.credit, MAX_EVT_RECORDS);
+                    const records = await (this.byteWindowed ? this.source.read(this.next, limit, { maxBytes: this.byteCredit }) : this.source.read(this.next, limit));
+                    if (this.closed || records.length === 0)
+                        break;
+                    const { batch, bytes } = this.fit(records);
+                    if (batch.length === 0) {
+                        this.end('record_too_large');
+                        return;
+                    }
+                    this.next = batch[batch.length - 1].o + 1;
+                    this.credit -= batch.length;
+                    if (this.byteWindowed)
+                        this.byteCredit -= bytes;
+                    this.connection.send({ t: 'evt', stream: this.stream, records: batch, ...(this.byteWindowed ? { bytes } : {}) });
+                }
+            }
+        }
+        catch {
+            this.end('source_error');
+        }
+        finally {
+            this.pumping = false;
+        }
+    }
+    /** Longest prefix of `records` whose evt frame stays under the outbound cap, and its size (each record's JSON line). */
+    fit(records) {
+        const budget = this.connection.limits.out - 256 - this.stream.length * 2;
+        let used = 0;
+        let count = 0;
+        for (const record of records) {
+            const size = Buffer.byteLength(JSON.stringify(record)) + 1;
+            if (used + size > budget)
+                break;
+            used += size;
+            count += 1;
+        }
+        return { batch: records.slice(0, count), bytes: used };
+    }
+    end(reason) {
+        this.close();
+        this.connection.removeSubscription(this.stream, this);
+        this.connection.send({ t: 'evt', stream: this.stream, records: [], closed: reason });
+    }
+}
+class Connection {
+    server;
+    socket;
+    id = randomUUID();
+    limits;
+    authenticated = false;
+    decoder;
+    subscriptions = new Map();
+    dead = false;
+    authTimer;
+    drainWaiters = [];
+    constructor(server, socket, unauthenticatedTimeoutMs) {
+        this.server = server;
+        this.socket = socket;
+        this.limits = frameLimitsOf(server.options.kind);
+        this.decoder = new LineDecoder(Math.min(PRE_HELLO_MAX_BYTES, this.limits.in)); // widened to the kind's cap once hello succeeds
+        socket.on('data', (chunk) => this.onData(chunk));
+        socket.on('error', () => this.destroy());
+        socket.on('close', () => this.destroy());
+        socket.on('drain', () => this.wake());
+        this.authTimer = setTimeout(() => { if (!this.authenticated)
+            this.destroy(); }, unauthenticatedTimeoutMs);
+        this.authTimer.unref();
+    }
+    get isDead() { return this.dead; }
+    pendingAnswers = 0;
+    /** No `res` frame of this connection is still waiting to leave the socket. Pushed `evt` data does not count: a stalled subscriber must not hold up a stop. */
+    get flushed() { return this.dead || this.pendingAnswers === 0; }
+    destroy() {
+        if (this.dead)
+            return;
+        this.dead = true;
+        clearTimeout(this.authTimer);
+        for (const sub of this.subscriptions.values())
+            sub.close();
+        this.subscriptions.clear();
+        this.socket.destroy();
+        this.wake();
+        this.server.forget(this);
+    }
+    wake() { const waiters = this.drainWaiters; this.drainWaiters = []; for (const w of waiters)
+        w(); }
+    /** Resolves true when the socket can take more data, false when the connection died meanwhile. */
+    async writable() {
+        while (!this.dead && this.socket.writableNeedDrain)
+            await new Promise((resolve) => this.drainWaiters.push(resolve));
+        return !this.dead;
+    }
+    /** Writes one frame. An oversized `res` is replaced by an error answer so the request still completes. */
+    send(frame, then) {
+        if (this.dead)
+            return;
+        let encoded;
+        try {
+            encoded = encodeFrame(frame);
+        }
+        catch {
+            encoded = frame.t === 'res' ? encodeFrame({ t: 'res', id: frame.id, ok: false, error: wireError(ErrorCode.internal, 'Result is not serialisable') }) : encodeFrame({ t: 'evt', stream: frame.stream, records: [], closed: 'source_error' });
+        }
+        if (encoded.length > this.limits.out) {
+            if (frame.t !== 'res')
+                return; // the pump never builds such an evt
+            encoded = encodeFrame({ t: 'res', id: frame.id, ok: false, error: wireError(ErrorCode.frameTooLarge, `Response exceeds ${this.limits.out} bytes`) });
+        }
+        if (frame.t === 'res') {
+            this.pendingAnswers += 1;
+            this.socket.write(encoded.data, () => { this.pendingAnswers -= 1; then?.(); });
+        }
+        else
+            this.socket.write(encoded.data);
+    }
+    removeSubscription(stream, sub) { if (this.subscriptions.get(stream) === sub)
+        this.subscriptions.delete(stream); }
+    starts = [];
+    subscribe(stream, source, from, credit, creditBytes) {
+        this.subscriptions.get(stream)?.close();
+        const sub = new Subscription(this, stream, source, from, credit, creditBytes);
+        this.subscriptions.set(stream, sub);
+        this.starts.push(sub);
+    }
+    discardStarts() { const starts = this.starts; this.starts = []; for (const sub of starts)
+        this.unsubscribeIf(sub); }
+    unsubscribeIf(sub) { if (this.subscriptions.get(sub.stream) === sub) {
+        sub.close();
+        this.subscriptions.delete(sub.stream);
+    } }
+    startSubscriptions() { const starts = this.starts; this.starts = []; for (const sub of starts)
+        if (this.subscriptions.get(sub.stream) === sub)
+            sub.start(); }
+    unsubscribe(stream) {
+        const sub = this.subscriptions.get(stream);
+        if (!sub)
+            return false;
+        sub.close();
+        this.subscriptions.delete(stream);
+        return true;
+    }
+    onData(chunk) {
+        let lines;
+        try {
+            lines = this.decoder.push(chunk);
+        }
+        catch {
+            this.protocolViolation();
+            return;
+        }
+        for (const line of lines) {
+            if (this.dead)
+                return;
+            let frame;
+            try {
+                frame = parseServerBound(line);
+            }
+            catch {
+                this.protocolViolation();
+                return;
+            }
+            this.dispatch(frame);
+        }
+    }
+    /** A malformed or oversized frame ends this connection (and only this one). */
+    protocolViolation() { this.destroy(); }
+    dispatch(frame) {
+        if (frame.t === 'credit') {
+            if (!this.authenticated) {
+                this.destroy();
+                return;
+            }
+            this.subscriptions.get(frame.stream)?.grant(frame.n, frame.bytes);
+            return;
+        }
+        const reply = (ok, payload, then) => this.send(ok ? { t: 'res', id: frame.id, ok: true, result: payload } : { t: 'res', id: frame.id, ok: false, error: payload }, then);
+        const fail = (code, message, then) => reply(false, wireError(code, message), then);
+        const { method } = frame;
+        const params = frame.params ?? {};
+        if (method === 'keeper.hello') { // synchronous on purpose: a pipelined request right behind it must see the authenticated state
+            const token = typeof params.token === 'string' ? Buffer.from(params.token) : Buffer.alloc(0);
+            const expected = Buffer.from(this.server.token);
+            if (token.length !== expected.length || !timingSafeEqual(token, expected)) {
+                fail(ErrorCode.badToken, 'Invalid keeper token', () => this.destroy());
+                return;
+            }
+            if (!isProtocolWindow(params.daemonProtocol)) {
+                fail(ErrorCode.badRequest, 'daemonProtocol must be [min, max]');
+                return;
+            }
+            this.authenticated = true;
+            clearTimeout(this.authTimer);
+            this.decoder.setMax(this.limits.in);
+            reply(true, this.server.hello());
+            return;
+        }
+        if (!this.authenticated && method !== 'keeper.ping') {
+            fail(ErrorCode.unauthenticated, 'keeper.hello required');
+            return;
+        }
+        const context = {
+            connection: { id: this.id },
+            subscribe: (stream, source, from, credit, creditBytes) => this.subscribe(stream, source, from, credit, creditBytes),
+            unsubscribe: (stream) => { this.unsubscribe(stream); },
+        };
+        void (async () => {
+            try {
+                const result = GENERIC.has(method) ? await this.server.generic(method, params, context, () => reply(true, {}, () => this.server.stopped(params.mode)))
+                    : await this.server.callHandler(method, params, context);
+                if (result !== NOT_REPLIED)
+                    reply(true, result);
+                this.startSubscriptions();
+            }
+            catch (error) {
+                this.discardStarts();
+                if (error instanceof KeeperMethodError)
+                    fail(error.code, error.message);
+                else
+                    fail(ErrorCode.internal, error instanceof Error ? error.message : String(error));
+            }
+        })();
+    }
+}
+const NOT_REPLIED = Symbol('replied');
+export class KeeperServer {
+    options;
+    startedAt = Date.now();
+    handlers = new Map();
+    streamResolver;
+    connections = new Set();
+    net = net.createServer((socket) => this.accept(socket));
+    tokenValue = '';
+    drainingFlag = false;
+    startsWork = new Set();
+    samplers = [];
+    inFlightCount = 0;
+    stopState;
+    constructor(options) {
+        this.options = options;
+        if (options.activity)
+            this.samplers.push(options.activity);
+        this.net.on('error', () => { });
+    }
+    get token() { return this.tokenValue; }
+    get connectionCount() { return this.connections.size; }
+    /** Connections that passed `keeper.hello` (a daemon, not a stray probe). */
+    get authenticatedConnectionCount() { let count = 0; for (const connection of this.connections)
+        if (connection.authenticated)
+            count += 1; return count; }
+    get draining() { return this.drainingFlag; }
+    /** Domain requests currently running (the `keeper.*` methods are not counted). */
+    get inFlight() { return this.inFlightCount; }
+    /** Registers a domain method. The `keeper.*` namespace belongs to the kit. */
+    register(method, handler, options = {}) {
+        if (method.startsWith('keeper.'))
+            throw new Error(`The keeper.* namespace is reserved: ${method}`);
+        if (this.handlers.has(method))
+            throw new Error(`Method already registered: ${method}`);
+        this.handlers.set(method, handler);
+        if (options.startsWork)
+            this.startsWork.add(method);
+    }
+    /** Adds a source for `keeper.activity`. Each sample asks every sampler afresh; nothing is accumulated. */
+    addActivitySampler(sampler) { this.samplers.push(sampler); }
+    report(error) { try {
+        this.options.onError?.(error);
+    }
+    catch { /* diagnostics must not matter */ } }
+    /** Stops accepting new work (methods registered with `startsWork`). Everything else keeps working. Idempotent. */
+    drain() {
+        if (this.drainingFlag)
+            return;
+        this.drainingFlag = true;
+        try {
+            this.options.onDrain?.();
+        }
+        catch (error) {
+            this.report(error);
+        }
+    }
+    /** How `keeper.subscribe{stream}` finds its source. */
+    setStreamResolver(resolver) { this.streamResolver = resolver; }
+    hello() {
+        const o = this.options;
+        return { kind: o.kind, protocolVersion: o.protocolVersion, buildVersion: o.buildVersion, gen: o.gen, pid: process.pid, startedAt: this.startedAt, capabilities: [...(o.capabilities ?? [])] };
+    }
+    /**
+     * Generates a fresh 32-byte token and listens on the 0600 socket. The token reaches `run/<gen>/token` (0600, atomic rename)
+     * only once this process actually owns the socket, before any connection is served: a second start for the same generation
+     * that fails with EADDRINUSE must leave the live keeper's token file alone.
+     */
+    async listen() {
+        const token = randomBytes(32).toString('hex');
+        const temp = `${this.options.tokenFile}.${randomUUID()}.tmp`;
+        await fs.writeFile(temp, token, { mode: 0o600 });
+        try {
+            await new Promise((resolve, reject) => {
+                const failed = (error) => reject(error);
+                this.net.once('error', failed);
+                this.net.listen(this.options.socket, () => {
+                    this.net.off('error', failed);
+                    try {
+                        renameSync(temp, this.options.tokenFile);
+                        this.tokenValue = token;
+                        resolve();
+                    } // synchronous: no connection is accepted before this
+                    catch (error) {
+                        this.net.close();
+                        reject(error);
+                    }
+                });
+            });
+        }
+        catch (error) {
+            try {
+                unlinkSync(temp);
+            }
+            catch { /* already gone */ }
+            throw error;
+        }
+        await fs.chmod(this.options.socket, 0o600);
+    }
+    async close() {
+        for (const connection of [...this.connections])
+            connection.destroy();
+        if (!this.net.listening)
+            return;
+        await new Promise((resolve) => this.net.close(() => resolve()));
+    }
+    forget(connection) { this.connections.delete(connection); }
+    accept(socket) {
+        if (this.connections.size >= (this.options.maxConnections ?? MAX_CONNECTIONS)) {
+            socket.destroy();
+            return;
+        }
+        this.connections.add(new Connection(this, socket, this.options.unauthenticatedTimeoutMs ?? 10_000));
+    }
+    async callHandler(method, params, context) {
+        const handler = this.handlers.get(method);
+        if (!handler)
+            throw new KeeperMethodError(ErrorCode.unknownMethod, `Unknown method: ${method}`);
+        if (this.drainingFlag && this.startsWork.has(method))
+            throw new KeeperMethodError(ErrorCode.draining, `Keeper is draining; ${method} starts new work`);
+        this.inFlightCount += 1;
+        try {
+            return await handler(params, context);
+        }
+        finally {
+            this.inFlightCount -= 1;
+        }
+    }
+    /**
+     * `keeper.stop` was accepted (before its reply is written): from now on `info.stopping` is set and new work is refused (both modes:
+     * a force stop's onStop may still be tearing things down, and no `startsWork` call may slip in meanwhile).
+     * A force request upgrades a pending graceful one.
+     */
+    beginStop(mode) {
+        if (!this.stopState)
+            this.stopState = { mode };
+        else if (mode === 'force')
+            this.stopState.mode = 'force';
+        this.drain();
+    }
+    /**
+     * The stop reply has been written. Graceful: wait until no request is in flight and every answer has left its socket, bounded by
+     * `stopGraceMs`, then stop (also when the grace ran out). Force: stop now, even over a pending graceful stop. The stop happens once.
+     */
+    stopped(mode) {
+        this.beginStop(mode);
+        const state = this.stopState;
+        if (state.finish) {
+            if (state.mode === 'force')
+                state.finish();
+            return;
+        }
+        let done = false;
+        let timer;
+        let poll;
+        const finish = () => {
+            if (done)
+                return;
+            done = true;
+            clearTimeout(timer);
+            clearTimeout(poll);
+            try {
+                if (this.options.onStop)
+                    this.options.onStop(state.mode);
+                else {
+                    for (const connection of [...this.connections])
+                        connection.destroy();
+                    if (this.net.listening)
+                        this.net.close();
+                    this.options.exit?.(0);
+                }
+            }
+            catch (error) {
+                this.report(error);
+                this.stopState = undefined;
+            } // the keeper stays up and a later stop can try again
+        };
+        state.finish = finish;
+        if (state.mode === 'force') {
+            finish();
+            return;
+        }
+        timer = setTimeout(finish, this.options.stopGraceMs ?? 3000);
+        timer.unref();
+        const check = () => {
+            if (done)
+                return;
+            if (this.inFlightCount === 0 && [...this.connections].every((c) => c.flushed)) {
+                finish();
+                return;
+            }
+            poll = setTimeout(check, 10);
+        };
+        check();
+    }
+    async generic(method, params, context, replyStop) {
+        const o = this.options;
+        switch (method) {
+            case 'keeper.ping': return { pid: process.pid, identity: o.identity, gen: o.gen, kind: o.kind, buildVersion: o.buildVersion, protocolVersion: o.protocolVersion };
+            case 'keeper.info': return { ...this.hello(), identity: o.identity, draining: this.drainingFlag, stopping: this.stopState?.mode ?? null, inFlight: this.inFlightCount, connections: this.connections.size, uptimeMs: Date.now() - this.startedAt };
+            case 'keeper.activity': return sampleActivity(this.samplers, { timeoutMs: o.activityTimeoutMs, onError: (error) => this.report(error) });
+            case 'keeper.drain':
+                this.drain();
+                return { draining: true };
+            case 'keeper.stop': {
+                if (params.mode !== 'graceful' && params.mode !== 'force')
+                    throw new KeeperMethodError(ErrorCode.badRequest, 'mode must be graceful or force');
+                this.beginStop(params.mode);
+                replyStop();
+                return NOT_REPLIED;
+            }
+            case 'keeper.subscribe': {
+                const { stream, creditBytes } = params;
+                const from = params.from ?? 0;
+                const credit = params.credit ?? 0;
+                if (typeof stream !== 'string' || !stream || !Number.isSafeInteger(from) || from < 0 || !Number.isSafeInteger(credit) || credit < 0)
+                    throw new KeeperMethodError(ErrorCode.badRequest, 'subscribe needs stream, from >= 0 and credit >= 0');
+                if (creditBytes !== undefined && (!Number.isSafeInteger(creditBytes) || creditBytes <= 0))
+                    throw new KeeperMethodError(ErrorCode.badRequest, 'creditBytes must be a positive integer');
+                const source = await this.streamResolver?.(stream);
+                if (!source)
+                    throw new KeeperMethodError(ErrorCode.unknownStream, `Unknown stream: ${stream}`);
+                context.subscribe(stream, source, from, Math.min(credit, MAX_CREDIT), creditBytes);
+                return { stream, from };
+            }
+            case 'keeper.unsubscribe': {
+                if (typeof params.stream !== 'string')
+                    throw new KeeperMethodError(ErrorCode.badRequest, 'unsubscribe needs stream');
+                context.unsubscribe(params.stream);
+                return {};
+            }
+            default: throw new KeeperMethodError(ErrorCode.unknownMethod, `Unknown method: ${method}`);
+        }
+    }
+}