Is it acceptable to use TaskCompletionSource as a WaitHandle substitute?

Question

My code handles a TCP connection to a remote host, with a ConcurrentQueue to store outgoing messages. It's intended to run in a single thread. The lifetime of the connection is contained within RunAsync while a separate object contains the "public state" of the connection:

class PublicState
{
    internal readonly ConcurrentQueue<Message> OutgoingMessageQueue = new ConcurrentQueue<Message>();
    internal TaskCompletionSource<Object> OutgoingMessageTcs = null;

    internal readonly TaskCompletionSource<Object> ConnectedTcs = new TaskCompletionSource<Object>();

    public void EnqueueMessages(IEnumerable<Message> messages)
    {
        foreach( Message m in messages ) this.OutgoingMessageQueue.Enqueue( m);
        if( this.OutgoingMessageTcs == null ) this.OutgoingMessageTcs = new TaskCompletionSource<Object>();
        this.OutgoingMessageTcs.SetResult( null );
    }
}

static async Task RunAsync(IPEndPoint endPoint, PublicState state)
{
    using( TcpClient tcp = new TcpClient() )
    {
        await tcp.ConnectAsync( endPoint.Address, endPoint.Port ).ConfigureAwait(false);

        Byte[] reusableBuffer = new Byte[ 4096 ];

        using( NetworkStream ns = tcp.GetStream() )
        {
            state.ConnectedTcs.SetResult( null );

            Task<Int32> nsReadTask = null;

            while( tcp.Connected )
            {
                if( !state.writeQueue.IsEmpty )
                {
                    await WriteMessagesAsync( ... ).ConfigureAwait( false );
                }

                if( ns.DataAvailable )
                {
                    await ReadMessagesAsync( ... ).ConfigureAwait( false );
                }

                // Wait for new data to arrive from remote host or for new messages to send:
                if( state.OutgoingMessageTcs == null ) state.OutgoingMessageTcs = new TaskCompletionSource<Object>();
                if( nsReadTask == null ) nsReadTask = ns.ReadAsync( reusableBuffer, 0, 0 ).ConfigureAwait( false );

                Task c = await Task.WhenAny( state.OutgoingMessageTcs,  nsReadTask ).ConfigureAwait( false );
                if( c == state.OutgoingMessageTcs.Task ) state.OutgoingMessageTcs = null;
                else if( c == nsReadTask ) nsReadTask = null;
            }
        }
    }
}

Used like this:

public async Task Main(String[] args)
{
    PublicState state = new PublicState();
    Task clientTask = Client.RunAsync( new IPEndPoint(args[0]), state );

    await state.ConnectedTcs.Task; // awaits until TCP connection is established

    state.EnqueueMessage( new Message("foo") );
    state.EnqueueMessage( new Message("bar") );
    state.EnqueueMessage( new Message("baz") );

    await clientTask; // awaits until the TCP connection is closed
}

This code works, but I don't like it: it feels like I'm using TaskCompletionSource which is meant to represent an actual Task or some kind of background operation, whereas I'm really using TaskCompletionSource as a kind of cheap EventWaitHandle. I'm not using EventWaitHandle because it's IDisposable (I don't want to risk leaking native resources) and it lacks a WaitAsync or WaitOneAsync method. I could use SemaphoreSlim (which is awaitable, but wraps an EventWaitHandle) but my code doesn't really represent a good use of a semaphore.

Is my use of TaskCompletionSource<T> acceptable, or is there a better way to "un-await" execution in RunAsync when an item is added to OutgoingMessageQueue?

Another reason I feel it's "wrong" is that TaskCompletionSource<T> can only be used once, then it needs replacing. I'm keen to avoid extraneous allocations.

@Fildor Another thread or elsewhere in the host program will append to the queue - the `await` inside the `while( tcp.Connected )` loop needs to "un-await" when that happens. Using `TryDequeue` doesn't help there - unless I poll it in a tight loop which defeats the point of async IO. — Dai, Apr 09 '18 at 07:11
Yep, just saw what you did there. But it feels like a pretty convoluted way of doing things to me. Still trying to wrap my head around how to improve. — Fildor, Apr 09 '18 at 07:13
Since you can both read and write to `NetworkStream` at the same time - why you insist on doing everything in one loop? — Evk, Apr 09 '18 at 07:16
@Evk Callers to `EnqueueMessage` may (but not necessarily) be on different threads (this probably reveals a major issue with my design now...), otherwise I'm keen to have this design work with a single thread. — Dai, Apr 09 '18 at 07:18
`TaskCompletionSource` is for async wrapping of an event, which can be unpredictable / user driven, but in this case it is predictable as it waits for the `NetworkStream`, which is a blocking operation. In my understanding it is not really a value add in this case, but as you said you are quite keen to save the thread and there's no WaitAsync on the event wait handle, therefore this looks like a fair workaround. Other points would be explicitly starting `RunAsync` task and also assign `SetException` for the `TaskCompletionSource`, else it might be a deadlock — Mrinal Kamboj, Apr 09 '18 at 07:27
You already have a `Task` that's suitable for this purpose, the one returned from `ConnectAsync`. If you didn't have one, though, `TaskCompletionSource` is perfectly fine for introducing synchronization points. — Ben Voigt, Apr 09 '18 at 07:28

