@shotkit/shotium 0.0.1 → 0.1.0

package/src/lib/pool.ts

@@ -0,0 +1,243 @@
+import {EventEmitter} from 'node:events';
+import fs from 'node:fs';
+import path from 'node:path';
+
+import type {ResolvedStartOptions} from './config.js';
+import {defaultCacheDir} from './config.js';
+import type {WireRequest} from './request.js';
+import type {WorkerResult} from './worker.js';
+import {Worker} from './worker.js';
+
+// A worker that exits sooner than this never really started, so its slot is
+// refilled on a doubling delay rather than immediately.
+const FAST_FAILURE_MS = 1000;
+const RESPAWN_DELAY_MS = 100;
+const MAX_RESPAWN_DELAY_MS = 5000;
+
+export interface SubmitOptions {
+  /** The supervisor's deadline, in milliseconds. */
+  timeout: number;
+  /** How many times to re-send after a crash or a timeout. */
+  retry: number;
+}
+
+interface Job {
+  request: WireRequest;
+  timeout: number;
+  attemptsLeft: number;
+  resolve: (result: WorkerResult) => void;
+  reject: (error: Error) => void;
+}
+
+// A fixed set of shotium.exe --serve processes, and a queue in front of them.
+//
+// The pool exists because blink is a process-wide singleton: one worker renders
+// one document at a time, so N concurrent screenshots means N processes. It is
+// also what makes a crash survivable -- a worker that dies takes its own
+// request down and nothing else, and the slot is refilled.
+//
+// Events:
+//   ready          {workers}                 the pool is up
+//   exit           {worker, code, signal}    a worker is gone
+//   crash          {worker, code, signal}    a worker died owing an answer
+//   timeout        {worker, timeout}         a request outlived its deadline
+//   worker-restart {worker, reason, delay}   a slot was refilled
+//   worker-error   {worker, error}           a worker could not be started
+//   stderr         {worker, line}            a diagnostic line from a worker
+class Pool extends EventEmitter {
+  private readonly binary: string;
+  private readonly size: number;
+  private readonly args: string[];
+  private readonly cacheDir: string|null;
+  private slots: Worker[] = [];
+  private failures: number[] = [];
+  private queue: Job[] = [];
+  private stopping = false;
+  private nextId = 0;
+
+  constructor(options: ResolvedStartOptions) {
+    super();
+    this.binary = options.binary;
+    this.size = options.workers;
+    this.args = options.args || [];
+    this.cacheDir = options.cacheDir || null;
+  }
+
+  start(): void {
+    if (this.slots.length > 0) {
+      return;
+    }
+    for (let slot = 0; slot < this.size; ++slot) {
+      this.slots[slot] = this.spawn(slot);
+    }
+    this.emit('ready', {workers: this.size});
+  }
+
+  private spawn(slot: number): Worker {
+    const id = this.nextId++;
+    const args = [...this.args];
+    if (this.cacheDir) {
+      // One directory per slot, not per process: the Simple backend takes an
+      // exclusive lock on its directory, so sharing one would leave every
+      // worker but the first running uncached. Keying on the slot rather than
+      // the worker id means a restarted worker inherits the warm cache its
+      // predecessor built.
+      const dir = path.join(this.cacheDir, `worker-${slot}`);
+      fs.mkdirSync(dir, {recursive: true});
+      args.push(`--cache-dir=${dir}`);
+    }
+
+    const startedAt = Date.now();
+    const worker = new Worker({id, binary: this.binary, args});
+    worker.on('stderr', (event) => this.emit('stderr', event));
+    worker.on('crash', (event) => this.emit('crash', event));
+    // A worker that could not be started at all -- a binary that is not there,
+    // a path that is not executable -- reports it here. Without a listener
+    // EventEmitter throws the error instead, which for a resident daemon means
+    // a typo in a path takes the whole pool down.
+    worker.on('error', (error) => this.emit('worker-error', {worker: id, error}));
+    worker.on('exit', (event) => {
+      this.emit('exit', event);
+      if (this.stopping || this.slots[slot] !== worker) {
+        return;
+      }
+      // A worker that died on the way up is not a crash to recover from, it is
+      // a configuration that does not work, and refilling the slot as fast as
+      // the loop allows would spin a core until someone noticed. Back off, but
+      // never give up: the binary may yet appear, and a pool that stopped
+      // trying would have to be restarted by hand.
+      //
+      // "On the way up" is answered nothing and did not last a second, in that
+      // order. Age alone would misread the ordinary case this design exists
+      // for -- a worker killed mid-request seconds after the pool started --
+      // as a startup failure, and delay the slot that the retry needs.
+      const started = worker.served > 0 ||
+          (Date.now() - startedAt) >= FAST_FAILURE_MS;
+      if (started) {
+        this.failures[slot] = 0;
+      }
+      const failures = this.failures[slot] || 0;
+      const delay = started ?
+          0 :
+          Math.min(MAX_RESPAWN_DELAY_MS, RESPAWN_DELAY_MS * 2 ** failures);
+      this.failures[slot] = failures + 1;
+      const refill = () => {
+        if (this.stopping || this.slots[slot] !== worker) {
+          return;
+        }
+        const replacement = this.spawn(slot);
+        this.slots[slot] = replacement;
+        this.emit('worker-restart',
+                  {worker: replacement.id, reason: 'exit', delay});
+        this.pump();
+      };
+      if (delay === 0) {
+        refill();
+        return;
+      }
+      const timer = setTimeout(refill, delay);
+      // An unref'd timer does not hold the process open: a pool whose workers
+      // all failed should not be the reason a program refuses to exit.
+      timer.unref();
+    });
+    return worker;
+  }
+
+  // Queues one request. `timeout` is the supervisor's deadline, which is longer
+  // than the worker's own: the worker fails a slow page by itself and answers,
+  // and this only fires when it has stopped answering at all.
+  submit(request: WireRequest, {timeout, retry}: SubmitOptions):
+      Promise<WorkerResult> {
+    return new Promise<WorkerResult>((resolve, reject) => {
+      this.queue.push({
+        request,
+        timeout,
+        attemptsLeft: Math.max(0, retry) + 1,
+        resolve,
+        reject,
+      });
+      this.pump();
+    });
+  }
+
+  private pump(): void {
+    while (this.queue.length > 0) {
+      const slot = this.slots.findIndex((w) => w && w.alive && !w.busy);
+      if (slot < 0) {
+        return;
+      }
+      this.dispatch(this.slots[slot]!, this.queue.shift()!);
+    }
+  }
+
+  private dispatch(worker: Worker, job: Job): void {
+    job.attemptsLeft -= 1;
+
+    let settled = false;
+    const timer = setTimeout(() => {
+      if (settled) {
+        return;
+      }
+      settled = true;
+      this.emit('timeout', {worker: worker.id, timeout: job.timeout});
+      // The worker is not answering, so the only way to get the slot back is to
+      // take the process down. The exit handler refills the slot.
+      worker.kill();
+      this.retryOrFail(
+          job,
+          new Error(`shotium: no answer within ${job.timeout}ms`));
+    }, job.timeout);
+
+    worker.send(job.request)
+        .then((result) => {
+          if (settled) {
+            return;
+          }
+          settled = true;
+          clearTimeout(timer);
+          job.resolve(result);
+          this.pump();
+        })
+        .catch((error: Error) => {
+          if (settled) {
+            return;
+          }
+          settled = true;
+          clearTimeout(timer);
+          this.retryOrFail(job, error);
+        });
+  }
+
+  private retryOrFail(job: Job, error: Error): void {
+    // A request rejected on its own merits -- a bad selector, an unreadable
+    // file -- would fail the same way every time, but the worker also rejects
+    // with the same shape when it dies. Retrying both is the safe direction:
+    // the cost of a pointless retry is one more render, and the cost of not
+    // retrying a crash is a failure the caller cannot do anything about.
+    if (job.attemptsLeft > 0 && !this.stopping) {
+      this.queue.unshift(job);
+      this.pump();
+      return;
+    }
+    job.reject(error);
+    this.pump();
+  }
+
+  async stop(): Promise<void> {
+    this.stopping = true;
+    for (const job of this.queue.splice(0)) {
+      job.reject(new Error('shotium: the runtime was stopped'));
+    }
+    await Promise.all(this.slots.map((worker) => new Promise<void>((resolve) => {
+      if (!worker || !worker.alive) {
+        resolve();
+        return;
+      }
+      worker.once('exit', () => resolve());
+      worker.stop();
+    })));
+    this.slots = [];
+  }
+}
+
+export {Pool, defaultCacheDir};

