add timeout option for datafetcher

rainlanguage · May 6, 2024 · 79ac39b · 79ac39b
1 parent 2092e08
commit 79ac39b
Show file tree

Hide file tree

Showing 4 changed files with 85 additions and 20 deletions.
diff --git a/packages/sushi/src/router/data-fetcher.ts b/packages/sushi/src/router/data-fetcher.ts
@@ -50,6 +50,7 @@ import { UniswapV3Provider } from './liquidity-providers/UniswapV3.js'
 import { VVSStandardProvider } from './liquidity-providers/VVSStandard.js'
 import { WagmiProvider } from './liquidity-providers/Wagmi.js'
 import type { PoolCode } from './pool-codes/index.js'
+import { promiseTimeout } from './timeout.js'
 
 // options for data fetching, such as pinning block number and memoize
 export type DataFetcherOptions = {
@@ -61,6 +62,8 @@ export type DataFetcherOptions = {
   blockNumber?: bigint
   /** Determines if memoizer should be used or not */
   memoize?: boolean
+  /** Determines a timeout (in ms) for fetching pools for a token pair */
+  fetchPoolsTimeout?: number,
 }
 
 // TODO: Should be a mode on the config for DataFetcher
@@ -233,12 +236,22 @@ export class DataFetcher {
       )
       if (provider) {
         try {
-          await provider.fetchPoolsForToken(
-            currency0.wrapped,
-            currency1.wrapped,
-            excludePools,
-            options,
-          )
+          options?.fetchPoolsTimeout
+            ? await promiseTimeout(
+              provider.fetchPoolsForToken(
+                currency0.wrapped,
+                currency1.wrapped,
+                excludePools,
+                options,
+              ),
+              options.fetchPoolsTimeout
+            )
+            : await provider.fetchPoolsForToken(
+              currency0.wrapped,
+              currency1.wrapped,
+              excludePools,
+              options,
+            )
         } catch {
           /**/
         }
@@ -249,11 +262,24 @@ export class DataFetcher {
         currency0.wrapped.sortsBefore(currency1.wrapped)
           ? [currency0.wrapped, currency1.wrapped]
           : [currency1.wrapped, currency0.wrapped]
-      await Promise.allSettled(
-        this.providers.map((p) =>
-          p.fetchPoolsForToken(token0, token1, excludePools, options),
-        ),
-      )
+        try {
+          options?.fetchPoolsTimeout
+            ? await promiseTimeout(
+              Promise.allSettled(
+                this.providers.map((p) =>
+                  p.fetchPoolsForToken(token0, token1, excludePools, options),
+                ),
+              ),
+              options.fetchPoolsTimeout
+            )
+            : await Promise.allSettled(
+              this.providers.map((p) =>
+                p.fetchPoolsForToken(token0, token1, excludePools, options),
+              ),
+            )
+        } catch {
+          /**/
+        }
     }
   }
 

diff --git a/packages/sushi/src/router/timeout.test.ts b/packages/sushi/src/router/timeout.test.ts
@@ -0,0 +1,36 @@
+import { describe, expect, it } from 'vitest'
+import { promiseTimeout } from './timeout.js'
+
+const sleep = async(ms: number, msg = "") => {
+  let _timeoutReference;
+  return new Promise(
+      resolve => _timeoutReference = setTimeout(() => resolve(msg), ms),
+  ).finally(
+      () => clearTimeout(_timeoutReference)
+  );
+};
+
+describe('Promise Timeout', async () => {
+  it('should timeout', async () => {
+    const mainPromise = sleep(1000);
+    await promiseTimeout(mainPromise, 500)
+      .then(() => {
+        throw "expected to reject, but resolved"
+      })
+      .catch(() => {
+        /**/
+      })
+
+  })
+  it('should NOT timeout', async () => {
+    const mainPromise = sleep(1000);
+    await promiseTimeout(mainPromise, 1500)
+      .then(() => {
+        /**/
+      })
+      .catch(() => {
+        throw "expected to resolve, but rejected"
+      })
+
+  })
+})
diff --git a/packages/sushi/src/router/timeout.ts b/packages/sushi/src/router/timeout.ts
@@ -0,0 +1,11 @@
+export async function promiseTimeout (promise: Promise<any>, time: number, exception?: any) {
+  let timer: any;
+  return Promise.race([
+      promise,
+      new Promise(
+          (_res, _rej) => timer = setTimeout(_rej, time, exception)
+      )
+  ]).finally(
+      () => clearTimeout(timer)
+  );
+};
diff --git a/protocols/route-processor/test/DataFetcher.test.ts b/protocols/route-processor/test/DataFetcher.test.ts
@@ -12,7 +12,7 @@ async function testDF(
 ) {
   if (!t0 || !t1) return
   const start = performance.now()
-  await dataFetcher.fetchPoolsForToken(t0, t1)
+  await dataFetcher.fetchPoolsForToken(t0, t1, undefined, { fetchPoolsTimeout: 10000 })
   const pools = dataFetcher.getCurrentPoolCodeMap(t0, t1)
   const time = Math.round(performance.now() - start)
   console.log(
@@ -48,14 +48,6 @@ async function runTest() {
           'WNATIVE',
           'USDC',
         )
-        await testDF(
-          chName,
-          dataFetcher,
-          SUSHI[chainId as keyof typeof SUSHI],
-          FRAX[chainId as keyof typeof FRAX],
-          'SUSHI',
-          'FRAX',
-        )
         await testDF(
           chName,
           dataFetcher,