Implement rate limiter for function invocations

Author: henrymercer

extensions/ql-vscode/src/helpers.ts

@@ -1,5 +1,5 @@
 import * as path from 'path';
-import { CancellationToken, ProgressOptions, window as Window, workspace } from 'vscode';
+import { CancellationToken, ExtensionContext, ProgressOptions, window as Window, workspace } from 'vscode';
 import { logger } from './logging';
 import { EvaluationInfo } from './queries';
 
@@ -134,3 +134,81 @@ export function getQueryName(info: EvaluationInfo) {
     return path.basename(info.query.program.queryPath);
   }
 }
+
+/**
+ * Provides a utility method to invoke a function only if a minimum time interval has elapsed since
+ * the last invocation of that function.
+ */
+export class InvocationRateLimiter<T> {
+  constructor(extensionContext: ExtensionContext, funcIdentifier: string, func: () => Promise<T>) {
+    this._extensionContext = extensionContext;
+    this._func = func;
+    this._funcIdentifier = funcIdentifier;
+  }
+
+  /**
+   * Invoke the function if `minSecondsSinceLastInvocation` seconds have elapsed since the last invocation.
+   */
+  public async invokeFunctionIfIntervalElapsed(minSecondsSinceLastInvocation: number): Promise<InvocationRateLimiterResult<T>> {
+    const updateCheckStartDate = new Date();
+    const lastInvocationDate = this.getLastInvocationDate();
+    if (minSecondsSinceLastInvocation && lastInvocationDate && lastInvocationDate <= updateCheckStartDate &&
+      lastInvocationDate.getTime() + minSecondsSinceLastInvocation * 1000 > updateCheckStartDate.getTime()) {
+      return createRateLimitedResult();
+    }
+    const result = await this._func();
+    await this.setLastInvocationDate(updateCheckStartDate);
+    return createInvokedResult(result);
+  }
+
+  private getLastInvocationDate(): Date | undefined {
+    const maybeDate: Date | undefined =
+      this._extensionContext.globalState.get(InvocationRateLimiter._invocationRateLimiterPrefix + this._funcIdentifier);
+    return maybeDate ? new Date(maybeDate) : undefined;
+  }
+
+  private async setLastInvocationDate(date: Date): Promise<void> {
+    return await this._extensionContext.globalState.update(InvocationRateLimiter._invocationRateLimiterPrefix + this._funcIdentifier, date);
+  }
+
+  private readonly _extensionContext: ExtensionContext;
+  private readonly _func: () => Promise<T>;
+  private readonly _funcIdentifier: string;
+
+  private static readonly _invocationRateLimiterPrefix = "invocationRateLimiter_lastInvocationDate_";
+}
+
+export enum InvocationRateLimiterResultKind {
+  Invoked,
+  RateLimited
+}
+
+/**
+ * The function was invoked and returned the value `result`.
+ */
+interface InvokedResult<T> {
+  kind: InvocationRateLimiterResultKind.Invoked,
+  result: T
+}
+
+/**
+ * The function was not invoked as the minimum interval since the last invocation had not elapsed.
+ */
+interface RateLimitedResult {
+  kind: InvocationRateLimiterResultKind.RateLimited
+}
+
+type InvocationRateLimiterResult<T> = InvokedResult<T> | RateLimitedResult;
+
+function createInvokedResult<T>(result: T): InvokedResult<T> {
+  return {
+    kind: InvocationRateLimiterResultKind.Invoked,
+    result
+  };
+}
+
+function createRateLimitedResult(): RateLimitedResult {
+  return {
+    kind: InvocationRateLimiterResultKind.RateLimited
+  };
+}

extensions/ql-vscode/src/vscode-tests/no-workspace/helpers.test.ts

