Unable to Cancel Socket.ReceiveAsync with CancellationTokenSource in Unity (Unity version: 2022.3.26

I’ve been experimenting with asynchronous programming in Unity for handling network communication recently, but I’ve encountered a tricky issue. I’m using the asynchronous method Socket.ReceiveAsync to receive data and I want to be able to cancel this operation using a CancellationTokenSource. However, no matter what I try, I can’t seem to cancel the ReceiveAsync operation successfully.

Here’s a simplified version of my code:

private async void Start()
{
    IPAddress ipAddress = IPAddress.Parse(ip);
    IPEndPoint remoteEndPoint = new IPEndPoint(ipAddress, port);

    var socket = new Socket(remoteEndPoint.AddressFamily, SocketType.Stream, ProtocolType.Tcp);
    Memory<byte> test = new Memory<byte>(new byte[2048]);
    var tokenSource = new CancellationTokenSource();

    InputSystem.onAnyButtonPress.Call(x =>
    {
        if (x is KeyControl keyboardCtl && keyboardCtl.keyCode == Key.E)
        {
            tokenSource.Cancel();
            Debug.Log("token cancel");
        }
    });

    Debug.Log("start receive");
    try
    {
        int count = await socket.ReceiveAsync(test, SocketFlags.None, tokenSource.Token);
        Debug.Log($"end receive : {count}");
    }
    catch (OperationCanceledException)
    {
        Debug.Log("ReceiveAsync cancelled");
    }
}

In this code snippet, I expect the ReceiveAsync operation to be cancelled when a certain button (e.g., the E key on the keyboard) is pressed. However, no matter what I try, I can’t seem to trigger the cancellation operation.

I suspect the issue may lie with the Socket.ReceiveAsync method, as it doesn’t seem to respond to the CancellationTokenSource passed to it. I expect the ReceiveAsync method to throw an OperationCanceledException when the cancellation token is triggered, but instead, it continues to block the thread

I discovered that I can indeed cancel the ReceiveAsync operation using the following approach:

int count = await socket.ReceiveAsync(test, SocketFlags.None, tokenSource.Token)
                    .AsUniTask<int>()
                    .AttachExternalCancellation(tokenSource.Token);

However, upon further consideration, I believe this approach may introduce some implicit issues. Using AttachExternalCancellation seems to work, but it feels like a workaround rather than a clean solution.

Have you tried your code with a regular dotnet framework application? In general if it works there but not in Unity it is a bug that we should investigate.

var receiveTask = _socket.ReceiveAsync(memory, SocketFlags.None);
if (_tcs.Task.IsCompleted || await Task.WhenAny((Task)receiveTask.AsTask(), _tcs.Task) == _tcs.Task)
break;