package/src/lib/protocol.ts

@@ -0,0 +1,53 @@
+// The wire format shotium.exe --serve speaks, in both directions: a 4-byte
+// little-endian length followed by that many bytes.
+//
+// Length-prefixed rather than line-delimited because the payload is binary and
+// a newline inside a PNG is not a message boundary. See shot/shot_server.h for
+// the same description from the other end.
+//
+//   ->  [len][{"file":"...","width":1248,...}]
+//   <-  [len][{"ok":true,"bytes":97756}]  [len][<PNG bytes>]
+//   <-  [len][{"ok":false,"error":"..."}] [0]
+
+const HEADER_BYTES = 4;
+
+function encodeFrame(payload: Buffer): Buffer {
+  const header = Buffer.allocUnsafe(HEADER_BYTES);
+  header.writeUInt32LE(payload.length, 0);
+  return Buffer.concat([header, payload]);
+}
+
+function encodeRequest(request: unknown): Buffer {
+  return encodeFrame(Buffer.from(JSON.stringify(request), 'utf8'));
+}
+
+// Reassembles frames out of whatever sizes the pipe hands over.
+//
+// A stream is not a sequence of messages: one read can carry half a header, or
+// three responses and the start of a fourth. Everything downstream assumes
+// whole frames, so this is the only place that has to know that.
+class FrameReader {
+  private buffer: Buffer = Buffer.alloc(0);
+
+  push(chunk: Buffer): void {
+    this.buffer = this.buffer.length === 0 ?
+        chunk :
+        Buffer.concat([this.buffer, chunk]);
+  }
+
+  // The next complete frame, or null when there is not one yet.
+  next(): Buffer|null {
+    if (this.buffer.length < HEADER_BYTES) {
+      return null;
+    }
+    const length = this.buffer.readUInt32LE(0);
+    if (this.buffer.length < HEADER_BYTES + length) {
+      return null;
+    }
+    const frame = this.buffer.subarray(HEADER_BYTES, HEADER_BYTES + length);
+    this.buffer = this.buffer.subarray(HEADER_BYTES + length);
+    return frame;
+  }
+}
+
+export {HEADER_BYTES, encodeFrame, encodeRequest, FrameReader};

