From 5e52fefd0fd99b7e01898ebfa77ad710d451f64f Mon Sep 17 00:00:00 2001 From: token Date: Mon, 31 Aug 2026 09:42:04 +0800 Subject: [PATCH] =?UTF-8?q?fix(tunnel):=20=E5=AE=8C=E5=96=84=E9=9A=A7?= =?UTF-8?q?=E9=81=93=E8=BF=9E=E6=8E=A5=E7=94=9F=E5=91=BD=E5=91=A8=E6=9C=9F?= =?UTF-8?q?=E6=B8=85=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/Core/CustomWebSocketStream.cs | 72 ++++++++++-- src/FastGateway/Cluster/ClusterTunnelHub.cs | 32 ++++-- .../Tunnels/AgentClientConnection.cs | 96 +++++++++++----- src/FastGateway/Tunnels/AgentClientManager.cs | 72 +++++++++--- .../Tunnels/AgentManagerMiddleware.cs | 11 +- .../Tunnels/AgentManagerTunnelMiddleware.cs | 9 +- src/FastGateway/Tunnels/AgentTunnelFactory.cs | 106 ++++++++++++++++-- src/FastGateway/Tunnels/HttpTunnel.cs | 50 +++++++-- .../Tunnels/TunnelClientFactory.cs | 9 +- src/FastGateway/Tunnels/WebSocketStream.cs | 103 +++++++++++------ 10 files changed, 434 insertions(+), 126 deletions(-) diff --git a/src/Core/CustomWebSocketStream.cs b/src/Core/CustomWebSocketStream.cs index bb9f8cd..ac2a835 100644 --- a/src/Core/CustomWebSocketStream.cs +++ b/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 buffer, CancellationToken cancellationToken = default) + public override async ValueTask WriteAsync(ReadOnlyMemory 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 ReadAsync(Memory 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(); + } } } \ No newline at end of file diff --git a/src/FastGateway/Cluster/ClusterTunnelHub.cs b/src/FastGateway/Cluster/ClusterTunnelHub.cs index 5bbbd10..4447197 100644 --- a/src/FastGateway/Cluster/ClusterTunnelHub.cs +++ b/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) diff --git a/src/FastGateway/Tunnels/AgentClientConnection.cs b/src/FastGateway/Tunnels/AgentClientConnection.cs index a371d91..77d0ade 100644 --- a/src/FastGateway/Tunnels/AgentClientConnection.cs +++ b/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 PingLine = "PING\r\n"u8.ToArray(); private static readonly ReadOnlyMemory 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 _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) { diff --git a/src/FastGateway/Tunnels/AgentClientManager.cs b/src/FastGateway/Tunnels/AgentClientManager.cs index 8e7f28a..55ce53c 100644 --- a/src/FastGateway/Tunnels/AgentClientManager.cs +++ b/src/FastGateway/Tunnels/AgentClientManager.cs @@ -9,10 +9,12 @@ namespace FastGateway.Tunnels; /// 客户端管理器 /// [DebuggerDisplay("Count = {Count}")] -public sealed class AgentClientManager : IEnumerable +public sealed class AgentClientManager : IEnumerable, IAsyncDisposable { private readonly AgentStateChannel _clientStateChannel; private readonly ConcurrentDictionary _dictionary = new(); + private readonly object _lifecycleSync = new(); + private int _disposed; public AgentClientManager(AgentStateChannel clientStateChannel) { @@ -42,15 +44,39 @@ public sealed class AgentClientManager : IEnumerable /// public async ValueTask 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 /// public async ValueTask 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(); + } } \ No newline at end of file diff --git a/src/FastGateway/Tunnels/AgentManagerMiddleware.cs b/src/FastGateway/Tunnels/AgentManagerMiddleware.cs index 8e1501b..c6cf5a8 100644 --- a/src/FastGateway/Tunnels/AgentManagerMiddleware.cs +++ b/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); diff --git a/src/FastGateway/Tunnels/AgentManagerTunnelMiddleware.cs b/src/FastGateway/Tunnels/AgentManagerTunnelMiddleware.cs index 97aed3c..3fb1f06 100644 --- a/src/FastGateway/Tunnels/AgentManagerTunnelMiddleware.cs +++ b/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(); + } } } } \ No newline at end of file diff --git a/src/FastGateway/Tunnels/AgentTunnelFactory.cs b/src/FastGateway/Tunnels/AgentTunnelFactory.cs index 9982bb0..78a67a7 100644 --- a/src/FastGateway/Tunnels/AgentTunnelFactory.cs +++ b/src/FastGateway/Tunnels/AgentTunnelFactory.cs @@ -9,8 +9,7 @@ namespace FastGateway.Tunnels; /// public partial class AgentTunnelFactory(ILogger logger) { - private readonly ConcurrentDictionary> _httpTunnelCompletionSources = - new(); + private readonly ConcurrentDictionary _httpTunnelCompletionSources = new(); /// /// 创建HttpTunnel @@ -23,23 +22,38 @@ public partial class AgentTunnelFactory(ILogger logger) CancellationToken cancellationToken) { var tunnelId = Guid.NewGuid(); - var httpTunnelSource = new TaskCompletionSource(); - 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 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 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 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 diff --git a/src/FastGateway/Tunnels/HttpTunnel.cs b/src/FastGateway/Tunnels/HttpTunnel.cs index c86c873..dbd1ac2 100644 --- a/src/FastGateway/Tunnels/HttpTunnel.cs +++ b/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; /// /// 等待HttpClient对其关闭 @@ -27,32 +32,53 @@ public sealed partial class HttpTunnel(Stream inner, Guid tunnelId, TransportPro /// 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() diff --git a/src/FastGateway/Tunnels/TunnelClientFactory.cs b/src/FastGateway/Tunnels/TunnelClientFactory.cs index 7a7e4f1..656134e 100644 --- a/src/FastGateway/Tunnels/TunnelClientFactory.cs +++ b/src/FastGateway/Tunnels/TunnelClientFactory.cs @@ -10,10 +10,10 @@ namespace FastGateway.Tunnels; /// 隧道出站工厂：每个网关实例共用一份 handler（ConnectCallback 绑定本实例的 Agent）。 /// 避免每条隧道路由各建一个连接池。 /// -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 DefaultConnectCallback( SocketsHttpConnectionContext context, CancellationToken cancellationToken) diff --git a/src/FastGateway/Tunnels/WebSocketStream.cs b/src/FastGateway/Tunnels/WebSocketStream.cs index 5ff5b6d..f45768b 100644 --- a/src/FastGateway/Tunnels/WebSocketStream.cs +++ b/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, ICloseable +internal sealed class WebSocketStream : Stream, IValueTaskSource, ICloseable { + private readonly WebSocket _ws; private readonly object _sync = new(); private ManualResetValueTaskSourceCore _tcs = new() { RunContinuationsAsynchronously = true }; + private int _disposed; + + public WebSocketStream(WebSocket ws) + { + _ws = ws; + GatewayResourceMetrics.WebSocketOpened(); + } internal ValueTask StreamCompleteTask => new(this, _tcs.Version); @@ -24,29 +33,24 @@ internal class WebSocketStream(WebSocket ws) : Stream, IValueTaskSource 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 return Task.CompletedTask; } - public override ValueTask WriteAsync(ReadOnlyMemory buffer, CancellationToken cancellationToken = default) + public override async ValueTask WriteAsync(ReadOnlyMemory 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 ReadAsync(Memory 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(); + } } } \ No newline at end of file