Add GroupRunner to limit concurrent RPCs (#457)

* Add GroupRunner to limit concurrent RPCs

* remove unnecessary async

* lint

* lint

* comment typo

* Remove type parameter in comment

Author: dskloetc

src/util/group-runner.ts

@@ -0,0 +1,69 @@
+type Callback<T> = () => Promise<T>
+type Resolve<T> = (arg: T | Promise<T>) => void
+
+// Runs tasks in groups of a fixed size.
+// The next group won't be started until the previous group is finished.
+//
+// Example usage:
+//
+// const fetchBalances(addresses: string[]): Promise<number[]> {
+//    // addresses can contain thousands of addresses
+//    const groupRunner = new GroupRunner(10)
+//    const getBalance = groupRunner.wrapFunction(fetchBalance)
+//    const balancePromises: Promise<number>[] = []
+//    for (const address of addresses) {
+//      // There will be at most 10 concurrent calls to fetchBalance.
+//      // fetchBalance might do an RPC and we don't want to get rate limited.
+//      balancePromises.push(getBalance(address)))
+//    }
+//    return Promise.all(balancePromises)
+// }
+//
+//
+// Implementation note:
+// Once the size has been reached, we wait for all previous tasks to finish
+// before running the new task.
+// Alternatively, we could run more tasks as soon as some (rather than all)
+// tasks have finished, to make progress sooner, but the former is what's
+// currently used in multiple places in the external-adapters-js repo so we
+// chose that behavior.
+export class GroupRunner {
+  private currentGroup: Promise<unknown>[] = []
+  private previousStartRunning: Promise<void> = Promise.resolve()
+
+  constructor(private groupSize: number) {}
+
+  // Calls the given callback eventually but makes sure any previous group of
+  // groupSize size has settled before calling and subsequent callbacks.
+  run<T>(callback: Callback<T>): Promise<T> {
+    return new Promise((resolve) => {
+      // This creates an implicit queue which guarantees that there are no
+      // concurrent calls into startRunning. This is necessary to avoid having
+      // currentGroup being cleared concurrently.
+      this.previousStartRunning = this.previousStartRunning.then(() => {
+        return this.startRunning(callback, resolve)
+      })
+    })
+  }
+
+  // Waits for a previous group to finish, if necessary, and then runs the
+  // given callback. When this method resolves, the callback has been called
+  // but not necessarily resolved.
+  async startRunning<T>(callback: Callback<T>, resolve: Resolve<T>) {
+    if (this.currentGroup.length >= this.groupSize) {
+      await Promise.allSettled(this.currentGroup)
+      this.currentGroup = []
+    }
+    const promise = callback()
+    this.currentGroup.push(promise)
+    resolve(promise)
+  }
+
+  wrapFunction<Args extends unknown[], Return>(
+    func: (...args: Args) => Promise<Return>,
+  ): (...args: Args) => Promise<Return> {
+    return (...args: Args) => {
+      return this.run(() => func(...args))
+    }
+  }
+}

src/util/index.ts

@@ -115,7 +115,7 @@ export const timeoutPromise = <T>(
   ]).finally(() => clearTimeout(timer))
 }
 
