<|start_filename|>examples/Container/Frontend/Balancer/ConfigurableResolverFactory.cs<|end_filename|>
﻿#region Copyright notice and license

// Copyright 2019 The gRPC Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#endregion

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Grpc.Net.Client;
using Grpc.Net.Client.Balancer;
using Grpc.Net.Client.Configuration;

namespace Frontend.Balancer
{
    public class ConfigurableResolverFactory : ResolverFactory
    {
        private readonly ResolverFactory _innerResolverFactory;
        private readonly BalancerConfiguration _balancerConfiguration;

        public ConfigurableResolverFactory(ResolverFactory innerResolverFactory, BalancerConfiguration balancerConfiguration)
        {
            _innerResolverFactory = innerResolverFactory;
            _balancerConfiguration = balancerConfiguration;
        }

        public override string Name => _innerResolverFactory.Name;

        public override Resolver Create(ResolverOptions options)
        {
            return new ConfigurableResolver(_innerResolverFactory.Create(options), _balancerConfiguration);
        }

        private class ConfigurableResolver : Resolver
        {
            private readonly Resolver _innerResolver;
            private readonly BalancerConfiguration _balancerConfiguration;

            private ResolverResult? _lastResult;
            private Action<ResolverResult>? _listener;

            public ConfigurableResolver(Resolver innerResolver, BalancerConfiguration balancerConfiguration)
            {
                _innerResolver = innerResolver;
                _balancerConfiguration = balancerConfiguration;
                _balancerConfiguration.Updated += OnConfigurationUpdated;
            }

            private void OnConfigurationUpdated(object? sender, EventArgs e)
            {
                // Can't just call RefreshAsync and get new results because of rate limiting.
                if (_listener != null && _lastResult != null)
                {
                    RaiseResult(_listener, _lastResult);
                }
            }

            public override Task RefreshAsync(CancellationToken cancellationToken)
            {
                return _innerResolver.RefreshAsync(cancellationToken);
            }

            public override void Start(Action<ResolverResult> listener)
            {
                _listener = listener;
                _innerResolver.Start(result =>
                {
                    _lastResult = result;

                    RaiseResult(_listener, result);
                });
            }

            private void RaiseResult(Action<ResolverResult> listener, ResolverResult result)
            {
                if (result.Addresses != null)
                {
                    var policyName = _balancerConfiguration.LoadBalancerPolicyName switch
                    {
                        LoadBalancerName.PickFirst => "pick_first",
                        LoadBalancerName.RoundRobin => "round_robin",
                        _ => throw new InvalidOperationException("Unexpected load balancer.")
                    };

                    var serviceConfig = new ServiceConfig
                    {
                        LoadBalancingConfigs = { new LoadBalancingConfig(policyName) }
                    };

                    // DNS results change order between refreshes.
                    // Explicitly order by host to keep result order consistent.
                    var orderedAddresses = result.Addresses!.OrderBy(a => a.EndPoint.Host).ToList();
                    listener(ResolverResult.ForResult(orderedAddresses, serviceConfig));
                }
                else
                {
                    listener(ResolverResult.ForFailure(result.Status));
                }
            }
        }
    }
}


<|start_filename|>src/Grpc.AspNetCore.Server.ClientFactory/GrpcServerHttpClientBuilderExtensions.cs<|end_filename|>
﻿#region Copyright notice and license

// Copyright 2019 The gRPC Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#endregion

using Grpc.AspNetCore.ClientFactory;
using Grpc.Core;
using Grpc.Net.ClientFactory;
using Microsoft.Extensions.DependencyInjection.Extensions;
using Microsoft.Extensions.Options;

namespace Microsoft.Extensions.DependencyInjection
{
    /// <summary>
    /// Extension methods for configuring an <see cref="IHttpClientBuilder"/>.
    /// </summary>
    public static class GrpcServerHttpClientBuilderExtensions
    {
        /// <summary>
        /// Configures the server to propagate values from a call's <see cref="ServerCallContext"/>
        /// onto the gRPC client.
        /// </summary>
        /// <param name="builder">The <see cref="IHttpClientBuilder"/>.</param>
        /// <returns>An <see cref="IHttpClientBuilder"/> that can be used to configure the client.</returns>
        public static IHttpClientBuilder EnableCallContextPropagation(this IHttpClientBuilder builder)
        {
            if (builder == null)
            {
                throw new ArgumentNullException(nameof(builder));
            }

            ValidateGrpcClient(builder);

            builder.Services.TryAddSingleton<ContextPropagationInterceptor>();
            builder.Services.AddHttpContextAccessor();
            builder.Services.Configure<GrpcClientFactoryOptions>(builder.Name, options =>
            {
                options.InterceptorRegistrations.Add(new InterceptorRegistration(
                    InterceptorScope.Channel,
                    s => s.GetRequiredService<ContextPropagationInterceptor>()));
            });

            return builder;
        }

        /// <summary>
        /// Configures the server to propagate values from a call's <see cref="ServerCallContext"/>
        /// onto the gRPC client.
        /// </summary>
        /// <param name="builder">The <see cref="IHttpClientBuilder"/>.</param>
        /// <param name="configureOptions">An <see cref="Action{GrpcContextPropagationOptions}"/> to configure the provided <see cref="GrpcContextPropagationOptions"/>.</param>
        /// <returns>An <see cref="IHttpClientBuilder"/> that can be used to configure the client.</returns>
        public static IHttpClientBuilder EnableCallContextPropagation(this IHttpClientBuilder builder, Action<GrpcContextPropagationOptions> configureOptions)
        {
            if (builder == null)
            {
                throw new ArgumentNullException(nameof(builder));
            }

            builder.Services.Configure(configureOptions);
            return builder.EnableCallContextPropagation();
        }

        private static void ValidateGrpcClient(IHttpClientBuilder builder)
        {
            // Validate the builder is for a gRPC client
            foreach (var service in builder.Services)
            {
                if (service.ServiceType == typeof(IConfigureOptions<GrpcClientFactoryOptions>))
                {
                    // Builder is from AddGrpcClient if options have been configured with the same name
                    if (service.ImplementationInstance is ConfigureNamedOptions<GrpcClientFactoryOptions> namedOptions && string.Equals(builder.Name, namedOptions.Name, StringComparison.Ordinal))
                    {
                        return;
                    }
                }
            }

            throw new InvalidOperationException($"{nameof(EnableCallContextPropagation)} must be used with a gRPC client.");
        }
    }
}


<|start_filename|>src/Grpc.Net.Client/Balancer/Internal/SocketConnectivitySubchannelTransport.cs<|end_filename|>
﻿#region Copyright notice and license

// Copyright 2019 The gRPC Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#endregion

#if SUPPORT_LOAD_BALANCING
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Net;
using System.Net.Http;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;
using Grpc.Core;
using Grpc.Shared;
using Microsoft.Extensions.Logging;

namespace Grpc.Net.Client.Balancer.Internal
{
#if NET5_0_OR_GREATER
    /// <summary>
    /// Transport that makes it possible to monitor connectivity state while using HttpClient.
    /// 
    /// Features:
    /// 1. When a connection is requested the transport creates a Socket and connects to the server.
    ///    The socket is used with the first stream created by SocketsHttpHandler.ConnectCallback.
    ///    The transport keeps track of the socket or the streams in use to determine if the connection
    ///    is ready. Connectivity API features require knowing whether there is a connection available.
    ///    A limitation of the .NET support is only socket connectivity to the server is tracked.
    ///    This transport is unable to check whether TLS and HTTP is succcessfully negotiated.
    /// 2. Transport supports multiple addresses. When connecting it will iterate through the addresses,
    ///    attempting to connect to each one.
    /// </summary>
    internal class SocketConnectivitySubchannelTransport : ISubchannelTransport, IDisposable
    {
        internal record struct ActiveStream(BalancerAddress Address, Socket Socket, Stream? Stream);

        private readonly ILogger _logger;
        private readonly Subchannel _subchannel;
        private readonly TimeSpan _socketPingInterval;
        private readonly List<ActiveStream> _activeStreams;
        private readonly Timer _socketConnectedTimer;

        private int _lastEndPointIndex;
        private Socket? _initialSocket;
        private BalancerAddress? _initialSocketAddress;
        private bool _disposed;
        private BalancerAddress? _currentAddress;

        public SocketConnectivitySubchannelTransport(Subchannel subchannel, TimeSpan socketPingInterval, ILoggerFactory loggerFactory)
        {
            _logger = loggerFactory.CreateLogger<SocketConnectivitySubchannelTransport>();
            _subchannel = subchannel;
            _socketPingInterval = socketPingInterval;
            _activeStreams = new List<ActiveStream>();
            _socketConnectedTimer = new Timer(OnCheckSocketConnection, state: null, Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan);
        }

        public object Lock => _subchannel.Lock;
        public BalancerAddress? CurrentAddress => _currentAddress;
        public bool HasStream { get; }

        // For testing. Take a copy under lock for thread-safety.
        internal IReadOnlyList<ActiveStream> GetActiveStreams()
        {
            lock (Lock)
            {
                return _activeStreams.ToList();
            }
        }

        public void Disconnect()
        {
            lock (Lock)
            {
                _initialSocket?.Dispose();
                _initialSocket = null;
                _initialSocketAddress = null;
                _lastEndPointIndex = 0;
                _socketConnectedTimer.Change(TimeSpan.Zero, TimeSpan.Zero);
                _currentAddress = null;
            }
            _subchannel.UpdateConnectivityState(ConnectivityState.Idle, "Disconnected.");
        }

