adding ability to manipulate capabilities (#1183)

Author: brettsam

samples/FunctionApp/Program.cs

@@ -1,9 +1,6 @@
 // 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.Diagnostics;
-using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Azure.Functions.Worker;
 using Microsoft.Extensions.DependencyInjection;

src/DotNetWorker.Core/Context/Features/IDictionaryExtensions.cs

@@ -8,49 +8,59 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.ComponentModel;
 
-namespace Microsoft.Azure.Functions.Worker
+internal static class IDictionaryExtensions
 {
-    internal static class IDictionaryExtensions
+    internal static bool TryAdd<TKey, TValue>(this IDictionary<TKey, TValue> dictionary, TKey key, TValue value)
     {
-        internal static bool TryAdd<TKey, TValue>(this IDictionary<TKey, TValue> dictionary, TKey key, TValue value)
+        if (key == null)
         {
-            if (key == null)
-            {
-                throw new ArgumentNullException(nameof(key));
-            }
-
-            if (dictionary.ContainsKey(key))
-            {
-                return false;
-            }
-
-            dictionary.Add(key, value);
-            return true;
+            throw new ArgumentNullException(nameof(key));
         }
+
+        if (dictionary.ContainsKey(key))
+        {
+            return false;
+        }
+
+        dictionary.Add(key, value);
+        return true;
     }
+}
 
-    internal static class ConcurrentDictionaryExtensions
+internal static class ConcurrentDictionaryExtensions
+{
+    public static TValue GetOrAdd<TKey, TValue, TArg>(this ConcurrentDictionary<TKey, TValue> dictionary, TKey key, Func<TKey, TArg, TValue> valueFactory, TArg factoryArgument)
     {
-        public static TValue GetOrAdd<TKey, TValue, TArg>(this ConcurrentDictionary<TKey, TValue> dictionary, TKey key, Func<TKey, TArg, TValue> valueFactory, TArg factoryArgument)
+        if (dictionary == null)
         {
-            if (dictionary == null)
-            {
-                throw new ArgumentNullException(nameof(dictionary));
-            }
-
-            if (key == null)
-            {
-                throw new ArgumentNullException(nameof(key));
-            }
-
-            if (valueFactory == null)
-            {
-                throw new ArgumentNullException(nameof(valueFactory));
-            }
-
-            return dictionary.GetOrAdd(key, k => valueFactory(k, factoryArgument));
+            throw new ArgumentNullException(nameof(dictionary));
         }
+
+        if (key == null)
+        {
+            throw new ArgumentNullException(nameof(key));
+        }
+
+        if (valueFactory == null)
+        {
+            throw new ArgumentNullException(nameof(valueFactory));
+        }
+
+        return dictionary.GetOrAdd(key, k => valueFactory(k, factoryArgument));
     }
 }
+
+internal static class KeyValuePairExtensions
+{
+    // Based on https://source.dot.net/#System.Private.CoreLib/src/libraries/System.Private.CoreLib/src/System/Collections/Generic/KeyValuePair.cs,aa57b8e336bf7f59    
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    public static void Deconstruct<TKey, TValue>(this KeyValuePair<TKey, TValue> pair, out TKey key, out TValue value)
+    {
+        key = pair.Key;
+        value = pair.Value;
+    }
+}
+
 #endif

src/DotNetWorker.Core/Hosting/WorkerOptions.cs

@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
+using System.Collections.Generic;
 using System.Text.Json;
 using Azure.Core.Serialization;
 
@@ -22,19 +23,56 @@ public class WorkerOptions
         /// </summary>
         public InputConverterCollection InputConverters { get; } = new InputConverterCollection();
 
+        /// <summary>
+        /// Gets the optional worker capabilities.
+        /// </summary>
+        public IDictionary<string, string> Capabilities { get; } = new Dictionary<string, string>()
+        {
+            // Enable these by default, although they are not strictly required and can be removed
+            { "HandlesWorkerTerminateMessage", bool.TrueString },
+            { "HandlesInvocationCancelMessage", bool.TrueString }
+        };
+
         /// <summary>
         /// Gets and sets the flag for opting in to unwrapping user-code-thrown
         /// exceptions when they are surfaced to the Host. 
         /// </summary>
-        public bool EnableUserCodeException { get; set; } = false;
-        
+        public bool EnableUserCodeException
+        {
+            get => GetBoolCapability(nameof(EnableUserCodeException));
+            set => SetBoolCapability(nameof(EnableUserCodeException), value);
+        }
+
         /// <summary>
         /// Gets or sets a value that determines if empty entries should be included in the function trigger message payload.
         /// For example, if a set of entries were sent to a messaging service such as Service Bus or Event Hub and your function
         /// app has a Service bus trigger or Event hub trigger, only the non-empty entries from the payload will be sent to the
         /// function code as trigger data when this setting value is <see langword="false"/>. When it is <see langword="true"/>,
         /// All entries will be sent to the function code as it is. Default value for this setting is <see langword="false"/>.
         /// </summary>
-        public bool IncludeEmptyEntriesInMessagePayload { get; set; }
+        public bool IncludeEmptyEntriesInMessagePayload
+        {
+            get => GetBoolCapability(nameof(IncludeEmptyEntriesInMessagePayload));
+            set => SetBoolCapability(nameof(IncludeEmptyEntriesInMessagePayload), value);
+        }
+
+        private bool GetBoolCapability(string name)
+        {
+            return Capabilities.TryGetValue(name, out string? value) && bool.TryParse(value, out bool b) && b;
+        }
+
+        // For false values, the host does not expect the capability to exist; there are some cases where this
+        // will be interpreted as "true" just because the key is there.
+        private void SetBoolCapability(string name, bool value)
+        {
+            if (value)
+            {
+                Capabilities[name] = bool.TrueString;
+            }
+            else
+            {
+                Capabilities.Remove(name);
+            }
+        }
     }
 }

