Repository: firebase/firebase-tools
Pull request #1378 (2019-06-07 18:49:05)

Update Throttler to support task master timeout
### Summary
Added the ability for Throttler to cancel tasks on timeout. This is useful for polling API requests, since we may want the polling job stops after a fixed amount of time.

### Change details
- Added `timeout?: number` parameter to `run()` and `add()` methods. If `timeout` is not specified, `Throttler` will retry without time limit
- If the task finally succeeds (after any number of retries), `run()` will be resolved with the resolved value of the task (same as before the update)
- If the task fails due to:
  - **timed out:** `run()` will reject with a `TimeoutError`. The task running thread will still run in the background but will not retry after that
  - **retries exhausted:** `run()` will reject with a `RetriesExhaustedError`, with which caller can inspect error from the last retry using `error.orginal`.
- Notes:
  - `TimeoutError` and `RetriesExhaustedError` are inherited from `TaskError` which indicates task permanent failure
  - `run()` no longer rejects with the error directly thrown from the task. 
  - If the task finishes early, the running timeout will be stopped to avoid it causing the program to keep running. (This is not true for the opposite, see below)
### Limitation
- If the task takes very long time to finishes compare to the specified timeout. `run()` and `wait()` will reject with `TimeoutError` as soon as timed out. This will not block the main program from running, but in the background the task execution thread is still running. Since the current implementation of `Promise` doesn't allow us to cancel a promise, if the task takes long time complete, it will prevent the program ends, even after the main thread exits.

### Implementation notes:
- Each task is always executed at least once (it may exceed timeout if the given timeout is short).
- If the task succeeds on the first try or other retries, it will return a resolved promise with resolved value
- If a retry attempt fails, the method will check if it's timed out or retries exhausted to stop the next attempt.
- The `Throttler.run()` should resolve to the task resolved value or rejected with the error on the last retry

- **Notes:** Unit tests is going to be updated.