        public async ValueTask<bool> TryConnectAsync(CancellationToken cancellationToken)
        {
            Debug.Assert(CurrentAddress == null);

            // Addresses could change while connecting. Make a copy of the subchannel's addresses.
            var addresses = _subchannel.GetAddresses();

            // Loop through endpoints and attempt to connect.
            Exception? firstConnectionError = null;

            for (var i = 0; i < addresses.Count; i++)
            {
                var currentIndex = (i + _lastEndPointIndex) % addresses.Count;
                var currentAddress = addresses[currentIndex];

                Socket socket;

                socket = new Socket(SocketType.Stream, ProtocolType.Tcp) { NoDelay = true };
                _subchannel.UpdateConnectivityState(ConnectivityState.Connecting, "Connecting to socket.");

                try
                {
                    SocketConnectivitySubchannelTransportLog.ConnectingSocket(_logger, currentAddress);
                    await socket.ConnectAsync(currentAddress.EndPoint, cancellationToken).ConfigureAwait(false);
                    SocketConnectivitySubchannelTransportLog.ConnectedSocket(_logger, currentAddress);

                    lock (Lock)
                    {
                        _currentAddress = currentAddress;
                        _lastEndPointIndex = currentIndex;
                        _initialSocket = socket;
                        _initialSocketAddress = currentAddress;
                        _socketConnectedTimer.Change(_socketPingInterval, _socketPingInterval);
                    }

                    _subchannel.UpdateConnectivityState(ConnectivityState.Ready, "Successfully connected to socket.");
                    return true;
                }
                catch (Exception ex)
                {
                    SocketConnectivitySubchannelTransportLog.ErrorConnectingSocket(_logger, currentAddress, ex);

                    if (firstConnectionError == null)
                    {
                        firstConnectionError = ex;
                    }
                }
            }

            // All connections failed
            _subchannel.UpdateConnectivityState(
                ConnectivityState.TransientFailure,
                new Status(StatusCode.Unavailable, "Error connecting to subchannel.", firstConnectionError));
            lock (Lock)
            {
                if (!_disposed)
                {
                    _socketConnectedTimer.Change(TimeSpan.Zero, TimeSpan.Zero);
                }
            }
            return false;
        }

        private async void OnCheckSocketConnection(object? state)
        {
            try
            {
                var socket = _initialSocket;
                if (socket != null)
                {
                    CompatibilityHelpers.Assert(_initialSocketAddress != null);

                    var closeSocket = false;
                    Exception? sendException = null;
                    try
                    {
                        // Check the socket is still valid by doing a zero byte send.
                        SocketConnectivitySubchannelTransportLog.CheckingSocket(_logger, _initialSocketAddress);
                        await socket.SendAsync(Array.Empty<byte>(), SocketFlags.None).ConfigureAwait(false);

                        // Also poll socket to check if it can be read from.
                        closeSocket = IsSocketInBadState(socket);
                    }
                    catch (Exception ex)
                    {
                        closeSocket = true;
                        sendException = ex;
                        SocketConnectivitySubchannelTransportLog.ErrorCheckingSocket(_logger, _initialSocketAddress, ex);
                    }

                    if (closeSocket)
                    {
                        lock (Lock)
                        {
                            if (_initialSocket == socket)
                            {
                                _initialSocket.Dispose();
                                _initialSocket = null;
                                _initialSocketAddress = null;
                                _currentAddress = null;
                                _lastEndPointIndex = 0;
                            }
                        }
                        _subchannel.UpdateConnectivityState(ConnectivityState.Idle, new Status(StatusCode.Unavailable, "Lost connection to socket.", sendException));
                    }
                }
            }
            catch (Exception ex)
            {
                SocketConnectivitySubchannelTransportLog.ErrorSocketTimer(_logger, ex);
            }
        }

        public async ValueTask<Stream> GetStreamAsync(BalancerAddress address, CancellationToken cancellationToken)
        {
            SocketConnectivitySubchannelTransportLog.CreatingStream(_logger, address);

            Socket? socket = null;
            lock (Lock)
            {
                if (_initialSocket != null &&
                    _initialSocketAddress != null &&
                    Equals(_initialSocketAddress, address))
                {
                    socket = _initialSocket;
                    _initialSocket = null;
                    _initialSocketAddress = null;
                }
            }

            // Check to see if we've received anything on the connection; if we have, that's
            // either erroneous data (we shouldn't have received anything yet) or the connection
            // has been closed; either way, we can't use it.
            if (socket != null)
            {
                if (IsSocketInBadState(socket))
                {
                    socket.Dispose();
                    socket = null;
                }
            }

            if (socket == null)
            {
                socket = new Socket(SocketType.Stream, ProtocolType.Tcp) { NoDelay = true };
                await socket.ConnectAsync(address.EndPoint, cancellationToken).ConfigureAwait(false);
            }

            var networkStream = new NetworkStream(socket, ownsSocket: true);

            // This stream wrapper intercepts dispose.
            var stream = new StreamWrapper(networkStream, OnStreamDisposed);

            lock (Lock)
            {
                _activeStreams.Add(new ActiveStream(address, socket, stream));
            }

            return stream;
        }

        private static bool IsSocketInBadState(Socket socket)
        {
            try
            {
                // Will return true if closed or there is pending data for some reason.
                return socket.Poll(0, SelectMode.SelectRead);
            }
            catch (Exception e) when (e is SocketException || e is ObjectDisposedException)
            {
                return false;
            }
        }

        private void OnStreamDisposed(Stream streamWrapper)
        {
            var disconnect = false;
            lock (Lock)
            {
                for (var i = _activeStreams.Count - 1; i >= 0; i--)
                {
                    var t = _activeStreams[i];
                    if (t.Stream == streamWrapper)
                    {
                        _activeStreams.RemoveAt(i);
                        SocketConnectivitySubchannelTransportLog.DisposingStream(_logger, t.Address);

                        // If the last active streams is removed then there is no active connection.
                        disconnect = _activeStreams.Count == 0;

                        break;
                    }
                }
            }

            if (disconnect)
            {
                // What happens after disconnect depends if the load balancer requests a new connection.
                // For example:
                // - Pick first will go into an idle state.
                // - Round-robin will reconnect to get back to a ready state.
                Disconnect();
            }
        }

        public void Dispose()
        {
            lock (Lock)
            {
                _socketConnectedTimer.Dispose();
                _disposed = true;
            }
        }

        public void OnRequestComplete(CompletionContext context)
        {
        }
    }

    internal static class SocketConnectivitySubchannelTransportLog
    {
        private static readonly Action<ILogger, BalancerAddress, Exception?> _connectingSocket =
            LoggerMessage.Define<BalancerAddress>(LogLevel.Trace, new EventId(1, "ConnectingSocket"), "Connecting socket to {Address}.");

        private static readonly Action<ILogger, BalancerAddress, Exception?> _connectedSocket =
            LoggerMessage.Define<BalancerAddress>(LogLevel.Debug, new EventId(2, "ConnectedSocket"), "Connected to socket {Address}.");

        private static readonly Action<ILogger, BalancerAddress, Exception?> _errorConnectingSocket =
            LoggerMessage.Define<BalancerAddress>(LogLevel.Error, new EventId(3, "ErrorConnectingSocket"), "Error connecting to socket {Address}.");

        private static readonly Action<ILogger, BalancerAddress, Exception?> _checkingSocket =
            LoggerMessage.Define<BalancerAddress>(LogLevel.Trace, new EventId(4, "CheckingSocket"), "Checking socket {Address}.");

        private static readonly Action<ILogger, BalancerAddress, Exception?> _errorCheckingSocket =
            LoggerMessage.Define<BalancerAddress>(LogLevel.Error, new EventId(5, "ErrorCheckingSocket"), "Error checking socket {Address}.");

        private static readonly Action<ILogger, Exception?> _errorSocketTimer =
            LoggerMessage.Define(LogLevel.Error, new EventId(6, "ErrorSocketTimer"), "Unexpected error in check socket timer.");

        private static readonly Action<ILogger, BalancerAddress, Exception?> _creatingStream =
            LoggerMessage.Define<BalancerAddress>(LogLevel.Trace, new EventId(7, "CreatingStream"), "Creating stream for {Address}.");

        private static readonly Action<ILogger, BalancerAddress, Exception?> _disposingStream =
            LoggerMessage.Define<BalancerAddress>(LogLevel.Trace, new EventId(8, "DisposingStream"), "Disposing stream for {Address}.");

        public static void ConnectingSocket(ILogger logger, BalancerAddress address)
        {
            _connectingSocket(logger, address, null);
        }

        public static void ConnectedSocket(ILogger logger, BalancerAddress address)
        {
            _connectedSocket(logger, address, null);
        }

        public static void ErrorConnectingSocket(ILogger logger, BalancerAddress address, Exception ex)
        {
            _errorConnectingSocket(logger, address, ex);
        }

        public static void CheckingSocket(ILogger logger, BalancerAddress address)
        {
            _checkingSocket(logger, address, null);
        }

        public static void ErrorCheckingSocket(ILogger logger, BalancerAddress address, Exception ex)
        {
            _errorCheckingSocket(logger, address, ex);
        }

        public static void ErrorSocketTimer(ILogger logger, Exception ex)
        {
            _errorSocketTimer(logger, ex);
        }

        public static void CreatingStream(ILogger logger, BalancerAddress address)
        {
            _creatingStream(logger, address, null);
        }

