improve queue

src/core/GitHubUtils.ts

@@ -1,5 +1,6 @@
 import { Octokit } from "@octokit/rest";
 import { throttling } from "@octokit/plugin-throttling";
+import { createClient, RedisClientType } from "redis";
 
 import AnonymousError from "./AnonymousError";
 import Repository from "./Repository";
@@ -53,6 +54,155 @@ function rateLimitDetail(err: OctokitRequestErrorLike): string {
 
 const ThrottledOctokit = Octokit.plugin(throttling);
 
+/**
+ * Per-token gate that blocks all callers when a rate limit is active.
+ * When any request for a given token hits a rate limit, the gate records
+ * the reset time and makes every subsequent caller wait until then —
+ * preventing a stampede of doomed requests.
+ */
+const tokenGates = new Map<string, { resetAt: number }>();
+
+function setTokenGate(token: string, retryAfterSec: number) {
+  const key = token.slice(-8);
+  const resetAt = Date.now() + retryAfterSec * 1000;
+  const existing = tokenGates.get(key);
+  if (!existing || resetAt > existing.resetAt) {
+    tokenGates.set(key, { resetAt });
+    logger.warn("rate limit gate set", {
+      code: "rate_limit_gate",
+      retryAfterSec,
+      resetAt: new Date(resetAt).toISOString(),
+    });
+    setRedisGate(retryAfterSec).catch(() => {});
+  }
+}
+
+export class RateLimitDelayError extends Error {
+  resetAt: number;
+  constructor(resetAt: number) {
+    const delaySec = Math.ceil((resetAt - Date.now()) / 1000);
+    super(`github_rate_limit_delay:${delaySec}s`);
+    this.name = "RateLimitDelayError";
+    this.resetAt = resetAt;
+  }
+}
+
+/**
+ * Check if a rate limit gate is active for a token.
+ * Returns the reset timestamp, or 0 if no gate is active.
+ */
+export function getTokenGateResetAt(token: string): number {
+  const key = token.slice(-8);
+  const gate = tokenGates.get(key);
+  if (!gate) return 0;
+  if (gate.resetAt <= Date.now()) {
+    tokenGates.delete(key);
+    return 0;
+  }
+  return gate.resetAt;
+}
+
+async function waitForTokenGate(token: string): Promise<void> {
+  const key = token.slice(-8);
+  const localGate = tokenGates.get(key);
+  let waitMs = 0;
+  let resetAt = 0;
+
+  if (localGate && localGate.resetAt > Date.now()) {
+    resetAt = localGate.resetAt;
+    waitMs = resetAt - Date.now();
+  }
+
+  const redisResetAt = await getRedisGateResetAt();
+  if (redisResetAt > resetAt) {
+    resetAt = redisResetAt;
+    waitMs = resetAt - Date.now();
+  }
+
+  if (waitMs <= 0) {
+    if (localGate) tokenGates.delete(key);
+    return;
+  }
+
+  logger.info("waiting for rate limit gate", {
+    code: "rate_limit_gate_wait",
+    waitMs,
+    resetAt: new Date(resetAt).toISOString(),
+  });
+  await new Promise((resolve) => setTimeout(resolve, waitMs));
+  if (localGate) tokenGates.delete(key);
+}
+
+const REDIS_GATE_PREFIX = "gh_rate_gate:";
+
+let redisGateDisabled = false;
+let redisGateReady: Promise<RedisClientType | null> | null = null;
+
+function ensureRedisGateClient(): Promise<RedisClientType | null> {
+  if (redisGateDisabled) return Promise.resolve(null);
+  if (redisGateReady) return redisGateReady;
+  redisGateReady = (async () => {
+    try {
+      const c = createClient({
+        socket: {
+          host: config.REDIS_HOSTNAME,
+          port: config.REDIS_PORT,
+          reconnectStrategy: () => false as any,
+        },
+      }) as RedisClientType;
+      c.on("error", () => {
+        redisGateDisabled = true;
+        c.disconnect().catch(() => {});
+        redisGateReady = null;
+      });
+      await c.connect();
+      return c;
+    } catch {
+      redisGateDisabled = true;
+      redisGateReady = null;
+      return null;
+    }
+  })();
+  return redisGateReady;
+}
+
+async function setRedisGate(retryAfterSec: number): Promise<void> {
+  const c = await ensureRedisGateClient();
+  if (!c || !c.isOpen) return;
+  const resetAt = Date.now() + retryAfterSec * 1000;
+  const ttl = Math.ceil(retryAfterSec) + 10;
+  try {
+    await c.set(REDIS_GATE_PREFIX + "global", String(resetAt), { EX: ttl });
+    logger.info("redis rate limit gate written", {
+      code: "redis_gate_set",
+      resetAt: new Date(resetAt).toISOString(),
+      ttl,
+    });
+  } catch {
+    // non-critical
+  }
+}
+
+export async function setRedisGateFromWorker(resetAt: number): Promise<void> {
+  const retryAfterSec = Math.max(0, (resetAt - Date.now()) / 1000);
+  if (retryAfterSec <= 0) return;
+  await setRedisGate(retryAfterSec);
+}
+
+export async function getRedisGateResetAt(): Promise<number> {
+  const c = await ensureRedisGateClient();
+  if (!c || !c.isOpen) return 0;
+  try {
+    const val = await c.get(REDIS_GATE_PREFIX + "global");
+    if (!val) return 0;
+    const resetAt = parseInt(val, 10);
+    if (isNaN(resetAt) || resetAt <= Date.now()) return 0;
+    return resetAt;
+  } catch {
+    return 0;
+  }
+}
+
 export function octokit(token: string) {
   const oct = new ThrottledOctokit({
     auth: token,
@@ -69,8 +219,7 @@ export function octokit(token: string) {
           retryAfter,
           retryCount,
         });
-        // Retry once; if GitHub is still throttling after that, surface the
-        // error to the caller so the UI shows github_rate_limit_exceeded.
+        setTokenGate(token, retryAfter);
         return retryCount < 1;
       },
       onSecondaryRateLimit: (retryAfter, options, _o, retryCount) => {
@@ -82,6 +231,7 @@ export function octokit(token: string) {
           retryAfter,
           retryCount,
         });
+        setTokenGate(token, retryAfter);
         return retryCount < 1;
       },
     },
