src/LanguageServer/Microsoft.CodeAnalysis.LanguageServer.UnitTests/Utilities/BrokeredServiceProxy.cs - main - githubmirror/roslyn

src/LanguageServer/Microsoft.CodeAnalysis.LanguageServer.UnitTests/Utilities/BrokeredServiceProxy.cs

76 строк3 KB

Tomáš Matoušek

Move LanguageServer out of Features directory (#73779)

30 май 2024, 18:49

30 май 2024, 18:4912f8968

// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

using Microsoft.VisualStudio.Threading;
using Nerdbank.Streams;
using StreamJsonRpc;

namespace Microsoft.CodeAnalysis.LanguageServer.UnitTests;

/// <summary>
/// A wrapper which takes a service but actually sends calls to it through JsonRpc to ensure we can actually use the service across a wire.
/// </summary>
internal sealed class BrokeredServiceProxy<T> : System.IAsyncDisposable where T : class
{
    /// <summary>
    /// A task that cane awaited to assert the rest of the fields in this class being assigned and non-null.
    /// </summary>
    private readonly Task _createConnectionTask;

    private JsonRpc? _serverRpc;
    private JsonRpc? _clientRpc;
    private T? _clientFactoryProxy;

    public BrokeredServiceProxy(T service)
    {
        var (serverStream, clientStream) = FullDuplexStream.CreatePair();

        _createConnectionTask = Task.WhenAll(CreateServerAsync(), CreateClientAsync());
        return;

        async Task CreateServerAsync()
        {
            // Always yield to ensure caller can proceed.
            await TaskScheduler.Default.SwitchTo(alwaysYield: true);

            var serverMultiplexingStream = await MultiplexingStream.CreateAsync(serverStream);
            var serverChannel = await serverMultiplexingStream.AcceptChannelAsync("");

            var serverFormatter = new MessagePackFormatter() { MultiplexingStream = serverMultiplexingStream };
            _serverRpc = new JsonRpc(new LengthHeaderMessageHandler(serverChannel, serverFormatter));

            _serverRpc.AddLocalRpcTarget(service, options: null);
            _serverRpc.StartListening();
        }

        async Task CreateClientAsync()
        {
            // Always yield to ensure caller can proceed.
            await TaskScheduler.Default.SwitchTo(alwaysYield: true);

            var clientMultiplexingStream = await MultiplexingStream.CreateAsync(clientStream);
            var clientChannel = await clientMultiplexingStream.OfferChannelAsync("");

            var clientFormatter = new MessagePackFormatter() { MultiplexingStream = clientMultiplexingStream };
            _clientRpc = new JsonRpc(new LengthHeaderMessageHandler(clientChannel, clientFormatter));

            _clientFactoryProxy = _clientRpc.Attach<T>();
            _clientRpc.StartListening();
        }
    }

    public async ValueTask DisposeAsync()
    {
        await _createConnectionTask;

        _serverRpc!.Dispose();
        _clientRpc!.Dispose();
    }

    public async Task<T> GetServiceAsync()
    {
        await _createConnectionTask;
        return _clientFactoryProxy!;
    }
}