diff --git a/dotnet/src/FfiRuntimeHost.cs b/dotnet/src/FfiRuntimeHost.cs index cc3bccae92..f9f586a385 100644 --- a/dotnet/src/FfiRuntimeHost.cs +++ b/dotnet/src/FfiRuntimeHost.cs @@ -35,14 +35,28 @@ namespace GitHub.Copilot; /// internal sealed partial class FfiRuntimeHost : IDisposable { + private enum NativeCleanupResult + { + Complete, + Retry, + Failed, + } + /// Logical name the native interop layer binds the cdylib to. private const string LibraryName = "copilot_runtime"; + private const int CleanupRetryDelayMilliseconds = 100; + private static readonly object QuarantineLock = new(); + private static readonly HashSet QuarantinedHosts = []; private readonly ILogger _logger; private readonly string? _cliEntrypoint; private readonly string _libraryPath; private readonly IReadOnlyDictionary? _environment; private readonly IReadOnlyList _args; + private readonly Func _connectionClose; + private readonly Func _hostShutdown; + private readonly Action _releaseNativeCallback; + private readonly object _lifecycleLock = new(); private readonly CallbackReceiveStream _receiveStream = new(); private CallbackSendStream? _sendStream; @@ -50,14 +64,31 @@ internal sealed partial class FfiRuntimeHost : IDisposable private uint _serverId; private uint _connectionId; private bool _disposed; + private bool _cleanupRetryScheduled; private FfiRuntimeHost(string libraryPath, string? cliEntrypoint, IReadOnlyDictionary? environment, IReadOnlyList args, ILogger logger) + : this(libraryPath, cliEntrypoint, environment, args, logger, null, null, null) + { + } + + private FfiRuntimeHost( + string libraryPath, + string? cliEntrypoint, + IReadOnlyDictionary? environment, + IReadOnlyList args, + ILogger logger, + Func? connectionClose, + Func? hostShutdown, + Action? releaseNativeCallback) { _libraryPath = libraryPath; _cliEntrypoint = cliEntrypoint; _environment = environment; _args = args; _logger = logger; + _connectionClose = connectionClose ?? NativeConnectionClose; + _hostShutdown = hostShutdown ?? NativeHostShutdown; + _releaseNativeCallback = releaseNativeCallback ?? DisposeNativeCallback; } /// The stream JSON-RPC reads server→client frames from. @@ -109,23 +140,36 @@ await Task.Run(() => 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 static byte[] BuildArgvJson(string? cliEntrypoint, IReadOnlyList private bool SendFrame(ReadOnlySpan 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 @@ private void FeedInbound(IntPtr bytesPtr, UIntPtr bytesLen) 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; + } + } + } } /// Length as the native pointer-sized unsigned integer the ABI expects. diff --git a/dotnet/test/Unit/FfiRuntimeHostLifetimeTests.cs b/dotnet/test/Unit/FfiRuntimeHostLifetimeTests.cs new file mode 100644 index 0000000000..483fd30f23 --- /dev/null +++ b/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 connectionClose = _ => + { + Interlocked.Increment(ref closeCalls); + return Volatile.Read(ref allowClose); + }; + Func 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(), + 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(), + NullLogger.Instance, + new Func(_ => + { + Interlocked.Increment(ref closeCalls); + return true; + }), + new Func(_ => + { + 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)); + } +} diff --git a/go/internal/ffihost/ffihost.go b/go/internal/ffihost/ffihost.go index 9a824881de..4add70cb5d 100644 --- a/go/internal/ffihost/ffihost.go +++ b/go/internal/ffihost/ffihost.go @@ -38,10 +38,12 @@ import ( "encoding/json" "fmt" "io" + "log" "runtime" "strings" "sync" "sync/atomic" + "time" "unsafe" "github.com/ebitengine/purego" @@ -145,12 +147,11 @@ type Host struct { lifecycleMu sync.Mutex // mu serializes disposal with native callbacks so the receive buffer cannot // be fed after it is closed. - mu sync.Mutex - serverID uint32 - connectionID uint32 - disposed bool - // activeCallbacks counts outbound native callbacks currently executing. - activeCallbacks int + mu sync.Mutex + serverID uint32 + connectionID uint32 + disposed bool + cleanupScheduled bool recv *receiveBuffer @@ -269,13 +270,9 @@ func (h *Host) onOutbound(bytesPtr uintptr, bytesLen uintptr) uintptr { h.mu.Unlock() return 0 } - h.activeCallbacks++ h.mu.Unlock() defer func() { - h.mu.Lock() - h.activeCallbacks-- - h.mu.Unlock() // Never let a panic unwind into native code. _ = recover() }() @@ -294,11 +291,18 @@ func (h *Host) onOutbound(bytesPtr uintptr, bytesLen uintptr) uintptr { } func (h *Host) writeFrame(frame []byte) (int, error) { + h.mu.Lock() + disposed := h.disposed + h.mu.Unlock() + if disposed { + return 0, fmt.Errorf("the in-process runtime connection is closed") + } + h.lifecycleMu.Lock() defer h.lifecycleMu.Unlock() h.mu.Lock() - disposed := h.disposed + disposed = h.disposed h.mu.Unlock() connID := h.connectionID if disposed || connID == 0 { @@ -315,9 +319,8 @@ func (h *Host) writeFrame(frame []byte) (int, error) { return len(frame), nil } -// Dispose closes the FFI connection, shuts down the native host, and releases -// resources. It is idempotent and waits for any in-flight outbound callback to -// finish before closing the receive buffer. +// Dispose closes the receive side immediately and releases native resources +// after connectionClose confirms that outbound callbacks are quiescent. func (h *Host) Dispose() { h.lifecycleMu.Lock() defer h.lifecycleMu.Unlock() @@ -327,45 +330,64 @@ func (h *Host) Dispose() { h.mu.Unlock() return } - // Publish disposed under the same lock onOutbound uses to check it, so no new - // callback can pass the check and increment activeCallbacks after the drain - // loop below observes zero. h.disposed = true - connID := h.connectionID - serverID := h.serverID - callbackToken := h.callbackToken - h.connectionID = 0 - h.serverID = 0 - h.callbackToken = 0 h.mu.Unlock() - if callbackToken != 0 { - outboundTargets.Delete(callbackToken) + h.recv.Close() + if !h.tryFinalizeCleanupLocked() { + h.scheduleCleanupRetryLocked() } +} + +func (h *Host) tryFinalizeCleanupLocked() bool { + connID := h.connectionID - // Stop accepting new callbacks and wait for in-flight ones to drain before - // closing the receive buffer they feed. - for { - h.mu.Lock() - if h.activeCallbacks == 0 { - h.mu.Unlock() - break + if connID != 0 { + if !h.lib.connectionClose(connID) { + return false } - h.mu.Unlock() - runtime.Gosched() + h.connectionID = 0 } - if connID != 0 { - h.lib.connectionClose(connID) + callbackToken := h.callbackToken + h.callbackToken = 0 + if callbackToken != 0 { + outboundTargets.Delete(callbackToken) } + + serverID := h.serverID if serverID != 0 { - h.lib.hostShutdown(serverID) + if !h.lib.hostShutdown(serverID) { + log.Printf("FfiRuntimeHost: host_shutdown did not recognize server %d", serverID) + } + h.serverID = 0 if h.cliEntrypoint != "" { // A legacy host may restore its saved SIGCHLD action during shutdown. rearmForeignSignalHandlers(h.lib.handle) } } - h.recv.Close() + return true +} + +func (h *Host) scheduleCleanupRetryLocked() { + if h.cleanupScheduled { + return + } + h.cleanupScheduled = true + go func() { + timer := time.NewTimer(100 * time.Millisecond) + defer timer.Stop() + for range timer.C { + h.lifecycleMu.Lock() + if h.tryFinalizeCleanupLocked() { + h.cleanupScheduled = false + h.lifecycleMu.Unlock() + return + } + h.lifecycleMu.Unlock() + timer.Reset(100 * time.Millisecond) + } + }() } // hostWriter adapts Host into the io.WriteCloser jsonrpc2 writes request frames to. diff --git a/go/internal/ffihost/ffihost_test.go b/go/internal/ffihost/ffihost_test.go index ccb48af419..3bb7555a6f 100644 --- a/go/internal/ffihost/ffihost_test.go +++ b/go/internal/ffihost/ffihost_test.go @@ -25,6 +25,69 @@ func TestDisposeUnregistersOutboundTarget(t *testing.T) { } } +func TestDisposeRetainsOutboundTargetUntilConnectionCloseSucceeds(t *testing.T) { + token := uintptr(nextOutboundToken.Add(1)) + var allowClose atomic.Bool + var closeCalls atomic.Int32 + var shutdownCalls atomic.Int32 + host := &Host{ + lib: &ffiLibrary{ + connectionClose: func(_ uint32) bool { + closeCalls.Add(1) + return allowClose.Load() + }, + hostShutdown: func(_ uint32) bool { + shutdownCalls.Add(1) + return true + }, + }, + recv: newReceiveBuffer(), + serverID: 11, + connectionID: 21, + callbackToken: token, + } + outboundTargets.Store(token, host) + + host.Dispose() + + if _, ok := outboundTargets.Load(token); !ok { + t.Fatal("Expected callback target to remain registered after connection close reported non-quiescence") + } + if got := shutdownCalls.Load(); got != 0 { + t.Fatalf("Expected host shutdown to be deferred, got %d calls", got) + } + + allowClose.Store(true) + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + _, registered := outboundTargets.Load(token) + if !registered && shutdownCalls.Load() == 1 { + break + } + time.Sleep(10 * time.Millisecond) + } + + if _, ok := outboundTargets.Load(token); ok { + t.Fatal("Expected callback target to be removed after connection close succeeded") + } + if got := closeCalls.Load(); got < 2 { + t.Fatalf("Expected connection close to be retried, got %d calls", got) + } + if got := shutdownCalls.Load(); got != 1 { + t.Fatalf("Expected exactly one host shutdown, got %d", got) + } + + closeCallsAfterCleanup := closeCalls.Load() + host.Dispose() + time.Sleep(150 * time.Millisecond) + if got := closeCalls.Load(); got != closeCallsAfterCleanup { + t.Fatalf("Expected repeated disposal to be a no-op, got %d additional close calls", got-closeCallsAfterCleanup) + } + if got := shutdownCalls.Load(); got != 1 { + t.Fatalf("Expected exactly one host shutdown after repeated disposal, got %d", got) + } +} + func TestBuildArgvWithoutEntrypointContainsOnlyManagedOptions(t *testing.T) { host := &Host{ args: []string{"--log-level", "debug", "--remote"}, diff --git a/java/sdk/src/main/java/com/github/copilot/ffi/FfiRuntimeHost.java b/java/sdk/src/main/java/com/github/copilot/ffi/FfiRuntimeHost.java index cb5bba1af1..e28a573bc8 100644 --- a/java/sdk/src/main/java/com/github/copilot/ffi/FfiRuntimeHost.java +++ b/java/sdk/src/main/java/com/github/copilot/ffi/FfiRuntimeHost.java @@ -14,6 +14,9 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -39,15 +42,15 @@ public final class FfiRuntimeHost implements AutoCloseable { private static final Logger LOG = Logger.getLogger(FfiRuntimeHost.class.getName()); private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final Set QUARANTINED_HOSTS = ConcurrentHashMap.newKeySet(); private final NativeBinding nativeBinding; private final QueueInputStream receiveStream; private final AtomicBoolean closing = new AtomicBoolean(false); private final AtomicBoolean disposed = new AtomicBoolean(false); + private final AtomicBoolean cleanupScheduled = new AtomicBoolean(false); private final AtomicInteger serverId = new AtomicInteger(0); private final AtomicInteger connectionId = new AtomicInteger(0); - private final AtomicInteger activeCallbacks = new AtomicInteger(0); - private final Object callbackDrainMonitor = new Object(); private final ReentrantLock operationLock = new ReentrantLock(); private final FfiOutputStream sendStream; private final String libraryPath; @@ -163,54 +166,66 @@ public void close() { closing.set(true); + try { + receiveStream.close(); + } catch (Throwable ignored) { + // never throw from close + } + + if (!tryFinalizeCleanup()) { + scheduleCleanupRetry(); + } + } + + private boolean tryFinalizeCleanup() { operationLock.lock(); try { - int connHandle = connectionId.getAndSet(0); + int connHandle = connectionId.get(); if (connHandle != 0) { try { - nativeBinding.connectionClose(connHandle); + if (!nativeBinding.connectionClose(connHandle)) { + return false; + } } catch (Throwable t) { LOG.log(Level.FINE, "Failed to close FFI connection", t); + QUARANTINED_HOSTS.add(this); + return true; } + connectionId.set(0); + callbackRef = null; + QUARANTINED_HOSTS.remove(this); } - } finally { - operationLock.unlock(); - } - - drainActiveCallbacks(); - int hostHandle = serverId.getAndSet(0); - if (hostHandle != 0) { - try { - nativeBinding.hostShutdown(hostHandle); - } catch (Throwable t) { - LOG.log(Level.FINE, "Failed to shut down FFI host", t); + int hostHandle = serverId.get(); + if (hostHandle != 0) { + try { + if (!nativeBinding.hostShutdown(hostHandle)) { + LOG.fine(() -> "FFI host shutdown did not recognize server " + hostHandle); + } + } catch (Throwable t) { + LOG.log(Level.FINE, "Failed to shut down FFI host", t); + } + serverId.set(0); } + return true; + } finally { + operationLock.unlock(); } + } - try { - receiveStream.close(); - } catch (Throwable ignored) { - // never throw from close + private void scheduleCleanupRetry() { + if (!cleanupScheduled.compareAndSet(false, true)) { + return; } - - callbackRef = null; + CompletableFuture.delayedExecutor(100, TimeUnit.MILLISECONDS).execute(this::retryCleanup); } - private void drainActiveCallbacks() { - while (activeCallbacks.get() > 0) { - synchronized (callbackDrainMonitor) { - if (activeCallbacks.get() == 0) { - return; - } - try { - callbackDrainMonitor.wait(10L); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - return; - } - } + private void retryCleanup() { + if (tryFinalizeCleanup()) { + cleanupScheduled.set(false); + return; } + CompletableFuture.delayedExecutor(100, TimeUnit.MILLISECONDS).execute(this::retryCleanup); } private OutboundCallback createOutboundCallback() { @@ -218,7 +233,6 @@ private OutboundCallback createOutboundCallback() { if (closing.get()) { return; } - activeCallbacks.incrementAndGet(); try { int length = len.intValue(); if (closing.get() || data == null || length <= 0) { @@ -230,12 +244,6 @@ private OutboundCallback createOutboundCallback() { } } catch (Throwable t) { LOG.log(Level.WARNING, "Exception in FFI outbound callback", t); - } finally { - if (activeCallbacks.decrementAndGet() == 0) { - synchronized (callbackDrainMonitor) { - callbackDrainMonitor.notifyAll(); - } - } } }; } diff --git a/java/sdk/src/main/java/com/github/copilot/ffi/JnaNativeBinding.java b/java/sdk/src/main/java/com/github/copilot/ffi/JnaNativeBinding.java index ba3c3c40a5..bf6c912f5f 100644 --- a/java/sdk/src/main/java/com/github/copilot/ffi/JnaNativeBinding.java +++ b/java/sdk/src/main/java/com/github/copilot/ffi/JnaNativeBinding.java @@ -41,17 +41,17 @@ *

