Fix in-process callback reclamation after connection close

dotnet/src/FfiRuntimeHost.cs

@@ -35,29 +35,60 @@
 /// </remarks>
 internal sealed partial class FfiRuntimeHost : IDisposable
 {
+    private enum NativeCleanupResult
+    {
+        Complete,
+        Retry,
+        Failed,
+    }
+
     /// <summary>Logical name the native interop layer binds the cdylib to.</summary>
     private const string LibraryName = "copilot_runtime";
+    private const int CleanupRetryDelayMilliseconds = 100;
+    private static readonly object QuarantineLock = new();
+    private static readonly HashSet<FfiRuntimeHost> QuarantinedHosts = [];
 
     private readonly ILogger _logger;
     private readonly string? _cliEntrypoint;
     private readonly string _libraryPath;
     private readonly IReadOnlyDictionary<string, string>? _environment;
     private readonly IReadOnlyList<string> _args;
+    private readonly Func<uint, bool> _connectionClose;
+    private readonly Func<uint, bool> _hostShutdown;
+    private readonly Action _releaseNativeCallback;
+    private readonly object _lifecycleLock = new();
 
     private readonly CallbackReceiveStream _receiveStream = new();
     private CallbackSendStream? _sendStream;
 
     private uint _serverId;
     private uint _connectionId;
     private bool _disposed;
+    private bool _cleanupRetryScheduled;
 
     private FfiRuntimeHost(string libraryPath, string? cliEntrypoint, IReadOnlyDictionary<string, string>? environment, IReadOnlyList<string> args, ILogger logger)
+        : this(libraryPath, cliEntrypoint, environment, args, logger, null, null, null)
+    {
+    }
+
+    private FfiRuntimeHost(
+        string libraryPath,
+        string? cliEntrypoint,
+        IReadOnlyDictionary<string, string>? environment,
+        IReadOnlyList<string> args,
+        ILogger logger,
+        Func<uint, bool>? connectionClose,
+        Func<uint, bool>? hostShutdown,
+        Action? releaseNativeCallback)
     {
         _libraryPath = libraryPath;
         _cliEntrypoint = cliEntrypoint;
         _environment = environment;
         _args = args;
         _logger = logger;
+        _connectionClose = connectionClose ?? NativeConnectionClose;
+        _hostShutdown = hostShutdown ?? NativeHostShutdown;
+        _releaseNativeCallback = releaseNativeCallback ?? DisposeNativeCallback;
     }
 
     /// <summary>The stream JSON-RPC reads server→client frames from.</summary>
@@ -109,23 +140,36 @@
             var argvJson = BuildArgvJson(_cliEntrypoint, _args);
             var envJson = BuildEnvJson(_environment);
 
-            _serverId = NativeHostStart(argvJson, envJson);
-            if (_serverId == 0)
+            var serverId = NativeHostStart(argvJson, envJson);
+            if (serverId == 0)
             {
                 throw new InvalidOperationException(
                     $"copilot_runtime_host_start failed (library '{_libraryPath}').");
             }
 
-            _connectionId = NativeOpenConnection(_serverId);
-            if (_connectionId == 0)
+            var connectionId = NativeOpenConnection(serverId);
+            if (connectionId == 0)
             {
-                DisposeNativeCallback();
-                NativeHostShutdown(_serverId);
-                _serverId = 0;
+                _releaseNativeCallback();
+                NativeHostShutdown(serverId);
                 throw new InvalidOperationException("copilot_runtime_connection_open failed.");
             }
 
-            _sendStream = new CallbackSendStream(SendFrame);
+            lock (_lifecycleLock)
+            {
+                _serverId = serverId;
+                _connectionId = connectionId;
+                _sendStream = new CallbackSendStream(SendFrame);
+                if (_disposed)
+                {
+                    if (TryFinalizeNativeCleanup() == NativeCleanupResult.Retry)
+                    {
+                        ScheduleNativeCleanupRetry();
+                    }
+                    throw new InvalidOperationException(
+                        "FfiRuntimeHost was disposed during startup.");
+                }
+            }
         }, cancellationToken).ConfigureAwait(false);
 
         if (_logger.IsEnabled(LogLevel.Debug))
