From 7bf622466a46caa241dfbb47d26389ee013fe0f0 Mon Sep 17 00:00:00 2001 From: minyitang Date: Tue, 25 Aug 2026 11:12:21 -0700 Subject: [PATCH] Fix .NET FFI callback lifetime Use tracked runtime connections and drain callbacks before releasing managed callback state. Preserve force-stop latency by retaining callback state when an immediate drain cannot complete. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b659bb88-f59c-49d3-afa7-c1b38d575f4b --- dotnet/src/Client.cs | 12 +++- dotnet/src/FfiRuntimeHost.cs | 94 +++++++++++++++++++++++-------- dotnet/test/E2E/ClientE2ETests.cs | 12 ++++ 3 files changed, 95 insertions(+), 23 deletions(-) diff --git a/dotnet/src/Client.cs b/dotnet/src/Client.cs index 86178ba74b..7b9161adbe 100644 --- a/dotnet/src/Client.cs +++ b/dotnet/src/Client.cs @@ -663,7 +663,17 @@ or IOException if (ctx.FfiHost is { } ffiHost) { - try { ffiHost.Dispose(); } + try + { + if (gracefulRuntimeShutdown) + { + ffiHost.Dispose(); + } + else + { + ffiHost.ForceDispose(); + } + } catch (Exception ex) { AddCleanupError(errors, ex, _logger); } _ffiHost = null; } diff --git a/dotnet/src/FfiRuntimeHost.cs b/dotnet/src/FfiRuntimeHost.cs index a838b9fd17..be74706710 100644 --- a/dotnet/src/FfiRuntimeHost.cs +++ b/dotnet/src/FfiRuntimeHost.cs @@ -39,6 +39,7 @@ internal sealed partial class FfiRuntimeHost : IDisposable { /// Logical name the native interop layer binds the cdylib to. private const string LibraryName = "copilot_runtime"; + private const uint ConnectionDrainTimeoutMilliseconds = 30_000; private readonly ILogger _logger; private readonly string _cliEntrypoint; @@ -53,6 +54,10 @@ internal sealed partial class FfiRuntimeHost : IDisposable private uint _connectionId; private bool _disposed; + // Roots this host while native code can invoke its outbound callback. It is + // released only after connection_close_and_wait confirms callbacks drained. + private GCHandle _selfHandle; + private FfiRuntimeHost(string libraryPath, string cliEntrypoint, IReadOnlyDictionary? environment, IReadOnlyList args, ILogger logger) { _libraryPath = libraryPath; @@ -224,7 +229,13 @@ private void FeedInbound(IntPtr bytesPtr, UIntPtr bytesLen) _receiveStream.Feed(buffer); } - public void Dispose() + public void Dispose() => + Dispose(ConnectionDrainTimeoutMilliseconds, throwOnDrainFailure: true); + + internal void ForceDispose() => + Dispose(timeoutMilliseconds: 0, throwOnDrainFailure: false); + + private void Dispose(uint timeoutMilliseconds, bool throwOnDrainFailure) { if (_disposed) { @@ -232,17 +243,31 @@ public void Dispose() } _disposed = true; + Exception? connectionCloseError = null; + var callbackDrained = _connectionId == 0; try { if (_connectionId != 0) { - NativeConnectionClose(_connectionId); + callbackDrained = NativeConnectionCloseAndWait( + _connectionId, + timeoutMilliseconds); + if (!callbackDrained) + { + connectionCloseError = new TimeoutException( + $"FfiRuntimeHost timed out after {timeoutMilliseconds} ms " + + "waiting for runtime callbacks to drain; callback state was retained."); + } _connectionId = 0; } } catch (Exception ex) { - _logger.LogDebug(ex, "FfiRuntimeHost: connection_close failed"); + connectionCloseError = new InvalidOperationException( + "FfiRuntimeHost failed to close and drain the runtime connection; " + + "callback state was retained.", + ex); + _connectionId = 0; } try @@ -259,7 +284,21 @@ public void Dispose() } _receiveStream.Complete(); - DisposeNativeCallback(); + if (callbackDrained) + { + DisposeNativeCallback(); + } + + if (connectionCloseError is not null) + { + if (throwOnDrainFailure) + { + throw connectionCloseError; + } + _logger.LogWarning( + connectionCloseError, + "FfiRuntimeHost force-closed before runtime callbacks drained; callback state was retained"); + } } /// Length as the native pointer-sized unsigned integer the ABI expects. @@ -272,10 +311,6 @@ public void Dispose() private static bool s_resolverRegistered; private static string? s_resolvedLibraryPath; - // A normal (non-pinned) handle to this instance, passed to the native side as - // the callback's user_data so the static outbound callback can route back here. - private GCHandle _selfHandle; - /// /// Registers (once) a process-wide /// that maps to the absolute runtime.node path so the @@ -332,7 +367,8 @@ private uint NativeOpenConnection(uint serverId) private static bool NativeConnectionWrite(uint connectionId, ReadOnlySpan frame) => ConnectionWrite(connectionId, frame, Len(frame.Length)); - private static bool NativeConnectionClose(uint connectionId) => ConnectionClose(connectionId); + private static bool NativeConnectionCloseAndWait(uint connectionId, uint timeoutMilliseconds) => + ConnectionCloseAndWait(connectionId, timeoutMilliseconds); private void DisposeNativeCallback() { @@ -366,7 +402,7 @@ private static partial uint HostStart( [return: MarshalAs(UnmanagedType.U1)] private static partial bool HostShutdown(uint serverId); - [LibraryImport(LibraryName, EntryPoint = "copilot_runtime_connection_open")] + [LibraryImport(LibraryName, EntryPoint = "copilot_runtime_connection_open_tracked")] [UnmanagedCallConv(CallConvs = new[] { typeof(CallConvCdecl) })] private static unsafe partial uint ConnectionOpen( uint serverId, @@ -381,10 +417,10 @@ private static unsafe partial uint ConnectionOpen( [return: MarshalAs(UnmanagedType.U1)] private static partial bool ConnectionWrite(uint connectionId, ReadOnlySpan bytes, nuint bytesLen); - [LibraryImport(LibraryName, EntryPoint = "copilot_runtime_connection_close")] + [LibraryImport(LibraryName, EntryPoint = "copilot_runtime_connection_close_and_wait")] [UnmanagedCallConv(CallConvs = new[] { typeof(CallConvCdecl) })] [return: MarshalAs(UnmanagedType.U1)] - private static partial bool ConnectionClose(uint connectionId); + private static partial bool ConnectionCloseAndWait(uint connectionId, uint timeoutMilliseconds); #else // ---- Legacy interop: delegate-based P/Invoke for netstandard2.0 ---- // netstandard2.0 has neither LibraryImport, NativeLibrary, nor UnmanagedCallersOnly, @@ -416,7 +452,7 @@ private delegate uint ConnectionOpenDelegate( [UnmanagedFunctionPointer(CallingConvention.Cdecl)] [return: MarshalAs(UnmanagedType.U1)] - private delegate bool ConnectionCloseDelegate(uint connectionId); + private delegate bool ConnectionCloseAndWaitDelegate(uint connectionId, uint timeoutMilliseconds); [UnmanagedFunctionPointer(CallingConvention.Cdecl)] private delegate void OutboundCallbackDelegate(IntPtr userData, IntPtr bytesPtr, UIntPtr bytesLen); @@ -428,7 +464,7 @@ private delegate uint ConnectionOpenDelegate( private static HostShutdownDelegate? s_hostShutdown; private static ConnectionOpenDelegate? s_connectionOpen; private static ConnectionWriteDelegate? s_connectionWrite; - private static ConnectionCloseDelegate? s_connectionClose; + private static ConnectionCloseAndWaitDelegate? s_connectionCloseAndWait; // Held for the connection's lifetime so the marshaled function pointer handed to the // native side is not collected while Rust may still invoke it. @@ -457,9 +493,11 @@ private static void PrepareNativeLibrary(string libraryPath) s_hostStart = Bind(handle, "copilot_runtime_host_start"); s_hostShutdown = Bind(handle, "copilot_runtime_host_shutdown"); - s_connectionOpen = Bind(handle, "copilot_runtime_connection_open"); + s_connectionOpen = Bind(handle, "copilot_runtime_connection_open_tracked"); s_connectionWrite = Bind(handle, "copilot_runtime_connection_write"); - s_connectionClose = Bind(handle, "copilot_runtime_connection_close"); + s_connectionCloseAndWait = Bind( + handle, + "copilot_runtime_connection_close_and_wait"); s_loaded = true; s_loadedPath = libraryPath; } @@ -480,11 +518,12 @@ private static uint NativeHostStart(byte[] argvJson, byte[]? env) => private uint NativeOpenConnection(uint serverId) { + _selfHandle = GCHandle.Alloc(this); _outboundDelegate = OnOutbound; return s_connectionOpen!( serverId, _outboundDelegate, - IntPtr.Zero, + GCHandle.ToIntPtr(_selfHandle), null, UIntPtr.Zero, null, UIntPtr.Zero, null, UIntPtr.Zero); @@ -500,17 +539,28 @@ private static unsafe bool NativeConnectionWrite(uint connectionId, ReadOnlySpan } } - private static bool NativeConnectionClose(uint connectionId) => s_connectionClose!(connectionId); + private static bool NativeConnectionCloseAndWait(uint connectionId, uint timeoutMilliseconds) => + s_connectionCloseAndWait!(connectionId, timeoutMilliseconds); - private void DisposeNativeCallback() => _outboundDelegate = null; + private void DisposeNativeCallback() + { + _outboundDelegate = null; + if (_selfHandle.IsAllocated) + { + _selfHandle.Free(); + } + } - private void OnOutbound(IntPtr userData, IntPtr bytesPtr, UIntPtr bytesLen) + private static void OnOutbound(IntPtr userData, IntPtr bytesPtr, UIntPtr bytesLen) { - if (bytesPtr == IntPtr.Zero || bytesLen == UIntPtr.Zero) + if (userData == IntPtr.Zero || bytesPtr == IntPtr.Zero || bytesLen == UIntPtr.Zero) { return; } - FeedInbound(bytesPtr, bytesLen); + if (GCHandle.FromIntPtr(userData).Target is FfiRuntimeHost self) + { + self.FeedInbound(bytesPtr, bytesLen); + } } /// diff --git a/dotnet/test/E2E/ClientE2ETests.cs b/dotnet/test/E2E/ClientE2ETests.cs index b6bdfd90fd..401dd24d4f 100644 --- a/dotnet/test/E2E/ClientE2ETests.cs +++ b/dotnet/test/E2E/ClientE2ETests.cs @@ -64,6 +64,18 @@ public async Task Should_Start_And_Connect_Over_InProcess_Ffi() } } + [Fact] + public async Task Should_Force_Stop_Over_InProcess_Ffi() + { + using var client = new CopilotClient(new CopilotClientOptions + { + Connection = RuntimeConnection.ForInProcess(), + }); + + await client.StartAsync(); + await client.ForceStopAsync(); + } + [Theory] [InlineData(true)] // stdio transport [InlineData(false)] // TCP transport