Support LastActiveTime for KestrelPipeConnection

kerryjiang · Mar 30, 2024 · 8eb89fe · 8eb89fe
1 parent 9f8ae75
commit 8eb89fe
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 12 deletions.
diff --git a/src/SuperSocket.Connection/PipeConnection.cs b/src/SuperSocket.Connection/PipeConnection.cs
@@ -3,9 +3,8 @@
 using System.Threading.Tasks;
 using System.Net.Sockets;
 using System.Buffers;
-using System.Collections.Generic;
-using SuperSocket.ProtoBase;
 using System.IO.Pipelines;
+using System.Runtime.InteropServices;
 
 namespace SuperSocket.Connection
 {
@@ -109,7 +108,7 @@ internal virtual async Task FillPipeAsync(PipeWriter writer, ISupplyController s
                         break;
                     }
 
-                    LastActiveTime = DateTimeOffset.Now;
+                    UpdateLastActiveTime();
 
                     // Tell the PipeWriter how much was read
                     writer.Advance(bytesRead);
@@ -164,7 +163,7 @@ protected async ValueTask<bool> ProcessOutputRead(PipeReader reader)
                 try
                 {
                     await SendOverIOAsync(buffer, CancellationToken.None).ConfigureAwait(false); ;
-                    LastActiveTime = DateTimeOffset.Now;
+                    UpdateLastActiveTime();
                 }
                 catch (Exception e)
                 {
@@ -184,6 +183,15 @@ protected async ValueTask<bool> ProcessOutputRead(PipeReader reader)
 
         protected abstract ValueTask<int> SendOverIOAsync(ReadOnlySequence<byte> buffer, CancellationToken cancellationToken);
 
+        protected internal ArraySegment<T> GetArrayByMemory<T>(ReadOnlyMemory<T> memory)
+        {
+            if (!MemoryMarshal.TryGetArray(memory, out var result))
+            {
+                throw new InvalidOperationException("Buffer backed by array was expected");
+            }
+
+            return result;
+        }
         protected override bool IsIgnorableException(Exception e)
         {
             if (base.IsIgnorableException(e))

diff --git a/src/SuperSocket.Connection/PipeConnectionBase.cs b/src/SuperSocket.Connection/PipeConnectionBase.cs
@@ -4,7 +4,6 @@
 using System.Threading.Tasks;
 using System.IO.Pipelines;
 using System.Collections.Generic;
-using System.Runtime.InteropServices;
 using Microsoft.Extensions.Logging;
 using SuperSocket.ProtoBase;
 
@@ -60,6 +59,11 @@ protected virtual Task StartTask<TPackageInfo>(IObjectPipe<TPackageInfo> package
             return StartInputPipeTask(packagePipe, _cts.Token);
         }
 
+        protected void UpdateLastActiveTime()
+        {
+            LastActiveTime = DateTimeOffset.Now;
+        }
+
         public async override IAsyncEnumerable<TPackageInfo> RunAsync<TPackageInfo>(IPipelineFilter<TPackageInfo> pipelineFilter)
         {
             var packagePipe = !Options.ReadAsDemand
@@ -211,14 +215,8 @@ protected void WritePackageWithEncoder<TPackage>(IBufferWriter<byte> writer, IPa
             packageEncoder.Encode(writer, package);
         }
 
-        protected internal ArraySegment<T> GetArrayByMemory<T>(ReadOnlyMemory<T> memory)
+        protected virtual void OnInputPipeRead(ReadResult result)
         {
-            if (!MemoryMarshal.TryGetArray(memory, out var result))
-            {
-                throw new InvalidOperationException("Buffer backed by array was expected");
-            }
-
-            return result;
         }
 
         protected async Task ReadPipeAsync<TPackageInfo>(PipeReader reader, IObjectPipe<TPackageInfo> packagePipe, CancellationToken cancellationToken)
@@ -232,6 +230,7 @@ protected async Task ReadPipeAsync<TPackageInfo>(PipeReader reader, IObjectPipe<
                 try
                 {
                     result = await reader.ReadAsync(cancellationToken).ConfigureAwait(false);
+                    OnInputPipeRead(result);
                 }
                 catch (Exception e)
                 {

diff --git a/src/SuperSocket.Kestrel/KestrelPipeConnection.cs b/src/SuperSocket.Kestrel/KestrelPipeConnection.cs
@@ -3,8 +3,10 @@
 using System;
 using System.IO.Pipelines;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.AspNetCore.Connections;
 using SuperSocket.Connection;
+using SuperSocket.ProtoBase;
 
 public class KestrelPipeConnection : PipeConnectionBase
 {
@@ -39,4 +41,30 @@ protected override void OnClosed()
 
         base.OnClosed();
     }
+
+    protected override void OnInputPipeRead(ReadResult result)
+    {
+        if (!result.IsCanceled && !result.IsCompleted)
+        {
+            UpdateLastActiveTime();
+        }
+    }
+
+    public override async ValueTask SendAsync(Action<PipeWriter> write)
+    {
+        await base.SendAsync(write);
+        UpdateLastActiveTime();
+    }
+
+    public override async ValueTask SendAsync(ReadOnlyMemory<byte> buffer)
+    {
+        await base.SendAsync(buffer);
+        UpdateLastActiveTime();
+    }
+
+    public override async ValueTask SendAsync<TPackage>(IPackageEncoder<TPackage> packageEncoder, TPackage package)
+    {
+        await base.SendAsync(packageEncoder, package);
+        UpdateLastActiveTime();
+    }
 }