Poller.remove socket disposed fix

src/NetMQ.Tests/NetMQPollerTest.cs

@@ -347,6 +347,8 @@ public void RemoveSocket()
                     // identity
                     e.Socket.SkipFrame();
 
+                    //**Note: bad to assert from worker thread!  
+                    // If it fails, the test will crash, not report failure!
                     Assert.Equal("Hello", e.Socket.ReceiveFrameString(out bool more));
                     Assert.False(more);
 
@@ -400,6 +402,150 @@ public void RemoveSocket()
             }
         }
 
+        [Fact]
+        public async void RemoveAndDisposeSocket()
+        {
+            //set up poller, start it
+            var patient = new NetMQPoller();
+            patient.RunAsync();
+
+            Assert.True(patient.IsRunning);
+
+            //create a pub-sub pair
+            var port = 55667;
+            var conn = $"tcp://127.0.0.1:{port}";
+
+            var pub = new PublisherSocket();
+            pub.Bind(conn);
+
+            var sub = new SubscriberSocket();
+            sub.Connect(conn);
+            sub.SubscribeToAnyTopic();
+
+            //handle callbacks from poller thread
+            sub.ReceiveReady += (s, e) =>
+            {
+                var msg = e.Socket.ReceiveFrameString();
+
+                Debug.WriteLine($"sub has data: {msg}");
+            };
+
+            //add the subscriber socket to poller
+            patient.Add(sub);
+
+            //set up pub on separate thread
+            var canceller = new CancellationTokenSource();
+
+            var pubAction = new Action(async () =>
+            {
+                var token = canceller.Token;
+
+                uint i = 0;
+
+                while (!token.IsCancellationRequested)
+                {
+                    pub.SendFrame($"Hello-{++i}");
+
+                    // send ~ 5Hz
+                    await Task.Delay(200);
+                }
+            });
+
+            var pubThread = Task.Run(pubAction);
+
+            //allow a little time to run
+            await Task.Delay(2000);
+
+            //now try to remove the sub from poller
+            patient.RemoveAndDispose(sub);
+
+            Assert.True(sub.IsDisposed);
+
+            //allow for poller to continue running
+            await Task.Delay(2000);
+
+            patient.Stop();
+            Assert.False(patient.IsRunning);
+
+            canceller.Cancel();
+
+            pub?.Dispose();
+            patient?.Dispose();
+        }
+
+        [Fact]
+        public async void DisposeSocketAfterRemoval()
+        {
+            //set up poller, start it
+            var patient = new NetMQPoller();
+            patient.RunAsync();
+
+            Assert.True(patient.IsRunning);
+
+            //create a pub-sub pair
+            var port = 55667;
+            var conn = $"tcp://127.0.0.1:{port}";
+
+            var pub = new PublisherSocket();
+            pub.Bind(conn);
+
+            var sub = new SubscriberSocket();
+            sub.Connect(conn);
+            sub.SubscribeToAnyTopic();
+
+            //handle callbacks from poller thread
+            sub.ReceiveReady += (s, e) =>
+            {
+                var msg = e.Socket.ReceiveFrameString();
+
+                Debug.WriteLine($"sub has data: {msg}");
+            };
+
+            //add the subscriber socket to poller
+            patient.Add(sub);
+
+            //set up pub on separate thread
+            var canceller = new CancellationTokenSource();
+
+            var pubAction = new Action(async () =>
+            {
+                var token = canceller.Token;
+
+                uint i = 0;
+
+                while(!token.IsCancellationRequested)
+                {
+                    pub.SendFrame($"Hello-{++i}");
+
+                    // send ~ 5Hz
+                    await Task.Delay(200);
+                }
+            });
+
+            var pubThread = Task.Run(pubAction);
+
+            //allow a little time to run
+            await Task.Delay(2000);
+
+            //now try to remove the sub from poller
+            patient.Remove(sub);
+
+            // dispose the sub (this will cause exception on poller's worker-thread) and it can't be caught!
+            sub.Dispose();
+            sub = null;
+
+            //allow for poller to continue running
+            await Task.Delay(2000);
+
+            patient.Stop();
+            Assert.False(patient.IsRunning);
+
+            canceller.Cancel();
+
+            pub?.Dispose();
+            patient?.Dispose();
+        }
+
         [Fact]
         public void AddThrowsIfSocketAlreadyDisposed()
         {
@@ -454,6 +600,11 @@ public void DisposeThrowsIfSocketAlreadyDisposed()
             // Dispose the socket.
             // It is incorrect to have a disposed socket in a poller.
             // Disposed sockets can throw into the poller's thread.
+
+            //**JASells: And what does that have to do with removing one?  Should check for disposed 
+            // socket on Add, not Remove!  Check only internally on to avoid accessing a potentially,
+            // disposed socket, but otherwise, removing it from the list is safe, and preferable to 
+            // throwing exception since it makes the poller more resilient to unintended mis-use!
             socket.Dispose();
 
             // Dispose throws if a polled socket is disposed

src/NetMQ/ISocketPollableCollection.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Threading.Tasks;
 using JetBrains.Annotations;
 using NetMQ.Monitoring;
 

src/NetMQ/NetMQ.csproj

@@ -130,4 +130,22 @@
     <PackageReference Include="SourceLink.Create.CommandLine" Version="2.7.4" PrivateAssets="All" />
   </ItemGroup>
 
+  <ItemGroup Condition="'$(TargetFramework)' == 'net47'">
+    <PackageReference Include="Microsoft.Bcl.Async">
+      <Version>1.0.168</Version>
+    </PackageReference>
+  </ItemGroup>
+
+  <ItemGroup Condition="'$(TargetFramework)' == 'net40'">
+    <PackageReference Include="Microsoft.Bcl.Async">
+      <Version>1.0.168</Version>
+    </PackageReference>
+  </ItemGroup>
+
+  <ItemGroup Condition="'$(TargetFramework)' == 'netstandard2.0'">
+    <PackageReference Include="Microsoft.Bcl.Async">
+      <Version>1.0.168</Version>
+    </PackageReference>
+  </ItemGroup>
+
 </Project>

src/NetMQ/NetMQPoller.cs

@@ -5,6 +5,7 @@
 using System.Linq;
 using System.Net.Sockets;
 using System.Threading;
+using System.Threading.Tasks;
 using JetBrains.Annotations;
 using NetMQ.Core.Utils;
 #if !NET35
@@ -112,21 +113,47 @@ protected override void QueueTask(Task task)
             m_tasksQueue.Enqueue(task);
         }
 
