refactor: improve code quality and documentation in loop-runner and capture

mohammedahmed18 · mohammedahmed18 · commit 017bde1c1f59 · 2026-02-03T21:16:26.000Z
Improvements to loop-runner.js:
- Extract isValidJestRunnerPath() helper to reduce code duplication
- Add comprehensive JSDoc comments for Jest version detection
- Improve error messages with more context about detected versions
- Add better documentation for runTests() method
- Add validation for TestRunner class availability in Jest 30

Improvements to capture.js:
- Extract _recordAsyncTiming() helper to reduce duplication
- Add comprehensive JSDoc for _capturePerfAsync() with all parameters
- Improve error handling in async looping (record timing before throwing)
- Enhance shouldStopStability() documentation with algorithm details
- Improve code organization with clearer comments

These changes improve maintainability and debugging without changing behavior.
diff --git a/packages/codeflash/runtime/capture.js b/packages/codeflash/runtime/capture.js
@@ -267,26 +267,40 @@ const results = [];
 let db = null;
 
 /**
- * Check if performance has stabilized (for internal looping).
- * Matches Python's pytest_plugin.should_stop() logic.
+ * Check if performance has stabilized, allowing early stopping of benchmarks.
+ * Matches Python's pytest_plugin.should_stop() logic for consistency.
+ *
+ * Performance is considered stable when BOTH conditions are met:
+ * 1. CENTER: All recent measurements are within ±10% of the median
+ * 2. SPREAD: The range (max-min) is within 10% of the minimum
+ *
+ * @param {Array<number>} runtimes - Array of runtime measurements in microseconds
+ * @param {number} window - Number of recent measurements to check
+ * @param {number} minWindowSize - Minimum samples required before checking
+ * @returns {boolean} True if performance has stabilized
  */
 function shouldStopStability(runtimes, window, minWindowSize) {
     if (runtimes.length < window || runtimes.length < minWindowSize) {
         return false;
     }
+
     const recent = runtimes.slice(-window);
     const recentSorted = [...recent].sort((a, b) => a - b);
     const mid = Math.floor(window / 2);
     const median = window % 2 ? recentSorted[mid] : (recentSorted[mid - 1] + recentSorted[mid]) / 2;
 
+    // Check CENTER: all recent points must be close to median
     for (const r of recent) {
         if (Math.abs(r - median) / median > STABILITY_CENTER_TOLERANCE) {
             return false;
         }
     }
+
+    // Check SPREAD: range must be small relative to minimum
     const rMin = recentSorted[0];
     const rMax = recentSorted[recentSorted.length - 1];
     if (rMin === 0) return false;
+
     return (rMax - rMin) / rMin <= STABILITY_SPREAD_TOLERANCE;
 }
 
@@ -775,11 +789,40 @@ function capturePerf(funcName, lineId, fn, ...args) {
     return lastReturnValue;
 }
 
