perf: parallelize credential migration owner analysis

src/core/migrate-credentials.ts

@@ -13,11 +13,14 @@ export interface MigrationOptions {
   recoverOwnerTokens?: boolean;
   identify?: (token: string) => Promise<IdentityResult>;
   batchSize?: number;
+  concurrency?: number;
   report?: (event: { collection: string; id: string; issue: string }) => void;
 }
 
 /** Run with all application writers stopped. Reruns never replace an existing credential. */
 export async function migrateCredentials(db: mongo.Db, cipher: Cipher, options: MigrationOptions = {}) {
+  const concurrency = options.concurrency ?? 10;
+  if (!Number.isInteger(concurrency) || concurrency < 1 || concurrency > 32) throw new Error("Concurrency must be 1..32");
   const credentials = db.collection("credentials");
   const users = db.collection("users");
   const batchSize = options.batchSize || 100;
@@ -28,7 +31,9 @@ export async function migrateCredentials(db: mongo.Db, cipher: Cipher, options:
     options.report?.({ collection, id: String(id), issue: reason });
   };
   if (options.apply) await credentials.createIndex({ ownerId: 1, provider: 1 }, { unique: true });
-  for await (const user of users.find({}, { projection: { accessTokens: 1, accessTokenDates: 1, status: 1, externalIDs: 1 } }).batchSize(batchSize)) {
+  const ownerIds = new Set<string>();
+  const processOwner = async (user: mongo.Document) => {
+    if (counts.halted) return;
     counts.owners++;
     const ownerId = user._id;
     const existing = await credentials.findOne({ ownerId, provider: "github" });
@@ -36,7 +41,7 @@ export async function migrateCredentials(db: mongo.Db, cipher: Cipher, options:
     let valid = true;
     if (existing) {
       try { selected = cipher.decrypt(existing.encryptedToken as EncryptedToken, String(ownerId), "github"); }
-      catch { issue("credentials", existing._id, "decryption_failed"); continue; }
+      catch { issue("credentials", existing._id, "decryption_failed"); return; }
     }
     const ownerToken = user.accessTokens?.github;
     const authoritative = !!(selected || (typeof ownerToken === "string" && ownerToken));
@@ -56,22 +61,23 @@ export async function migrateCredentials(db: mongo.Db, cipher: Cipher, options:
     inspect(ownerToken, "users", ownerId);
     for (const name of resources) {
       for await (const row of db.collection(name).find({ owner: ownerId, ...legacyQuery }, {
-        projection: { source: 1, accessToken: 1 },
+        projection: { "source.accessToken": 1, accessToken: 1 },
       }).batchSize(batchSize)) {
         inspect(row.source?.accessToken, name, row._id);
         inspect(row.accessToken, name, row._id);
       }
     }
-    if (!valid) continue;
+    if (!valid) return;
     if (recovering && candidates.size) {
       const result = await recover(user.externalIDs?.github, candidates);
       if ("issue" in result) {
         issue("users", ownerId, result.issue);
-        if (result.halt) { counts.halted = true; return counts; }
-        continue;
+        if (result.halt) { counts.halted = true; return; }
+        return;
       }
       selected = result.token;
     }
+    if (counts.halted) return;
     // Removed accounts must never regain credentials during backfill.
     if (user.status === "removed") {
       if (options.apply && options.removeLegacy) await credentials.deleteMany({ ownerId });
@@ -84,7 +90,7 @@ export async function migrateCredentials(db: mongo.Db, cipher: Cipher, options:
         } }, { upsert: true });
         const stored = await credentials.findOne({ ownerId, provider: "github" });
         if (!stored || cipher.decrypt(stored.encryptedToken as EncryptedToken, String(ownerId), "github") !== selected) {
-          issue("users", ownerId, "credential_changed_retry"); continue;
+          issue("users", ownerId, "credential_changed_retry"); return;
         }
       }
       counts.created++;
@@ -100,11 +106,34 @@ export async function migrateCredentials(db: mongo.Db, cipher: Cipher, options:
         counts.removed += result.modifiedCount;
       }
     }
+  };
+  // Keep one cursor reader and at most concurrency owner jobs in flight.
+  const pending = new Set<Promise<void>>();
+  let failure: unknown;
+  try {
+    for await (const user of users.find({}, { projection: {
+      accessTokens: 1, accessTokenDates: 1, status: 1, externalIDs: 1,
+    } }).batchSize(batchSize)) {
+      if (counts.halted) break;
+      ownerIds.add(String(user._id));
+      const job = processOwner(user).catch(error => {
+        failure = error;
+        counts.halted = true;
+      });
+      pending.add(job);
+      void job.then(() => pending.delete(job));
+      if (pending.size >= concurrency) await Promise.race(pending);
+    }
+  } finally {
+    // Drain before returning or letting the CLI disconnect, including cursor failures.
+    await Promise.all(pending);
   }
+  if (failure) throw failure;
+  if (counts.halted) return counts;
   // Credentials attached to missing owners cannot be assigned safely.
   for (const name of resources) {
     for await (const row of db.collection(name).find(legacyQuery, { projection: { owner: 1 } }).batchSize(batchSize)) {
-      if (!row.owner || !(await users.findOne({ _id: row.owner }, { projection: { _id: 1 } }))) {
+      if (!row.owner || !ownerIds.has(String(row.owner))) {
         issue(name, row._id, "missing_owner");
       }
     }

src/core/recover-owner-credential.ts

@@ -5,23 +5,35 @@ type Result = { token: string } | { issue: string; halt?: boolean };
 
 /** A resource token is usable only when it authenticates the already recorded owner. */
 export function createOwnerCredentialRecovery(identify = identifyGitHubToken) {
-  const cache = new Map<string, IdentityResult>();
-  let nextRequest = 0;
+  const cache = new Map<string, Promise<IdentityResult>>();
+  let gate = Promise.resolve();
+  let halted: IdentityResult | undefined;
+  const lookup = (token: string) => {
+    const hash = createHash("sha256").update(token).digest("hex");
+    const cached = cache.get(hash);
+    if (cached) return cached;
+    const turn = gate;
+    gate = turn.then(() => new Promise<void>(resolve => setTimeout(resolve, 250)));
+    const request = (async (): Promise<IdentityResult> => {
+      await turn;
+      if (halted) return halted;
+      let identity: IdentityResult;
+      try { identity = await identify(token); }
+      catch { identity = { issue: "github_request_failed", halt: true }; }
+      if ("issue" in identity && identity.halt) halted = identity;
+      return identity;
+    })();
+    // Cache pending requests too, so owners sharing a token share one lookup.
+    if (cache.size >= 10000) cache.delete(cache.keys().next().value!);
+    cache.set(hash, request);
+    return request;
+  };
   return async (githubId: unknown, candidates: Iterable<string>): Promise<Result> => {
     const expected = String(githubId ?? "");
     if (!/^[1-9][0-9]*$/.test(expected)) return { issue: "missing_or_invalid_owner_github_id" };
     let selected: string | undefined;
     for (const token of new Set(candidates)) {
-      const hash = createHash("sha256").update(token).digest("hex");
-      let identity = cache.get(hash);
-      if (!identity) {
-        await new Promise(resolve => setTimeout(resolve, Math.max(0, nextRequest - Date.now())));
-        nextRequest = Date.now() + 250;
-        try { identity = await identify(token); }
-        catch { return { issue: "github_request_failed", halt: true }; }
-        if (cache.size >= 10000) cache.delete(cache.keys().next().value!);
-        cache.set(hash, identity);
-      }
+      const identity = await lookup(token);
       if ("issue" in identity) {
         if (identity.issue === "invalid_or_revoked_token") continue;
         return identity;

src/scripts/migrate-credentials.ts

@@ -6,7 +6,12 @@ import { migrateCredentials, verifyCredentials, enforceCredentialStorage } from
 
 async function main() {
   const args = new Set(process.argv.slice(2));
+  const concurrencyArgs = [...args].filter(arg => arg.startsWith("--concurrency="));
+  if (concurrencyArgs.length > 1) throw new Error("Specify concurrency once");
+  const concurrency = concurrencyArgs.length ? Number(concurrencyArgs[0].slice("--concurrency=".length)) : 10;
+  if (!Number.isInteger(concurrency) || concurrency < 1 || concurrency > 32) throw new Error("Concurrency must be 1..32");
   for (const arg of args) {
+    if (concurrencyArgs.includes(arg)) continue;
     if (!["--apply", "--remove-legacy", "--prefer-owner-token", "--recover-owner-tokens", "--maintenance", "--verify", "--enforce"].includes(arg)) {
       throw new Error("Unknown migration option");
     }
@@ -26,6 +31,7 @@ async function main() {
     if (result.legacy) process.exitCode = 1;
   } else {
     const result = await migrateCredentials(db, cipher, {
+      concurrency,
       apply: args.has("--apply"), removeLegacy: args.has("--remove-legacy"),
       preferOwnerToken: args.has("--prefer-owner-token"),
       recoverOwnerTokens: args.has("--recover-owner-tokens"),