Fix error log in custom handler streaming for $return scenarios (#11262)

Author: satvu

release_notes.md

@@ -8,4 +8,5 @@
 - Moving to version 1.5.7 of Microsoft.Azure.AppService.Middleware.Functions (https://github.com/Azure/azure-functions-host/pull/11231)
 - Refactor code to move the logic to search for WorkerConfigs to a default worker configuration resolver (#11219)
 - Update Node.js Worker Version to [3.11.0](https://github.com/Azure/azure-functions-nodejs-worker/releases/tag/v3.11.0)
-- Adding restart reason to the logs (#11191)
+- Adding restart reason to the logs (#11191)
+- Fix custom handler streaming bug where `$return` output binding scenarios result in error logs (#11262)

src/WebJobs.Script/Description/Workers/WorkerFunctionInvoker.cs

@@ -1,8 +1,9 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// 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.Collections.Immutable;
 using System.Collections.ObjectModel;
 using System.Diagnostics;
 using System.IO;
@@ -211,6 +212,11 @@ private CancellationToken HandleCancellationTokenParameter(object input)
 
         private void HandleReturnParameter(ScriptInvocationResult result)
         {
+            if (result.Outputs is IImmutableDictionary<string, object> immutableOutputs)
+            {
+                return;
+            }
+
             result.Outputs[ScriptConstants.SystemReturnParameterBindingName] = result.Return;
         }
 

src/WebJobs.Script/Workers/Http/DefaultHttpWorkerService.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
@@ -101,6 +101,13 @@ internal async Task ProxyInvocationRequest(ScriptInvocationContext scriptInvocat
                 _httpProxyService.StartForwarding(scriptInvocationContext, _destinationPrefix);
 
                 await _httpProxyService.EnsureSuccessfulForwardingAsync(scriptInvocationContext); // this will throw if forwarding is unsuccessful
+
+                // this needs to be removed to avoid downstream errors since the HttpResponse is already returned by proxying
+                if (scriptInvocationContext.BindingData.Any(p => string.Equals(ScriptConstants.SystemReturnParameterBindingName, p.Key, StringComparison.OrdinalIgnoreCase)))
+                {
+                    scriptInvocationContext.BindingData[ScriptConstants.SystemReturnParameterBindingName] = null;
+                }
+
                 scriptInvocationContext.ResultSource.SetResult(ScriptInvocationResult.Success);
             }
             catch (Exception exc)

test/WebJobs.Script.Tests/Description/Worker/WorkerFunctionInvokerTests.cs

@@ -1,15 +1,13 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// 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.ObjectModel;
-using System.Diagnostics.Metrics;
 using System.Threading.Tasks;
 using Microsoft.Azure.WebJobs.Script.Binding;
 using Microsoft.Azure.WebJobs.Script.Description;
 using Microsoft.Azure.WebJobs.Script.Metrics;
 using Microsoft.Azure.WebJobs.Script.Workers;
-using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging.Abstractions;
@@ -50,7 +48,9 @@ public WorkerFunctionInvokerTests()
             var sc = host.GetScriptHost();
 
             FunctionMetadata metaData = new FunctionMetadata();
-            _testFunctionInvoker = new TestWorkerFunctionInvoker(sc, null, metaData, NullLoggerFactory.Instance, null, new Collection<FunctionBinding>(),
+            BindingMetadata bindingMetadata = new BindingMetadata();
+            bindingMetadata.Name = "TestName";
+            _testFunctionInvoker = new TestWorkerFunctionInvoker(sc, bindingMetadata, metaData, NullLoggerFactory.Instance, new Collection<FunctionBinding>(), new Collection<FunctionBinding>(),
                 _mockFunctionInvocationDispatcher.Object, _applicationLifetime.Object, TimeSpan.FromSeconds(5));
         }
 
@@ -104,5 +104,44 @@ public async Task InvokeInitialized_DoesNotCallShutdown()
             }
             _applicationLifetime.Verify(a => a.StopApplication(), Times.Never);
         }
+
+        [Fact]
+        public async Task InvokeCore_DoesNotAddReturn_WhenOutputsIsImmutableDictionary()
+        {
+            // Arrange
+            var mockBinder = new Mock<Binder>();
+            var immutableOutputs = System.Collections.Immutable.ImmutableDictionary<string, object>.Empty;
+            var invocationResult = new ScriptInvocationResult
+            {
+                Outputs = immutableOutputs,
+            };
+
+            // Setup the dispatcher to complete the invocation with our result
+            _mockFunctionInvocationDispatcher
+                .Setup(d => d.InvokeAsync(It.IsAny<ScriptInvocationContext>()))
+                .Callback<ScriptInvocationContext>(ctx =>
+                {
+                    ctx.ResultSource.SetResult(invocationResult);
+                })
+                .Returns(Task.CompletedTask);
+
+            _mockFunctionInvocationDispatcher.Setup(a => a.State).Returns(FunctionInvocationDispatcherState.Initialized);
+
+            // Act
+            var result = await _testFunctionInvoker.InvokeCore(
+                new object[] { null, null, null, null, default(System.Threading.CancellationToken) },
+                new FunctionInvocationContext
+                {
+                    Binder = mockBinder.Object,
+                    ExecutionContext = new ExecutionContext
+                    {
+                        InvocationId = Guid.NewGuid()
+                    }
+                });
+
+            // Assert
+            Assert.IsType<System.Collections.Immutable.ImmutableDictionary<string, object>>(invocationResult.Outputs);
+            Assert.False(invocationResult.Outputs.ContainsKey("$return"));
+        }
     }
 }