Tag Archives: TaskCompletionSource

Abstracting RabbitMQ RPC with TaskCompletionSource

I recently wrote about TaskCompletionSource, a little-known tool in .NET that is great for transforming arbitrary asynchrony into the Task-Based Asynchronous Pattern. That means you can hide the whole thing behind a simple and elegant async/await.

In this article, we’ll see this in practice as we implement the Remote Procedure Call (RPC) pattern in RabbitMQ. This is a fancy way of saying request/response, except that it all happens asynchronously! That’s right. No blocking.

The source code for this article is in the RabbitMqRpc folder at the Gigi Labs BitBucket Repository.

The RabbitMQ.Client NuGet package is necessary to make this code work. The client is written using an asynchronous Main() method, which requires at least C# 7.1 to compile.

RabbitMQ RPC Overview

You can think of RPC as request/response communication. We have a client asking a server to process some input and return the output in its response. However, this all happens asynchronously. The client sends the request on a request queue and forgets about it, rather than waiting for the response. Eventually, the server will (hopefully) process the request and send a response message back on a response queue.

The request and response can be matched on the client side by attaching a CorellationId to both the request and the response.

In this context, we don’t really talk about publishers and consumers, as is typical when talking about messaging frameworks. That’s because in order to make this work, both the client and the server must have both a publisher and a consumer.

Client: Main Program

For our client application, we’ll have the following main program code. We will implement an RpcClient that will hide the request/response plumbing behind a simple Task that we then await:

        static async Task Main(string[] args)
        {
            Console.Title = "RabbitMQ RPC Client";

            using (var rpcClient = new RpcClient())
            {
                Console.WriteLine("Press ENTER or Ctrl+C to exit.");

                while (true)
                {
                    string message = null;

                    Console.Write("Enter a message to send: ");
                    using (var colour = new ScopedConsoleColour(ConsoleColor.Blue))
                        message = Console.ReadLine();

                    if (string.IsNullOrWhiteSpace(message))
                        break;
                    else
                    {
                        var response = await rpcClient.SendAsync(message);

                        Console.Write("Response was: ");
                        using (var colour = new ScopedConsoleColour(ConsoleColor.Green))
                            Console.WriteLine(response);
                    }
                }
            }
        }

The program continuously asks for input, and sends that input as the request message. The server will process this message and return a response. Note that we are using the ScopedConsoleColour class from my “Scope Bound Resource Management in C#” article to colour certain sections of the output. Here is a taste of what it will look like:

While this console application only allows us to send one request at a time, the underlying approach is really powerful with APIs that can concurrently serve a multitude of clients. It is asynchronous and can scale pretty well, yet the consuming code sees none of the underlying complexity.

Client: Request Sending

The heart of this abstraction is the RpcClient class. In the constructor, we set up the typical plumbing: create a connection, channel, queues, and a consumer.

    public class RpcClient : IDisposable
    {
        private bool disposed = false;
        private IConnection connection;
        private IModel channel;
        private EventingBasicConsumer consumer;
        private ConcurrentDictionary<string,
            TaskCompletionSource<string>> pendingMessages;

        private const string requestQueueName = "requestqueue";
        private const string responseQueueName = "responsequeue";
        private const string exchangeName = ""; // default exchange

        public RpcClient()
        {
            var factory = new ConnectionFactory() { HostName = "localhost" };

            this.connection = factory.CreateConnection();
            this.channel = connection.CreateModel();

            this.channel.QueueDeclare(requestQueueName, true, false, false, null);
            this.channel.QueueDeclare(responseQueueName, true, false, false, null);

            this.consumer = new EventingBasicConsumer(this.channel);
            this.consumer.Received += Consumer_Received;
            this.channel.BasicConsume(responseQueueName, true, consumer);

            this.pendingMessages = new ConcurrentDictionary<string,
                TaskCompletionSource<string>>();
        }

        // ...
    }

A few other things to notice here:

  1. We are keeping a dictionary that allow us to match responses with the requests that generated them, based on a CorrelationId. We have already seen this approach in “TaskCompletionSource by Example“.
  2. This class implements IDisposable, as it has several resources that need to be cleaned up. While I don’t show the code for this for brevity’s sake, you can find it in the source code.
  3. We are not using exchanges here, so using an empty string for the exchange name allows us to use the default exchange and publish directly to the queue.

The SendAsync() method, which we saw being used in the main program, is implemented as follows:

        public Task<string> SendAsync(string message)
        {
            var tcs = new TaskCompletionSource<string>();
            var correlationId = Guid.NewGuid().ToString();

            this.pendingMessages[correlationId] = tcs;

            this.Publish(message, correlationId);

            return tcs.Task;
        }