@@ -189,11 +233,18 @@
 
     private bool SendFrame(ReadOnlySpan<byte> frame)
     {
-        if (_disposed || _connectionId == 0)
+        if (Volatile.Read(ref _disposed))
         {
             return false;
         }
-        return NativeConnectionWrite(_connectionId, frame);
+        lock (_lifecycleLock)
+        {
+            if (_disposed || _connectionId == 0)
+            {
+                return false;
+            }
+            return NativeConnectionWrite(_connectionId, frame);
+        }
     }
 
     private void FeedInbound(IntPtr bytesPtr, UIntPtr bytesLen)
@@ -206,40 +257,100 @@
 
     public void Dispose()
     {
-        if (_disposed)
+        lock (_lifecycleLock)
         {
-            return;
+            if (_disposed)
+            {
+                return;
+            }
+            _disposed = true;
         }
-        _disposed = true;
 
-        try
+        _receiveStream.Complete();
+
+        lock (_lifecycleLock)
         {
-            if (_connectionId != 0)
+            if (TryFinalizeNativeCleanup() == NativeCleanupResult.Retry)
             {
-                NativeConnectionClose(_connectionId);
-                _connectionId = 0;
+                ScheduleNativeCleanupRetry();
             }
         }
-        catch (Exception ex)
+    }
+
+    private NativeCleanupResult TryFinalizeNativeCleanup()
+    {
+        if (_connectionId != 0)
         {
-            _logger.LogDebug(ex, "FfiRuntimeHost: connection_close failed");
+            bool closed;
+            try
+            {
+                closed = _connectionClose(_connectionId);
+            }
+            catch (Exception ex)
+            {
+                _logger.LogDebug(ex, "FfiRuntimeHost: connection_close failed");
+                lock (QuarantineLock)
+                {
+                    QuarantinedHosts.Add(this);
+                }
+                return NativeCleanupResult.Failed;
+            }
+            if (!closed)
+            {
+                return NativeCleanupResult.Retry;
+            }
+
+            _connectionId = 0;
+            _releaseNativeCallback();
         }
 
-        try
+        if (_serverId != 0)
         {
-            if (_serverId != 0)
+            try
             {
-                NativeHostShutdown(_serverId);
-                _serverId = 0;
+                if (!_hostShutdown(_serverId) && _logger.IsEnabled(LogLevel.Debug))
+                {
+                    _logger.LogDebug(
+                        "FfiRuntimeHost: host_shutdown did not recognize server {ServerId}",
+                        _serverId);
+                }
             }
+            catch (Exception ex)
+            {
+                _logger.LogDebug(ex, "FfiRuntimeHost: host_shutdown failed");
+            }
+
+            _serverId = 0;
         }
-        catch (Exception ex)
+
+        return NativeCleanupResult.Complete;
+    }
+
+    private void ScheduleNativeCleanupRetry()
+    {
+        if (_cleanupRetryScheduled)
         {
-            _logger.LogDebug(ex, "FfiRuntimeHost: host_shutdown failed");
+            return;
         }
+        _cleanupRetryScheduled = true;
+        _ = RetryNativeCleanupAsync();
+    }
 
-        _receiveStream.Complete();
-        DisposeNativeCallback();
+    private async Task RetryNativeCleanupAsync()
+    {
+        while (true)
+        {
+            await Task.Delay(CleanupRetryDelayMilliseconds).ConfigureAwait(false);
+            lock (_lifecycleLock)
+            {
+                var result = TryFinalizeNativeCleanup();
+                if (result != NativeCleanupResult.Retry)
+                {
+                    _cleanupRetryScheduled = false;
+                    return;
+                }
+            }
+        }
     }
 
     /// <summary>Length as the native pointer-sized unsigned integer the ABI expects.</summary>

dotnet/test/Unit/FfiRuntimeHostLifetimeTests.cs