package/src/lib/request.ts

@@ -0,0 +1,108 @@
+import type {Clip, PageGotoParams, ScreenshotOptions} from '../types.js';
+
+const DEFAULT_TIMEOUT_MS = 30000;
+// How much longer than the page's own deadline a supervisor waits before
+// deciding the worker is not going to answer at all. The worker fails a slow
+// page by itself and replies; this margin covers process startup and the
+// encode, and firing it means something worse than a slow page.
+const SUPERVISOR_MARGIN_MS = 10000;
+
+// What actually goes down the pipe. It is ScreenshotOptions with the viewport
+// flattened and `retry` taken out -- see toRequest below for why each.
+export interface WireRequest {
+  file: string;
+  type?: 'png'|'jpeg'|'webp';
+  fullPage?: boolean;
+  selector?: string;
+  quality?: number;
+  scale?: number;
+  omitBackground?: boolean;
+  path?: string;
+  pageGotoParams?: PageGotoParams;
+  clip?: Clip;
+  allowFileAccess?: boolean;
+  width?: number;
+  height?: number;
+}
+
+// Everything the worker understands, and nothing else. An unknown field is a
+// typo, and a typo that is silently dropped is a screenshot that quietly
+// ignored what was asked for -- so this rejects rather than filters.
+//
+// It is a runtime check even though the argument has a type, because the
+// argument having a type says nothing about a caller who is not compiled
+// against it: a JavaScript program, or a JSON body from somewhere else.
+const WIRE_FIELDS = new Set([
+  'file',
+  'type',
+  'fullPage',
+  'selector',
+  'quality',
+  'scale',
+  'omitBackground',
+  'path',
+  'pageGotoParams',
+  'clip',
+  'viewport',
+  'allowFileAccess',
+]);
+
+// One ScreenshotOptions, checked and flattened into what goes on the wire.
+//
+// It lives here rather than in index.ts because the in-process pool and the
+// daemon both send it: a request that is valid through one entry point and
+// rejected through the other would be a difference nobody asked for.
+function toRequest(options: ScreenshotOptions): WireRequest {
+  if (!options || typeof options !== 'object') {
+    throw new TypeError('shotium: screenshot(options) needs an object');
+  }
+  if (typeof options.file !== 'string' || options.file.length === 0) {
+    throw new TypeError('shotium: options.file is required');
+  }
+
+  const request: Record<string, unknown> = {};
+  for (const [key, value] of Object.entries(options)) {
+    if (value === undefined) {
+      continue;
+    }
+    // retry is the supervisor's, not the worker's: it decides how many times a
+    // request is re-sent, which is not something the worker could act on.
+    if (key === 'retry') {
+      continue;
+    }
+    if (!WIRE_FIELDS.has(key)) {
+      throw new TypeError(`shotium: unknown option "${key}"`);
+    }
+    request[key] = value;
+  }
+
+  // The viewport is flattened because the worker takes width and height at the
+  // top level: it is one screenshot's frame, not a nested object on the wire.
+  if (request.viewport) {
+    const {width, height} = request.viewport as {
+      width?: number,
+      height?: number,
+    };
+    delete request.viewport;
+    if (width !== undefined) {
+      request.width = width;
+    }
+    if (height !== undefined) {
+      request.height = height;
+    }
+  }
+  return request as unknown as WireRequest;
+}
+
+function timeoutFor(options: ScreenshotOptions): number {
+  const timeout = options.pageGotoParams && options.pageGotoParams.timeout;
+  return typeof timeout === 'number' ? timeout : DEFAULT_TIMEOUT_MS;
+}
+
+export {
+  DEFAULT_TIMEOUT_MS,
+  SUPERVISOR_MARGIN_MS,
+  WIRE_FIELDS,
+  timeoutFor,
+  toRequest,
+};