Here, we are generating GUID to use as a CorrelationId, and we are adding an entry in the dictionary for this request. This dictionary maps the CorrelationId to a corresponding TaskCompletionSource. When the response arrives, it will set the result on this TaskCompletionSource, which enables the underlying task to complete. We return this underlying task, and that’s what the main program awaits. The main program will not be able to continue until the response is received.

In this method, we are also calling a private Publish() method, which takes care of the details of publishing to the request queue on RabbitMQ:

        private void Publish(string message, string correlationId)
        {
            var props = this.channel.CreateBasicProperties();
            props.CorrelationId = correlationId;
            props.ReplyTo = responseQueueName;

            byte[] messageBytes = Encoding.UTF8.GetBytes(message);
            this.channel.BasicPublish(exchangeName, requestQueueName, props, messageBytes);

            using (var colour = new ScopedConsoleColour(ConsoleColor.Yellow))
                Console.WriteLine($"Sent: {message} with CorrelationId {correlationId}");
        }

While this publishing code is for the most part pretty standard, we are using two particular properties that are especially suited for the RPC pattern. The first is CorrelationId, where we store the CorrelationId we generated earlier, and which the server will copy and send back as part of the response, enabling this whole orchestration. The second is the ReplyTo property, which is used to indicate to the server on which queue it should send the response. We don’t need it for this simple example since we are always using the same response queue, but this property enables the server to dynamically route responses where they are needed.

Server

The request eventually reaches a server which has a consumer waiting on the request queue. Its Main() method is mostly plumbing that enables this consumer to work:

        private static IModel channel;

        static void Main(string[] args)
        {
            Console.Title = "RabbitMQ RPC Server";

            var factory = new ConnectionFactory() { HostName = "localhost" };

            using (var connection = factory.CreateConnection())
            {
                using (channel = connection.CreateModel())
                {
                    const string requestQueueName = "requestqueue";
                    channel.QueueDeclare(requestQueueName, true, false, false, null);

                    // consumer

                    var consumer = new EventingBasicConsumer(channel);
                    consumer.Received += Consumer_Received;
                    channel.BasicConsume(requestQueueName, true, consumer);

                    Console.WriteLine("Waiting for messages...");
                    Console.WriteLine("Press ENTER to exit.");
                    Console.WriteLine();
                    Console.ReadLine();
                }
            }
        }

When a message is received, the Consumer_Received event handler processes the message:

        private static void Consumer_Received(object sender, BasicDeliverEventArgs e)
        {
            var requestMessage = Encoding.UTF8.GetString(e.Body);
            var correlationId = e.BasicProperties.CorrelationId;
            string responseQueueName = e.BasicProperties.ReplyTo;

            Console.WriteLine($"Received: {requestMessage} with CorrelationId {correlationId}");

            var responseMessage = Reverse(requestMessage);
            Publish(responseMessage, correlationId, responseQueueName);
        }

In this example, the server’s job is to reverse whatever messages it receives. Thus, each response will contain the same message as in the corresponding request, but backwards. This reversal code is taken from this Stack Overflow answer. Although trivial to implement, this serves as a reminder that there’s no need to reinvent the wheel if somebody already implemented the same thing (and quite well, at that) before you.

        public static string Reverse(string s)
        {
            char[] charArray = s.ToCharArray();
            Array.Reverse(charArray);
            return new string(charArray);
        }

Having computed the reverse of the request message, and extracted both the CorrelationId and ReplyTo properties, these are all passed to the Publish() method which sends back the response:

        private static void Publish(string responseMessage, string correlationId,
            string responseQueueName)
        {
            byte[] responseMessageBytes = Encoding.UTF8.GetBytes(responseMessage);

            const string exchangeName = ""; // default exchange
            var responseProps = channel.CreateBasicProperties();
            responseProps.CorrelationId = correlationId;

            channel.BasicPublish(exchangeName, responseQueueName, responseProps, responseMessageBytes);

            Console.WriteLine($"Sent: {responseMessage} with CorrelationId {correlationId}");
            Console.WriteLine();
        }

The response is sent back on the queue specified in the ReplyTo property of the request message. The response is also given the same CorrelationId as the request; that way the client will know that this response is for that particular request.

Client: Response Handling

When the response arrives, the client’s own consumer event handler will run to process it:

        private void Consumer_Received(object sender, BasicDeliverEventArgs e)
        {
            var correlationId = e.BasicProperties.CorrelationId;
            var message = Encoding.UTF8.GetString(e.Body);

            using (var colour = new ScopedConsoleColour(ConsoleColor.Yellow))
                Console.WriteLine($"Received: {message} with CorrelationId {correlationId}");

            this.pendingMessages.TryRemove(correlationId, out var tcs);
            if (tcs != null)
                tcs.SetResult(message);
        }

