Return IAsyncEnumerable from grpc with timeout try-catch

Question

I have a gRPC client and I want to have a method that simplifies its use. The method should return IAsyncEnumerable of items being streamed from the gRPC server. I have a specified timeout for the streaming not to exceed. If the timeout occurs, I want to just walk away with all the items I managed to fetch so far.

Here's what I tried to do:

    public async IAsyncEnumerable<Item> Search(
        SearchParameters parameters, 
        CancellationToken cancellationToken, 
        IDictionary<string, string> headers = null)
    {
        try
        {
            await _client.Search(
                    MapInput(parameters),
                    cancellationToken: cancellationToken,
                    deadline: DateTime.UtcNow.Add(_configuration.Timeout),
                    headers: MapHeaders(headers))
                .ResponseStream.ForEachAsync(item =>
                {
                    yield return MapSingleItem(item); // compilation error
                });
        }
        catch (RpcException ex) when (ex.StatusCode == StatusCode.DeadlineExceeded)
        {
            _logger.LogWarning("Steam finished due to timeout, a limited number of items has been returned");
        }
    }

Logically, that should work. However, the yield keyword is not supported within lambdas, so it does not compile. Is there any other way I could write it?

Answer 1

You need an intermediate buffer to hold the items, because the consumer of the IAsyncEnumerable<Item> can enumerate it at its own pace. An excellent asynchronous buffer for this purpose is the Channel<T> class.

Another thing that you might want to consider is what happens if the consumer abandons the enumeration of the IAsyncEnumerable<Item> prematurely, either deliberately by breaking or returning, or unwillingly because it suffered an exception. You need to watch for this occurrence, and the best way to do it is to cancel a linked CancellationTokenSource in the finally block of your iterator.

Putting everything together:

public async IAsyncEnumerable<Item> Search(
    SearchParameters parameters, 
    [EnumeratorCancellation] CancellationToken cancellationToken = default,
    IDictionary<string, string> headers = null)
{
    Channel<Item> channel = Channel.CreateUnbounded<Item>();
    using CancellationTokenSource linkedCTS = CancellationTokenSource
        .CreateLinkedTokenSource(cancellationToken);

    Task producer = Task.Run(async () =>
    {
        try
        {
            await _client.Search(
                    MapInput(parameters),
                    cancellationToken: linkedCTS.Token,
                    deadline: DateTime.UtcNow.Add(_configuration.Timeout),
                    headers: MapHeaders(headers))
                .ResponseStream.ForEachAsync(item =>
                {
                    channel.Writer.TryWrite(item);
                }).ConfigureAwait(false);
        }
        finally { channel.Writer.Complete(); }
    });

    try
    {
        await foreach (Item item in channel.Reader.ReadAllAsync()
            .ConfigureAwait(false))
        {
            yield return item;
        }
        await producer.ConfigureAwait(false); // Propagate possible error
    }
    finally
    {
        // Prevent fire-and-forget in case the enumeration is abandoned
        if (!producer.IsCompleted)
        {
            linkedCTS.Cancel();
            await Task.WhenAny(producer).ConfigureAwait(false);
        }
    }
}

Most likely the resulting IAsyncEnumerable<Item> will complete with an OperationCanceledException when the token is canceled. If you prefer your token to have stopping semantics, you should at first rename it to stoppingToken, and then handle accordingly a OperationCanceledException exception inside the producer task.

Answer 2

With Rx.net you can do this with the .Debounce operator and the .TakeUntil operator.

var input = Observable.Create<Item>((observer, cancellationToken) => 
    Task.Factory.StartNew(() =>
    {
        try
        {
           var items = _client.Search(
                MapInput(parameters),
                cancellationToken: cancellationToken,
                deadline: DateTime.UtcNow.Add(_configuration.Timeout),
                headers: MapHeaders(headers));
           foreach(var item in items)
                observer.OnNext(item);
        }
        catch(Exception ex)
        {
            observer.OnError(ex);
        }
     }, TaskCreationOptions.LongRunning)
);

var inputObservable = input
      .Publish()
      .RefCount();


var timeout = inputObs
     .Throttle(TimeSpan.FromSeconds(10));
var outputObs = inputObservable
    .TakeUntil(timeout);
  

return outputObs
     .ToAsyncEnumerable()
     .ToListAsync();

Edited: assuming _client.Search returns an IEnumerable.