-type DeferredResolve<T> = (value: T) => void
+type DeferredResolve<T> = (value: T | PromiseLike<T>) => void
 type DeferredReject = (reason?: unknown) => void
 
 /**

test/util/group-runner.test.ts

@@ -0,0 +1,242 @@
+import { deferredPromise, sleep } from '../../src/util'
+import { GroupRunner } from '../../src/util/group-runner'
+import test from 'ava'
+
+type Task<T> = {
+  promise: Promise<T>
+  resolve: (arg: T) => void
+  reject: (arg: unknown) => void
+  getCallCount: () => number
+  callback: () => Promise<T>
+}
+
+const createTask = <T>(): Task<T> => {
+  const [promise, resolve, reject] = deferredPromise<T>()
+  let callCount = 0
+  return {
+    promise,
+    resolve,
+    reject,
+    getCallCount: () => callCount,
+    callback: () => {
+      callCount++
+      return promise
+    },
+  }
+}
+
+test('should run tasks', async (t) => {
+  const runner = new GroupRunner(2)
+  const result = 'result'
+  const callback = () => Promise.resolve(result)
+  t.is(await runner.run(callback), result)
+})
+
+test('should wait after group size is reached', async (t) => {
+  const runner = new GroupRunner(2)
+
+  const tasks: Task<void>[] = []
+  for (let i = 0; i < 3; i++) {
+    tasks.push(createTask())
+  }
+
+  runner.run(tasks[0].callback)
+  runner.run(tasks[1].callback)
+  runner.run(tasks[2].callback)
+  await sleep(0)
+  t.is(tasks[0].getCallCount(), 1)
+  t.is(tasks[1].getCallCount(), 1)
+  t.is(tasks[2].getCallCount(), 0)
+})
+
+test('should continue after group has finished', async (t) => {
+  const runner = new GroupRunner(2)
+
+  const tasks: Task<void>[] = []
+  for (let i = 0; i < 3; i++) {
+    tasks.push(createTask())
+  }
+
+  runner.run(tasks[0].callback)
+  runner.run(tasks[1].callback)
+  runner.run(tasks[2].callback)
+
+  await sleep(0)
+  t.is(tasks[2].getCallCount(), 0)
+
+  tasks[0].resolve(undefined)
+
+  await sleep(0)
+  t.is(tasks[2].getCallCount(), 0)
+
+  tasks[1].resolve(undefined)
+
+  await sleep(0)
+  t.is(tasks[2].getCallCount(), 1)
+})
+
+test('should not clear current group concurrently', async (t) => {
+  const runner = new GroupRunner(2)
+
+  const tasks: Task<void>[] = []
+  for (let i = 0; i < 4; i++) {
+    tasks.push(createTask())
+    runner.run(tasks[i].callback)
+  }
+  // Tasks 3 and 4 will both be waiting for tasks 1 and 2 to finish.
+  // When they do, they should not both clear the current group resulting in a
+  // group with only task 4.
+  tasks[0].resolve(undefined)
+  tasks[1].resolve(undefined)
+
+  // If we did end up with a group of 1 after concurrent clearing, then task 5
+  // will run immediately, which it shouldn't because task 3 and 4 are not
+  // finished yet.
+  const task5: Task<void> = createTask()
+  runner.run(task5.callback)
+
+  await sleep(0)
+  t.is(task5.getCallCount(), 0)
+})
+
+test('multiple groups', async (t) => {
+  const runner = new GroupRunner(3)
+
+  const tasks: Task<void>[] = []
+  for (let i = 0; i < 5; i++) {
+    tasks.push(createTask())
+    runner.run(tasks[i].callback)
+  }
+
+  await sleep(0)
+  t.is(tasks[0].getCallCount(), 1)
+  t.is(tasks[1].getCallCount(), 1)
+  t.is(tasks[2].getCallCount(), 1)
+  t.is(tasks[3].getCallCount(), 0)
+  t.is(tasks[4].getCallCount(), 0)
+
+  tasks[0].resolve(undefined)
+  tasks[1].resolve(undefined)
+  tasks[2].resolve(undefined)
+
+  await sleep(0)
+  t.is(tasks[3].getCallCount(), 1)
+  t.is(tasks[4].getCallCount(), 1)
+
+  // 5 more tasks for a total of 10:
+  for (let i = 5; i < 10; i++) {
+    tasks.push(createTask())
+    runner.run(tasks[i].callback)
+  }
+
+  await sleep(0)
+  t.is(tasks[5].getCallCount(), 1)
+  t.is(tasks[6].getCallCount(), 0)
+  t.is(tasks[7].getCallCount(), 0)
+  t.is(tasks[8].getCallCount(), 0)
+  t.is(tasks[9].getCallCount(), 0)
+
+  tasks[3].resolve(undefined)
+  tasks[4].resolve(undefined)
+  tasks[5].resolve(undefined)
+
+  await sleep(0)
+  t.is(tasks[6].getCallCount(), 1)
+  t.is(tasks[7].getCallCount(), 1)
+  t.is(tasks[8].getCallCount(), 1)
+  t.is(tasks[9].getCallCount(), 0)
+
+  tasks[6].resolve(undefined)
+  tasks[7].resolve(undefined)
+  tasks[8].resolve(undefined)
+
+  await sleep(0)
+  t.is(tasks[9].getCallCount(), 1)
+})
+
+test('multiple return values', async (t) => {
+  const runner = new GroupRunner(3)
+
+  const tasks: Task<number>[] = []
+  const promises: Promise<number>[] = []
+  for (let i = 0; i < 10; i++) {
+    tasks.push(createTask())
+    promises.push(runner.run(tasks[i].callback))
+  }
+
+  await sleep(0)
+  for (let i = 0; i < 10; i++) {
+    tasks[i].resolve(i)
+  }
+
+  await sleep(0)
+  for (let i = 0; i < 10; i++) {
+    t.is(await promises[i], i)
+  }
+})
+
+test('rejecting promises', async (t) => {
+  const runner = new GroupRunner(2)
+
+  const tasks: Task<void>[] = []
+  for (let i = 0; i < 3; i++) {
+    tasks.push(createTask())
+  }
+
+  runner.run(tasks[0].callback).catch(() => {
+    /* Ignore */
+  })
+  runner.run(tasks[1].callback).catch(() => {
+    /* Ignore */
+  })
+  runner.run(tasks[2].callback)
+
+  await sleep(0)
+  t.is(tasks[2].getCallCount(), 0)
+
+  tasks[0].reject(undefined)
+
+  await sleep(0)
+  t.is(tasks[2].getCallCount(), 0)
+
+  await sleep(0)
+  tasks[1].reject(undefined)
+
+  await sleep(0)
+  t.is(tasks[2].getCallCount(), 1)
+})
+
+test('wrap function', async (t) => {
+  const runner = new GroupRunner(2)
+
+  const tasks: Task<number>[] = []
+  for (let i = 0; i < 3; i++) {
+    tasks.push(createTask())
+  }
+
+  const f = runner.wrapFunction((i: number) => tasks[i].callback())
+
+  const promise0 = f(0)
+  const promise1 = f(1)
+  const promise2 = f(2)
+
+  await sleep(0)
+  t.is(tasks[2].getCallCount(), 0)
+
+  tasks[0].resolve(0)
+
+  await sleep(0)
+  t.is(tasks[2].getCallCount(), 0)
+
+  await sleep(0)
+  tasks[1].resolve(1)
+
+  await sleep(0)
+  t.is(tasks[2].getCallCount(), 1)
+
+  tasks[2].resolve(2)
+
+  t.is(await promise0, 0)
+  t.is(await promise1, 1)
+  t.is(await promise2, 2)
+})