How can I await my custom data flow block completion?

Question

I'm working with ITargetBlock and have somewhat successfully created a custom data flow block. However, I can't wrap my head around implementing the "Completion" task. I currently define this as:

public Task Completion { get; }

Once I build my data flow line, I attempt to await this completion task but get a "null" object referenced. This is due to the fact that I'm not assigning the completion task to anything.

What should I assign this to? I would think it would automatically be associated with the data flow block but it defaults to nothing...

What should I assign Completion to?

Answer 1

When you implement your code Completion should complete once you have called Complete() and all "in flight" work finishes.

This can be done with a TaskCompletionSource

public Example<TInput>() : ITargetBlock 
{
    private TaskCompletionSource<Object> tcs = new TaskCompletionSource<Object>()
    public Task Completion { get; }

    public Example()
    {
        Completion = tcs.Task;
    }

    public void Complete()
    {
        // We run this on a background thread because we don't want the call to Complete be blocking.
        Task.Run(() => {
            // Wait here for any currently executing async work your dataflow block does to finish.
            // ...

            tcs.TrySetResult(null);
        }
    }

    public void Fault (Exception exception)
    {
        // Cancel here any running work.
        // ...

        tcs.TrySetException(exception);
    }

    public DataflowMessageStatus OfferMessage(DataflowMessageHeader messageHeader, TInput messageValue, ISourceBlock<in TInput> source, bool consumeToAccept)
    {
        if (Completion.IsCompleted)
        {
            return DataflowMessageStatus.DecliningPermanently;
        }
        // ...
    }
}