package/src/lib/worker.ts

@@ -0,0 +1,220 @@
+import {spawn} from 'node:child_process';
+import type {ChildProcess, StdioOptions} from 'node:child_process';
+import {EventEmitter} from 'node:events';
+
+import {FrameReader, encodeRequest} from './protocol.js';
+import type {WireRequest} from './request.js';
+
+export interface WorkerOptions {
+  id: number;
+  binary: string;
+  args?: string[];
+}
+
+// The header frame the worker answers with, followed by the image frame.
+export interface ResponseHeader {
+  ok: boolean;
+  error?: string;
+  bytes?: number;
+  path?: string;
+}
+
+export interface WorkerResult {
+  header: ResponseHeader;
+  image: Buffer|null;
+}
+
+interface Pending {
+  resolve: (result: WorkerResult) => void;
+  reject: (error: Error) => void;
+}
+
+// One shotium.exe --serve process.
+//
+// Exactly one request is in flight at a time, and that is not a simplification:
+// blink is a process-wide singleton bound to the worker's main thread, so a
+// second request could not be rendered concurrently even if the protocol
+// allowed it. Concurrency is the pool's job, and it gets it by running more
+// processes.
+//
+// Events:
+//   ready                       the process has started
+//   exit   {code, signal}       it is gone, for any reason
+//   crash  {code, signal}       it is gone while it owed an answer
+//   stderr {line}               a diagnostic line, useful when a render is wrong
+class Worker extends EventEmitter {
+  readonly id: number;
+  // How many requests this process has answered, either way. The pool reads
+  // it to tell a worker that was working and then died from one that never
+  // came up at all.
+  served = 0;
+
+  private readonly binary: string;
+  private readonly args: string[];
+  private process: ChildProcess|null = null;
+  private pending: Pending|null = null;
+  private stopping = false;
+  private reader = new FrameReader();
+  private header: ResponseHeader|null = null;
+  private stderr = '';
+
+  constructor(options: WorkerOptions) {
+    super();
+    this.id = options.id;
+    this.binary = options.binary;
+    this.args = options.args || [];
+    this.start();
+  }
+
+  get busy(): boolean {
+    return this.pending !== null;
+  }
+
+  get alive(): boolean {
+    return this.process !== null && this.process.exitCode === null &&
+        !this.stopping;
+  }
+
+  private start(): void {
+    const stdio: StdioOptions = ['pipe', 'pipe', 'pipe'];
+    const child = spawn(this.binary, ['--serve', ...this.args], {
+      stdio,
+      windowsHide: true,
+      // Detached, which on Windows means DETACHED_PROCESS: no console, and so
+      // no conhost.exe beside every worker. Four of those cost 40 MB of
+      // working set for a console nothing writes to -- the worker's output is
+      // three pipes.
+      //
+      // It does not outlive its supervisor despite the name: the worker exits
+      // when its stdin closes, and stdin closes when this process dies.
+      detached: true,
+    });
+    this.process = child;
+
+    child.stdout?.on('data', (chunk: Buffer) => this.onStdout(chunk));
+    child.stderr?.on('data', (chunk: Buffer) => this.onStderr(chunk));
+    child.on('error', (error) => this.onGone(null, null, error));
+    child.on('exit', (code, signal) => this.onGone(code, signal, null));
+
+    // The process is up as soon as spawn resolves the executable; there is no
+    // handshake in the protocol, and adding one would only move the failure --
+    // a binary that cannot start fails the first request just as visibly.
+    this.emit('ready');
+  }
+
+  private onStdout(chunk: Buffer): void {
+    this.reader.push(chunk);
+    for (;;) {
+      const frame = this.reader.next();
+      if (frame === null) {
+        return;
+      }
+      if (this.header === null) {
+        // Frame one of two: the JSON header.
+        try {
+          this.header = JSON.parse(frame.toString('utf8')) as ResponseHeader;
+        } catch {
+          this.fail(new Error(
+              `shotium: worker ${this.id} sent a header that is not JSON`));
+          return;
+        }
+        continue;
+      }
+      // Frame two: the image, empty when the header reported a failure or when
+      // the worker was asked to write the file itself.
+      const header = this.header;
+      this.header = null;
+      this.settle(header, frame);
+    }
+  }
+
+  private onStderr(chunk: Buffer): void {
+    this.stderr += chunk.toString('utf8');
+    const lines = this.stderr.split(/\r?\n/);
+    this.stderr = lines.pop() ?? '';
+    for (const line of lines) {
+      if (line.length > 0) {
+        this.emit('stderr', {worker: this.id, line});
+      }
+    }
+  }
+
+  private onGone(
+      code: number|null, signal: NodeJS.Signals|null,
+      error: Error|null): void {
+    const wasOwed = this.pending !== null;
+    this.process = null;
+    if (wasOwed) {
+      // A worker that dies mid-request is indistinguishable from one that never
+      // answered, which is the point: the retry path does not have to tell a
+      // crash from a hang.
+      this.fail(
+          error ||
+          new Error(`shotium: worker ${this.id} exited (code ${code}, signal ${
+              signal}) with a request in flight`));
+      this.emit('crash', {worker: this.id, code, signal});
+    } else if (error) {
+      this.emit('error', error);
+    }
+    this.emit('exit', {worker: this.id, code, signal});
+  }
+
+  private settle(header: ResponseHeader, payload: Buffer): void {
+    const pending = this.pending;
+    if (!pending) {
+      return;
+    }
+    this.pending = null;
+    this.served += 1;
+    if (header.ok) {
+      pending.resolve({header, image: header.path ? null : payload});
+    } else {
+      pending.reject(new Error(header.error || 'shotium: request failed'));
+    }
+  }
+
+  private fail(error: Error): void {
+    const pending = this.pending;
+    if (!pending) {
+      return;
+    }
+    this.pending = null;
+    pending.reject(error);
+  }
+
+  // Sends one request. Rejects if the worker dies before answering; there is no
+  // timeout here, because the deadline belongs to whoever owns the retry
+  // policy.
+  send(request: WireRequest): Promise<WorkerResult> {
+    if (this.pending) {
+      return Promise.reject(
+          new Error(`shotium: worker ${this.id} is already busy`));
+    }
+    if (!this.alive) {
+      return Promise.reject(new Error(`shotium: worker ${this.id} is not up`));
+    }
+    return new Promise<WorkerResult>((resolve, reject) => {
+      this.pending = {resolve, reject};
+      this.process?.stdin?.write(encodeRequest(request), (error) => {
+        if (error) {
+          this.fail(error);
+        }
+      });
+    });
+  }
+
+  // Closing stdin is the shutdown message: the worker sees the frame stream end
+  // on a frame boundary and exits 0. kill() is for when it has stopped
+  // listening.
+  stop(): void {
+    this.stopping = true;
+    this.process?.stdin?.end();
+  }
+
+  kill(): void {
+    this.stopping = true;
+    this.process?.kill();
+  }
+}
+
+export {Worker};