feat: close traffic timeout (#104)

Feat: AONE#54728515

fixtures/baseline/aworker_stream/stream_with_delay.js

@@ -0,0 +1,22 @@
+'use strict';
+
+addEventListener('fetch', event => {
+  event.respondWith(
+    Promise.resolve().then(async () => {
+      const time = await event.request.text();
+      const body = new ReadableStream({
+        start(controller) {
+          controller.enqueue('foobar');
+          setTimeout(() => {
+            controller.enqueue('end');
+            controller.close();
+          }, parseInt(time, 10));
+        },
+      });
+
+      return new Response(body, {
+        status: 200,
+      });
+    })
+  );
+});

src/config/default.ts

@@ -26,6 +26,7 @@ export default {
     daprAdaptorModulePath: undefined,
     daprAdaptorModuleOptions: undefined,
     beaconHostModulePath: undefined,
+    closeTrafficTimeout: 30_000,
   },
 
   dirs: {

src/config/index.ts

@@ -141,6 +141,10 @@ export interface DatePlaneConfig {
    * Beacon Host 模块路径
    */
   beaconHostModulePath?: string;
+  /**
+   * 关闭流量超时时间，单位毫秒，默认 30s
+   */
+  closeTrafficTimeout: number;
 }
 
 export interface DirConfig {

src/data_plane/__test__/e2e/close_traffic.test.ts

@@ -0,0 +1,122 @@
+import * as common from '#self/test/common';
+import { baselineDir } from '#self/test/common';
+import assert from 'assert';
+import { bufferFromStream, sleep } from '#self/lib/util';
+import { DefaultEnvironment } from '#self/test/env/environment';
+
+describe(common.testName(__filename), function () {
+  // Debug version of Node.js may take longer time to bootstrap.
+  this.timeout(60_000);
+
+  const env = new DefaultEnvironment({
+    config: common.extendDefaultConfig({
+      controlPlane: {
+        useEmaScaling: true,
+        dumpWorkerTrafficStats: true,
+      },
+      dataPlane: {
+        closeTrafficTimeout: 2000,
+      },
+    }),
+  });
+
+  it('should close worker', async () => {
+    await env.agent.setFunctionProfile([
+      {
+        name: 'aworker_stream',
+        runtime: 'aworker',
+        url: `file://${baselineDir}/aworker_stream`,
+        sourceFile: 'stream_with_delay.js',
+        signature: 'md5:234234',
+        worker: {
+          maxActivateRequests: 1,
+          disposable: true,
+        },
+        resourceLimit: {
+          memory: 200 * 1024 * 1024,
+        },
+      },
+    ]);
+
+    await request('aworker_stream', env);
+
+    await sleep(1000);
+
+    const broker = env.data.dataFlowController.brokers.get(
+      'aworker_stream$$noinspect'
+    );
+
+    assert.strictEqual(broker?.workerCount, 0);
+  });
+
+  it('should close worker with timeout', async () => {
+    await env.agent.setFunctionProfile([
+      {
+        name: 'aworker_stream_delay',
+        runtime: 'aworker',
+        url: `file://${baselineDir}/aworker_stream`,
+        sourceFile: 'stream_with_delay.js',
+        signature: 'md5:234234',
+        worker: {
+          maxActivateRequests: 1,
+          concurrencyShrinkThreshold: 0.6,
+        },
+        resourceLimit: {
+          memory: 200 * 1024 * 1024,
+        },
+      },
+    ]);
+
+    request('aworker_stream_delay', env, 30000).catch(error => {
+      assert(error);
+      assert(error.message.includes('Peer connection closed'));
+    });
+
+    await sleep(1000);
+
+    let broker = env.data.dataFlowController.brokers.get(
+      'aworker_stream_delay$$noinspect'
+    );
+
+    const worker = Array.from(
+      broker?.['_workerMap'].values() || []
+    ).pop() as any;
+
+    env.data.dataFlowController.closeTraffic([
+      {
+        functionName: 'aworker_stream_delay',
+        inspector: false,
+        workers: [
+          {
+            credential: worker!.worker!.credential,
+          },
+        ],
+      },
+    ]);
+
+    await sleep(3000);
+
+    broker = env.data.dataFlowController.brokers.get(
+      'aworker_stream_delay$$noinspect'
+    );
+
+    assert.strictEqual(broker?.workerCount, 0);
+  });
+});
+
+async function request(
+  functionName: string,
+  env: DefaultEnvironment,
+  timeout = 0
+) {
+  const data = Buffer.from(timeout + '');
+
+  const response = await env.agent.invoke(functionName, data, {
+    method: 'POST',
+    timeout: 30000,
+  });
+
+  assert.strictEqual(response.status, 200);
+
+  return await bufferFromStream(response);
+}

src/data_plane/data_flow_controller.ts

@@ -241,7 +241,7 @@ export class DataFlowController extends BaseOf(EventEmitter) {
     const tag = worker
       ? (worker as Worker).name || `{${credential}}`
       : `{${credential}}`;
-    logger.info('worker(%s) disconnected %s', tag);
+    logger.info('worker(%s) disconnected', tag);
 
     broker.removeWorker(credential);
     this.credentialBrokerMap.delete(credential);
@@ -314,7 +314,12 @@ export class DataFlowController extends BaseOf(EventEmitter) {
         code: RpcStatus.NOT_FOUND,
       });
     }
-    const broker = new WorkerBroker(this, profile, options);
+    const broker = new WorkerBroker(
+      this,
+      profile,
+      options,
+      this.config.dataPlane.closeTrafficTimeout
+    );
     this.brokers.set(key, broker);
     return broker;
   }

src/data_plane/worker_broker.ts

@@ -128,20 +128,38 @@ export class Worker extends EventEmitter {
   /**
    * Close this worker's traffic.
    */
-  async closeTraffic() {
+  async closeTraffic(timeout: number) {
     this.trafficOff = true;
+    this.logger.info(
+      'traffic closed, waiting for activeRequestCount down to 0...'
+    );
 
     if (this.activeRequestCount <= 0) {
       return Promise.resolve(true);
     }
 
     const { promise, resolve } = utils.createDeferred<boolean>();
+
+    let timer: NodeJS.Timeout | undefined = undefined;
+
     const downToZero: (...args: any[]) => void = () => {
+      clearTimeout(timer);
+      this.logger.debug('activeRequestCount down to 0.');
       resolve(true);
     };
 
     this.once('downToZero', downToZero);
 
+    timer = setTimeout(() => {
+      this.logger.warn(
+        `waiting for activeRequestCount down to 0 timeout, activeRequestCount is ${this.activeRequestCount}, will force kill.`
+      );
+      resolve(true);
+      this.removeListener('downToZero', downToZero);
+    }, timeout);
+
+    timer.unref();
+
     return promise;
   }
 
@@ -277,7 +295,8 @@ export class WorkerBroker extends Base implements DispatcherDelegate {
   constructor(
     public dataFlowController: DataFlowController,
     private _profile: RawWithDefaultsFunctionProfile,
-    public options: BrokerOptions = {}
+    public options: BrokerOptions = {},
+    private _closeTrafficTimeout = 10_000
   ) {
     super();
 
@@ -354,7 +373,7 @@ export class WorkerBroker extends Base implements DispatcherDelegate {
   async closeTraffic(worker: Worker) {
     try {
       this._dispatcher.unregisterWorker(worker);
-      await worker.closeTraffic();
+      await worker.closeTraffic(this._closeTrafficTimeout);
 
       // 同步 RequestDrained
       this.host.broadcastContainerStatusReport({