threadforge 0.1.0

package/src/core/ThreadAllocator.js

@@ -0,0 +1,196 @@
+import os from "node:os";
+
+/**
+ * ThreadAllocator
+ *
+ * Distributes available CPU cores across services based on
+ * fixed thread counts and weighted auto-allocation.
+ */
+export class ThreadAllocator {
+  constructor(options = {}) {
+    // Allow override via environment variable or config
+    let envCpus = null;
+    if (process.env.FORGE_CPUS) {
+      const parsed = parseInt(process.env.FORGE_CPUS, 10);
+      if (!isNaN(parsed) && parsed > 0) {
+        envCpus = parsed;
+      } else {
+        console.warn(
+          `[ThreadForge] Warning: Invalid FORGE_CPUS value "${process.env.FORGE_CPUS}". ` +
+          `Expected a positive integer. Ignoring.`
+        );
+      }
+    }
+    
+    const rawCpus = options.cpus ?? envCpus ?? os.cpus().length;
+    
+    // Clamp to minimum of 1 to handle constrained environments where os.cpus() returns 0
+    this.totalCpus = Math.max(1, rawCpus);
+    
+    // Log warning if clamping occurred
+    if (rawCpus < 1) {
+      console.warn(
+        `[ThreadForge] Warning: CPU count detected as ${rawCpus}, clamping to minimum of 1. ` +
+        `This may occur in constrained sandbox environments. ` +
+        `Override with FORGE_CPUS environment variable if needed.`
+      );
+    }
+    
+    this.reservedForSupervisor = options.reserved ?? 1;
+  }
+
+  /**
+   * Calculate thread allocation for all services.
+   *
+   * @param {Object} services - Normalized service configs (keyed by name)
+   * @returns {Map<string, number>} - Service name → allocated thread count
+   */
+  allocate(services) {
+    const available = this.totalCpus - this.reservedForSupervisor;
+    const entries = Object.values(services);
+
+    if (entries.length === 0) {
+      return new Map();
+    }
+
+    // Phase 1: Honor fixed thread counts
+    let remaining = available;
+    const fixedServices = entries.filter((s) => typeof s.threads === "number");
+    const autoServices = entries.filter((s) => s.threads === "auto");
+
+    let fixedTotal = 0;
+    for (const s of fixedServices) {
+      fixedTotal += s.threads;
+    }
+
+    if (fixedTotal > available) {
+      console.warn(
+        `[ThreadForge] Warning: Fixed thread allocations (${fixedTotal}) exceed ` +
+          `available cores (${available}). Services may contend for resources.`,
+      );
+    }
+
+    remaining = Math.max(0, remaining - fixedTotal);
+
+    // Phase 2: Distribute remaining cores to 'auto' services by weight
+    const totalWeight = autoServices.reduce((sum, s) => sum + Math.max(1, s.weight), 0);
+
+    const allocation = new Map();
+
+    for (const s of fixedServices) {
+      allocation.set(s.name, s.threads);
+    }
+
+    if (autoServices.length > 0 && remaining > 0) {
+      // First pass: proportional allocation
+      let allocated = 0;
+      const provisional = [];
+
+      for (const s of autoServices) {
+        const share = Math.max(1, Math.round((s.weight / totalWeight) * remaining));
+        provisional.push({ name: s.name, threads: share });
+        allocated += share;
+      }
+
+      // Adjust if we over/under allocated due to rounding
+      const diff = remaining - allocated;
+      if (diff !== 0) {
+        // Sort by weight descending, give/take from highest weight services
+        provisional.sort((a, b) => {
+          const wA = services[a.name]?.weight ?? 1;
+          const wB = services[b.name]?.weight ?? 1;
+          return wB - wA;
+        });
+
+        let adjust = Math.abs(diff);
+        const direction = diff > 0 ? 1 : -1;
+        let i = 0;
+
+        let safetyCounter = 0;
+        const maxIterations = provisional.length * 2;
+        while (adjust > 0 && safetyCounter < maxIterations) {
+          safetyCounter++;
+          const newVal = provisional[i].threads + direction;
+          if (newVal >= 1) {
+            provisional[i].threads = newVal;
+            adjust--;
+          }
+          i++;
+          if (i >= provisional.length) i = 0;
+        }
+
+        if (adjust > 0) {
+          console.warn(`[ThreadAllocator] Could not fully adjust allocation (${adjust} threads remaining)`);
+        }
+      }
+
+      // Second pass: if total still exceeds remaining, reduce lowest-weight services
+      let totalProvisional = provisional.reduce((sum, p) => sum + p.threads, 0);
+      if (totalProvisional > remaining) {
+        // Sort by weight ascending so we reduce lowest-weight first
+        const byWeightAsc = [...provisional].sort((a, b) => {
+          const wA = services[a.name]?.weight ?? 1;
+          const wB = services[b.name]?.weight ?? 1;
+          return wA - wB;
+        });
+
+        for (const p of byWeightAsc) {
+          if (totalProvisional <= remaining) break;
+          const excess = totalProvisional - remaining;
+          const canReduce = p.threads - 1; // keep at least 1
+          if (canReduce <= 0) continue;
+          const reduction = Math.min(canReduce, excess);
+          p.threads -= reduction;
+          totalProvisional -= reduction;
+        }
+
+        if (totalProvisional > remaining) {
+          console.warn(
+            `[ThreadAllocator] Auto services allocated ${totalProvisional} threads ` +
+              `but only ${remaining} available. Some services may contend for resources.`,
+          );
+        }
+      }
+
+      for (const p of provisional) {
+        allocation.set(p.name, p.threads);
+      }
+    } else if (autoServices.length > 0) {
+      // No remaining cores, give each auto service 1 thread minimum
+      for (const s of autoServices) {
+        allocation.set(s.name, 1);
+      }
+    }
+
+    return allocation;
+  }
+
+  /**
+   * Get a summary of the allocation for display.
+   */
+  summary(services) {
+    const allocation = this.allocate(services);
+    let totalAllocated = 0;
+    const rows = [];
+
+    for (const [name, threads] of allocation) {
+      totalAllocated += threads;
+      const svc = services[name];
+      rows.push({
+        name,
+        port: svc.port,
+        threads,
+        mode: svc.mode,
+        weight: svc.weight,
+      });
+    }
+
+    return {
+      totalCpus: this.totalCpus,
+      reserved: this.reservedForSupervisor,
+      allocated: totalAllocated,
+      idle: Math.max(0, this.totalCpus - this.reservedForSupervisor - totalAllocated),
+      services: rows,
+    };
+  }
+}