tywrap 0.1.2 → 0.2.1

package/src/runtime/worker-pool.ts

@@ -0,0 +1,498 @@
+/**
+ * WorkerPool - Manages multiple Transport instances for concurrent request handling.
+ *
+ * Provides semaphore-based concurrency control with configurable limits per worker
+ * and a wait queue for callers when all workers are at capacity.
+ *
+ * @see https://github.com/bbopen/tywrap/issues/149
+ */
+
+import { BoundedContext } from './bounded-context.js';
+import { BridgeTimeoutError, BridgeExecutionError, BridgeProtocolError } from './errors.js';
+import type { Transport } from './transport.js';
+
+// =============================================================================
+// TYPES
+// =============================================================================
+
+/**
+ * Configuration options for the WorkerPool.
+ */
+export interface WorkerPoolOptions {
+  /** Factory function to create transports */
+  createTransport: () => Transport;
+
+  /** Maximum number of workers in the pool */
+  maxWorkers: number;
+
+  /** Minimum number of workers to pre-spawn during init. Default: 0 (lazy) */
+  minWorkers?: number;
+
+  /** Timeout for waiting in queue (ms). Default: 30000 */
+  queueTimeoutMs?: number;
+
+  /** Maximum concurrent requests per worker. Default: 1 */
+  maxConcurrentPerWorker?: number;
+
+  /**
+   * Callback invoked after each worker is created and initialized.
+   * Use this for per-worker warmup (e.g., importing modules, running setup).
+   */
+  onWorkerReady?: (worker: PooledWorker) => Promise<void>;
+}
+
+/**
+ * A pooled worker with its transport and current in-flight request count.
+ */
+export interface PooledWorker {
+  /** The underlying transport instance */
+  transport: Transport;
+
+  /** Number of requests currently being processed by this worker */
+  inFlightCount: number;
+}
+
+/**
+ * Internal representation of a waiter in the queue.
+ */
+interface QueuedWaiter {
+  /** Resolve function to fulfill the promise with a worker */
+  resolve: (worker: PooledWorker) => void;
+
+  /** Reject function to reject the promise with an error */
+  reject: (error: Error) => void;
+
+  /** Timeout timer for queue timeout */
+  timer: NodeJS.Timeout;
+}
+
+// =============================================================================
+// WORKER POOL
+// =============================================================================
+
+/**
+ * Pool of Transport workers with semaphore-based concurrency control.
+ *
+ * Features:
+ * - Lazy worker creation (workers created on demand)
+ * - Configurable concurrency per worker
+ * - Wait queue with timeout for callers when pool is at capacity
+ * - Automatic cleanup of timers and workers on disposal
+ *
+ * @example
+ * ```typescript
+ * const pool = new WorkerPool({
+ *   createTransport: () => new ProcessIO({ pythonPath: 'python3' }),
+ *   maxWorkers: 4,
+ *   maxConcurrentPerWorker: 2,
+ *   queueTimeoutMs: 5000,
+ * });
+ *
+ * await pool.init();
+ *
+ * // Use withWorker for automatic acquire/release
+ * const result = await pool.withWorker(async (worker) => {
+ *   return worker.transport.send(message, timeout);
+ * });
+ *
+ * await pool.dispose();
+ * ```
+ */
+export class WorkerPool extends BoundedContext {
+  private readonly options: Omit<Required<WorkerPoolOptions>, 'onWorkerReady'> & {
+    onWorkerReady?: (worker: PooledWorker) => Promise<void>;
+  };
+  private readonly workers: PooledWorker[] = [];
+  private readonly waitQueue: QueuedWaiter[] = [];
+  /** Tracks workers being created to prevent race condition in acquire() */
+  private pendingCreations = 0;
+
+  /**
+   * Create a new WorkerPool.
+   *
+   * @param options - Pool configuration options
+   */
+  constructor(options: WorkerPoolOptions) {
+    super();
+
+    // Validate required options
+    if (typeof options.createTransport !== 'function') {
+      throw new BridgeExecutionError('createTransport must be a function');
+    }
+    if (typeof options.maxWorkers !== 'number' || options.maxWorkers < 1) {
+      throw new BridgeExecutionError('maxWorkers must be a positive number');
+    }
+
+    const minWorkers = options.minWorkers ?? 0;
+    if (minWorkers > options.maxWorkers) {
+      throw new BridgeExecutionError('minWorkers cannot exceed maxWorkers');
+    }
+
+    this.options = {
+      createTransport: options.createTransport,
+      maxWorkers: options.maxWorkers,
+      minWorkers,
+      queueTimeoutMs: options.queueTimeoutMs ?? 30000,
+      maxConcurrentPerWorker: options.maxConcurrentPerWorker ?? 1,
+      onWorkerReady: options.onWorkerReady,
+    };
+  }
+
+  // ===========================================================================
+  // LIFECYCLE
+  // ===========================================================================
+
+  /**
+   * Initialize the pool.
+   *
+   * If minWorkers > 0, pre-spawns workers during initialization.
+   * Otherwise, workers are created lazily on demand.
+   */
+  protected async doInit(): Promise<void> {
+    // Pre-spawn minimum workers if configured
+    if (this.options.minWorkers > 0) {
+      const spawns: Promise<PooledWorker>[] = [];
+      for (let i = 0; i < this.options.minWorkers; i++) {
+        spawns.push(this.createWorker());
+      }
+      await Promise.all(spawns);
+    }
+  }
+
+  /**
+   * Dispose the pool and all workers.
+   *
+   * - Rejects all waiters in the queue
+   * - Disposes all transport instances
+   * - Clears internal state
+   */
+  protected async doDispose(): Promise<void> {
+    // Reject all waiters in the queue
+    for (const waiter of this.waitQueue) {
+      clearTimeout(waiter.timer);
+      waiter.reject(new BridgeExecutionError('Pool disposed'));
+    }
+    this.waitQueue.length = 0;
+
+    // Dispose all workers
+    const errors: Error[] = [];
+    for (const worker of this.workers) {
+      try {
+        await worker.transport.dispose();
+      } catch (e) {
+        errors.push(e instanceof Error ? e : new Error(String(e)));
+      }
+    }
+    this.workers.length = 0;
+
+    // Report errors if any
+    if (errors.length === 1) {
+      throw errors[0];
+    }
+    if (errors.length > 1) {
+      throw new AggregateError(errors, 'Multiple errors during worker disposal');
+    }
+  }
+
+  // ===========================================================================
+  // WORKER MANAGEMENT
+  // ===========================================================================
+
+  /**
+   * Acquire a worker from the pool.
+   *
+   * This method:
+   * - Returns an available worker if one exists (inFlightCount < maxConcurrentPerWorker)
+   * - Creates a new worker if under the maxWorkers limit
+   * - Waits in queue if all workers are at capacity
+   *
+   * @returns Promise resolving to a pooled worker
+   * @throws BridgeTimeoutError if queue timeout expires
+   * @throws BridgeExecutionError if pool is disposed while waiting
+   */
+  async acquire(): Promise<PooledWorker> {
+    // Check for disposed state
+    if (this.isDisposed) {
+      throw new BridgeExecutionError('Pool has been disposed');
+    }
+
+    // Find an available worker (one with capacity)
+    const availableWorker = this.findAvailableWorker();
+    if (availableWorker) {
+      availableWorker.inFlightCount++;
+      return availableWorker;
+    }
+
+    // Create a new worker if under the limit
+    // Include pendingCreations to prevent race condition where multiple
+    // concurrent acquire() calls all pass the length check before any
+    // worker is actually added to the array
+    if (this.workers.length + this.pendingCreations < this.options.maxWorkers) {
+      this.pendingCreations++;
+      try {
+        const newWorker = await this.createWorker();
+        newWorker.inFlightCount++;
+        return newWorker;
+      } finally {
+        this.pendingCreations--;
+      }
+    }
+
+    // All workers at capacity - wait in queue
+    return this.waitForWorker();
+  }
+
+  /**
+   * Release a worker back to the pool.
+   *
+   * Decrements the worker's in-flight count and notifies any waiters
+   * that a worker may be available.
+   *
+   * @param worker - The worker to release
+   */
+  release(worker: PooledWorker): void {
+    // Validate the worker belongs to this pool
+    if (!this.workers.includes(worker)) {
+      return;
+    }
+
+    // Decrement in-flight count (minimum 0)
+    worker.inFlightCount = Math.max(0, worker.inFlightCount - 1);
+
+    // If there are waiters and this worker has capacity, fulfill the first waiter
+    if (this.waitQueue.length > 0 && worker.inFlightCount < this.options.maxConcurrentPerWorker) {
+      const waiter = this.waitQueue.shift();
+      if (waiter) {
+        clearTimeout(waiter.timer);
+        worker.inFlightCount++;
+        waiter.resolve(worker);
+      }
+    }
+  }
+
+  /**
+   * Execute a function with an acquired worker, automatically releasing afterward.
+   *
+   * This is the recommended way to use the pool, as it ensures proper cleanup
+   * even if the function throws an error.
+   *
+   * @param fn - Async function to execute with the worker
+   * @returns Promise resolving to the function's return value
+   *
+   * @example
+   * ```typescript
+   * const result = await pool.withWorker(async (worker) => {
+   *   return worker.transport.send(message, timeout);
+   * });
+   * ```
+   */
+  async withWorker<T>(fn: (worker: PooledWorker) => Promise<T>): Promise<T> {
+    const worker = await this.acquire();
+    let workerRemoved = false;
+
+    try {
+      return await fn(worker);
+    } catch (error) {
+      // If this is a fatal error indicating the worker is dead, remove it from the pool
+      if (this.isFatalWorkerError(error)) {
+        this.removeWorker(worker);
+        workerRemoved = true;
+      }
+      throw error;
+    } finally {
+      // Only release if worker wasn't removed due to fatal error
+      if (!workerRemoved) {
+        this.release(worker);
+      }
+    }
+  }
+
+  // ===========================================================================
+  // WORKER HEALTH
+  // ===========================================================================
+
+  /**
+   * Check if an error indicates the worker is dead and should be removed.
+   *
+   * Fatal errors include:
+   * - Process not running
+   * - Process exited unexpectedly
+   * - Pipe errors (EPIPE)
+   * - Connection reset errors (ECONNRESET)
+   */
+  private isFatalWorkerError(error: unknown): boolean {
+    if (error instanceof BridgeProtocolError) {
+      const msg = error.message.toLowerCase();
+      return (
+        msg.includes('not running') ||
+        msg.includes('process exited') ||
+        msg.includes('epipe') ||
+        msg.includes('econnreset')
+      );
+    }
+    return false;
+  }
+
+  /**
+   * Remove a worker from the pool.
+   *
+   * This is called when a worker is detected as dead (crashed, pipe error, etc.).
+   * The worker's transport is disposed in the background.
+   */
+  private removeWorker(worker: PooledWorker): void {
+    const index = this.workers.indexOf(worker);
+    if (index !== -1) {
+      this.workers.splice(index, 1);
+      // Dispose transport in background - don't await to avoid blocking
+      worker.transport.dispose().catch(() => {
+        // Ignore disposal errors for dead workers
+      });
+    }
+  }
+
+  // ===========================================================================
+  // POOL STATISTICS
+  // ===========================================================================
+
+  /**
+   * Current number of workers in the pool.
+   */
+  get workerCount(): number {
+    return this.workers.length;
+  }
+
+  /**
+   * Number of callers waiting in the queue.
+   */
+  get queueLength(): number {
+    return this.waitQueue.length;
+  }
+
+  /**
+   * Total number of in-flight requests across all workers.
+   */
+  get totalInFlight(): number {
+    return this.workers.reduce((sum, w) => sum + w.inFlightCount, 0);
+  }
+
+  // ===========================================================================
+  // PRIVATE HELPERS
+  // ===========================================================================
+
+  /**
+   * Find an available worker with capacity for another request.
+   */
+  private findAvailableWorker(): PooledWorker | undefined {
+    return this.workers.find(w => w.inFlightCount < this.options.maxConcurrentPerWorker);
+  }
+
+  /**
+   * Create a new worker and add it to the pool.
+   *
+   * If onWorkerReady is configured, calls it after the transport is initialized.
+   * This is useful for per-worker warmup (importing modules, running setup).
+   */
+  private async createWorker(): Promise<PooledWorker> {
+    const transport = this.options.createTransport();
+
+    // Initialize the transport
+    await transport.init();
+
+    const worker: PooledWorker = {
+      transport,
+      inFlightCount: 0,
+    };
+
+    this.workers.push(worker);
+
+    // Call onWorkerReady callback if provided
+    if (this.options.onWorkerReady) {
+      await this.options.onWorkerReady(worker);
+    }
+
+    return worker;
+  }
+
+  /**
+   * Wait in queue for a worker to become available.
+   */
+  private waitForWorker(): Promise<PooledWorker> {
+    return new Promise<PooledWorker>((resolve, reject) => {
+      const timer = setTimeout(() => {
+        // Remove this waiter from the queue
+        const index = this.waitQueue.findIndex(w => w.timer === timer);
+        if (index !== -1) {
+          this.waitQueue.splice(index, 1);
+        }
+        reject(
+          new BridgeTimeoutError(
+            `Timed out waiting for available worker after ${this.options.queueTimeoutMs}ms`
+          )
+        );
+      }, this.options.queueTimeoutMs);
+
+      // Unref the timer so it doesn't keep the Node.js process alive
+      if (typeof timer.unref === 'function') {
+        timer.unref();
+      }
+
+      this.waitQueue.push({ resolve, reject, timer });
+    });
+  }
+
+  // ===========================================================================
+  // RUNTIME EXECUTION (Not implemented - WorkerPool is just for worker management)
+  // ===========================================================================
+
+  /**
+   * Not implemented - WorkerPool does not execute Python calls directly.
+   * Use the BridgeProtocol layer with a pooled worker's transport.
+   */
+  async call<T = unknown>(
+    _module: string,
+    _functionName: string,
+    _args: unknown[],
+    _kwargs?: Record<string, unknown>
+  ): Promise<T> {
+    throw new BridgeExecutionError(
+      'WorkerPool does not implement call() - use withWorker() to get a transport'
+    );
+  }
+
+  /**
+   * Not implemented - WorkerPool does not execute Python calls directly.
+   */
+  async instantiate<T = unknown>(
+    _module: string,
+    _className: string,
+    _args: unknown[],
+    _kwargs?: Record<string, unknown>
+  ): Promise<T> {
+    throw new BridgeExecutionError(
+      'WorkerPool does not implement instantiate() - use withWorker() to get a transport'
+    );
+  }
+
+  /**
+   * Not implemented - WorkerPool does not execute Python calls directly.
+   */
+  async callMethod<T = unknown>(
+    _handle: string,
+    _methodName: string,
+    _args: unknown[],
+    _kwargs?: Record<string, unknown>
+  ): Promise<T> {
+    throw new BridgeExecutionError(
+      'WorkerPool does not implement callMethod() - use withWorker() to get a transport'
+    );
+  }
+
+  /**
+   * Not implemented - WorkerPool does not execute Python calls directly.
+   */
+  async disposeInstance(_handle: string): Promise<void> {
+    throw new BridgeExecutionError(
+      'WorkerPool does not implement disposeInstance() - use withWorker() to get a transport'
+    );
+  }
+}

