Update capabilities and log worker metadata after specialization (#9020)

* Update capabilities and log worker metadata in env reload request.

* Updated tests

* Release notes

* Updated release notes to include proto upgrade change

Author: kshyju

release_notes.md

@@ -16,4 +16,6 @@
 
 **Release sprint:** Sprint 132
 [ [bugs](https://github.com/Azure/azure-functions-host/issues?q=is%3Aissue+milestone%3A%22Functions+Sprint+132%22+label%3Abug+is%3Aclosed) | [features](https://github.com/Azure/azure-functions-host/issues?q=is%3Aissue+milestone%3A%22Functions+Sprint+132%22+label%3Afeature+is%3Aclosed) ]
-- Update App Service Authentication/Authorization on Linux Consumption from 1.4.17 to 1.5.1 [Release Note](https://github.com/Azure/app-service-announcements/issues/406)
+- Update App Service Authentication/Authorization on Linux Consumption from 1.4.17 to 1.5.1 [Release Note](https://github.com/Azure/app-service-announcements/issues/406)
+- Upgraded proto files to v1.7.1-protofile release (#9023)
+- Refresh capabilities & log worker metadata after specialization (#9020)

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

@@ -78,7 +78,7 @@ internal class GrpcWorkerChannel : IRpcWorkerChannel, IDisposable
         private TaskCompletionSource<List<RawFunctionMetadata>> _functionsIndexingTask = new TaskCompletionSource<List<RawFunctionMetadata>>(TaskCreationOptions.RunContinuationsAsynchronously);
         private TimeSpan _functionLoadTimeout = TimeSpan.FromMinutes(1);
         private bool _isSharedMemoryDataTransferEnabled;
-        private bool _cancelCapabilityEnabled;
+        private bool? _cancelCapabilityEnabled;
         private bool _isWorkerApplicationInsightsLoggingEnabled;
 
         private System.Timers.Timer _timer;
@@ -357,6 +357,11 @@ internal WorkerInitRequest GetWorkerInitRequest()
         internal void FunctionEnvironmentReloadResponse(FunctionEnvironmentReloadResponse res, IDisposable latencyEvent)
         {
             _workerChannelLogger.LogDebug("Received FunctionEnvironmentReloadResponse from WorkerProcess with Pid: '{0}'", _rpcWorkerProcess.Id);
+
+            LogWorkerMetadata(res.WorkerMetadata);
+            UpdateCapabilities(res.Capabilities);
+            _cancelCapabilityEnabled ??= !string.IsNullOrEmpty(_workerCapabilities.GetCapabilityState(RpcWorkerConstants.HandlesInvocationCancelMessage));
+
             if (res.Result.IsFailure(out Exception reloadEnvironmentVariablesException))
             {
                 _workerChannelLogger.LogError(reloadEnvironmentVariablesException, "Failed to reload environment variables");
@@ -375,13 +380,10 @@ internal void WorkerInitResponse(GrpcEvent initEvent)
             _initMessage = initEvent.Message.WorkerInitResponse;
             _workerChannelLogger.LogDebug("Worker capabilities: {capabilities}", _initMessage.Capabilities);
 
-            if (_initMessage.WorkerMetadata != null)
-            {
-                _initMessage.UpdateWorkerMetadata(_workerConfig);
-                var workerMetadata = _initMessage.WorkerMetadata.ToString();
-                _metricsLogger.LogEvent(MetricEventNames.WorkerMetadata, functionName: null, workerMetadata);
-                _workerChannelLogger.LogDebug("Worker metadata: {workerMetadata}", workerMetadata);
-            }
+            // In placeholder scenario, the capabilities and worker metadata will not be available
+            // until specialization is done (env reload request). So these can be removed from worker init response code path.
+            // to do to track this: https://github.com/Azure/azure-functions-host/issues/9019
+            LogWorkerMetadata(_initMessage.WorkerMetadata);
 
             if (_initMessage.Result.IsFailure(out Exception exc))
             {
@@ -395,7 +397,7 @@ internal void WorkerInitResponse(GrpcEvent initEvent)
             UpdateCapabilities(_initMessage.Capabilities);
 
             _isSharedMemoryDataTransferEnabled = IsSharedMemoryDataTransferEnabled();
-            _cancelCapabilityEnabled = !string.IsNullOrEmpty(_workerCapabilities.GetCapabilityState(RpcWorkerConstants.HandlesInvocationCancelMessage));
+            _cancelCapabilityEnabled ??= !string.IsNullOrEmpty(_workerCapabilities.GetCapabilityState(RpcWorkerConstants.HandlesInvocationCancelMessage));
 
             if (!_isSharedMemoryDataTransferEnabled)
             {
@@ -413,6 +415,19 @@ internal void WorkerInitResponse(GrpcEvent initEvent)
             _workerInitTask.TrySetResult(true);
         }
 
+        private void LogWorkerMetadata(WorkerMetadata workerMetadata)
+        {
+            if (workerMetadata == null)
+            {
+                return;
+            }
+
+            workerMetadata.UpdateWorkerMetadata(_workerConfig);
+            var workerMetadataString = workerMetadata.ToString();
+            _metricsLogger.LogEvent(MetricEventNames.WorkerMetadata, functionName: null, workerMetadataString);
+            _workerChannelLogger.LogDebug("Worker metadata: {workerMetadata}", workerMetadataString);
+        }
+
         // Allow tests to add capabilities, even if not directly supported by the worker.
         internal virtual void UpdateCapabilities(IDictionary<string, string> fields)
         {
@@ -668,7 +683,7 @@ await SendStreamingMessageAsync(new StreamingMessage
                     InvocationRequest = invocationRequest
                 });
 
-                if (_cancelCapabilityEnabled)
+                if (_cancelCapabilityEnabled != null && _cancelCapabilityEnabled.Value)
                 {
                     context.CancellationToken.Register(() => SendInvocationCancel(invocationRequest.InvocationId));
                 }

src/WebJobs.Script.Grpc/MessageExtensions/GrpcMessageExtensionUtilities.cs

@@ -78,12 +78,12 @@ public static Tuple<string, string, CookieOptions> RpcHttpCookieConverter(RpcHtt
             return new Tuple<string, string, CookieOptions>(cookie.Name, cookie.Value, cookieOptions);
         }
 
-        internal static void UpdateWorkerMetadata(this WorkerInitResponse initResponse, RpcWorkerConfig workerConfig)
+        internal static void UpdateWorkerMetadata(this WorkerMetadata workerMetadata, RpcWorkerConfig workerConfig)
         {
-            initResponse.WorkerMetadata.RuntimeName = string.IsNullOrEmpty(initResponse.WorkerMetadata.RuntimeName)
-                                            ? workerConfig.Description.Language : initResponse.WorkerMetadata.RuntimeName;
-            initResponse.WorkerMetadata.RuntimeVersion = string.IsNullOrEmpty(initResponse.WorkerMetadata.RuntimeVersion)
-                                            ? workerConfig.Description.DefaultRuntimeVersion : initResponse.WorkerMetadata.RuntimeVersion;
+            workerMetadata.RuntimeName = string.IsNullOrEmpty(workerMetadata.RuntimeName)
+                                            ? workerConfig.Description.Language : workerMetadata.RuntimeName;
+            workerMetadata.RuntimeVersion = string.IsNullOrEmpty(workerMetadata.RuntimeVersion)
+                                            ? workerConfig.Description.DefaultRuntimeVersion : workerMetadata.RuntimeVersion;
         }
 
         private static SameSiteMode RpcSameSiteEnumConverter(RpcHttpCookie.Types.SameSite sameSite)

test/WebJobs.Script.Tests/Workers/Rpc/GrpcWorkerChannelTests.cs

@@ -702,7 +702,7 @@ public async Task SendSendFunctionEnvironmentReloadRequest_PublishesOutboundEven
                 Environment.SetEnvironmentVariable("TestNull", null);
                 Environment.SetEnvironmentVariable("TestEmpty", string.Empty);
                 Environment.SetEnvironmentVariable("TestValid", "TestValue");
-                _testFunctionRpcService.AutoReply(StreamingMessage.ContentOneofCase.FunctionEnvironmentReloadRequest);
+                _testFunctionRpcService.AutoReply(StreamingMessage.ContentOneofCase.FunctionEnvironmentReloadRequest, workerSupportsSpecialization: true);
                 var pending = _workerChannel.SendFunctionEnvironmentReloadRequest();
                 await Task.Delay(500);
                 await pending; // this can timeout
@@ -718,6 +718,14 @@ public async Task SendSendFunctionEnvironmentReloadRequest_PublishesOutboundEven
             ShowOutput(traces);
             var functionLoadLogs = traces.Where(m => string.Equals(m.FormattedMessage, "Sending FunctionEnvironmentReloadRequest to WorkerProcess with Pid: '910'"));
             Assert.Equal(1, functionLoadLogs.Count());
+
+            // for specialization use case, env reload response include worker metadata and capabilities.
+            var metatadataLog = traces.Where(m => string.Equals(m.FormattedMessage,
+                @"Worker metadata: { ""runtimeName"": "".NET"", ""runtimeVersion"": ""7.0"", ""workerVersion"": ""1.0.0"", ""workerBitness"": ""x64"" }"));
+            var capabilityUpdateLog = traces.Where(m => string.Equals(m.FormattedMessage,
+                @"Updating capabilities: { ""RpcHttpBodyOnly"": ""True"", ""TypedDataCollection"": ""True"" }"));
+            Assert.Single(metatadataLog);
+            Assert.Single(capabilityUpdateLog);
         }
 
         [Fact]

test/WebJobs.Script.Tests/Workers/Rpc/TestFunctionRpcService.cs

@@ -11,9 +11,7 @@
 using Microsoft.Azure.WebJobs.Script.Eventing;
 using Microsoft.Azure.WebJobs.Script.Grpc.Eventing;
 using Microsoft.Azure.WebJobs.Script.Grpc.Messages;
-using Microsoft.Azure.WebJobs.Script.Workers.Rpc;
 using Microsoft.Extensions.Logging;
-using Newtonsoft.Json.Linq;
 
 namespace Microsoft.Azure.WebJobs.Script.Tests.Workers.Rpc
 {
@@ -41,12 +39,12 @@ public TestFunctionRpcService(IScriptEventManager eventManager, string workerId,
         public void OnMessage(StreamingMessage.ContentOneofCase messageType, Action<OutboundGrpcEvent> callback)
             => _handlers.AddOrUpdate(messageType, callback, (messageType, oldValue) => oldValue + callback);
 
-        public void AutoReply(StreamingMessage.ContentOneofCase messageType)
+        public void AutoReply(StreamingMessage.ContentOneofCase messageType, bool workerSupportsSpecialization = false)
         {
             // apply standard default responses
             Action<OutboundGrpcEvent> callback = messageType switch
             {
-                StreamingMessage.ContentOneofCase.FunctionEnvironmentReloadRequest => _ => PublishFunctionEnvironmentReloadResponseEvent(),
+                StreamingMessage.ContentOneofCase.FunctionEnvironmentReloadRequest => _ => PublishFunctionEnvironmentReloadResponseEvent(workerSupportsSpecialization),
                 _ => null,
             };
             if (callback is not null)
@@ -178,13 +176,27 @@ public void PublishFunctionLoadResponsesEvent(List<string> functionIds, StatusRe
             Write(responseMessage);
         }
 
-        public void PublishFunctionEnvironmentReloadResponseEvent()
+        public void PublishFunctionEnvironmentReloadResponseEvent(bool workerSupportsSpecialization)
         {
-            FunctionEnvironmentReloadResponse relaodEnvResponse = GetTestFunctionEnvReloadResponse();
+            FunctionEnvironmentReloadResponse reloadEnvResponse = GetTestFunctionEnvReloadResponse();
             StreamingMessage responseMessage = new StreamingMessage()
             {
-                FunctionEnvironmentReloadResponse = relaodEnvResponse
+                FunctionEnvironmentReloadResponse = reloadEnvResponse
             };
+
+            if (workerSupportsSpecialization)
+            {
+                responseMessage.FunctionEnvironmentReloadResponse.WorkerMetadata = new()
+                {
+                    RuntimeName = ".NET",
+                    RuntimeVersion = "7.0",
+                    WorkerVersion = "1.0.0",
+                    WorkerBitness = "x64"
+                };
+                responseMessage.FunctionEnvironmentReloadResponse.Capabilities.Add("RpcHttpBodyOnly", bool.TrueString);
+                responseMessage.FunctionEnvironmentReloadResponse.Capabilities.Add("TypedDataCollection", bool.TrueString);
+            }
+
             Write(responseMessage);
         }