AsyncDiagnosticChannelWriter.cs

Go to the documentation of this file.

using System;
using System.Buffers;
using System.Threading;
using System.Threading.Channels;
using System.Threading.Tasks;
 
namespace IKVM.Tools.Core.Diagnostics
{
 
    class AsyncDiagnosticChannelWriter : IDiagnosticChannelWriter, IDisposable
    {
 
        static readonly UnboundedChannelOptions unboundedChannelOptions = new UnboundedChannelOptions()
        {
            SingleWriter = false,
            SingleReader = true,
            AllowSynchronousContinuations = true
        };
 
        readonly IBufferWriter<byte> _writer;
        readonly Channel<(IMemoryOwner<byte> Owner, int Length)> _channel = Channel.CreateUnbounded<(IMemoryOwner<byte> Owner, int Length)>(unboundedChannelOptions);
 
        Task? _task;
        CancellationTokenSource? _stop;
 
        public AsyncDiagnosticChannelWriter(IBufferWriter<byte> writer)
        {
            _writer = writer ?? throw new ArgumentNullException(nameof(writer));
        }
 
        public void Write(IMemoryOwner<byte> owner, int length)
        {
            if (owner is null)
                throw new ArgumentNullException(nameof(owner));
 
            if (_stop == null || _task == null)
            {
                lock (this)
                {
                    if (_stop == null || _task == null)
                    {
                        _stop = new CancellationTokenSource();
                        _task = DequeueLoop(_stop.Token);
                    }
                }
            }
 
            _channel.Writer.TryWrite((owner, length));
        }
 
        async Task DequeueLoop(CancellationToken cancellationToken)
        {
            while (cancellationToken.IsCancellationRequested == false)
            {
                try
                {
                    while (await _channel.Reader.WaitToReadAsync(cancellationToken))
                        while (_channel.Reader.TryRead(out var item))
                            WriteData(item.Owner, item.Length);
                }
                catch (OperationCanceledException)
                {
                    // ignore
                }
            }
        }
 
        void WriteData(IMemoryOwner<byte> owner, int length)
        {
            // allocate memory and copy data
            var buffer = _writer.GetMemory(length);
            owner.Memory.Slice(0, length).CopyTo(buffer);
            _writer.Advance(length);
 
            // we are finished with the owner
            owner.Dispose();
        }
 
        public void Dispose()
        {
            lock (this)
            {
                // stop dequeue task
                _stop?.Cancel();
                _task?.GetAwaiter().GetResult();
                _stop = null;
                _task = null;
 
                // empty the channel to pick up any straggling items not caught by the dequeue thread
                while (_channel.Reader.TryRead(out var item))
                    item.Owner.Dispose();
 
                // complete the channel
                _channel.Writer.TryComplete();
            }
        }
 
    }
 
}