Dispatch multiple publishers to single handler with result

Question

This code defines a MessageBus used to handle messages from many producers to a single consumer. The bus allows multiple publishers to send messages asynchronously to a single message handler. This is not a Pub-Sub model, instead, it's a many publisher and single consumer with the ability to send a result back to the publisher. The publisher can await for the message to be sent using the PostAsync methods. Or await for the completion or optional result of the handler by using the SendAsync methods. All messages and handlers are treated asynchronously regardless of their implementation, however, a synchronous implementation will not block the publisher. This setup centralizes who actually handles messages instead of keeping track of many subscriptions.

Underneath the hood this is using TPL-Dataflow to flow messages and SimpleInjector for IoC and wiring up handlers to messages.

First, the interface definition of the example message bus:

public interface IMessageBus
{
    /// <summary>
    /// Posts a message to the bus and returns a Task representing the acceptance of the message.
    /// </summary>
    /// <typeparam name="TMessage"></typeparam>
    /// <param name="message"></param>
    /// <returns></returns>
    Task PostAsync<TMessage>(TMessage message) where TMessage : IMessage;
    /// <summary>
    /// Posts a message to the bus and returns a Task representing the acceptance of the message with an error handling delegate.
    /// </summary>
    /// <typeparam name="TMessage"></typeparam>
    /// <param name="message"></param>
    /// <param name="OnError"></param>
    /// <returns></returns>
    Task PostAsync<TMessage>(TMessage message, Action<Exception> OnError) where TMessage : IMessage;   
    /// <summary>
    /// Sends a message to the bus and returns a Task representing completion of handling the message.
    /// </summary>
    /// <typeparam name="TMessage"></typeparam>
    /// <param name="message"></param>
    /// <returns></returns>
    Task SendAsync<TMessage>(TMessage message) where TMessage : IMessage;
    /// <summary>
    /// Sends a message to the bus and returns a Task representing completion of handling the message and yields the result of the handling.
    /// </summary>
    /// <typeparam name="TMessage"></typeparam>
    /// <typeparam name="TReply"></typeparam>
    /// <param name="message"></param>
    /// <returns></returns>
    Task<TReply> SendAsync<TMessage, TReply>(TMessage message) where TMessage : IMessage<TReply>;
}

And the implementation:

public class MessageBus : IMessageBus
{
    private IContainer Container { get; }

    public MessageBus(IContainer container) => Container = container ?? throw new ArgumentNullException(nameof(container));

    public Task PostAsync<TMessage>(TMessage message) where TMessage : IMessage => PostAsync(message, null);

    public Task PostAsync<TMessage>(TMessage message, Action<Exception> OnError) where TMessage : IMessage
    {
        var block = GetHandlerBlock<TMessage>(async msg =>
        {
            try
            {
                var handler = Container.GetInstance<IMessageHandler<TMessage>>();
                await handler.HandleAsync(msg);
            }
            catch (Exception ex)
            {
                if (OnError == null) return;
                OnError(ex);
            }
        });
        return block.SendAsync(message);
    }

    public Task SendAsync<TMessage>(TMessage message) where TMessage : IMessage
    {
        var tcs = new TaskCompletionSource<bool>();
        var block = GetHandlerBlock<TMessage>(async msg =>
        {
            try
            {
                var handler = Container.GetInstance<IMessageHandler<TMessage>>();
                await handler.HandleAsync(msg);
                tcs.TrySetResult(true);
            }
            catch (Exception ex)
            {
                tcs.TrySetException(ex);
            }
        });
        block.Post(message);
        return tcs.Task;
    }

    public Task<TReply> SendAsync<TMessage, TReply>(TMessage message) where TMessage : IMessage<TReply>
    {
        var tcs = new TaskCompletionSource<TReply>();
        var block = GetHandlerBlock<TMessage>(async msg =>
        {
            try
            {
                var handler = Container.GetInstance<IMessageHandler<TMessage, TReply>>();
                var reply = await handler.HandleAsync(msg);
                tcs.TrySetResult(reply);
            }
            catch (Exception ex)
            {
                tcs.TrySetException(ex);
            }
        });
        block.Post(message);
        return tcs.Task;
    }

    private ITargetBlock<TMessage> GetHandlerBlock<TMessage>(Func<TMessage, Task> action) => new ActionBlock<TMessage>(action);
}

The publisher has a few options:

1. Post a message, awaiting only for the message to be accepted.
2. Post a message with an error handling delegate, awaiting only for the message to be accepted.
3. Send a message awaiting completion of processing.
4. Send a message awaiting a reply from the handler.

And the supporting interfaces:

public interface IMessage { }
public interface IMessage<TReply> : IMessage { }
public interface IMessageHandler<TMessage> where TMessage : IMessage
{
    Task HandleAsync(TMessage message);
}
public interface IMessageHandler<TMessage, TReply> where TMessage : IMessage<TReply>
{
    Task<TReply> HandleAsync(TMessage message);
}
public interface IContainer
{
    TInstance GetInstance<TInstance>() where TInstance : class;
}

Answer 1

A few comments:

var handler = Container.GetInstance<IMessageHandler<TMessage, TReply>>();

You're looking up the IMessageHandler inside the handler block, which means you're looking it up for every message. I'm not that familiar with SimpleInjector, but you may want to memoize this or somehow cache the handler for a given message type.

private ITargetBlock<TMessage> GetHandlerBlock<TMessage>(Func<TMessage, Task> action) => new ActionBlock<TMessage>(action);

You're using ActionBlock directly, which might be fine, but message busses often use some form of queueing and FIFO processing internally. It might be useful to have a BufferBlock where the messages are enqueued when posted. This StackOverflow question elaborates on the advantages and appropriate use cases of BufferBlock.

return block.SendAsync(message);

Your PostAsync method returns the Task from ActionBlock.SendAsync. I think this will complete when the message is processed, not when it is enqueued. Take a look at the description of the return value on the MSDN page. Also, the result of this Task will be true or false based on whether the message was accepted and consumed or declined by the Target. Is that what you wanted? If so, I would clarify that in your documentation comments.