feat(internal): Rate limiting in ymax planner when fetching pending transactions (#12634)

Author: usmanmani1122

services/ymax-planner/src/engine.ts

@@ -67,6 +67,7 @@ import type {
   HandlePendingTxOpts,
 } from './pending-tx-manager.ts';
 import { handlePendingTx } from './pending-tx-manager.ts';
+import rateLimitedSource from './rate-limited-source.ts';
 import type { BalanceQueryPowers } from './plan-deposit.ts';
 import {
   getNonDustBalances,
@@ -173,6 +174,7 @@ export type Powers = {
   console?: Pick<Console, 'debug' | 'info' | 'log' | 'warn' | 'error'>;
   evmCtx: Omit<EvmContext, 'signingSmartWalletKit' | 'fetch'>;
   rpc: CosmosRPCClient;
+  setTimeout: typeof globalThis.setTimeout;
   spectrumBlockchain: SpectrumBlockchainSdk;
   spectrumChainIds: Partial<Record<SupportedChain, string>>;
   evmTokenAddresses: Partial<Record<InstrumentId, EvmAddress>>;
@@ -801,34 +803,48 @@ export const startEngine = async (
   console.warn(`Found ${pendingTxKeys.length} pending transactions`);
 
   const initialPendingTxData: PendingTxRecord[] = [];
-  await makeWorkPool(pendingTxKeys, undefined, async (txId: TxId) => {
-    const path = `${pendingTxPathPrefix}.${txId}`;
-    await null;
-    let streamCellJson;
-    let data;
-    try {
-      const metaResponse = await readStorageMeta(query.vstorage, path, 'data', {
-        retries: 4,
-      });
-      streamCellJson = metaResponse.result.value;
-      const streamCell = parseStreamCell(streamCellJson, path);
-      const marshalledData = parseStreamCellValue(streamCell, -1, path);
-      data = marshaller.fromCapData(marshalledData);
-      if (
-        data?.status !== TxStatus.PENDING ||
-        !RESOLVER_SUPPORTED_TRANSACTIONS.includes(data.type)
-      )
-        return;
-      mustMatch(harden(data), PublishedTxShape, path);
-      initialPendingTxData.push({
-        blockHeight: BigInt(streamCell.blockHeight),
-        tx: { txId, ...data },
-      });
-    } catch (err) {
-      const errLabel = `🚨 Failed to read old pending tx ${path}`;
-      console.error(errLabel, data || streamCellJson, err);
-    }
-  }).done;
+  const capacity = 10;
+  const throttledPendingTxKeys = rateLimitedSource({
+    policy: { quota: capacity, windowMs: 1000 },
+    powers: { now: powers.now, setTimeout: powers.setTimeout },
+    source: pendingTxKeys as Array<string>,
+  });
+
+  await makeWorkPool(
+    throttledPendingTxKeys,
+    { capacity },
+    async (txId: TxId) => {
+      const path = `${pendingTxPathPrefix}.${txId}`;
+      await null;
+      let streamCellJson;
+      let data;
+      try {
+        const metaResponse = await readStorageMeta(
+          query.vstorage,
+          path,
+          'data',
+          { retries: 4 },
+        );
+        streamCellJson = metaResponse.result.value;
+        const streamCell = parseStreamCell(streamCellJson, path);
+        const marshalledData = parseStreamCellValue(streamCell, -1, path);
+        data = marshaller.fromCapData(marshalledData);
+        if (
+          data?.status !== TxStatus.PENDING ||
+          !RESOLVER_SUPPORTED_TRANSACTIONS.includes(data.type)
+        )
+          return;
+        mustMatch(harden(data), PublishedTxShape, path);
+        initialPendingTxData.push({
+          blockHeight: BigInt(streamCell.blockHeight),
+          tx: { txId, ...data },
+        });
+      } catch (err) {
+        const errLabel = `🚨 Failed to read old pending tx ${path}`;
+        console.error(errLabel, data || streamCellJson, err);
+      }
+    },
+  ).done;
 
   if (initialPendingTxData.length > 0) {
     // Process initial transactions in lookback mode upon planner startup
@@ -931,4 +947,5 @@ export const startEngine = async (
   }
   Fail`⚠️ rpc.subscribeAll finished`;
 };
+
 harden(startEngine);

services/ymax-planner/src/main.ts

@@ -295,6 +295,7 @@ export const main = async (
       ...evmCtx,
     },
     rpc,
+    setTimeout,
     spectrumChainIds,
     evmTokenAddresses,
     spectrumBlockchain,
@@ -327,4 +328,5 @@ export const main = async (
     });
   });
 };
