Merge pull request #804 from tdurieux/t3code/parallel-credential-migration

perf: parallelize credential migration owner analysis

docs/credential-encryption.md

@@ -409,8 +409,9 @@ docker compose run --rm --no-deps -T --entrypoint node anonymous_github \
 Then rerun the preview and follow the legacy-removal, verification and enforcement
 steps above. Do not regenerate encryption keys between runs.
 
-Recovery contacts GitHub even in preview mode. Requests are sequential, paced at
-least 250 ms apart, with a bounded cache keyed by token hashes. Network failures,
+Recovery contacts GitHub even in preview mode. Request starts are paced at
+least 250 ms apart across all workers, with a bounded cache keyed by token hashes
+that also shares in-flight requests. Network failures,
 rate limits and server errors halt the run (`halted: true`); earlier completed
 owners may already have been migrated in apply mode. Rerunning is safe.
 
@@ -426,3 +427,27 @@ These owners retain their legacy tokens, including with `--remove-legacy`.
 Have the owner sign in again to establish a fresh authoritative credential, or
 review their records manually. Do not bulk-archive owned repositories solely
 because credential recovery failed.
+
+
+### Parallel migration analysis
+
+Migration processes 10 owners concurrently by default. Set `--concurrency=1..32`
+to tune database load; `1` restores sequential owner processing. For example:
+
+```bash
+docker compose run --rm --no-deps -T --entrypoint node anonymous_github \
+  build/scripts/migrate-credentials.js \
+  --prefer-owner-token --recover-owner-tokens --concurrency=20
+```
+
+The same option works with `--apply --maintenance` and `--remove-legacy`.
+Each owner is handled by one job, with token validation before writes or cleanup.
+On failure, new jobs stop and existing jobs drain before disconnecting. Writes
+already in progress may finish; reruns remain safe. Reports can arrive out of
+owner order. All writers must remain stopped for apply runs.
+
+The orphan scan reuses the user IDs read during migration instead of querying
+MongoDB for each resource. Only IDs are retained for this check (memory grows
+with the number of users). Existing resource `owner` indexes and the credentials
+`(ownerId, provider)` index should be present for efficient lookups. Increasing
+concurrency does not bypass GitHub pacing or rate-limit handling.

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"),

test/recover-owner-credential.test.js

@@ -110,3 +110,70 @@ describe("migration owner recovery decisions", function () {
     expect(cipher.decrypt(f.credential().encryptedToken, "owner", "github")).to.equal("authoritative");
   });
 });
+
+describe("parallel credential migration", function () {
+  const { setTimeout } = require("timers/promises");
+  const cipher = createTokenCipher(JSON.stringify({ test: Buffer.alloc(32, 7).toString("base64") }), "test");
+  it("bounds owner jobs and avoids per-resource owner lookups", async () => {
+    const owners = Array.from({ length: 12 }, (_, i) => ({ _id: String(i) }));
+    let active = 0, peak = 0, completed = 0;
+    const cursor = values => ({ batchSize() { return this; }, async *[Symbol.asyncIterator]() { yield* values; } });
+    const events = [];
+    const db = { collection(name) { return {
+      find: query => cursor(name === "users" ? owners : name === "anonymizedrepositories" && !query.owner
+        ? [{ _id: "known", owner: "1" }, { _id: "orphan", owner: "missing" }] : []),
+      findOne: async () => {
+        expect(name).to.equal("credentials");
+        active++; peak = Math.max(peak, active);
+        await setTimeout(5);
+        active--; completed++;
+        return null;
+      },
+    }; } };
+    const result = await migrateCredentials(db, cipher, { concurrency: 3, report: e => events.push(e) });
+    expect(peak).to.equal(3);
+    expect(active).to.equal(0);
+    expect(completed).to.equal(12);
+    expect(result.owners).to.equal(12);
+    expect(events).to.deep.equal([{ collection: "anonymizedrepositories", id: "orphan", issue: "missing_owner" }]);
+  });
+  it("drains outstanding jobs before propagating a database failure", async () => {
+    let active = 0, calls = 0;
+    const cursor = values => ({ batchSize() { return this; }, async *[Symbol.asyncIterator]() { yield* values; } });
+    const db = { collection(name) { return {
+      find: () => cursor(name === "users" ? [{ _id: "1" }, { _id: "2" }, { _id: "3" }] : []),
+      findOne: async () => {
+        const call = ++calls;
+        active++;
+        await setTimeout(call === 1 ? 5 : 20);
+        active--;
+        if (call === 1) throw new Error("database failed");
+        return null;
+      },
+    }; } };
+    let error;
+    try { await migrateCredentials(db, cipher, { concurrency: 2 }); } catch (e) { error = e; }
+    expect(error.message).to.equal("database failed");
+    expect(active).to.equal(0);
+    expect(calls).to.equal(2);
+  });
+  it("shares in-flight identities and spaces requests across concurrent owners", async () => {
+    const starts = [];
+    const recover = createOwnerCredentialRecovery(async () => {
+      starts.push(Date.now());
+      await setTimeout(30);
+      return { githubId: "42" };
+    });
+    const results = await Promise.all([recover("42", ["a"]), recover("42", ["a"]), recover("42", ["b"])]);
+    expect(results).to.deep.equal([{ token: "a" }, { token: "a" }, { token: "b" }]);
+    expect(starts).to.have.length(2);
+    expect(starts[1] - starts[0]).to.be.at.least(240);
+  });
+  it("suppresses queued GitHub requests after a rate limit", async () => {
+    let calls = 0;
+    const recover = createOwnerCredentialRecovery(async () => { calls++; return { issue: "github_http_429", halt: true }; });
+    const results = await Promise.all([recover("42", ["a"]), recover("42", ["b"])]);
+    expect(calls).to.equal(1);
+    expect(results.every(r => r.halt)).to.equal(true);
+  });
+});