From d1156a2e3e6f702be777926e401a3abe30370af6 Mon Sep 17 00:00:00 2001
From: Marco Stephan <marc.stephan96@hotmail.de>
Date: Sat, 14 May 2022 22:32:47 +0200
Subject: [PATCH] Add option to TCP client/session to send remaining bytes when
 disconnecting

---
 source/NetCoreServer/TcpClient.cs  | 21 +++++++++++++++++++++
 source/NetCoreServer/TcpSession.cs | 21 +++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/source/NetCoreServer/TcpClient.cs b/source/NetCoreServer/TcpClient.cs
index 87d8b1ca..0725d94a 100644
--- a/source/NetCoreServer/TcpClient.cs
+++ b/source/NetCoreServer/TcpClient.cs
@@ -126,6 +126,14 @@ private TcpClient(EndPoint endpoint, string address, int port)
         /// </summary>
         public int OptionSendBufferSize { get; set; } = 8192;
 
+        /// <summary>
+        /// Option: whether to send all remaining bytes when disconnecting
+        /// </summary>
+        /// <remarks>
+        /// This only affects bytes end via SendAsync(...)
+        /// </remarks>
+        public bool OptionSendRemainingBytesOnDisconnect { get; set; } = false;
+
         #region Connect/Disconnect client
 
         private SocketAsyncEventArgs _connectEventArg;
@@ -270,6 +278,12 @@ public virtual bool Disconnect()
             if (IsConnecting)
                 Socket.CancelConnectAsync(_connectEventArg);
 
+            // Wait for pending asynchronous send operations if desired
+            if (OptionSendRemainingBytesOnDisconnect)
+            {
+                _sendBuffersEmpty.WaitOne();
+            }
+
             // Reset event args
             _connectEventArg.Completed -= OnAsyncCompleted;
             _receiveEventArg.Completed -= OnAsyncCompleted;
@@ -413,6 +427,7 @@ public virtual bool ReconnectAsync()
         private Buffer _sendBufferFlush;
         private SocketAsyncEventArgs _sendEventArg;
         private long _sendBufferFlushOffset;
+        private readonly ManualResetEvent _sendBuffersEmpty = new ManualResetEvent(true);
 
         /// <summary>
         /// Send data to the server (synchronous)
@@ -496,6 +511,7 @@ public virtual bool SendAsync(byte[] buffer, long offset, long size)
                 }
 
                 // Fill the main send buffer
+                _sendBuffersEmpty.Reset();
                 _sendBufferMain.Append(buffer, offset, size);
 
                 // Update statistic
@@ -659,6 +675,8 @@ private void TrySend()
                 // Call the empty send buffer handler
                 if (empty)
                 {
+                    _sendBuffersEmpty.Set();
+
                     OnEmpty();
                     return;
                 }
@@ -867,6 +885,8 @@ private bool ProcessSend(SocketAsyncEventArgs e)
                 return true;
             else
             {
+                _sendBuffersEmpty.Set();
+
                 SendError(e.SocketError);
                 DisconnectAsync();
                 return false;
@@ -992,6 +1012,7 @@ protected virtual void Dispose(bool disposingManagedResources)
                 {
                     // Dispose managed resources here...
                     DisconnectAsync();
+                    _sendBuffersEmpty.Dispose();
                 }
 
                 // Dispose unmanaged resources here...
diff --git a/source/NetCoreServer/TcpSession.cs b/source/NetCoreServer/TcpSession.cs
index d3842e99..38f912d4 100644
--- a/source/NetCoreServer/TcpSession.cs
+++ b/source/NetCoreServer/TcpSession.cs
@@ -71,6 +71,14 @@ public TcpSession(TcpServer server)
         /// </summary>
         public int OptionSendBufferSize { get; set; } = 8192;
 
+        /// <summary>
+        /// Option: whether to send all remaining bytes when disconnecting
+        /// </summary>
+        /// <remarks>
+        /// This only affects bytes end via SendAsync(...)
+        /// </remarks>
+        public bool OptionSendRemainingBytesOnDisconnect { get; set; } = false;
+
         #region Connect/Disconnect session
 
         /// <summary>
@@ -154,6 +162,12 @@ public virtual bool Disconnect()
             if (!IsConnected)
                 return false;
 
+            // Wait for pending asynchronous send operations if desired
+            if (OptionSendRemainingBytesOnDisconnect)
+            {
+                _sendBuffersEmpty.WaitOne();
+            }
+
             // Reset event args
             _receiveEventArg.Completed -= OnAsyncCompleted;
             _sendEventArg.Completed -= OnAsyncCompleted;
@@ -225,6 +239,7 @@ public virtual bool Disconnect()
         private Buffer _sendBufferFlush;
         private SocketAsyncEventArgs _sendEventArg;
         private long _sendBufferFlushOffset;
+        private readonly ManualResetEvent _sendBuffersEmpty = new ManualResetEvent(true);
 
         /// <summary>
         /// Send data to the client (synchronous)
@@ -309,6 +324,7 @@ public virtual bool SendAsync(byte[] buffer, long offset, long size)
                 }
 
                 // Fill the main send buffer
+                _sendBuffersEmpty.Reset();
                 _sendBufferMain.Append(buffer, offset, size);
 
                 // Update statistic
@@ -473,6 +489,8 @@ private void TrySend()
                 // Call the empty send buffer handler
                 if (empty)
                 {
+                    _sendBuffersEmpty.Set();
+
                     OnEmpty();
                     return;
                 }
@@ -628,6 +646,8 @@ private bool ProcessSend(SocketAsyncEventArgs e)
                 return true;
             else
             {
+                _sendBuffersEmpty.Set();
+
                 SendError(e.SocketError);
                 Disconnect();
                 return false;
@@ -753,6 +773,7 @@ protected virtual void Dispose(bool disposingManagedResources)
                 {
                     // Dispose managed resources here...
                     Disconnect();
+                    _sendBuffersEmpty.Dispose();
                 }
 
                 // Dispose unmanaged resources here...