Evk · Answer 1 · 2018-04-09T07:57:27.323

If I understood you correctly - TPL BufferBlock might be what you need. Analog of current Enqueue is Post, and you can receive next value via ReceiveAsync extension method.

So with BufferBlock your code becomes something like this:

class PublicState {
    internal readonly BufferBlock<Message> OutgoingMessageQueue = new BufferBlock<Message>();
    internal readonly TaskCompletionSource<Object> ConnectedTcs = new TaskCompletionSource<Object>();

    public void EnqueueMessage(Message message) {
        this.OutgoingMessageQueue.Post(message);
    }
}

static async Task RunAsync(IPEndPoint endPoint, PublicState state) {
    using (TcpClient tcp = new TcpClient()) {
        await tcp.ConnectAsync(endPoint.Address, endPoint.Port).ConfigureAwait(false);

        Byte[] reusableBuffer = new Byte[4096];

        using (NetworkStream ns = tcp.GetStream()) {
            state.ConnectedTcs.SetResult(null);

            Task<Int32> nsReadTask = null;
            Task<Message> newMessageTask = null;
            while (tcp.Connected) {
                // Wait for new data to arrive from remote host or for new messages to send:
                if (nsReadTask == null)
                    nsReadTask = ns.ReadAsync(reusableBuffer, 0, 0);
                if (newMessageTask == null)
                    newMessageTask = state.OutgoingMessageQueue.ReceiveAsync();
                var completed = await Task.WhenAny(nsReadTask, newMessageTask).ConfigureAwait(false);
                if (completed == newMessageTask) {
                    var result = await newMessageTask;
                    // do stuff
                    newMessageTask = null;
                }
                else {
                    var bytesRead = await nsReadTask;
                    nsReadTask = null;
                }
            }
        }
    }
}

As a bonus, this version is (I think) thread-safe, while your current version is not, because you are doing non-thread-safe things with OutgoingMessageTcs from potentially multiple threads (thread of RunAsync and thread of EnqueueMessages caller).

If for some reason you don't like BufferBlock - you can use AsyncCollection from Nito.AsyncEx nuget package in exactly the same way. Initialization becomes:

internal readonly AsyncCollection<Message> OutgoingMessageQueue = new AsyncCollection<Message>(new ConcurrentQueue<Message>());

And fetching:

if (newMessageTask == null)
   newMessageTask = state.OutgoingMessageQueue.TakeAsync();

Out of curiosity, I poked around the definition of `BufferBlock.ReceiveAsync` and I see that it creates a new instance of `DataflowBlock.ReceiveTarget.Task` which derives from `TaskCompletionSource` - so ultimately your code does the same as mine, just in a more round-about way. — Dai, Apr 09 '18 at 07:39
Maybe, but I'd argue that it does so in a more clean and robust way, since I believe that your current way is not thread safe. — Evk, Apr 09 '18 at 07:44
And anyway, using `TaskCompletionSource` for asynchronous waiting for certain event is perfectly fine. Even `SemaphoreSlim.WaitAsync` will do similar thing (not `TaskCompletionSource` directly, but still very close). — Evk, Apr 09 '18 at 07:52

score 1 · Answer 2 · answered Jun 06 '18 at 22:44

To back up what others have mentioned, it does look like Microsoft's documentation mentions and even encourages developing a Semaphore class which is written on top of the Task objects here:

You can also build an asynchronous semaphore that does not rely on wait handles and instead works completely with tasks. To do this, you can use techniques such as those discussed in Consuming the Task-based Asynchronous Pattern for building data structures on top of Task.

This does make me wonder why such a prepackaged class does not already exist, but it certainly shows that this is fine.

Is it acceptable to use TaskCompletionSource as a WaitHandle substitute?

2 Answers2

Linked