-        public void Run([NotNull] Action action)
+        public Task Run([NotNull] Action action)
         {
+            Task t;
+
             if (!IsRunning || CanExecuteTaskInline)
+            {
                 action();
+
+                t = FromResult<object>(null);
+            }
             else
-                new Task(action).Start(this);
+            {
+                t = new Task(action);
+                t.Start(this);
+            }
+
+            return t;
         }
+
+        /// <summary>
+        /// Provides a completed task with the result for a syncronously run action.
+        /// this only needed for .NET40.  Depricated by <see cref="Task.FromResult()"/> in 4.5+
+        /// </summary>
+        /// <typeparam name="TResult"></typeparam>
+        /// <param name="result"></param>
+        /// <returns></returns>
+        private static Task<TResult> FromResult<TResult>(TResult result)
+        {
+            var tcs = new TaskCompletionSource<TResult>();
+            tcs.SetResult(result);
+            return tcs.Task;
+        }
+
 #else
         private void Run(Action action)
         {
             action();
         }
 #endif
 
-        #endregion
+#endregion
 
         public NetMQPoller()
         {
@@ -225,6 +252,11 @@ public void Remove(ISocketPollable socket)
         {
             if (socket == null)
                 throw new ArgumentNullException(nameof(socket));
+
+            // JASells: not sure I agree with this thow.
+            // If trying to remove a disposed socket, why complain?  It *might* work.  The issue
+            // is if the poller's thread tries to actually service the socket before the remove call...
+
             if (socket.IsDisposed)
                 throw new ArgumentException("Must not be disposed.", nameof(socket));
             CheckDisposed();
@@ -243,33 +275,18 @@ public void Remove(ISocketPollable socket)
                 socket.Socket.EventsChanged -= OnSocketEventsChanged;
                 m_sockets.Remove(socket.Socket);
                 m_isPollSetDirty = true;
-            });
+            })
+            .Wait(); 
+            // keep the API syncronous by blocking the calling thread via Wait(), else RemoveThrowsIfSocketAlreadyDisposed() test fails
         }
 
         public void RemoveAndDispose<T>(T socket) where T : ISocketPollable, IDisposable
         {
-            if (socket == null)
-                throw new ArgumentNullException(nameof(socket));
-            if (socket.IsDisposed)
-                throw new ArgumentException("Must not be disposed.", nameof(socket));
-            CheckDisposed();
-
-            Run(() =>
-            {
-                // Ensure the socket wasn't disposed while this code was waiting to be run on the poller thread
-                if (socket.IsDisposed)
-                    throw new InvalidOperationException(
-                        $"{nameof(NetMQPoller)}.{nameof(RemoveAndDispose)} was called from a non-poller thread, " +
-                        "so ran asynchronously. " +
-                        $"The {socket.GetType().Name} being removed was disposed while the remove " +
-                        $"operation waited to start on the poller thread. When using {nameof(RemoveAndDispose)} " +
-                        "you should not dispose the pollable object .");
+            //call the remove method
+            Remove(socket);
 
-                socket.Socket.EventsChanged -= OnSocketEventsChanged;
-                m_sockets.Remove(socket.Socket);
-                m_isPollSetDirty = true;
-                socket.Dispose();
-            });
+            //dispose of socket
+            socket.Dispose();
         }
 
         public void Remove([NotNull] NetMQTimer timer)
@@ -280,7 +297,7 @@ public void Remove([NotNull] NetMQTimer timer)
 
             timer.When = -1;
 
-            Run(() => m_timers.Remove(timer));
+            Run(() => m_timers.Remove(timer)).Wait();
         }
 
         public void Remove([NotNull] Socket socket)
@@ -293,7 +310,8 @@ public void Remove([NotNull] Socket socket)
             {
                 m_pollinSockets.Remove(socket);
                 m_isPollSetDirty = true;
-            });
+            })
+            .Wait();
         }
 
         #endregion