+/**
+ * Helper to record async timing and update state.
+ * @private
+ */
+function _recordAsyncTiming(startTime, testStdoutTag, durationNs, runtimes) {
+    console.log(`!######${testStdoutTag}:${durationNs}######!`);
+    sharedPerfState.totalLoopsCompleted++;
+    if (durationNs > 0) {
+        runtimes.push(durationNs / 1000);
+    }
+}
+
 /**
  * Async helper for capturePerf to handle async function looping.
  * This function awaits promises and continues the benchmark loop properly.
  *
  * @private
+ * @param {string} funcName - Name of the function being benchmarked
+ * @param {string} lineId - Line identifier for this capture point
+ * @param {Function} fn - The async function to benchmark
+ * @param {Array} args - Arguments to pass to fn
+ * @param {Promise} firstPromise - The first promise that was already started
+ * @param {number} firstStartTime - Start time of the first execution
+ * @param {string} firstTestStdoutTag - Timing marker tag for the first execution
+ * @param {string} safeModulePath - Sanitized module path
+ * @param {string|null} testClassName - Test class name (if any)
+ * @param {string} safeTestFunctionName - Sanitized test function name
+ * @param {string} invocationKey - Unique key for this invocation
+ * @param {Array<number>} runtimes - Array to collect runtimes for stability checking
+ * @param {number} batchSize - Number of iterations per batch
+ * @param {number} startBatchIndex - Index where async looping started
+ * @param {boolean} shouldLoop - Whether to continue looping
+ * @param {Function} getStabilityWindow - Function to get stability window size
+ * @returns {Promise} The last return value from fn
  */
 async function _capturePerfAsync(
     funcName, lineId, fn, args,
@@ -796,61 +839,52 @@ async function _capturePerfAsync(
         lastReturnValue = await firstPromise;
         const asyncEndTime = getTimeNs();
         const asyncDurationNs = getDurationNs(firstStartTime, asyncEndTime);
-        console.log(`!######${firstTestStdoutTag}:${asyncDurationNs}######!`);
-        sharedPerfState.totalLoopsCompleted++;
-        if (asyncDurationNs > 0) {
-            runtimes.push(asyncDurationNs / 1000);
-        }
+        _recordAsyncTiming(firstStartTime, firstTestStdoutTag, asyncDurationNs, runtimes);
     } catch (err) {
         const asyncEndTime = getTimeNs();
         const asyncDurationNs = getDurationNs(firstStartTime, asyncEndTime);
-        console.log(`!######${firstTestStdoutTag}:${asyncDurationNs}######!`);
-        sharedPerfState.totalLoopsCompleted++;
-        throw err;
+        _recordAsyncTiming(firstStartTime, firstTestStdoutTag, asyncDurationNs, runtimes);
+        lastError = err;
+        // Don't throw yet - we want to record the timing first
+    }
+
+    // If first iteration failed, stop and throw
+    if (lastError) {
+        throw lastError;
     }
 
     // Continue looping for remaining iterations
     for (let batchIndex = startBatchIndex + 1; batchIndex < batchSize; batchIndex++) {
-        // Check shared time limit
+        // Check exit conditions before starting next iteration
         if (shouldLoop && checkSharedTimeLimit()) {
             break;
         }
 
-        // Check if this invocation has already reached stability
         if (getPerfStabilityCheck() && sharedPerfState.stableInvocations[invocationKey]) {
             break;
         }
 
-        // Get the global loop index for this invocation
         const loopIndex = getInvocationLoopIndex(invocationKey);
-
-        // Check if we've exceeded max loops
         if (loopIndex > getPerfLoopCount()) {
             break;
         }
 
-        // Get invocation index for the timing marker
+        // Generate timing marker identifiers
         const testId = `${safeModulePath}:${testClassName}:${safeTestFunctionName}:${lineId}:${loopIndex}`;
         const invocationIndex = getInvocationIndex(testId);
         const invocationId = `${lineId}_${invocationIndex}`;
-
-        // Format stdout tag
         const testStdoutTag = `${safeModulePath}:${testClassName ? testClassName + '.' : ''}${safeTestFunctionName}:${funcName}:${loopIndex}:${invocationId}`;
 
+        // Execute and time the function
         try {
             const startTime = getTimeNs();
             lastReturnValue = await fn(...args);
             const endTime = getTimeNs();
             const durationNs = getDurationNs(startTime, endTime);
 
-            console.log(`!######${testStdoutTag}:${durationNs}######!`);
-            sharedPerfState.totalLoopsCompleted++;
-
-            if (durationNs > 0) {
-                runtimes.push(durationNs / 1000);
-            }
+            _recordAsyncTiming(startTime, testStdoutTag, durationNs, runtimes);
 
-            // Check stability
+            // Check if we've reached performance stability
             if (getPerfStabilityCheck() && runtimes.length >= getPerfMinLoops()) {
                 const window = getStabilityWindow();
                 if (shouldStopStability(runtimes, window, getPerfMinLoops())) {
diff --git a/packages/codeflash/runtime/loop-runner.js b/packages/codeflash/runtime/loop-runner.js
@@ -34,10 +34,26 @@ const { createRequire } = require('module');
 const path = require('path');
 const fs = require('fs');
 
+/**
+ * Validates that a jest-runner path is valid by checking for package.json.
+ * @param {string} jestRunnerPath - Path to check
+ * @returns {boolean} True if valid jest-runner package
+ */
+function isValidJestRunnerPath(jestRunnerPath) {
+    if (!fs.existsSync(jestRunnerPath)) {
+        return false;
+    }
+    const packageJsonPath = path.join(jestRunnerPath, 'package.json');
+    return fs.existsSync(packageJsonPath);
+}
+
 /**
  * Resolve jest-runner with monorepo support.
  * Uses CODEFLASH_MONOREPO_ROOT environment variable if available,
  * otherwise walks up the directory tree looking for node_modules/jest-runner.
+ *
+ * @returns {string} Path to jest-runner package
+ * @throws {Error} If jest-runner cannot be found
  */
 function resolveJestRunner() {
     // Try standard resolution first (works in simple projects)
@@ -51,11 +67,8 @@ function resolveJestRunner() {
     const monorepoRoot = process.env.CODEFLASH_MONOREPO_ROOT;
     if (monorepoRoot) {
         const jestRunnerPath = path.join(monorepoRoot, 'node_modules', 'jest-runner');
-        if (fs.existsSync(jestRunnerPath)) {
-            const packageJsonPath = path.join(jestRunnerPath, 'package.json');
-            if (fs.existsSync(packageJsonPath)) {
-                return jestRunnerPath;
-            }
+        if (isValidJestRunnerPath(jestRunnerPath)) {
+            return jestRunnerPath;
         }
     }
 
@@ -71,11 +84,8 @@ function resolveJestRunner() {
 
         // Try node_modules/jest-runner at this level
         const jestRunnerPath = path.join(currentDir, 'node_modules', 'jest-runner');
-        if (fs.existsSync(jestRunnerPath)) {
-            const packageJsonPath = path.join(jestRunnerPath, 'package.json');
-            if (fs.existsSync(packageJsonPath)) {
-                return jestRunnerPath;
-            }
+        if (isValidJestRunnerPath(jestRunnerPath)) {
+            return jestRunnerPath;
         }
 
         // Check if this is a workspace root (has monorepo markers)
@@ -91,11 +101,18 @@ function resolveJestRunner() {
         currentDir = path.dirname(currentDir);
     }
 
-    throw new Error('jest-runner not found');
+    throw new Error(
+        'jest-runner not found. Please install jest-runner in your project: npm install --save-dev jest-runner'
+    );
 }
 
-// Try to load jest-runner from the PROJECT's node_modules, not from codeflash package
-// This ensures we use the same version of jest-runner that the project uses
+/**
+ * Jest runner components - loaded dynamically from project's node_modules.
+ * This ensures we use the same version that the project uses.
+ *
+ * Jest 30+ uses TestRunner class with event-based architecture.
+ * Jest 29 uses runTest function for direct test execution.
+ */
 let TestRunner;
 let runTest;
 let jestRunnerAvailable = false;
@@ -110,19 +127,24 @@ try {
     TestRunner = jestRunner.default || jestRunner.TestRunner;
 
     if (TestRunner && TestRunner.prototype && typeof TestRunner.prototype.runTests === 'function') {
-        // Jest 30+ - use TestRunner class
+        // Jest 30+ - use TestRunner class with event emitter pattern
         jestVersion = 30;
         jestRunnerAvailable = true;
     } else {
         // Try Jest 29 style import
         try {
             runTest = internalRequire('./runTest').default;
             if (typeof runTest === 'function') {
+                // Jest 29 - use direct runTest function
                 jestVersion = 29;
                 jestRunnerAvailable = true;
             }
         } catch (e29) {
             // Neither Jest 29 nor 30 style import worked
+            const errorMsg = `Found jest-runner at ${jestRunnerPath} but could not load it. ` +
+                `This may indicate an unsupported Jest version. ` +
+                `Supported versions: Jest 29.x and Jest 30.x`;
+            console.error(errorMsg);
             jestRunnerAvailable = false;
         }
     }
@@ -203,15 +225,22 @@ class CodeflashLoopRunner {
                 'codeflash/loop-runner requires jest-runner to be installed.\n' +
                 'Please install it: npm install --save-dev jest-runner\n\n' +
                 'If you are using Vitest, the loop-runner is not needed - ' +
-                'Vitest projects use external looping handled by the Python runner.'
+                'Vitest projects use internal looping handled by capturePerf().'
             );
         }
+
         this._globalConfig = globalConfig;
         this._context = context || {};
         this._eventEmitter = new SimpleEventEmitter();
 
         // For Jest 30+, create an instance of the base TestRunner for delegation
-        if (jestVersion >= 30 && TestRunner) {
+        if (jestVersion >= 30) {
+            if (!TestRunner) {
+                throw new Error(
+                    `Jest ${jestVersion} detected but TestRunner class not available. ` +
+                    `This indicates an internal error in loop-runner initialization.`
+                );
+            }
             this._baseRunner = new TestRunner(globalConfig, context);
         }
     }
@@ -229,7 +258,17 @@ class CodeflashLoopRunner {
     }
 
     /**
-     * Run tests with batched looping for fair distribution.
+     * Run tests with batched looping for fair distribution across all test invocations.
+     *
+     * This implements the batched looping strategy:
+     *   Batch 1: Test1(N loops) → Test2(N loops) → Test3(N loops)
+     *   Batch 2: Test1(N loops) → Test2(N loops) → Test3(N loops)
+     *   ...until time budget exhausted or max batches reached
+     *
+     * @param {Array} tests - Jest test objects to run
+     * @param {Object} watcher - Jest watcher for interrupt handling
+     * @param {Object} options - Jest runner options
+     * @returns {Promise<void>}
      */
     async runTests(tests, watcher, options) {
         const startTime = Date.now();
@@ -238,7 +277,7 @@ class CodeflashLoopRunner {
         let allConsoleOutput = '';
 
         // Time limit check - must use local time tracking because Jest runs tests
-        // in worker processes, so shared state from capture.js isn't accessible here
+        // in isolated worker processes where shared state from capture.js isn't accessible
         const checkTimeLimit = () => {
             const elapsed = Date.now() - startTime;
             return elapsed >= TARGET_DURATION_MS && batchCount >= MIN_BATCHES;
