Add CancellationToken overloads to async methods in System.Net.Sockets

[EgoPingvina]

Good day. I work with UdpClient and have wrapper upon it.

For reading I have asynchronous method:

    private async Task<byte[]> Receive(UdpClient client, CancellationToken breakToken)
    {
        // Выход из async, если произошёл CancellationRequest
        breakToken.ThrowIfCancellationRequested();

        UdpReceiveResult result;
        try
        {
            result = await client.ReceiveAsync().WithCancellation(breakToken);
        }
        catch(OperationCanceledException)
        {
            // Штатная ситуация ручной остановки Task-а
        }
        
        return result.Buffer;
    }

Where WithCancellation is my extension-method for early termination:

    public static async Task<T> WithCancellation<T>(
        this Task<T> task, CancellationToken cancellationToken)
    {
        var tcs = new TaskCompletionSource<bool>();

        using (cancellationToken.Register(
            s => ((TaskCompletionSource<bool>)s).TrySetResult(true),
            tcs))
            if (task != await Task.WhenAny(task, tcs.Task))
                throw new OperationCanceledException(cancellationToken);

        return await task;
    }

And after manual reading stop, when I call Dispose, System.ObjectDisposedException occurs. CallStack:

    >	System.dll!System.Net.Sockets.UdpClient.EndReceive(System.IAsyncResult asyncResult, ref System.Net.IPEndPoint remoteEP)	Unknown
 	System.dll!System.Net.Sockets.UdpClient.ReceiveAsync.AnonymousMethod__64_1(System.IAsyncResult ar)	Unknown
 	mscorlib.dll!System.Threading.Tasks.TaskFactory<System.Net.Sockets.UdpReceiveResult>.FromAsyncCoreLogic(System.IAsyncResult iar, System.Func<System.IAsyncResult, System.Net.Sockets.UdpReceiveResult> endFunction, System.Action<System.IAsyncResult> endAction, System.Threading.Tasks.Task<System.Net.Sockets.UdpReceiveResult> promise, bool requiresSynchronization)	Unknown
 	mscorlib.dll!System.Threading.Tasks.TaskFactory<System.Net.Sockets.UdpReceiveResult>.FromAsyncImpl.AnonymousMethod__0(System.IAsyncResult iar)	Unknown
 	System.dll!System.Net.LazyAsyncResult.Complete(System.IntPtr userToken)	Unknown
 	System.dll!System.Net.ContextAwareResult.CompleteCallback(object state)	Unknown
 	mscorlib.dll!System.Threading.ExecutionContext.RunInternal(System.Threading.ExecutionContext executionContext, System.Threading.ContextCallback callback, object state, bool preserveSyncCtx)	Unknown
 	mscorlib.dll!System.Threading.ExecutionContext.Run(System.Threading.ExecutionContext executionContext, System.Threading.ContextCallback callback, object state, bool preserveSyncCtx)	Unknown
 	mscorlib.dll!System.Threading.ExecutionContext.Run(System.Threading.ExecutionContext executionContext, System.Threading.ContextCallback callback, object state)	Unknown
 	System.dll!System.Net.ContextAwareResult.Complete(System.IntPtr userToken)	Unknown
 	System.dll!System.Net.LazyAsyncResult.ProtectedInvokeCallback(object result, System.IntPtr userToken)	Unknown
 	System.dll!System.Net.Sockets.BaseOverlappedAsyncResult.CompletionPortCallback(uint errorCode, uint numBytes, System.Threading.NativeOverlapped* nativeOverlapped)	Unknown
 	mscorlib.dll!System.Threading._IOCompletionCallback.PerformIOCompletionCallback(uint errorCode, uint numBytes, System.Threading.NativeOverlapped* pOVERLAP)	Unknown

If I understood correctly, root of error in ReceiveAsync, in my method of stop it to be exact.

After nearly a week of suffering, I have found the reason and solution.

At first, I looked at the UdpClient source code. The ReceiveAsync method:

    [HostProtection(ExternalThreading = true)]
    public Task<UdpReceiveResult> ReceiveAsync()
    {
        return Task<UdpReceiveResult>.Factory.FromAsync((callback, state) => BeginReceive(callback, state), (ar)=>
            {
                IPEndPoint remoteEP = null;
                Byte[] buffer = EndReceive(ar, ref remoteEP);
                return new UdpReceiveResult(buffer, remoteEP);
 
            }, null);
    }

