Conform IAsyncEnumerable to Dataflow ISourceBlock

Viewed 91

I have an existing transformation function Func<IAsyncEnumerable<TIn>, IAsyncEnumerable<TOut>> that I want to use inside a Dataflow pipeline. In order to link this transformation function, I need the input side of this function to map to ITargetBlock<TIn> and the output side to ISourceBlock<TOut>.

To handle the input-side (target), I've used a BufferBlock<TIn> with ToAsyncEnumerable() extension shown here.

However I can't seem to get the output-side (source) in a shape that I want. I currently have the following:

IPropagatorBlock<TIn, TOut> Adapt<TIn, TOut>(
    Func<IAsyncEnumerable<TIn>, IAsyncEnumerable<TOut>> transformation)
{
    var target = new BufferBlock<TIn>(new DataflowBlockOptions { BoundedCapacity = 1024 });
    var generator = transformation(target.ToAsyncEnumerable());
    var source = new BufferBlock<TOut>(new DataflowBlockOptions { BoundedCapacity = 1024 });
    var task = Task.Run(async () =>
    {
        await foreach (var item in generator)
        {
            await source.SendAsync(item);
        }
        source.Complete()
    });
    return DataflowBlock.Encapsulate(target, source);
}

My goal here is that the transformation function can execute on an enumerated dataset (millions of database rows) without loading these entries in memory. I have currently implemented that using BoundedCapacity, but performance takes a hit when doing so. I would prefer a solution without having to introduce BufferBlocks.

What I do not like the additional BufferBlock and the spawning of a Task. Reading through Dataflow's source code, it seems like a lot of effort was put into avoiding creation of tasks. However that seems the primitive I need to iterate over the enumerable and make the entries flow through my transformation function.

Is there another obvious way to adapt my transformation function such that I can use it as a Dataflow Block (both consumer + producer)?

2 Answers

I think there's a fundamental problem with what you are trying to do. Both Dataflow and AsyncEnumerable are distinct Task-based asynchronous programming methods, which don't combine well.

Dataflow blocks work on data. But an AsyncEnumerable is not data: it represents an interface to asynchronously retrieve data, but it itself is not data. So using an AsyncEnumerable within a Dataflow block doesn't make much sense. This is what you encounter in your solution: in order to get the data you have to enummerate the AsyncEnumerable, process it, and pass it on as AsyncEnumerable. You would always need to have some inner Task to handles asynchronous enumeration for you. And in the end you're not using any of the advantage/features the Dataflow approach, so it's redundant. It actually only adds unnecessary noise and ballast.

Better just to use the transformation directly on the AsyncEnumerable.

However, If you still really want to use Dataflow and the source is an AsyncEnumerable, you would have to make a AsyncEnumerable-to-Dataflow adapter (the opposite block of what's in your link), process the data item-by-item in the Dataflow and finally use the Dataflow-to-AsyncEnumerable adapter in your link.

I don't think that it's possible to implement the Adapt method using only one dataflow block. The ITargetBlock<TIn> and the ISourceBlock<TOut> facades could be the same block only if the TPL Dataflow library offered a built-in block with the exact functionality that you are trying to implement. But there is no built-in block that can be initialized with a Func<IAsyncEnumerable<TIn>, IAsyncEnumerable<TOut>> lambda. The closest block available will be the TransformManyBlock<TIn, TOut> with Func<TIn, IAsyncEnumerable<TOut>> parameter, which is expected with the upcoming .NET 7, and could not be used for this purpose either.

My suggestion is to take care of propagating the exception that might be received by the ITargetBlock<TIn> facade, originated from the linked dataflow block upstream. For this purpose you can use the Fault method.

Also I would suggest to replace the fire-and-forget Task.Run with an async void operation that is invoked on the ThreadPool. This way if the unthinkable happens and the code in the catch block fails, the error will be surfaced immediately as an unhandled exception that will crash the process, instead of just leaving the process in a hanging state.

IPropagatorBlock<TIn, TOut> Adapt<TIn, TOut>(
    Func<IAsyncEnumerable<TIn>, IAsyncEnumerable<TOut>> transformation)
{
    DataflowBlockOptions options = new() { BoundedCapacity = 1024 };
    BufferBlock<TIn> target = new(options);
    IAsyncEnumerable<TOut> output = transformation(target.ToAsyncEnumerable());
    BufferBlock<TOut> source = new(options);
    ThreadPool.QueueUserWorkItem(async _ =>
    {
        try
        {
            await foreach (TOut item in output)
            {
                await source.SendAsync(item);
            }
            source.Complete();
        }
        catch (Exception ex)
        {
            ((IDataflowBlock)source).Fault(ex);
        }
    });
    return DataflowBlock.Encapsulate(target, source);
}

As a side-note, the TPL Dataflow library is intended for "coarse-grained dataflow and pipelining tasks" (citation). The overhead of passing messages from block to block is not negligible, so if your workload is too granular (lightweight) it is advised to chunkify it by using the Chunk LINQ operator (or other similar means) at the entry point of the pipeline.

Related