Implement host configuration property for handling pre-cancelled invocation requests (#9523) (#9826)

Co-authored-by: Lilian Kasem <[email protected]>

Author: v-smanchem

release_notes.md

@@ -8,3 +8,7 @@
 - Add function grouping information (https://github.com/Azure/azure-functions-host/pull/9735)
 - Bump `Microsoft.IdentityModel.Tokens`, `Microsoft.IdentityModel.Protocols.OpenIdConnect`, and
   `System.IdentityModel.Tokens.Jwt` from 6.32.0 to 6.35.0 (#9793)
+- Implement host configuration property for handling pre-cancelled invocation requests (#9523)
+  - If a worker supports CancellationTokens, cancelled invocations will now be sent to the worker by default
+    - Customers can opt-out of this behavior by setting `SendCanceledInvocationsToWorker` to `false` in host.json
+  - If a worker does not support CancellationTokens, cancelled invocations will not be sent to the worker

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

@@ -29,6 +29,7 @@
 using Microsoft.Azure.WebJobs.Script.Workers;
 using Microsoft.Azure.WebJobs.Script.Workers.Rpc;
 using Microsoft.Azure.WebJobs.Script.Workers.SharedMemoryDataTransfer;
+using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using Yarp.ReverseProxy.Forwarder;
@@ -42,6 +43,7 @@ namespace Microsoft.Azure.WebJobs.Script.Grpc
     internal partial class GrpcWorkerChannel : IRpcWorkerChannel, IDisposable
     {
         private readonly IScriptEventManager _eventManager;
+        private readonly IScriptHostManager _scriptHostManager;
         private readonly RpcWorkerConfig _workerConfig;
         private readonly string _runtime;
         private readonly IEnvironment _environment;
@@ -81,16 +83,18 @@ internal partial 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 _isHandlesInvocationCancelMessageCapabilityEnabled;
         private bool _isWorkerApplicationInsightsLoggingEnabled;
         private IHttpProxyService _httpProxyService;
         private Uri _httpProxyEndpoint;
         private System.Timers.Timer _timer;
         private bool _functionMetadataRequestSent = false;
+        private IOptions<ScriptJobHostOptions> _scriptHostOptions;
 
         internal GrpcWorkerChannel(
             string workerId,
             IScriptEventManager eventManager,
+            IScriptHostManager hostManager,
             RpcWorkerConfig workerConfig,
             IWorkerProcess rpcWorkerProcess,
             ILogger logger,
@@ -105,6 +109,7 @@ internal GrpcWorkerChannel(
         {
             _workerId = workerId;
             _eventManager = eventManager;
+            _scriptHostManager = hostManager;
             _workerConfig = workerConfig;
             _runtime = workerConfig.Description.Language;
             _rpcWorkerProcess = rpcWorkerProcess;
@@ -140,6 +145,10 @@ internal GrpcWorkerChannel(
 
             // Temporary switch to allow fully testing new algorithm in production
             _messageDispatcherFactory = GetProcessorFactory();
+
+            _scriptHostManager.ActiveHostChanged += HandleActiveHostChange;
+
+            LoadScriptJobHostOptions(_scriptHostManager as IServiceProvider);
         }
 
         private bool IsHttpProxyingWorker => _httpProxyEndpoint is not null;
@@ -152,6 +161,30 @@ internal GrpcWorkerChannel(
 
         public RpcWorkerConfig WorkerConfig => _workerConfig;
 
+        public IOptions<ScriptJobHostOptions> JobHostOptions => _scriptHostOptions;
+
+        private void HandleActiveHostChange(object sender, ActiveHostChangedEventArgs e)
+        {
+            if (e.NewHost is null)
+            {
+                return;
+            }
+
+            LoadScriptJobHostOptions(e.NewHost.Services);
+        }
+
+        private void LoadScriptJobHostOptions(IServiceProvider provider)
+        {
+            if (provider?.GetService(typeof(IOptions<ScriptJobHostOptions>)) is IOptions<ScriptJobHostOptions> scriptHostOptions)
+            {
+                _scriptHostOptions = scriptHostOptions;
+            }
+            else
+            {
+                _workerChannelLogger.LogDebug("Unable to resolve ScriptJobHostOptions");
+            }
+        }
+
         // Temporary switch that allows us to move between the "old" ThreadPool-only processor
         // and a "new" Channel processor (for proper ordering of messages).
         private IInvocationMessageDispatcherFactory GetProcessorFactory()
@@ -511,7 +544,7 @@ internal void ApplyCapabilities(IDictionary<string, string> capabilities, GrpcCa
             UpdateCapabilities(capabilities, strategy);
 
             _isSharedMemoryDataTransferEnabled = IsSharedMemoryDataTransferEnabled();
-            _cancelCapabilityEnabled ??= !string.IsNullOrEmpty(_workerCapabilities.GetCapabilityState(RpcWorkerConstants.HandlesInvocationCancelMessage));
+            _isHandlesInvocationCancelMessageCapabilityEnabled = !string.IsNullOrEmpty(_workerCapabilities.GetCapabilityState(RpcWorkerConstants.HandlesInvocationCancelMessage));
 
             if (!_isSharedMemoryDataTransferEnabled)
             {
@@ -816,12 +849,21 @@ internal async Task SendInvocationRequest(ScriptInvocationContext context)
                     return;
                 }
 
-                // do not send an invocation request if cancellation has been requested
                 if (context.CancellationToken.IsCancellationRequested)
                 {
-                    _workerChannelLogger.LogWarning("Cancellation has been requested. The invocation request with id '{invocationId}' is canceled and will not be sent to the worker.", invocationId);
-                    context.ResultSource.TrySetCanceled();
-                    return;
+                    _workerChannelLogger.LogDebug("Cancellation was requested prior to the invocation request ('{invocationId}') being sent to the worker.", invocationId);
+
+                    // If the worker does not support handling InvocationCancel grpc messages, or if cancellation is supported and the customer opts-out
+                    // of sending cancelled invocations to the worker, we will cancel the result source and not send the invocation to the worker.
+                    if (!_isHandlesInvocationCancelMessageCapabilityEnabled || !JobHostOptions.Value.SendCanceledInvocationsToWorker)
+                    {
+                        _workerChannelLogger.LogInformation("Cancelling invocation '{invocationId}' due to cancellation token being signaled. "
+                            + "This invocation was not sent to the worker. Read more about this here: https://aka.ms/azure-functions-cancellations", invocationId);
+
+                        // This will result in an invocation failure with a "FunctionInvocationCanceled" exception.
+                        context.ResultSource.TrySetCanceled();
+                        return;
+                    }
                 }
 
                 var invocationRequest = await context.ToRpcInvocationRequest(_workerChannelLogger, _workerCapabilities, _isSharedMemoryDataTransferEnabled, _sharedMemoryManager);
@@ -834,9 +876,10 @@ await SendStreamingMessageAsync(new StreamingMessage
                     InvocationRequest = invocationRequest
                 });
 
-                if (_cancelCapabilityEnabled != null && _cancelCapabilityEnabled.Value)
+                if (_isHandlesInvocationCancelMessageCapabilityEnabled)
                 {
-                    context.CancellationToken.Register(() => SendInvocationCancel(invocationRequest.InvocationId));
+                    var cancellationCtr = context.CancellationToken.Register(() => SendInvocationCancel(invocationRequest.InvocationId));
+                    context.Properties.Add(ScriptConstants.CancellationTokenRegistration, cancellationCtr);
                 }
 
                 if (IsHttpProxyingWorker && context.FunctionMetadata.IsHttpTriggerFunction())
@@ -1038,6 +1081,13 @@ internal async Task InvokeResponse(InvocationResponse invokeResponse)
             if (_executingInvocations.TryRemove(invokeResponse.InvocationId, out var invocation))
             {
                 var context = invocation.Context;
+
+                if (context.Properties.TryGetValue(ScriptConstants.CancellationTokenRegistration, out CancellationTokenRegistration ctr))
+                {
+                    await ctr.DisposeAsync();
+                    context.Properties.Remove(ScriptConstants.CancellationTokenRegistration);
+                }
+
                 if (invokeResponse.Result.IsInvocationSuccess(context.ResultSource, capabilityEnabled))
                 {
                     _metricsLogger.LogEvent(string.Format(MetricEventNames.WorkerInvokeSucceeded, Id));
@@ -1355,6 +1405,7 @@ protected virtual void Dispose(bool disposing)
                     _startLatencyMetric?.Dispose();
                     _workerInitTask?.TrySetCanceled();
                     _timer?.Dispose();
+                    _scriptHostManager.ActiveHostChanged -= HandleActiveHostChange;
 
                     // unlink function inputs
                     if (_inputLinks is not null)

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

@@ -21,18 +21,20 @@ public class GrpcWorkerChannelFactory : IRpcWorkerChannelFactory
         private readonly ILoggerFactory _loggerFactory = null;
         private readonly IRpcWorkerProcessFactory _rpcWorkerProcessFactory = null;
         private readonly IScriptEventManager _eventManager = null;
+        private readonly IScriptHostManager _hostManager = null;
         private readonly IEnvironment _environment = null;
         private readonly IOptionsMonitor<ScriptApplicationHostOptions> _applicationHostOptions = null;
         private readonly ISharedMemoryManager _sharedMemoryManager = null;
         private readonly IOptions<WorkerConcurrencyOptions> _workerConcurrencyOptions;
         private readonly IOptions<FunctionsHostingConfigOptions> _hostingConfigOptions;
         private readonly IHttpProxyService _httpProxyService;
 
-        public GrpcWorkerChannelFactory(IScriptEventManager eventManager, IEnvironment environment, ILoggerFactory loggerFactory,
+        public GrpcWorkerChannelFactory(IScriptEventManager eventManager, IScriptHostManager hostManager, IEnvironment environment, ILoggerFactory loggerFactory,
             IOptionsMonitor<ScriptApplicationHostOptions> applicationHostOptions, IRpcWorkerProcessFactory rpcWorkerProcessManager, ISharedMemoryManager sharedMemoryManager,
             IOptions<WorkerConcurrencyOptions> workerConcurrencyOptions, IOptions<FunctionsHostingConfigOptions> hostingConfigOptions, IHttpProxyService httpProxyService)
         {
             _eventManager = eventManager;
+            _hostManager = hostManager;
             _loggerFactory = loggerFactory;
             _rpcWorkerProcessFactory = rpcWorkerProcessManager;
             _environment = environment;
@@ -55,17 +57,18 @@ public IRpcWorkerChannel Create(string scriptRootPath, string runtime, IMetricsL
             ILogger workerLogger = _loggerFactory.CreateLogger($"Worker.LanguageWorkerChannel.{runtime}.{workerId}");
             IWorkerProcess rpcWorkerProcess = _rpcWorkerProcessFactory.Create(workerId, runtime, scriptRootPath, languageWorkerConfig);
 
-            return CreateInternal(workerId, _eventManager, languageWorkerConfig, rpcWorkerProcess, workerLogger, metricsLogger, attemptCount,
+            return CreateInternal(workerId, _eventManager, _hostManager, languageWorkerConfig, rpcWorkerProcess, workerLogger, metricsLogger, attemptCount,
                 _environment, _applicationHostOptions, _sharedMemoryManager, _workerConcurrencyOptions, _hostingConfigOptions, _httpProxyService);
         }
 
-        internal virtual IRpcWorkerChannel CreateInternal(string workerId, IScriptEventManager eventManager, RpcWorkerConfig languageWorkerConfig, IWorkerProcess rpcWorkerProcess,
+        internal virtual IRpcWorkerChannel CreateInternal(string workerId, IScriptEventManager eventManager, IScriptHostManager hostManager, RpcWorkerConfig languageWorkerConfig, IWorkerProcess rpcWorkerProcess,
             ILogger workerLogger, IMetricsLogger metricsLogger, int attemptCount, IEnvironment environment, IOptionsMonitor<ScriptApplicationHostOptions> applicationHostOptions,
             ISharedMemoryManager sharedMemoryManager, IOptions<WorkerConcurrencyOptions> workerConcurrencyOptions, IOptions<FunctionsHostingConfigOptions> hostingConfigOptions, IHttpProxyService httpProxyService)
         {
             return new GrpcWorkerChannel(
                          workerId,
                          eventManager,
+                         hostManager,
                          languageWorkerConfig,
                          rpcWorkerProcess,
                          workerLogger,

src/WebJobs.Script/Config/ConfigurationSectionNames.cs

@@ -23,5 +23,6 @@ public static class ConfigurationSectionNames
         public const string EasyAuth = "easyauth";
         public const string Retry = "retry";
         public const string SequentialJobHostRestart = JobHost + ":sequentialRestart";
+        public const string SendCanceledInvocationsToWorker = "sendCanceledInvocationsToWorker";
     }
 }

src/WebJobs.Script/Config/HostJsonFileConfigurationSource.cs

@@ -51,7 +51,7 @@ public class HostJsonFileConfigurationProvider : ConfigurationProvider
             {
                 "version", "functionTimeout", "retry", "functions", "http", "watchDirectories", "watchFiles", "queues", "serviceBus",
                 "eventHub", "singleton", "logging", "aggregator", "healthMonitor", "extensionBundle", "managedDependencies",
-                "customHandler", "httpWorker", "extensions", "concurrency"
+                "customHandler", "httpWorker", "extensions", "concurrency", ConfigurationSectionNames.SendCanceledInvocationsToWorker
             };
 
             private static readonly string[] CredentialNameFragments = new[] { "password", "pwd", "key", "secret", "token", "sas" };

src/WebJobs.Script/Config/ScriptJobHostOptions.cs

@@ -122,5 +122,12 @@ public string RootScriptPath
         /// Gets or sets a value indicating whether the ScriptHost is in standby mode.
         /// </summary>
         public bool IsStandbyConfiguration { get; set; }
+
+        /// <summary>
+        /// Gets or sets a value indicating whether the host sends cancelled invocations to the worker.
+        /// This defaults to true, meaning if cancellation is signalled we will still send the pre-cancelled
+        /// invocation to the worker.
+        /// </summary>
+        public bool SendCanceledInvocationsToWorker { get; set; } = true;
     }
 }

src/WebJobs.Script/Config/ScriptJobHostOptionsSetup.cs

@@ -51,6 +51,7 @@ public void Configure(ScriptJobHostOptions options)
                 {
                     options.FileLoggingMode = fileLoggingMode.Value;
                 }
+
                 Utility.ValidateRetryOptions(options.Retry);
             }
 

src/WebJobs.Script/ScriptConstants.cs

@@ -243,5 +243,7 @@ public static class ScriptConstants
         public static readonly string HttpProxyTask = "HttpProxyTask";
 
         public static readonly string OperationNameKey = "OperationName";
+
+        public static readonly string CancellationTokenRegistration = "CancellationTokenRegistration";
     }
 }

test/WebJobs.Script.Tests.Integration/ApplicationInsights/ApplicationInsightsTestFixture.cs

@@ -97,24 +97,24 @@ public void Dispose()
 
         private class TestGrpcWorkerChannelFactory : GrpcWorkerChannelFactory
         {
-            public TestGrpcWorkerChannelFactory(IScriptEventManager eventManager, IEnvironment environment, ILoggerFactory loggerFactory, IOptionsMonitor<ScriptApplicationHostOptions> applicationHostOptions, IRpcWorkerProcessFactory rpcWorkerProcessManager, ISharedMemoryManager sharedMemoryManager, IOptions<WorkerConcurrencyOptions> workerConcurrencyOptions, IOptions<FunctionsHostingConfigOptions> hostingConfigOptions, IHttpProxyService httpProxyService)
-                : base(eventManager, environment, loggerFactory, applicationHostOptions, rpcWorkerProcessManager, sharedMemoryManager, workerConcurrencyOptions, hostingConfigOptions, httpProxyService)
+            public TestGrpcWorkerChannelFactory(IScriptEventManager eventManager, IScriptHostManager hostManager, IEnvironment environment, ILoggerFactory loggerFactory, IOptionsMonitor<ScriptApplicationHostOptions> applicationHostOptions, IRpcWorkerProcessFactory rpcWorkerProcessManager, ISharedMemoryManager sharedMemoryManager, IOptions<WorkerConcurrencyOptions> workerConcurrencyOptions, IOptions<FunctionsHostingConfigOptions> hostingConfigOptions, IHttpProxyService httpProxyService)
+                : base(eventManager, hostManager, environment, loggerFactory, applicationHostOptions, rpcWorkerProcessManager, sharedMemoryManager, workerConcurrencyOptions, hostingConfigOptions, httpProxyService)
             {
             }
 
-            internal override IRpcWorkerChannel CreateInternal(string workerId, IScriptEventManager eventManager, RpcWorkerConfig languageWorkerConfig, IWorkerProcess rpcWorkerProcess,
+            internal override IRpcWorkerChannel CreateInternal(string workerId, IScriptEventManager eventManager, IScriptHostManager hostManager, RpcWorkerConfig languageWorkerConfig, IWorkerProcess rpcWorkerProcess,
             ILogger workerLogger, IMetricsLogger metricsLogger, int attemptCount, IEnvironment environment, IOptionsMonitor<ScriptApplicationHostOptions> applicationHostOptions,
             ISharedMemoryManager sharedMemoryManager, IOptions<WorkerConcurrencyOptions> workerConcurrencyOptions, IOptions<FunctionsHostingConfigOptions> hostingConfigOptions, IHttpProxyService httpProxyService)
             {
-                return new TestGrpcWorkerChannel(workerId, eventManager, languageWorkerConfig, rpcWorkerProcess, workerLogger, metricsLogger,
+                return new TestGrpcWorkerChannel(workerId, eventManager, hostManager, languageWorkerConfig, rpcWorkerProcess, workerLogger, metricsLogger,
                     attemptCount, environment, applicationHostOptions, sharedMemoryManager, workerConcurrencyOptions, hostingConfigOptions, httpProxyService);
             }
 
             private class TestGrpcWorkerChannel : GrpcWorkerChannel
             {
-                internal TestGrpcWorkerChannel(string workerId, IScriptEventManager eventManager, RpcWorkerConfig workerConfig, IWorkerProcess rpcWorkerProcess, ILogger logger, IMetricsLogger metricsLogger, int attemptCount, IEnvironment environment, IOptionsMonitor<ScriptApplicationHostOptions> applicationHostOptions, ISharedMemoryManager sharedMemoryManager, IOptions<WorkerConcurrencyOptions> workerConcurrencyOptions, IOptions<FunctionsHostingConfigOptions> hostingConfigOptions, IHttpProxyService httpProxyService)
-                    : base(workerId, eventManager, workerConfig, rpcWorkerProcess, logger, metricsLogger, attemptCount, environment, applicationHostOptions, sharedMemoryManager, workerConcurrencyOptions, hostingConfigOptions, httpProxyService)
-                {   
+                internal TestGrpcWorkerChannel(string workerId, IScriptEventManager eventManager, IScriptHostManager hostManager, RpcWorkerConfig workerConfig, IWorkerProcess rpcWorkerProcess, ILogger logger, IMetricsLogger metricsLogger, int attemptCount, IEnvironment environment, IOptionsMonitor<ScriptApplicationHostOptions> applicationHostOptions, ISharedMemoryManager sharedMemoryManager, IOptions<WorkerConcurrencyOptions> workerConcurrencyOptions, IOptions<FunctionsHostingConfigOptions> hostingConfigOptions, IHttpProxyService httpProxyService)
+                    : base(workerId, eventManager, hostManager, workerConfig, rpcWorkerProcess, logger, metricsLogger, attemptCount, environment, applicationHostOptions, sharedMemoryManager, workerConcurrencyOptions, hostingConfigOptions, httpProxyService)
+                {
                 }
 
                 internal override void UpdateCapabilities(IDictionary<string, string> fields, GrpcCapabilitiesUpdateStrategy strategy)