@@ -99,12 +249,20 @@ export function octokit(token: string) {
   return oct;
 }
 
+export { waitForTokenGate };
+
 export async function checkToken(token: string) {
   const oct = octokit(token);
   try {
     await oct.users.getAuthenticated();
     return true;
-  } catch {
+  } catch (err) {
+    if (
+      err instanceof AnonymousError &&
+      err.message === "github_rate_limit_exceeded"
+    ) {
+      throw err;
+    }
     return false;
   }
 }

src/core/Repository.ts

@@ -19,7 +19,7 @@ import {
   getRepositoryFromGitHub,
   GitHubRepository,
 } from "./source/GitHubRepository";
-import { getToken } from "./GitHubUtils";
+import { getToken, getRedisGateResetAt } from "./GitHubUtils";
 import config from "../config";
 import FileModel from "./model/files/files.model";
 import AnonymizedRepositoryModel from "./model/anonymizedRepositories/anonymizedRepositories.model";
@@ -234,14 +234,30 @@ export default class Repository {
         httpStatus: 410,
       });
     }
+    const redisGateReset = await getRedisGateResetAt();
+    if (redisGateReset > 0) {
+      throw new AnonymousError("rate_limited", {
+        httpStatus: 425,
+        object: { resetAt: redisGateReset },
+      });
+    }
+
     const fiveMinuteAgo = new Date();
     fiveMinuteAgo.setMinutes(fiveMinuteAgo.getMinutes() - 5);
 
     if (
       this.status == RepositoryStatus.PREPARING ||
+      this.status == RepositoryStatus.QUEUE ||
       (this.status == RepositoryStatus.DOWNLOAD &&
         this._model.statusDate > fiveMinuteAgo)
     ) {
+      const rlMatch = (this._model.statusMessage || "").match(/^rate_limited:(\d+)$/);
+      if (rlMatch) {
+        throw new AnonymousError("rate_limited", {
+          httpStatus: 425,
+          object: { resetAt: parseInt(rlMatch[1], 10) },
+        });
+      }
       throw new AnonymousError("repository_not_ready", {
         object: this,
         httpStatus: 425,

src/core/source/GitHubStream.ts

@@ -11,7 +11,7 @@ import { basename, dirname } from "path";
 import * as stream from "stream";
 import AnonymousError from "../AnonymousError";
 import { FILE_TYPE } from "../storage/Storage";
-import { octokit } from "../GitHubUtils";
+import { octokit, waitForTokenGate } from "../GitHubUtils";
 import FileModel from "../model/files/files.model";
 import { IFile } from "../model/files/files.types";
 import { createLogger, serializeError } from "../logger";
@@ -20,6 +20,27 @@ import config from "../../config";
 
 const logger = createLogger("gh-stream");
 
+const GH_API_CONCURRENCY = 6;
+
+async function pMap<T, R>(
+  items: T[],
+  fn: (item: T, index: number) => Promise<R>,
+  concurrency: number
+): Promise<R[]> {
+  const results: R[] = new Array(items.length);
+  let next = 0;
+  async function worker() {
+    while (next < items.length) {
+      const i = next++;
+      results[i] = await fn(items[i], i);
+    }
+  }
+  await Promise.all(
+    Array.from({ length: Math.min(concurrency, items.length) }, () => worker())
+  );
+  return results;
+}
+
 export function githubRawFileUrl(
   owner: string,
   repo: string,
@@ -354,11 +375,13 @@ export default class GitHubStream extends GitHubBase {
   }
 
   private async getGHTree(
+    oct: ReturnType<typeof octokit>,
+    token: string,
     sha: string,
     count = { request: 0, file: 0 },
     opt = { recursive: true, callback: () => {} }
   ) {
-    const oct = octokit(await this.data.getToken());
+    await waitForTokenGate(token);
     const ghRes = await oct.git.getTree({
       owner: this.data.organization,
       repo: this.data.repoName,
@@ -378,6 +401,8 @@ export default class GitHubStream extends GitHubBase {
     progress?: (status: string) => void,
     parentPath: string = ""
   ) {
+    const token = await this.data.getToken();
+    const oct = octokit(token);
     const count = {
       request: 0,
       file: 0,
@@ -385,7 +410,7 @@ export default class GitHubStream extends GitHubBase {
     const output: IFile[] = [];
     let data;
     try {
-      data = await this.getGHTree(sha, count, {
+      data = await this.getGHTree(oct, token, sha, count, {
         recursive: false,
         callback: () => {
           if (progress) {
@@ -423,29 +448,33 @@ export default class GitHubStream extends GitHubBase {
         cause: error as Error,
       });
     }
-    const promises: ReturnType<GitHubStream["getGHTree"]>[] = [];
-    const parentPaths: string[] = [];
+    const subtrees: { sha: string; parentPath: string }[] = [];
     for (const file of data.tree) {
       if (file.type == "tree" && file.path && file.sha) {
-        const elementPath = path.join(parentPath, file.path);
-        parentPaths.push(elementPath);
-        promises.push(
-          this.getGHTree(file.sha, count, {
-            recursive: true,
-            callback: () => {
-              if (progress) {
-                progress("List file: " + count.file);
-              }
-            },
-          })
-        );
+        subtrees.push({
+          sha: file.sha,
+          parentPath: path.join(parentPath, file.path),
+        });
       }
     }
-    (await Promise.all(promises)).forEach((data, i) => {
+    const results = await pMap(
+      subtrees,
+      async (entry) =>
+        this.getGHTree(oct, token, entry.sha, count, {
+          recursive: true,
+          callback: () => {
+            if (progress) {
+              progress("List file: " + count.file);
+            }
+          },
+        }),
+      GH_API_CONCURRENCY
+    );
+    results.forEach((data, i) => {
       if (data.truncated) {
-        this._truncatedFolders.push(parentPaths[i]);
+        this._truncatedFolders.push(subtrees[i].parentPath);
       }
-      output.push(...this.tree2Tree(data.tree, parentPaths[i]));
+      output.push(...this.tree2Tree(data.tree, subtrees[i].parentPath));
     });
     return output;
   }