+
 harden(main);

services/ymax-planner/src/rate-limited-source.ts

@@ -0,0 +1,76 @@
+/**
+ * Rate-limited async iterator wrapper.
+ *
+ * Releases items from `source` at a rate of at most `quota` items per any
+ * `windowMs` sliding window, with bursts of up to `quota` allowed at the head
+ * of an idle window. Naming follows the IETF `RateLimit` header draft's
+ * quota/window vocabulary (q, w) but uses `windowMs` for sub-second precision.
+ *
+ * Time and timer access are passed in as powers — no ambient clock or
+ * `setTimeout`.
+ *
+ * @see https://datatracker.ietf.org/doc/draft-ietf-httpapi-ratelimit-headers/
+ */
+
+type RateLimitPolicy = {
+  /** Maximum items released per window. */
+  quota: number;
+  /** Sliding window length in milliseconds. */
+  windowMs: number;
+};
+
+type RateLimitPowers = {
+  now: () => number;
+  setTimeout: typeof globalThis.setTimeout;
+};
+
+const rateLimitedSource = <T>({
+  policy,
+  powers,
+  source,
+}: {
+  policy: RateLimitPolicy;
+  powers: RateLimitPowers;
+  source: Iterable<T> | AsyncIterable<T>;
+}) => {
+  const { quota, windowMs } = policy;
+  const { now, setTimeout } = powers;
+
+  if (!(Number.isInteger(quota) && quota > 0)) {
+    throw RangeError('quota must be a positive integer');
+  }
+
+  if (!(Number.isFinite(windowMs) && windowMs > 0)) {
+    throw RangeError('windowMs must be a positive finite number');
+  }
+
+  const delay = (ms: number) =>
+    new Promise<void>(resolve => setTimeout(resolve, ms));
+
+  async function* generate(): AsyncGenerator<T, void, void> {
+    const releases: number[] = [];
+
+    for await (const item of source) {
+      const remainingCapacity = quota - releases.length;
+      const t = now();
+
+      if (remainingCapacity > 0) {
+        releases.push(t);
+      } else {
+        if (remainingCapacity !== 0) throw Error('internal: quota exceeded');
+
+        const earliest = releases.shift()!;
+        const effective = Math.max(t, earliest + windowMs);
+        releases.push(effective);
+        const wait = effective - t;
+        if (wait) await delay(wait);
+      }
+
+      yield item;
+    }
+  }
+
+  return generate();
+};
+
+export default rateLimitedSource;

services/ymax-planner/test/mocks.ts