Active-callback tracking

*

* The {@link #activeCallbacks} counter is incremented when the native runtime - * enters the outbound callback and decremented when the callback returns. - * Callers (e.g. {@code FfiRuntimeHost}) must drain this counter to zero before - * calling {@link #connectionClose} or {@link #hostShutdown}. + * enters the outbound callback and decremented when the callback returns. It is + * retained for diagnostics and tests; a successful {@link #connectionClose} is + * the authoritative callback-quiescence barrier. * *

Callback lifetime

*

- * The native runtime can invoke an outbound callback after connection close and - * host shutdown return. Each JNA callback wrapper is therefore retained for the - * lifetime of the JVM. After host shutdown, its Java delegate is detached so a - * late native invocation safely becomes a no-op without retaining the complete - * host object graph. + * The native runtime can still be inside an outbound callback when + * {@link #connectionClose} returns {@code false}. Each JNA callback wrapper is + * therefore retained for the lifetime of the JVM. After connection close + * reports quiescence, its Java delegate is detached so the wrapper no longer + * retains the complete host object graph. * *

GraalVM Native Image

*

@@ -133,18 +133,16 @@ int copilot_runtime_connection_open(int serverId, OutboundCallback callback, Poi */ private final CopilotRuntimeLibrary lib; - /** - * Count of callbacks currently executing on native threads. Must reach zero - * before {@link #connectionClose} or {@link #hostShutdown} is called. - */ + /** Count of callbacks currently executing on native threads. */ final AtomicInteger activeCallbacks = new AtomicInteger(0); /** * Callback registrations keyed by connection handle. *

* Registrations remain here through connection close because native callbacks - * can still arrive. Successful host shutdown detaches their Java delegates; the - * wrappers themselves remain rooted by {@link #RETAINED_CALLBACKS}. + * can still arrive while close reports non-quiescence. Successful connection + * close detaches their Java delegates; the wrappers themselves remain rooted by + * {@link #RETAINED_CALLBACKS}. */ private final Map callbackRegistrations = new ConcurrentHashMap<>(); @@ -269,7 +267,14 @@ public boolean connectionWrite(int connectionId, byte[] data, int dataLen) { @Override public boolean connectionClose(int connectionId) { - return lib.copilot_runtime_connection_close(connectionId) != 0; + boolean closed = lib.copilot_runtime_connection_close(connectionId) != 0; + if (closed) { + CallbackRegistration registration = callbackRegistrations.remove(connectionId); + if (registration != null) { + registration.detach(); + } + } + return closed; } // ------------------------------------------------------------------------- diff --git a/java/sdk/src/test/java/com/github/copilot/ffi/FfiRuntimeHostTest.java b/java/sdk/src/test/java/com/github/copilot/ffi/FfiRuntimeHostTest.java index 545b316c8c..ffb4eb3dfd 100644 --- a/java/sdk/src/test/java/com/github/copilot/ffi/FfiRuntimeHostTest.java +++ b/java/sdk/src/test/java/com/github/copilot/ffi/FfiRuntimeHostTest.java @@ -6,7 +6,6 @@ import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -330,11 +329,14 @@ public boolean connectionClose(int connectionId) { } @Test - void closeDrainsActiveCallbacksBeforeHostShutdown() throws Exception { + void closeRetriesUntilCallbackQuiescenceIsReported() throws Exception { CountDownLatch callbackEntered = new CountDownLatch(1); CountDownLatch allowCallbackToReturn = new CountDownLatch(1); + CountDownLatch shutdownCalled = new CountDownLatch(1); AtomicBoolean shutdownObservedAfterCallbackReturn = new AtomicBoolean(false); AtomicBoolean callbackFinished = new AtomicBoolean(false); + AtomicInteger closeCalls = new AtomicInteger(0); + AtomicInteger shutdownCalls = new AtomicInteger(0); AtomicReference callbackRef = new AtomicReference<>(); NativeBinding binding = new NativeBinding() { @@ -345,7 +347,9 @@ public int hostStart(byte[] argvJson, int argvJsonLen, byte[] envJson, int envJs @Override public boolean hostShutdown(int serverId) { + shutdownCalls.incrementAndGet(); shutdownObservedAfterCallbackReturn.set(callbackFinished.get()); + shutdownCalled.countDown(); return true; } @@ -363,7 +367,8 @@ public boolean connectionWrite(int connectionId, byte[] data, int dataLen) { @Override public boolean connectionClose(int connectionId) { - return true; + closeCalls.incrementAndGet(); + return callbackFinished.get(); } }; @@ -393,11 +398,17 @@ void enqueue(byte[] bytes) { assertTrue(callbackEntered.await(2, TimeUnit.SECONDS)); CompletableFuture closeFuture = CompletableFuture.runAsync(host::close); - Thread.sleep(150); - assertFalse(closeFuture.isDone(), "close should wait for active callback to drain"); + closeFuture.get(5, TimeUnit.SECONDS); + assertEquals(0, shutdownCalls.get(), "host shutdown must wait for a successful connection close"); + allowCallbackToReturn.countDown(); callbackFuture.get(5, TimeUnit.SECONDS); - closeFuture.get(5, TimeUnit.SECONDS); + assertTrue(shutdownCalled.await(5, TimeUnit.SECONDS), "deferred cleanup should retry connection close"); + assertTrue(closeCalls.get() >= 2, "connection close should be retried after reporting non-quiescence"); assertTrue(shutdownObservedAfterCallbackReturn.get(), "host_shutdown should run after callback drains"); + + host.close(); + Thread.sleep(150); + assertEquals(1, shutdownCalls.get(), "host shutdown should happen exactly once"); } } diff --git a/java/sdk/src/test/java/com/github/copilot/ffi/JnaNativeBindingTest.java b/java/sdk/src/test/java/com/github/copilot/ffi/JnaNativeBindingTest.java index e37a7cfc82..f0ec519d47 100644 --- a/java/sdk/src/test/java/com/github/copilot/ffi/JnaNativeBindingTest.java +++ b/java/sdk/src/test/java/com/github/copilot/ffi/JnaNativeBindingTest.java @@ -230,7 +230,7 @@ void activeCallbacksStartsAtZero() { } @Test - void callbackWrapperRemainsReachableAfterConnectionClose() throws InterruptedException { + void callbackWrapperRemainsReachableAndIsDetachedAfterConnectionClose() throws InterruptedException { StubRuntimeLibrary stub = new StubRuntimeLibrary(); stub.connectionOpenReturn = 99; JnaNativeBinding binding = new JnaNativeBinding(stub); @@ -244,7 +244,7 @@ void callbackWrapperRemainsReachableAfterConnectionClose() throws InterruptedExc OutboundCallback callback = callbackReference.get(); assertNotNull(callback, "Callback wrapper must remain strongly reachable after connection close"); callback.invoke(Pointer.NULL, Pointer.NULL, new SizeT(0)); - assertEquals(1, invocations.get(), "A callback queued before close must remain safely invocable"); + assertEquals(0, invocations.get(), "Successful connection close must detach the Java delegate"); } private static WeakReference openAndCloseConnection(JnaNativeBinding binding, diff --git a/nodejs/src/ffiRuntimeHost.ts b/nodejs/src/ffiRuntimeHost.ts index 4795e325ce..ee71c44717 100644 --- a/nodejs/src/ffiRuntimeHost.ts +++ b/nodejs/src/ffiRuntimeHost.ts @@ -25,6 +25,7 @@ const SYMBOL_PREFIX = "copilot_runtime_"; // A long, referenced no-op timer keeps the Node event loop alive while the in-process // connection is open (see start()); the exact interval is irrelevant. const KEEP_ALIVE_INTERVAL_MS = 1 << 30; +const CLEANUP_RETRY_INTERVAL_MS = 100; type KoffiFunction = ReturnType["func"]>; type KoffiType = ReturnType; @@ -122,12 +123,17 @@ function buildEnvJson(environment?: Record): Buffer } export class FfiRuntimeHost { + private static readonly quarantinedHosts = new Set(); + private readonly lib: FfiLibrary; private serverId = 0; private connectionId = 0; private disposed = false; + private starting = false; private outboundCallback: KoffiRegisteredCallback | undefined; private keepAliveTimer: ReturnType | undefined; + private cleanupRetryTimer: ReturnType | undefined; + private cleanupInProgress = false; /** The stream JSON-RPC reads server→client frames from. */ readonly receiveStream: PassThrough; @@ -179,64 +185,80 @@ export class FfiRuntimeHost { /** Starts the in-process Rust runtime and opens the FFI JSON-RPC connection. */ async start(): Promise { + if (this.disposed) { + throw new Error("The in-process runtime host is disposed."); + } + this.starting = true; const argvJson = buildArgvJson(this.cliEntrypoint, this.args); const envJson = buildEnvJson(this.environment); - // The native host has no cwd parameter, so it uses this process's cwd. A custom - // working directory is intentionally - // unsupported for the in-process transport (rejected by the client constructor) - // rather than mutating the shared process-global cwd here. - - // host_start constructs the native engine synchronously; run it as an async FFI - // call so the Node event loop isn't blocked. - this.serverId = await new Promise((resolvePromise, rejectPromise) => { - this.lib.hostStart.async( - argvJson, - argvJson.length, - envJson, - envJson ? envJson.length : 0, - (error: Error | null, result: number) => { - if (error) { - rejectPromise(error); - } else { - resolvePromise(result); + try { + // The native host has no cwd parameter, so it uses this process's cwd. A custom + // working directory is intentionally + // unsupported for the in-process transport (rejected by the client constructor) + // rather than mutating the shared process-global cwd here. + + // host_start constructs the native engine synchronously; run it as an async FFI + // call so the Node event loop isn't blocked. + this.serverId = await new Promise((resolvePromise, rejectPromise) => { + this.lib.hostStart.async( + argvJson, + argvJson.length, + envJson, + envJson ? envJson.length : 0, + (error: Error | null, result: number) => { + if (error) { + rejectPromise(error); + } else { + resolvePromise(result); + } } - } + ); + }); + if (!this.serverId) { + throw new Error( + `copilot_runtime_host_start failed (library '${this.libraryPath}').` + ); + } + if (this.disposed) { + throw new Error("The in-process runtime host was disposed during startup."); + } + + this.outboundCallback = koffi.register( + (_userData: unknown, bytesPtr: unknown, bytesLen: number | bigint) => + this.feedInbound(bytesPtr, bytesLen), + this.lib.outboundCallbackType ); - }); - if (!this.serverId) { - throw new Error(`copilot_runtime_host_start failed (library '${this.libraryPath}').`); - } - this.outboundCallback = koffi.register( - (_userData: unknown, bytesPtr: unknown, bytesLen: number | bigint) => - this.feedInbound(bytesPtr, bytesLen), - this.lib.outboundCallbackType - ); + this.connectionId = this.lib.connectionOpen( + this.serverId, + this.outboundCallback, + null, + null, + 0, + null, + 0, + null, + 0 + ); + if (!this.connectionId) { + this.unregisterCallback(); + this.lib.hostShutdown(this.serverId); + this.serverId = 0; + throw new Error("copilot_runtime_connection_open failed."); + } - this.connectionId = this.lib.connectionOpen( - this.serverId, - this.outboundCallback, - null, - null, - 0, - null, - 0, - null, - 0 - ); - if (!this.connectionId) { - this.unregisterCallback(); - this.lib.hostShutdown(this.serverId); - this.serverId = 0; - throw new Error("copilot_runtime_connection_open failed."); + // The in-process transport has no socket/pipe handle to keep the Node event loop + // alive while the SDK is idle awaiting a server→client frame. koffi delivers the + // outbound callback on the loop but does not reference it, so hold one referenced + // timer for the lifetime of the connection. + this.keepAliveTimer = setInterval(() => {}, KEEP_ALIVE_INTERVAL_MS); + } finally { + this.starting = false; + if (this.disposed) { + this.tryFinalizeCleanup(); + } } - - // The in-process transport has no socket/pipe handle to keep the Node event loop - // alive while the SDK is idle awaiting a server→client frame. koffi delivers the - // outbound callback on the loop but does not reference it, so hold one referenced - // timer for the lifetime of the connection. - this.keepAliveTimer = setInterval(() => {}, KEEP_ALIVE_INTERVAL_MS); } private writeFrame(frame: Buffer): void { @@ -283,50 +305,104 @@ export class FfiRuntimeHost { } } - private unregisterCallback(): void { + private unregisterCallback(): boolean { if (this.outboundCallback === undefined) { - return; + return true; } const callback = this.outboundCallback; - this.outboundCallback = undefined; try { koffi.unregister(callback); - } catch { - // Ignore teardown failures. + this.outboundCallback = undefined; + return true; + } catch (error) { + console.error( + `Failed to unregister in-process FFI callback: ${error instanceof Error ? (error.stack ?? error.message) : String(error)}` + ); + FfiRuntimeHost.quarantinedHosts.add(this); + return false; } } - /** Closes the FFI connection, shuts down the native host, and releases resources. */ - dispose(): void { - if (this.disposed) { + private scheduleCleanupRetry(): void { + if (this.cleanupRetryTimer !== undefined) { return; } - this.disposed = true; + this.cleanupRetryTimer = setTimeout(() => { + this.cleanupRetryTimer = undefined; + this.tryFinalizeCleanup(); + }, CLEANUP_RETRY_INTERVAL_MS); + } - if (this.keepAliveTimer !== undefined) { - clearInterval(this.keepAliveTimer); - this.keepAliveTimer = undefined; + private tryFinalizeCleanup(): void { + if (this.cleanupInProgress) { + this.scheduleCleanupRetry(); + return; } + this.cleanupInProgress = true; try { if (this.connectionId) { - this.lib.connectionClose(this.connectionId); + let closed = false; + try { + closed = Boolean(this.lib.connectionClose(this.connectionId)); + } catch (error) { + console.error( + `Failed to close in-process FFI connection: ${error instanceof Error ? (error.stack ?? error.message) : String(error)}` + ); + FfiRuntimeHost.quarantinedHosts.add(this); + if (this.keepAliveTimer !== undefined) { + clearInterval(this.keepAliveTimer); + this.keepAliveTimer = undefined; + } + return; + } + if (!closed) { + this.scheduleCleanupRetry(); + return; + } this.connectionId = 0; } - } catch { - // Ignore teardown failures. - } + const callbackUnregistered = this.unregisterCallback(); + + // The referenced timer is part of the callback lifetime. Clearing it + // before connection_close reports quiescence can let the process exit + // while native code still owns the Koffi registration. + if (this.keepAliveTimer !== undefined) { + clearInterval(this.keepAliveTimer); + this.keepAliveTimer = undefined; + } - try { if (this.serverId) { - this.lib.hostShutdown(this.serverId); + try { + if (!this.lib.hostShutdown(this.serverId)) { + console.error( + `In-process FFI host shutdown did not recognize server ${this.serverId}.` + ); + } + } catch (error) { + console.error( + `Failed to shut down in-process FFI host: ${error instanceof Error ? (error.stack ?? error.message) : String(error)}` + ); + } this.serverId = 0; } - } catch { - // Ignore teardown failures. + if (callbackUnregistered) { + FfiRuntimeHost.quarantinedHosts.delete(this); + } + } finally { + this.cleanupInProgress = false; } + } + /** Closes the FFI connection, shuts down the native host, and releases resources. */ + dispose(): void { + if (this.disposed) { + return; + } + this.disposed = true; this.receiveStream.end(); - this.unregisterCallback(); + if (!this.starting) { + this.tryFinalizeCleanup(); + } } } diff --git a/nodejs/test/ffiRuntimeHost.test.ts b/nodejs/test/ffiRuntimeHost.test.ts new file mode 100644 index 0000000000..d538b45385 --- /dev/null +++ b/nodejs/test/ffiRuntimeHost.test.ts @@ -0,0 +1,177 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +const ffi = vi.hoisted(() => { + let registeredCallback: + | ((userData: unknown, bytesPtr: unknown, bytesLen: number) => void) + | undefined; + const callbackToken = {}; + const hostStart = Object.assign(vi.fn(), { + async: vi.fn( + ( + _argv: Buffer, + _argvLength: number, + _env: Buffer | null, + _envLength: number, + callback: (error: Error | null, result: number) => void + ) => callback(null, 11) + ), + }); + const hostShutdown = vi.fn(() => true); + const connectionOpen = vi.fn(() => 21); + const connectionWrite = vi.fn(() => true); + const connectionClose = vi.fn<() => boolean>(); + const register = vi.fn( + (callback: (userData: unknown, bytesPtr: unknown, bytesLen: number) => void) => { + registeredCallback = callback; + return callbackToken; + } + ); + const unregister = vi.fn(); + + return { + callbackToken, + connectionClose, + connectionOpen, + connectionWrite, + getRegisteredCallback: () => registeredCallback, + hostShutdown, + hostStart, + register, + unregister, + }; +}); + +vi.mock("node:fs", () => ({ + existsSync: vi.fn(() => true), +})); + +vi.mock("koffi", () => ({ + default: { + array: vi.fn(() => ({})), + decode: vi.fn(() => new Uint8Array([1])), + load: vi.fn(() => ({ + func: vi.fn((name: string) => { + if (name.endsWith("host_start")) return ffi.hostStart; + if (name.endsWith("host_shutdown")) return ffi.hostShutdown; + if (name.endsWith("connection_open")) return ffi.connectionOpen; + if (name.endsWith("connection_write")) return ffi.connectionWrite; + if (name.endsWith("connection_close")) return ffi.connectionClose; + throw new Error(`Unexpected FFI symbol: ${name}`); + }), + })), + pointer: vi.fn(() => ({})), + proto: vi.fn(() => ({})), + register: ffi.register, + unregister: ffi.unregister, + }, +})); + +import { FfiRuntimeHost } from "../src/ffiRuntimeHost.js"; + +describe("FfiRuntimeHost callback cleanup", () => { + beforeEach(() => { + vi.useFakeTimers(); + ffi.connectionClose.mockReset(); + ffi.connectionOpen.mockClear(); + ffi.hostShutdown.mockClear(); + ffi.hostStart.mockClear(); + ffi.hostStart.async.mockClear(); + ffi.register.mockClear(); + ffi.unregister.mockClear(); + }); + + afterEach(() => { + vi.clearAllTimers(); + vi.useRealTimers(); + }); + + it("retains the callback and keepalive until a retry closes the connection", async () => { + ffi.connectionClose.mockReturnValueOnce(false).mockReturnValueOnce(true); + const host = FfiRuntimeHost.create("runtime.node", undefined, undefined, []); + await host.start(); + + host.dispose(); + + expect(ffi.connectionClose).toHaveBeenCalledTimes(1); + expect(ffi.unregister).not.toHaveBeenCalled(); + expect(ffi.hostShutdown).not.toHaveBeenCalled(); + expect((host as any).outboundCallback).toBe(ffi.callbackToken); + expect((host as any).keepAliveTimer).toBeDefined(); + + await vi.advanceTimersByTimeAsync(100); + + expect(ffi.connectionClose).toHaveBeenCalledTimes(2); + expect(ffi.unregister).toHaveBeenCalledTimes(1); + expect(ffi.hostShutdown).toHaveBeenCalledTimes(1); + expect((host as any).outboundCallback).toBeUndefined(); + expect((host as any).keepAliveTimer).toBeUndefined(); + expect(vi.getTimerCount()).toBe(0); + + host.dispose(); + await vi.advanceTimersByTimeAsync(100); + expect(ffi.connectionClose).toHaveBeenCalledTimes(2); + expect(ffi.unregister).toHaveBeenCalledTimes(1); + expect(ffi.hostShutdown).toHaveBeenCalledTimes(1); + }); + + it("defers reclamation when dispose is called from the outbound callback", async () => { + ffi.connectionClose.mockReturnValueOnce(false).mockReturnValueOnce(true); + const host = FfiRuntimeHost.create("runtime.node", undefined, undefined, []); + await host.start(); + vi.spyOn(host as any, "feedInbound").mockImplementation(() => host.dispose()); + + ffi.getRegisteredCallback()?.(null, {}, 1); + + expect(ffi.unregister).not.toHaveBeenCalled(); + expect((host as any).outboundCallback).toBe(ffi.callbackToken); + + await vi.advanceTimersByTimeAsync(100); + + expect(ffi.unregister).toHaveBeenCalledTimes(1); + expect(ffi.hostShutdown).toHaveBeenCalledTimes(1); + expect(vi.getTimerCount()).toBe(0); + }); + + it("does not retry a terminal host shutdown failure", async () => { + ffi.connectionClose.mockReturnValue(true); + ffi.hostShutdown.mockReturnValueOnce(false); + const error = vi.spyOn(console, "error").mockImplementation(() => {}); + const host = FfiRuntimeHost.create("runtime.node", undefined, undefined, []); + await host.start(); + + host.dispose(); + await vi.advanceTimersByTimeAsync(500); + + expect(ffi.connectionClose).toHaveBeenCalledTimes(1); + expect(ffi.unregister).toHaveBeenCalledTimes(1); + expect(ffi.hostShutdown).toHaveBeenCalledTimes(1); + expect(vi.getTimerCount()).toBe(0); + error.mockRestore(); + }); + + it("retains the callback token when Koffi unregistration fails", async () => { + ffi.connectionClose.mockReturnValue(true); + ffi.unregister.mockImplementationOnce(() => { + throw new Error("unregister failed"); + }); + const error = vi.spyOn(console, "error").mockImplementation(() => {}); + const host = FfiRuntimeHost.create("runtime.node", undefined, undefined, []); + await host.start(); + + host.dispose(); + await vi.advanceTimersByTimeAsync(500); + + expect(ffi.connectionClose).toHaveBeenCalledTimes(1); + expect(ffi.unregister).toHaveBeenCalledTimes(1); + expect(ffi.hostShutdown).toHaveBeenCalledTimes(1); + expect((host as any).outboundCallback).toBe(ffi.callbackToken); + expect((FfiRuntimeHost as any).quarantinedHosts.has(host)).toBe(true); + expect((host as any).keepAliveTimer).toBeUndefined(); + expect(vi.getTimerCount()).toBe(0); + + host.dispose(); + expect(ffi.unregister).toHaveBeenCalledTimes(1); + error.mockRestore(); + }); +}); diff --git a/nodejs/tsconfig.test.json b/nodejs/tsconfig.test.json index 03d23317f7..cc659e9a3a 100644 --- a/nodejs/tsconfig.test.json +++ b/nodejs/tsconfig.test.json @@ -5,6 +5,11 @@ "emitDeclarationOnly": false, "types": ["node"] }, - "include": ["src/**/*", "test/session-event-types.test.ts", "test/message-source.test.ts"], + "include": [ + "src/**/*", + "test/ffiRuntimeHost.test.ts", + "test/session-event-types.test.ts", + "test/message-source.test.ts" + ], "exclude": ["node_modules", "dist"] } diff --git a/python/copilot/_ffi_runtime_host.py b/python/copilot/_ffi_runtime_host.py index 511cbae9c4..8674ed6ebc 100644 --- a/python/copilot/_ffi_runtime_host.py +++ b/python/copilot/_ffi_runtime_host.py @@ -41,13 +41,13 @@ import os import sys import threading -import time from collections.abc import Sequence from pathlib import Path logger = logging.getLogger("copilot.ffi") _SYMBOL_PREFIX = "copilot_runtime_" +_CLEANUP_RETRY_INTERVAL_SECONDS = 0.1 # The C ABI outbound callback: void(void *user_data, uint8 *bytes, size_t len). _OutboundCallback = ctypes.CFUNCTYPE( @@ -337,6 +337,8 @@ class FfiRuntimeHost: :class:`JsonRpcClient`, and call :meth:`dispose` to tear everything down. """ + _quarantined_hosts: set[FfiRuntimeHost] = set() + def __init__( self, library_path: str, @@ -354,15 +356,14 @@ def __init__( self._connection_id = 0 self._disposed = False self._dispose_lock = threading.Lock() + self._operation_lock = threading.Lock() + self._cleanup_timer: threading.Timer | None = None + self._starting = False self._receive_buffer = _ReceiveBuffer() # Keep a strong reference to the ctypes callback for its whole lifetime; # dropping it while native code can still invoke it is a use-after-free. self._outbound_callback: ctypes._FuncPointer | None = None - # Serializes teardown against in-flight native callbacks. - self._active_callbacks = 0 - self._callback_lock = threading.Lock() - self._process = _FfiProcessAdapter(self) @property @@ -415,32 +416,47 @@ def start_blocking(self) -> None: Must be run off the event loop (e.g. via :func:`asyncio.to_thread`). """ + with self._dispose_lock: + if self._disposed: + raise RuntimeError("The in-process runtime host is disposed.") + self._starting = True + argv = self._build_argv() env = self._build_env() - self._server_id = self._lib.host_start(argv, len(argv), env, len(env) if env else 0) - if not self._server_id: - raise RuntimeError( - f"copilot_runtime_host_start failed (library '{self._library_path}')." + try: + self._server_id = self._lib.host_start(argv, len(argv), env, len(env) if env else 0) + if not self._server_id: + raise RuntimeError( + f"copilot_runtime_host_start failed (library '{self._library_path}')." + ) + with self._dispose_lock: + if self._disposed: + raise RuntimeError("The in-process runtime host was disposed during startup.") + + self._outbound_callback = _OutboundCallback(self._on_outbound) + self._connection_id = self._lib.connection_open( + self._server_id, + self._outbound_callback, + None, + None, + 0, + None, + 0, + None, + 0, ) - - self._outbound_callback = _OutboundCallback(self._on_outbound) - self._connection_id = self._lib.connection_open( - self._server_id, - self._outbound_callback, - None, - None, - 0, - None, - 0, - None, - 0, - ) - if not self._connection_id: - self._outbound_callback = None - self._lib.host_shutdown(self._server_id) - self._server_id = 0 - raise RuntimeError("copilot_runtime_connection_open failed.") + if not self._connection_id: + self._outbound_callback = None + self._lib.host_shutdown(self._server_id) + self._server_id = 0 + raise RuntimeError("copilot_runtime_connection_open failed.") + finally: + with self._dispose_lock: + self._starting = False + disposed = self._disposed + if disposed: + self._try_finalize_cleanup() def _on_outbound( self, @@ -454,62 +470,81 @@ def _on_outbound( out before returning. Exceptions must not cross the FFI boundary, so everything is caught and logged. """ - with self._callback_lock: - if self._disposed: - return - self._active_callbacks += 1 + if self._disposed: + return try: if bytes_ptr and bytes_len > 0: data = ctypes.string_at(bytes_ptr, bytes_len) self._receive_buffer.feed(data) except Exception: # noqa: BLE001 logger.error("In-process FFI inbound callback failed", exc_info=True) - finally: - with self._callback_lock: - self._active_callbacks -= 1 def _write_frame(self, frame: bytes) -> None: - if self._disposed or not self._connection_id: + if self._disposed: raise RuntimeError("The in-process runtime connection is closed.") - ok = self._lib.connection_write(self._connection_id, frame, len(frame)) - if not ok: - raise RuntimeError("Failed to write a frame to the in-process runtime connection.") + with self._operation_lock: + if self._disposed or not self._connection_id: + raise RuntimeError("The in-process runtime connection is closed.") + ok = self._lib.connection_write(self._connection_id, frame, len(frame)) + if not ok: + raise RuntimeError("Failed to write a frame to the in-process runtime connection.") def dispose(self) -> None: """Close the FFI connection, shut down the native host, release resources. - Idempotent. Waits for any in-flight outbound callback to finish before - dropping the callback reference to avoid a use-after-free. + Idempotent. Callback state remains rooted until connection_close reports + that native callbacks are quiescent. """ with self._dispose_lock: if self._disposed: return self._disposed = True - # Stop accepting new callbacks and wait for in-flight ones to drain. - with self._callback_lock: - pass # _disposed is set; new callbacks bail out immediately. - while True: - with self._callback_lock: - if self._active_callbacks == 0: - break - time.sleep(0.001) - - try: - if self._connection_id: - self._lib.connection_close(self._connection_id) - self._connection_id = 0 - except Exception: # noqa: BLE001 - logger.debug("Error closing in-process FFI connection", exc_info=True) - - try: - if self._server_id: - self._lib.host_shutdown(self._server_id) - self._server_id = 0 - except Exception: # noqa: BLE001 - logger.debug("Error shutting down in-process FFI host", exc_info=True) - self._receive_buffer.close() - # Safe to drop now: no native code can invoke the callback after - # connection_close, and all in-flight callbacks have drained. - self._outbound_callback = None + with self._dispose_lock: + starting = self._starting + if not starting: + self._try_finalize_cleanup() + + def _try_finalize_cleanup(self) -> None: + with self._dispose_lock: + if self._cleanup_timer is not None: + return + with self._operation_lock: + if self._connection_id: + try: + closed = self._lib.connection_close(self._connection_id) + except Exception: # noqa: BLE001 + logger.debug("Error closing in-process FFI connection", exc_info=True) + self._quarantined_hosts.add(self) + return + if not closed: + self._schedule_cleanup_retry() + return + self._connection_id = 0 + self._outbound_callback = None + self._quarantined_hosts.discard(self) + + if self._server_id: + try: + if not self._lib.host_shutdown(self._server_id): + logger.debug( + "In-process FFI host shutdown did not recognize server %s", + self._server_id, + ) + except Exception: # noqa: BLE001 + logger.debug("Error shutting down in-process FFI host", exc_info=True) + self._server_id = 0 + + def _schedule_cleanup_retry(self) -> None: + if self._cleanup_timer is not None: + return + timer = threading.Timer(_CLEANUP_RETRY_INTERVAL_SECONDS, self._run_cleanup_retry) + timer.daemon = True + self._cleanup_timer = timer + timer.start() + + def _run_cleanup_retry(self) -> None: + with self._dispose_lock: + self._cleanup_timer = None + self._try_finalize_cleanup() diff --git a/python/test_ffi_runtime_host.py b/python/test_ffi_runtime_host.py new file mode 100644 index 0000000000..5806afa1e1 --- /dev/null +++ b/python/test_ffi_runtime_host.py @@ -0,0 +1,56 @@ +import threading +import time +from unittest.mock import patch + +from copilot._ffi_runtime_host import FfiRuntimeHost + + +class _TestLibrary: + def __init__(self) -> None: + self.allow_close = False + self.close_calls = 0 + self.shutdown_calls = 0 + self.shutdown = threading.Event() + + def connection_close(self, _connection_id: int) -> bool: + self.close_calls += 1 + return self.allow_close + + def host_shutdown(self, _server_id: int) -> bool: + self.shutdown_calls += 1 + self.shutdown.set() + return True + + +def test_dispose_retains_callback_until_connection_close_succeeds(): + library = _TestLibrary() + with ( + patch("copilot._ffi_runtime_host._load_library", return_value=library), + patch("copilot._ffi_runtime_host._CLEANUP_RETRY_INTERVAL_SECONDS", 0.01), + ): + host = FfiRuntimeHost("test-runtime", None) + callback = object() + host._server_id = 11 + host._connection_id = 21 + host._outbound_callback = callback + + host.dispose() + + assert host._outbound_callback is callback + assert host._connection_id == 21 + assert library.close_calls == 1 + assert library.shutdown_calls == 0 + + library.allow_close = True + assert library.shutdown.wait(5), "Deferred native cleanup did not complete" + + assert host._outbound_callback is None + assert host._connection_id == 0 + assert library.close_calls >= 2 + assert library.shutdown_calls == 1 + + close_calls_after_cleanup = library.close_calls + host.dispose() + time.sleep(0.05) + assert library.close_calls == close_calls_after_cleanup + assert library.shutdown_calls == 1 diff --git a/rust/src/ffi.rs b/rust/src/ffi.rs index a25990062c..02fce3f030 100644 --- a/rust/src/ffi.rs +++ b/rust/src/ffi.rs @@ -12,14 +12,14 @@ use std::collections::HashMap; use std::ffi::c_void; use std::path::{Path, PathBuf}; use std::pin::Pin; -use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicU32, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicU32, Ordering}; use std::sync::{Arc, OnceLock}; use std::task::{Context, Poll}; use libloading::Library; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; use tokio::sync::mpsc; -use tracing::debug; +use tracing::{debug, warn}; use crate::{Error, ErrorKind}; @@ -45,7 +45,6 @@ type ConnectionCloseFn = unsafe extern "C" fn(u32) -> bool; /// route inbound frames back to the reader. struct CallbackState { tx: mpsc::UnboundedSender>, - active_callbacks: AtomicUsize, closing: AtomicBool, } @@ -54,14 +53,11 @@ extern "C" fn on_outbound(user_data: *mut c_void, bytes: *const u8, len: usize) return; } let state = unsafe { &*(user_data as *const CallbackState) }; - state.active_callbacks.fetch_add(1, Ordering::SeqCst); if state.closing.load(Ordering::SeqCst) { - state.active_callbacks.fetch_sub(1, Ordering::SeqCst); return; } let slice = unsafe { std::slice::from_raw_parts(bytes, len) }; let _ = state.tx.send(slice.to_vec()); - state.active_callbacks.fetch_sub(1, Ordering::SeqCst); } /// Bound exports and connection lifecycle state, shared between the @@ -98,24 +94,56 @@ impl FfiShared { if !state.is_null() { unsafe { &*state }.closing.store(true, Ordering::SeqCst); } - let conn = self.connection_id.swap(0, Ordering::SeqCst); + let conn = self.connection_id.load(Ordering::SeqCst); if conn != 0 { - unsafe { (self.connection_close)(conn) }; + let quiesced = unsafe { (self.connection_close)(conn) }; + if !quiesced { + let conn = self.connection_id.swap(0, Ordering::SeqCst); + let server = self.server_id.swap(0, Ordering::SeqCst); + let state = + self.callback_state + .swap(std::ptr::null_mut(), Ordering::SeqCst) as usize; + let connection_close = self.connection_close; + let host_shutdown = self.host_shutdown; + let library_path = self.library_path.clone(); + if let Err(error) = std::thread::Builder::new() + .name("copilot-ffi-cleanup".to_owned()) + .spawn(move || { + while !unsafe { connection_close(conn) } { + std::thread::sleep(std::time::Duration::from_millis(100)); + } + release_callback_state(state); + if server != 0 && !unsafe { host_shutdown(server) } { + warn!( + library = %library_path.display(), + server_id = server, + "FFI runtime host shutdown did not recognize server" + ); + } + debug!(library = %library_path.display(), "FFI runtime connection closed"); + }) + { + warn!( + error = %error, + library = %self.library_path.display(), + "failed to start deferred FFI cleanup thread; callback state retained" + ); + } + return; + } + self.connection_id.store(0, Ordering::SeqCst); } let server = self.server_id.swap(0, Ordering::SeqCst); - if server != 0 { - unsafe { (self.host_shutdown)(server) }; - } - // Free the callback state only after the connection is closed and the - // host is shut down, so native can no longer invoke the callback. let state = self .callback_state - .swap(std::ptr::null_mut(), Ordering::SeqCst); - if !state.is_null() { - while unsafe { &*state }.active_callbacks.load(Ordering::SeqCst) != 0 { - std::thread::yield_now(); - } - drop(unsafe { Box::from_raw(state) }); + .swap(std::ptr::null_mut(), Ordering::SeqCst) as usize; + release_callback_state(state); + if server != 0 && !unsafe { (self.host_shutdown)(server) } { + warn!( + library = %self.library_path.display(), + server_id = server, + "FFI runtime host shutdown did not recognize server" + ); } debug!(library = %self.library_path.display(), "FFI runtime connection closed"); } @@ -125,6 +153,7 @@ impl FfiShared { if self.closed.load(Ordering::SeqCst) { return false; } + let conn = self.connection_id.load(Ordering::SeqCst); if conn == 0 { return false; @@ -133,6 +162,14 @@ impl FfiShared { } } +fn release_callback_state(state: usize) { + if state == 0 { + return; + } + let state = state as *mut CallbackState; + drop(unsafe { Box::from_raw(state) }); +} + impl Drop for FfiShared { fn drop(&mut self) { self.close(); @@ -316,7 +353,6 @@ impl FfiHost { let (tx, rx) = mpsc::unbounded_channel::>(); let state_ptr = Box::into_raw(Box::new(CallbackState { tx, - active_callbacks: AtomicUsize::new(0), closing: AtomicBool::new(false), })); let connection_id = unsafe { @@ -559,8 +595,35 @@ fn build_env_json(environment: &[(String, String)]) -> Option> { #[cfg(test)] mod tests { + use std::sync::Mutex; + use std::sync::atomic::AtomicUsize; + use std::time::{Duration, Instant}; + use super::*; + static FFI_LIFECYCLE_TEST_LOCK: Mutex<()> = Mutex::new(()); + static TEST_ALLOW_CLOSE: AtomicBool = AtomicBool::new(false); + static TEST_CLOSE_CALLS: AtomicUsize = AtomicUsize::new(0); + static TEST_SHUTDOWN_CALLS: AtomicUsize = AtomicUsize::new(0); + + unsafe extern "C" fn test_host_shutdown(_server_id: u32) -> bool { + TEST_SHUTDOWN_CALLS.fetch_add(1, Ordering::SeqCst); + true + } + + unsafe extern "C" fn test_connection_write( + _connection_id: u32, + _bytes: *const u8, + _length: usize, + ) -> bool { + true + } + + unsafe extern "C" fn test_connection_close(_connection_id: u32) -> bool { + TEST_CLOSE_CALLS.fetch_add(1, Ordering::SeqCst); + TEST_ALLOW_CLOSE.load(Ordering::SeqCst) + } + #[test] fn argv_without_entrypoint_contains_only_client_options() { let argv: Vec = serde_json::from_slice(&build_argv_json( @@ -620,4 +683,62 @@ mod tests { }) ); } + + #[test] + fn callback_state_is_retained_until_connection_close_succeeds() { + let _guard = FFI_LIFECYCLE_TEST_LOCK.lock().unwrap(); + TEST_ALLOW_CLOSE.store(false, Ordering::SeqCst); + TEST_CLOSE_CALLS.store(0, Ordering::SeqCst); + TEST_SHUTDOWN_CALLS.store(0, Ordering::SeqCst); + + let (tx, mut rx) = mpsc::unbounded_channel(); + let state_ptr = Box::into_raw(Box::new(CallbackState { + tx, + closing: AtomicBool::new(false), + })); + let shared = FfiShared { + host_shutdown: test_host_shutdown, + connection_write: test_connection_write, + connection_close: test_connection_close, + server_id: AtomicU32::new(11), + connection_id: AtomicU32::new(21), + callback_state: AtomicPtr::new(state_ptr), + closed: AtomicBool::new(false), + operation_lock: parking_lot::Mutex::new(()), + library_path: PathBuf::from("test-runtime"), + }; + + shared.close(); + + assert_eq!(TEST_CLOSE_CALLS.load(Ordering::SeqCst), 1); + assert_eq!(TEST_SHUTDOWN_CALLS.load(Ordering::SeqCst), 0); + assert_eq!(shared.connection_id.load(Ordering::SeqCst), 0); + assert!(shared.callback_state.load(Ordering::SeqCst).is_null()); + + assert!(matches!( + rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + + TEST_ALLOW_CLOSE.store(true, Ordering::SeqCst); + let deadline = Instant::now() + Duration::from_secs(5); + while Instant::now() < deadline && TEST_SHUTDOWN_CALLS.load(Ordering::SeqCst) == 0 { + std::thread::sleep(Duration::from_millis(10)); + } + + assert!(TEST_CLOSE_CALLS.load(Ordering::SeqCst) >= 2); + assert_eq!(TEST_SHUTDOWN_CALLS.load(Ordering::SeqCst), 1); + assert!(matches!( + rx.try_recv(), + Err(mpsc::error::TryRecvError::Disconnected) + )); + + let close_calls_after_cleanup = TEST_CLOSE_CALLS.load(Ordering::SeqCst); + shared.close(); + assert_eq!( + TEST_CLOSE_CALLS.load(Ordering::SeqCst), + close_calls_after_cleanup + ); + assert_eq!(TEST_SHUTDOWN_CALLS.load(Ordering::SeqCst), 1); + } }