Fix draining behaviour (#1180)

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Author: lukasIO

.changeset/plenty-laws-hug.md

@@ -0,0 +1,5 @@
+---
+"@livekit/agents": patch
+---
+
+Fix worker draining behaviour

agents/src/cli.ts

@@ -40,7 +40,11 @@ const runServer = async (args: CliArgs) => {
       process.exit(130); // SIGINT exit code
     });
     if (args.production) {
-      await server.drain();
+      try {
+        await server.drain();
+      } catch (e) {
+        logger.error(e);
+      }
     }
     await server.close();
     logger.debug('worker closed due to SIGINT.');
@@ -50,7 +54,11 @@ const runServer = async (args: CliArgs) => {
   process.once('SIGTERM', async () => {
     logger.debug('SIGTERM received in CLI.');
     if (args.production) {
-      await server.drain();
+      try {
+        await server.drain();
+      } catch (e) {
+        logger.error(e);
+      }
     }
     await server.close();
     logger.debug('worker closed due to SIGTERM.');

agents/src/ipc/job_executor.ts

@@ -1,6 +1,7 @@
 // SPDX-FileCopyrightText: 2024 LiveKit, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
+import type { Throws } from '@livekit/throws-transformer/throws';
 import type { RunningJobInfo } from '../job.js';
 
 export interface JobExecutor {
@@ -9,11 +10,11 @@ export interface JobExecutor {
   runningJob: RunningJobInfo | undefined;
   status: JobStatus;
 
-  start(): Promise<void>;
-  join(): Promise<void>;
-  initialize(): Promise<void>;
-  close(): Promise<void>;
-  launchJob(info: RunningJobInfo): Promise<void>;
+  start(): Promise<Throws<void, Error>>;
+  join(): Promise<Throws<void, Error>>;
+  initialize(): Promise<Throws<void, Error>>;
+  close(): Promise<Throws<void, Error>>;
+  launchJob(info: RunningJobInfo): Promise<Throws<void, Error>>;
 }
 
 export enum JobStatus {

agents/src/ipc/proc_pool.ts

@@ -2,6 +2,7 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 import { MultiMutex, Mutex } from '@livekit/mutex';
+import type { Throws } from '@livekit/throws-transformer/throws';
 import type { RunningJobInfo } from '../job.js';
 import { Queue } from '../utils.js';
 import type { InferenceExecutor } from './inference_executor.js';
@@ -53,7 +54,7 @@ export class ProcPool {
     return this.executors.find((x) => x.runningJob && x.runningJob.job.id === id) || null;
   }
 
-  async launchJob(info: RunningJobInfo) {
+  async launchJob(info: RunningJobInfo): Promise<Throws<void, Error>> {
     let proc: JobExecutor;
     if (this.procMutex) {
       const entry = await this.warmedProcQueue.get();

agents/src/utils.ts

@@ -837,7 +837,7 @@ export function waitUntilTimeout<T, E extends Error = IdleTimeoutError>(
   promise: Promise<T>,
   timeoutMs: number,
   throwError?: () => E,
-): Promise<Throws<T, E>> {
+): Promise<Throws<T, E | IdleTimeoutError>> {
   let timer: ReturnType<typeof setTimeout> | undefined;
   return Promise.race([
     promise,
@@ -968,11 +968,20 @@ export async function waitForAbort(signal: AbortSignal) {
     abortFuture.resolve();
     signal.removeEventListener('abort', handler);
   };
-
+  if (signal.aborted) {
+    return;
+  }
   signal.addEventListener('abort', handler, { once: true });
   return await abortFuture.await;
 }
 
+export async function rejectOnAbort(signal: AbortSignal): Promise<never> {
+  if (signal.aborted) throw signal.reason;
+  const abortFuture = new Future<never>();
+  signal.addEventListener('abort', () => abortFuture.reject(signal.reason), { once: true });
+  return abortFuture.await;
+}
+
 /**
  * Combines two abort signals into a single abort signal.
  * @param a - The first abort signal.

agents/src/worker.ts

@@ -10,6 +10,7 @@ import {
   WorkerMessage,
   WorkerStatus,
 } from '@livekit/protocol';
+import type { Throws } from '@livekit/throws-transformer/throws';
 import type { ParticipantInfo } from 'livekit-server-sdk';
 import { AccessToken, RoomServiceClient } from 'livekit-server-sdk';
 import { EventEmitter } from 'node:events';
@@ -22,7 +23,7 @@ import { ProcPool } from './ipc/proc_pool.js';
 import type { JobAcceptArguments, JobProcess, RunningJobInfo } from './job.js';
 import { JobRequest } from './job.js';
 import { log } from './log.js';
-import { Future } from './utils.js';
+import { Future, rejectOnAbort } from './utils.js';
 import { version } from './version.js';
 
 const MAX_RECONNECT_ATTEMPTS = 10;
@@ -428,7 +429,7 @@ export class AgentServer {
   }
 
   /** @throws {@link WorkerError} if worker did not drain in time */
-  async drain(timeout?: number) {
+  async drain(timeout?: number): Promise<Throws<void, WorkerError>> {
     if (this.#draining) {
       return;
     }
@@ -450,7 +451,7 @@ export class AgentServer {
 
     const joinJobs = async () => {
       return Promise.all(
-        this.#procPool.processes.map((proc) => {
+        this.#procPool.processes.map((proc): Promise<Throws<void, Error>> => {
           if (!proc.runningJob) {
             proc.close();
           }
@@ -459,17 +460,16 @@ export class AgentServer {
       );
     };
 
-    let timer: NodeJS.Timeout | undefined;
+    const promises = [joinJobs()];
+
     if (timeout) {
-      timer = setTimeout(() => {
-        throw new WorkerError('timed out draining');
-      }, timeout);
+      promises.push(
+        rejectOnAbort(AbortSignal.timeout(timeout)).catch(() => {
+          throw new WorkerError('timed out draining');
+        }),
+      );
     }
-    await joinJobs().then(() => {
-      if (timeout) {
-        clearTimeout(timer);
-      }
-    });
+    await Promise.race(promises);
   }
 
   async simulateJob(roomName: string, participantIdentity?: string) {
@@ -629,6 +629,25 @@ export class AgentServer {
     const loadMonitor = setInterval(() => {
       if (closingWS) clearInterval(loadMonitor);
 
+      if (this.#draining) {
+        if (currentStatus !== WorkerStatus.WS_FULL) {
+          currentStatus = WorkerStatus.WS_FULL;
+          this.event.emit(
+            'worker_msg',
+            new WorkerMessage({
+              message: {
+                case: 'updateWorker',
+                value: {
+                  load: 1,
+                  status: WorkerStatus.WS_FULL,
+                },
+              },
+            }),
+          );
+        }
+        return;
+      }
+
       const oldStatus = currentStatus;
       this.#opts
         .loadFunc(this)
@@ -708,6 +727,7 @@ export class AgentServer {
       );
 
       this.#pending[req.id] = new PendingAssignment();
+
       const timer = setTimeout(() => {
         this.#logger.child({ req }).warn(`assignment for job ${req.id} timed out`);
         return;
@@ -718,13 +738,17 @@ export class AgentServer {
       });
 
       if (asgn) {
-        await this.#procPool.launchJob({
-          acceptArguments: args,
-          job: msg.job!,
-          url: asgn.url || this.#opts.wsURL,
-          token: asgn.token,
-          workerId: this.id,
-        });
+        try {
+          await this.#procPool.launchJob({
+            acceptArguments: args,
+            job: msg.job!,
+            url: asgn.url || this.#opts.wsURL,
+            token: asgn.token,
+            workerId: this.id,
+          });
+        } catch (e) {
+          this.#logger.child({ requestId: req.id }).error(e, 'error launching job');
+        }
       } else {
         this.#logger.child({ requestId: req.id }).warn('pending assignment not found');
       }
@@ -735,6 +759,14 @@ export class AgentServer {
       .child({ jobId: msg.job?.id, resuming: msg.resuming, agentName: this.#opts.agentName })
       .info('received job request');
 
+    if (this.#draining) {
+      this.#logger
+        .child({ jobId: msg.job?.id, resuming: msg.resuming, agentName: this.#opts.agentName })
+        .info('Worker is draining and no longer available, rejecting job');
+      await req.reject();
+      return;
+    }
+
     const jobRequestTask = async () => {
       try {
         await this.#opts.requestFunc(req);
@@ -770,7 +802,7 @@ export class AgentServer {
       // safe to ignore
       return;
     }
-    await proc.close();
+    await proc.close().catch((e) => this.#logger.error(e, 'Error terminating job'));
   }
 
   async close() {