Fix race condition on Unix pooling canceled Socket operations (#27866)

Author: stephentoub

src/System.Net.Sockets/src/System/Net/Sockets/SocketAsyncContext.Unix.cs

@@ -208,11 +208,6 @@ public void SetWaiting()
                 Volatile.Write(ref _state, (int)State.Waiting);
             }
 
-            public void DoCallback()
-            {
-                InvokeCallback();
-            }
-
             public bool TryCancel()
             {
                 Trace("Enter");
@@ -267,8 +262,11 @@ public bool TryCancel()
 #if DEBUG
                     Debug.Assert(Interlocked.CompareExchange(ref _callbackQueued, 1, 0) == 0, $"Unexpected _callbackQueued: {_callbackQueued}");
 #endif
-
-                    ThreadPool.QueueUserWorkItem(o => ((AsyncOperation)o).InvokeCallback(), this);
+                    // We've marked the operation as canceled, and so should invoke the callback, but
+                    // we can't pool the object, as ProcessQueue may still have a reference to it, due to
+                    // using a pattern whereby it takes the lock to grab an item, but then releases the lock
+                    // to do further processing on the item that's still in the list.
+                    ThreadPool.QueueUserWorkItem(o => ((AsyncOperation)o).InvokeCallback(allowPooling: false), this);
                 }
 
                 Trace("Exit");
@@ -289,7 +287,7 @@ public void DoAbort()
 
             protected abstract bool DoTryComplete(SocketAsyncContext context);
 
-            protected abstract void InvokeCallback();
+            public abstract void InvokeCallback(bool allowPooling);
 
             [Conditional("SOCKETASYNCCONTEXT_TRACE")]
             public void Trace(string message, [CallerMemberName] string memberName = null)
@@ -332,7 +330,7 @@ public Action<int, byte[], int, SocketFlags, SocketError> Callback
                 set => CallbackOrEvent = value;
             }
 
-            protected override void InvokeCallback() =>
+            public override void InvokeCallback(bool allowPooling) =>
                 ((Action<int, byte[], int, SocketFlags, SocketError>)CallbackOrEvent)(BytesTransferred, SocketAddress, SocketAddressLen, SocketFlags.None, ErrorCode);
         }
 
@@ -348,15 +346,18 @@ protected override bool DoTryComplete(SocketAsyncContext context)
                 return SocketPal.TryCompleteSendTo(context._socket, Buffer.Span, null, ref bufferIndex, ref Offset, ref Count, Flags, SocketAddress, SocketAddressLen, ref BytesTransferred, out ErrorCode);
             }
 
-            protected override void InvokeCallback()
+            public override void InvokeCallback(bool allowPooling)
             {
                 var cb = (Action<int, byte[], int, SocketFlags, SocketError>)CallbackOrEvent;
                 int bt = BytesTransferred;
                 byte[] sa = SocketAddress;
                 int sal = SocketAddressLen;
                 SocketError ec = ErrorCode;
 
-                AssociatedContext.ReturnOperation(this);
+                if (allowPooling)
+                {
+                    AssociatedContext.ReturnOperation(this);
+                }
 
                 cb(bt, sa, sal, SocketFlags.None, ec);
             }
@@ -374,15 +375,18 @@ protected override bool DoTryComplete(SocketAsyncContext context)
                 return SocketPal.TryCompleteSendTo(context._socket, default(ReadOnlySpan<byte>), Buffers, ref BufferIndex, ref Offset, ref Count, Flags, SocketAddress, SocketAddressLen, ref BytesTransferred, out ErrorCode);
             }
 
-            protected override void InvokeCallback()
+            public override void InvokeCallback(bool allowPooling)
             {
                 var cb = (Action<int, byte[], int, SocketFlags, SocketError>)CallbackOrEvent;
                 int bt = BytesTransferred;
                 byte[] sa = SocketAddress;
                 int sal = SocketAddressLen;
                 SocketError ec = ErrorCode;
 
-                AssociatedContext.ReturnOperation(this);
+                if (allowPooling)
+                {
+                    AssociatedContext.ReturnOperation(this);
+                }
 
                 cb(bt, sa, sal, SocketFlags.None, ec);
             }
@@ -416,7 +420,7 @@ public Action<int, byte[], int, SocketFlags, SocketError> Callback
                 set => CallbackOrEvent = value;
             }
 
-            protected override void InvokeCallback() =>
+            public override void InvokeCallback(bool allowPooling) =>
                 ((Action<int, byte[], int, SocketFlags, SocketError>)CallbackOrEvent)(
                     BytesTransferred, SocketAddress, SocketAddressLen, ReceivedFlags, ErrorCode);
         }
@@ -430,7 +434,7 @@ public BufferMemoryReceiveOperation(SocketAsyncContext context) : base(context)
             protected override bool DoTryComplete(SocketAsyncContext context) =>
                 SocketPal.TryCompleteReceiveFrom(context._socket, Buffer.Span, null, Flags, SocketAddress, ref SocketAddressLen, out BytesTransferred, out ReceivedFlags, out ErrorCode);
 
