From dc789c325023646409fda1d1ac6708ea4bb83081 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Thu, 20 Nov 2025 11:39:58 +0000 Subject: [PATCH] fix: implement streaming for large files to prevent memory exhaustion - Add stream-reader utility for memory-efficient file reading - Implement streaming in scanner, parser, and file-watcher for files > 10MB - Add memory monitoring to detect and prevent OOM conditions - Files are now streamed with max 50MB in memory at once - Large lines are automatically chunked to prevent buffer overflow This addresses the issue where indexing 10M+ lines of code causes the process to hang due to memory exhaustion. --- src/services/code-index/orchestrator.ts | 49 ++++++++ .../processors/__tests__/parser.spec.ts | 6 +- .../code-index/processors/file-watcher.ts | 25 +++- src/services/code-index/processors/parser.ts | 25 +++- src/services/code-index/processors/scanner.ts | 32 +++++- .../code-index/processors/stream-reader.ts | 107 ++++++++++++++++++ 6 files changed, 229 insertions(+), 15 deletions(-) create mode 100644 src/services/code-index/processors/stream-reader.ts diff --git a/src/services/code-index/orchestrator.ts b/src/services/code-index/orchestrator.ts index a07bfefe46..90d0497dff 100644 --- a/src/services/code-index/orchestrator.ts +++ b/src/services/code-index/orchestrator.ts @@ -1,5 +1,6 @@ import * as vscode from "vscode" import * as path from "path" +import * as os from "os" import { CodeIndexConfigManager } from "./config-manager" import { CodeIndexStateManager, IndexingState } from "./state-manager" import { IFileWatcher, IVectorStore, BatchProcessingSummary } from "./interfaces" @@ -93,6 +94,29 @@ export class CodeIndexOrchestrator { * Updates the status of a file in the state manager. */ + /** + * Gets the current memory usage percentage + */ + private getMemoryUsagePercent(): number { + const totalMem = os.totalmem() + const freeMem = os.freemem() + const usedMem = totalMem - freeMem + return (usedMem / totalMem) * 100 + } + + /** + * Checks if memory usage is at a critical level + */ + private isMemoryCritical(): boolean { + const memUsage = this.getMemoryUsagePercent() + const MEMORY_THRESHOLD = 90 // 90% memory usage is critical + if (memUsage > MEMORY_THRESHOLD) { + console.warn(`[CodeIndexOrchestrator] High memory usage detected: ${memUsage.toFixed(2)}%`) + return true + } + return false + } + /** * Initiates the indexing process (initial scan and starts watcher). */ @@ -144,6 +168,16 @@ export class CodeIndexOrchestrator { this._stopProgressMonitor() }, INDEXING_TIMEOUT_MS) + // Check initial memory before starting + if (this.isMemoryCritical()) { + this.stateManager.setSystemState( + "Error", + "Not enough memory available to start indexing. Please close other applications and try again.", + ) + this._isProcessing = false + return + } + // Start progress monitoring to detect stuck state this._startProgressMonitor() const collectionCreated = await this.vectorStore.initialize() @@ -445,6 +479,21 @@ export class CodeIndexOrchestrator { this._progressMonitorInterval = setInterval(() => { const timeSinceLastProgress = Date.now() - this._lastProgressUpdate + // Check memory usage + if (this.isMemoryCritical()) { + console.error( + `[CodeIndexOrchestrator] Critical memory usage detected during indexing - stopping process`, + ) + this.stateManager.setSystemState( + "Error", + "Indexing stopped due to high memory usage. Try indexing smaller portions or increasing available memory.", + ) + this._isProcessing = false + this.stopWatcher() + this._stopProgressMonitor() + return + } + if (this._isProcessing && timeSinceLastProgress > STUCK_THRESHOLD) { console.error( `[CodeIndexOrchestrator] Indexing appears stuck - no progress for ${timeSinceLastProgress / 1000} seconds`, diff --git a/src/services/code-index/processors/__tests__/parser.spec.ts b/src/services/code-index/processors/__tests__/parser.spec.ts index 68b523c7df..45edd366cd 100644 --- a/src/services/code-index/processors/__tests__/parser.spec.ts +++ b/src/services/code-index/processors/__tests__/parser.spec.ts @@ -3,7 +3,7 @@ import { CodeParser, codeParser } from "../parser" import { loadRequiredLanguageParsers } from "../../../tree-sitter/languageParser" import { parseMarkdown } from "../../../tree-sitter/markdownParser" -import { readFile } from "fs/promises" +import { readFile, stat } from "fs/promises" import { Node } from "web-tree-sitter" // Mock TelemetryService @@ -23,6 +23,7 @@ vi.mock("fs/promises", () => ({ mkdir: vi.fn(), access: vi.fn(), rename: vi.fn(), + stat: vi.fn(), constants: {}, }, readFile: vi.fn(), @@ -30,6 +31,7 @@ vi.mock("fs/promises", () => ({ mkdir: vi.fn(), access: vi.fn(), rename: vi.fn(), + stat: vi.fn(), })) vi.mock("../../../tree-sitter/languageParser") @@ -63,6 +65,8 @@ describe("CodeParser", () => { ;(loadRequiredLanguageParsers as any).mockResolvedValue(mockLanguageParser as any) // Set up default fs.readFile mock return value vi.mocked(readFile).mockResolvedValue("// default test content") + // Set up default fs.stat mock return value (small file size) + vi.mocked(stat).mockResolvedValue({ size: 1024 } as any) }) describe("parseFile", () => { diff --git a/src/services/code-index/processors/file-watcher.ts b/src/services/code-index/processors/file-watcher.ts index 1e5ebcbceb..6cf15b77f6 100644 --- a/src/services/code-index/processors/file-watcher.ts +++ b/src/services/code-index/processors/file-watcher.ts @@ -27,6 +27,7 @@ import { TelemetryService } from "@roo-code/telemetry" import { TelemetryEventName } from "@roo-code/types" import { sanitizeErrorMessage } from "../shared/validation-helpers" import { Package } from "../../../shared/package" +import { streamReadFile } from "./stream-reader" /** * Implementation of the file watcher interface @@ -540,12 +541,26 @@ export class FileWatcher implements IFileWatcher { } } - // Read file content - const fileContent = await vscode.workspace.fs.readFile(vscode.Uri.file(filePath)) - const content = fileContent.toString() + // Read file content with streaming for large files + const STREAM_THRESHOLD = 10 * 1024 * 1024 // 10MB + let content: string + let newHash: string - // Calculate hash - const newHash = createHash("sha256").update(content).digest("hex") + if (fileStat.size > STREAM_THRESHOLD) { + console.log( + `[FileWatcher] Using streaming for large file: ${filePath} (${(fileStat.size / 1024 / 1024).toFixed(2)}MB)`, + ) + const streamResult = await streamReadFile(filePath, 10000, 50) + content = streamResult.lines.join("\n") + newHash = streamResult.hash + if (streamResult.truncated) { + console.warn(`[FileWatcher] File truncated for processing: ${filePath}`) + } + } else { + const fileContent = await vscode.workspace.fs.readFile(vscode.Uri.file(filePath)) + content = fileContent.toString() + newHash = createHash("sha256").update(content).digest("hex") + } // Check if file has changed if (this.cacheManager.getHash(filePath) === newHash) { diff --git a/src/services/code-index/processors/parser.ts b/src/services/code-index/processors/parser.ts index 8611884ade..929297a151 100644 --- a/src/services/code-index/processors/parser.ts +++ b/src/services/code-index/processors/parser.ts @@ -1,4 +1,4 @@ -import { readFile } from "fs/promises" +import { readFile, stat } from "fs/promises" import { createHash } from "crypto" import * as path from "path" import { Node } from "web-tree-sitter" @@ -10,6 +10,7 @@ import { MAX_BLOCK_CHARS, MIN_BLOCK_CHARS, MIN_CHUNK_REMAINDER_CHARS, MAX_CHARS_ import { TelemetryService } from "@roo-code/telemetry" import { TelemetryEventName } from "@roo-code/types" import { sanitizeErrorMessage } from "../shared/validation-helpers" +import { streamReadFile } from "./stream-reader" /** * Implementation of the code parser interface @@ -50,8 +51,26 @@ export class CodeParser implements ICodeParser { fileHash = options.fileHash || this.createFileHash(content) } else { try { - content = await readFile(filePath, "utf8") - fileHash = this.createFileHash(content) + // Check file size for streaming decision + const stats = await stat(filePath) + const STREAM_THRESHOLD = 10 * 1024 * 1024 // 10MB + + if (stats.size > STREAM_THRESHOLD) { + // Use streaming for large files + console.log( + `[CodeParser] Using streaming for large file: ${filePath} (${(stats.size / 1024 / 1024).toFixed(2)}MB)`, + ) + const streamResult = await streamReadFile(filePath, 10000, 50) + content = streamResult.lines.join("\n") + fileHash = streamResult.hash + if (streamResult.truncated) { + console.warn(`[CodeParser] File truncated for processing: ${filePath}`) + } + } else { + // Small file - read normally + content = await readFile(filePath, "utf8") + fileHash = this.createFileHash(content) + } } catch (error) { console.error(`Error reading file ${filePath}:`, error) TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { diff --git a/src/services/code-index/processors/scanner.ts b/src/services/code-index/processors/scanner.ts index 92a7d77c27..0b2a3e6333 100644 --- a/src/services/code-index/processors/scanner.ts +++ b/src/services/code-index/processors/scanner.ts @@ -30,6 +30,7 @@ import { TelemetryService } from "@roo-code/telemetry" import { TelemetryEventName } from "@roo-code/types" import { sanitizeErrorMessage } from "../shared/validation-helpers" import { Package } from "../../../shared/package" +import { streamReadFile, streamProcessFile } from "./stream-reader" export class DirectoryScanner implements IDirectoryScanner { private readonly batchSegmentThreshold: number @@ -134,13 +135,32 @@ export class DirectoryScanner implements IDirectoryScanner { return } - // Read file content - const content = await vscode.workspace.fs - .readFile(vscode.Uri.file(filePath)) - .then((buffer) => Buffer.from(buffer).toString("utf-8")) + // Read file content with streaming for large files (> 10MB) + const STREAM_THRESHOLD = 10 * 1024 * 1024 // 10MB + let content: string + let currentFileHash: string + + if (stats.size > STREAM_THRESHOLD) { + // Use streaming for large files to avoid memory issues + console.log( + `[DirectoryScanner] Using streaming for large file: ${filePath} (${(stats.size / 1024 / 1024).toFixed(2)}MB)`, + ) + const streamResult = await streamReadFile(filePath, 10000, 50) // Max 50MB in memory + content = streamResult.lines.join("\n") + currentFileHash = streamResult.hash + if (streamResult.truncated) { + console.warn( + `[DirectoryScanner] File truncated for processing: ${filePath} (total: ${streamResult.totalSize} bytes)`, + ) + } + } else { + // Small file - read normally + content = await vscode.workspace.fs + .readFile(vscode.Uri.file(filePath)) + .then((buffer) => Buffer.from(buffer).toString("utf-8")) + currentFileHash = createHash("sha256").update(content).digest("hex") + } - // Calculate current hash - const currentFileHash = createHash("sha256").update(content).digest("hex") processedFiles.add(filePath) // Check against cache diff --git a/src/services/code-index/processors/stream-reader.ts b/src/services/code-index/processors/stream-reader.ts new file mode 100644 index 0000000000..a15ed47b4d --- /dev/null +++ b/src/services/code-index/processors/stream-reader.ts @@ -0,0 +1,107 @@ +import * as fs from "fs" +import * as readline from "readline" +import { createHash } from "crypto" +import { MAX_FILE_SIZE_BYTES } from "../constants" + +export interface StreamReadResult { + hash: string + lines: string[] + truncated: boolean + totalSize: number +} + +/** + * Reads a file using streaming to avoid loading entire file into memory. + * For very large files, returns a truncated version for processing. + */ +export async function streamReadFile( + filePath: string, + maxLinesPerChunk: number = 10000, + maxMemoryMB: number = 50, +): Promise { + const maxBytes = maxMemoryMB * 1024 * 1024 + const lines: string[] = [] + let totalSize = 0 + let truncated = false + let currentMemoryUsage = 0 + + const hash = createHash("sha256") + const fileStream = fs.createReadStream(filePath, { encoding: "utf8" }) + const rl = readline.createInterface({ + input: fileStream, + crlfDelay: Infinity, + }) + + return new Promise((resolve, reject) => { + rl.on("line", (line) => { + // Update hash with full line content + hash.update(line + "\n") + totalSize += line.length + 1 + + // Only store line if within memory limits + if (!truncated && currentMemoryUsage < maxBytes) { + lines.push(line) + currentMemoryUsage += line.length + 1 + + // Check if we should truncate + if (lines.length >= maxLinesPerChunk || currentMemoryUsage >= maxBytes) { + truncated = totalSize > currentMemoryUsage + } + } + }) + + rl.on("close", () => { + resolve({ + hash: hash.digest("hex"), + lines, + truncated, + totalSize, + }) + }) + + rl.on("error", reject) + fileStream.on("error", reject) + }) +} + +/** + * Process a large file in chunks without loading it entirely into memory + */ +export async function* streamProcessFile( + filePath: string, + chunkSizeLines: number = 5000, +): AsyncGenerator<{ lines: string[]; startLine: number; endLine: number }> { + const fileStream = fs.createReadStream(filePath, { encoding: "utf8" }) + const rl = readline.createInterface({ + input: fileStream, + crlfDelay: Infinity, + }) + + let lines: string[] = [] + let currentLine = 0 + let chunkStartLine = 1 + + for await (const line of rl) { + currentLine++ + lines.push(line) + + if (lines.length >= chunkSizeLines) { + yield { + lines: [...lines], + startLine: chunkStartLine, + endLine: currentLine, + } + lines = [] + chunkStartLine = currentLine + 1 + } + } + + // Yield remaining lines if any + if (lines.length > 0) { + yield { + lines, + startLine: chunkStartLine, + endLine: currentLine, + } + } +}