fix(tunnel): 完善隧道连接生命周期清理

src/Core/CustomWebSocketStream.cs

@@ -4,6 +4,8 @@ namespace Core;
 
 public class CustomWebSocketStream(WebSocket webSocket) : Stream
 {
+    private int _disposed;
+
     public override bool CanRead => true;
     public override bool CanSeek => false;
     public override bool CanWrite => true;
@@ -24,15 +26,38 @@ public class CustomWebSocketStream(WebSocket webSocket) : Stream
     public override void SetLength(long value) => throw new NotSupportedException();
     public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException();
 
-    public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
+    public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
     {
-        return webSocket.SendAsync(buffer, WebSocketMessageType.Binary, endOfMessage: false, cancellationToken);
+        try
+        {
+            await webSocket.SendAsync(buffer, WebSocketMessageType.Binary, endOfMessage: false, cancellationToken)
+                .ConfigureAwait(false);
+        }
+        catch
+        {
+            await DisposeAsync().ConfigureAwait(false);
+            throw;
+        }
     }
 
     public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
     {
-        var result = await webSocket.ReceiveAsync(buffer, cancellationToken);
-        return result.MessageType == WebSocketMessageType.Close ? 0 : result.Count;
+        try
+        {
+            var result = await webSocket.ReceiveAsync(buffer, cancellationToken).ConfigureAwait(false);
+            if (result.MessageType == WebSocketMessageType.Close)
+            {
+                await DisposeAsync().ConfigureAwait(false);
+                return 0;
+            }
+
+            return result.Count;
+        }
+        catch
+        {
+            await DisposeAsync().ConfigureAwait(false);
+            throw;
+        }
     }
 
     public override Task FlushAsync(CancellationToken cancellationToken)
@@ -42,18 +67,41 @@ public class CustomWebSocketStream(WebSocket webSocket) : Stream
 
     public override async ValueTask DisposeAsync()
     {
-        if (webSocket.State == WebSocketState.Open)
-        {
-            using var timeoutTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(1d));
-            await webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, null, timeoutTokenSource.Token)
-                .ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
-        }
+        if (Interlocked.Exchange(ref _disposed, 1) != 0)
+            return;
 
-        webSocket.Dispose();
+        try
+        {
+            if (webSocket.State == WebSocketState.Open)
+            {
+                using var timeoutTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(1d));
+                await webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, null, timeoutTokenSource.Token)
+                    .ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
+            }
+        }
+        finally
+        {
+            webSocket.Dispose();
+        }
     }
 
     protected override void Dispose(bool disposing)
     {
-        this.DisposeAsync().ConfigureAwait(false).GetAwaiter().GetResult();
+        if (Interlocked.Exchange(ref _disposed, 1) != 0)
+            return;
+
+        try
+        {
+            if (webSocket.State == WebSocketState.Open)
+            {
+                using var timeoutTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(1d));
+                webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, null, timeoutTokenSource.Token)
+                    .ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing).GetAwaiter().GetResult();
+            }
+        }
+        finally
+        {
+            webSocket.Dispose();
+        }
     }
 }

src/FastGateway/Cluster/ClusterTunnelHub.cs

@@ -45,9 +45,14 @@ public static class ClusterTunnelHub
             var httpTunnel = new HttpTunnel(tunnelStream, tunnelId, feature.Protocol, logger);
 
             if (Tunnels.SetResult(httpTunnel))
-                await httpTunnel.Closed;
+            {
+                try { await httpTunnel.Closed; }
+                finally { await httpTunnel.DisposeAsync(); }
+            }
             else
-                httpTunnel.Dispose();
+            {
+                await httpTunnel.DisposeAsync();
+            }
 
             return;
         }
@@ -73,14 +78,23 @@ public static class ClusterTunnelHub
             var connection = new AgentClientConnection(host, stream, new ConnectionConfig(), logger);
 
             await using var client = new AgentClient(connection, Tunnels, context);
