threadforge 0.1.0

package/src/core/RoutingStrategy.js

@@ -0,0 +1,316 @@
+/**
+ * Routing Strategies
+ *
+ * When a service has multiple workers, the proxy needs to decide
+ * which worker handles each request. Different strategies suit
+ * different service patterns.
+ *
+ * Strategies:
+ *
+ *   round-robin     Default. Cycles through workers sequentially.
+ *                   Best for: stateless services, uniform workloads.
+ *
+ *   hash            Routes by a key (e.g., userId). Same key always
+ *                   goes to the same worker. Uses consistent hashing
+ *                   so adding/removing workers only remaps ~1/N keys.
+ *                   Best for: stateful services, in-memory caches,
+ *                   services that benefit from data locality.
+ *
+ *   least-pending   Routes to the worker with fewest in-flight requests.
+ *                   Best for: variable-latency services (DB queries,
+ *                   external API calls).
+ *
+ *   broadcast       Sends to ALL workers. Returns the first response.
+ *                   Best for: search/aggregation across sharded data.
+ *
+ *   primary         Always routes to worker 0. Others are standbys.
+ *                   Best for: singleton services (cron scheduler,
+ *                   leader election).
+ */
+
+import { createHmac, randomBytes } from "node:crypto";
+
+const HASH_SALT = randomBytes(16).toString("hex");
+
+/**
+ * @typedef {Object} WorkerEntry
+ * @property {string} key    - "serviceName:workerId"
+ * @property {Object} socket - UDS socket or local reference
+ * @property {number} pending - Number of in-flight requests
+ */
+
+// ─── Round-Robin ────────────────────────────────────────────
+
+export class RoundRobinStrategy {
+  constructor() {
+    this.index = 0;
+  }
+
+  /**
+   * @param {WorkerEntry[]} workers
+   * @param {Object} _callContext - Ignored by round-robin
+   * @returns {WorkerEntry}
+   */
+  pick(workers, _callContext) {
+    if (workers.length === 0) return null;
+    const worker = workers[this.index % workers.length];
+    this.index = (this.index + 1) % 1_000_000_000;
+    return worker;
+  }
+
+  get name() {
+    return "round-robin";
+  }
+}
+
+// ─── Hash-Based (Consistent Hashing) ────────────────────────
+
+export class HashStrategy {
+  /**
+   * @param {Object} options
+   * @param {string} options.key - The field name to hash on (e.g., 'userId')
+   *                               If the call payload has this field, it's used.
+   *                               Otherwise falls back to round-robin.
+   * @param {number} [options.vnodes=150] - Virtual nodes per worker for
+   *                                        consistent hash ring.
+   */
+  constructor(options = {}) {
+    this.keyField = options.key ?? null;
+    this.vnodes = options.vnodes ?? 150;
+    this._ring = null;
+    this._workerCacheKey = '';
+  }
+
+  pick(workers, callContext) {
+    if (workers.length === 0) return null;
+    if (workers.length === 1) return workers[0];
+
+    // Extract the hash key from the call context
+    const hashKey = this._extractKey(callContext);
+    if (!hashKey) {
+      // No key available — fall back to round-robin-ish
+      return workers[Math.abs(this._simpleHash(Date.now().toString())) % workers.length];
+    }
+
+    // Rebuild ring when worker identities change (not just count).
+    // This ensures the ring is correct when workers restart with new keys.
+    const currentCacheKey = workers.map(w => w.key).sort().join(',');
+    if (this._workerCacheKey !== currentCacheKey) {
+      this._buildRing(workers);
+    }
+
+    // Find the worker on the ring
+    const hash = this._simpleHash(hashKey);
+    return this._findOnRing(hash, workers);
+  }
+
+  _extractKey(callContext) {
+    if (!callContext) return null;
+
+    // If a key field is specified, look for it in the call args
+    if (this.keyField) {
+      const args = callContext.__forge_args ?? [];
+      // Check first arg (usually the main parameter)
+      if (args.length > 0) {
+        if (typeof args[0] === "string" || typeof args[0] === "number") {
+          return String(args[0]);
+        }
+        if (typeof args[0] === "object" && args[0]?.[this.keyField]) {
+          return String(args[0][this.keyField]);
+        }
+      }
+    }
+
+    // Try to use __forge_hash if explicitly set by the caller
+    if (callContext.__forge_hash) {
+      return String(callContext.__forge_hash);
+    }
+
+    return null;
+  }
+
+  _buildRing(workers) {
+    // P5: Pre-compute all virtual node hashes and store sorted for binary search
+    const totalNodes = workers.length * this.vnodes;
+    const hashes = new Uint32Array(totalNodes);
+    const indices = new Uint16Array(totalNodes);
+    let pos = 0;
+    for (let i = 0; i < workers.length; i++) {
+      for (let v = 0; v < this.vnodes; v++) {
+        hashes[pos] = this._computeHash(`${workers[i].key}:${v}`);
+        indices[pos] = i;
+        pos++;
+      }
+    }
+    // Sort by hash value, keeping indices aligned
+    const sortOrder = Array.from({ length: totalNodes }, (_, i) => i);
+    sortOrder.sort((a, b) => hashes[a] - hashes[b]);
+    this._ringHashes = new Uint32Array(totalNodes);
+    this._ringIndices = new Uint16Array(totalNodes);
+    for (let i = 0; i < totalNodes; i++) {
+      this._ringHashes[i] = hashes[sortOrder[i]];
+      this._ringIndices[i] = indices[sortOrder[i]];
+    }
+    this._ring = true; // sentinel
+    this._workerCacheKey = workers.map(w => w.key).sort().join(',');
+  }
+
+  _findOnRing(hash, workers) {
+    const hashes = this._ringHashes;
+    // Binary search for the first ring entry >= hash
+    let lo = 0,
+      hi = hashes.length;
+    while (lo < hi) {
+      const mid = (lo + hi) >>> 1;
+      if (hashes[mid] < hash) lo = mid + 1;
+      else hi = mid;
+    }
+    // Wrap around
+    const idx = lo < hashes.length ? lo : 0;
+    return workers[this._ringIndices[idx]];
+  }
+
+  /**
+   * Pre-compute a keyed SHA-256 hash for ring building (called at build time).
+   */
+  _computeHash(str) {
+    return createHmac("sha256", HASH_SALT).update(str).digest().readUInt32BE(0);
+  }
+
+  /**
+   * Keyed SHA-256 hash — collision-resistant, prevents hash flooding attacks.
+   */
+  _simpleHash(str) {
+    return createHmac("sha256", HASH_SALT).update(str).digest().readUInt32BE(0);
+  }
+
+  get name() {
+    return `hash(${this.keyField ?? "auto"})`;
+  }
+}
+
+// ─── Least-Pending ──────────────────────────────────────────
+
+export class LeastPendingStrategy {
+  constructor() {
+    /** @type {Map<string, number>} key → pending count */
+    this._pending = new Map();
+  }
+
+  /**
+   * Pick the worker with the fewest in-flight requests.
+   */
+  pick(workers, _callContext) {
+    if (workers.length === 0) return null;
+    if (workers.length === 1) return workers[0];
+
+    // Prune stale entries for workers no longer in the current set
+    if (this._pending.size > 0) {
+      const activeKeys = new Set(workers.map(w => w.key));
+      for (const key of this._pending.keys()) {
+        if (!activeKeys.has(key)) {
+          this._pending.delete(key);
+        }
+      }
+    }
+
+    let best = workers[0];
+    let bestPending = this._pending.get(best.key) ?? 0;
+    for (let i = 1; i < workers.length; i++) {
+      const p = this._pending.get(workers[i].key) ?? 0;
+      if (p < bestPending) {
+        best = workers[i];
+        bestPending = p;
+      }
+    }
+    return best;
+  }
+
+  /** Call when dispatching a request to a worker */
+  acquire(workerKey) {
+    this._pending.set(workerKey, (this._pending.get(workerKey) ?? 0) + 1);
+  }
+
+  /** Call when a request to a worker completes */
+  release(workerKey) {
+    const count = this._pending.get(workerKey) ?? 0;
+    if (count <= 0) {
+      if (process.env.NODE_ENV !== 'production') {
+        console.warn('[LeastPendingStrategy] Spurious release()');
+      }
+      this._pending.delete(workerKey);
+      return;
+    }
+    if (count <= 1) {
+      this._pending.delete(workerKey);
+    } else {
+      this._pending.set(workerKey, count - 1);
+    }
+  }
+
+  get name() {
+    return "least-pending";
+  }
+}
+
+// ─── Broadcast ──────────────────────────────────────────────
+
+export class BroadcastStrategy {
+  /**
+   * Returns ALL workers. The caller is responsible for sending
+   * to all of them and handling multiple responses.
+   */
+  pick(workers, _callContext) {
+    return workers; // Return the full array — caller detects this
+  }
+
+  get name() {
+    return "broadcast";
+  }
+
+  get isBroadcast() {
+    return true;
+  }
+}
+
+// ─── Primary ────────────────────────────────────────────────
+
+export class PrimaryStrategy {
+  /**
+   * Always routes to worker index 0. If worker 0 is down,
+   * routes to the next available.
+   */
+  pick(workers, _callContext) {
+    if (workers.length === 0) return null;
+    return workers[0]; // Primary is always first
+  }
+
+  get name() {
+    return "primary";
+  }
+}
+
+// ─── Strategy Factory ───────────────────────────────────────
+
+const STRATEGIES = {
+  "round-robin": (_opts) => new RoundRobinStrategy(),
+  hash: (opts) => new HashStrategy(opts),
+  "least-pending": (_opts) => new LeastPendingStrategy(),
+  broadcast: (_opts) => new BroadcastStrategy(),
+  primary: (_opts) => new PrimaryStrategy(),
+};
+
+/**
+ * Create a routing strategy by name.
+ *
+ * @param {string} name - Strategy name
+ * @param {Object} [options] - Strategy-specific options
+ * @returns {Object} Strategy instance with a pick() method
+ */
+export function createStrategy(name, options = {}) {
+  const factory = STRATEGIES[name];
+  if (!factory) {
+    throw new Error(`Unknown routing strategy: "${name}". Available: ${Object.keys(STRATEGIES).join(", ")}`);
+  }
+  return factory(options);
+}