B! Fix inconsistency

pixonic · May 31, 2021 · 738b449 · 738b449
1 parent 351148a
commit 738b449
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 11 deletions.
diff --git a/Pixockets/ThreadSmartSock.cs b/Pixockets/ThreadSmartSock.cs
@@ -1,4 +1,5 @@
-using System.Net;
+using System.Collections.Generic;
+using System.Net;
 using System.Threading;
 using Pixockets.Pools;
 
@@ -12,7 +13,7 @@ public class ThreadSmartSock : SmartReceiverBase
         private readonly ThreadSafeQueue<SmartPacketToSend> _sendQueue = new ThreadSafeQueue<SmartPacketToSend>();
         private readonly ThreadSafeQueue<ReceivedSmartPacket> _recvQueue = new ThreadSafeQueue<ReceivedSmartPacket>();
         private readonly ThreadSafeQueue<IPEndPoint> _connectQueue = new ThreadSafeQueue<IPEndPoint>();
-        private readonly ThreadSafeQueue<IPEndPoint> _disconnectQueue = new ThreadSafeQueue<IPEndPoint>();
+        private readonly ThreadSafeQueue<KeyValuePair<IPEndPoint, DisconnectReason>> _disconnectQueue = new ThreadSafeQueue<KeyValuePair<IPEndPoint, DisconnectReason>>();
         private readonly BufferPoolBase _buffersPool;
         private readonly SmartReceiverBase _callbacks;
 
@@ -49,8 +50,8 @@ public void Listen(int port)
 
         public void Tick()
         {
-            while (_disconnectQueue.TryTake(out var disconnectEndPoint))
-                _callbacks.OnDisconnect(disconnectEndPoint);
+            while (_disconnectQueue.TryTake(out var disconnectPair))
+                _callbacks.OnDisconnect(disconnectPair.Key, disconnectPair.Value);
 
             while (_connectQueue.TryTake(out var connectEndPoint))
                 _callbacks.OnConnect(connectEndPoint);
@@ -132,9 +133,9 @@ public override void OnConnect(IPEndPoint endPoint)
             _connectQueue.Add(endPoint);
         }
 
-        public override void OnDisconnect(IPEndPoint endPoint)
+        public override void OnDisconnect(IPEndPoint endPoint, DisconnectReason reason)
         {
-            _disconnectQueue.Add(endPoint);
+            _disconnectQueue.Add(new KeyValuePair<IPEndPoint, DisconnectReason>(endPoint, reason));
         }
     }
 }
diff --git a/UnitTests/ThreadSmartSockTests.cs b/UnitTests/ThreadSmartSockTests.cs
@@ -13,7 +13,6 @@ public class ThreadSmartSockTests
         private MockSmartCallbacks _cbs;
         private MockBufferPool _bufferPool;
         private MockSock _bareSock;
-        private SmartSock _smartSock;
         private ThreadSmartSock _threadSmartSock;
 
         [SetUp]
@@ -22,8 +21,7 @@ public void Setup()
             _cbs = new MockSmartCallbacks();
             _bareSock = new MockSock();
             _bufferPool = new MockBufferPool();
-            _smartSock = new SmartSock(_bufferPool, _bareSock, _cbs);
-            _threadSmartSock = new ThreadSmartSock(_smartSock, _bufferPool);
+            _threadSmartSock = new ThreadSmartSock(_bufferPool, _bareSock, _cbs);
         }
 
         [TearDown]
@@ -44,7 +42,7 @@ public void ThreadSmartSockRedirects()
             Thread.Sleep(50);
 
             Assert.AreEqual(PixocketState.Connecting, _threadSmartSock.State);
-            Assert.AreEqual(_bareSock.RemoteEndPoint, endPoint); 
+            Assert.AreEqual(_bareSock.RemoteEndPoint, endPoint);
 
             _threadSmartSock.Close();
 
@@ -89,7 +87,7 @@ public void ThreadSmartSockReceives()
             {
                 if (_threadSmartSock.Receive(ref receivedPacket))
                     break;
-                
+
                 Thread.Sleep(1);
             }