feat(worker-proxy): relay extension RPC streams

src/Functions.WorkerProxy/ExtensionRpc/ExtensionCall.cs

@@ -0,0 +1,130 @@
+// 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.Collections.Generic;
+using System.Threading;
+using System.Threading.Channels;
+using System.Threading.Tasks;
+using Microsoft.Azure.WebJobs.Script.Grpc.Messages;
+
+namespace Azure.Functions.WorkerProxy.ExtensionRpc;
+
+/// <summary>
+/// Represents one worker-originated gRPC call multiplexed over the shared extension RPC stream.
+/// </summary>
+/// <param name="stream">The physical extension RPC stream carrying the call.</param>
+/// <param name="callId">The identifier used to correlate lifecycle messages for the call.</param>
+/// <param name="ready">The negotiated transport settings for the stream.</param>
+internal sealed class ExtensionCall(ExtensionRpcStream stream, string callId, ExtensionRpcReady ready)
+    : IAsyncDisposable
+{
+    private const int InboundQueueCapacity = 32;
+
+    private readonly Channel<ExtensionRpcMessage> _inbound = Channel.CreateBounded<ExtensionRpcMessage>(
+        new BoundedChannelOptions(InboundQueueCapacity)
+        {
+            SingleReader = true,
+            SingleWriter = true,
+            AllowSynchronousContinuations = false,
+            FullMode = BoundedChannelFullMode.Wait,
+        });
+
+    private readonly ExtensionRpcCreditWindow _requestCredits = new(ready.InitialReceiveWindowBytes);
+    private int _disposed;
+
+    /// <summary>
+    /// Gets the identifier used to correlate lifecycle messages for the call.
+    /// </summary>
+    public string CallId { get; } = callId;
+
+    /// <summary>
+    /// Gets the identifier of the physical extension RPC stream carrying this call.
+    /// </summary>
+    public string StreamId => stream.StreamId;
+
+    /// <summary>
+    /// Gets the transport settings negotiated for the stream.
+    /// </summary>
+    public ExtensionRpcReady Ready { get; } = ready;
+
+    /// <summary>
+    /// Gets the token that is cancelled when the physical stream closes.
+    /// </summary>
+    public CancellationToken CancellationToken => stream.CancellationToken;
+
+    /// <summary>
+    /// Gets a snapshot of the logical calls currently registered with the physical stream.
+    /// </summary>
+    public int ActiveCallCount => stream.ActiveCallCount;
+
+    /// <summary>
+    /// Reads ordered response lifecycle messages received from the host.
+    /// </summary>
+    /// <param name="cancellationToken">A token that cancels response enumeration.</param>
+    /// <returns>The ordered response messages for this call.</returns>
+    public IAsyncEnumerable<ExtensionRpcMessage> ReadAllAsync(CancellationToken cancellationToken)
+    {
+        return _inbound.Reader.ReadAllAsync(cancellationToken);
+    }
+
+    /// <summary>
+    /// Writes a request lifecycle message to the shared extension RPC stream.
+    /// </summary>
+    /// <param name="message">The message to write.</param>
+    /// <param name="cancellationToken">A token that cancels the write.</param>
+    /// <returns>A task that represents the asynchronous write.</returns>
+    public async ValueTask WriteAsync(ExtensionRpcMessage message, CancellationToken cancellationToken)
+    {
+        using CancellationTokenSource cancellationTokenSource =
+            CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, CancellationToken);
+        cancellationToken = cancellationTokenSource.Token;
+        if (message.ContentCase is ExtensionRpcMessage.ContentOneofCase.Data)
+        {
+            await _requestCredits.ReserveAsync((ulong)message.Data.Payload.Length, cancellationToken);
+        }
+
+        await stream.WriteExtensionMessageAsync(CallId, message, cancellationToken);
+    }
+
+    /// <summary>
+    /// Processes a response lifecycle message routed from the shared extension RPC stream.
+    /// </summary>
+    /// <param name="message">The inbound message to process.</param>
+    /// <param name="cancellationToken">A token that cancels queueing the message.</param>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    public async ValueTask HandleInboundAsync(ExtensionRpcMessage message, CancellationToken cancellationToken)
+    {
+        if (message.ContentCase is ExtensionRpcMessage.ContentOneofCase.WindowUpdate)
+        {
+            _requestCredits.Add(message.WindowUpdate.ByteCount);
+            return;
+        }
+
+        await _inbound.Writer.WriteAsync(message, cancellationToken);
+        if (message.ContentCase is ExtensionRpcMessage.ContentOneofCase.Complete)
+        {
+            await DisposeAsync();
+        }
+    }
+
+    /// <summary>
+    /// Completes the inbound response queue because its physical stream has ended.
+    /// </summary>
+    public void Complete()
+    {
+        _inbound.Writer.TryComplete();
+    }
+
+    /// <inheritdoc />
+    public ValueTask DisposeAsync()
+    {
+        if (Interlocked.Exchange(ref _disposed, 1) is 0)
+        {
+            stream.RemoveCall(CallId);
+            _inbound.Writer.TryComplete();
+        }
+
+        return ValueTask.CompletedTask;
+    }
+}