        public static void DisposingStream(ILogger logger, BalancerAddress address)
        {
            _disposingStream(logger, address, null);
        }
    }
#endif
}
#endif


<|start_filename|>src/Grpc.Net.Client/Balancer/SubchannelsLoadBalancer.cs<|end_filename|>
﻿#region Copyright notice and license

// Copyright 2019 The gRPC Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#endregion

#if SUPPORT_LOAD_BALANCING
using System;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using Grpc.Core;
using Grpc.Net.Client.Balancer.Internal;
using Microsoft.Extensions.Logging;

namespace Grpc.Net.Client.Balancer
{
    /// <summary>
    /// An abstract <see cref="LoadBalancer"/> that manages creating <see cref="Subchannel"/> instances
    /// from addresses. It is designed to make it easy to implement a custom picking policy by overriding
    /// <see cref="CreatePicker(IReadOnlyList{Subchannel})"/> and returning a custom <see cref="SubchannelPicker"/>.
    /// <para>
    /// Note: Experimental API that can change or be removed without any prior notice.
    /// </para>
    /// </summary>
    public abstract class SubchannelsLoadBalancer : LoadBalancer
    {
        /// <summary>
        /// Gets the controller.
        /// </summary>
        protected IChannelControlHelper Controller { get; }

        /// <summary>
        /// Gets the connectivity state.
        /// </summary>
        protected ConnectivityState State { get; private set; }

        private readonly List<AddressSubchannel> _addressSubchannels;
        private readonly ILogger _logger;

        /// <summary>
        /// Initializes a new instance of the <see cref="SubchannelsLoadBalancer"/> class.
        /// </summary>
        /// <param name="controller">The controller.</param>
        /// <param name="loggerFactory">The logger factory.</param>
        protected SubchannelsLoadBalancer(IChannelControlHelper controller, ILoggerFactory loggerFactory)
        {
            _addressSubchannels = new List<AddressSubchannel>();
            Controller = controller;
            _logger = loggerFactory.CreateLogger(GetType());
        }

        private void ResolverError(Status status)
        {
            // If balancer doesn't have a ready subchannel then remove any current subchannels
            // and update channel state with resolver error.
            switch (State)
            {
                case ConnectivityState.Idle:
                case ConnectivityState.Connecting:
                case ConnectivityState.TransientFailure:
                    foreach (var addressSubchannel in _addressSubchannels)
                    {
                        RemoveSubchannel(addressSubchannel.Subchannel);
                    }
                    Controller.UpdateState(new BalancerState(ConnectivityState.TransientFailure, new ErrorPicker(status)));
                    break;
            }
        }

        private int? FindSubchannelByAddress(List<AddressSubchannel> addressSubchannels, BalancerAddress address)
        {
            for (var i = 0; i < addressSubchannels.Count; i++)
            {
                var s = addressSubchannels[i];
                if (Equals(s.Address, address))
                {
                    return i;
                }
            }

            return null;
        }

        private int? FindSubchannel(List<AddressSubchannel> addressSubchannels, Subchannel subchannel)
        {
            for (var i = 0; i < addressSubchannels.Count; i++)
            {
                var s = addressSubchannels[i];
                if (Equals(s.Subchannel, subchannel))
                {
                    return i;
                }
            }

            return null;
        }

        /// <inheritdoc />
        public override void UpdateChannelState(ChannelState state)
        {
            if (state.Status.StatusCode != StatusCode.OK)
            {
                ResolverError(state.Status);
                return;
            }
            if (state.Addresses == null || state.Addresses.Count == 0)
            {
                ResolverError(new Status(StatusCode.Unavailable, "Resolver returned no addresses."));
                return;
            }

            var allUpdatedSubchannels = new List<AddressSubchannel>();
            var newSubchannels = new List<Subchannel>();
            var currentSubchannels = _addressSubchannels.ToList();

            // The state's addresses is the new authoritative list of addresses.
            // However, we want to keep existing subchannels when possible.
            foreach (var address in state.Addresses)
            {
                // Check existing subchannels for a match.
                var i = FindSubchannelByAddress(currentSubchannels, address);

                AddressSubchannel newOrCurrentSubConnection;
                if (i != null)
                {
                    // There is a match so take current subchannel.
                    newOrCurrentSubConnection = currentSubchannels[i.GetValueOrDefault()];

                    // Remove from current collection because any subchannels
                    // remaining in this collection at the end will be disposed.
                    currentSubchannels.RemoveAt(i.GetValueOrDefault());
                }
                else
                {
                    // No match so create a new subchannel.
                    var c = Controller.CreateSubchannel(new SubchannelOptions(new[] { address }));
                    c.OnStateChanged(s => UpdateSubchannelState(c, s));

                    newSubchannels.Add(c);
                    newOrCurrentSubConnection = new AddressSubchannel(c, address);
                }

                allUpdatedSubchannels.Add(newOrCurrentSubConnection);
            }

            // Any sub-connections still in this collection are no longer returned by the resolver.
            // This can all be removed.
            var removedSubConnections = currentSubchannels;

            if (removedSubConnections.Count == 0 && newSubchannels.Count == 0)
            {
                SubchannelsLoadBalancerLog.ConnectionsUnchanged(_logger);
                return;
            }

            foreach (var removedSubConnection in removedSubConnections)
            {
                RemoveSubchannel(removedSubConnection.Subchannel);
            }

            _addressSubchannels.Clear();
            _addressSubchannels.AddRange(allUpdatedSubchannels);

            // Start new connections after collection on balancer has been updated.
            foreach (var c in newSubchannels)
            {
                c.RequestConnection();
            }

            UpdateBalancingState(state.Status);
        }

        private void UpdateBalancingState(Status status)
        {
            var readySubchannels = new List<Subchannel>();
            for (var i = 0; i < _addressSubchannels.Count; i++)
            {
                var addressSubchannel = _addressSubchannels[i];
                if (addressSubchannel.Subchannel.State == ConnectivityState.Ready)
                {
                    readySubchannels.Add(addressSubchannel.Subchannel);
                }
            }

            if (readySubchannels.Count == 0)
            {
                // No READY subchannels, determine aggregate state and error status
                var isConnecting = false;
                foreach (var subchannel in _addressSubchannels)
                {
                    var state = subchannel.Subchannel.State;

                    if (state == ConnectivityState.Connecting || state == ConnectivityState.Idle)
                    {
                        isConnecting = true;
                        break;
                    }
                }

                if (isConnecting)
                {
                    UpdateChannelState(ConnectivityState.Connecting, EmptyPicker.Instance);
                }
                else
                {
                    // Only care about status if it is non-OK.
                    // Pass it to the picker so that it is reported to the caller on pick.
                    var errorStatus = status.StatusCode != StatusCode.OK
                        ? status
                        : new Status(StatusCode.Internal, "Unknown error.");

                    UpdateChannelState(ConnectivityState.TransientFailure, new ErrorPicker(errorStatus));
                }
            }
            else
            {
                UpdateChannelState(ConnectivityState.Ready, CreatePicker(readySubchannels));
            }
        }

        private void UpdateChannelState(ConnectivityState state, SubchannelPicker subchannelPicker)
        {
            State = state;
            Controller.UpdateState(new BalancerState(state, subchannelPicker));
        }

        private void UpdateSubchannelState(Subchannel subchannel, SubchannelState state)
        {
            var index = FindSubchannel(_addressSubchannels, subchannel);
            if (index == null)
            {
                SubchannelsLoadBalancerLog.IgnoredSubchannelStateChange(_logger, subchannel.Id);
                return;
            }

            SubchannelsLoadBalancerLog.ProcessingSubchannelStateChanged(_logger, subchannel.Id, state.State, state.Status);

            UpdateBalancingState(state.Status);

            if (state.State == ConnectivityState.TransientFailure || state.State == ConnectivityState.Idle)
            {
                SubchannelsLoadBalancerLog.RefreshingResolverForSubchannel(_logger, subchannel.Id, state.State);
                Controller.RefreshResolver();
            }
            if (state.State == ConnectivityState.Idle)
            {
                SubchannelsLoadBalancerLog.RequestingConnectionForSubchannel(_logger, subchannel.Id, state.State);
                subchannel.RequestConnection();
            }
        }

        private void RemoveSubchannel(Subchannel subchannel)
        {
            subchannel.Dispose();
        }

        /// <inheritdoc />
        public override void RequestConnection()
        {
            foreach (var addressSubchannel in _addressSubchannels)
            {
                addressSubchannel.Subchannel.RequestConnection();
            }
        }

        /// <inheritdoc />
        protected override void Dispose(bool disposing)
        {
            base.Dispose(disposing);

            foreach (var addressSubchannel in _addressSubchannels)
            {
                RemoveSubchannel(addressSubchannel.Subchannel);
            }
            _addressSubchannels.Clear();
        }

        /// <summary>
        /// Creates a <see cref="SubchannelPicker"/> for the specified <see cref="Subchannel"/> instances.
        /// This method can be overriden to return new a <see cref="SubchannelPicker"/> implementation
        /// with custom load balancing logic.
        /// </summary>
        /// <param name="readySubchannels">A collection of ready subchannels.</param>
        /// <returns>A subchannel picker.</returns>
        protected abstract SubchannelPicker CreatePicker(IReadOnlyList<Subchannel> readySubchannels);

        private record AddressSubchannel(Subchannel Subchannel, BalancerAddress Address);
    }

