Search code examples
c#.netunit-testingtpl-dataflow

TPL Dataflow Broadcastblock discards last message


I have a pretty simple problem. I need a way to easily perform some processing on messages that take some time. While processing, new requests might be entered, but all requests except the last one can be discarded.

So I thought that the TPL Broadcastblock should do just that, looking at the documentation and posts on, well, StackExchange for example. I created the following solution and added some unit tests for it, but in the unit tests, sometimes the last item is not sent.

This is not what I expected. If it should drop anything I'd say it should drop the first item, since it should overwrite it's buffer of 1 if it can't process a message. Can anyone see what it is?
Any help would be greatly appreciated!

Here's the code for the block:

/// <summary>
/// This block will take items and perform the specified action on it. Any incoming messages while the action is being performed
/// will be discarded.
/// </summary>
public class DiscardWhileBusyActionBlock<T> : ITargetBlock<T>
{
    private readonly BroadcastBlock<T> broadcastBlock;

    private readonly ActionBlock<T> actionBlock;

    /// <summary>
    /// Initializes a new instance of the <see cref="DiscardWhileBusyActionBlock{T}"/> class.
    /// Constructs a SyncFilterTarget{TInput}.
    /// </summary>
    /// <param name="actionToPerform">Thing to do.</param>
    public DiscardWhileBusyActionBlock(Action<T> actionToPerform)
    {
        if (actionToPerform == null)
        {
            throw new ArgumentNullException(nameof(actionToPerform));
        }

        this.broadcastBlock = new BroadcastBlock<T>(item => item);
        this.actionBlock = new ActionBlock<T>(actionToPerform, new ExecutionDataflowBlockOptions { BoundedCapacity = 1, MaxDegreeOfParallelism = 1 });
        this.broadcastBlock.LinkTo(this.actionBlock);
        this.broadcastBlock.Completion.ContinueWith(task => this.actionBlock.Complete());
    }

    public DataflowMessageStatus OfferMessage(DataflowMessageHeader messageHeader, T messageValue, ISourceBlock<T> source, bool consumeToAccept)
    {
        return ((ITargetBlock<T>)this.broadcastBlock).OfferMessage(messageHeader, messageValue, source, consumeToAccept);
    }

    public void Complete()
    {
        this.broadcastBlock.Complete();
    }

    public void Fault(Exception exception)
    {
        ((ITargetBlock<T>)this.broadcastBlock).Fault(exception);
    }

    public Task Completion => this.actionBlock.Completion;
}

And here's the code for the test:

[TestClass]
public class DiscardWhileBusyActionBlockTest
{
    [TestMethod]
    public void PostToConnectedBuffer_ActionNotBusy_MessageConsumed()
    {
        var actionPerformer = new ActionPerformer();

        var block = new DiscardWhileBusyActionBlock<int>(actionPerformer.Perform);
        var buffer = DiscardWhileBusyActionBlockTest.SetupBuffer(block);

        buffer.Post(1);

        DiscardWhileBusyActionBlockTest.WaitForCompletion(buffer, block);

        var expectedMessages = new[] { 1 };
        actionPerformer.LastReceivedMessage.Should().BeEquivalentTo(expectedMessages);
    }

    [TestMethod]
    public void PostToConnectedBuffer_ActionBusy_MessagesConsumedWhenActionBecomesAvailable()
    {
        var actionPerformer = new ActionPerformer();

        var block = new DiscardWhileBusyActionBlock<int>(actionPerformer.Perform);
        var buffer = DiscardWhileBusyActionBlockTest.SetupBuffer(block);

        actionPerformer.SetBusy();

        // 1st message will set the actionperformer to busy, 2nd message should be sent when
        // it becomes available.
        buffer.Post(1);
        buffer.Post(2);

        actionPerformer.SetAvailable();

        DiscardWhileBusyActionBlockTest.WaitForCompletion(buffer, block);

        var expectedMessages = new[] { 1, 2 };
        actionPerformer.LastReceivedMessage.Should().BeEquivalentTo(expectedMessages);
    }

    [TestMethod]
    public void PostToConnectedBuffer_ActionBusy_DiscardMessagesInBetweenAndProcessOnlyLastMessage()
    {
        var actionPerformer = new ActionPerformer();

        var block = new DiscardWhileBusyActionBlock<int>(actionPerformer.Perform);
        var buffer = DiscardWhileBusyActionBlockTest.SetupBuffer(block);

        actionPerformer.SetBusy();

        buffer.Post(1);
        buffer.Post(2);
        buffer.Post(3);
        buffer.Post(4);
        buffer.Post(5);

        actionPerformer.SetAvailable();

        DiscardWhileBusyActionBlockTest.WaitForCompletion(buffer, block);

        var expectedMessages = new[] { 1, 5 };
        actionPerformer.LastReceivedMessage.Should().BeEquivalentTo(expectedMessages);
    }

    private static void WaitForCompletion(IDataflowBlock source, IDataflowBlock target)
    {
        source.Complete();
        target.Completion.Wait(TimeSpan.FromSeconds(1));
    }

    private static BufferBlock<int> SetupBuffer(ITargetBlock<int> block)
    {
        var buffer = new BufferBlock<int>();
        buffer.LinkTo(block);
        buffer.Completion.ContinueWith(task => block.Complete());
        return buffer;
    }

    private class ActionPerformer
    {
        private readonly ManualResetEvent resetEvent = new ManualResetEvent(true);

        public List<int> LastReceivedMessage { get; } = new List<int>();

        public void Perform(int message)
        {
            this.resetEvent.WaitOne(TimeSpan.FromSeconds(3));
            this.LastReceivedMessage.Add(message);
        }

        public void SetBusy()
        {
            this.resetEvent.Reset();
        }

        public void SetAvailable()
        {
            this.resetEvent.Set();
        }
    }
}

Solution

  • When you do level the BoundedCapacity of the action block to 1, it means that, if it does processing, and already has item in it's queue, it will discard the message, which will run out of scope. So basically what happening is that your block do it's job, rejecting the new message while the buffer is full. After that the broadcast block is done, as whole messages are sent to the recipients, and it calls the Completion, which finishes whole pipeline.

    You need either to check the returned boolean value of Post for the last messages, or, more probably, store the last message in some variable, ensuring that it will go to the pipeline. Looks like you better not to use the BroadcastBlock, as its purpose to provide a copy of the message to the number of linked blocks, and just write your logic by yourself. Maybe you can use a simple BufferBlock instead.

    Update: OfferMessagemethod also do provide the information about the message being offered. I think you don't need a buffer block at all, as you have to deal with non-default logic for your pipeline. More easy for to have a field like _lastMessage, storing the last value in it, and erasing it when the request is accepted by actionBlock. You may even completely remove the dataflow dependency, as all you do is calling method for the request.

    Side notes: you can link blocks with completion propagation set in options:

    var linkOptions = new DataflowLinkOptions { PropagateCompletion = true };
    this.broadcastBlock.LinkTo(this.actionBlock, linkOptions);
    

    This can remove some of your code with potentially dangerous ContinueWith usage. Also you can await broadcastBlock.SendAsync() instead of Post, if you need the asynchronous behavior.