src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcCreditWindow.cs

@@ -0,0 +1,76 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Azure.Functions.WorkerProxy.ExtensionRpc;
+
+/// <summary>
+/// Coordinates byte-based flow-control credits for one direction of an extension RPC call.
+/// </summary>
+/// <param name="initialCredits">The number of bytes initially available to the sender.</param>
+internal sealed class ExtensionRpcCreditWindow(ulong initialCredits)
+{
+    private readonly Lock _syncLock = new();
+    private ulong _available = initialCredits;
+    private TaskCompletionSource _changed = CreateChangedSource();
+
+    /// <summary>
+    /// Adds byte credits granted by the receiver.
+    /// </summary>
+    /// <param name="credits">The number of credits to add.</param>
+    public void Add(ulong credits)
+    {
+        if (credits is 0)
+        {
+            return;
+        }
+
+        TaskCompletionSource changed;
+        lock (_syncLock)
+        {
+            _available = ulong.MaxValue - _available < credits ? ulong.MaxValue : _available + credits;
+            changed = _changed;
+            _changed = CreateChangedSource();
+        }
+
+        changed.TrySetResult();
+    }
+
+    /// <summary>
+    /// Waits for and reserves the requested number of byte credits.
+    /// </summary>
+    /// <param name="credits">The number of credits to reserve.</param>
+    /// <param name="cancellationToken">A token that cancels the wait.</param>
+    /// <returns>A task that completes when the credits have been reserved.</returns>
+    public async ValueTask ReserveAsync(ulong credits, CancellationToken cancellationToken)
+    {
+        if (credits is 0)
+        {
+            return;
+        }
+
+        while (true)
+        {
+            Task waitTask;
+            lock (_syncLock)
+            {
+                if (_available >= credits)
+                {
+                    _available -= credits;
+                    return;
+                }
+
+                waitTask = _changed.Task;
+            }
+
+            await waitTask.WaitAsync(cancellationToken);
+        }
+    }
+
+    private static TaskCompletionSource CreateChangedSource()
+    {
+        return new(TaskCreationOptions.RunContinuationsAsynchronously);
+    }
+}

src/Functions.WorkerProxy/ExtensionRpc/ExtensionRpcRelay.cs

