DapResponseRouter.cs

using System;
using System.Collections.Concurrent;
using System.Diagnostics.CodeAnalysis;
using System.Threading;
using System.Threading.Tasks;
using MediatR;
using Newtonsoft.Json.Linq;
using OmniSharp.Extensions.DebugAdapter.Protocol.Requests;
using OmniSharp.Extensions.JsonRpc;
using OmniSharp.Extensions.JsonRpc.Client;

namespace OmniSharp.Extensions.DebugAdapter.Shared
{
    public class DapResponseRouter : IResponseRouter
    {
        internal readonly IOutputHandler OutputHandler;
        internal readonly ISerializer Serializer;

        internal readonly ConcurrentDictionary<object, (string method, TaskCompletionSource<JToken> pendingTask)> Requests =
            new ConcurrentDictionary<object, (string method, TaskCompletionSource<JToken> pendingTask)>();

        internal static readonly ConcurrentDictionary<Type, string> MethodCache =
            new ConcurrentDictionary<Type, string>();

        public DapResponseRouter(IOutputHandler outputHandler, ISerializer serializer)
        {
            OutputHandler = outputHandler;
            Serializer = serializer;
        }

        public void SendNotification(string method)
        {
            OutputHandler.Send(
                new OutgoingNotification
                {
                    Method = method
                }
            );
        }

        public void SendNotification<T>(string method, T @params)
        {
            OutputHandler.Send(
                new OutgoingNotification
                {
                    Method = method,
                    Params = @params
                }
            );
        }

        public void SendNotification(IRequest<Unit> @params)
        {
            SendNotification(GetMethodName(@params.GetType()), @params);
        }

        public Task<TResponse> SendRequest<TResponse>(IRequest<TResponse> @params, CancellationToken cancellationToken)
        {
            return SendRequest(GetMethodName(@params.GetType()), @params).Returning<TResponse>(cancellationToken);
        }

        public IResponseRouterReturns SendRequest(string method)
        {
            return new ResponseRouterReturnsImpl(this, method, new object());
        }

        public IResponseRouterReturns SendRequest<T>(string method, T @params)
        {
            return new ResponseRouterReturnsImpl(this, method, @params);
        }

        public bool TryGetRequest(object id, [NotNullWhen(true)] out string? method, [NotNullWhen(true)] out TaskCompletionSource<JToken>? pendingTask)
        {
            var result = Requests.TryGetValue(id, out var source);
            method = source.method;
            pendingTask = source.pendingTask;
            return result;
        }

        private string GetMethodName(Type type)
        {
            if (!MethodCache.TryGetValue(type, out var methodName))
            {
                var attribute = MethodAttribute.From(type);
                if (attribute == null)
                {
                    throw new NotSupportedException($"Unable to infer method name for type {type.FullName}");
                }

                methodName = attribute.Method;
                MethodCache.TryAdd(type, methodName);
            }

            return methodName;
        }

        private class ResponseRouterReturnsImpl : IResponseRouterReturns
        {
            private readonly DapResponseRouter _router;
            private readonly string _method;
            private readonly object? _params;

            public ResponseRouterReturnsImpl(DapResponseRouter router, string method, object? @params)
            {
                _router = router;
                _method = method;
                _params = @params;
            }

            public async Task<TResponse> Returning<TResponse>(CancellationToken cancellationToken)
            {
                var nextId = _router.Serializer.GetNextId();
                var tcs = new TaskCompletionSource<JToken>();
                _router.Requests.TryAdd(nextId, ( _method, tcs ));

                cancellationToken.ThrowIfCancellationRequested();

                _router.OutputHandler.Send(
                    new OutgoingRequest
                    {
                        Method = _method,
                        Params = _params,
                        Id = nextId
                    }
                );
                if (_method != RequestNames.Cancel)
                {
                    cancellationToken.Register(
                        () =>
                        {
                            if (tcs.Task.IsCompleted) return;
#pragma warning disable VSTHRD110
                            _router.SendRequest(RequestNames.Cancel, new { requestId = nextId }).Returning<CancelArguments>(CancellationToken.None);
#pragma warning restore VSTHRD110
                        }
                    );
                }

                try
                {
                    var result = await tcs.Task.ConfigureAwait(false);
                    if (typeof(TResponse) == typeof(Unit))
                    {
                        return (TResponse)(object)Unit.Value;
                    }

                    return result.ToObject<TResponse>(_router.Serializer.JsonSerializer);
                }
                finally
                {
                    _router.Requests.TryRemove(nextId, out _);
                }
            }

            public async Task ReturningVoid(CancellationToken cancellationToken)
            {
                await Returning<Unit>(cancellationToken).ConfigureAwait(false);
            }
        }
    }
}