fix: retain concurrently published path indexes for cleanup

src/core/Repository.ts

@@ -38,6 +38,23 @@ import AnonymizedRepositoryModel from "./model/anonymizedRepositories/anonymized
 import { createLogger, serializeError } from "./logger";
 
 const logger = createLogger("repository");
+/** Capture and retire the stored key while clearing it, including concurrent index publication. */
+async function updateRetiringPathIndex(
+  filter: () => FilterQuery<IAnonymizedRepositoryDocument>,
+  update: UpdateQuery<IAnonymizedRepositoryDocument>
+): Promise<{ matchedCount: number }> {
+  for (let attempt = 0; attempt < 5; attempt++) {
+    const current = await AnonymizedRepositoryModel.findOne(filter()).select("pathIndexKey").lean().exec();
+    if (!current) return { matchedCount: 0 };
+    const result = await AnonymizedRepositoryModel.updateOne({ ...filter(),
+      pathIndexKey: current.pathIndexKey === undefined ? { $exists: false } : current.pathIndexKey,
+    }, { ...update, ...(current.pathIndexKey ? { $addToSet: {
+      ...update.$addToSet, retiredPathIndexKeys: current.pathIndexKey,
+    } } : {}) }).exec();
+    if (result.matchedCount) return result;
+  }
+  return { matchedCount: 0 };
+}
 const pathIndexes = new Map<string, Promise<void>>();
 async function buildPathIndex(repoId: string, key: string, terms: string[], generation?: string) {
   const cacheKey = `${repoId}:${key}`;
@@ -69,7 +86,7 @@ async function buildPathIndex(repoId: string, key: string, terms: string[], gene
 }
 import { IFile, IFileDocument } from "./model/files/files.types";
 import AnonymizedFile from "./AnonymizedFile";
-import { FilterQuery } from "mongoose";
+import { FilterQuery, UpdateQuery } from "mongoose";
 function anonymizeTreeRecursive(
   tree: IFile[],
   terms: string[],
@@ -296,20 +313,19 @@ export default class Repository {
           throw error;
         });
         if (isConnected) {
-          const activated = await AnonymizedRepositoryModel.updateOne({ _id: this.model._id,
+          const activated = await updateRetiringPathIndex(() => ({ _id: this.model._id,
             treeGeneration: previousGeneration || { $exists: false }, ...this.refreshFilter(),
             status: this.model.status, statusDate: this.model.statusDate,
             "githubAccess.revision": this.model.githubAccess?.revision || { $exists: false },
             stagedFileTrees: { $elemMatch: { generation, until: { $gt: new Date() } } },
-          }, { $set: { treeGeneration: generation, truncatedFolders, size: { storage: 0, file: 0 },
+          }), { $set: { treeGeneration: generation, truncatedFolders, size: { storage: 0, file: 0 },
             ...(streamed && !this.model.contentCacheVersion ? { contentCacheVersion: 1, legacyContentCleanupPending: true } : {}),
             ...(!fileCount ? { emptyTreeGeneration: generation } : {}),
           }, $unset: { pathIndexKey: "", pathIndexBuiltAt: "", sizeComputedAt: "",
             ...(fileCount ? { emptyTreeGeneration: "" } : {}),
           }, $addToSet: { retiredTreeGenerations: previousGeneration || "",
             ...(streamed ? { retiredContentPrefixes: previousPrefix } : {}),
-            ...(this.model.pathIndexKey ? { retiredPathIndexKeys: this.model.pathIndexKey } : {}),
-          }, $pull: { stagedFileTrees: { generation } } }).exec();
+          }, $pull: { stagedFileTrees: { generation } } });
           if (!activated.matchedCount) {
             await deleteRowsInBatches(FileModel.collection, { repoId: this.repoId, treeGeneration: generation });
             throw new AnonymousError("repository_changed", { httpStatus: 409 });
@@ -1042,14 +1058,13 @@ export default class Repository {
     const generation = defer ? randomUUID() : previousGeneration;
     this.model.emptyTreeGeneration = undefined;
     if (isConnected) {
-      const result = await AnonymizedRepositoryModel.updateOne({ _id: this.model._id,
+      const result = await updateRetiringPathIndex(() => ({ _id: this.model._id,
       treeGeneration: this.model.treeGeneration || { $exists: false },
-      }, { $set: { fileMetadataRevision: revision, ...(defer ? { treeGeneration: generation } : {}) },
+      }), { $set: { fileMetadataRevision: revision, ...(defer ? { treeGeneration: generation } : {}) },
         ...(defer ? { $addToSet: { retiredTreeGenerations: previousGeneration || "",
-          ...(this.model.pathIndexKey ? { retiredPathIndexKeys: this.model.pathIndexKey } : {}),
         } } : {}),
         $unset: { emptyTreeGeneration: "", pathIndexKey: "", pathIndexBuiltAt: "" },
-      }).exec();
+      });
       if (!result.matchedCount) throw new AnonymousError("repository_changed", { httpStatus: 409 });
     }
     this.model.treeGeneration = generation;

test/large-repository-tree.test.js

@@ -179,6 +179,46 @@ const uri = process.env.PERFORMANCE_MONGO_URI;
     expect(await File.countDocuments({ repoId: row.repoId, treeStaged: true })).to.equal(500);
   });
 
+  it("retires an index published between reading and clearing the persisted key", async () => {
+    const row = await repository({ status: "ready", treeGeneration: "old" }), resetter = new Repository(row);
+    await rows(2);
+    const searcher = new Repository(await Repo.findById(row._id)), updateOne = Repo.updateOne;
+    let attempts = 0, publishedKey;
+    Repo.updateOne = function (filter, update, ...args) {
+      const query = updateOne.call(this, filter, update, ...args);
+      if (update.$unset && Object.hasOwn(update.$unset, "pathIndexKey") && Object.hasOwn(filter, "pathIndexKey")) {
+        const exec = query.exec;
+        query.exec = async function (...execArgs) {
+          if (++attempts === 1) {
+            expect(await searcher.searchFiles("file")).to.have.length(2);
+            publishedKey = (await Repo.findById(row._id)).pathIndexKey;
+          }
+          return exec.apply(this, execArgs);
+        };
+      }
+      return query;
+    };
+    try { await resetter.resetSate("preparing"); } finally { Repo.updateOne = updateOne; }
+    expect(attempts).to.equal(2);
+    expect((await Repo.findById(row._id)).retiredPathIndexKeys).to.deep.equal([publishedKey]);
+    await resetter.cleanupRetiredFileTrees();
+    expect(await Path.countDocuments({ repoId: row.repoId })).to.equal(0);
+  });
+
+  it("retires a path index published after the activating worker loaded its model", async () => {
+    const row = await repository({ status: "ready", treeGeneration: "old" }), worker = new Repository(row);
+    await rows(2);
+    const searcher = new Repository(await Repo.findById(row._id));
+    Object.defineProperty(worker, "source", { get: () => ({ getFiles: async () => {
+      expect(await searcher.searchFiles("file")).to.have.length(2);
+      expect((await Repo.findById(row._id)).pathIndexKey).to.be.a("string");
+      return [{ name: "replacement.txt", path: "", size: 1, sha: "c".repeat(40) }];
+    } }) });
+    await worker.files({ force: true });
+    expect(await Path.countDocuments({ repoId: row.repoId })).to.equal(0);
+    expect((await File.find({ repoId: row.repoId })).map(file => file.name)).to.deep.equal(["replacement.txt"]);
+  });
+
   it("retains retirement markers when foreground cleanup reaches its batch budget", async () => {
     const row = await repository({ treeGeneration: "active", retiredTreeGenerations: ["old"] }), repo = new Repository(row);
     await rows(501);
@@ -218,8 +258,12 @@ const uri = process.env.PERFORMANCE_MONGO_URI;
 
   it("preserves rows that stop matching the deletion filter after selection", async () => {
     await rows(2);
+    let checks = 0;
     const complete = await deleteRowsInBatches(File.collection, { repoId: "large", treeGeneration: "old" }, {
-      current: async () => { await File.updateMany({ repoId: "large" }, { $set: { treeGeneration: "replacement" } }); return true; },
+      current: async () => {
+        if (++checks === 2) await File.updateMany({ repoId: "large" }, { $set: { treeGeneration: "replacement" } });
+        return true;
+      },
     });
     expect(complete).to.equal(true); expect(await File.countDocuments({ repoId: "large" })).to.equal(2);
   });

test/production-regressions.test.js

@@ -206,7 +206,8 @@ describe("production regressions", function () {
   it("retains truncation warnings from the source instance that fetched the tree", async function () {
     const FileModel = require("../src/core/model/files/files.model").default;
     stub(db, "isConnected", false);
-    for (const method of ["exists", "deleteMany", "find"]) stub(FileModel, method, () => ({ exec: async () => method === "find" ? [] : null }));
+    for (const method of ["exists", "deleteMany"]) stub(FileModel, method, () => ({ exec: async () => null }));
+    stub(FileModel, "find", () => ({ select() { return this; }, limit() { return this; }, lean() { return this; }, exec: async () => [] }));
     stub(FileModel, "insertMany", async () => []);
     const repo = new Repository(new RepoModel({ repoId: "repo", options: {} }));
     repo.computeSize = async () => {};