From 343dbb99b83eba91b5c760ee7f72dd6f639d8322 Mon Sep 17 00:00:00 2001 From: Widthdom Date: Fri, 5 Jun 2026 10:49:07 +0900 Subject: [PATCH] Fix post-extraction hook timeout isolation (#3241) --- DEVELOPER_GUIDE.md | 8 +- changelog.d/unreleased/3241.security.md | 19 + src/CodeIndex/Cli/ProgramRunner.cs | 4 + .../Hooks/PostExtractionHookCallbackWorker.cs | 562 ++++++++++++++++++ .../Indexer/Hooks/PostExtractionHooks.cs | 89 ++- src/CodeIndex/Mcp/McpToolHandlers.cs | 2 +- .../PostExtractionHookTests.cs | 210 ++++++- 7 files changed, 848 insertions(+), 46 deletions(-) create mode 100644 changelog.d/unreleased/3241.security.md create mode 100644 src/CodeIndex/Indexer/Hooks/PostExtractionHookCallbackWorker.cs diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index 01629f2612..53acc157ae 100644 --- a/DEVELOPER_GUIDE.md +++ b/DEVELOPER_GUIDE.md @@ -177,7 +177,7 @@ trust metadata before the dependent rows are written. Out-of-tree post-extraction hooks can implement `CodeIndex.Indexer.Hooks.IPostExtractionHook` in a `.dll` placed under `~/.config/cdidx/hooks/` (or the directory named by `CDIDX_HOOKS_DIR`). Hook discovery examines at most `CDIDX_HOOK_DISCOVERY_MAX_DLLS` DLL candidates (default: 128), requires each candidate to be no larger than `CDIDX_HOOK_DISCOVERY_MAX_BYTES` bytes (default: 67108864), then loads the bounded candidate set in path order. Each concrete hook type is instantiated with a public parameterless constructor, then called after built-in symbol extraction and again after built-in reference extraction, before rows are persisted. Hooks receive a `FileContext` plus mutable `IList` / `IList` values, so they can annotate extracted records, add synthetic symbols, or add domain-specific references. -Hook failures are isolated to that hook invocation: assembly load, construction, and callback exceptions are captured as diagnostics and indexing continues. Each callback runs against a scratch copy with a bounded wall-clock budget controlled by `CDIDX_HOOK_CALLBACK_BUDGET_MS` (default: 5000 ms). A timed-out callback contributes no mutations, emits an index warning, and disables that hook for the remainder of the current index run. `status --json` and MCP `status` expose loaded hooks under `hooks` with `name`, `assembly_path`, `type_name`, and `callback_budget_ms` so users can confirm which extensions are active and what timeout is being enforced. +Hook failures are isolated to that hook invocation: assembly load, construction, and callback exceptions are captured as diagnostics and indexing continues. Each loaded hook runs in an isolated worker process, and callbacks run against scratch copies with a bounded wall-clock budget controlled by `CDIDX_HOOK_CALLBACK_BUDGET_MS` (default: 5000 ms). The first callback budget covers worker startup and callback execution. A timed-out callback kills the worker process tree, contributes no mutations, emits an index warning, and disables that hook for the remainder of the current index run. `status --json` and MCP `status` expose loaded hooks under `hooks` with `name`, `assembly_path`, `type_name`, and `callback_budget_ms` so users can confirm which extensions are active and what timeout is being enforced. ### Ignore file parsing @@ -2269,8 +2269,10 @@ hook は `FileContext` と mutable な `IList` / `IList? Symbols, + List? References); + +internal sealed class PostExtractionHookCallbackWorkerClient : IDisposable +{ + private readonly PostExtractionHookInfo hook; + private readonly object gate = new(); + private Process? process; + private StringBuilder stderr = new(); + private bool disposed; + + internal PostExtractionHookCallbackWorkerClient(PostExtractionHookInfo hook) + { + this.hook = hook; + } + + internal PostExtractionHookCallbackResult Invoke( + PostExtractionHookCallbackKind kind, + string callback, + FileContext context, + IReadOnlyList? symbols, + IReadOnlyList? references, + TimeSpan callbackBudget) + { + lock (gate) + { + ObjectDisposedException.ThrowIf(disposed, this); + var stopwatch = Stopwatch.StartNew(); + if (!EnsureStarted(out var startError)) + { + stopwatch.Stop(); + return Failure(startError, stopwatch.ElapsedMilliseconds); + } + + var request = new PostExtractionHookCallbackWorker.WorkerRequest( + callback, + context, + symbols?.ToList(), + references?.ToList()); + var requestJson = JsonSerializer.Serialize(request, PostExtractionHookCallbackWorker.JsonOptions); + var waitMilliseconds = GetRemainingWaitMilliseconds(stopwatch, callbackBudget); + if (waitMilliseconds <= 0) + { + KillWorker(); + stopwatch.Stop(); + return TimedOut(stopwatch.ElapsedMilliseconds); + } + + Task responseTask; + Task sendTask; + try + { + responseTask = process!.StandardOutput.ReadLineAsync(); + sendTask = SendRequestAsync(process.StandardInput, requestJson); + } + catch (Exception ex) + { + KillWorker(); + stopwatch.Stop(); + return Failure($"failed to send worker request: {ex.Message}", stopwatch.ElapsedMilliseconds); + } + + if (!WaitForTask(sendTask, waitMilliseconds, out var sendException)) + { + KillWorker(); + stopwatch.Stop(); + return TimedOut(stopwatch.ElapsedMilliseconds); + } + + if (sendException != null) + { + KillWorker(); + stopwatch.Stop(); + return Failure($"failed to send worker request: {sendException.Message}", stopwatch.ElapsedMilliseconds); + } + + waitMilliseconds = GetRemainingWaitMilliseconds(stopwatch, callbackBudget); + if (waitMilliseconds <= 0 || !WaitForTask(responseTask, waitMilliseconds, out var responseException)) + { + KillWorker(); + stopwatch.Stop(); + return TimedOut(stopwatch.ElapsedMilliseconds); + } + + if (responseException != null) + { + KillWorker(); + stopwatch.Stop(); + return Failure($"failed to read worker response: {responseException.Message}", stopwatch.ElapsedMilliseconds); + } + + stopwatch.Stop(); + var responseJson = responseTask.GetAwaiter().GetResult(); + if (responseJson == null) + { + var workerError = BuildWorkerExitError(process, stderr.ToString(), "worker exited before returning a response."); + ClearExitedWorker(); + return Failure(workerError, stopwatch.ElapsedMilliseconds); + } + + PostExtractionHookCallbackWorker.WorkerResponse? response; + try + { + response = JsonSerializer.Deserialize( + responseJson, + PostExtractionHookCallbackWorker.JsonOptions); + } + catch (JsonException ex) + { + KillWorker(); + return Failure($"worker returned invalid JSON: {ex.Message}", stopwatch.ElapsedMilliseconds); + } + + if (response == null) + return Failure("worker returned an empty response.", stopwatch.ElapsedMilliseconds); + if (!string.IsNullOrWhiteSpace(response.WorkerError)) + return Failure(response.WorkerError, stopwatch.ElapsedMilliseconds); + if (kind == PostExtractionHookCallbackKind.Symbols && response.Symbols == null) + return Failure("worker response omitted symbols.", stopwatch.ElapsedMilliseconds); + if (kind == PostExtractionHookCallbackKind.References && response.References == null) + return Failure("worker response omitted references.", stopwatch.ElapsedMilliseconds); + + return new PostExtractionHookCallbackResult( + Success: true, + TimedOut: false, + WorkerError: null, + CallbackError: response.CallbackError, + DurationMs: stopwatch.ElapsedMilliseconds, + Symbols: response.Symbols, + References: response.References); + } + } + + public void Dispose() + { + lock (gate) + { + if (disposed) + return; + + disposed = true; + if (process == null) + return; + + try + { + process.StandardInput.Close(); + } + catch + { + // Best effort: disposal should not throw after indexing has completed. + } + + if (!WaitForWorkerExit(process, 1000)) + KillWorker(); + else + ClearExitedWorker(); + } + } + + private bool EnsureStarted(out string error) + { + if (process is { HasExited: false }) + { + error = string.Empty; + return true; + } + + ClearExitedWorker(); + stderr = new StringBuilder(); + if (!PostExtractionHookCallbackWorker.TryCreateStartInfo(hook, out var startInfo, out error)) + return false; + + var next = new Process { StartInfo = startInfo, EnableRaisingEvents = true }; + next.ErrorDataReceived += (_, eventArgs) => + { + if (eventArgs.Data != null) + stderr.AppendLine(eventArgs.Data); + }; + + try + { + if (!next.Start()) + { + error = "worker process did not start."; + next.Dispose(); + return false; + } + + next.BeginErrorReadLine(); + process = next; + error = string.Empty; + return true; + } + catch (Exception ex) + { + error = $"failed to start worker process: {ex.Message}"; + next.Dispose(); + return false; + } + } + + private void KillWorker() + { + if (process == null) + return; + + PostExtractionHookCallbackWorker.TryKillProcess(process); + ClearExitedWorker(); + } + + private void ClearExitedWorker() + { + if (process == null) + return; + + process.Dispose(); + process = null; + } + + private static PostExtractionHookCallbackResult Failure(string? message, long durationMs) + => new( + Success: false, + TimedOut: false, + WorkerError: string.IsNullOrWhiteSpace(message) ? "isolated hook callback worker failed." : message, + CallbackError: null, + DurationMs: Math.Max(0, durationMs), + Symbols: null, + References: null); + + private static PostExtractionHookCallbackResult TimedOut(long durationMs) + => new( + Success: false, + TimedOut: true, + WorkerError: null, + CallbackError: null, + DurationMs: Math.Max(0, durationMs), + Symbols: null, + References: null); + + private static async Task SendRequestAsync(TextWriter input, string requestJson) + { + await input.WriteLineAsync(requestJson).ConfigureAwait(false); + await input.FlushAsync().ConfigureAwait(false); + } + + private static bool WaitForTask(Task task, int milliseconds, out Exception? exception) + { + try + { + if (!task.Wait(milliseconds)) + { + exception = null; + return false; + } + + exception = null; + return true; + } + catch (AggregateException ex) + { + exception = ex.GetBaseException(); + return true; + } + catch (Exception ex) + { + exception = ex; + return true; + } + } + + private static int GetRemainingWaitMilliseconds(Stopwatch stopwatch, TimeSpan callbackBudget) + { + var remainingMilliseconds = callbackBudget.TotalMilliseconds - stopwatch.Elapsed.TotalMilliseconds; + if (remainingMilliseconds <= 0) + return 0; + + return Math.Max(1, (int)Math.Ceiling(Math.Min(remainingMilliseconds, int.MaxValue))); + } + + private static string BuildWorkerExitError(Process? process, string stderr, string fallback) + { + var exitCodeText = process == null + ? "unknown" + : process.ExitCode.ToString(System.Globalization.CultureInfo.InvariantCulture); + var detail = !string.IsNullOrWhiteSpace(stderr) ? stderr.Trim() : fallback; + return $"worker exited with code {exitCodeText}: {detail}"; + } + + private static bool WaitForWorkerExit(Process process, int milliseconds) + { + try + { + return process.WaitForExit(milliseconds); + } + catch + { + return false; + } + } +} + +internal static class PostExtractionHookCallbackWorker +{ + internal const string CommandName = "__cdidx-post-extraction-hook-callback"; + internal const int WorkerKillWaitMilliseconds = 5000; + internal static readonly JsonSerializerOptions JsonOptions = new() + { + PropertyNameCaseInsensitive = true, + }; + + internal static bool TryRunCommand( + string[] args, + TextReader input, + TextWriter output, + TextWriter error, + out int exitCode) + { + if (args.Length == 0 || !StringComparer.Ordinal.Equals(args[0], CommandName)) + { + exitCode = 0; + return false; + } + + exitCode = RunCommand(args, input, output, error); + return true; + } + + internal static bool TryCreateStartInfo( + PostExtractionHookInfo hook, + out ProcessStartInfo startInfo, + out string error) + { + var runnerAssemblyPath = typeof(PostExtractionHookCallbackWorker).Assembly.Location; + if (string.IsNullOrWhiteSpace(runnerAssemblyPath)) + { + startInfo = new ProcessStartInfo(); + error = "could not resolve the cdidx assembly path for isolated hook callback execution."; + return false; + } + + startInfo = new ProcessStartInfo + { + FileName = ResolveDotnetHostPath(), + UseShellExecute = false, + RedirectStandardInput = true, + RedirectStandardOutput = true, + RedirectStandardError = true, + StandardInputEncoding = new UTF8Encoding(encoderShouldEmitUTF8Identifier: false), + StandardOutputEncoding = new UTF8Encoding(encoderShouldEmitUTF8Identifier: false), + StandardErrorEncoding = new UTF8Encoding(encoderShouldEmitUTF8Identifier: false), + CreateNoWindow = true, + }; + startInfo.ArgumentList.Add(runnerAssemblyPath); + startInfo.ArgumentList.Add(CommandName); + startInfo.ArgumentList.Add(hook.AssemblyPath); + startInfo.ArgumentList.Add(hook.TypeName); + ApplyCurrentRuntimeRollForward(startInfo); + + error = string.Empty; + return true; + } + + internal static void TryKillProcess(Process process) + { + try + { + if (!process.HasExited) + process.Kill(entireProcessTree: true); + } + catch + { + // Best effort: timeout reporting must not fail because cleanup failed. + } + + try + { + process.WaitForExit(WorkerKillWaitMilliseconds); + } + catch + { + // Best effort: the parent continues with the timeout diagnostic. + } + } + + private static int RunCommand(string[] args, TextReader input, TextWriter output, TextWriter error) + { + if (args.Length != 3) + { + error.WriteLine("post-extraction hook callback worker requires assembly path and type name."); + return 2; + } + + var hookAssemblyPath = args[1]; + var hookTypeName = args[2]; + try + { + IPostExtractionHook? hook = null; + string? requestJson; + while ((requestJson = input.ReadLine()) != null) + { + WorkerResponse response; + try + { + var request = JsonSerializer.Deserialize(requestJson, JsonOptions) + ?? throw new InvalidOperationException("worker request was empty."); + hook ??= CreateHook(hookAssemblyPath, hookTypeName); + response = InvokeInsideWorker(hook, request); + } + catch (Exception ex) + { + response = new WorkerResponse(null, null, null, ex.Message); + } + + output.WriteLine(JsonSerializer.Serialize(response, JsonOptions)); + output.Flush(); + } + + return 0; + } + catch (Exception ex) + { + error.WriteLine(ex.Message); + return 1; + } + } + + private static IPostExtractionHook CreateHook(string hookAssemblyPath, string hookTypeName) + { + var assembly = AssemblyLoadContext.Default.LoadFromAssemblyPath(Path.GetFullPath(hookAssemblyPath)); + var type = assembly.GetType(hookTypeName, throwOnError: true) + ?? throw new InvalidOperationException($"hook type `{hookTypeName}` was not found."); + return Activator.CreateInstance(type) as IPostExtractionHook + ?? throw new InvalidOperationException($"hook type `{hookTypeName}` could not be instantiated as `{nameof(IPostExtractionHook)}`."); + } + + private static WorkerResponse InvokeInsideWorker(IPostExtractionHook hook, WorkerRequest request) + { + var originalOut = Console.Out; + var originalError = Console.Error; + using var capturedOut = new StringWriter(); + using var capturedError = new StringWriter(); + Exception? callbackFailure = null; + try + { + Console.SetOut(capturedOut); + Console.SetError(capturedError); + if (request.Callback == nameof(IPostExtractionHook.OnSymbolsExtracted)) + { + if (request.Symbols == null) + throw new InvalidOperationException("symbol callback request omitted symbols."); + hook.OnSymbolsExtracted(request.Context, request.Symbols); + } + else if (request.Callback == nameof(IPostExtractionHook.OnReferencesExtracted)) + { + if (request.References == null) + throw new InvalidOperationException("reference callback request omitted references."); + hook.OnReferencesExtracted(request.Context, request.References); + } + else + { + throw new InvalidOperationException($"unknown hook callback `{request.Callback}`."); + } + } + catch (Exception ex) + { + callbackFailure = ex is TargetInvocationException { InnerException: not null } ? ex.InnerException : ex; + } + finally + { + Console.SetOut(originalOut); + Console.SetError(originalError); + } + + return new WorkerResponse(request.Symbols, request.References, callbackFailure?.Message, null); + } + + private static string ResolveDotnetHostPath() + { + var dotnetHostPath = Environment.GetEnvironmentVariable("DOTNET_HOST_PATH"); + if (!string.IsNullOrWhiteSpace(dotnetHostPath)) + return dotnetHostPath; + + var processPath = Environment.ProcessPath; + if (!string.IsNullOrWhiteSpace(processPath) + && string.Equals(Path.GetFileNameWithoutExtension(processPath), "dotnet", StringComparison.OrdinalIgnoreCase)) + { + return processPath; + } + + return "dotnet"; + } + + private static void ApplyCurrentRuntimeRollForward(ProcessStartInfo startInfo) + { + var targetMajor = GetRunnerTargetFrameworkMajor(); + if (targetMajor.HasValue && Environment.Version.Major > targetMajor.Value) + startInfo.Environment["DOTNET_ROLL_FORWARD"] = "LatestMajor"; + } + + private static int? GetRunnerTargetFrameworkMajor() + { + var frameworkName = typeof(PostExtractionHookCallbackWorker) + .Assembly + .GetCustomAttribute() + ?.FrameworkName; + if (string.IsNullOrWhiteSpace(frameworkName)) + return null; + + const string versionPrefix = "Version=v"; + var versionIndex = frameworkName.IndexOf(versionPrefix, StringComparison.OrdinalIgnoreCase); + if (versionIndex < 0) + return null; + + var majorStart = versionIndex + versionPrefix.Length; + var majorEnd = frameworkName.IndexOf('.', majorStart); + var majorText = majorEnd < 0 + ? frameworkName[majorStart..] + : frameworkName[majorStart..majorEnd]; + return int.TryParse( + majorText, + System.Globalization.NumberStyles.Integer, + System.Globalization.CultureInfo.InvariantCulture, + out var major) + ? major + : null; + } + + internal sealed record WorkerRequest( + string Callback, + FileContext Context, + List? Symbols, + List? References); + + internal sealed record WorkerResponse( + List? Symbols, + List? References, + string? CallbackError, + string? WorkerError); +} diff --git a/src/CodeIndex/Indexer/Hooks/PostExtractionHooks.cs b/src/CodeIndex/Indexer/Hooks/PostExtractionHooks.cs index 30f2180fac..94297bb133 100644 --- a/src/CodeIndex/Indexer/Hooks/PostExtractionHooks.cs +++ b/src/CodeIndex/Indexer/Hooks/PostExtractionHooks.cs @@ -125,13 +125,17 @@ public static PostExtractionHookRunner Discover(string? hooksDirectory) try { - if (Activator.CreateInstance(type) is not IPostExtractionHook hook) + if (type.GetConstructor(Type.EmptyTypes) == null) + { + runner.EnqueueDiagnostic(dllPath, type.FullName, "Failed to instantiate hook: public parameterless constructor not found."); continue; + } + var info = new PostExtractionHookInfo(type.Name, Path.GetFullPath(dllPath), type.FullName ?? type.Name); loaded.Add(new LoadedPostExtractionHook( - hook, - new PostExtractionHookInfo(type.Name, Path.GetFullPath(dllPath), type.FullName ?? type.Name), - AssemblyLoadContext.GetLoadContext(type.Assembly))); + info, + AssemblyLoadContext.GetLoadContext(type.Assembly), + new PostExtractionHookCallbackWorkerClient(info))); } catch (Exception) { @@ -277,8 +281,11 @@ public void OnSymbolsExtracted(FileContext context, IList symbols) var workingSymbols = CloneSymbols(symbols); if (InvokeHookWithBudget( hook, + PostExtractionHookCallbackKind.Symbols, nameof(IPostExtractionHook.OnSymbolsExtracted), - () => hook.Instance.OnSymbolsExtracted(context, workingSymbols))) + context, + workingSymbols, + null)) { ReplaceList(symbols, workingSymbols); } @@ -294,59 +301,72 @@ public void OnReferencesExtracted(FileContext context, IList re var workingReferences = CloneReferences(references); if (InvokeHookWithBudget( hook, + PostExtractionHookCallbackKind.References, nameof(IPostExtractionHook.OnReferencesExtracted), - () => hook.Instance.OnReferencesExtracted(context, workingReferences))) + context, + null, + workingReferences)) { ReplaceList(references, workingReferences); } } } - private bool InvokeHookWithBudget(LoadedPostExtractionHook hook, string callback, Action invoke) + private bool InvokeHookWithBudget( + LoadedPostExtractionHook hook, + PostExtractionHookCallbackKind kind, + string callback, + FileContext context, + List? symbols, + List? references) { if (disabledHooks.ContainsKey(hook.Info.TypeName)) return false; - var stopwatch = Stopwatch.StartNew(); - Exception? failure = null; - var task = Task.Run(() => + var result = hook.Worker.Invoke( + kind, + callback, + context, + symbols, + references, + callbackBudget); + if (result.TimedOut) { - try - { - lock (hook.Instance) - invoke(); - } - catch (Exception ex) - { - failure = ex; - } - }); + disabledHooks.TryAdd(hook.Info.TypeName, 0); + EnqueueDiagnostic( + hook.Info.AssemblyPath, + hook.Info.TypeName, + $"{callback} exceeded the {callbackBudget.TotalMilliseconds:0} ms callback budget; hook disabled for this index run.", + callback, + result.DurationMs); + return false; + } - if (!task.Wait(callbackBudget)) + if (!result.Success) { - stopwatch.Stop(); - var timeoutDurationMs = Math.Max( - stopwatch.ElapsedMilliseconds, - (long)Math.Ceiling(callbackBudget.TotalMilliseconds)); disabledHooks.TryAdd(hook.Info.TypeName, 0); EnqueueDiagnostic( hook.Info.AssemblyPath, hook.Info.TypeName, - $"{callback} exceeded the {callbackBudget.TotalMilliseconds:0} ms callback budget; hook disabled for this index run.", + $"{callback} failed in isolated worker.", callback, - timeoutDurationMs); + result.DurationMs); return false; } - stopwatch.Stop(); - if (failure != null) + if (result.Symbols != null && symbols != null) + ReplaceList(symbols, result.Symbols); + if (result.References != null && references != null) + ReplaceList(references, result.References); + + if (result.CallbackError != null) { EnqueueDiagnostic( hook.Info.AssemblyPath, hook.Info.TypeName, $"{callback} failed.", callback, - stopwatch.ElapsedMilliseconds); + result.DurationMs); } return true; @@ -482,6 +502,11 @@ public void Dispose() return; disposed = true; + foreach (var hook in hooks) + { + hook.Worker.Dispose(); + } + var loadContexts = hooks .Select(hook => hook.LoadContext) .Where(loadContext => loadContext is { IsCollectible: true }) @@ -496,7 +521,7 @@ public void Dispose() } private sealed record LoadedPostExtractionHook( - IPostExtractionHook Instance, PostExtractionHookInfo Info, - AssemblyLoadContext? LoadContext); + AssemblyLoadContext? LoadContext, + PostExtractionHookCallbackWorkerClient Worker); } diff --git a/src/CodeIndex/Mcp/McpToolHandlers.cs b/src/CodeIndex/Mcp/McpToolHandlers.cs index ae247af257..78e09abfd3 100644 --- a/src/CodeIndex/Mcp/McpToolHandlers.cs +++ b/src/CodeIndex/Mcp/McpToolHandlers.cs @@ -3684,7 +3684,7 @@ private async Task ExecuteIndexAsync(JsonNode? id, JsonNode? args, Jso var writer = new DbWriter(db); var indexer = new FileIndexer(projectPath, GitHelper.ResolveIgnoreCase(projectPath), GitHelper.TryGetRepositoryRoot(projectPath) ?? Path.GetFullPath(projectPath), maxFileBytes); - var postExtractionHooks = PostExtractionHookRunner.DiscoverDefault(); + using var postExtractionHooks = PostExtractionHookRunner.DiscoverDefault(); var requestToken = _currentRequestToken.Value; requestToken.ThrowIfCancellationRequested(); var currentHotspotFamilyMarkerFingerprints = GetHotspotFamilyMarkerFingerprints(indexer, requestToken); diff --git a/tests/CodeIndex.Tests/PostExtractionHookTests.cs b/tests/CodeIndex.Tests/PostExtractionHookTests.cs index 49cfd50962..72a03bb3ef 100644 --- a/tests/CodeIndex.Tests/PostExtractionHookTests.cs +++ b/tests/CodeIndex.Tests/PostExtractionHookTests.cs @@ -9,6 +9,9 @@ public class PostExtractionHookTests { internal const string SlowHookDelayEnvironmentVariable = "CDIDX_TEST_SLOW_POST_EXTRACTION_HOOK_MS"; internal const string SlowHookCompletionPathEnvironmentVariable = "CDIDX_TEST_SLOW_POST_EXTRACTION_HOOK_DONE_PATH"; + internal const string SlowConstructorHookDelayEnvironmentVariable = "CDIDX_TEST_SLOW_CTOR_POST_EXTRACTION_HOOK_MS"; + internal const string StatefulHookEnvironmentVariable = "CDIDX_TEST_STATEFUL_POST_EXTRACTION_HOOK"; + internal const string ThrowingConstructorHookEnvironmentVariable = "CDIDX_TEST_THROWING_CTOR_POST_EXTRACTION_HOOK"; [Fact] public void Discover_LoadsHooksAndAllowsSymbolAndReferenceMutation() @@ -78,7 +81,79 @@ public void CallbackExceptions_AreDiagnosticsAndDoNotBlockOtherHooks() } [Fact] - public void CallbackBudgetExceeded_AddsDiagnosticAndSkipsTimedOutMutation() + public void WorkerConstructionFailure_DisablesHookForCurrentRun() + { + var projectRoot = TestProjectHelper.CreateTempProject("post-extraction-hook-ctor-failure"); + lock (TestConsoleLock.Gate) + { + using var env = EnvironmentVariableScope.Capture(ThrowingConstructorHookEnvironmentVariable); + try + { + env.Set(ThrowingConstructorHookEnvironmentVariable, "1"); + var hooksDir = Path.Combine(projectRoot, "hooks"); + Directory.CreateDirectory(hooksDir); + File.Copy(Assembly.GetExecutingAssembly().Location, Path.Combine(hooksDir, "CodeIndex.Tests.dll")); + + { + using var runner = PostExtractionHookRunner.Discover(hooksDir); + var context = new FileContext(projectRoot, "src/App.cs", Path.Combine(projectRoot, "src", "App.cs"), "csharp"); + var symbols = new List(); + var references = new List(); + + runner.OnSymbolsExtracted(context, symbols); + runner.OnReferencesExtracted(context, references); + + var diagnostic = Assert.Single( + runner.Diagnostics, + diagnostic => diagnostic.TypeName == typeof(ThrowingConstructorPostExtractionHook).FullName); + Assert.Contains("isolated worker", diagnostic.Message, StringComparison.Ordinal); + Assert.DoesNotContain("ctor boom", diagnostic.Message, StringComparison.Ordinal); + } + CollectUnloadedHookAssemblies(); + } + finally + { + TestProjectHelper.DeleteDirectory(projectRoot); + } + } + } + + [Fact] + public void Callbacks_ReuseIsolatedWorkerHookInstance() + { + var projectRoot = TestProjectHelper.CreateTempProject("post-extraction-hook-state"); + lock (TestConsoleLock.Gate) + { + using var env = EnvironmentVariableScope.Capture(StatefulHookEnvironmentVariable); + try + { + env.Set(StatefulHookEnvironmentVariable, "1"); + var hooksDir = Path.Combine(projectRoot, "hooks"); + Directory.CreateDirectory(hooksDir); + File.Copy(Assembly.GetExecutingAssembly().Location, Path.Combine(hooksDir, "CodeIndex.Tests.dll")); + + { + using var runner = PostExtractionHookRunner.Discover(hooksDir); + var context = new FileContext(projectRoot, "src/App.cs", Path.Combine(projectRoot, "src", "App.cs"), "csharp"); + var symbols = new List(); + var references = new List(); + + runner.OnSymbolsExtracted(context, symbols); + runner.OnReferencesExtracted(context, references); + + Assert.Contains(references, reference => reference.SymbolName == "StatefulHookSawSymbols"); + } + CollectUnloadedHookAssemblies(); + } + finally + { + TestProjectHelper.DeleteDirectory(projectRoot); + } + } + } + + [Fact] + public void CallbackBudgetExceeded_KillsWorkerAndSkipsTimedOutMutation() { var projectRoot = TestProjectHelper.CreateTempProject("post-extraction-hook-budget"); lock (TestConsoleLock.Gate) @@ -103,16 +178,18 @@ public void CallbackBudgetExceeded_AddsDiagnosticAndSkipsTimedOutMutation() var symbols = new List(); runner.OnSymbolsExtracted(context, symbols); - WaitForSlowHookCompletion(completionPath); + AssertFileDoesNotAppear(completionPath, TimeSpan.FromMilliseconds(750)); Assert.DoesNotContain(symbols, symbol => symbol.Name == "SlowHookTag"); var diagnostic = Assert.Single( runner.Diagnostics, item => item.TypeName == typeof(SlowPostExtractionHook).FullName && item.Callback == nameof(IPostExtractionHook.OnSymbolsExtracted)); - Assert.Contains("exceeded", diagnostic.Message, StringComparison.Ordinal); - // Task.Wait can time out at the budget boundary before ElapsedMilliseconds - // rounds up to the full budget on some CI hosts. + Assert.True( + diagnostic.Message.Contains("exceeded", StringComparison.Ordinal), + diagnostic.Message); + // The worker wait can time out at the budget boundary before + // ElapsedMilliseconds rounds up to the full budget on some CI hosts. Assert.True(diagnostic.DurationMs > 0); Assert.Equal(50, (long)Math.Round(runner.CallbackBudget.TotalMilliseconds, MidpointRounding.AwayFromZero)); } @@ -126,6 +203,58 @@ public void CallbackBudgetExceeded_AddsDiagnosticAndSkipsTimedOutMutation() } } + [Fact] + public void CallbackBudgetExceeded_KillsSlowConstructorAfterLargeRequestIsSent() + { + var projectRoot = TestProjectHelper.CreateTempProject("post-extraction-hook-slow-ctor"); + lock (TestConsoleLock.Gate) + { + using var env = EnvironmentVariableScope.Capture(SlowConstructorHookDelayEnvironmentVariable); + var originalBudget = PostExtractionHookRunner.CallbackBudgetForTesting; + try + { + env.Set(SlowConstructorHookDelayEnvironmentVariable, "200"); + PostExtractionHookRunner.CallbackBudgetForTesting = () => TimeSpan.FromMilliseconds(50); + var hooksDir = Path.Combine(projectRoot, "hooks"); + Directory.CreateDirectory(hooksDir); + File.Copy(Assembly.GetExecutingAssembly().Location, Path.Combine(hooksDir, "CodeIndex.Tests.dll")); + + { + using var runner = PostExtractionHookRunner.Discover(hooksDir); + var context = new FileContext(projectRoot, "src/App.cs", Path.Combine(projectRoot, "src", "App.cs"), "csharp"); + var symbols = Enumerable + .Range(0, 1000) + .Select(index => new SymbolRecord + { + FileId = 10, + Kind = "function", + Name = $"LargePayloadSymbol{index}", + Line = index + 1, + StartLine = index + 1, + EndLine = index + 1, + Signature = new string('x', 512), + }) + .ToList(); + + runner.OnSymbolsExtracted(context, symbols); + + var diagnostic = Assert.Single( + runner.Diagnostics, + item => item.TypeName == typeof(SlowConstructorPostExtractionHook).FullName + && item.Callback == nameof(IPostExtractionHook.OnSymbolsExtracted)); + Assert.Contains("exceeded", diagnostic.Message, StringComparison.Ordinal); + Assert.True(diagnostic.DurationMs > 0); + } + CollectUnloadedHookAssemblies(); + } + finally + { + PostExtractionHookRunner.CallbackBudgetForTesting = originalBudget; + TestProjectHelper.DeleteDirectory(projectRoot); + } + } + } + [Fact] public void CallbackBudget_NormalizesInvalidAndTooLargeValues() { @@ -230,13 +359,13 @@ private static void CollectUnloadedHookAssemblies() GC.Collect(); } - private static void WaitForSlowHookCompletion(string completionPath) + private static void AssertFileDoesNotAppear(string path, TimeSpan duration) { - var deadline = DateTimeOffset.UtcNow.AddSeconds(5); - while (!File.Exists(completionPath)) + var deadline = DateTimeOffset.UtcNow.Add(duration); + while (DateTimeOffset.UtcNow < deadline) { - if (DateTimeOffset.UtcNow >= deadline) - throw new TimeoutException("Timed out waiting for the slow post-extraction hook to finish."); + if (File.Exists(path)) + throw new InvalidOperationException("The timed-out post-extraction hook continued running after the callback returned."); Thread.Sleep(25); } @@ -282,6 +411,67 @@ public void OnReferencesExtracted(FileContext context, IList re => throw new InvalidOperationException("boom"); } +public sealed class ThrowingConstructorPostExtractionHook : IPostExtractionHook +{ + public ThrowingConstructorPostExtractionHook() + { + if (Environment.GetEnvironmentVariable(PostExtractionHookTests.ThrowingConstructorHookEnvironmentVariable) == "1") + throw new InvalidOperationException("ctor boom"); + } + + public void OnSymbolsExtracted(FileContext context, IList symbols) + { + } + + public void OnReferencesExtracted(FileContext context, IList references) + { + } +} + +public sealed class SlowConstructorPostExtractionHook : IPostExtractionHook +{ + public SlowConstructorPostExtractionHook() + { + var raw = Environment.GetEnvironmentVariable(PostExtractionHookTests.SlowConstructorHookDelayEnvironmentVariable); + if (int.TryParse(raw, out var milliseconds) && milliseconds > 0) + Thread.Sleep(milliseconds); + } + + public void OnSymbolsExtracted(FileContext context, IList symbols) + { + } + + public void OnReferencesExtracted(FileContext context, IList references) + { + } +} + +public sealed class StatefulPostExtractionHook : IPostExtractionHook +{ + private bool sawSymbols; + + public void OnSymbolsExtracted(FileContext context, IList symbols) + { + if (Environment.GetEnvironmentVariable(PostExtractionHookTests.StatefulHookEnvironmentVariable) == "1") + sawSymbols = true; + } + + public void OnReferencesExtracted(FileContext context, IList references) + { + if (!sawSymbols || Environment.GetEnvironmentVariable(PostExtractionHookTests.StatefulHookEnvironmentVariable) != "1") + return; + + references.Add(new ReferenceRecord + { + SymbolName = "StatefulHookSawSymbols", + ReferenceKind = "domain_reference", + Line = 1, + Column = 1, + Context = context.Path, + }); + } +} + public sealed class SlowPostExtractionHook : IPostExtractionHook { public void OnSymbolsExtracted(FileContext context, IList symbols)