@@ -105,6 +105,7 @@ export const createMockEnginePowers = (): EnginePowers => ({
   walletStore: {} as any,
   getWalletInvocationUpdate: async () => undefined,
   now: () => NaN,
+  setTimeout: globalThis.setTimeout,
   gasEstimator: {} as any,
   usdcTokensByChain: {},
   chainNameToChainIdMap: CaipChainIds.testnet,

services/ymax-planner/test/rate-limited-source.test.ts

@@ -0,0 +1,190 @@
+import test from 'ava';
+
+import rateLimitedSource from '../src/rate-limited-source.ts';
+
+const makeFakeClock = () => {
+  let clockNow = 0;
+
+  return {
+    advance: (ms: number) => {
+      clockNow += ms;
+    },
+    setTimeout: ((cb: () => void, ms?: number) => {
+      clockNow += ms ?? 0;
+      cb();
+      return 0;
+    }) as typeof globalThis.setTimeout,
+    now: () => clockNow,
+  };
+};
+
+test('rateLimitedSource: bursts up to quota, then waits one window', async t => {
+  const clock = makeFakeClock();
+  const items = [0, 1, 2, 3, 4, 5, 6, 7, 8];
+  const releases: Array<{ item: number; at: number }> = [];
+  for await (const item of rateLimitedSource({
+    source: items,
+    policy: { quota: 3, windowMs: 100 },
+    powers: clock,
+  })) {
+    releases.push({ item, at: clock.now() });
+  }
+  t.deepEqual(releases, [
+    { item: 0, at: 0 },
+    { item: 1, at: 0 },
+    { item: 2, at: 0 },
+    { item: 3, at: 100 },
+    { item: 4, at: 100 },
+    { item: 5, at: 100 },
+    { item: 6, at: 200 },
+    { item: 7, at: 200 },
+    { item: 8, at: 200 },
+  ]);
+});
+
+test('rateLimitedSource: an idle period grants the next item without delay', async t => {
+  const clock = makeFakeClock();
+  let setTimeoutCalls = 0;
+  const powers = {
+    now: clock.now,
+    setTimeout: ((cb: () => void, ms?: number) => {
+      setTimeoutCalls += 1;
+      return clock.setTimeout(cb, ms);
+    }) as typeof globalThis.setTimeout,
+  };
+  const iter = rateLimitedSource({
+    source: [0, 1, 2],
+    policy: { quota: 1, windowMs: 100 },
+    powers,
+  });
+
+  t.is((await iter.next()).value, 0);
+  t.is(clock.now(), 0);
+  t.is(setTimeoutCalls, 0);
+
+  clock.advance(150);
+  t.is((await iter.next()).value, 1);
+  t.is(clock.now(), 150);
+  t.is(setTimeoutCalls, 0);
+
+  t.is((await iter.next()).value, 2);
+  t.is(clock.now(), 250);
+  t.is(setTimeoutCalls, 1);
+});
+
+test('rateLimitedSource: sliding window spaces by oldest in-window release', async t => {
+  const clock = makeFakeClock();
+  const iter = rateLimitedSource({
+    source: [0, 1, 2, 3],
+    policy: { quota: 2, windowMs: 100 },
+    powers: clock,
+  });
+
+  await iter.next();
+  t.is(clock.now(), 0);
+
+  clock.advance(50);
+  await iter.next();
+  t.is(clock.now(), 50);
+
+  await iter.next();
+  t.is(clock.now(), 100);
+
+  await iter.next();
+  t.is(clock.now(), 150);
+});
+
+test('rateLimitedSource: empty source completes without invoking setTimeout', async t => {
+  const clock = makeFakeClock();
+  let setTimeoutCalls = 0;
+  const powers = {
+    now: clock.now,
+    setTimeout: ((cb: () => void, ms?: number) => {
+      setTimeoutCalls += 1;
+      return clock.setTimeout(cb, ms);
+    }) as typeof globalThis.setTimeout,
+  };
+  for await (const _ of rateLimitedSource({
+    source: [],
+    policy: { quota: 1, windowMs: 100 },
+    powers,
+  })) {
+    t.fail('should not emit');
+  }
+  t.is(setTimeoutCalls, 0);
+  t.is(clock.now(), 0);
+});
+
+test('rateLimitedSource: accepts async iterables', async t => {
+  const clock = makeFakeClock();
+  async function* gen() {
+    yield 'a';
+    yield 'b';
+    yield 'c';
+  }
+  const releases: Array<{ item: string; at: number }> = [];
+  for await (const item of rateLimitedSource({
+    source: gen(),
+    policy: { quota: 1, windowMs: 100 },
+    powers: clock,
+  })) {
+    releases.push({ item, at: clock.now() });
+  }
+  t.deepEqual(releases, [
+    { item: 'a', at: 0 },
+    { item: 'b', at: 100 },
+    { item: 'c', at: 200 },
+  ]);
+});
+
+test('rateLimitedSource: validates quota', t => {
+  const clock = makeFakeClock();
+  for (const bad of [0, -1, 1.5, NaN, Infinity]) {
+    t.throws(
+      () =>
+        rateLimitedSource({
+          source: [],
+          policy: { quota: bad, windowMs: 100 },
+          powers: clock,
+        }),
+      { instanceOf: RangeError, message: /quota/ },
+      `quota=${bad}`,
+    );
+  }
+});
+
+test('rateLimitedSource: validates windowMs', t => {
+  const clock = makeFakeClock();
+  for (const bad of [0, -1, NaN, Infinity]) {
+    t.throws(
+      () =>
+        rateLimitedSource({
+          source: [],
+          policy: { quota: 1, windowMs: bad },
+          powers: clock,
+        }),
+      { instanceOf: RangeError, message: /windowMs/ },
+      `windowMs=${bad}`,
+    );
+  }
+});
+
+test('rateLimitedSource: composes with worker pool style consumers', async t => {
+  const clock = makeFakeClock();
+  const items = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];
+  const starts: Array<{ item: number; at: number }> = [];
+  for await (const item of rateLimitedSource({
+    source: items,
+    policy: { quota: 4, windowMs: 60 },
+    powers: clock,
+  })) {
+    starts.push({ item, at: clock.now() });
+    clock.advance(80);
+  }
+  // Quota=4 in 60ms, but each "task" advances 80ms — so the window is always
+  // already empty by the time we pull again. Every item starts immediately.
+  t.deepEqual(
+    starts.map(s => s.at),
+    [0, 80, 160, 240, 320, 400, 480, 560, 640, 720, 800, 880],
+  );
+});