feat: one-shot mode timeout

Author: doomsower

src/core/app/fetcher.ts

@@ -1,6 +1,5 @@
 import moment from "moment";
 import type { Address } from "viem";
-
 import type { PoolExternalAPYResult, PoolPointsResult } from "../../pools";
 import { getPoolExternalAPY, getPoolPoints } from "../../pools";
 import { getPoolExtraAPY } from "../../pools/extraAPY/apy";
@@ -52,8 +51,10 @@ import type { TokenExtraRewardsResult } from "../../tokens/tokenExtraRewards";
 import { getTokenExtraRewards } from "../../tokens/tokenExtraRewards";
 import type { NetworkType } from "../chains";
 import { getChainId, getNetworkType, supportedChains } from "../chains";
+import { ONE_SHOT_TIMEOUT } from "../config";
 import type { IOutputWriter, OneShotOutput } from "../output";
 import { captureException } from "../sentry";
+import { withTimeout } from "../utils";
 import { POOL_APY_TASK_INTERVAL, TOKEN_APY_TASK_INTERVAL } from "./constants";
 
 export type ApyDetails = Apy & { lastUpdated: string };
@@ -319,7 +320,7 @@ export class Fetcher {
     setInterval(quarterTask, POOL_APY_TASK_INTERVAL);
   }
 
-  public async oneShot(outputWriter: IOutputWriter): Promise<void> {
+  async #oneShot(outputWriter: IOutputWriter): Promise<void> {
     console.log("[SYSTEM]: Starting one-shot mode");
 
     await this.runNetworkRewards();
@@ -384,6 +385,17 @@ export class Fetcher {
       `[SYSTEM]: One-shot mode completed. Chains: ${output.metadata.totalChains}, Successful: ${output.metadata.successfulChains}, Failed: ${output.metadata.failedChains}`,
     );
   }
+
+  public async oneShot(outputWriter: IOutputWriter): Promise<void> {
+    if (ONE_SHOT_TIMEOUT) {
+      await withTimeout(
+        () => this.#oneShot(outputWriter),
+        ONE_SHOT_TIMEOUT * 1000,
+      );
+    } else {
+      await this.#oneShot(outputWriter);
+    }
+  }
 }
 
 interface LogRewardsProps {

src/core/config.ts

@@ -9,6 +9,16 @@ export const OUTPUT_JSON = process.env.OUTPUT_JSON || "output.json";
 export const SENTRY_DSN = process.env.SENTRY_DSN;
 
 export const S3_BUCKET = process.env.S3_BUCKET;
+/**
+ * Cache TTL in seconds
+ */
 export const CACHE_TTL = process.env.CACHE_TTL
   ? parseInt(process.env.CACHE_TTL, 10)
   : undefined;
+
+/**
+ * One-shot mode timeout in seconds
+ */
+export const ONE_SHOT_TIMEOUT = process.env.ONE_SHOT_TIMEOUT
+  ? parseInt(process.env.ONE_SHOT_TIMEOUT, 10)
+  : 0;

src/core/utils/index.ts

@@ -1,3 +1,5 @@
+import { setTimeout } from "node:timers/promises";
+
 export const json_stringify = (
   o: any,
   space?: number,
@@ -28,3 +30,15 @@ export const json_stringify = (
 };
 
 export type PartialRecord<K extends keyof any, T> = { [P in K]?: T };
+
+export async function timeout(ms: number): Promise<never> {
+  await setTimeout(ms);
+  throw new Error("The operation was timed out");
+}
+
+export async function withTimeout<T>(
+  run: () => Promise<T>,
+  ms: number,
+): Promise<T> {
+  return Promise.race([run(), timeout(ms)]);
+}