coveralls: 
[![Coverage Status](https://coveralls.io/builds/23866621/badge)](https://coveralls.io/builds/23866621)

Coverage increased (+0.1%) to 62.506% when pulling **58c406de2eff19706ac989cb0593d169c80c222f on tc-throttler** into **3264213624d2208e275f9fd2626d51425c502f31 on master**.

TrCaM: > Also, please add documentation to `Throttler#close()` that indicates that it's used to signal that no more tasks will be added to the throttler, but any tasks already added the throttler will continue to run.

Added documentation
TrCaM: Requested review from @fredzqm since he is the author of the original `Throttler` code

Patch:
diff --git a/src/throttler/errors/retries-exhausted-error.ts b/src/throttler/errors/retries-exhausted-error.ts
new file mode 100644
index 000000000..3889d1e96
--- /dev/null
+++ b/src/throttler/errors/retries-exhausted-error.ts
@@ -0,0 +1,9 @@
+import TaskError from "./task-error";
+
+export default class RetriesExhaustedError extends TaskError {
+  constructor(taskName: string, totalRetries: number, lastTrialError: Error) {
+    super(taskName, `retries exhausted after ${totalRetries + 1} attempts`, {
+      original: lastTrialError,
+    });
+  }
+}
diff --git a/src/throttler/errors/task-error.ts b/src/throttler/errors/task-error.ts
new file mode 100644
index 000000000..1c7a8e7e4
--- /dev/null
+++ b/src/throttler/errors/task-error.ts
@@ -0,0 +1,7 @@
+import * as FirebaseError from "../../error";
+
+export default abstract class TaskError extends FirebaseError {
+  constructor(taskName: string, message: string, options: object = {}) {
+    super(`Task ${taskName} failed: ${message}`, options);
+  }
+}
diff --git a/src/throttler/errors/timeout-error.ts b/src/throttler/errors/timeout-error.ts
new file mode 100644
index 000000000..37774cddc
--- /dev/null
+++ b/src/throttler/errors/timeout-error.ts
@@ -0,0 +1,7 @@
+import TaskError from "./task-error";
+
+export default class TimeoutError extends TaskError {
+  constructor(taskName: string, timeout: number) {
+    super(taskName, `timed out after ${timeout}ms.`);
+  }
+}
diff --git a/src/throttler/throttler.ts b/src/throttler/throttler.ts
index 94f9d1968..c42e408b4 100644
--- a/src/throttler/throttler.ts
+++ b/src/throttler/throttler.ts
@@ -1,4 +1,7 @@
 import * as logger from "../logger";
+import RetriesExhaustedError from "./errors/retries-exhausted-error";
+import TimeoutError from "./errors/timeout-error";
+import TaskError from "./errors/task-error";
 
 function backoff(retryNumber: number, delay: number): Promise<void> {
   return new Promise((resolve: () => void) => {
@@ -34,7 +37,10 @@ export interface ThrottlerStats {
 interface TaskData<T, R> {
   task: T;
   retryCount: number;
-  wait?: { resolve: (R: any) => void; reject: (err: Error) => void };
+  wait?: { resolve: (R: any) => void; reject: (err: TaskError) => void };
+  timeoutMillis?: number;
+  timeoutId?: NodeJS.Timeout;
+  isTimedOut: boolean;
 }
 
 /**
@@ -112,20 +118,24 @@ export abstract class Throttler<T, R> {
    * When the task is completed, resolve will be called with handler's result.
    * If this task fails after retries, reject will be called with the error.
    */
-  add(task: T): void {
-    this.addHelper(task);
+  add(task: T, timeoutMillis?: number): void {
+    this.addHelper(task, timeoutMillis);
   }
 
   /**
    * Add the task to the throttler and return a promise of handler's result.
    * If the task failed, both the promised returned by throttle and wait will reject.
    */
-  run(task: T): Promise<R> {
+  run(task: T, timeoutMillis?: number): Promise<R> {
     return new Promise((resolve, reject) => {
-      this.addHelper(task, { resolve, reject });
+      this.addHelper(task, timeoutMillis, { resolve, reject });
     });
   }
 
+  /**
+   * Signal that no more tasks will be added to the throttler, but any tasks already added to the
+   * throttler will continue to run
+   */
   close(): boolean {
     this.closed = true;
     return this.finishIfIdle();
@@ -145,53 +155,25 @@ export abstract class Throttler<T, R> {
     if (!taskData) {
       throw new Error(`taskData.get(${cursorIndex}) does not exist`);
     }
-    const task = taskData.task;
-    const tname = this.taskName(cursorIndex);
-    const t0 = Date.now();
+    const promises = [this.executeTask(cursorIndex)];
+    if (taskData.timeoutMillis) {
+      promises.push(this.initializeTimeout(cursorIndex));
+    }
 
+    let result;
     try {
-      const result = await this.handler(task);
-      const dt = Date.now() - t0;
-      if (dt < this.min) {
-        this.min = dt;
-      }
-      if (dt > this.max) {
-        this.max = dt;
-      }
-      this.avg = (this.avg * this.complete + dt) / (this.complete + 1);
-
-      this.success++;
-      this.complete++;
-      this.active--;
-      if (taskData.wait) {
-        taskData.wait.resolve(result);
-      }
-      this.taskDataMap.delete(cursorIndex);
-      this.process();
+      result = await Promise.race(promises);
     } catch (err) {
-      if (this.retries > 0) {
-        if (taskData.retryCount < this.retries) {
-          taskData.retryCount++;
-          this.retried++;
-          await backoff(taskData.retryCount, this.backoff);
-          logger.debug(`[${this.name}] Retrying task`, tname);
-          return this.handle(cursorIndex);
-        }
-      }
-
       this.errored++;
       this.complete++;
       this.active--;
-      if (taskData.retryCount > 0) {
-        logger.debug(`[${this.name}] Retries exhausted for task ${tname}:`, err);
-      } else {
-        logger.debug(`[${this.name}] Error on task ${tname}:`, err);
-      }
-      if (taskData.wait) {
-        taskData.wait.reject(err);
-      }
-      this.finish(err);
+      this.onTaskFailed(err, cursorIndex);
+      return;
     }
+    this.success++;
+    this.complete++;
+    this.active--;
+    this.onTaskFulfilled(result, cursorIndex);
   }
 
   stats(): ThrottlerStats {
@@ -219,6 +201,7 @@ export abstract class Throttler<T, R> {
 
   private addHelper(
     task: T,
+    timeoutMillis?: number,
     wait?: { resolve: (result: R) => void; reject: (err: Error) => void }
   ): void {
     if (this.closed) {
@@ -230,7 +213,9 @@ export abstract class Throttler<T, R> {
     this.taskDataMap.set(this.total, {
       task,
       wait,
+      timeoutMillis,
       retryCount: 0,
+      isTimedOut: false,
     });
     this.total++;
     this.process();
@@ -238,14 +223,14 @@ export abstract class Throttler<T, R> {
 
   private finishIfIdle(): boolean {
     if (this.closed && !this.hasWaitingTask() && this.active === 0) {
-      this.finish(null);
+      this.finish();
       return true;
     }
 
     return false;
   }
 
-  private finish(err: Error | null): void {
+  private finish(err?: TaskError): void {
     this.waits.forEach((p) => {
       if (err) {
         return p.reject(err);
@@ -254,4 +239,76 @@ export abstract class Throttler<T, R> {
       return p.resolve();
     });
   }
+
+  private initializeTimeout(cursorIndex: number): Promise<void> {
+    const taskData = this.taskDataMap.get(cursorIndex)!;
+    const timeoutMillis = taskData.timeoutMillis!;
+    const timeoutPromise = new Promise<void>((_, reject) => {
+      taskData.timeoutId = setTimeout(() => {
+        taskData.isTimedOut = true;
+        reject(new TimeoutError(this.taskName(cursorIndex), timeoutMillis));
+      }, timeoutMillis);
+    });
+
+    return timeoutPromise;
+  }
+
+  private async executeTask(cursorIndex: number): Promise<any> {
+    const taskData = this.taskDataMap.get(cursorIndex)!;
+    const t0 = Date.now();
+    let result;
+    try {
+      result = await this.handler(taskData.task);
+    } catch (err) {
+      if (taskData.retryCount === this.retries) {
+        throw new RetriesExhaustedError(this.taskName(cursorIndex), this.retries, err);
+      }
+      await backoff(taskData.retryCount + 1, this.backoff);
+      if (taskData.isTimedOut) {
+        throw new TimeoutError(this.taskName(cursorIndex), taskData.timeoutMillis!);
+      }
+      this.retried++;
+      taskData.retryCount++;
+      logger.debug(`[${this.name}] Retrying task`, this.taskName(cursorIndex));
+      return this.executeTask(cursorIndex);
+    }
+
+    if (taskData.isTimedOut) {
+      throw new TimeoutError(this.taskName(cursorIndex), taskData.timeoutMillis!);
+    }
+    const dt = Date.now() - t0;
+    this.min = Math.min(dt, this.min);
+    this.max = Math.max(dt, this.max);
+    this.avg = (this.avg * this.complete + dt) / (this.complete + 1);
+
+    return result;
+  }
+
+  private onTaskFulfilled(result: any, cursorIndex: number): void {
+    const taskData = this.taskDataMap.get(cursorIndex)!;
+    if (taskData.wait) {
+      taskData.wait.resolve(result);
+    }
+    this.cleanupTask(cursorIndex);
+    this.process();
+  }
+
+  private onTaskFailed(error: TaskError, cursorIndex: number): void {
+    const taskData = this.taskDataMap.get(cursorIndex)!;
+    logger.debug(error);
+
+    if (taskData.wait) {
+      taskData.wait.reject(error);
+    }
+    this.cleanupTask(cursorIndex);
+    this.finish(error);
+  }
+
+  private cleanupTask(cursorIndex: number): void {
+    const { timeoutId } = this.taskDataMap.get(cursorIndex)!;
+    if (timeoutId) {
+      clearTimeout(timeoutId);
+    }
+    this.taskDataMap.delete(cursorIndex);
+  }
 }


Test patch:
diff --git a/src/test/throttler/throttler.spec.ts b/src/test/throttler/throttler.spec.ts
index d2660c49d..111c5c36e 100644
--- a/src/test/throttler/throttler.spec.ts
+++ b/src/test/throttler/throttler.spec.ts
@@ -4,6 +4,9 @@ import { expect } from "chai";
 import Queue from "../../throttler/queue";
 import Stack from "../../throttler/stack";
 import { Throttler, ThrottlerOptions } from "../../throttler/throttler";
+import TaskError from "../../throttler/errors/task-error";
+import TimeoutError from "../../throttler/errors/timeout-error";
+import RetriesExhaustedError from "../../throttler/errors/retries-exhausted-error";
 
 const TEST_ERROR = new Error("foobar");
 
@@ -103,8 +106,10 @@ const throttlerTest = (throttlerConstructor: ThrottlerConstructor) => {
       .then(() => {
         throw new Error("handler should have rejected");
       })
-      .catch((err: Error) => {
-        expect(err).to.equal(TEST_ERROR);
+      .catch((err: TaskError) => {
+        expect(err).to.be.an.instanceof(RetriesExhaustedError);
+        expect(err.original).to.equal(TEST_ERROR);
+        expect(err.message).to.equal("Task index 0 failed: retries exhausted after 1 attempts");
       })
       .then(() => {
         expect(handler.callCount).to.equal(1);
@@ -132,8 +137,10 @@ const throttlerTest = (throttlerConstructor: ThrottlerConstructor) => {
       .then(() => {
         throw new Error("handler should have rejected");
       })
-      .catch((err: Error) => {
-        expect(err).to.equal(TEST_ERROR);
+      .catch((err: TaskError) => {
+        expect(err).to.be.an.instanceof(RetriesExhaustedError);
+        expect(err.original).to.equal(TEST_ERROR);
+        expect(err.message).to.equal("Task index 0 failed: retries exhausted after 4 attempts");
       })
       .then(() => {
         expect(handler.callCount).to.equal(4);
@@ -236,6 +243,162 @@ const throttlerTest = (throttlerConstructor: ThrottlerConstructor) => {
     expect(q.run(2)).to.eventually.to.equal("result: 2");
     expect(q.run(3)).to.eventually.to.equal("result: 3");
   });
+
+  it("should resolve if task finishes before timeout", async () => {
+    const handler = (task: number) => {
+      return Promise.resolve(`result: ${task}`);
+    };
+
+    const q = new Queue({
+      handler,
+    });
+
+    expect(await q.run(2, 20000000)).to.equal("result: 2");
+    expect(q.complete).to.equal(1);
+    expect(q.success).to.equal(1);
+    expect(q.errored).to.equal(0);
+    expect(q.retried).to.equal(0);
+    expect(q.total).to.equal(1);
+  });
+
+  it("should reject if timeout", async () => {
+    const handler = (task: number) =>
+      new Promise((resolve) => {
+        setTimeout(() => {
+          resolve(`result: ${task}`);
+        }, 150);
+      });
+
+    const q = new Queue({
+      handler,
+    });
+
+    let err;
+    try {
+      await q.run(2, 100);
+    } catch (e) {
+      err = e;
+    }
+    expect(err).to.be.instanceOf(TimeoutError);
+    expect(err.message).to.equal("Task index 0 failed: timed out after 100ms.");
+  });
+
+  it("should reject with RetriesExhaustedError if last trial is rejected before timeout", async () => {
+    const handler = sinon.stub().rejects(TEST_ERROR);
+
+    const q = new Queue({
+      handler,
+      retries: 2,
+      backoff: 10,
+    });
+
+    let err;
+    try {
+      await q.run(2, 200);
+    } catch (e) {
+      err = e;
+    }
+    expect(err).to.be.instanceOf(RetriesExhaustedError);
+    expect(err.original).to.equal(TEST_ERROR);
+    expect(err.message).to.equal("Task index 0 failed: retries exhausted after 3 attempts");
+    expect(handler.callCount).to.equal(3);
+    expect(q.complete).to.equal(1);
+    expect(q.success).to.equal(0);
+    expect(q.errored).to.equal(1);
+    expect(q.retried).to.equal(2);
+    expect(q.total).to.equal(1);
+  });
+
+  it("should reject with TimeoutError if timeout while retrying", async () => {
+    const handler = sinon.stub().rejects(TEST_ERROR);
+
+    const q = new Queue({
+      handler,
+      retries: 1000,
+      backoff: 5,
+    });
+
+    let err;
+    try {
+      await q.run(2, 100);
+    } catch (e) {
+      err = e;
+    }
+    expect(err).to.be.instanceOf(TimeoutError);
+    expect(handler.callCount).to.be.at.least(2);
+    expect(err.message).to.equal("Task index 0 failed: timed out after 100ms.");
+    expect(q.complete).to.equal(1);
+    expect(q.success).to.equal(0);
+    expect(q.errored).to.equal(1);
+    expect(q.retried).to.be.at.least(3);
+    expect(q.total).to.equal(1);
+  });
+
+  it("should reject with TimeoutError when waiting", async () => {
+    const handler = sinon
+      .stub()
+      .rejects(TEST_ERROR)
+      .onFirstCall()
+      .resolves(0);
+
+    const q = new Queue({
+      handler,
+      retries: 4,
+      backoff: 20,
+    });
+
+    q.add(2);
+    q.add(3, 10); // This fails first due to very short timeout
+    q.add(4, 500);
+    q.close();
+
+    let err;
+    try {
+      await q.wait();
+    } catch (e) {
+      err = e;
+    }
+    expect(err).to.be.instanceOf(TimeoutError);
+    expect(err.message).to.equal("Task index 1 failed: timed out after 10ms.");
+    expect(handler.callCount).to.equal(3);
+    expect(q.complete).to.equal(2);
+    expect(q.success).to.equal(1);
+    expect(q.errored).to.equal(1);
+    expect(q.total).to.equal(3);
+  });
+
+  it("should reject with RetriesExhaustedError when waiting", async () => {
+    const handler = sinon
+      .stub()
+      .rejects(TEST_ERROR)
+      .onFirstCall()
+      .resolves(0);
+
+    const q = new Queue({
+      handler,
+      retries: 1,
+      backoff: 10,
+    });
+
+    q.add(2);
+    q.add(3, 100); // This fails due to retries exhausted since we only retry once
+    q.close();
+
+    let err;
+    try {
+      await q.wait();
+    } catch (e) {
+      err = e;
+    }
+    expect(err).to.be.instanceOf(RetriesExhaustedError);
+    expect(err.message).to.equal("Task index 1 failed: retries exhausted after 2 attempts");
+    expect(handler.callCount).to.equal(3);
+    expect(q.complete).to.equal(2);
+    expect(q.success).to.equal(1);
+    expect(q.errored).to.equal(1);
+    expect(q.retried).to.equal(1);
+    expect(q.total).to.equal(2);
+  });
 };
 
 describe("Throttler", () => {