package/src/utils/codec.ts

@@ -243,6 +243,101 @@ function isPromiseLike(value: unknown): value is PromiseLike<unknown> {
   );
 }
 
+/**
+ * Convert a typed array (Int32Array, Float64Array, BigInt64Array, etc.) to a plain JS array.
+ *
+ * Why: Arrow's column.toArray() returns typed arrays, but we need plain arrays for
+ * JSON-compatible output and proper nested array reshaping.
+ *
+ * @param arr - Typed array or plain array
+ * @returns Plain JavaScript array with values converted (BigInt → Number where safe)
+ */
+function typedArrayToPlain(arr: unknown): unknown[] {
+  if (Array.isArray(arr)) {
+    return arr;
+  }
+  // Handle typed arrays (Int32Array, Float64Array, BigInt64Array, etc.)
+  if (ArrayBuffer.isView(arr) && 'length' in arr) {
+    const typedArr = arr as unknown as { length: number; [index: number]: unknown };
+    const result: unknown[] = [];
+    for (let i = 0; i < typedArr.length; i++) {
+      const val = typedArr[i];
+      // Convert BigInt to Number if within safe integer range
+      if (typeof val === 'bigint') {
+        if (val >= BigInt(Number.MIN_SAFE_INTEGER) && val <= BigInt(Number.MAX_SAFE_INTEGER)) {
+          result.push(Number(val));
+        } else {
+          result.push(val); // Keep as BigInt if too large
+        }
+      } else {
+        result.push(val);
+      }
+    }
+    return result;
+  }
+  // Fallback: check if iterable before converting
+  if (arr !== null && arr !== undefined && typeof arr === 'object' && Symbol.iterator in arr) {
+    return Array.from(arr as Iterable<unknown>);
+  }
+  // Non-iterable: return empty array (shouldn't happen with valid Arrow data)
+  return [];
+}
+
+/**
+ * Extract values from an Arrow table as a plain JavaScript array.
+ *
+ * Why: Arrow decoding returns Table objects, not raw arrays. We need to extract
+ * the column values and convert any typed arrays to plain arrays.
+ */
+function extractArrowValues(data: unknown): unknown[] | null {
+  if (Array.isArray(data)) {
+    return data;
+  }
+  // Arrow table - extract values from first column
+  const table = data as ArrowTable & { getChildAt?: (i: number) => { toArray?: () => unknown } };
+  if (typeof table.getChildAt === 'function') {
+    const column = table.getChildAt(0);
+    if (column && typeof column.toArray === 'function') {
+      return typedArrayToPlain(column.toArray());
+    }
+  }
+  return null;
+}
+
+/**
+ * Reshape a flat array into a multi-dimensional nested array.
+ *
+ * Why: PyArrow's pa.array() only handles 1D arrays, so we flatten multi-dimensional
+ * arrays before Arrow encoding and reshape after decoding. This maintains Arrow's
+ * binary efficiency while working with current arrow-js (which doesn't yet support
+ * FixedShapeTensorArray). See: https://github.com/apache/arrow-js/issues/115
+ *
+ * @param flat - Flat array of values (must be a plain array, not typed array)
+ * @param shape - Target shape, e.g., [2, 3] for a 2x3 matrix
+ * @returns Nested array with the specified shape
+ */
+function reshapeArray(flat: unknown[], shape: readonly number[]): unknown {
+  if (shape.length === 0) {
+    return flat[0];
+  }
+  if (shape.length === 1) {
+    return flat;
+  }
+
+  // shape.length >= 2, so first is always defined
+  const first = shape[0]!;
+  const rest = shape.slice(1);
+  const chunkSize = rest.reduce((a, b) => a * b, 1);
+  const result: unknown[] = [];
+
+  for (let i = 0; i < first; i++) {
+    const chunk = flat.slice(i * chunkSize, (i + 1) * chunkSize);
+    result.push(reshapeArray(chunk, rest));
+  }
+
+  return result;
+}
+
 // Why: decoding needs to reject incompatible envelopes before we attempt to interpret payloads.
 const CODEC_VERSION = 1;
 
