gRPC keeping response streams open for subscriptions

Question

I've tried to define a gRPC service where client can subscribe to receive broadcasted messages and they can also send them.

syntax = "proto3";

package Messenger;

service MessengerService {
    rpc SubscribeForMessages(User) returns (stream Message) {}
    rpc SendMessage(Message) returns (Close) {}
}

message User {
    string displayName = 1;
}

message Message {
    User from = 1;
    string message = 2;
}

message Close {}

My idea was that when a client requests to subscribe to the messages, the response stream would be added to a collection of response streams, and when a message is sent, the message is sent through all the response streams.

However, when my server attempts to write to the response streams, I get an exception System.InvalidOperationException: 'Response stream has already been completed.'

Is there any way to tell the server to keep the streams open so that new messages can be sent through them? Or is this not something that gRPC was designed for and a different technology should be used?

The end goal service would be allows multiple types of subscriptions (could be to new messages, weather updates, etc...) through different clients written in different languages (C#, Java, etc...). The different languages part is mainly the reason I chose gRPC to try this, although I intend on writing the server in C#.

---

Implementation example

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Grpc.Core;
using Messenger;

namespace SimpleGrpcTestStream
{
    /*
    Dependencies
Install-Package Google.Protobuf
Install-Package Grpc
Install-Package Grpc.Tools
Install-Package System.Interactive.Async
Install-Package System.Linq.Async

    */
    internal static class Program
    {
        private static void Main()
        {
            var messengerServer = new MessengerServer();
            messengerServer.Start();

            var channel = Common.GetNewInsecureChannel();
            var client = new MessengerService.MessengerServiceClient(channel);
            var clientUser = Common.GetUser("Client");
            var otherUser = Common.GetUser("Other");

            var cancelClientSubscription = AddCancellableMessageSubscription(client, clientUser);
            var cancelOtherSubscription = AddCancellableMessageSubscription(client, otherUser);

            client.SendMessage(new Message { From = clientUser, Message_ = "Hello" });
            client.SendMessage(new Message { From = otherUser, Message_ = "World" });
            client.SendMessage(new Message { From = clientUser, Message_ = "Whoop" });

            cancelClientSubscription.Cancel();
            cancelOtherSubscription.Cancel();
            channel.ShutdownAsync().Wait();
            messengerServer.ShutDown().Wait();
        }

        private static CancellationTokenSource AddCancellableMessageSubscription(
            MessengerService.MessengerServiceClient client,
            User user)
        {
            var cancelMessageSubscription = new CancellationTokenSource();

            var messages = client.SubscribeForMessages(user);

            var messageSubscription = messages
                .ResponseStream
                .ToAsyncEnumerable()
                .Finally(() => messages.Dispose());

            messageSubscription.ForEachAsync(
                message => Console.WriteLine($"New Message: {message.Message_}"),
                cancelMessageSubscription.Token);

            return cancelMessageSubscription;
        }
    }

    public static class Common
    {
        private const int Port = 50051;

        private const string Host = "localhost";

        private static readonly string ChannelAddress = $"{Host}:{Port}";

        public static User GetUser(string name) => new User { DisplayName = name };

        public static readonly User ServerUser = GetUser("Server");

        public static readonly Close EmptyClose = new Close();

        public static Channel GetNewInsecureChannel() => new Channel(ChannelAddress, ChannelCredentials.Insecure);

        public static ServerPort GetNewInsecureServerPort() => new ServerPort(Host, Port, ServerCredentials.Insecure);
    }

    public sealed class MessengerServer : MessengerService.MessengerServiceBase
    {
        private readonly Server _server;

        public MessengerServer()
        {
            _server = new Server
            {
                Ports = { Common.GetNewInsecureServerPort() },
                Services = { MessengerService.BindService(this) },
            };
        }

        public void Start()
        {
            _server.Start();
        }

        public async Task ShutDown()
        {
            await _server.ShutdownAsync().ConfigureAwait(false);
        }

        private readonly ConcurrentDictionary<User, IServerStreamWriter<Message>> _messageSubscriptions = new ConcurrentDictionary<User, IServerStreamWriter<Message>>();

        public override async Task<Close> SendMessage(Message request, ServerCallContext context)
        {
            await Task.Run(() =>
            {
                foreach (var (_, messageStream) in _messageSubscriptions)
                {
                    messageStream.WriteAsync(request);
                }
            }).ConfigureAwait(false);

            return await Task.FromResult(Common.EmptyClose).ConfigureAwait(false);
        }

        public override async Task SubscribeForMessages(User request, IServerStreamWriter<Message> responseStream, ServerCallContext context)
        {
            await Task.Run(() =>
            {
                responseStream.WriteAsync(new Message
                {
                    From = Common.ServerUser,
                    Message_ = $"{request.DisplayName} is listening for messages!",
                });
                _messageSubscriptions.TryAdd(request, responseStream);
            }).ConfigureAwait(false);
        }
    }

    public static class AsyncStreamReaderExtensions
    {
        public static IAsyncEnumerable<T> ToAsyncEnumerable<T>(this IAsyncStreamReader<T> asyncStreamReader)
        {
            if (asyncStreamReader is null) { throw new ArgumentNullException(nameof(asyncStreamReader)); }

            return new ToAsyncEnumerableEnumerable<T>(asyncStreamReader);
        }

        private sealed class ToAsyncEnumerableEnumerable<T> : IAsyncEnumerable<T>
        {
            public IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = default)
                => new ToAsyncEnumerator<T>(_asyncStreamReader, cancellationToken);

            private readonly IAsyncStreamReader<T> _asyncStreamReader;

            public ToAsyncEnumerableEnumerable(IAsyncStreamReader<T> asyncStreamReader)
            {
                _asyncStreamReader = asyncStreamReader;
            }

            private sealed class ToAsyncEnumerator<TEnumerator> : IAsyncEnumerator<TEnumerator>
            {
                public TEnumerator Current => _asyncStreamReader.Current;

                public async ValueTask<bool> MoveNextAsync() => await _asyncStreamReader.MoveNext(_cancellationToken);

                public ValueTask DisposeAsync() => default;

                private readonly IAsyncStreamReader<TEnumerator> _asyncStreamReader;
                private readonly CancellationToken _cancellationToken;

                public ToAsyncEnumerator(IAsyncStreamReader<TEnumerator> asyncStreamReader, CancellationToken cancellationToken)
                {
                    _asyncStreamReader = asyncStreamReader;
                    _cancellationToken = cancellationToken;
                }
            }
        }
    }
}