@@ -0,0 +1,116 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *--------------------------------------------------------------------------------------------*/
+
+using Microsoft.Extensions.Logging.Abstractions;
+using System.Reflection;
+using Xunit;
+
+namespace GitHub.Copilot.Test.Unit;
+
+public sealed class FfiRuntimeHostLifetimeTests
+{
+    [Fact]
+    public void Dispose_Retains_Callback_Until_Connection_Close_Succeeds()
+    {
+        var allowClose = false;
+        var closeCalls = 0;
+        var shutdownCalls = 0;
+        var callbackReleaseCalls = 0;
+        Func<uint, bool> connectionClose = _ =>
+        {
+            Interlocked.Increment(ref closeCalls);
+            return Volatile.Read(ref allowClose);
+        };
+        Func<uint, bool> hostShutdown = _ =>
+        {
+            Interlocked.Increment(ref shutdownCalls);
+            return true;
+        };
+        Action releaseCallback = () => Interlocked.Increment(ref callbackReleaseCalls);
+
+        var hostType = typeof(CopilotClient).Assembly.GetType("GitHub.Copilot.FfiRuntimeHost", throwOnError: true)!;
+        var constructor = hostType.GetConstructors(BindingFlags.Instance | BindingFlags.NonPublic)
+            .Single(candidate => candidate.GetParameters().Length == 8);
+        using var host = (IDisposable)constructor.Invoke(
+            [
+                "test-runtime",
+                null,
+                null,
+                Array.Empty<string>(),
+                NullLogger.Instance,
+                connectionClose,
+                hostShutdown,
+                releaseCallback,
+            ]);
+
+        hostType.GetField("_connectionId", BindingFlags.Instance | BindingFlags.NonPublic)!
+            .SetValue(host, (uint)21);
+        hostType.GetField("_serverId", BindingFlags.Instance | BindingFlags.NonPublic)!
+            .SetValue(host, (uint)11);
+
+        host.Dispose();
+
+        Assert.Equal(0, Volatile.Read(ref callbackReleaseCalls));
+        Assert.Equal(0, Volatile.Read(ref shutdownCalls));
+
+        Volatile.Write(ref allowClose, true);
+        Assert.True(
+            SpinWait.SpinUntil(
+                () => Volatile.Read(ref callbackReleaseCalls) == 1
+                    && Volatile.Read(ref shutdownCalls) == 1,
+                TimeSpan.FromSeconds(5)),
+            "Deferred native cleanup did not complete.");
+
+        var closeCallsAfterCleanup = Volatile.Read(ref closeCalls);
+        host.Dispose();
+        Thread.Sleep(50);
+
+        Assert.Equal(closeCallsAfterCleanup, Volatile.Read(ref closeCalls));
+        Assert.Equal(1, Volatile.Read(ref callbackReleaseCalls));
+        Assert.Equal(1, Volatile.Read(ref shutdownCalls));
+    }
+
+    [Fact]
+    public void Dispose_Does_Not_Retry_Terminal_Host_Shutdown_Failure()
+    {
+        var closeCalls = 0;
+        var shutdownCalls = 0;
+        var callbackReleaseCalls = 0;
+
+        var hostType = typeof(CopilotClient).Assembly.GetType("GitHub.Copilot.FfiRuntimeHost", throwOnError: true)!;
+        var constructor = hostType.GetConstructors(BindingFlags.Instance | BindingFlags.NonPublic)
+            .Single(candidate => candidate.GetParameters().Length == 8);
+        using var host = (IDisposable)constructor.Invoke(
+            [
+                "test-runtime",
+                null,
+                null,
+                Array.Empty<string>(),
+                NullLogger.Instance,
+                new Func<uint, bool>(_ =>
+                {
+                    Interlocked.Increment(ref closeCalls);
+                    return true;
+                }),
+                new Func<uint, bool>(_ =>
+                {
+                    Interlocked.Increment(ref shutdownCalls);
+                    return false;
+                }),
+                new Action(() => Interlocked.Increment(ref callbackReleaseCalls)),
+            ]);
+
+        hostType.GetField("_connectionId", BindingFlags.Instance | BindingFlags.NonPublic)!
+            .SetValue(host, (uint)21);
+        hostType.GetField("_serverId", BindingFlags.Instance | BindingFlags.NonPublic)!
+            .SetValue(host, (uint)11);
+
+        host.Dispose();
+        Thread.Sleep(250);
+
+        Assert.Equal(1, Volatile.Read(ref closeCalls));
+        Assert.Equal(1, Volatile.Read(ref callbackReleaseCalls));
+        Assert.Equal(1, Volatile.Read(ref shutdownCalls));
+    }
+}