-            if (await Clients.AddAsync(client, default))
+            var added = await Clients.AddAsync(client, CancellationToken.None);
+            try
             {
-                logger.LogInformation("集群隧道节点已注册：{Host}", host);
-
-                await connection.WaitForCloseAsync();
-                await Clients.RemoveAsync(client, default);
-
-                logger.LogInformation("集群隧道节点已断开：{Host}", host);
+                if (added)
+                {
+                    logger.LogInformation("集群隧道节点已注册：{Host}", host);
+                    await connection.WaitForCloseAsync();
+                }
+            }
+            finally
+            {
+                if (added)
+                {
+                    try { await Clients.RemoveAsync(client, CancellationToken.None); }
+                    catch { /* ignored */ }
+                    logger.LogInformation("集群隧道节点已断开：{Host}", host);
+                }
             }
         }
         catch (Exception ex)

src/FastGateway/Tunnels/AgentClientConnection.cs

@@ -1,5 +1,6 @@
 using System.Buffers;
 using System.Text;
+using FastGateway.Infrastructure;
 
 namespace FastGateway.Tunnels;
 
@@ -11,17 +12,23 @@ public sealed partial class AgentClientConnection : IAsyncDisposable
     private static readonly ReadOnlyMemory<byte> PingLine = "PING\r\n"u8.ToArray();
     private static readonly ReadOnlyMemory<byte> PongLine = "PONG\r\n"u8.ToArray();
     private readonly CancellationTokenSource _disposeTokenSource = new();
+    private readonly CancellationToken _disposedToken;
     private readonly TimeSpan _keepAliveTimeout;
     private readonly Timer? _keepAliveTimer;
     private readonly ILogger _logger;
     private readonly Stream _stream;
+    private readonly object _lifecycleSync = new();
+    private readonly HashSet<HttpTunnel> _httpTunnels = [];
     private int _httpTunnelCount;
+    private int _disposed;
 
     public AgentClientConnection(string clientId, Stream stream, ConnectionConfig config, ILogger logger)
     {
         ClientId = clientId;
         _stream = stream;
         _logger = logger;
+        _disposedToken = _disposeTokenSource.Token;
+        GatewayResourceMetrics.ControlConnectionOpened();
 
         var keepAliveInterval = config.KeepAliveInterval;
         if (config.KeepAlive && keepAliveInterval > TimeSpan.Zero)
@@ -37,33 +44,54 @@ public sealed partial class AgentClientConnection : IAsyncDisposable
 
     public string ClientId { get; }
 
-    public int HttpTunnelCount => _httpTunnelCount;
+    internal CancellationToken DisposedToken => _disposedToken;
+
+    public int HttpTunnelCount => Volatile.Read(ref _httpTunnelCount);
 
     public async ValueTask DisposeAsync()
     {
-        if (!_disposeTokenSource.IsCancellationRequested)
+        if (Interlocked.Exchange(ref _disposed, 1) != 0)
+            return;
+
+        // 先取消未完成的远端隧道创建，避免控制连接关闭后等待者继续占用请求资源。
+        try
         {
-            // 取消回调会关闭底层 WebSocket/HttpContext；若对端已先断开，这些对象可能已释放，
-            // Cancel 会把 ObjectDisposedException 从回调透传出来，这里吞掉保证清理幂等
-            try
-            {
-                _disposeTokenSource.Cancel();
-            }
-            catch
-            {
-                // ignored
-            }
+            _disposeTokenSource.Cancel();
+        }
+        catch
+        {
+            // ignored
+        }
 
-            _disposeTokenSource.Dispose();
+        _keepAliveTimer?.Dispose();
+        GatewayResourceMetrics.ControlConnectionClosed();
 
-            try
-            {
-                await _stream.DisposeAsync();
-            }
-            catch
-            {
-                // ignored
-            }
+        HttpTunnel[] tunnels;
+        lock (_lifecycleSync)
+        {
+            tunnels = _httpTunnels.ToArray();
+        }
+
+        foreach (var tunnel in tunnels)
+        {
+            try { await tunnel.DisposeAsync(); }
+            catch { /* ignored */ }
+        }
+
+        lock (_lifecycleSync)
+        {
+            _httpTunnels.Clear();
+        }
+
+        _disposeTokenSource.Dispose();
+
+        try
+        {
+            await _stream.DisposeAsync();
+        }
+        catch
+        {
+            // ignored
         }
     }
 
@@ -87,6 +115,9 @@ public sealed partial class AgentClientConnection : IAsyncDisposable
     public async Task CreateHttpTunnelAsync(Guid tunnelId,
         CancellationToken cancellationToken)
     {
+        using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _disposedToken);
+        cancellationToken = linkedCts.Token;
+
         const int size = 64;
         var tunnelIdLine = $"{tunnelId}\r\n";
 
@@ -97,22 +128,33 @@ public sealed partial class AgentClientConnection : IAsyncDisposable
         await _stream.WriteAsync(buffer, cancellationToken);
     }
 