@@ -0,0 +1,104 @@
+import { expect } from "chai";
+import "mocha";
+import { ExtensionContext, Memento } from "vscode";
+import { InvocationRateLimiter } from "../../helpers";
+
+describe("Invocation rate limiter", () => {
+  function createInvocationRateLimiter<T>(funcIdentifier: string, func: () => Promise<T>): InvocationRateLimiter<T> {
+    return new InvocationRateLimiter(new MockExtensionContext(), funcIdentifier, func);
+  }
+
+  it("initially invokes function", async () => {
+    let numTimesFuncCalled = 0;
+    const invocationRateLimiter = createInvocationRateLimiter("funcid", async () => {
+      numTimesFuncCalled++;
+    });
+    await invocationRateLimiter.invokeFunctionIfIntervalElapsed(100);
+    expect(numTimesFuncCalled).to.equal(1);
+  });
+
+  it("doesn't invoke function within time period", async () => {
+    let numTimesFuncCalled = 0;
+    const invocationRateLimiter = createInvocationRateLimiter("funcid", async () => {
+      numTimesFuncCalled++;
+    });
+    await invocationRateLimiter.invokeFunctionIfIntervalElapsed(100);
+    await invocationRateLimiter.invokeFunctionIfIntervalElapsed(100);
+    expect(numTimesFuncCalled).to.equal(1);
+  });
+
+  it("invoke function again after 0s time period has elapsed", async () => {
+    let numTimesFuncCalled = 0;
+    const invocationRateLimiter = createInvocationRateLimiter("funcid", async () => {
+      numTimesFuncCalled++;
+    });
+    await invocationRateLimiter.invokeFunctionIfIntervalElapsed(0);
+    await invocationRateLimiter.invokeFunctionIfIntervalElapsed(0);
+    expect(numTimesFuncCalled).to.equal(2);
+  });
+
+  it("invoke function again after 1s time period has elapsed", async () => {
+    let numTimesFuncCalled = 0;
+    const invocationRateLimiter = createInvocationRateLimiter("funcid", async () => {
+      numTimesFuncCalled++;
+    });
+    await invocationRateLimiter.invokeFunctionIfIntervalElapsed(1);
+    await new Promise((resolve, _reject) => setTimeout(() => resolve(), 1000));
+    await invocationRateLimiter.invokeFunctionIfIntervalElapsed(1);
+    expect(numTimesFuncCalled).to.equal(2);
+  });
+
+  it("invokes functions with different rate limiters", async () => {
+    let numTimesFuncACalled = 0;
+    const invocationRateLimiterA = createInvocationRateLimiter("funcid", async () => {
+      numTimesFuncACalled++;
+    });
+    let numTimesFuncBCalled = 0;
+    const invocationRateLimiterB = createInvocationRateLimiter("funcid", async () => {
+      numTimesFuncBCalled++;
+    });
+    await invocationRateLimiterA.invokeFunctionIfIntervalElapsed(100);
+    await invocationRateLimiterB.invokeFunctionIfIntervalElapsed(100);
+    expect(numTimesFuncACalled).to.equal(1);
+    expect(numTimesFuncBCalled).to.equal(1);
+  });
+});
+
+class MockExtensionContext implements ExtensionContext {
+  subscriptions: { dispose(): unknown; }[] = [];
+  workspaceState: Memento = new MockMemento();
+  globalState: Memento = new MockMemento();
+  extensionPath: string = "";
+  asAbsolutePath(_relativePath: string): string {
+    throw new Error("Method not implemented.");
+  }
+  storagePath: string = "";
+  globalStoragePath: string = "";
+  logPath: string = "";
+}
+
+class MockMemento implements Memento {
+  map = new Map<any, any>();
+
+  /**
+   * Return a value.
+   *
+   * @param key A string.
+   * @param defaultValue A value that should be returned when there is no
+   * value (`undefined`) with the given key.
+   * @return The stored value or the defaultValue.
+   */
+  get<T>(key: string, defaultValue?: T): T {
+    return this.map.has(key) ? this.map.get(key) : defaultValue;
+  }
+
+  /**
+   * Store a value. The value must be JSON-stringifyable.
+   *
+   * @param key A string.
+   * @param value A value. MUST not contain cyclic references.
+   */
+  async update(key: string, value: any): Promise<void> {
+    this.map.set(key, value);
+  }
+}