ShutdownAwareTopicEventReceiver.cs

Go to the documentation of this file.

using System;
using System.Collections.Concurrent;
using System.Threading;
using System.Threading.Tasks;
 
using HotChocolate.Execution;
using HotChocolate.Subscriptions;
 
using Microsoft.Extensions.Hosting;
 
namespace Tgstation.Server.Host.GraphQL.Subscriptions
{
    sealed class ShutdownAwareTopicEventReceiver : ITopicEventReceiver, IAsyncDisposable
    {
        readonly IHostApplicationLifetime hostApplicationLifetime;
 
        readonly HotChocolate.Subscriptions.ITopicEventReceiver hotChocolateReceiver;
 
        readonly ConcurrentBag<CancellationTokenRegistration> registrations;
 
        readonly ConcurrentBag<Task> disposeTasks;
 
        public ShutdownAwareTopicEventReceiver(
            IHostApplicationLifetime hostApplicationLifetime,
            HotChocolate.Subscriptions.ITopicEventReceiver hotChocolateReceiver)
        {
            this.hostApplicationLifetime = hostApplicationLifetime ?? throw new ArgumentNullException(nameof(hostApplicationLifetime));
            this.hotChocolateReceiver = hotChocolateReceiver ?? throw new ArgumentNullException(nameof(hotChocolateReceiver));
 
            registrations = new ConcurrentBag<CancellationTokenRegistration>();
            disposeTasks = new ConcurrentBag<Task>();
        }
 
        public async ValueTask DisposeAsync()
        {
            foreach (var registration in registrations)
            {
                registration.Dispose();
            }
 
            await Task.WhenAll(disposeTasks);
        }
 
        public ValueTask<ISourceStream<TMessage>> SubscribeAsync<TMessage>(string topicName, CancellationToken cancellationToken)
            => WrapWithApplicationLifetimeCancellation(
                hotChocolateReceiver.SubscribeAsync<TMessage>(topicName, cancellationToken));
 
        public ValueTask<ISourceStream<TMessage>> SubscribeAsync<TMessage>(string topicName, int? bufferCapacity, TopicBufferFullMode? bufferFullMode, CancellationToken cancellationToken)
            => WrapWithApplicationLifetimeCancellation(
                hotChocolateReceiver.SubscribeAsync<TMessage>(topicName, bufferCapacity, bufferFullMode, cancellationToken));
 
        async ValueTask<ISourceStream<TMessage>> WrapWithApplicationLifetimeCancellation<TMessage>(ValueTask<ISourceStream<TMessage>> sourceStreamTask)
        {
            var sourceStream = await sourceStreamTask;
            registrations.Add(
                hostApplicationLifetime.ApplicationStopping.Register(
                    () => disposeTasks.Add(
                        sourceStream.DisposeAsync().AsTask())));
            return sourceStream;
        }
    }
}