Merge pull request #940 from tdurieux/perf/covered-folder-counts

perf: count folder files from a covering index

src/core/folder-counts.ts

@@ -21,7 +21,7 @@ export async function folderCounts(repo: Repository, requestedPath: unknown): Pr
   }
   const prefix = originalPath ? originalPath + "/" : "";
   const escapedPrefix = prefix.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
-  const prefixLength = Array.from(prefix).length;
+  const childDepth = originalPath ? originalPath.split("/").length : 0;
   const rows = await FileModel.aggregate<{ _id: string; count: number }>([
     { $match: {
       repoId: repo.repoId,
@@ -29,12 +29,11 @@ export async function folderCounts(repo: Repository, requestedPath: unknown): Pr
       size: { $ne: null },
       ...(originalPath ? { $or: [{ path: originalPath }, { path: { $regex: "^" + escapedPrefix } }] } : {}),
     } },
-    { $project: { child: { $cond: [
+    // Group directly by the immediate child using only indexed fields.
+    { $group: { _id: originalPath ? { $cond: [
       { $eq: ["$path", { $literal: originalPath }] }, "",
-      { $arrayElemAt: [{ $split: [{ $substrCP: ["$path", prefixLength,
-        { $subtract: [{ $strLenCP: "$path" }, prefixLength] }] }, "/"] }, 0] },
-    ] } } },
-    { $group: { _id: "$child", count: { $sum: 1 } } },
+      { $arrayElemAt: [{ $split: ["$path", "/"] }, childDepth] },
+    ] } : { $arrayElemAt: [{ $split: ["$path", "/"] }, 0] }, count: { $sum: 1 } } },
     { $limit: FILE_LISTING_LIMIT + 1 },
   ]).option({ maxTimeMS: 2000, allowDiskUse: true }).exec();
   if (rows.length > FILE_LISTING_LIMIT) throw new AnonymousError("file_listing_too_large", { httpStatus: 413 });

src/core/model/files/files.schema.ts

@@ -25,6 +25,8 @@ FileSchema.index({ metadataPending: 1, repoId: 1, treeGeneration: 1 },
 FileSchema.index({ treeStaged: 1, repoId: 1, treeGeneration: 1 },
   { partialFilterExpression: { treeStaged: true } });
 FileSchema.index({ repoId: 1, treeGeneration: 1, path: 1, name: 1 });
+// Cover generation-scoped count aggregations without fetching every file.
+FileSchema.index({ repoId: 1, treeGeneration: 1, path: 1, size: 1 }, { name: "file_counts_covering" });
 FileSchema.index(
   { repoId: 1, size: 1, path: 1 },
   { name: "repoId_1_size_1_path_1" }

src/scripts/create-performance-indexes.ts

@@ -14,6 +14,7 @@ async function main() {
   // Additive only. Never use syncIndexes here, which can drop deployed indexes.
   await Repository.collection.createIndex({ "coauthors.githubId": 1 });
   await File.collection.createIndex({ repoId: 1, treeGeneration: 1, path: 1, name: 1 });
+  await File.collection.createIndex({ repoId: 1, treeGeneration: 1, path: 1, size: 1 }, { name: "file_counts_covering" });
   await File.collection.createIndex({ metadataPending: 1, repoId: 1, treeGeneration: 1 },
     { partialFilterExpression: { metadataPending: true } });
   await Repository.collection.createIndex({ retiredTreeGenerations: 1 }, { sparse: true });

test/folder-counts.test.js

@@ -46,6 +46,32 @@ describe("folder count request validation", () => {
     expect(await folderCounts(repo, "src")).to.deep.equal({ src: 12000, "src/deep": 12000 });
   });
 
+  it("covers root counts without fetching documents and counts every file in a directory", async () => {
+    await insert(Array.from({ length: 2000 }, (_, i) => ({ name: `file${i}`, path: `src/folder${i % 20}`, size: i % 2 })));
+    await insert([{ name: "excluded-directory", path: "" },
+      { name: "old", path: "old", size: 1, treeGeneration: "retired" }]);
+    let pipeline;
+    const aggregate = File.aggregate;
+    File.aggregate = function (stages, ...args) { pipeline = stages; return aggregate.call(this, stages, ...args); };
+    try { expect(await folderCounts(repo, "")).to.deep.equal({ "": 2000, src: 2000 }); }
+    finally { File.aggregate = aggregate; }
+    const explained = await File.aggregate(pipeline).option({ hint: "file_counts_covering" }).explain("executionStats");
+    const cursor = explained.stages?.find(stage => stage.$cursor)?.$cursor || explained;
+    expect(cursor.executionStats.totalDocsExamined).to.equal(0);
+    expect(cursor.executionStats.totalKeysExamined).to.be.at.least(2000);
+  });
+
+  it("preserves counts for legacy repositories without a generation", async () => {
+    repo.model.treeGeneration = undefined; await repo.model.save();
+    await File.collection.insertMany([
+      { repoId: "counts", name: "zero", path: "", size: 0 },
+      { repoId: "counts", name: "one", path: "src", size: 1 },
+      { repoId: "counts", name: "folder", path: "" },
+      { repoId: "counts", name: "staged", path: "src", size: 1, treeGeneration: "future" },
+    ]);
+    expect(await folderCounts(repo, "")).to.deep.equal({ "": 2, src: 1 });
+  });
+
   it("escapes regex punctuation and handles Unicode scopes and empty directories", async () => {
     await insert([{ path: "", name: "猫.[x]" }, { path: "猫.[x]", name: "empty" },
       { path: "猫.[x]/child", name: "a", size: 0 }, { path: "猫.[x]", name: "b", size: 2 },