At second, I found this post with perfect answer:http://stackoverflow.com/questions/4662553/how-to-abort-sockets-beginreceive, in which said:

To cancel a pending call to the BeginConnect() method, close the Socket. When the Close() method is called while an asynchronous operation is in progress, the callback provided to the BeginConnect() method is called. A subsequent call to the EndConnect(IAsyncResult) method will throw an ObjectDisposedException to indicate that the operation has been cancelled.

And, as we can see, the original ReceiveAsync method return us the ObjectDisposedException, because IOOperation has not been completed after Close invoking.

To overcome this problem I have done like this:

New ReceiveAsync realization:

    /// <summary>
    /// Асинхронный запрос на ожидание приёма данных с возможностью досрочного выхода
    /// (для выхода из ожидания вызовите метод Disconnect())
    /// </summary>
    /// <param name="client">Рабочий экземпляр класса UdpClient</param>
    /// <param name="breakToken">Признак досрочного завершения</param>
    /// <returns>Если breakToken произошёл до вызова данного метода или в режиме ожидания
    /// ответа, вернёт пустой UdpReceiveResult; при удачном получении ответа-результат
    /// асинхронной операции чтения</returns>
    public Task<UdpReceiveResult> ReceiveAsync(UdpClient client, CancellationToken breakToken)
        => breakToken.IsCancellationRequested
            ? Task<UdpReceiveResult>.Run(() => new UdpReceiveResult())
            : Task<UdpReceiveResult>.Factory.FromAsync(
                (callback, state) => client.BeginReceive(callback, state),
                (ar) =>
                    {
                        /// Предотвращение <exception cref="ObjectDisposedException"/>
                        if (breakToken.IsCancellationRequested)
                            return new UdpReceiveResult();
    
                        IPEndPoint remoteEP = null;
                        var buffer = client.EndReceive(ar, ref remoteEP);
                        return new UdpReceiveResult(buffer, remoteEP);
                    },
                null);

New Dispose realization:

    protected virtual void Dispose(bool disposing)
    {
        if (disposing)
        {
            this.cancelReceive?.Cancel();
            this.client?.Close();
            this.cancelReceive?.Dispose();
        }
    }

Is my solution OK(will be cool if you add it to framework in next version) and if not can you say me how to fix it differently?

[i3arnon]

@EgoPingvina

Is my solution OK?

Generaly, I think so: Not All Beginnings Must Have An End

[CIPop]

@EgoPingvina I believe that the ObjectDisposedException is what .Net Framework returns as well.
Wouldn't your solution change this behavior and potentially break existing applications relying on this exception?

/cc @stephentoub

[EgoPingvina]

@CIPop UdpClient has ReceiveAsync method and if overloading with 2 parameters(like I done it) will be added, I think this will not break the existing applications, but will give me and many other people convenient mechanism for interrupting the asynchronous reading

[CIPop]

overloading with 2 parameters

Sorry for missing that: I definitely agree with having new APIs that use CancellationToken. We should consider adding new API not only for UdpClient but for all other async networking APIs that are missing it.
For Windows, I also suggest prototyping cancellation via CancellIoEx instead of just closing the handle.

The work requires following the API review process.
The CancelIoEx route may be quite large while the proposed implementation above should be a medium-large item.

[EgoPingvina]

@CIPop Yes, you are absolutely right.

[CIPop]

Sounds good. I'm changing this item to track adding CancellationToken in Sockets.
I'll keep #16031 to track the same feature in Security (SslStream, NegotiateStream).

[felipepessoto]

Related: #17711, #16236

[AnderssonPeter]

2 years later and we still don't have a Overload that accepts CancellationToken when using
UdpClient.ReceiveAsync, it seems to be the same case with .net core 3.0 preview6.
Any news when this will be added?

Socket.ReceiveAsync seems to have support for CancellationToken but it does not seem to be usable with UDP as i can't get the sender when using it.

[karelz]

Duplicate of #921