    internal static class SubchannelsLoadBalancerLog
    {
        private static readonly Action<ILogger, int, ConnectivityState, string, Exception?> _processingSubchannelStateChanged =
            LoggerMessage.Define<int, ConnectivityState, string>(LogLevel.Trace, new EventId(1, "ProcessingSubchannelStateChanged"), "Processing subchannel id '{SubchannelId}' state changed to {State}. Detail: '{Detail}'.");

        private static readonly Action<ILogger, int, Exception?> _ignoredSubchannelStateChange =
            LoggerMessage.Define<int>(LogLevel.Trace, new EventId(2, "IgnoredSubchannelStateChange"), "Ignored state change because of unknown subchannel id '{SubchannelId}'.");

        private static readonly Action<ILogger, Exception?> _connectionsUnchanged =
            LoggerMessage.Define(LogLevel.Trace, new EventId(3, "ConnectionsUnchanged"), "Connections unchanged.");

        private static readonly Action<ILogger, int, ConnectivityState, Exception?> _refreshingResolverForSubchannel =
            LoggerMessage.Define<int, ConnectivityState>(LogLevel.Trace, new EventId(4, "RefreshingResolverForSubchannel"), "Refreshing resolver because subchannel id '{SubchannelId}' is in state {State}.");

        private static readonly Action<ILogger, int, ConnectivityState, Exception?> _requestingConnectionForSubchannel =
            LoggerMessage.Define<int, ConnectivityState>(LogLevel.Trace, new EventId(5, "RequestingConnectionForSubchannel"), "Requesting connection for subchannel id '{SubchannelId}' because it is in state {State}.");

        public static void ProcessingSubchannelStateChanged(ILogger logger, int subchannelId, ConnectivityState state, Status status)
        {
            _processingSubchannelStateChanged(logger, subchannelId, state, status.Detail, status.DebugException);
        }

        public static void IgnoredSubchannelStateChange(ILogger logger, int subchannelId)
        {
            _ignoredSubchannelStateChange(logger, subchannelId, null);
        }

        public static void ConnectionsUnchanged(ILogger logger)
        {
            _connectionsUnchanged(logger, null);
        }

        public static void RefreshingResolverForSubchannel(ILogger logger, int subchannelId, ConnectivityState state)
        {
            _refreshingResolverForSubchannel(logger, subchannelId, state, null);
        }

        public static void RequestingConnectionForSubchannel(ILogger logger, int subchannelId, ConnectivityState state)
        {
            _requestingConnectionForSubchannel(logger, subchannelId, state, null);
        }
    }
}
#endif


<|start_filename|>src/Grpc.Net.Client/Internal/DefaultChannelCredentialsConfigurator.cs<|end_filename|>
﻿#region Copyright notice and license

// Copyright 2019 The gRPC Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#endregion

using Grpc.Core;

namespace Grpc.Net.Client.Internal
{
    internal class DefaultChannelCredentialsConfigurator : ChannelCredentialsConfiguratorBase
    {
        public bool? IsSecure { get; private set; }
        public List<CallCredentials>? CallCredentials { get; private set; }

        public override void SetCompositeCredentials(object state, ChannelCredentials channelCredentials, CallCredentials callCredentials)
        {
            channelCredentials.InternalPopulateConfiguration(this, null);

            if (callCredentials != null)
            {
                if (CallCredentials == null)
                {
                    CallCredentials = new List<CallCredentials>();
                }

                CallCredentials.Add(callCredentials);
            }
        }

        public override void SetInsecureCredentials(object state) => IsSecure = false;

        public override void SetSslCredentials(object state, string rootCertificates, KeyCertificatePair keyCertificatePair, VerifyPeerCallback verifyPeerCallback)
        {
            if (!string.IsNullOrEmpty(rootCertificates) ||
                keyCertificatePair != null ||
                verifyPeerCallback != null)
            {
                throw new InvalidOperationException(
                    $"{nameof(SslCredentials)} with non-null arguments is not supported by {nameof(GrpcChannel)}. " +
                    $"{nameof(GrpcChannel)} uses HttpClient to make gRPC calls and HttpClient automatically loads root certificates from the operating system certificate store. " +
                    $"Client certificates should be configured on HttpClient. See https://aka.ms/aspnet/grpc/certauth for details.");
            }

            IsSecure = true;
        }
    }
}


<|start_filename|>test/FunctionalTests/Balancer/DnsResolverTests.cs<|end_filename|>
﻿#region Copyright notice and license

// Copyright 2019 The gRPC Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#endregion

#if SUPPORT_LOAD_BALANCING

using System;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;
using Grpc.Core;
using Grpc.Net.Client.Balancer;
using Grpc.Net.Client.Internal;
using Grpc.Tests.Shared;
using Microsoft.Extensions.Logging;
using NUnit.Framework;

namespace Grpc.AspNetCore.FunctionalTests.Balancer
{
    [TestFixture]
    public class DnsResolverTests : FunctionalTestBase
    {
        [Test]
        public async Task Refresh_HasStarted_HasResult()
        {
            // Arranged
            var tcs = new TaskCompletionSource<ResolverResult>(TaskCreationOptions.RunContinuationsAsynchronously);

            var dnsResolver = CreateDnsResolver(new Uri("dns:///localhost"));
            dnsResolver.Start(r =>
            {
                tcs.SetResult(r);
            });

            // Act
            dnsResolver.Refresh();

            // Assert
            var result = await tcs.Task.DefaultTimeout();
            Assert.Greater(result.Addresses!.Count, 0);
        }

        private DnsResolver CreateDnsResolver(Uri address, int? defaultPort = null, TimeSpan? refreshInterval = null)
        {
            return new DnsResolver(address, defaultPort ?? 80, LoggerFactory, refreshInterval ?? Timeout.InfiniteTimeSpan);
        }

        [Test]
        public async Task Start_IntervalSet_MultipleCallbacks()
        {
            // Arrange
            const int waitForCallCount = 10;
            var currentCallCount = 0;
            var tcs = new TaskCompletionSource<ResolverResult>(TaskCreationOptions.RunContinuationsAsynchronously);
            var testSystemClock = new TestSystemClock(DateTime.UtcNow);

            var dnsResolver = CreateDnsResolver(new Uri("dns:///localhost"), refreshInterval: TimeSpan.FromSeconds(0.05));
            dnsResolver.SystemClock = testSystemClock;

            // Act
            dnsResolver.Start(r =>
            {
                // Avoid rate limit.
                testSystemClock.UtcNow = testSystemClock.UtcNow + TimeSpan.FromSeconds(20);

                if (Interlocked.Increment(ref currentCallCount) >= waitForCallCount)
                {
                    tcs.SetResult(r);
                }
            });

            // Assert
            var result = await tcs.Task.DefaultTimeout();
            Assert.NotNull(result);
            Assert.Greater(result!.Addresses!.Count, 0);
        }

        [Test]
        public async Task RefreshAsync_NoAddressPath_Error()
        {
            // Arrange
            SetExpectedErrorsFilter(writeContext =>
            {
                if (writeContext.State.ToString() == "Error querying DNS hosts for ''." &&
                    writeContext.Exception!.Message == "Resolver address 'dns://localhost/' is not valid. Please use dns:/// for DNS provider.")
                {
                    return true;
                }

                return false;
            });

            var tcs = new TaskCompletionSource<ResolverResult>(TaskCreationOptions.RunContinuationsAsynchronously);
            var dnsResolver = CreateDnsResolver(address: new Uri("dns://localhost"));
            dnsResolver.Start(r =>
            {
                tcs.SetResult(r);
            });

            // Act
            dnsResolver.Refresh();

            // Assert
            var result = await tcs.Task.DefaultTimeout();

            Assert.AreEqual(StatusCode.Unavailable, result.Status.StatusCode);
            Assert.AreEqual("Error getting DNS hosts for address ''. InvalidOperationException: Resolver address 'dns://localhost/' is not valid. Please use dns:/// for DNS provider.", result!.Status.Detail);
            Assert.AreEqual("Resolver address 'dns://localhost/' is not valid. Please use dns:/// for DNS provider.", result!.Status.DebugException.Message);
        }

        [Test]
        public async Task RefreshAsync_MultipleCalls_HitRateLimit()
        {
            // Arrange
            var tcs = new TaskCompletionSource<ResolverResult>(TaskCreationOptions.RunContinuationsAsynchronously);
            var dnsResolver = CreateDnsResolver(new Uri("dns:///localhost"));
            dnsResolver.Start(r =>
            {
                tcs.SetResult(r);
            });

            // Act
            dnsResolver.Refresh();

            // Assert
            var result = await tcs.Task.DefaultTimeout();
            Assert.Greater(result.Addresses!.Count, 0);

            // Wait for the internal resolve task to be completed before triggering refresh again
            await dnsResolver._resolveTask.DefaultTimeout();
            tcs = new TaskCompletionSource<ResolverResult>(TaskCreationOptions.RunContinuationsAsynchronously);
            dnsResolver.Refresh();

            var completedTask = await Task.WhenAny(tcs.Task, Task.Delay(TimeSpan.FromSeconds(0.5))).DefaultTimeout();
            if (completedTask == tcs.Task)
            {
                Assert.Fail("Expected refresh to be delayed.");
            }
        }