The client extracts the CorrelationId from the response, and uses it to get the TaskCompletionSource for the corresponding request. If the TaskCompletionSource is found, then its result is set to the content of the response. This causes the underlying task to complete, and thus the caller awaiting that task will be able to resume and work with the result.

If the TaskCompletionSource is not found, then we ignore the response, and there is a reason for this:

“You may ask, why should we ignore unknown messages in the callback queue, rather than failing with an error? It’s due to a possibility of a race condition on the server side. Although unlikely, it is possible that the RPC server will die just after sending us the answer, but before sending an acknowledgment message for the request. If that happens, the restarted RPC server will process the request again. That’s why on the client we must handle the duplicate responses gracefully, and the RPC should ideally be idempotent.” — RabbitMQ RPC tutorial

Demo

If we run both the client and server, we can enter messages in the client, one by one. The client publishes each message on the request queue and waits for the response, at which point it allows the main program to continue by setting the result of that request’s TaskCompletionSource.

Summary

What we have seen in this article is the same material I had explained in “TaskCompletionSource by Example“, but with a real application to RabbitMQ.

A TaskCompletionSource has an underlying Task that can represent a pending request. By giving each request an ID, you can keep track of it as the corresponding response should carry the same ID. A mapping between request IDs and TaskCompletionSource can easily be kept in a dictionary. When a response arrives, its corresponding entry in the dictionary can be found, and the Task can be completed. Any client code awaiting this Task may then resume.

TaskCompletionSource by Example

In this article, we’ll learn how to use TaskCompletionSource. It’s one of those tools which you will rarely need to use, but when you do, you’ll be glad that you knew about it. Let’s dive right into it.

Basic Usage

The source code for this section is in the TaskCompletionSource1 folder at the Gigi Labs BitBucket Repository.

Let’s create a new console application, and in Main(), we’ll have my usual workaround for running asynchronous code in a console application:

        static void Main(string[] args)
        {
            Run();
            Console.ReadLine();
        }

In the Run() method, we have a simple example showing how TaskCompletionSource works:

        static async void Run()
        {
            var tcs = new TaskCompletionSource<bool>();

            var fireAndForgetTask = Task.Delay(5000)
                                        .ContinueWith(task => tcs.SetResult(true));

            await tcs.Task;
        }

TaskCompletionSource is just a wrapper for a Task, giving you control over its completion. Thus, a TaskCompletionSource<bool> will contain a Task<bool>, and you can set the bool result based on your own logic.

Here, we are using TaskCompletionSource as a synchronization mechanism. Our main thread spawns off an operation and waits for its result, using the Task in the TaskCompletionSource. Even if the operation is not Task-based, it can set the result of the Task in the TaskCompletionSource, allowing the main thread to resume its execution.

Let’s add some diagnostic code so that we can understand what’s going on from the output:

        static async void Run()
        {
            var stopwatch = Stopwatch.StartNew();

            var tcs = new TaskCompletionSource<bool>();

            Console.WriteLine($"Starting... (after {stopwatch.ElapsedMilliseconds}ms)");

            var fireAndForgetTask = Task.Delay(5000)
                                        .ContinueWith(task => tcs.SetResult(true));

            Console.WriteLine($"Waiting...  (after {stopwatch.ElapsedMilliseconds}ms)");

            await tcs.Task;

            Console.WriteLine($"Done.       (after {stopwatch.ElapsedMilliseconds}ms)");

            stopwatch.Stop();
        }

And here is the output:

Starting... (after 0ms)
Waiting...  (after 41ms)
Done.       (after 5072ms)

As you can see, the main thread waited until tcs.SetResult(true) was called; this triggered completion of the TaskCompletionSource’s underlying task (which the main thread was awaiting), and allowed the main thread to resume.

Aside from SetResult(), TaskCompletionSource offers methods to cancel a task or fault it with an exception. There are also safe Try...() equivalents:

SDK Example

The source code for this section is in the TaskCompletionSource2 folder at the Gigi Labs BitBucket Repository.

One scenario where I found TaskCompletionSource to be extremely well-suited is when you are given a third-party SDK which exposes events. Imagine this: you submit an order via an SDK method, and it gives you an ID for that order, but not the result. The SDK goes off and does what it has to do to perhaps talk to an external service and have the order processed. When this eventually happens, the SDK fires an event to notify the calling application on whether the order was placed successfully.

We’ll use this as an example of the SDK code:

    public class MockSdk
    {
        public event EventHandler<OrderOutcome> OnOrderCompleted;

        public Guid SubmitOrder(decimal price)
        {
            var orderId = Guid.NewGuid();

            // do a REST call over the network or something
            Task.Delay(3000).ContinueWith(task => OnOrderCompleted(this,
                new OrderOutcome(orderId, true)));

            return orderId;
        }
    }