src/DotNetWorker.Grpc/GrpcWorker.cs

@@ -163,24 +163,20 @@ internal static WorkerInitResponse WorkerInitRequestHandler(WorkerInitRequest re
 
             response.WorkerMetadata.CustomProperties.Add("Worker.Grpc.Version", typeof(GrpcWorker).Assembly.GetName().Version?.ToString());
 
-            response.Capabilities.Add("RpcHttpBodyOnly", bool.TrueString);
-            response.Capabilities.Add("RawHttpBodyBytes", bool.TrueString);
-            response.Capabilities.Add("RpcHttpTriggerMetadataRemoved", bool.TrueString);
-            response.Capabilities.Add("UseNullableValueDictionaryForHttp", bool.TrueString);
-            response.Capabilities.Add("TypedDataCollection", bool.TrueString);
-            response.Capabilities.Add("WorkerStatus", bool.TrueString);
-            response.Capabilities.Add("HandlesWorkerTerminateMessage", bool.TrueString);
-            response.Capabilities.Add("HandlesInvocationCancelMessage", bool.TrueString);
-
-            if (workerOptions.EnableUserCodeException)
+            // Add additional capabilities defined by WorkerOptions
+            foreach ((string key, string value) in workerOptions.Capabilities)
             {
-                response.Capabilities.Add("EnableUserCodeException", bool.TrueString);
-            }
-            if (workerOptions.IncludeEmptyEntriesInMessagePayload)
-            {
-                response.Capabilities.Add("IncludeEmptyEntriesInMessagePayload", bool.TrueString);
+                response.Capabilities[key] = value;
             }
 
+            // Add required capabilities; these cannot be modified and will override anything from WorkerOptions
+            response.Capabilities["RpcHttpBodyOnly"] = bool.TrueString;
+            response.Capabilities["RawHttpBodyBytes"] = bool.TrueString;
+            response.Capabilities["RpcHttpTriggerMetadataRemoved"] = bool.TrueString;
+            response.Capabilities["UseNullableValueDictionaryForHttp"] = bool.TrueString;
+            response.Capabilities["TypedDataCollection"] = bool.TrueString;
+            response.Capabilities["WorkerStatus"] = bool.TrueString;
+
             return response;
         }
 

test/DotNetWorkerTests/GrpcWorkerTests.cs

@@ -16,7 +16,10 @@
 using Microsoft.Azure.Functions.Worker.Handlers;
 using Microsoft.Azure.Functions.Worker.Invocation;
 using Microsoft.Azure.Functions.Worker.OutputBindings;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Options;
 using Moq;
 using Xunit;
 
@@ -41,7 +44,8 @@ public GrpcWorkerTests()
 
             _mockApplication
                 .Setup(m => m.CreateContext(It.IsAny<IInvocationFeatures>(), It.IsAny<CancellationToken>()))
-                .Returns((IInvocationFeatures f, CancellationToken ct) => {
+                .Returns((IInvocationFeatures f, CancellationToken ct) =>
+                {
                     _context = new TestFunctionContext(f, ct);
                     return _context;
                 });
@@ -184,6 +188,7 @@ public void InitRequest_ReturnsExpectedCapabilities_BasedOnWorkerOptions(
             string expectedCapabilityValue = null)
         {
             var workerOptions = new WorkerOptions();
+
             // Update boolean property values of workerOption based on test input parameters.
             workerOptions.GetType().GetProperty(booleanPropertyName)?.SetValue(workerOptions, booleanPropertyValue);
 
@@ -210,6 +215,36 @@ public void InitRequest_ReturnsExpectedCapabilities_BasedOnWorkerOptions(
             }
         }
 
+        [Fact]
+        public void WorkerOptions_CanChangeOptionalCapabilities()
+        {
+            var host = new HostBuilder()
+                .ConfigureFunctionsWorkerDefaults((WorkerOptions options) =>
+                {
+                    options.Capabilities.Remove("HandlesWorkerTerminateMessage");
+                    options.Capabilities.Add("SomeNewCapability", bool.TrueString);
+                }).Build();
+
+            var workerOptions = host.Services.GetService<IOptions<WorkerOptions>>().Value;
+            var response = GrpcWorker.WorkerInitRequestHandler(new(), workerOptions);
+
+            void AssertKeyAndValue(KeyValuePair<string, string> kvp, string expectedKey, string expectedValue)
+            {
+                Assert.Same(expectedKey, kvp.Key);
+                Assert.Same(expectedValue, kvp.Value);
+            }
+
+            Assert.Collection(response.Capabilities.OrderBy(p => p.Key),
+                c => AssertKeyAndValue(c, "HandlesInvocationCancelMessage", bool.TrueString),
+                c => AssertKeyAndValue(c, "RawHttpBodyBytes", bool.TrueString),
+                c => AssertKeyAndValue(c, "RpcHttpBodyOnly", bool.TrueString),
+                c => AssertKeyAndValue(c, "RpcHttpTriggerMetadataRemoved", bool.TrueString),
+                c => AssertKeyAndValue(c, "SomeNewCapability", bool.TrueString),
+                c => AssertKeyAndValue(c, "TypedDataCollection", bool.TrueString),
+                c => AssertKeyAndValue(c, "UseNullableValueDictionaryForHttp", bool.TrueString),
+                c => AssertKeyAndValue(c, "WorkerStatus", bool.TrueString));
+        }
+
         [Fact]
         public async Task Invoke_ReturnsSuccess()
         {