feat: 增加接收数据逻辑

Author: ArgoZhang

src/BootstrapBlazor/Services/TcpSocket/DefaultTcpSocketClient.cs

@@ -69,8 +69,7 @@ public async ValueTask<bool> ConnectAsync(IPEndPoint endPoint, CancellationToken
 
             if (IsAutoReceive)
             {
-                // 开始接收数据
-                _ = Task.Run(ReceiveAsync, token);
+                _ = Task.Run(AutoReceiveAsync);
             }
             ret = true;
         }
@@ -140,24 +139,57 @@ public async ValueTask<bool> SendAsync(ReadOnlyMemory<byte> data, CancellationTo
 
     public async ValueTask<Memory<byte>> ReceiveAsync(CancellationToken token = default)
     {
+        if (_client == null || !_client.Connected)
+        {
+            throw new InvalidOperationException($"TCP Socket is not connected {LocalEndPoint}");
+        }
+
         if (IsAutoReceive)
         {
             throw new InvalidOperationException("Cannot call ReceiveAsync when IsAutoReceive is true. Use the auto-receive mechanism instead.");
         }
 
-        if (_client == null || !_client.Connected)
+        using var block = MemoryPool<byte>.Shared.Rent(ReceiveBufferSize);
+        var buffer = block.Memory;
+        var len = await ReceiveCoreAsync(_client, buffer, token);
+        return buffer[0..len];
+    }
+
+    private async ValueTask AutoReceiveAsync()
+    {
+        _receiveCancellationTokenSource ??= new();
+        while (_receiveCancellationTokenSource is { IsCancellationRequested: false })
         {
-            throw new InvalidOperationException($"TCP Socket is not connected {LocalEndPoint}");
+            if (_client is not { Connected: true })
+            {
+                throw new InvalidOperationException($"TCP Socket is not connected {LocalEndPoint}");
+            }
+
+            using var block = MemoryPool<byte>.Shared.Rent(ReceiveBufferSize);
+            var buffer = block.Memory;
+            var len = await ReceiveCoreAsync(_client, buffer, _receiveCancellationTokenSource.Token);
+            if (len == 0)
+            {
+                break;
+            }
         }
+    }
 
-        var ret = Memory<byte>.Empty;
+    private async ValueTask<int> ReceiveCoreAsync(TcpClient client, Memory<byte> buffer, CancellationToken token)
+    {
+        var len = 0;
         try
         {
-            _receiveCancellationTokenSource ??= new();
-            using var block = MemoryPool<byte>.Shared.Rent(ReceiveBufferSize);
-            var buffer = block.Memory;
-            var stream = _client.GetStream();
-            var len = await stream.ReadAsync(buffer, _receiveCancellationTokenSource.Token);
+            var stream = client.GetStream();
+
+            var receiveToken = token;
+            if (ReceiveTimeout > 0)
+            {
+                // 设置接收超时时间
+                var receiveTokenSource = new CancellationTokenSource(ReceiveTimeout);
+                receiveToken = CancellationTokenSource.CreateLinkedTokenSource(receiveToken, receiveTokenSource.Token).Token;
+            }
+            len = await stream.ReadAsync(buffer, receiveToken);
             if (len == 0)
             {
                 // 远端主机关闭链路
@@ -174,69 +206,26 @@ public async ValueTask<Memory<byte>> ReceiveAsync(CancellationToken token = defa
 
                 if (_dataPackageHandler != null)
                 {
-                    await _dataPackageHandler.ReceiveAsync(buffer);
+                    await _dataPackageHandler.ReceiveAsync(buffer, receiveToken);
                 }
-
-                ret = buffer;
             }
         }
         catch (OperationCanceledException ex)
         {
-            Logger.LogWarning(ex, "TCP Socket receive operation was canceled from {LocalEndPoint} to {RemoteEndPoint}", LocalEndPoint, _remoteEndPoint);
-        }
-        catch (Exception ex)
-        {
-            Logger.LogError(ex, "TCP Socket receive failed from {LocalEndPoint} to {RemoteEndPoint}", LocalEndPoint, _remoteEndPoint);
-        }
-        return ret;
-    }
-
-    private async ValueTask ReceiveAsync()
-    {
-        _receiveCancellationTokenSource ??= new();
-        while (_receiveCancellationTokenSource is { IsCancellationRequested: false })
-        {
-            if (_client is not { Connected: true })
-            {
-                throw new InvalidOperationException($"TCP Socket is not connected {LocalEndPoint}");
-            }
-
-            try
-            {
-                using var block = MemoryPool<byte>.Shared.Rent(ReceiveBufferSize);
-                var buffer = block.Memory;
-                var stream = _client.GetStream();
-                var len = await stream.ReadAsync(buffer, _receiveCancellationTokenSource.Token);
-                if (len == 0)
-                {
-                    // 远端主机关闭链路
-                    Logger.LogInformation("TCP Socket {LocalEndPoint} received 0 data closed by {RemoteEndPoint}", LocalEndPoint, _remoteEndPoint);
-                    break;
-                }
-                else
-                {
-                    buffer = buffer[..len];
-
-                    if (ReceivedCallBack != null)
-                    {
-                        await ReceivedCallBack(buffer);
-                    }
-
-                    if (_dataPackageHandler != null)
-                    {
-                        await _dataPackageHandler.ReceiveAsync(buffer);
-                    }
-                }
-            }
-            catch (OperationCanceledException ex)
+            if (token.IsCancellationRequested)
             {
-                Logger.LogWarning(ex, "TCP Socket receive operation was canceled from {LocalEndPoint} to {RemoteEndPoint}", LocalEndPoint, _remoteEndPoint);
+                Logger.LogWarning(ex, "TCP Socket receive operation canceled from {LocalEndPoint} to {RemoteEndPoint}", LocalEndPoint, _remoteEndPoint);
             }
-            catch (Exception ex)
+            else
             {
-                Logger.LogError(ex, "TCP Socket receive failed from {LocalEndPoint} to {RemoteEndPoint}", LocalEndPoint, _remoteEndPoint);
+                Logger.LogWarning(ex, "TCP Socket receive operation timed out from {LocalEndPoint} to {RemoteEndPoint}", LocalEndPoint, _remoteEndPoint);
             }
         }
+        catch (Exception ex)
+        {
+            Logger.LogError(ex, "TCP Socket receive failed from {LocalEndPoint} to {RemoteEndPoint}", LocalEndPoint, _remoteEndPoint);
+        }
+        return len;
     }
 
     public void Close()