        [Test]
        public async Task RefreshAsync_MultipleCallsThenCancellation_CallCanceled()
        {
            SetExpectedErrorsFilter(writeContext =>
            {
                if (writeContext.State.ToString() == "Error querying DNS hosts for 'localhost'." &&
                    writeContext.Exception!.Message == "A task was canceled.")
                {
                    return true;
                }

                return false;
            });

            // Arrange
            var tcs = new TaskCompletionSource<ResolverResult>(TaskCreationOptions.RunContinuationsAsynchronously);
            var dnsResolver = CreateDnsResolver(new Uri("dns:///localhost"));
            dnsResolver.Start(r =>
            {
                tcs.SetResult(r);
            });

            // Act
            dnsResolver.Refresh();

            // Assert
            var result = await tcs.Task.DefaultTimeout();
            Assert.Greater(result.Addresses!.Count, 0);

            // Wait for the internal resolve task to be completed before triggering refresh again
            await dnsResolver._resolveTask.DefaultTimeout();
            tcs = new TaskCompletionSource<ResolverResult>(TaskCreationOptions.RunContinuationsAsynchronously);
            dnsResolver.Refresh();

            dnsResolver.Dispose();

            result = await tcs.Task.DefaultTimeout();
            Assert.AreEqual(StatusCode.Unavailable, result.Status.StatusCode);
            Assert.AreEqual("Error getting DNS hosts for address 'localhost'. TaskCanceledException: A task was canceled.", result.Status.Detail);
            Assert.AreEqual("A task was canceled.", result.Status.DebugException.Message);
        }

        [Test]
        public async Task DNS_Port_Works()
        {
            const int defaultPort = 1234; // Will be ignored because port is specified in address
            const int specifiedPort = 8080;

            var tcs = new TaskCompletionSource<ResolverResult>(TaskCreationOptions.RunContinuationsAsynchronously);
            var dnsResolver = CreateDnsResolver(new Uri($"dns:///localhost:{specifiedPort}"), defaultPort);
            dnsResolver.Start(r =>
            {
                tcs.SetResult(r);
            });

            // Act
            dnsResolver.Refresh();

            // Assert
            Assert.False(HasLogException((ex) =>
            {
                return ex is SocketException;
            }));

            var result = await tcs.Task.DefaultTimeout();
            Assert.AreEqual(StatusCode.OK, result.Status.StatusCode);

            var addresses = result.Addresses!;
            Logger.LogInformation($"Got {addresses.Count} addresses.");

            // Note: addresses returned from localhost depend on operating system.
            // Validate whatever we get back has the specified port.
            for (var i = 0; i < addresses.Count; i++)
            {
                var address = addresses[i];

                Logger.LogInformation($"Address: {address}");
                Assert.AreEqual(specifiedPort, address.EndPoint.Port);
            }
        }

        [Test]
        public async Task DefaultPort_Localhost_AddressesHaveDefaultPort()
        {
            const int defaultPort = 8081;

            var tcs = new TaskCompletionSource<ResolverResult>(TaskCreationOptions.RunContinuationsAsynchronously);
            var dnsResolver = CreateDnsResolver(new Uri("dns:///localhost"), defaultPort);
            dnsResolver.Start(r =>
            {
                tcs.SetResult(r);
            });

            // Act
            dnsResolver.Refresh();

            // Assert
            var result = await tcs.Task.DefaultTimeout();
            Assert.AreEqual(StatusCode.OK, result.Status.StatusCode);

            var addresses = result.Addresses!;
            Logger.LogInformation($"Got {addresses.Count} addresses.");

            // Note: addresses returned from localhost depend on operating system.
            // Validate whatever we get back has the specified port.
            for (var i = 0; i < addresses.Count; i++)
            {
                var address = addresses[i];

                Logger.LogInformation($"Address: {address}");
                Assert.AreEqual(defaultPort, address.EndPoint.Port);
            }
        }

        private class TestSystemClock : ISystemClock
        {
            public TestSystemClock(DateTime utcNow)
            {
                UtcNow = utcNow;
            }

            public DateTime UtcNow { get; set; }
        }
    }
}

#endif

<|start_filename|>test/Grpc.Net.Client.Tests/Balancer/ResolverTests.cs<|end_filename|>
﻿#region Copyright notice and license

// Copyright 2019 The gRPC Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#endregion

#if SUPPORT_LOAD_BALANCING
using System;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Threading;
using System.Threading.Tasks;
using Grpc.Core;
using Grpc.Net.Client.Balancer;
using Grpc.Net.Client.Balancer.Internal;
using Grpc.Net.Client.Configuration;
using Grpc.Net.Client.Tests.Infrastructure;
using Grpc.Net.Client.Tests.Infrastructure.Balancer;
using Grpc.Tests.Shared;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.DependencyInjection.Extensions;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Logging.Testing;
using NUnit.Framework;

namespace Grpc.Net.Client.Tests.Balancer
{
    [TestFixture]
    public class ResolverTests
    {
        [Test]
        public async Task Resolver_ResolveNameFromServices_Success()
        {
            // Arrange
            var services = new ServiceCollection();

            var resolver = new TestResolver();
            resolver.UpdateAddresses(new List<BalancerAddress>
            {
                new BalancerAddress("localhost", 80)
            });

            services.AddSingleton<ResolverFactory>(new TestResolverFactory(resolver));
            services.AddSingleton<ISubchannelTransportFactory>(new TestSubchannelTransportFactory());

            var handler = new TestHttpMessageHandler((r, ct) => default!);
            var channelOptions = new GrpcChannelOptions
            {
                Credentials = ChannelCredentials.Insecure,
                ServiceProvider = services.BuildServiceProvider(),
                HttpHandler = handler
            };

            // Act
            var channel = GrpcChannel.ForAddress("test:///localhost", channelOptions);
            await channel.ConnectAsync();

            // Assert
            var subchannels = channel.ConnectionManager.GetSubchannels();
            Assert.AreEqual(1, subchannels.Count);
        }

        [Test]
        public async Task Resolver_WaitForRefreshAsync_Success()
        {
            // Arrange
            var services = new ServiceCollection();
            var tcs = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);

            var resolver = new TestResolver(() => tcs.Task);
            resolver.UpdateAddresses(new List<BalancerAddress>
            {
                new BalancerAddress("localhost", 80)
            });

            services.AddSingleton<ResolverFactory>(new TestResolverFactory(resolver));
            services.AddSingleton<ISubchannelTransportFactory>(new TestSubchannelTransportFactory());

            var handler = new TestHttpMessageHandler((r, ct) => default!);
            var channelOptions = new GrpcChannelOptions
            {
                Credentials = ChannelCredentials.Insecure,
                ServiceProvider = services.BuildServiceProvider(),
                HttpHandler = handler
            };

            // Act
            var channel = GrpcChannel.ForAddress("test:///localhost", channelOptions);
            var connectTask = channel.ConnectAsync();

            // Assert
            Assert.IsFalse(connectTask.IsCompleted);

            tcs.SetResult(null);

            await connectTask.DefaultTimeout();

            var subchannels = channel.ConnectionManager.GetSubchannels();
            Assert.AreEqual(1, subchannels.Count);
        }

        [Test]
        public async Task Resolver_NoServiceConfigInResult_LoadBalancerUnchanged()
        {
            await Resolver_ServiceConfigInResult(resolvedServiceConfig: null);
        }

        [Test]
        public async Task Resolver_EmptyServiceConfigInResult_LoadBalancerUnchanged()
        {
            await Resolver_ServiceConfigInResult(resolvedServiceConfig: new ServiceConfig());
        }

        [Test]
        public async Task Resolver_MatchingPolicyInResult_LoadBalancerUnchanged()
        {
            await Resolver_ServiceConfigInResult(resolvedServiceConfig: new ServiceConfig
            {
                LoadBalancingConfigs = { new PickFirstConfig() }
            });
        }

        [Test]
        public void ResolverResult_OkStatusWithNoResolver_Error()
        {
            Assert.Throws<ArgumentException>(() => ResolverResult.ForResult(new List<BalancerAddress> { new BalancerAddress("localhost", 80) }, null, Status.DefaultSuccess));
        }

        private async Task Resolver_ServiceConfigInResult(ServiceConfig? resolvedServiceConfig)
        {
            // Arrange
            var services = new ServiceCollection();
            var tcs = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);

            var resolver = new TestResolver(() => tcs.Task);
            var result = ResolverResult.ForResult(new List<BalancerAddress> { new BalancerAddress("localhost", 80) }, resolvedServiceConfig, serviceConfigStatus: null);
            resolver.UpdateResult(result);

            var createdCount = 0;
            var test = new TestLoadBalancerFactory(
                LoadBalancingConfig.PickFirstPolicyName,
                c =>
                {
                    createdCount++;
                    return new PickFirstBalancer(c, NullLoggerFactory.Instance);
                });

            services.AddSingleton<ResolverFactory>(new TestResolverFactory(resolver));
            services.AddSingleton<ISubchannelTransportFactory>(new TestSubchannelTransportFactory());
            services.TryAddEnumerable(ServiceDescriptor.Singleton<LoadBalancerFactory>(test));

            var handler = new TestHttpMessageHandler((r, ct) => default!);
            var channelOptions = new GrpcChannelOptions
            {
                Credentials = ChannelCredentials.Insecure,
                ServiceProvider = services.BuildServiceProvider(),
                HttpHandler = handler
            };

            // Act
            var channel = GrpcChannel.ForAddress("test:///localhost", channelOptions);
            var connectTask = channel.ConnectAsync();

            // Assert
            Assert.IsFalse(connectTask.IsCompleted);

            tcs.SetResult(null);

            await connectTask.DefaultTimeout();

            var subchannels = channel.ConnectionManager.GetSubchannels();
            Assert.AreEqual(1, subchannels.Count);

