remove RX as the primary pipe for gRPC communications

src/WebJobs.Script.Grpc/Channel/GrpcWorkerChannelFactory.cs

@@ -7,6 +7,7 @@
 using System.Reactive.Linq;
 using Microsoft.Azure.WebJobs.Script.Diagnostics;
 using Microsoft.Azure.WebJobs.Script.Eventing;
+using Microsoft.Azure.WebJobs.Script.Grpc.Eventing;
 using Microsoft.Azure.WebJobs.Script.Workers;
 using Microsoft.Azure.WebJobs.Script.Workers.Rpc;
 using Microsoft.Azure.WebJobs.Script.Workers.SharedMemoryDataTransfer;
@@ -47,6 +48,7 @@ public IRpcWorkerChannel Create(string scriptRootPath, string runtime, IMetricsL
                 throw new InvalidOperationException($"WorkerCofig for runtime: {runtime} not found");
             }
             string workerId = Guid.NewGuid().ToString();
+            _eventManager.AddGrpcChannels(workerId); // prepare the inbound/outbound dedicated channels
             ILogger workerLogger = _loggerFactory.CreateLogger($"Worker.LanguageWorkerChannel.{runtime}.{workerId}");
             IWorkerProcess rpcWorkerProcess = _rpcWorkerProcessFactory.Create(workerId, runtime, scriptRootPath, languageWorkerConfig);
             return new GrpcWorkerChannel(

src/WebJobs.Script.Grpc/Eventing/GrpcEventExtensions.cs

@@ -0,0 +1,67 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Threading.Channels;
+using Microsoft.Azure.WebJobs.Script.Eventing;
+
+namespace Microsoft.Azure.WebJobs.Script.Grpc.Eventing;
+
+internal static class GrpcEventExtensions
+{
+    // flow here is:
+    // 1) external request is proxied to the the GrpcWorkerChannel via one of the many Send* APIs, which writes
+    //    to outbound-writer; this means we can have concurrent writes to outbound
+    // 2) if an out-of-process function is connected, a FunctionRpcService-EventStream will consume
+    //    from outbound-reader (we'll allow for the multi-stream possibility, hence concurrent), and push it via gRPC
+    // 3) when the out-of-process function provides a response to FunctionRpcService-EventStream, it is written to
+    //    inbound-writer (note we will allow for multi-stream possibility)
+    // 4) the GrpcWorkerChannel has a single dedicated consumer of inbound-reader, which it then marries to
+    //    in-flight operations
+    internal static readonly UnboundedChannelOptions InboundOptions = new UnboundedChannelOptions
+    {
+        SingleReader = true, // see 4
+        SingleWriter = false, // see 3
+        AllowSynchronousContinuations = false,
+    };
+
+    internal static readonly UnboundedChannelOptions OutboundOptions = new UnboundedChannelOptions
+    {
+        SingleReader = false, // see 2
+        SingleWriter = false, // see 1
+        AllowSynchronousContinuations = false,
+    };
+
+    public static void AddGrpcChannels(this IScriptEventManager manager, string workerId)
+    {
+        var inbound = Channel.CreateUnbounded<InboundGrpcEvent>(InboundOptions);
+        if (manager.TryAddWorkerState(workerId, inbound))
+        {
+            var outbound = Channel.CreateUnbounded<OutboundGrpcEvent>(OutboundOptions);
+            if (manager.TryAddWorkerState(workerId, outbound))
+            {
+                return; // successfully added both
+            }
+            // we added the inbound but not the outbound; revert
+            manager.TryRemoveWorkerState(workerId, out inbound);
+        }
+         // this is not anticipated, so don't panic abount the allocs above
+        throw new ArgumentException("Duplicate worker id: " + workerId, nameof(workerId));
+    }
+
+    public static bool TryGetGrpcChannels(this IScriptEventManager manager, string workerId, out Channel<InboundGrpcEvent> inbound, out Channel<OutboundGrpcEvent> outbound)
+        => manager.TryGetWorkerState(workerId, out inbound) & manager.TryGetWorkerState(workerId, out outbound);
+
+    public static void RemoveGrpcChannels(this IScriptEventManager manager, string workerId)
+    {
+        // remove any channels, and shut them down
+        if (manager.TryGetWorkerState<Channel<InboundGrpcEvent>>(workerId, out var inbound))
+        {
+            inbound.Writer.TryComplete();
+        }
+        if (manager.TryGetWorkerState<Channel<OutboundGrpcEvent>>(workerId, out var outbound))
+        {
+            outbound.Writer.TryComplete();
+        }
+    }
+}

src/WebJobs.Script.Grpc/Server/FunctionRpcService.cs

@@ -2,10 +2,9 @@
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
-using System.Collections.Generic;
-using System.Reactive.Concurrency;
 using System.Reactive.Linq;
 using System.Threading;
+using System.Threading.Channels;
 using System.Threading.Tasks;
 using Grpc.Core;
 using Microsoft.Azure.WebJobs.Script.Eventing;
@@ -21,7 +20,6 @@ namespace Microsoft.Azure.WebJobs.Script.Grpc
     // TODO: move to WebJobs.Script.Grpc package and provide event stream abstraction
     internal class FunctionRpcService : FunctionRpc.FunctionRpcBase
     {
-        private readonly SemaphoreSlim _writeLock = new SemaphoreSlim(1, 1);
         private readonly IScriptEventManager _eventManager;
         private readonly ILogger _logger;
 
@@ -33,64 +31,48 @@ public FunctionRpcService(IScriptEventManager eventManager, ILogger<FunctionRpcS
 
         public override async Task EventStream(IAsyncStreamReader<StreamingMessage> requestStream, IServerStreamWriter<StreamingMessage> responseStream, ServerCallContext context)
         {
-            var cancelSource = new TaskCompletionSource<bool>();
-            IDictionary<string, IDisposable> outboundEventSubscriptions = new Dictionary<string, IDisposable>();
-
+            var cancelSource = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var cts = CancellationTokenSource.CreateLinkedTokenSource(context.CancellationToken);
+            CancellationTokenRegistration ctr = cts.Token.Register(static state => ((TaskCompletionSource<bool>)state).TrySetResult(false), cancelSource);
             try
             {
-                context.CancellationToken.Register(() => cancelSource.TrySetResult(false));
-                Func<Task<bool>> messageAvailable = async () =>
+                static Task<Task<bool>> MoveNextAsync(IAsyncStreamReader<StreamingMessage> requestStream, TaskCompletionSource<bool> cancelSource)
                 {
                     // GRPC does not accept cancellation tokens for individual reads, hence wrapper
                     var requestTask = requestStream.MoveNext(CancellationToken.None);
-                    var completed = await Task.WhenAny(cancelSource.Task, requestTask);
-                    return completed.Result;
-                };
+                    return Task.WhenAny(cancelSource.Task, requestTask);
+                }
 
-                if (await messageAvailable())
+                if (await await MoveNextAsync(requestStream, cancelSource))
                 {
-                    string workerId = requestStream.Current.StartStream.WorkerId;
-                    _logger.LogDebug("Established RPC channel. WorkerId: {workerId}", workerId);
-                    outboundEventSubscriptions.Add(workerId, _eventManager.OfType<OutboundGrpcEvent>()
-                        .Where(evt => evt.WorkerId == workerId)
-                        .ObserveOn(NewThreadScheduler.Default)
-                        .Subscribe(async evt =>
+                    var currentMessage = requestStream.Current;
+                    // expect first operation (and only the first; we don't support re-registration) to be StartStream
+                    if (currentMessage.ContentCase == MsgType.StartStream)
+                    {
+                        var workerId = currentMessage.StartStream?.WorkerId;
+                        if (!string.IsNullOrEmpty(workerId) && _eventManager.TryGetGrpcChannels(workerId, out var inbound, out var outbound))
                         {
-                            try
-                            {
-                                if (evt.MessageType == MsgType.InvocationRequest)
-                                {
-                                    _logger.LogTrace("Writing invocation request invocationId: {invocationId} to workerId: {workerId}", evt.Message.InvocationRequest.InvocationId, workerId);
-                                }
+                            // send work
+                            _ = PushFromOutboundToGrpc(workerId, responseStream, outbound.Reader, cts.Token);
 
-                                try
+                            // this loop is "pull from gRPC and push to inbound"
+                            do
+                            {
+                                currentMessage = requestStream.Current;
+                                if (currentMessage.ContentCase == MsgType.InvocationResponse && !string.IsNullOrEmpty(currentMessage.InvocationResponse?.InvocationId))
                                 {
-                                    // WriteAsync only allows one pending write at a time, so we
-                                    // serialize access to the stream for each subscription
-                                    await _writeLock.WaitAsync();
-                                    await responseStream.WriteAsync(evt.Message);
+                                    _logger.LogTrace("Received invocation response for invocationId: {invocationId} from workerId: {workerId}", currentMessage.InvocationResponse.InvocationId, workerId);
                                 }
-                                finally
+                                var newInbound = new InboundGrpcEvent(workerId, currentMessage);
+                                if (!inbound.Writer.TryWrite(newInbound))
                                 {
-                                    _writeLock.Release();
+                                    await inbound.Writer.WriteAsync(newInbound);
                                 }
+                                currentMessage = null; // allow old messages to be collected while we wait
                             }
-                            catch (Exception subscribeEventEx)
-                            {
-                                _logger.LogError(subscribeEventEx, "Error writing message type {messageType} to workerId: {workerId}", evt.MessageType, workerId);
-                            }
-                        }));
-
-                    do
-                    {
-                        var currentMessage = requestStream.Current;
-                        if (currentMessage.InvocationResponse != null && !string.IsNullOrEmpty(currentMessage.InvocationResponse.InvocationId))
-                        {
-                            _logger.LogTrace("Received invocation response for invocationId: {invocationId} from workerId: {workerId}", currentMessage.InvocationResponse.InvocationId, workerId);
+                            while (await await MoveNextAsync(requestStream, cancelSource));
                         }
-                        _eventManager.Publish(new InboundGrpcEvent(workerId, currentMessage));
                     }
-                    while (await messageAvailable());
                 }
             }
             catch (Exception rpcException)
@@ -101,14 +83,47 @@ public override async Task EventStream(IAsyncStreamReader<StreamingMessage> requ
             }
             finally
             {
-                foreach (var sub in outboundEventSubscriptions)
-                {
-                    sub.Value?.Dispose();
-                }
+                cts.Cancel();
+                ctr.Dispose();
 
                 // ensure cancellationSource task completes
                 cancelSource.TrySetResult(false);
             }
         }
+
+        private async Task PushFromOutboundToGrpc(string workerId, IServerStreamWriter<StreamingMessage> responseStream, ChannelReader<OutboundGrpcEvent> source, CancellationToken cancellationToken)
+        {
+            try
+            {
+                _logger.LogDebug("Established RPC channel. WorkerId: {workerId}", workerId);
+                await Task.Yield(); // free up the caller
+                while (await source.WaitToReadAsync(cancellationToken))
+                {
+                    while (source.TryRead(out var evt))
+                    {
+                        if (evt.MessageType == MsgType.InvocationRequest)
+                        {
+                            _logger.LogTrace("Writing invocation request invocationId: {invocationId} to workerId: {workerId}", evt.Message.InvocationRequest.InvocationId, workerId);
+                        }
+                        try
+                        {
+                            await responseStream.WriteAsync(evt.Message);
+                        }
+                        catch (Exception subscribeEventEx)
+                        {
+                            _logger.LogError(subscribeEventEx, "Error writing message type {messageType} to workerId: {workerId}", evt.MessageType, workerId);
+                        }
+                    }
+                }
+            }
+            catch (OperationCanceledException oce) when (oce.CancellationToken == cancellationToken)
+            {
+                // that's fine; normaly exit through cancellation
+            }
+            catch (Exception ex)
+            {
+                _logger.LogError(ex, "Error pushing from outbound to gRPC");
+            }
+        }
     }
 }

src/WebJobs.Script.Grpc/WebJobs.Script.Grpc.csproj

@@ -30,6 +30,7 @@
     <PackageReference Include="StyleCop.Analyzers" Version="1.1.0-beta004">
       <PrivateAssets>all</PrivateAssets>
     </PackageReference>
+    <PackageReference Include="System.Threading.Channels" Version="6.0.0" />
   </ItemGroup>
 
   <ItemGroup>

src/WebJobs.Script/Eventing/IScriptEventManager.cs

@@ -8,5 +8,11 @@ namespace Microsoft.Azure.WebJobs.Script.Eventing
     public interface IScriptEventManager : IObservable<ScriptEvent>
     {
         void Publish(ScriptEvent scriptEvent);
+
+        bool TryAddWorkerState<T>(string workerId, T state);
+
+        bool TryGetWorkerState<T>(string workerId, out T state);
+
+        bool TryRemoveWorkerState<T>(string workerId, out T state);
     }
 }

src/WebJobs.Script/Eventing/ScriptEventManager.cs

@@ -2,13 +2,16 @@
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Concurrent;
 using System.Reactive.Subjects;
 
 namespace Microsoft.Azure.WebJobs.Script.Eventing
 {
-    public sealed class ScriptEventManager : IScriptEventManager, IDisposable
+    public class ScriptEventManager : IScriptEventManager, IDisposable
     {
         private readonly Subject<ScriptEvent> _subject = new Subject<ScriptEvent>();
+        private readonly ConcurrentDictionary<(string, Type), object> _workerState = new ();
+
         private bool _disposed = false;
 
         public void Publish(ScriptEvent scriptEvent)
@@ -47,5 +50,35 @@ private void Dispose(bool disposing)
         }
 
         public void Dispose() => Dispose(true);
+
+        bool IScriptEventManager.TryAddWorkerState<T>(string workerId, T state)
+        {
+            var key = (workerId, typeof(T));
+            return _workerState.TryAdd(key, state);
+        }
+
+        bool IScriptEventManager.TryGetWorkerState<T>(string workerId, out T state)
+        {
+            var key = (workerId, typeof(T));
+            if (_workerState.TryGetValue(key, out var tmp) && tmp is T typed)
+            {
+                state = typed;
+                return true;
+            }
+            state = default;
+            return false;
+        }
+
+        bool IScriptEventManager.TryRemoveWorkerState<T>(string workerId, out T state)
+        {
+            var key = (workerId, typeof(T));
+            if (_workerState.TryRemove(key, out var tmp) && tmp is T typed)
+            {
+                state = typed;
+                return true;
+            }
+            state = default;
+            return false;
+        }
     }
 }

src/WebJobs.Script/Workers/Rpc/RpcWorkerConstants.cs

@@ -62,6 +62,7 @@ public static class RpcWorkerConstants
 
         // Host Capabilities
         public const string V2Compatable = "V2Compatable";
+        public const string MultiStream = nameof(MultiStream);
 
         // dotnet executable file path components
         public const string DotNetExecutableName = "dotnet";