@@ -326,13 +421,43 @@ function decodeEnvelopeCore<T>(
 
   if (marker === 'ndarray') {
     const encoding = (value as { encoding?: unknown }).encoding;
+    const shapeValue = (value as { shape?: unknown }).shape;
+    const shape = isNumberArray(shapeValue) ? shapeValue : undefined;
+
     if (encoding === 'arrow') {
       const b64 = (value as { b64?: unknown }).b64;
       if (typeof b64 !== 'string') {
         throw new Error('Invalid ndarray envelope: missing b64');
       }
       const bytes = fromBase64(b64);
-      return decodeArrow(bytes);
+      const decoded = decodeArrow(bytes);
+
+      // Extract values from Arrow table and reshape if needed
+      // Arrow only handles 1D arrays, so we flatten on encode and reshape here
+      // Reshape for: scalars (shape.length === 0) and multi-dim (shape.length > 1)
+      // Skip reshape for: 1D arrays (shape.length === 1) - return as-is
+      if (isPromiseLike(decoded)) {
+        return decoded.then(data => {
+          const values = extractArrowValues(data);
+          if (!values) {
+            return data; // Fallback: return raw data if extraction fails
+          }
+          // Reshape scalars and multi-dimensional arrays, but not 1D
+          if (shape && shape.length !== 1) {
+            return reshapeArray(values, shape);
+          }
+          return values;
+        });
+      }
+      const values = extractArrowValues(decoded);
+      if (!values) {
+        return decoded; // Fallback: return raw data if extraction fails
+      }
+      // Reshape scalars and multi-dimensional arrays, but not 1D
+      if (shape && shape.length !== 1) {
+        return reshapeArray(values, shape);
+      }
+      return values;
     }
     if (encoding === 'json') {
       if (!('data' in (value as object))) {