AsyncClientStreamingCall <TRequest, TResponse>(Method <TRequest, TResponse> method, string host, CallOptions options)
{
// No channel affinity feature for client streaming call.
ChannelRef channelRef = GetChannelRef();
var callDetails = new CallInvocationDetails <TRequest, TResponse>(channelRef.Channel, method, host, options);
var originalCall = Calls.AsyncClientStreamingCall(callDetails);
// Decrease the active streams count once async response finishes.
var gcpResponseAsync = DecrementCountAndPropagateResult(originalCall.ResponseAsync);
// Create a wrapper of the original AsyncClientStreamingCall.
return(new AsyncClientStreamingCall <TRequest, TResponse>(
originalCall.RequestStream,
gcpResponseAsync,
originalCall.ResponseHeadersAsync,
() => originalCall.GetStatus(),
() => originalCall.GetTrailers(),
() => originalCall.Dispose()));
async Task <TResponse> DecrementCountAndPropagateResult(Task <TResponse> task)
{
try
{
return(await task.ConfigureAwait(false));
}
finally
{
channelRef.ActiveStreamCountDecr();
}
}
}
private ChannelRef PreProcess <TRequest>(AffinityConfig affinityConfig, TRequest request)
{
// Gets the affinity bound key if required in the request method.
string boundKey = null;
if (affinityConfig != null && affinityConfig.Command == AffinityConfig.Types.Command.Bound)
{
boundKey = GetAffinityKeysFromProto(affinityConfig.AffinityKey, (IMessage)request).SingleOrDefault();
}
ChannelRef channelRef = GetChannelRef(boundKey);
channelRef.ActiveStreamCountIncr();
return(channelRef);
}
private void Bind(ChannelRef channelRef, string affinityKey)
{
if (!string.IsNullOrEmpty(affinityKey))
{
lock (thisLock)
{
// TODO: What should we do if the dictionary already contains this key, but for a different channel ref?
if (!channelRefByAffinityKey.Keys.Contains(affinityKey))
{
channelRefByAffinityKey.Add(affinityKey, channelRef);
}
channelRefByAffinityKey[affinityKey].AffinityCountIncr();
}
}
}
// Note: response may be default(TResponse) in the case of a failure. We only expect to be called from
// protobuf-based calls anyway, so it will always be a class type, and will never be null for success cases.
// We can therefore check for nullity rather than having a separate "success" parameter.
private void PostProcess <TResponse>(AffinityConfig affinityConfig, ChannelRef channelRef, string boundKey, TResponse response)
{
channelRef.ActiveStreamCountDecr();
// Process BIND or UNBIND if the method has affinity feature enabled, but only for successful calls.
if (affinityConfig != null && response != null)
{
if (affinityConfig.Command == AffinityConfig.Types.Command.Bind)
{
Bind(channelRef, GetAffinityKeyFromProto(affinityConfig.AffinityKey, (IMessage)response));
}
else if (affinityConfig.Command == AffinityConfig.Types.Command.Unbind)
{
Unbind(boundKey);
}
}
}
private ChannelRef GetChannelRef(string affinityKey = null)
{
// TODO(fengli): Supports load reporting.
lock (thisLock)
{
if (!string.IsNullOrEmpty(affinityKey))
{
// Finds the gRPC channel according to the affinity key.
if (channelRefByAffinityKey.TryGetValue(affinityKey, out ChannelRef channelRef))
{
return(channelRef);
}
// TODO(fengli): Affinity key not found, log an error.
}
// TODO(fengli): Creates new gRPC channels on demand, depends on the load reporting.
IOrderedEnumerable <ChannelRef> orderedChannelRefs =
channelRefs.OrderBy(channelRef => channelRef.ActiveStreamCount);
foreach (ChannelRef channelRef in orderedChannelRefs)
{
if (channelRef.ActiveStreamCount < apiConfig.ChannelPool.MaxConcurrentStreamsLowWatermark)
{
// If there's a free channel, use it.
return(channelRef);
}
else
{
// If all channels are busy, break.
break;
}
}
int count = channelRefs.Count;
if (count < apiConfig.ChannelPool.MaxSize)
{
// Creates a new gRPC channel.
GrpcEnvironment.Logger.Info("Grpc.Gcp creating new channel");
Channel channel = new Channel(target, credentials,
options.Concat(new[] { new ChannelOption(ClientChannelId, Interlocked.Increment(ref clientChannelIdCounter)) }));
ChannelRef channelRef = new ChannelRef(channel, count);
channelRefs.Add(channelRef);
return(channelRef);
}
// If all channels are overloaded and the channel pool is full already,
// return the channel with least active streams.
return(orderedChannelRefs.First());
}
}
private Tuple <ChannelRef, string> PreProcess <TRequest>(AffinityConfig affinityConfig, TRequest request)
{
// Gets the affinity bound key if required in the request method.
string boundKey = null;
if (affinityConfig != null)
{
if (affinityConfig.Command == AffinityConfig.Types.Command.Bound ||
affinityConfig.Command == AffinityConfig.Types.Command.Unbind)
{
boundKey = GetAffinityKeyFromProto(affinityConfig.AffinityKey, (IMessage)request);
}
}
ChannelRef channelRef = GetChannelRef(boundKey);
channelRef.ActiveStreamCountIncr();
return(new Tuple <ChannelRef, string>(channelRef, boundKey));
}
BlockingUnaryCall <TRequest, TResponse>(Method <TRequest, TResponse> method, string host, CallOptions options, TRequest request)
{
affinityByMethod.TryGetValue(method.FullName, out AffinityConfig affinityConfig);
ChannelRef channelRef = PreProcess(affinityConfig, request);
var callDetails = new CallInvocationDetails <TRequest, TResponse>(channelRef.Channel, method, host, options);
TResponse response = default(TResponse);
try
{
response = Calls.BlockingUnaryCall <TRequest, TResponse>(callDetails, request);
return(response);
}
finally
{
PostProcess(affinityConfig, channelRef, request, response);
}
}
AsyncDuplexStreamingCall <TRequest, TResponse>(Method <TRequest, TResponse> method, string host, CallOptions options)
{
// No channel affinity feature for duplex streaming call.
ChannelRef channelRef = GetChannelRef();
var callDetails = new CallInvocationDetails <TRequest, TResponse>(channelRef.Channel, method, host, options);
var originalCall = Calls.AsyncDuplexStreamingCall(callDetails);
// Decrease the active streams count once the streaming response finishes its final batch.
var gcpResponseStream = new GcpClientResponseStream <TRequest, TResponse>(
originalCall.ResponseStream,
(resp) => channelRef.ActiveStreamCountDecr());
// Create a wrapper of the original AsyncDuplexStreamingCall.
return(new AsyncDuplexStreamingCall <TRequest, TResponse>(
originalCall.RequestStream,
gcpResponseStream,
originalCall.ResponseHeadersAsync,
() => originalCall.GetStatus(),
() => originalCall.GetTrailers(),
() => originalCall.Dispose()));
}
AsyncServerStreamingCall <TRequest, TResponse>(Method <TRequest, TResponse> method, string host, CallOptions options, TRequest request)
{
affinityByMethod.TryGetValue(method.FullName, out AffinityConfig affinityConfig);
ChannelRef channelRef = PreProcess(affinityConfig, request);
var callDetails = new CallInvocationDetails <TRequest, TResponse>(channelRef.Channel, method, host, options);
var originalCall = Calls.AsyncServerStreamingCall(callDetails, request);
// Executes affinity postprocess once the streaming response finishes its final batch.
var gcpResponseStream = new GcpClientResponseStream <TRequest, TResponse>(
originalCall.ResponseStream,
(resp) => PostProcess(affinityConfig, channelRef, request, resp));
// Create a wrapper of the original AsyncServerStreamingCall.
return(new AsyncServerStreamingCall <TResponse>(
gcpResponseStream,
originalCall.ResponseHeadersAsync,
() => originalCall.GetStatus(),
() => originalCall.GetTrailers(),
() => originalCall.Dispose()));
}
AsyncUnaryCall <TRequest, TResponse>(Method <TRequest, TResponse> method, string host, CallOptions options, TRequest request)
{
affinityByMethod.TryGetValue(method.FullName, out AffinityConfig affinityConfig);
Tuple <ChannelRef, string> tupleResult = PreProcess(affinityConfig, request);
ChannelRef channelRef = tupleResult.Item1;
string boundKey = tupleResult.Item2;
var callDetails = new CallInvocationDetails <TRequest, TResponse>(channelRef.Channel, method, host, options);
var originalCall = Calls.AsyncUnaryCall(callDetails, request);
// Executes affinity postprocess once the async response finishes.
var gcpResponseAsync = PostProcessPropagateResult(originalCall.ResponseAsync);
// Create a wrapper of the original AsyncUnaryCall.
return(new AsyncUnaryCall <TResponse>(
gcpResponseAsync,
originalCall.ResponseHeadersAsync,
() => originalCall.GetStatus(),
() => originalCall.GetTrailers(),
() => originalCall.Dispose()));
async Task <TResponse> PostProcessPropagateResult(Task <TResponse> task)
{
TResponse response = default(TResponse);
try
{
response = await task.ConfigureAwait(false);
return(response);
}
finally
{
PostProcess(affinityConfig, channelRef, boundKey, response);
}
}
}