            Assert.AreEqual(1, createdCount);

            resolver.UpdateResult(result);

            Assert.AreEqual(1, createdCount);
        }

        [Test]
        public async Task Resolver_ServiceConfigInResult()
        {
            // Arrange
            var services = new ServiceCollection();
            var tcs = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);

            var resolver = new TestResolver(() => tcs.Task);
            var result = ResolverResult.ForResult(new List<BalancerAddress> { new BalancerAddress("localhost", 80) }, serviceConfig: null, serviceConfigStatus: null);
            resolver.UpdateResult(result);

            TestLoadBalancer? firstLoadBalancer = null;
            var firstLoadBalancerCreatedCount = 0;
            var firstLoadBalancerFactory = new TestLoadBalancerFactory(
                LoadBalancingConfig.PickFirstPolicyName,
                c =>
                {
                    firstLoadBalancerCreatedCount++;
                    firstLoadBalancer = new TestLoadBalancer(new PickFirstBalancer(c, NullLoggerFactory.Instance));
                    return firstLoadBalancer;
                });

            TestLoadBalancer? secondLoadBalancer = null;
            var secondLoadBalancerCreatedCount = 0;
            var secondLoadBalancerFactory = new TestLoadBalancerFactory(
                "custom",
                c =>
                {
                    secondLoadBalancerCreatedCount++;
                    secondLoadBalancer = new TestLoadBalancer(new PickFirstBalancer(c, NullLoggerFactory.Instance));
                    return secondLoadBalancer;
                });

            var syncPoint = new SyncPoint(runContinuationsAsynchronously: true);
            var currentConnectivityState = ConnectivityState.Ready;

            services.AddSingleton<ResolverFactory>(new TestResolverFactory(resolver));
            services.AddSingleton<ISubchannelTransportFactory>(new TestSubchannelTransportFactory(async (s, c) =>
            {
                await syncPoint.WaitToContinue();
                return currentConnectivityState;
            }));
            services.Add(ServiceDescriptor.Singleton<LoadBalancerFactory>(firstLoadBalancerFactory));
            services.Add(ServiceDescriptor.Singleton<LoadBalancerFactory>(secondLoadBalancerFactory));

            var handler = new TestHttpMessageHandler((r, ct) => default!);
            var channelOptions = new GrpcChannelOptions
            {
                Credentials = ChannelCredentials.Insecure,
                ServiceProvider = services.BuildServiceProvider(),
                HttpHandler = handler
            };

            // Act
            var channel = GrpcChannel.ForAddress("test:///localhost", channelOptions);
            var connectTask = channel.ConnectAsync();

            // Assert
            Assert.IsFalse(connectTask.IsCompleted);

            tcs.SetResult(null);

            // Ensure that channel has processed results
            await resolver.HasResolvedTask.DefaultTimeout();

            var subchannels = channel.ConnectionManager.GetSubchannels();
            Assert.AreEqual(1, subchannels.Count);
            Assert.AreEqual(ConnectivityState.Connecting, subchannels[0].State);

            Assert.AreEqual(1, firstLoadBalancerCreatedCount);

            syncPoint!.Continue();

            var pick = await channel.ConnectionManager.PickAsync(new PickContext(), true, CancellationToken.None).AsTask().DefaultTimeout();
            Assert.AreEqual(80, pick.Address.EndPoint.Port);

            // Create new SyncPoint so new load balancer is waiting to connect
            syncPoint = new SyncPoint(runContinuationsAsynchronously: false);

            result = ResolverResult.ForResult(
                new List<BalancerAddress> { new BalancerAddress("localhost", 81) },
                serviceConfig: new ServiceConfig
                {
                    LoadBalancingConfigs = { new LoadBalancingConfig("custom") }
                },
                serviceConfigStatus: null);
            resolver.UpdateResult(result);

            Assert.AreEqual(1, firstLoadBalancerCreatedCount);
            Assert.AreEqual(1, secondLoadBalancerCreatedCount);

            // Old address is still used because new load balancer is connecting
            pick = await channel.ConnectionManager.PickAsync(new PickContext(), true, CancellationToken.None).AsTask().DefaultTimeout();
            Assert.AreEqual(80, pick.Address.EndPoint.Port);

            Assert.IsFalse(firstLoadBalancer!.Disposed);

            syncPoint!.Continue();

            // New address is used
            pick = await channel.ConnectionManager.PickAsync(new PickContext(), true, CancellationToken.None).AsTask().DefaultTimeout();
            Assert.AreEqual(81, pick.Address.EndPoint.Port);

            Assert.IsTrue(firstLoadBalancer!.Disposed);

            await connectTask.DefaultTimeout();
        }

        [TestCase(true)]
        [TestCase(false)]
        public async Task ResolverOptions_ResolveServiceConfig_LoadBalancerChangedIfNotDisabled(bool disabled)
        {
            // Arrange
            var services = new ServiceCollection();

            var resolver = new TestResolver();
            resolver.UpdateAddresses(
                new List<BalancerAddress> { new BalancerAddress("localhost", 80) },
                new ServiceConfig
                {
                    LoadBalancingConfigs = { new RoundRobinConfig() }
                });

            ResolverOptions? resolverOptions = null;
            services.AddSingleton<ResolverFactory>(new TestResolverFactory(o =>
            {
                resolverOptions = o;
                return resolver;
            }));
            services.AddSingleton<ISubchannelTransportFactory>(new TestSubchannelTransportFactory());

            var handler = new TestHttpMessageHandler((r, ct) => default!);
            var channelOptions = new GrpcChannelOptions
            {
                Credentials = ChannelCredentials.Insecure,
                ServiceProvider = services.BuildServiceProvider(),
                HttpHandler = handler,
                DisableResolverServiceConfig = disabled
            };

            // Act
            var channel = GrpcChannel.ForAddress("test:///localhost", channelOptions);
            await channel.ConnectAsync();

            // Assert
            var subchannels = channel.ConnectionManager.GetSubchannels();
            Assert.AreEqual(1, subchannels.Count);

            Assert.AreEqual(disabled, resolverOptions!.DisableServiceConfig);

            if (disabled)
            {
                Assert.IsNotNull(GetInnerLoadBalancer<PickFirstBalancer>(channel));
            }
            else
            {
                Assert.IsNotNull(GetInnerLoadBalancer<RoundRobinBalancer>(channel));
            }
        }

        [Test]
        public async Task ResolveServiceConfig_UnknownPolicyName_LoadBalancerUnchanged()
        {
            // Arrange
            var services = new ServiceCollection();

            var resolver = new TestResolver();
            resolver.UpdateAddresses(
                new List<BalancerAddress> { new BalancerAddress("localhost", 80) },
                new ServiceConfig
                {
                    LoadBalancingConfigs = { new LoadBalancingConfig("unknown!") }
                });

            services.AddSingleton<ResolverFactory>(new TestResolverFactory(resolver));
            services.AddSingleton<ISubchannelTransportFactory>(new TestSubchannelTransportFactory());

            var handler = new TestHttpMessageHandler((r, ct) => default!);
            var channelOptions = new GrpcChannelOptions
            {
                Credentials = ChannelCredentials.Insecure,
                ServiceProvider = services.BuildServiceProvider(),
                HttpHandler = handler
            };

            // Act
            var channel = GrpcChannel.ForAddress("test:///localhost", channelOptions);
            await channel.ConnectAsync();

            // Assert
            Assert.IsNotNull(GetInnerLoadBalancer<PickFirstBalancer>(channel));
        }

        [Test]
        public async Task ResolveServiceConfig_ErrorOnFirstResolve_PickError()
        {
            // Arrange
            var services = new ServiceCollection();

            var resolver = new TestResolver();
            resolver.UpdateAddresses(
                new List<BalancerAddress> { new BalancerAddress("localhost", 80) },
                serviceConfig: null,
                serviceConfigStatus: new Status(StatusCode.Internal, "An error!"));

            services.AddSingleton<ResolverFactory>(new TestResolverFactory(resolver));
            services.AddSingleton<ISubchannelTransportFactory>(new TestSubchannelTransportFactory());

            var handler = new TestHttpMessageHandler((r, ct) => default!);
            var channelOptions = new GrpcChannelOptions
            {
                Credentials = ChannelCredentials.Insecure,
                ServiceProvider = services.BuildServiceProvider(),
                HttpHandler = handler
            };

            // Act
            var channel = GrpcChannel.ForAddress("test:///localhost", channelOptions);
            await channel.ConnectionManager.ConnectAsync(waitForReady: false, CancellationToken.None).DefaultTimeout();

            var ex = await ExceptionAssert.ThrowsAsync<RpcException>(async () =>
            {
                await channel.ConnectionManager.PickAsync(new PickContext(), waitForReady: false, CancellationToken.None);
            }).DefaultTimeout();

            // Assert
            Assert.AreEqual(StatusCode.Internal, ex.StatusCode);
            Assert.AreEqual("An error!", ex.Status.Detail);
        }

        [Test]
        public async Task ResolveServiceConfig_ErrorOnSecondResolve_PickSuccess()
        {
            // Arrange
            var services = new ServiceCollection();

            var resolver = new TestResolver();
            resolver.UpdateAddresses(
                new List<BalancerAddress> { new BalancerAddress("localhost", 80) },
                serviceConfig: null,
                serviceConfigStatus: null);

            services.AddSingleton<ResolverFactory>(new TestResolverFactory(resolver));
            services.AddSingleton<ISubchannelTransportFactory>(new TestSubchannelTransportFactory());

            var testSink = new TestSink();
            services.AddLogging(b =>
            {
                b.AddProvider(new TestLoggerProvider(testSink));
            });
            services.AddNUnitLogger();

            var handler = new TestHttpMessageHandler((r, ct) => default!);
            var channelOptions = new GrpcChannelOptions
            {
                Credentials = ChannelCredentials.Insecure,
                ServiceProvider = services.BuildServiceProvider(),
                HttpHandler = handler
            };

            // Act
            var channel = GrpcChannel.ForAddress("test:///localhost", channelOptions);
            await channel.ConnectionManager.ConnectAsync(waitForReady: false, CancellationToken.None).DefaultTimeout();

            resolver.UpdateAddresses(
                new List<BalancerAddress> { new BalancerAddress("localhost", 80) },
                serviceConfig: null,
                serviceConfigStatus: new Status(StatusCode.Internal, "An error!"));

            var result = await channel.ConnectionManager.PickAsync(new PickContext(), waitForReady: false, CancellationToken.None);

            // Assert
            Assert.AreEqual("localhost", result.Address.EndPoint.Host);
            Assert.AreEqual(80, result.Address.EndPoint.Port);

            var pickStartedCount = testSink.Writes.Count(w => w.EventId.Name == "ResolverServiceConfigFallback");
            Assert.AreEqual(1, pickStartedCount);
        }

        public static T? GetInnerLoadBalancer<T>(GrpcChannel channel) where T : LoadBalancer
        {
            var balancer = (ChildHandlerLoadBalancer)channel.ConnectionManager._balancer!;
            return (T?)balancer._current?.LoadBalancer;
        }
    }
}
#endif