Answer 1

The problem you're experiencing is due to the fact that MessengerServer.SubscribeForMessages returns immediately. Once that method returns, the stream is closed.

You'll need an implementation similar to this to keep the stream alive:

public class MessengerService : MessengerServiceBase
{
    private static readonly ConcurrentDictionary<User, IServerStreamWriter<Message>> MessageSubscriptions =
        new Dictionary<User, IServerStreamWriter<Message>>();

    public override async Task SubscribeForMessages(User request, IServerStreamWriter<ReferralAssignment> responseStream, ServerCallContext context)
    {
        if (!MessageSubscriptions.TryAdd(request))
        {
            // User is already subscribed
            return;
        }

        // Keep the stream open so we can continue writing new Messages as they are pushed
        while (!context.CancellationToken.IsCancellationRequested)
        {
            // Avoid pegging CPU
            await Task.Delay(100);
        }

        // Cancellation was requested, remove the stream from stream map
        MessageSubscriptions.TryRemove(request);
    }
}

As far as unsubscribing / cancellation goes, there are two possible approaches:

1. The client can hold onto a CancellationToken and call Cancel() when it wants to disconnect
2. The server can hold onto a CancellationToken which you would then store along with the IServerStreamWriter in the MessageSubscriptions dictionary via a Tuple or similar. Then, you could introduce an Unsubscribe method on the server which looks up the CancellationToken by User and calls Cancel on it server-side