The OrderOutcome class is just a simple DTO:

    public class OrderOutcome
    {
        public Guid OrderId { get; set; }
        public bool Success { get; set; }

        public OrderOutcome(Guid orderId, bool success)
        {
            this.OrderId = orderId;
            this.Success = success;
        }
    }

Notice how MockSdk‘s SubmitOrder does not return any form of Task, and we can’t await it. This doesn’t necessarily mean that it’s blocking; it might be using another form of asynchrony such as the Asynchronous Programming Model or a messaging framework with a request-response fashion (such as RPC over RabbitMQ).

At the end of the day, this is still asynchrony, and we can use TaskCompletionSource to build a Task-based Asynchronous Pattern abstraction over it (allowing the application to simply await the call).

First, we start building a simple proxy class that wraps the SDK:

    public class SdkProxy
    {
        private MockSdk sdk;

        public SdkProxy()
        {
            this.sdk = new MockSdk();
            this.sdk.OnOrderCompleted += Sdk_OnOrderCompleted;
        }

        private void Sdk_OnOrderCompleted(object sender, OrderOutcome e)
        {
            // TODO
        }
    }

We then add a dictionary, which allows us to relate each OrderId to its corresponding TaskCompletionSource. Using a ConcurrentDictionary instead of a normal Dictionary helps deal with multithreading scenarios without needing to lock:

        private ConcurrentDictionary<Guid,
            TaskCompletionSource<bool>> pendingOrders;
        private MockSdk sdk;

        public SdkProxy()
        {
            this.pendingOrders = new ConcurrentDictionary<Guid,
                TaskCompletionSource<bool>>();

            this.sdk = new MockSdk();
            this.sdk.OnOrderCompleted += Sdk_OnOrderCompleted;
        }

The proxy class exposes a SubmitOrderAsync() method:

        public Task SubmitOrderAsync(decimal price)
        {
            var orderId = sdk.SubmitOrder(price);

            Console.WriteLine($"OrderId {orderId} submitted with price {price}");

            var tcs = new TaskCompletionSource<bool>();
            this.pendingOrders.TryAdd(orderId, tcs);

            return tcs.Task;
        }

This method calls the underlying SubmitOrder() in the SDK, and uses the returned OrderId to add a new TaskCompletionSource in the dictionary. The TaskCompletionSource’s underlying Task is returned, so that the application can await it.

        private void Sdk_OnOrderCompleted(object sender, OrderOutcome e)
        {
            string successStr = e.Success ? "was successful" : "failed";
            Console.WriteLine($"OrderId {e.OrderId} {successStr}");

            this.pendingOrders.TryRemove(e.OrderId, out var tcs);
            tcs.SetResult(e.Success);
        }

When the SDK fires a completion event, the proxy will remove the TaskCompletionSource from the pending order and set its result. The application awaiting the underlying task will resume and take a decision depending on the logic.

We can test this with the following program code in a console application:

        static void Main(string[] args)
        {
            Run();
            Console.ReadLine();
        }

        static async void Run()
        {
            var sdkProxy = new SdkProxy();

            await sdkProxy.SubmitOrderAsync(10);
            await sdkProxy.SubmitOrderAsync(20);
            await sdkProxy.SubmitOrderAsync(5);
            await sdkProxy.SubmitOrderAsync(15);
            await sdkProxy.SubmitOrderAsync(4);
        }

The output shows that the program did indeed wait for each order to complete before starting the next one:

OrderId 3e2d4577-8bbb-46b7-a5df-2efec23bae6b submitted with price 10
OrderId 3e2d4577-8bbb-46b7-a5df-2efec23bae6b was successful
OrderId e22425b9-3aa3-48db-a40f-8b8cfbdcd3af submitted with price 20
OrderId e22425b9-3aa3-48db-a40f-8b8cfbdcd3af was successful
OrderId 3b5a2602-a5d2-4225-bbdb-10642a63f7bc submitted with price 5
OrderId 3b5a2602-a5d2-4225-bbdb-10642a63f7bc was successful
OrderId ffd61cea-343e-4a9c-a76f-889598a45993 submitted with price 15
OrderId ffd61cea-343e-4a9c-a76f-889598a45993 was successful
OrderId b443462c-f949-49b9-a6f0-08bbbb82fe7e submitted with price 4
OrderId b443462c-f949-49b9-a6f0-08bbbb82fe7e was successful

Summary

Use TaskCompletionSource to adapt an arbitrary form of asynchrony to use Tasks, and enable elegant async/await usage.

Do not use it simply expose an asynchronous wrapper for a synchronous method. You should either not do that at all, or use Task.FromResult() instead.

If you’re concerned that the the asynchronous call might never resume, consider adding a timeout.