Fix memory leak and consolidate worker initialization #748
Conversation
TokenCounter instances were not being properly freed when worker threads were terminated by Tinypool's idle timeout. This caused memory leaks when using runCli as a library. Changes: - Add SIGTERM/SIGINT handlers to fileMetricsWorker and outputMetricsWorker - Add freeTokenCounters function with proper cleanup and debug logging - Convert all worker usage to consistent taskRunner pattern with cleanup - Add cleanupWorkerPool function for explicit worker pool termination - Update all related tests to match new taskRunner interface The fix ensures TokenCounter resources are properly freed when workers terminate, preventing memory accumulation during library usage. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
Add generic initTaskRunner function to processConcurrency.ts to eliminate duplicate initialization logic across multiple modules. This reduces code duplication and provides consistent worker pool management with proper type safety through generic parameters. - Add TaskRunner<T, R> interface and initTaskRunner function - Remove duplicate createTaskRunner wrappers from 5 modules - Update all deps parameters to use shared initTaskRunner directly - Maintain type safety with explicit generic type parameters - Update corresponding test mocks to match new signature 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the WalkthroughThis change standardizes concurrency handling by introducing a generic Changes
Sequence Diagram(s)sequenceDiagram
participant Caller
participant TaskRunner (from initTaskRunner)
participant WorkerPool
Caller->>TaskRunner: run(task)
TaskRunner->>WorkerPool: execute task
WorkerPool-->>TaskRunner: result
TaskRunner-->>Caller: result
Caller->>TaskRunner: cleanup()
TaskRunner->>WorkerPool: terminate pool
WorkerPool-->>TaskRunner: cleanup complete
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
✨ Finishing Touches🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
|
Claude finished @yamadashy's task —— View job PR Review: Fix memory leak and consolidate worker initialization ✅Overall Assessment: Excellent - This PR demonstrates strong software engineering practices with significant improvements to memory management and code maintainability. Memory Leak Fix Analysis ✅The memory leak fix in ✅ Good practices:
✅ Cleanup integration: The cleanup is properly hooked into worker termination via Code Consolidation Analysis ✅The generic ✅ Strong design decisions:
✅ Consistent usage pattern: All modules now follow the same pattern: const taskRunner = deps.initTaskRunner<TaskType, ResultType>(numTasks, workerPath);
try {
// ... use taskRunner.run(task)
} finally {
await taskRunner.cleanup(); // Always executed
}Type Safety & Error Handling ✅Excellent type safety implementation:
Robust error handling:
Test Coverage & Quality ✅Well-maintained test coverage:
Architecture & Design Excellence ✅Benefits achieved:
Recommendations & Minor SuggestionsMinor enhancements to consider:
No blocking issues identified - this PR is ready for merge. ConclusionThis PR exemplifies excellent software engineering practices:
The consolidation approach is architecturally sound and will make future maintenance much easier. Great work! 🎉 |
There was a problem hiding this comment.
Pull Request Overview
This PR addresses memory leaks in worker threads and consolidates duplicate worker initialization logic across multiple modules.
- Fixes TokenCounter memory leak by properly freeing resources on worker termination
- Consolidates duplicate worker initialization code into a generic
initTaskRunnerfunction - Adds proper worker pool cleanup with
finallyblocks across all modules
Reviewed Changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| src/shared/processConcurrency.ts | Adds generic initTaskRunner and cleanupWorkerPool functions |
| src/core/metrics/tokenCounterFactory.ts | Renames function and adds debug logging for resource cleanup |
| src/core/metrics/workers/*.ts | Updates worker termination handlers to use renamed cleanup function |
| src/core//.ts | Replaces module-specific task runners with generic implementation |
| tests//.test.ts | Updates test mocks to match new function signatures |
There was a problem hiding this comment.
Summary of Changes
Hello @yamadashy, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request primarily focuses on improving the stability and maintainability of the codebase by addressing a critical memory leak in worker threads and refactoring the worker initialization logic. It introduces a centralized, generic approach for managing worker pools, ensuring proper resource cleanup and reducing code duplication across various modules.
Highlights
- Memory Leak Fix: Addressed a memory leak related to
TokenCounterinstances in worker threads by ensuring resources are properly freed upon worker termination. This includes renamingfreeTokenCountertofreeTokenCountersfor clarity and adding debug logging for cleanup. - Worker Initialization Consolidation: Introduced a new generic
initTaskRunner<T, R>function inprocessConcurrency.tsto centralize and standardize the creation and management of worker pools, reducing boilerplate code across multiple modules. - Consistent Worker Pool Cleanup: Implemented
finallyblocks in all modules utilizing worker pools to guarantee that worker resources are consistently cleaned up, even if errors occur during task execution. - Enhanced Type Safety: The new generic
initTaskRunnermaintains type safety for worker tasks and results, improving code reliability and developer experience. - Test Suite Updates: Modified existing test mocks to align with the new
TaskRunnerinterface, ensuring continued test coverage and correctness.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request refactors the worker initialization logic, significantly reducing code duplication by introducing a generic initTaskRunner. The addition of finally blocks to ensure worker pool cleanup improves robustness. The memory leak fix in TokenCounter is also well-addressed.
I've left a couple of minor suggestions for improvement, one regarding an unused import and another on enhancing error logging. Overall, this is a high-quality contribution that improves maintainability and reliability.
| @@ -1,7 +1,8 @@ | |||
| import { parentPort } from 'node:worker_threads'; | |||
There was a problem hiding this comment.
This parentPort import seems to be unused in the file. It's best to remove it to keep the code clean and avoid confusion.
src/shared/processConcurrency.ts
Outdated
| await pool.destroy(); | ||
| logger.debug('Worker pool cleaned up successfully'); | ||
| } catch (error) { | ||
| logger.debug(`Error during worker pool cleanup: ${error}`); |
There was a problem hiding this comment.
Using string interpolation for logging errors (logger.debug(...${error})) can sometimes hide useful information like the stack trace. It's generally better to pass the error object as a separate argument to the logger, like logger.debug('message', error). This allows the logging library to format it properly, often including the stack trace which is very helpful for debugging.
| logger.debug(`Error during worker pool cleanup: ${error}`); | |
| logger.debug('Error during worker pool cleanup:', error); |
Codecov ReportAttention: Patch coverage is
❌ Your patch check has failed because the patch coverage (76.05%) is below the target coverage (80.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #748 +/- ##
==========================================
+ Coverage 89.68% 89.77% +0.09%
==========================================
Files 101 101
Lines 5806 5838 +32
Branches 1203 1205 +2
==========================================
+ Hits 5207 5241 +34
+ Misses 599 597 -2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Remove unused parentPort import from fileMetricsWorker.ts that was introduced during the memory leak fix but is not actually used. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
Deploying repomix with
|
| Latest commit: |
aad3d81
|
| Status: | ✅ Deploy successful! |
| Preview URL: | https://a6ec3f3f.repomix.pages.dev |
| Branch Preview URL: | https://feat-memory.repomix.pages.dev |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
src/core/metrics/tokenCounterFactory.ts (1)
2-2: LGTM! Good improvements for debugging memory leaks.The function renaming to
freeTokenCounters(plural) better reflects its behavior of freeing multiple token counters. The added debug logging provides valuable visibility into resource cleanup.Consider making the debug log message more specific by including the encoding information:
- logger.debug('Freed TokenCounter resources for encoding.'); + logger.debug(`Freed TokenCounter resources for encoding: ${encoding}.`);However, this would require capturing the encoding from the Map entries, which might not be worth the complexity for a debug message.
Also applies to: 25-31
tests/integration-tests/packager.test.ts (1)
27-36: LGTM! Well-implemented refactor to match new task runner interface.The mock has been correctly updated to return an object with
runandcleanupmethods, aligning with the new standardized task runner pattern. The implementation preserves the original behavior while adding the required cleanup method as a no-op for testing.The type casting in the
runmethod (task as FileCollectTaskandas R) maintains functionality but reduces type safety. This is acceptable for test mocks, but consider if more specific typing could be used if the mock is only used with known types.src/shared/processConcurrency.ts (1)
53-63: Consider re-throwing errors in cleanup for better error visibility.While logging errors during cleanup is good, consider whether errors should be re-thrown after logging to ensure callers are aware of cleanup failures. This could be important for debugging memory leaks.
export const cleanupWorkerPool = async (pool: Tinypool): Promise<void> => { try { logger.debug('Cleaning up worker pool...'); // Terminate pool which should trigger worker cleanup await pool.destroy(); logger.debug('Worker pool cleaned up successfully'); } catch (error) { logger.debug(`Error during worker pool cleanup: ${error}`); + // Re-throw to ensure callers are aware of cleanup failures + throw error; } };
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
src/core/file/fileCollect.ts(4 hunks)src/core/file/fileProcess.ts(4 hunks)src/core/metrics/calculateAllFileMetrics.ts(7 hunks)src/core/metrics/calculateOutputMetrics.ts(5 hunks)src/core/metrics/tokenCounterFactory.ts(2 hunks)src/core/metrics/workers/fileMetricsWorker.ts(2 hunks)src/core/metrics/workers/outputMetricsWorker.ts(2 hunks)src/core/security/securityCheck.ts(4 hunks)src/shared/processConcurrency.ts(2 hunks)tests/core/file/fileCollect.test.ts(1 hunks)tests/core/file/fileProcess.test.ts(1 hunks)tests/core/metrics/calculateAllFileMetrics.test.ts(1 hunks)tests/core/metrics/calculateOutputMetrics.test.ts(5 hunks)tests/core/security/securityCheck.test.ts(3 hunks)tests/integration-tests/packager.test.ts(1 hunks)tests/shared/processConcurrency.test.ts(3 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (15)
src/core/metrics/workers/outputMetricsWorker.ts (1)
src/core/metrics/tokenCounterFactory.ts (1)
freeTokenCounters(25-31)
tests/core/security/securityCheck.test.ts (1)
src/core/security/workers/securityCheckWorker.ts (2)
SecurityCheckTask(13-17)SecurityCheckTask(25-42)
src/core/file/fileProcess.ts (2)
src/core/file/workers/fileProcessWorker.ts (2)
FileProcessTask(10-13)FileProcessTask(15-21)src/core/file/fileTypes.ts (1)
ProcessedFile(6-9)
src/shared/processConcurrency.ts (1)
src/shared/logger.ts (1)
logger(89-89)
tests/integration-tests/packager.test.ts (1)
src/core/file/workers/fileCollectWorker.ts (2)
FileCollectTask(9-13)FileCollectTask(15-27)
tests/core/metrics/calculateOutputMetrics.test.ts (1)
src/core/metrics/workers/outputMetricsWorker.ts (2)
OutputMetricsTask(9-13)OutputMetricsTask(15-26)
src/core/security/securityCheck.ts (1)
src/core/security/workers/securityCheckWorker.ts (3)
SecurityCheckTask(13-17)SecurityCheckTask(25-42)SuspiciousFileResult(19-23)
src/core/file/fileCollect.ts (2)
src/core/file/workers/fileCollectWorker.ts (2)
FileCollectTask(9-13)FileCollectTask(15-27)src/core/file/fileTypes.ts (1)
RawFile(1-4)
tests/core/file/fileProcess.test.ts (1)
src/core/file/workers/fileProcessWorker.ts (2)
FileProcessTask(10-13)FileProcessTask(15-21)
src/core/metrics/calculateAllFileMetrics.ts (2)
src/core/metrics/workers/fileMetricsWorker.ts (2)
FileMetricsTask(12-17)FileMetricsTask(19-28)src/core/metrics/workers/types.ts (1)
FileMetrics(1-5)
src/core/metrics/workers/fileMetricsWorker.ts (1)
src/core/metrics/tokenCounterFactory.ts (1)
freeTokenCounters(25-31)
tests/core/metrics/calculateAllFileMetrics.test.ts (1)
src/core/metrics/workers/fileMetricsWorker.ts (2)
FileMetricsTask(12-17)FileMetricsTask(19-28)
tests/shared/processConcurrency.test.ts (1)
src/shared/processConcurrency.ts (2)
createWorkerPool(26-51)initTaskRunner(70-76)
tests/core/file/fileCollect.test.ts (1)
src/core/file/workers/fileCollectWorker.ts (2)
FileCollectTask(9-13)FileCollectTask(15-27)
src/core/metrics/tokenCounterFactory.ts (1)
src/shared/logger.ts (1)
logger(89-89)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (13)
- GitHub Check: Test (windows-latest, 24.x)
- GitHub Check: Test (macos-latest, 22.x)
- GitHub Check: Test (windows-latest, 23.x)
- GitHub Check: Test (macos-latest, 23.x)
- GitHub Check: Test (windows-latest, 22.x)
- GitHub Check: Test (windows-latest, 21.x)
- GitHub Check: Test (windows-latest, 20.x)
- GitHub Check: Build and run (windows-latest, 22.x)
- GitHub Check: Test coverage
- GitHub Check: Build and run with Bun (ubuntu-latest, latest)
- GitHub Check: Build and run with Bun (windows-latest, latest)
- GitHub Check: Build and run with Bun (macos-latest, latest)
- GitHub Check: auto-review
🔇 Additional comments (32)
src/core/metrics/workers/outputMetricsWorker.ts (1)
3-3: LGTM! Consistent function renaming applied.The import and usage have been correctly updated from
freeTokenCountertofreeTokenCounters, maintaining consistency with the renamed function intokenCounterFactory.ts.Also applies to: 30-30
src/core/metrics/workers/fileMetricsWorker.ts (1)
5-5: LGTM! Consistent function renaming applied.The import and usage have been correctly updated from
freeTokenCountertofreeTokenCounters, maintaining consistency with the renamed function intokenCounterFactory.ts.Also applies to: 43-43
tests/shared/processConcurrency.test.ts (2)
4-9: LGTM! Test correctly updated for renamed function.The test has been properly updated to test
createWorkerPoolinstead of the previous function name, maintaining the same validation logic for Tinypool initialization.Also applies to: 68-89
91-112: LGTM! Good test coverage for the new task runner interface.The new test suite properly validates that
initTaskRunnerreturns an object with the expectedrunandcleanupmethods, ensuring the interface contract is maintained. The mocking setup appropriately simulates the Tinypool behavior.src/core/metrics/calculateOutputMetrics.ts (4)
3-3: LGTM: Import consolidation implemented correctly.The import of
initTaskRunnerfrom the shared concurrency module aligns with the PR objective to consolidate duplicate worker initialization logic.
19-22: LGTM: Task runner initialization follows the new shared interface correctly.The generic type parameters
<OutputMetricsTask, number>correctly match the worker's input/output types, and the worker script URL is properly resolved usingimport.meta.url.
42-42: LGTM: Task execution updated to use the new TaskRunner interface.The calls to
taskRunner.run(task)correctly replace the previous direct function calls, maintaining the same functionality while using the standardized interface.Also applies to: 54-54
65-68: Excellent: Critical memory leak fix implemented.The
finallyblock ensures thattaskRunner.cleanup()is always called, preventing memory leaks by properly freeing worker pool resources regardless of success or failure. This directly addresses the memory leak issue mentioned in the PR objectives.src/core/file/fileCollect.ts (4)
4-4: LGTM: Consistent import consolidation.The import aligns with the shared concurrency module pattern used across other files in this refactor.
18-21: LGTM: Task runner initialization properly typed.The generic parameters
<FileCollectTask, RawFile | null>correctly reflect the worker's interface, and the number of tasks is appropriately set to the number of file paths.
40-40: LGTM: Task execution correctly updated.The change to
taskRunner.run(task)maintains the same functionality while using the standardized TaskRunner interface.
57-60: Excellent: Consistent resource cleanup implementation.The
finallyblock withtaskRunner.cleanup()ensures proper worker pool cleanup, preventing memory leaks and maintaining consistency across all modules.tests/core/file/fileCollect.test.ts (1)
24-33: LGTM: Test mock properly adapted to new TaskRunner interface.The mock function correctly:
- Accepts generic type parameters and the required parameters
- Returns an object with
runandcleanupmethods matching the TaskRunner interface- Delegates to the actual worker function while maintaining type safety
- Provides a no-op cleanup method appropriate for testing
This ensures tests continue to work with the new concurrency abstraction.
tests/core/file/fileProcess.test.ts (1)
22-31: LGTM: Consistent test mock implementation.The mock implementation follows the same pattern as other test files, correctly adapting to the new TaskRunner interface while maintaining proper delegation to the
fileProcessWorkerfor accurate testing.src/core/file/fileProcess.ts (4)
4-4: LGTM: Import consolidation consistent with other files.The import change aligns with the shared concurrency module pattern used throughout this refactor.
24-27: LGTM: Task runner initialization with correct type parameters.The generic parameters
<FileProcessTask, ProcessedFile>accurately reflect the worker's input and output types, and the task count is appropriately set to the number of raw files.
45-45: LGTM: Task execution follows standardized pattern.The use of
taskRunner.run(task)is consistent with the new TaskRunner interface implemented across all modules.
62-65: Excellent: Memory leak prevention implemented consistently.The
finallyblock ensures proper cleanup of worker pool resources, completing the memory leak fix across all modules that use worker threads.tests/core/metrics/calculateAllFileMetrics.test.ts (1)
15-24: LGTM! Mock correctly implements the new task runner interface.The mock properly returns an object with
runandcleanupmethods, aligning with the newTaskRunnerinterface introduced in the production code.src/core/metrics/calculateAllFileMetrics.ts (2)
56-59: Excellent addition of cleanup in finally block.The
finallyblock ensures that the worker pool is always cleaned up, even if an error occurs during metrics calculation. This directly addresses the memory leak issue mentioned in the PR objectives.
116-119: Consistent cleanup pattern applied.Good to see the same cleanup pattern applied in
calculateSelectiveFileMetrics. This ensures consistent resource management across both functions.src/shared/processConcurrency.ts (2)
26-51: Good rename for clarity.Renaming
initWorkertocreateWorkerPoolbetter reflects the function's purpose of creating a Tinypool instance rather than initializing a single worker.
70-76: Clean implementation of the generic task runner.The generic
initTaskRunnerfunction effectively consolidates the duplicate worker initialization logic mentioned in the PR objectives. The use of generics provides type safety while maintaining flexibility.tests/core/security/securityCheck.test.ts (2)
42-51: Mock correctly implements the new interface.The mock task runner properly returns an object with
runandcleanupmethods, matching the productionTaskRunnerinterface.
79-88: Good error handling test coverage.The error handling test correctly verifies that errors are propagated from the task runner's
runmethod while still maintaining the cleanup interface.src/core/security/securityCheck.ts (1)
86-89: Proper resource cleanup implemented.The
finallyblock ensures that the worker pool is cleaned up regardless of success or failure, preventing memory leaks as mentioned in the PR objectives.tests/core/metrics/calculateOutputMetrics.test.ts (6)
9-18: LGTM! Clean implementation of the new TaskRunner interface.The mock correctly implements the new interface pattern with proper generic type parameters and appropriate test-specific behavior (no-op cleanup).
49-58: Error handling mock correctly implemented.The mock properly throws errors in the
runmethod while maintaining the same interface structure as other mocks.
99-110: Parallel processing mock implementation looks good.The mock correctly tracks chunk processing and returns consistent values for parallel processing test scenarios.
125-134: Parallel error handling mock correctly implemented.The mock appropriately throws errors for testing parallel processing error scenarios while maintaining interface consistency.
150-161: Chunk tracking mock implementation is effective.The mock properly collects processed chunks and returns content length, enabling comprehensive validation of chunking behavior.
1-176: Excellent refactoring to align with the new TaskRunner interface.The test file has been systematically updated to work with the new generic
initTaskRunnerinterface. All mock implementations are consistent, maintain type safety, and preserve the original test logic while adapting to the new worker pool pattern with proper cleanup methods.
The issue where Bun runtime causes process termination when calling pool.destroy() on Tinypool is a known compatibility problem. This implementation skips the problematic destroy() call in Bun while maintaining proper cleanup in Node.js environments. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
- Add encoding name to TokenCounter cleanup debug messages - Improves debugging experience by showing which encoding was freed - Addresses PR review feedback for better observability 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
Summary
This PR addresses memory leaks in worker threads and consolidates duplicate worker initialization logic across multiple modules.
Changes
Memory Leak Fix:
freeTokenCountertofreeTokenCountersfor clarityCode Consolidation:
initTaskRunner<T, R>function toprocessConcurrency.tsinitTaskRunnerimplementations from 5 modules:fileCollect.tsfileProcess.tssecurityCheck.tscalculateAllFileMetrics.tscalculateOutputMetrics.tsfinallyblocksBenefits
Test plan
npm run test- All 694 tests passnpm run lint- No lint errors🤖 Generated with Claude Code