<|start_filename|>src/Grpc.Net.Client/Balancer/Internal/ISubchannelTransport.cs<|end_filename|>
﻿#region Copyright notice and license

// Copyright 2019 The gRPC Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#endregion

#if SUPPORT_LOAD_BALANCING
using System;
using System.IO;
using System.Net;
using System.Threading;
using System.Threading.Tasks;

namespace Grpc.Net.Client.Balancer.Internal
{
    /// <summary>
    /// An abstraction for subchannels to create a transport and connect to the server.
    /// This abstraction allows the connection to be customized. Used in unit tests.
    /// Might be made public in the future to support using load balancing with non-socket transports.
    /// </summary>
    internal interface ISubchannelTransport : IDisposable
    {
        void OnRequestComplete(CompletionContext context);
        BalancerAddress? CurrentAddress { get; }

#if NET5_0_OR_GREATER
        ValueTask<Stream> GetStreamAsync(BalancerAddress address, CancellationToken cancellationToken);
#endif

#if !NETSTANDARD2_0
        ValueTask<bool>
#else
        Task<bool>
#endif
            TryConnectAsync(CancellationToken cancellationToken);

        void Disconnect();
    }

    internal interface ISubchannelTransportFactory
    {
        ISubchannelTransport Create(Subchannel subchannel);
    }
}
#endif


<|start_filename|>test/Grpc.AspNetCore.Server.Tests/HealthChecks/GrpcHealthChecksPublisherTests.cs<|end_filename|>
﻿#region Copyright notice and license

// Copyright 2019 The gRPC Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#endregion

using Grpc.AspNetCore.HealthChecks;
using Grpc.AspNetCore.Server.Tests.Infrastructure;
using Grpc.Core;
using Grpc.Health.V1;
using Grpc.HealthCheck;
using Grpc.Tests.Shared;
using Microsoft.Extensions.Diagnostics.HealthChecks;
using NUnit.Framework;

namespace Grpc.AspNetCore.Server.Tests.Reflection
{
    [TestFixture]
    public class GrpcHealthChecksPublisherTests
    {
        [Test]
        public async Task PublishAsync_Check_ChangingStatus()
        {
            // Arrange
            var healthService = new HealthServiceImpl();
            var publisher = new GrpcHealthChecksPublisher(healthService);

            HealthCheckResponse response;

            // Act 1
            var ex = await ExceptionAssert.ThrowsAsync<RpcException>(() => healthService.Check(new HealthCheckRequest { Service = "" }, context: null));

            // Assert 1
            Assert.AreEqual(StatusCode.NotFound, ex.StatusCode);

            // Act 2
            var report = CreateSimpleHealthReport(HealthStatus.Healthy);
            await publisher.PublishAsync(report, CancellationToken.None);

            response = await healthService.Check(new HealthCheckRequest { Service = "" }, context: null);

            // Assert 2
            Assert.AreEqual(HealthCheckResponse.Types.ServingStatus.Serving, response.Status);

            // Act 3
            report = CreateSimpleHealthReport(HealthStatus.Unhealthy);
            await publisher.PublishAsync(report, CancellationToken.None);

            response = await healthService.Check(new HealthCheckRequest { Service = "" }, context: null);

            // Act 3
            Assert.AreEqual(HealthCheckResponse.Types.ServingStatus.NotServing, response.Status);
        }

        private static HealthReport CreateSimpleHealthReport(HealthStatus healthStatus)
        {
            return new HealthReport(
                new Dictionary<string, HealthReportEntry>
                {
                    [""] = new HealthReportEntry(healthStatus, "Description!", TimeSpan.Zero, exception: null, data: null)
                },
                TimeSpan.Zero);
        }

        [Test]
        public async Task PublishAsync_Check_MapStatuses()
        {
            // Arrange
            var healthService = new HealthServiceImpl();
            var publisher = new GrpcHealthChecksPublisher(healthService);

            // Act
            HealthCheckResponse response;

            var report = new HealthReport(
                new Dictionary<string, HealthReportEntry>
                {
                    [nameof(HealthStatus.Healthy)] = new HealthReportEntry(HealthStatus.Healthy, "Description!", TimeSpan.Zero, exception: null, data: null),
                    [nameof(HealthStatus.Degraded)] = new HealthReportEntry(HealthStatus.Degraded, "Description!", TimeSpan.Zero, exception: null, data: null),
                    [nameof(HealthStatus.Unhealthy)] = new HealthReportEntry(HealthStatus.Unhealthy, "Description!", TimeSpan.Zero, exception: null, data: null)
                },
                TimeSpan.Zero);
            await publisher.PublishAsync(report, CancellationToken.None);

            // Assert
            response = await healthService.Check(new HealthCheckRequest { Service = nameof(HealthStatus.Healthy) }, context: null);
            Assert.AreEqual(HealthCheckResponse.Types.ServingStatus.Serving, response.Status);

            response = await healthService.Check(new HealthCheckRequest { Service = nameof(HealthStatus.Degraded) }, context: null);
            Assert.AreEqual(HealthCheckResponse.Types.ServingStatus.Serving, response.Status);

            response = await healthService.Check(new HealthCheckRequest { Service = nameof(HealthStatus.Unhealthy) }, context: null);
            Assert.AreEqual(HealthCheckResponse.Types.ServingStatus.NotServing, response.Status);
        }

        [Test]
        public async Task PublishAsync_Watch_ChangingStatus()
        {
            // Arrange
            var healthService = new HealthServiceImpl();
            var publisher = new GrpcHealthChecksPublisher(healthService);
            var responseStream = new TestServerStreamWriter<HealthCheckResponse>();
            var cts = new CancellationTokenSource();
            var serverCallContext = new TestServerCallContext(DateTime.MinValue, cts.Token);

            // Act 1
            var call = healthService.Watch(new HealthCheckRequest { Service = "" }, responseStream, serverCallContext);

            // Assert 1
            await TestHelpers.AssertIsTrueRetryAsync(() => responseStream.Responses.Count == 1, "Unexpected response count.");
            Assert.AreEqual(HealthCheckResponse.Types.ServingStatus.ServiceUnknown, responseStream.Responses.Last().Status);

            // Act 2
            await publisher.PublishAsync(CreateSimpleHealthReport(HealthStatus.Healthy), CancellationToken.None);

            // Assert 2
            await TestHelpers.AssertIsTrueRetryAsync(() => responseStream.Responses.Count == 2, "Unexpected response count.");
            Assert.AreEqual(HealthCheckResponse.Types.ServingStatus.Serving, responseStream.Responses.Last().Status);

            // Act 3
            await publisher.PublishAsync(CreateSimpleHealthReport(HealthStatus.Degraded), CancellationToken.None);

            // Act 3
            await TestHelpers.AssertIsTrueRetryAsync(() => responseStream.Responses.Count == 2, "Unexpected response count.");
            Assert.AreEqual(HealthCheckResponse.Types.ServingStatus.Serving, responseStream.Responses.Last().Status);

            // Act 4
            await publisher.PublishAsync(CreateSimpleHealthReport(HealthStatus.Unhealthy), CancellationToken.None);

            // Act 4
            await TestHelpers.AssertIsTrueRetryAsync(() => responseStream.Responses.Count == 3, "Unexpected response count.");
            Assert.AreEqual(HealthCheckResponse.Types.ServingStatus.NotServing, responseStream.Responses.Last().Status);

            // End call
            cts.Cancel();
            await call.DefaultTimeout();

            // Act 4
            await publisher.PublishAsync(CreateSimpleHealthReport(HealthStatus.Healthy), CancellationToken.None);

            // Act 4
            await TestHelpers.AssertIsTrueRetryAsync(() => responseStream.Responses.Count == 3, "Unexpected response count.");
            Assert.AreEqual(HealthCheckResponse.Types.ServingStatus.NotServing, responseStream.Responses.Last().Status);
        }
    }
}