-    public int IncrementHttpTunnelCount()
+    public bool TryBindHttpTunnel(HttpTunnel tunnel)
     {
-        return Interlocked.Increment(ref _httpTunnelCount);
+        lock (_lifecycleSync)
+        {
+            if (_disposed != 0) return false;
+
+            _httpTunnelCount++;
+            _httpTunnels.Add(tunnel);
+            return true;
+        }
     }
 
-    public int DecrementHttpTunnelCount()
+    public int DecrementHttpTunnelCount(HttpTunnel tunnel)
     {
-        return Interlocked.Decrement(ref _httpTunnelCount);
+        lock (_lifecycleSync)
+        {
+            _httpTunnels.Remove(tunnel);
+            if (_httpTunnelCount <= 0) return 0;
+            return --_httpTunnelCount;
+        }
     }
 
     public async Task WaitForCloseAsync()
     {
         try
         {
-            var cancellationToken = _disposeTokenSource.Token;
-            await HandleConnectionAsync(cancellationToken);
+            await HandleConnectionAsync(_disposedToken);
         }
         catch (Exception)
         {

src/FastGateway/Tunnels/AgentClientManager.cs

@@ -9,10 +9,12 @@ namespace FastGateway.Tunnels;
 ///     客户端管理器
 /// </summary>
 [DebuggerDisplay("Count = {Count}")]
-public sealed class AgentClientManager : IEnumerable
+public sealed class AgentClientManager : IEnumerable, IAsyncDisposable
 {
     private readonly AgentStateChannel _clientStateChannel;
     private readonly ConcurrentDictionary<string, AgentClient> _dictionary = new();
+    private readonly object _lifecycleSync = new();
+    private int _disposed;
 
     public AgentClientManager(AgentStateChannel clientStateChannel)
     {
@@ -42,15 +44,39 @@ public sealed class AgentClientManager : IEnumerable
     /// <returns></returns>
     public async ValueTask<bool> AddAsync(AgentClient client, CancellationToken cancellationToken)
     {
-        var clientId = client.Id;
-        if (_dictionary.TryRemove(clientId.ToLowerInvariant(), out var existClient)) await existClient.DisposeAsync();
-
-        if (_dictionary.TryAdd(clientId.ToLowerInvariant(), client))
+        var clientId = client.Id.ToLowerInvariant();
+        AgentClient? existClient;
+        lock (_lifecycleSync)
         {
-            await _clientStateChannel.WriteAsync(client, true, cancellationToken);
-            return true;
+            if (_disposed != 0) return false;
+            _dictionary.TryRemove(clientId, out existClient);
+            if (!_dictionary.TryAdd(clientId, client)) return false;
         }
 
+        try
+        {
+            if (existClient is not null) await existClient.DisposeAsync();
+            await _clientStateChannel.WriteAsync(client, true, cancellationToken);
+
+            lock (_lifecycleSync)
+            {
+                if (_disposed == 0 && _dictionary.TryGetValue(clientId, out var current) &&
+                    ReferenceEquals(current, client))
+                    return true;
+            }
+        }
+        catch
+        {
+        }
+
+        lock (_lifecycleSync)
+        {
+            if (_dictionary.TryGetValue(clientId, out var current) && ReferenceEquals(current, client))
+                _dictionary.TryRemove(clientId, out _);
+        }
+
+        try { await client.DisposeAsync(); }
+        catch { /* ignored */ }
         return false;
     }
 
@@ -62,19 +88,17 @@ public sealed class AgentClientManager : IEnumerable
     /// <returns></returns>
     public async ValueTask<bool> RemoveAsync(AgentClient client, CancellationToken cancellationToken)
     {
-        var clientId = client.Id;
-        if (_dictionary.TryRemove(clientId.ToLowerInvariant(), out var existClient))
+        var clientId = client.Id.ToLowerInvariant();
+        lock (_lifecycleSync)
         {
-            if (ReferenceEquals(existClient, client))
-            {
-                await _clientStateChannel.WriteAsync(client, false, cancellationToken);
-                return true;
-            }
-
-            _dictionary.TryAdd(clientId, existClient);
+            if (!_dictionary.TryGetValue(clientId, out var existClient) ||
+                !ReferenceEquals(existClient, client) ||
+                !_dictionary.TryRemove(clientId, out _))
+                return false;
         }
 
-        return false;
+        await _clientStateChannel.WriteAsync(client, false, cancellationToken);
+        return true;
     }
 
 
@@ -83,4 +107,18 @@ public sealed class AgentClientManager : IEnumerable
     {
         foreach (var keyValue in _dictionary) yield return keyValue.Value;
     }
+
+    public async ValueTask DisposeAsync()
+    {
+        AgentClient[] clients;
+        lock (_lifecycleSync)
+        {
+            if (Interlocked.Exchange(ref _disposed, 1) != 0) return;
+            clients = _dictionary.Values.ToArray();
+            _dictionary.Clear();
+        }
+
+        foreach (var client in clients)
+            await client.DisposeAsync();
+    }
 }

src/FastGateway/Tunnels/AgentManagerMiddleware.cs

@@ -64,6 +64,7 @@ internal partial class AgentManagerMiddleware(
         var host = "node_" + nodeName;
 
         AgentClient? client = null;
+        var added = false;
         try
         {
             // 创建连接
@@ -77,14 +78,12 @@ internal partial class AgentManagerMiddleware(
             await using var _ = client;
             if (await agentClientManager.AddAsync(client, default))
             {
+                added = true;
                 Log.LogNodeConnected(logger, nodeName, heartbeatSeconds);
                 TunnelClientProxy.OnConnected(nodeName, client, server);
 
                 // 等待连接关闭
                 await connection.WaitForCloseAsync();
-
-                // 从客户管理器中移除客户端对象
-                await agentClientManager.RemoveAsync(client, default);
             }
         }
         catch (Exception e)
@@ -96,6 +95,12 @@ internal partial class AgentManagerMiddleware(
         {
             if (client != null)
             {
+                if (added)
+                {
+                    try { await agentClientManager.RemoveAsync(client, CancellationToken.None); }
+                    catch { /* ignored */ }
+                }
+
                 // 更新离线状态与最近连接时间，并清理该节点的 YARP 路由
                 TunnelClientProxy.OnDisconnected(nodeName, client);
                 Log.LogNodeDisconnected(logger, nodeName);

src/FastGateway/Tunnels/AgentManagerTunnelMiddleware.cs

@@ -35,9 +35,14 @@ public class AgentManagerTunnelMiddleware(
             var httpTunnel = new HttpTunnel(stream, tunnelId, requiredFeature.Protocol, logger);
 
             if (agentTunnelFactory.SetResult(httpTunnel))
-                await httpTunnel.Closed;
+            {
+                try { await httpTunnel.Closed; }
+                finally { await httpTunnel.DisposeAsync(); }
+            }
             else
-                httpTunnel.Dispose();
+            {
+                await httpTunnel.DisposeAsync();
+            }
         }
     }
 }

src/FastGateway/Tunnels/AgentTunnelFactory.cs

@@ -9,8 +9,7 @@ namespace FastGateway.Tunnels;
 /// <param name="logger"></param>
 public partial class AgentTunnelFactory(ILogger<AgentTunnelFactory> logger)
 {
-    private readonly ConcurrentDictionary<Guid, TaskCompletionSource<HttpTunnel>> _httpTunnelCompletionSources =
-        new();
+    private readonly ConcurrentDictionary<Guid, PendingHttpTunnel> _httpTunnelCompletionSources = new();
 
     /// <summary>
     ///     创建HttpTunnel
@@ -23,23 +22,38 @@ public partial class AgentTunnelFactory(ILogger<AgentTunnelFactory> logger)
         CancellationToken cancellationToken)
     {
         var tunnelId = Guid.NewGuid();
-        var httpTunnelSource = new TaskCompletionSource<HttpTunnel>();
-        if (!_httpTunnelCompletionSources.TryAdd(tunnelId, httpTunnelSource))
+        var pendingTunnel = new PendingHttpTunnel();
+        if (!_httpTunnelCompletionSources.TryAdd(tunnelId, pendingTunnel))
             throw new SystemException($"系统中已存在{tunnelId}的tunnelId");
 
+        HttpTunnel? ownedTunnel = null;
         try
         {
+            using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(
+                cancellationToken, connection.DisposedToken);
+            cancellationToken = linkedCts.Token;
+
             var stopwatch = Stopwatch.StartNew();
             Log.LogTunnelCreating(logger, connection.ClientId, tunnelId);
             await connection.CreateHttpTunnelAsync(tunnelId, cancellationToken);
-            var httpTunnel = await httpTunnelSource.Task.WaitAsync(cancellationToken);
+            var httpTunnel = await pendingTunnel.Source.Task.WaitAsync(cancellationToken);
+            cancellationToken.ThrowIfCancellationRequested();
 
-            var httpTunnelCount = connection.IncrementHttpTunnelCount();
-            httpTunnel.BindConnection(connection);
+            if (!pendingTunnel.TryClaim(httpTunnel))
+            {
+                await httpTunnel.DisposeAsync();
+                throw new OperationCanceledException("控制连接已关闭");
+            }
+
+            ownedTunnel = httpTunnel;
+            cancellationToken.ThrowIfCancellationRequested();
+            if (!httpTunnel.BindConnection(connection))
+                throw new OperationCanceledException("控制连接已关闭");
 
             stopwatch.Stop();
             Log.LogTunnelCreateSuccess(logger, connection.ClientId, httpTunnel.Protocol, tunnelId,
-                stopwatch.Elapsed, httpTunnelCount);
+                stopwatch.Elapsed, connection.HttpTunnelCount);
+            ownedTunnel = null;
             return httpTunnel;
         }
         catch (OperationCanceledException)
@@ -54,7 +68,14 @@ public partial class AgentTunnelFactory(ILogger<AgentTunnelFactory> logger)
         }
         finally
         {
+            if (ownedTunnel is not null)
+            {
+                try { await ownedTunnel.DisposeAsync(); }
+                catch { /* ignored */ }
+            }
+
             _httpTunnelCompletionSources.TryRemove(tunnelId, out _);
+            pendingTunnel.Abandon();
         }
     }
 
@@ -65,8 +86,73 @@ public partial class AgentTunnelFactory(ILogger<AgentTunnelFactory> logger)
 
     public bool SetResult(HttpTunnel httpTunnel)
     {
-        return _httpTunnelCompletionSources.TryRemove(httpTunnel.Id, out var source) &&
-               source.TrySetResult(httpTunnel);
+        if (_httpTunnelCompletionSources.TryRemove(httpTunnel.Id, out var pending) &&
+            pending.TrySetResult(httpTunnel))
+            return true;
+
+        httpTunnel.Dispose();
+        return false;
+    }
+
+    private sealed class PendingHttpTunnel
+    {
+        private readonly object _sync = new();
+        private HttpTunnel? _tunnel;
+        private bool _abandoned;
+        private bool _claimed;
+        private bool _completed;
+
+        public TaskCompletionSource<HttpTunnel> Source { get; } =
+            new(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        public bool TrySetResult(HttpTunnel tunnel)
+        {
+            lock (_sync)
+            {
+                if (_abandoned || _completed || _claimed)
+                {
+                    tunnel.Dispose();
+                    return false;
+                }
+
+                _completed = true;
+                _tunnel = tunnel;
+                if (Source.TrySetResult(tunnel)) return true;
+
+                _tunnel = null;
+            }
+
+            tunnel.Dispose();
+            return false;
+        }
+
+        public bool TryClaim(HttpTunnel tunnel)
+        {
+            lock (_sync)
+            {
+                if (_abandoned || !_completed || !ReferenceEquals(_tunnel, tunnel)) return false;
+
+                _claimed = true;
+                _tunnel = null;
+                return true;
+            }
+        }
+
+        public void Abandon()
+        {
+            HttpTunnel? tunnel;
+            lock (_sync)
+            {
+                if (_abandoned) return;
+
+                _abandoned = true;
+                tunnel = _completed && !_claimed ? _tunnel : null;
+                _tunnel = null;
+                if (!_completed) Source.TrySetCanceled();
+            }
+
+            tunnel?.Dispose();
+        }
     }
 
     private static partial class Log

src/FastGateway/Tunnels/HttpTunnel.cs

@@ -1,4 +1,5 @@
 using Core;
+using FastGateway.Infrastructure;
 
 namespace FastGateway.Tunnels;
 
@@ -10,7 +11,11 @@ public sealed partial class HttpTunnel(Stream inner, Guid tunnelId, TransportPro
 {
     private readonly TaskCompletionSource _closeTaskCompletionSource = new();
     private readonly long _tickCout = Environment.TickCount64;
+    private readonly object _lifecycleSync = new();
     private AgentClientConnection? _connection;
+    private bool _counted;
+    private bool _closed;
+    private int _disposed;
 
     /// <summary>
     ///     等待HttpClient对其关闭
@@ -27,32 +32,53 @@ public sealed partial class HttpTunnel(Stream inner, Guid tunnelId, TransportPro
     /// </summary>
     public TransportProtocol Protocol { get; } = protocol;
 
-    public void BindConnection(AgentClientConnection connection)
+    public bool BindConnection(AgentClientConnection connection)
     {
-        _connection = connection;
+        lock (_lifecycleSync)
+        {
+            if (_closed || _counted || !connection.TryBindHttpTunnel(this)) return false;
+
+            _connection = connection;
+            _counted = true;
+            GatewayResourceMetrics.HttpTunnelOpened();
+            return true;
+        }
     }
 
-    public override ValueTask DisposeAsync()
+    public override async ValueTask DisposeAsync()
     {
-        SetClosedResult();
-        return Inner.DisposeAsync();
+        if (!SetClosedResult()) return;
+        await Inner.DisposeAsync();
     }
 
     protected override void Dispose(bool disposing)
     {
-        SetClosedResult();
+        if (!SetClosedResult()) return;
         Inner.Dispose();
     }
 
-    private void SetClosedResult()
+    private bool SetClosedResult()
     {
-        if (_closeTaskCompletionSource.TrySetResult())
+        int? httpTunnelCount;
+        AgentClientConnection? connection;
+        lock (_lifecycleSync)
         {
-            var httpTunnelCount = _connection?.DecrementHttpTunnelCount();
-            var lifeTime = TimeSpan.FromMilliseconds(Environment.TickCount64 - _tickCout);
-            Log.LogTunnelClosed(logger, _connection?.ClientId, Protocol, Id, lifeTime,
-                httpTunnelCount);
+            if (_disposed != 0) return false;
+
+            _disposed = 1;
+            _closed = true;
+            connection = _connection;
+            httpTunnelCount = _counted && connection is not null
+                ? connection.DecrementHttpTunnelCount(this)
+                : null;
+            if (_counted) GatewayResourceMetrics.HttpTunnelClosed();
+            _counted = false;
         }
+
+        _closeTaskCompletionSource.TrySetResult();
+        var lifeTime = TimeSpan.FromMilliseconds(Environment.TickCount64 - _tickCout);
+        Log.LogTunnelClosed(logger, connection?.ClientId, Protocol, Id, lifeTime, httpTunnelCount);
+        return true;
     }
 
     public override string ToString()

src/FastGateway/Tunnels/TunnelClientFactory.cs

@@ -10,10 +10,10 @@ namespace FastGateway.Tunnels;
 ///     隧道出站工厂：每个网关实例共用一份 handler（ConnectCallback 绑定本实例的 Agent）。
 ///     避免每条隧道路由各建一个连接池。
 /// </summary>
-internal class TunnelClientFactory(
+internal sealed class TunnelClientFactory(
     AgentClientManager agentClientManager,
     AgentTunnelFactory agentTunnelFactory)
-    : IForwarderHttpClientFactory
+    : IForwarderHttpClientFactory, IDisposable
 {
     private readonly SocketsHttpHandler _handler = CreateHandler(agentClientManager, agentTunnelFactory);
 
@@ -71,6 +71,11 @@ internal class TunnelClientFactory(
         return handler;
     }
 
+    public void Dispose()
+    {
+        _handler.Dispose();
+    }
+
     private static async ValueTask<Stream> DefaultConnectCallback(
         SocketsHttpConnectionContext context,
         CancellationToken cancellationToken)

src/FastGateway/Tunnels/WebSocketStream.cs

@@ -1,12 +1,21 @@
 using System.Net.WebSockets;
 using System.Threading.Tasks.Sources;
+using FastGateway.Infrastructure;
 
 namespace FastGateway.Tunnels;
 
-internal class WebSocketStream(WebSocket ws) : Stream, IValueTaskSource<object?>, ICloseable
+internal sealed class WebSocketStream : Stream, IValueTaskSource<object?>, ICloseable
 {
+    private readonly WebSocket _ws;
     private readonly object _sync = new();
     private ManualResetValueTaskSourceCore<object?> _tcs = new() { RunContinuationsAsynchronously = true };
+    private int _disposed;
+
+    public WebSocketStream(WebSocket ws)
+    {
+        _ws = ws;
+        GatewayResourceMetrics.WebSocketOpened();
+    }
 
     internal ValueTask<object?> StreamCompleteTask => new(this, _tcs.Version);
 
@@ -24,29 +33,24 @@ internal class WebSocketStream(WebSocket ws) : Stream, IValueTaskSource<object?>
         set => throw new NotSupportedException();
     }
 
-    public bool IsClosed => ws.State != WebSocketState.Open;
+    public bool IsClosed
+    {
+        get
+        {
+            try
+            {
+                return _ws.State != WebSocketState.Open;
+            }
+            catch (ObjectDisposedException)
+            {
+                return true;
+            }
+        }
+    }
 
     public void Abort()
     {
-        // Debug.Assert(!Thread.CurrentThread.IsThreadPoolThread);
-        ws.Abort();
-
-        // The shutdown path is currently synchronous but at least we're not blocking a threadpool thread
-        // Attempt a graceful close
-        //using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
-        //_ws.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "", timeout.Token).GetAwaiter().GetResult();
-        // Wait for closed to be sent back otherwise abort
-        //if (_ws.State != WebSocketState.Closed)
-        //{
-        //    _ws.Abort();
-        //}
-
-        lock (_sync)
-        {
-            if (GetStatus(_tcs.Version) != ValueTaskSourceStatus.Pending) return;
-
-            _tcs.SetResult(null);
-        }
+        CompleteAndDispose();
     }
 
     public object? GetResult(short token)
@@ -95,35 +99,70 @@ internal class WebSocketStream(WebSocket ws) : Stream, IValueTaskSource<object?>
         return Task.CompletedTask;
     }
 
-    public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
+    public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
     {
-        return ws.SendAsync(buffer, WebSocketMessageType.Binary, false, cancellationToken);
+        try
+        {
+            await _ws.SendAsync(buffer, WebSocketMessageType.Binary, false, cancellationToken);
+        }
+        catch
+        {
+            CompleteAndDispose();
+            throw;
+        }
     }
 
     public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
     {
-        var result = await ws.ReceiveAsync(buffer, cancellationToken);
+        ValueWebSocketReceiveResult result;
+        try
+        {
+            result = await _ws.ReceiveAsync(buffer, cancellationToken);
+        }
+        catch
+        {
+            CompleteAndDispose();
+            throw;
+        }
 
-        if (result.MessageType == WebSocketMessageType.Close) return 0;
+        if (result.MessageType == WebSocketMessageType.Close)
+        {
+            CompleteAndDispose();
+            return 0;
+        }
 
         return result.Count;
     }
 
     protected override void Dispose(bool disposing)
     {
+        CompleteAndDispose();
+    }
+
+    private void CompleteAndDispose()
+    {
+        if (Interlocked.Exchange(ref _disposed, 1) == 0)
+        {
+            try { _ws.Abort(); }
+            catch { /* ignore */ }
+            try { _ws.Dispose(); }
+            catch { /* ignore */ }
+            GatewayResourceMetrics.WebSocketClosed();
+        }
+
         lock (_sync)
         {
-            if (GetStatus(_tcs.Version) != ValueTaskSourceStatus.Pending) return;
-
-            // This might seem evil but we're using dispose to know if the stream
-            // has been given discarded by http client. We trigger the continuation and take back ownership
-            // of it here.
-            _tcs.SetResult(null);
+            if (_tcs.GetStatus(_tcs.Version) == ValueTaskSourceStatus.Pending)
+                _tcs.SetResult(null);
         }
     }
 
     public void Reset()
     {
-        _tcs.Reset();
+        lock (_sync)
+        {
+            if (_disposed != 0) return;
+            _tcs.Reset();
+        }
     }
 }