From 394c89a36edf3a44bc46ddf05a22725e7b4fc8b5 Mon Sep 17 00:00:00 2001
From: Peter Giacomo Lombardo <peter.lombardo@hivemq.com>
Date: Thu, 28 Sep 2023 16:13:53 +0200
Subject: [PATCH] Extended AfterDisconnect Support (#93)

---
 Examples/Reconnect/Program.cs                          | 3 ++-
 Examples/Reconnect/Reconnect.csproj                    | 2 +-
 Source/HiveMQtt/Client/HiveMQClient.cs                 | 4 ++--
 Source/HiveMQtt/Client/HiveMQClientSocket.cs           | 4 ++--
 Source/HiveMQtt/Client/HiveMQClientTrafficProcessor.cs | 3 +--
 5 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/Examples/Reconnect/Program.cs b/Examples/Reconnect/Program.cs
index 4bd4e77a..4d9d38ec 100644
--- a/Examples/Reconnect/Program.cs
+++ b/Examples/Reconnect/Program.cs
@@ -135,5 +135,6 @@
     }
 }
 
-Console.WriteLine("Disconnecting gracefully...");
+Console.WriteLine("Disconnecting gracefully and waiting 5 seconds...");
 await client.DisconnectAsync().ConfigureAwait(false);
+await Task.Delay(5000).ConfigureAwait(false);
diff --git a/Examples/Reconnect/Reconnect.csproj b/Examples/Reconnect/Reconnect.csproj
index f0ab0a94..0e7e705f 100644
--- a/Examples/Reconnect/Reconnect.csproj
+++ b/Examples/Reconnect/Reconnect.csproj
@@ -14,7 +14,7 @@
 
   <!-- Update the version to match -->
   <ItemGroup>
-    <PackageReference Include="HiveMQtt" Version="0.4.0" />
+    <PackageReference Include="HiveMQtt" Version="0.4.1" />
   </ItemGroup>
 
 </Project>
diff --git a/Source/HiveMQtt/Client/HiveMQClient.cs b/Source/HiveMQtt/Client/HiveMQClient.cs
index f07be098..767bf04b 100644
--- a/Source/HiveMQtt/Client/HiveMQClient.cs
+++ b/Source/HiveMQtt/Client/HiveMQClient.cs
@@ -158,11 +158,11 @@ public async Task<bool> DisconnectAsync(DisconnectOptions? options = null)
             this.OnDisconnectSent -= eventHandler;
         }
 
-        this.connectState = ConnectState.Disconnected;
-
         // Close the socket
         this.CloseSocket();
 
+        this.connectState = ConnectState.Disconnected;
+
         // Clear the send queue
         this.sendQueue.Clear();
         return true;
diff --git a/Source/HiveMQtt/Client/HiveMQClientSocket.cs b/Source/HiveMQtt/Client/HiveMQClientSocket.cs
index 401f9481..7fd49ebb 100644
--- a/Source/HiveMQtt/Client/HiveMQClientSocket.cs
+++ b/Source/HiveMQtt/Client/HiveMQClientSocket.cs
@@ -137,8 +137,6 @@ internal async Task<bool> ConnectSocketAsync()
 
     internal bool CloseSocket(bool? shutdownPipeline = true)
     {
-        this.cancellationSource.Cancel();
-
         if (shutdownPipeline == true)
         {
             // Shutdown the pipeline
@@ -150,6 +148,8 @@ internal bool CloseSocket(bool? shutdownPipeline = true)
         this.socket?.Shutdown(SocketShutdown.Both);
         this.socket?.Close();
 
+        this.cancellationSource.Cancel();
+
         return true;
     }
 }
diff --git a/Source/HiveMQtt/Client/HiveMQClientTrafficProcessor.cs b/Source/HiveMQtt/Client/HiveMQClientTrafficProcessor.cs
index 3a3c4382..c95af7a3 100644
--- a/Source/HiveMQtt/Client/HiveMQClientTrafficProcessor.cs
+++ b/Source/HiveMQtt/Client/HiveMQClientTrafficProcessor.cs
@@ -353,6 +353,7 @@ private Task<bool> TrafficInflowProcessorAsync(CancellationToken cancellationTok
             switch (this.connectState)
             {
                 case ConnectState.Disconnecting:
+                case ConnectState.Disconnected:
                     // This disconnect was either user or broker initiated.
                     // Launch the AfterDisconnect event with a clean disconnect set to true.
                     this.AfterDisconnectEventLauncher(true);
@@ -364,8 +365,6 @@ private Task<bool> TrafficInflowProcessorAsync(CancellationToken cancellationTok
                     break;
                 case ConnectState.Connecting:
                     break;
-                case ConnectState.Disconnected:
-                    break;
                 default:
                     break;
             }