<|start_filename|>test/FunctionalTests/Server/ServerStreamingMethodTests.cs<|end_filename|>
#region Copyright notice and license

// Copyright 2019 The gRPC Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#endregion

using System.IO.Pipelines;
using Greet;
using Grpc.AspNetCore.FunctionalTests.Infrastructure;
using Grpc.Core;
using Grpc.Tests.Shared;
using Microsoft.Extensions.Logging;
using NUnit.Framework;

namespace Grpc.AspNetCore.FunctionalTests.Server
{
    [TestFixture]
    public class ServerStreamingMethodTests : FunctionalTestBase
    {
        [Test]
        public async Task NoBuffering_SuccessResponsesStreamed()
        {
            using var httpEventSource = new HttpEventSourceListener(LoggerFactory);

            var methodWrapper = new MethodWrapper
            {
                Logger = Logger,
                SyncPoint = new SyncPoint(runContinuationsAsynchronously: true)
            };

            async Task SayHellos(HelloRequest request, IServerStreamWriter<HelloReply> responseStream, ServerCallContext context)
            {
                // Explicitly send the response headers before any streamed content
                Metadata responseHeaders = new Metadata();
                responseHeaders.Add("test-response-header", "value");
                await context.WriteResponseHeadersAsync(responseHeaders);

                await methodWrapper.SayHellosAsync(request, responseStream);
            }

            // Arrange
            var method = Fixture.DynamicGrpc.AddServerStreamingMethod<HelloRequest, HelloReply>(SayHellos);

            var requestMessage = new HelloRequest
            {
                Name = "World"
            };

            var requestStream = new MemoryStream();
            MessageHelpers.WriteMessage(requestStream, requestMessage);

            var httpRequest = GrpcHttpHelper.Create(method.FullName);
            httpRequest.Content = new GrpcStreamContent(requestStream);

            // Act
            var response = await Fixture.Client.SendAsync(httpRequest, HttpCompletionOption.ResponseHeadersRead).DefaultTimeout();

            // Assert
            response.AssertIsSuccessfulGrpcRequest();

            var responseStream = await response.Content.ReadAsStreamAsync().DefaultTimeout();
            var pipeReader = PipeReader.Create(responseStream);

            for (var i = 0; i < 3; i++)
            {
                var greetingTask = MessageHelpers.AssertReadStreamMessageAsync<HelloReply>(pipeReader);

                Assert.IsFalse(greetingTask.IsCompleted);

                await methodWrapper.SyncPoint.WaitForSyncPoint().DefaultTimeout();

                var currentSyncPoint = methodWrapper.SyncPoint;
                methodWrapper.SyncPoint = new SyncPoint(runContinuationsAsynchronously: true);
                currentSyncPoint.Continue();

                var greeting = (await greetingTask.DefaultTimeout())!;

                Assert.AreEqual($"How are you World? {i}", greeting.Message);
            }

            var goodbyeTask = MessageHelpers.AssertReadStreamMessageAsync<HelloReply>(pipeReader);
            Assert.False(goodbyeTask.IsCompleted);

            await methodWrapper.SyncPoint.WaitForSyncPoint().DefaultTimeout();
            methodWrapper.SyncPoint.Continue();

            Assert.AreEqual("Goodbye World!", (await goodbyeTask.DefaultTimeout())!.Message);

            var finishedTask = MessageHelpers.AssertReadStreamMessageAsync<HelloReply>(pipeReader);
            Assert.IsNull(await finishedTask.DefaultTimeout());
        }

        [Test]
        public async Task WriteResponseHeadersAsyncCore_FlushesHeadersToClient()
        {
            var methodWrapper = new MethodWrapper
            {
                Logger = Logger,
                SyncPoint = new SyncPoint(runContinuationsAsynchronously: true)
            };

            async Task SayHellosSendHeadersFirst(HelloRequest request, IServerStreamWriter<HelloReply> responseStream, ServerCallContext context)
            {
                await context.WriteResponseHeadersAsync(null);

                await methodWrapper.SayHellosAsync(request, responseStream);
            }

            // Arrange
            var method = Fixture.DynamicGrpc.AddServerStreamingMethod<HelloRequest, HelloReply>(SayHellosSendHeadersFirst);

            var requestMessage = new HelloRequest
            {
                Name = "World"
            };

            var requestStream = new MemoryStream();
            MessageHelpers.WriteMessage(requestStream, requestMessage);

            var httpRequest = GrpcHttpHelper.Create(method.FullName);
            httpRequest.Content = new GrpcStreamContent(requestStream);

            // Act
            var response = await Fixture.Client.SendAsync(httpRequest, HttpCompletionOption.ResponseHeadersRead).DefaultTimeout();

            // Assert
            response.AssertIsSuccessfulGrpcRequest();

            Logger.LogInformation("Headers received. Starting to read stream");

            var responseStream = await response.Content.ReadAsStreamAsync().DefaultTimeout();
            var pipeReader = PipeReader.Create(responseStream);

            for (var i = 0; i < 3; i++)
            {
                Logger.LogInformation($"Reading message {i}.");

                var greetingTask = MessageHelpers.AssertReadStreamMessageAsync<HelloReply>(pipeReader);

                // The headers are already sent
                // All responses are streamed
                Logger.LogInformation($"Message task completed: {greetingTask.IsCompleted}");
                Assert.IsFalse(greetingTask.IsCompleted);

                await methodWrapper.SyncPoint.WaitForSyncPoint().DefaultTimeout();

                var currentSyncPoint = methodWrapper.SyncPoint;
                methodWrapper.SyncPoint = new SyncPoint(runContinuationsAsynchronously: true);
                currentSyncPoint.Continue();

                var greeting = (await greetingTask.DefaultTimeout())!;

                Logger.LogInformation($"Received message {i}: {greeting.Message}");

                Assert.AreEqual($"How are you World? {i}", greeting.Message);
            }

            var goodbyeTask = MessageHelpers.AssertReadStreamMessageAsync<HelloReply>(pipeReader);
            Assert.False(goodbyeTask.IsCompleted);

            await methodWrapper.SyncPoint.WaitForSyncPoint().DefaultTimeout();
            methodWrapper.SyncPoint.Continue();

            Assert.AreEqual("Goodbye World!", (await goodbyeTask.DefaultTimeout())!.Message);

            var finishedTask = MessageHelpers.AssertReadStreamMessageAsync<HelloReply>(pipeReader);
            Assert.IsNull(await finishedTask.DefaultTimeout());
        }

        [Test]
        public async Task Buffering_SuccessResponsesStreamed()
        {
            var methodWrapper = new MethodWrapper
            {
                Logger = Logger,
                SyncPoint = new SyncPoint(runContinuationsAsynchronously: true)
            };

            Task SayHellosBufferHint(HelloRequest request, IServerStreamWriter<HelloReply> responseStream, ServerCallContext context)
            {
                context.WriteOptions = new WriteOptions(WriteFlags.BufferHint);

                return methodWrapper.SayHellosAsync(request, responseStream);
            }

            // Arrange
            var method = Fixture.DynamicGrpc.AddServerStreamingMethod<HelloRequest, HelloReply>(SayHellosBufferHint);

            var requestMessage = new HelloRequest
            {
                Name = "World"
            };

            var requestStream = new MemoryStream();
            MessageHelpers.WriteMessage(requestStream, requestMessage);

            var httpRequest = GrpcHttpHelper.Create(method.FullName);
            httpRequest.Content = new GrpcStreamContent(requestStream);

            // Act
            var responseTask = Fixture.Client.SendAsync(httpRequest, HttpCompletionOption.ResponseHeadersRead).DefaultTimeout();

            // Assert
            Assert.IsFalse(responseTask.IsCompleted, "Server should wait for first message from client");

            await methodWrapper.SyncPoint.WaitForSyncPoint().DefaultTimeout();
            methodWrapper.SyncPoint.Continue();

            var response = await responseTask.DefaultTimeout();
            response.AssertIsSuccessfulGrpcRequest();

            var responseStream = await response.Content.ReadAsStreamAsync().DefaultTimeout();
            var pipeReader = PipeReader.Create(responseStream);

            for (var i = 0; i < 3; i++)
            {
                var greeting = (await MessageHelpers.AssertReadStreamMessageAsync<HelloReply>(pipeReader).DefaultTimeout())!;

                Assert.AreEqual($"How are you World? {i}", greeting.Message);
            }

            var goodbye = await MessageHelpers.AssertReadStreamMessageAsync<HelloReply>(pipeReader).DefaultTimeout();
            Assert.AreEqual("Goodbye World!", goodbye!.Message);

            Assert.IsNull(await MessageHelpers.AssertReadStreamMessageAsync<HelloReply>(pipeReader).DefaultTimeout());

            response.AssertTrailerStatus();
        }

        public class MethodWrapper
        {
            public SyncPoint SyncPoint { get; set; } = default!;
            public ILogger Logger { get; set; } = default!;

            public async Task SayHellosAsync(HelloRequest request, IServerStreamWriter<HelloReply> responseStream)
            {
                for (var i = 0; i < 3; i++)
                {
                    await SyncPoint.WaitToContinue();

                    var message = $"How are you {request.Name}? {i}";

                    Logger.LogInformation("Sending message");
                    await responseStream.WriteAsync(new HelloReply { Message = message });
                }

                await SyncPoint.WaitToContinue();

                Logger.LogInformation("Sending message");
                await responseStream.WriteAsync(new HelloReply { Message = $"Goodbye {request.Name}!" });
            }
        }
    }
}
