Remove task allocation in unary call handler

perf/Grpc.AspNetCore.Microbenchmarks/DefaultCoreConfig.cs

@@ -44,7 +44,7 @@ public DefaultCoreConfig()
             Add(JitOptimizationsValidator.FailOnError);
 
             Add(Job.Core
-                .With(CsProjCoreToolchain.From(new NetCoreAppSettings("netcoreapp3.0", null, ".NET Core 3.0")))
+                .With(CsProjCoreToolchain.From(new NetCoreAppSettings("netcoreapp3.1", null, ".NET Core 3.1")))
                 .With(new GcMode { Server = true })
                 .With(RunStrategy.Throughput));
         }

src/Shared/Server/UnaryServerMethodInvoker.cs

@@ -16,6 +16,8 @@
 
 #endregion
 
+using System;
+using System.Runtime.ExceptionServices;
 using System.Threading.Tasks;
 using Grpc.AspNetCore.Server;
 using Grpc.AspNetCore.Server.Model;
@@ -86,33 +88,89 @@ private async Task<TResponse> ResolvedInterceptorInvoker(TRequest resolvedReques
         /// <param name="request">The <typeparamref name="TRequest"/> message.</param>
         /// <returns>A <see cref="Task{TResponse}"/> that represents the asynchronous method. The <see cref="Task{TResponse}.Result"/>
         /// property returns the <typeparamref name="TResponse"/> message.</returns>
-        public async Task<TResponse> Invoke(HttpContext httpContext, ServerCallContext serverCallContext, TRequest request)
+        public Task<TResponse> Invoke(HttpContext httpContext, ServerCallContext serverCallContext, TRequest request)
         {
             if (_pipelineInvoker == null)
             {
                 GrpcActivatorHandle<TService> serviceHandle = default;
+                Task<TResponse>? invokerTask = null;
                 try
                 {
                     serviceHandle = ServiceActivator.Create(httpContext.RequestServices);
-                    return await _invoker(
+                    invokerTask = _invoker(
                         serviceHandle.Instance,
                         request,
                         serverCallContext);
                 }
-                finally
+                catch (Exception ex)
                 {
+                    // Invoker calls user code. User code may throw an exception instead
+                    // of a faulted task. We need to catch the exception, ensure cleanup
+                    // runs and convert exception into a faulted task.
                     if (serviceHandle.Instance != null)
                     {
-                        await ServiceActivator.ReleaseAsync(serviceHandle);
+                        var releaseTask = ServiceActivator.ReleaseAsync(serviceHandle);
+                        if (!releaseTask.IsCompletedSuccessfully)
+                        {
+                            // Capture the current exception state so we can rethrow it after awaiting
+                            // with the same stack trace.
+                            var exceptionDispatchInfo = ExceptionDispatchInfo.Capture(ex);
+                            return AwaitServiceReleaseAndThrow(releaseTask, exceptionDispatchInfo);
+                        }
                     }
+
+                    return Task.FromException<TResponse>(ex);
+                }
+
+                if (invokerTask.IsCompletedSuccessfully && serviceHandle.Instance != null)
+                {
+                    var releaseTask = ServiceActivator.ReleaseAsync(serviceHandle);
+                    if (!releaseTask.IsCompletedSuccessfully)
+                    {
+                        return AwaitServiceReleaseAndReturn(invokerTask.Result, serviceHandle);
+                    }
+
+                    return invokerTask;
                 }
+
+                return AwaitInvoker(invokerTask, serviceHandle);
             }
             else
             {
-                return await _pipelineInvoker(
+                return _pipelineInvoker(
                     request,
                     serverCallContext);
             }
         }
+
+        private async Task<TResponse> AwaitInvoker(Task<TResponse> invokerTask, GrpcActivatorHandle<TService> serviceHandle)
+        {
+            try
+            {
+                return await invokerTask;
+            }
+            finally
+            {
+                if (serviceHandle.Instance != null)
+                {
+                    await ServiceActivator.ReleaseAsync(serviceHandle);
+                }
+            }
+        }
+
+        private async Task<TResponse> AwaitServiceReleaseAndThrow(ValueTask releaseTask, ExceptionDispatchInfo ex)
+        {
+            await releaseTask;
+            ex.Throw();
+
+            // Should never reach here
+            return null;
+        }
+
+        private async Task<TResponse> AwaitServiceReleaseAndReturn(TResponse invokerResult, GrpcActivatorHandle<TService> serviceHandle)
+        {
+            await ServiceActivator.ReleaseAsync(serviceHandle);
+            return invokerResult;
+        }
     }
 }

test/Grpc.AspNetCore.Server.Tests/TestObjects/TestGrpcServiceActivator.cs

@@ -23,13 +23,16 @@ namespace Grpc.AspNetCore.Server.Tests.TestObjects
 {
     internal class TestGrpcServiceActivator<TGrpcService> : IGrpcServiceActivator<TGrpcService> where TGrpcService : class, new()
     {
+        public bool Released { get; private set; }
+
         public GrpcActivatorHandle<TGrpcService> Create(IServiceProvider serviceProvider)
         {
             return new GrpcActivatorHandle<TGrpcService>(new TGrpcService(), false, null);
         }
 
         public ValueTask ReleaseAsync(GrpcActivatorHandle<TGrpcService> service)
         {
+            Released = true;
             return default;
         }
     }

test/Grpc.AspNetCore.Server.Tests/UnaryServerCallHandlerTests.cs

@@ -0,0 +1,163 @@
+#region Copyright notice and license
+
+// Copyright 2019 The gRPC Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#endregion
+
+using System;
+using System.Threading.Tasks;
+using Grpc.AspNetCore.Server.Tests.TestObjects;
+using Grpc.Core;
+using Grpc.Shared.Server;
+using Grpc.Tests.Shared;
+using NUnit.Framework;
+
+namespace Grpc.AspNetCore.Server.Tests
+{
+    [TestFixture]
+    public class UnaryServerCallHandlerTests
+    {
+        private static readonly Marshaller<TestMessage> _marshaller = new Marshaller<TestMessage>((message, context) => { context.Complete(Array.Empty<byte>()); }, context => new TestMessage());
+
+        [Test]
+        public void Invoke_ThrowException_ReleaseCalledAndErrorThrown()
+        {
+            // Arrange
+            var serviceActivator = new TestGrpcServiceActivator<TestService>();
+            var ex = new Exception("Exception!");
+            var invoker = new UnaryServerMethodInvoker<TestService, TestMessage, TestMessage>(
+                (service, reader, context) => throw ex,
+                new Method<TestMessage, TestMessage>(MethodType.Unary, "test", "test", _marshaller, _marshaller),
+                HttpContextServerCallContextHelper.CreateMethodOptions(),
+                serviceActivator);
+            var httpContext = HttpContextHelpers.CreateContext();
+
+            // Act
+            var task = invoker.Invoke(httpContext, HttpContextServerCallContextHelper.CreateServerCallContext(), new TestMessage());
+
+            // Assert
+            Assert.True(serviceActivator.Released);
+            Assert.True(task.IsFaulted);
+            Assert.AreEqual(ex, task.Exception!.InnerException);
+        }
+
+        [Test]
+        public async Task Invoke_ThrowExceptionAwaitedRelease_ReleaseCalledAndErrorThrown()
+        {
+            // Arrange
+            var releaseTcs = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var serviceActivator = new TcsGrpcServiceActivator<TestService>(releaseTcs);
+            var thrownException = new Exception("Exception!");
+            var invoker = new UnaryServerMethodInvoker<TestService, TestMessage, TestMessage>(
+                (service, reader, context) => throw thrownException,
+                new Method<TestMessage, TestMessage>(MethodType.Unary, "test", "test", _marshaller, _marshaller),
+                HttpContextServerCallContextHelper.CreateMethodOptions(),
+                serviceActivator);
+            var httpContext = HttpContextHelpers.CreateContext();
+
+            // Act
+            var task = invoker.Invoke(httpContext, HttpContextServerCallContextHelper.CreateServerCallContext(), new TestMessage());
+            Assert.False(task.IsCompleted);
+
+            //await Task.Delay(1000);
+            releaseTcs.SetResult(null);
+
+            try
+            {
+                await task;
+                Assert.Fail();
+            }
+            catch (Exception ex)
+            {
+                // Assert
+                Assert.True(serviceActivator.Released);
+                Assert.AreEqual(thrownException, ex);
+            }
+        }
+
+        [Test]
+        public async Task Invoke_SuccessAwaitedRelease_ReleaseCalled()
+        {
+            // Arrange
+            var releaseTcs = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var serviceActivator = new TcsGrpcServiceActivator<TestService>(releaseTcs);
+            var invoker = new UnaryServerMethodInvoker<TestService, TestMessage, TestMessage>(
+                (service, reader, context) => Task.FromResult(new TestMessage()),
+                new Method<TestMessage, TestMessage>(MethodType.Unary, "test", "test", _marshaller, _marshaller),
+                HttpContextServerCallContextHelper.CreateMethodOptions(),
+                serviceActivator);
+            var httpContext = HttpContextHelpers.CreateContext();
+
+            // Act
+            var task = invoker.Invoke(httpContext, HttpContextServerCallContextHelper.CreateServerCallContext(), new TestMessage());
+            Assert.False(task.IsCompleted);
+
+            releaseTcs.SetResult(null);
+            await task;
+
+            // Assert
+            Assert.True(serviceActivator.Released);
+        }
+
+        [Test]
+        public async Task Invoke_AwaitedSuccess_ReleaseCalled()
+        {
+            // Arrange
+            var methodTcs = new TaskCompletionSource<TestMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var methodResult = new TestMessage();
+            var serviceActivator = new TestGrpcServiceActivator<TestService>();
+            var invoker = new UnaryServerMethodInvoker<TestService, TestMessage, TestMessage>(
+                (service, reader, context) => methodTcs.Task,
+                new Method<TestMessage, TestMessage>(MethodType.Unary, "test", "test", _marshaller, _marshaller),
+                HttpContextServerCallContextHelper.CreateMethodOptions(),
+                serviceActivator);
+            var httpContext = HttpContextHelpers.CreateContext();
+
+            // Act
+            var task = invoker.Invoke(httpContext, HttpContextServerCallContextHelper.CreateServerCallContext(), new TestMessage());
+            Assert.False(task.IsCompleted);
+
+            methodTcs.SetResult(methodResult);
+            var awaitedResult = await task;
+
+            // Assert
+            Assert.AreEqual(methodResult, awaitedResult);
+            Assert.True(serviceActivator.Released);
+        }
+
+        private class TcsGrpcServiceActivator<TGrpcService> : IGrpcServiceActivator<TGrpcService> where TGrpcService : class, new()
+        {
+            private readonly TaskCompletionSource<object?> _tcs;
+
+            public bool Released { get; private set; }
+
+            public TcsGrpcServiceActivator(TaskCompletionSource<object?> tcs)
+            {
+                _tcs = tcs;
+            }
+
+            public GrpcActivatorHandle<TGrpcService> Create(IServiceProvider serviceProvider)
+            {
+                return new GrpcActivatorHandle<TGrpcService>(new TGrpcService(), false, null);
+            }
+
+            public ValueTask ReleaseAsync(GrpcActivatorHandle<TGrpcService> service)
+            {
+                Released = true;
+                return new ValueTask(_tcs.Task);
+            }
+        }
+    }
+}