@@ -0,0 +1,119 @@
+// 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;
+using System.Threading.Tasks;
+using Grpc.Core;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Azure.WebJobs.Script.Grpc.Messages;
+using ExtensionRpcService = Microsoft.Azure.WebJobs.Script.Grpc.Messages.ExtensionRpc;
+using GrpcRpcException = Grpc.Core.RpcException;
+
+namespace Azure.Functions.WorkerProxy.ExtensionRpc;
+
+/// <summary>
+/// Implements the runtime-facing extension RPC service and attaches its stream to the coordinator.
+/// </summary>
+/// <param name="endpoints">The WorkerProxy listener configuration.</param>
+/// <param name="streamCoordinator">The coordinator that owns the active extension RPC stream.</param>
+internal sealed class ExtensionRpcRelay(
+    WorkerProxyEndpointConfiguration endpoints, ExtensionRpcStreamCoordinator streamCoordinator)
+    : ExtensionRpcService.ExtensionRpcBase
+{
+    /// <inheritdoc />
+    public override Task EventStream(
+        IAsyncStreamReader<ExtensionRpcMessage> requestStream,
+        IServerStreamWriter<ExtensionRpcMessage> responseStream,
+        ServerCallContext context)
+    {
+        HttpContext httpContext = context.GetHttpContext();
+        if (!endpoints.TryGetRelaySide(httpContext.Connection.LocalPort, out FunctionRpcRelaySide side))
+        {
+            throw new GrpcRpcException(
+                new Status(StatusCode.Unimplemented, "ExtensionRpc is unavailable on this listener."));
+        }
+
+        if (side is not FunctionRpcRelaySide.Runtime)
+        {
+            throw new GrpcRpcException(
+                new Status(StatusCode.PermissionDenied, "ExtensionRpc is only available on the runtime gRPC port."));
+        }
+
+        return RelayAsync(requestStream, responseStream, context.CancellationToken);
+    }
+
+    /// <summary>
+    /// Relays inbound and outbound lifecycle messages for one physical extension RPC stream.
+    /// </summary>
+    /// <param name="requestStream">Messages received from the host runtime.</param>
+    /// <param name="responseStream">Messages sent to the host runtime.</param>
+    /// <param name="cancellationToken">A token that is cancelled when the stream ends.</param>
+    /// <returns>A task that represents the stream lifetime.</returns>
+    internal async Task RelayAsync(
+        IAsyncStreamReader<ExtensionRpcMessage> requestStream,
+        IServerStreamWriter<ExtensionRpcMessage> responseStream,
+        CancellationToken cancellationToken)
+    {
+        using CancellationTokenSource cancellationTokenSource =
+            CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        await using ExtensionRpcStreamLease lease = OpenStream(cancellationTokenSource.Token);
+        Task readTask = ReadInboundAsync(lease.Stream, requestStream, cancellationTokenSource.Token);
+        Task writeTask = WriteOutboundAsync(lease.Stream, responseStream, cancellationTokenSource.Token);
+
+        try
+        {
+            Task completedTask = await Task.WhenAny(readTask, writeTask);
+            await completedTask;
+        }
+        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+        {
+        }
+        finally
+        {
+            cancellationTokenSource.Cancel();
+            await Task.WhenAll(
+                ObserveCompletionAsync(readTask),
+                ObserveCompletionAsync(writeTask));
+        }
+    }
+
+    private ExtensionRpcStreamLease OpenStream(CancellationToken cancellationToken)
+    {
+        try
+        {
+            return streamCoordinator.Open(cancellationToken);
+        }
+        catch (InvalidOperationException exception)
+        {
+            throw new GrpcRpcException(new Status(StatusCode.AlreadyExists, exception.Message));
+        }
+    }
+
+    private static async Task ReadInboundAsync(
+        ExtensionRpcStream stream,
+        IAsyncStreamReader<ExtensionRpcMessage> requestStream,
+        CancellationToken cancellationToken)
+    {
+        while (await requestStream.MoveNext(cancellationToken))
+        {
+            await stream.HandleInboundAsync(requestStream.Current, cancellationToken);
+        }
+    }
+
+    private static async Task WriteOutboundAsync(
+        ExtensionRpcStream stream,
+        IServerStreamWriter<ExtensionRpcMessage> responseStream,
+        CancellationToken cancellationToken)
+    {
+        await foreach (ExtensionRpcMessage message in stream.Outbound.ReadAllAsync(cancellationToken))
+        {
+            await responseStream.WriteAsync(message, cancellationToken);
+        }
+    }
+
+    private static async Task ObserveCompletionAsync(Task task)
+    {
+        await task.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
+    }
+}