Make indexing functional in large workspaces by chunking segments of the workspace instead of workign on the entire workspace at once (#1876)

Co-authored-by: Rob Leidle <[email protected]>

Author: spew

core/indexing/CodebaseIndexer.ts

@@ -7,7 +7,7 @@ import { LanceDbIndex } from "./LanceDbIndex.js";
 import { ChunkCodebaseIndex } from "./chunk/ChunkCodebaseIndex.js";
 import { getComputeDeleteAddRemove } from "./refreshIndex.js";
 import { CodebaseIndex, IndexResultType } from "./types.js";
-import { walkDir } from "./walkDir.js";
+import { walkDirAsync } from "./walkDir.js";
 
 export class PauseToken {
   constructor(private _paused: boolean) {}
@@ -62,25 +62,7 @@ export class CodebaseIndexer {
     }
     const branch = await this.ide.getBranch(workspaceDir);
     const repoName = await this.ide.getRepoName(workspaceDir);
-    const stats = await this.ide.getLastModified([file]);
-    const indexesToBuild = await this.getIndexesToBuild();
-    for (const codebaseIndex of indexesToBuild) {
-      const tag: IndexTag = {
-        directory: workspaceDir,
-        branch,
-        artifactId: codebaseIndex.artifactId,
-      };
-      const [results, lastUpdated, markComplete] = await getComputeDeleteAddRemove(
-        tag,
-        { ...stats },
-        (filepath) => this.ide.readFile(filepath),
-        repoName,
-      );
-      for await (const _ of codebaseIndex.update(tag, results, markComplete, repoName)) {
-        lastUpdated.forEach((lastUpdated, path) => {
-          markComplete([lastUpdated], IndexResultType.UpdateLastUpdated);
-        });
-      }
+    for await (const updateDesc of this.indexFiles(workspaceDir, branch, repoName, [file])) {
     }
   }
 
@@ -131,33 +113,40 @@ export class CodebaseIndexer {
       desc: "Starting indexing...",
       status: "loading",
     };
+    const beginTime = Date.now();
 
     for (const directory of workspaceDirs) {
-      const files = await walkDir(directory, this.ide);
-      const stats = await this.ide.getLastModified(files);
+      const dirBasename = await this.basename(directory);
+      yield {
+        progress,
+        desc: `Discovering files in ${dirBasename}...`,
+        status: "indexing"
+      };
+      // compute the number of files in this directory to display an accurate progress bar
+      let totalFileCount = 0;
+      for await (const p of walkDirAsync(directory, this.ide)) {
+        totalFileCount += 1;
+        if (abortSignal.aborted) {
+          yield {
+            progress: 1,
+            desc: "Indexing cancelled",
+            status: "disabled",
+          };
+          return;
+        }
+        if (this.pauseToken.paused) {
+          yield *this.yieldUpdateAndPause();
+        }
+      }
+
       const branch = await this.ide.getBranch(directory);
       const repoName = await this.ide.getRepoName(directory);
-      let completedRelativeExpectedTime = 0;
-
-      for (const codebaseIndex of indexesToBuild) {
-        // TODO: IndexTag type should use repoName rather than directory
-        const tag: IndexTag = {
-          directory,
-          branch,
-          artifactId: codebaseIndex.artifactId,
-        };
-        const [results, lastUpdated, markComplete] = await getComputeDeleteAddRemove(
-          tag,
-          { ...stats },
-          (filepath) => this.ide.readFile(filepath),
-          repoName,
-        );
+      const batchSize = this.getBatchSize(totalFileCount);
+      let completedFileCount = 0;
 
+      for await (const files of this.walkDirInBatches(directory, batchSize)) {
         try {
-          for await (let {
-            progress: indexProgress,
-            desc,
-          } of codebaseIndex.update(tag, results, markComplete, repoName)) {
+          for await (const updateDesc of this.indexFiles(directory, branch, repoName, files)) {
             // Handle pausing in this loop because it's the only one really taking time
             if (abortSignal.aborted) {
               yield {
@@ -167,77 +156,125 @@ export class CodebaseIndexer {
               };
               return;
             }
-
             if (this.pauseToken.paused) {
-              yield {
-                progress,
-                desc: "Paused",
-                status: "paused",
-              };
-              while (this.pauseToken.paused) {
-                await new Promise((resolve) => setTimeout(resolve, 100));
-              }
+              yield *this.yieldUpdateAndPause();
             }
-
-            progress =
-              (completedDirs +
-                (completedRelativeExpectedTime +
-                  Math.min(1.0, indexProgress) *
-                    codebaseIndex.relativeExpectedTime) /
-                  totalRelativeExpectedTime) /
-              workspaceDirs.length;
             yield {
-              progress,
-              desc,
+              progress: progress,
+              desc: updateDesc,
               status: "indexing",
             };
           }
+        } catch (err) {
+          yield this.handleErrorAndGetProgressUpdate(err);
+          return;
+        }
+        completedFileCount += files.length;
+        progress = completedFileCount / totalFileCount / workspaceDirs.length + completedDirs / workspaceDirs.length;
+        this.logProgress(beginTime, completedFileCount, progress);
+      }
+      completedDirs += 1;
+    }
+    yield {
+      progress: 100,
+      desc: "Indexing Complete",
+      status: "done",
+    };
+  }
 
-          lastUpdated.forEach((lastUpdated, path) => {
-            markComplete([lastUpdated], IndexResultType.UpdateLastUpdated);
-          });
+  private handleErrorAndGetProgressUpdate(err: unknown): IndexingProgressUpdate {
+    console.log("error when indexing: ", err);
+    if (err instanceof Error) {
+      return this.errorToProgressUpdate(err);
+    }
+    return {
+      progress: 0,
+      desc: `Indexing failed: ${err}`,
+      status: "failed",
+    };
+  }
 
-          completedRelativeExpectedTime += codebaseIndex.relativeExpectedTime;
-          yield {
-            progress:
-              (completedDirs +
-                completedRelativeExpectedTime / totalRelativeExpectedTime) /
-              workspaceDirs.length,
-            desc: "Completed indexing " + codebaseIndex.artifactId,
-            status: "indexing",
-          };
-        } catch (e: any) {
-          let errMsg = `${e}`;
+  private errorToProgressUpdate(err: Error): IndexingProgressUpdate {
+    const errorRegex =
+      /Invalid argument error: Values length (\d+) is less than the length \((\d+)\) multiplied by the value size \(\d+\)/;
+    const match = err.message.match(errorRegex);
+    let errMsg: string;
+    if (match) {
+      const [_, valuesLength, expectedLength] = match;
+      errMsg = `Generated embedding had length ${valuesLength} but was expected to be ${expectedLength}. This may be solved by deleting ~/.continue/index and refreshing the window to re-index.`;
+    } else {
+      errMsg = `${err}`;
+    }
+    return {
+      progress: 0,
+      desc: errMsg,
+      status: "failed",
+    };
+  }
 
-          const errorRegex =
-            /Invalid argument error: Values length (\d+) is less than the length \((\d+)\) multiplied by the value size \(\d+\)/;
-          const match = e.message.match(errorRegex);
+  private logProgress(beginTime: number, completedFileCount: number, progress: number) {
+    const timeTaken = Date.now() - beginTime;
+    const seconds = Math.round(timeTaken / 1000);
+    const progressPercentage = (progress * 100).toFixed(1);
+    const filesPerSec = (completedFileCount / seconds).toFixed(2);
+    console.log(`Indexing: ${progressPercentage}% complete, elapsed time: ${seconds}s, ${filesPerSec} file/sec`);
+  }
 
-          if (match) {
-            const [_, valuesLength, expectedLength] = match;
-            errMsg = `Generated embedding had length ${valuesLength} but was expected to be ${expectedLength}. This may be solved by deleting ~/.continue/index and refreshing the window to re-index.`;
-          }
+  private async* yieldUpdateAndPause(): AsyncGenerator<IndexingProgressUpdate> {
+    yield {
+      progress: 0,
+      desc: "Indexing Paused",
+      status: "paused",
+    };
+    while (this.pauseToken.paused) {
+      await new Promise((resolve) => setTimeout(resolve, 100));
+    }
+  }
 
-          yield {
-            progress: 0,
-            desc: errMsg,
-            status: "failed",
-          };
+  private getBatchSize(workspaceSize: number): number {
+    // at least 10 and as much as 100 (in a repository with 10000 files)
+    return Math.min(100, Math.max(10, Math.floor(workspaceSize / 100)));
+  }
 
-          console.warn(
-            `Error updating the ${codebaseIndex.artifactId} index: ${e}`,
-          );
-          return;
-        }
+  /*
+   * enables the indexing operation to be completed in small batches, this is important in large
+   * repositories where indexing can quickly use up all the memory available
+   */
+  private async* walkDirInBatches(directory: string, batchSize: number): AsyncGenerator<string[]> {
+    let results = [];
+    for await (const p of walkDirAsync(directory, this.ide)) {
+      results.push(p);
+      if (results.length === batchSize) {
+        yield results;
+        results = [];
       }
+    }
+    if (results.length > 0) {
+      yield results;
+    }
+  }
 
-      completedDirs++;
-      progress = completedDirs / workspaceDirs.length;
-      yield {
-        progress,
-        desc: "Indexing Complete",
-        status: "done",
+  private async* indexFiles(workspaceDir: string, branch: string, repoName: string | undefined, filePaths: string[]): AsyncGenerator<string> {
+    const stats = await this.ide.getLastModified(filePaths);
+    const indexesToBuild = await this.getIndexesToBuild();
+    for (const codebaseIndex of indexesToBuild) {
+      const tag: IndexTag = {
+        directory: workspaceDir,
+        branch,
+        artifactId: codebaseIndex.artifactId,
       };
+      const [results, lastUpdated, markComplete] = await getComputeDeleteAddRemove(
+        tag,
+        { ...stats },
+        (filepath) => this.ide.readFile(filepath),
+        repoName,
+      );
+      for await (const { desc } of codebaseIndex.update(tag, results, markComplete, repoName)) {
+        lastUpdated.forEach((lastUpdated, path) => {
+          markComplete([lastUpdated], IndexResultType.UpdateLastUpdated);
+        });
+        yield desc;
+      }
     }
   }
 
@@ -250,4 +287,10 @@ export class CodebaseIndexer {
     }
     return undefined;
   }
+
+  private async basename(filepath: string): Promise<string> {
+    const pathSep = await this.ide.pathSep();
+    const path = filepath.split(pathSep);
+    return path[path.length - 1];
+  }
 }

core/indexing/walkDir.ts

@@ -46,6 +46,7 @@ class DFSWalker {
 
   // walk is a depth-first search implementation
   public async *walk(): AsyncGenerator<string> {
+    const fixupFunc = await this.newPathFixupFunc(this.options.returnRelativePaths ? "" : this.path, this.ide);
     const root: WalkContext = {
       walkableEntry: {
         relPath: "",
@@ -73,10 +74,10 @@ class DFSWalker {
           });
           if (this.options.onlyDirs) {
             // when onlyDirs is enabled the walker will only return directory names
-            yield w.relPath;
+            yield fixupFunc(w.relPath);
           }
         } else {
-          yield w.relPath;
+          yield fixupFunc(w.relPath);
         }
       }
     }
@@ -155,6 +156,24 @@ class DFSWalker {
   private entryIsSymlink(entry: Entry) {
     return entry[1] === (64 as FileType.SymbolicLink);
   }
+
+  // returns a function which will optionally prefix a root path and fixup the paths for the appropriate OS filesystem (i.e. windows)
+  // the reason to construct this function once is to avoid the need to call ide.pathSep() multiple times
+  private async newPathFixupFunc(rootPath: string, ide: IDE): Promise<(relPath: string) => string> {
+    const pathSep = await ide.pathSep();
+    const prefix = rootPath === "" ? "" : rootPath + pathSep;
+    if (pathSep === "/") {
+      if (rootPath === "") {
+        // return a no-op function in this case to avoid unnecessary string concatentation
+        return (relPath: string) => relPath;
+      }
+      return (relPath: string) => prefix + relPath;
+    }
+    // this serves to 'fix-up' the path on Windows
+    return (relPath: string) => {
+      return prefix + relPath.split("/").join(pathSep);
+    };
+  }
 }
 
 const defaultOptions: WalkerOptions = {
@@ -167,18 +186,18 @@ export async function walkDir(
   ide: IDE,
   _options?: WalkerOptions,
 ): Promise<string[]> {
-  let entries: string[] = [];
-  const options = { ...defaultOptions, ..._options };
-  const dfsWalker = new DFSWalker(path, ide, options);
-  let relativePaths: string[] = [];
-  for await (const e of dfsWalker.walk()) {
-    relativePaths.push(e);
+  let paths: string[] = [];
+  for await (const p of walkDirAsync(path, ide, _options)) {
+    paths.push(p);
   }
-  const pathSep = await ide.pathSep();
-  const prefix = options.returnRelativePaths ? "" : path + pathSep;
+  return paths;
+}
 
-  if (pathSep === "/") {
-    return relativePaths.map((p) => prefix + p);
-  }
-  return relativePaths.map((p) => prefix + p.split("/").join(pathSep));
+export async function* walkDirAsync(
+  path: string,
+  ide: IDE,
+  _options?: WalkerOptions,
+): AsyncGenerator<string> {
+  const options = { ...defaultOptions, ..._options };
+  yield* new DFSWalker(path, ide, options).walk();
 }