-            protected override void InvokeCallback()
+            public override void InvokeCallback(bool allowPooling)
             {
                 var cb = (Action<int, byte[], int, SocketFlags, SocketError>)CallbackOrEvent;
                 int bt = BytesTransferred;
@@ -439,7 +443,10 @@ protected override void InvokeCallback()
                 SocketFlags rf = ReceivedFlags;
                 SocketError ec = ErrorCode;
 
-                AssociatedContext.ReturnOperation(this);
+                if (allowPooling)
+                {
+                    AssociatedContext.ReturnOperation(this);
+                }
 
                 cb(bt, sa, sal, rf, ec);
             }
@@ -454,7 +461,7 @@ public BufferListReceiveOperation(SocketAsyncContext context) : base(context) {
             protected override bool DoTryComplete(SocketAsyncContext context) =>
                 SocketPal.TryCompleteReceiveFrom(context._socket, default(Span<byte>), Buffers, Flags, SocketAddress, ref SocketAddressLen, out BytesTransferred, out ReceivedFlags, out ErrorCode);
 
-            protected override void InvokeCallback()
+            public override void InvokeCallback(bool allowPooling)
             {
                 var cb = (Action<int, byte[], int, SocketFlags, SocketError>)CallbackOrEvent;
                 int bt = BytesTransferred;
@@ -463,7 +470,10 @@ protected override void InvokeCallback()
                 SocketFlags rf = ReceivedFlags;
                 SocketError ec = ErrorCode;
 
-                AssociatedContext.ReturnOperation(this);
+                if (allowPooling)
+                {
+                    AssociatedContext.ReturnOperation(this);
+                }
 
                 cb(bt, sa, sal, rf, ec);
             }
@@ -504,7 +514,7 @@ public Action<int, byte[], int, SocketFlags, IPPacketInformation, SocketError> C
             protected override bool DoTryComplete(SocketAsyncContext context) =>
                 SocketPal.TryCompleteReceiveMessageFrom(context._socket, Buffer.Span, Buffers, Flags, SocketAddress, ref SocketAddressLen, IsIPv4, IsIPv6, out BytesTransferred, out ReceivedFlags, out IPPacketInformation, out ErrorCode);
 
-            protected override void InvokeCallback() =>
+            public override void InvokeCallback(bool allowPooling) =>
                 ((Action<int, byte[], int, SocketFlags, IPPacketInformation, SocketError>)CallbackOrEvent)(
                     BytesTransferred, SocketAddress, SocketAddressLen, ReceivedFlags, IPPacketInformation, ErrorCode);
         }
@@ -530,15 +540,18 @@ protected override bool DoTryComplete(SocketAsyncContext context)
                 return completed;
             }
 
-            protected override void InvokeCallback()
+            public override void InvokeCallback(bool allowPooling)
             {
                 var cb = (Action<IntPtr, byte[], int, SocketError>)CallbackOrEvent;
                 IntPtr fd = AcceptedFileDescriptor;
                 byte[] sa = SocketAddress;
                 int sal = SocketAddressLen;
                 SocketError ec = ErrorCode;
 
-                AssociatedContext.ReturnOperation(this);
+                if (allowPooling)
+                {
+                    AssociatedContext.ReturnOperation(this);
+                }
 
                 cb(fd, sa, sal, ec);
             }
@@ -562,7 +575,7 @@ protected override bool DoTryComplete(SocketAsyncContext context)
                 return result;
             }
 
-            protected override void InvokeCallback() =>
+            public override void InvokeCallback(bool allowPooling) =>
                 ((Action<SocketError>)CallbackOrEvent)(ErrorCode);
         }
 
@@ -582,7 +595,7 @@ public Action<long, SocketError> Callback
                 set => CallbackOrEvent = value;
             }
 
-            protected override void InvokeCallback() =>
+            public override void InvokeCallback(bool allowPooling) =>
                 ((Action<long, SocketError>)CallbackOrEvent)(BytesTransferred, ErrorCode);
 
             protected override bool DoTryComplete(SocketAsyncContext context) =>
@@ -944,7 +957,10 @@ public void ProcessQueue(SocketAsyncContext context)
                         ThreadPool.QueueUserWorkItem(s_processingCallback, context);
                     }
 
-                    op.DoCallback();
+                    // At this point, the operation has completed and it's no longer
+                    // in the queue / no one else has a reference to it.  We can invoke
+                    // the callback and let it pool the object if appropriate.
+                    op.InvokeCallback(allowPooling: true);
                 }
                 else
                 {

src/System.Net.Sockets/tests/FunctionalTests/Accept.cs

@@ -276,6 +276,30 @@ public async Task Accept_WithInUseTargetSocket_Fails()
                 Assert.Throws<InvalidOperationException>(() => { AcceptAsync(listener, server); });
             }
         }
+
+        [Fact]
+        public async Task AcceptAsync_MultipleAcceptsThenDispose_AcceptsThrowAfterDispose()
+        {
+            if (UsesSync)
+            {
+                return;
+            }
+
+            for (int i = 0; i < 100; i++)
+            {
+                using (var listener = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp))
+                {
+                    listener.Bind(new IPEndPoint(IPAddress.Loopback, 0));
+                    listener.Listen(2);
+
+                    Task accept1 = AcceptAsync(listener);
+                    Task accept2 = AcceptAsync(listener);
+                    listener.Dispose();
+                    await Assert.ThrowsAnyAsync<Exception>(() => accept1);
+                    await Assert.ThrowsAnyAsync<Exception>(() => accept2);
+                }
+            }
+        }
     }
 
     public sealed class AcceptSync : Accept<SocketHelperArraySync> { }