Merge pull request #793 from Cysharp/feature/ReduceAllocationStreamin…

…gHubClient2 Reduce allocation on StreamingHubClient method calls
Cysharp · Jun 21, 2024 · 1a04afd · 1a04afd
2 parents 1192202 + f732e56
commit 1a04afd
Show file tree

Hide file tree

Showing 61 changed files with 768 additions and 569 deletions.
diff --git a/perf/BenchmarkApp/PerformanceTest.Client/Program.cs b/perf/BenchmarkApp/PerformanceTest.Client/Program.cs
@@ -143,7 +143,9 @@ async Task<PerformanceResult> RunScenarioAsync(ScenarioType scenario, ScenarioCo
         ScenarioType.UnaryLargePayload32K => () => new UnaryLargePayload32KScenario(),
         ScenarioType.UnaryLargePayload64K => () => new UnaryLargePayload64KScenario(),
         ScenarioType.StreamingHub => () => new StreamingHubScenario(),
+        ScenarioType.StreamingHubValueTask => () => new StreamingHubValueTaskScenario(),
         ScenarioType.StreamingHubComplex => () => new StreamingHubComplexScenario(),
+        ScenarioType.StreamingHubComplexValueTask => () => new StreamingHubComplexValueTaskScenario(),
         ScenarioType.StreamingHubLargePayload1K => () => new StreamingHubLargePayload1KScenario(),
         ScenarioType.StreamingHubLargePayload2K => () => new StreamingHubLargePayload2KScenario(),
         ScenarioType.StreamingHubLargePayload4K => () => new StreamingHubLargePayload4KScenario(),

diff --git a/perf/BenchmarkApp/PerformanceTest.Client/Properties/launchSettings.json b/perf/BenchmarkApp/PerformanceTest.Client/Properties/launchSettings.json
@@ -2,7 +2,7 @@
   "profiles": {
     "PerformanceTest.Client": {
       "commandName": "Project",
-      "commandLineArgs": "-u http://localhost:5000 -s streaminghub --channels 1 --streams 1"
+      "commandLineArgs": "-u http://localhost:5000 -s StreamingHubValueTask --channels 1 --streams 1"
     }
   }
-}
+}
diff --git a/perf/BenchmarkApp/PerformanceTest.Client/ScenarioType.cs b/perf/BenchmarkApp/PerformanceTest.Client/ScenarioType.cs
@@ -12,7 +12,9 @@ public enum ScenarioType
     UnaryLargePayload64K,
 
     StreamingHub,
+    StreamingHubValueTask,
     StreamingHubComplex,
+    StreamingHubComplexValueTask,
     StreamingHubLargePayload1K,
     StreamingHubLargePayload2K,
     StreamingHubLargePayload4K,

diff --git a/perf/BenchmarkApp/PerformanceTest.Client/StreamingHubScenario.cs b/perf/BenchmarkApp/PerformanceTest.Client/StreamingHubScenario.cs
@@ -21,6 +21,25 @@ public async ValueTask RunAsync(PerformanceTestRunningContext ctx, CancellationT
     }
 }
 
+public class StreamingHubValueTaskScenario : IScenario, IPerfTestHubReceiver
+{
+    IPerfTestHub client = default!;
+
+    public async ValueTask PrepareAsync(GrpcChannel channel)
+    {
+        this.client = await StreamingHubClient.ConnectAsync<IPerfTestHub, IPerfTestHubReceiver>(channel, this);
+    }
+
+    public async ValueTask RunAsync(PerformanceTestRunningContext ctx, CancellationToken cancellationToken)
+    {
+        while (!cancellationToken.IsCancellationRequested)
+        {
+            await client.CallMethodValueTaskAsync("FooBarBaz🚀こんにちは世界", 123, 4567, 891011);
+            ctx.Increment();
+        }
+    }
+}
+
 public class StreamingHubComplexScenario : IScenario, IPerfTestHubReceiver
 {
     IPerfTestHub client = default!;
@@ -39,3 +58,22 @@ public async ValueTask RunAsync(PerformanceTestRunningContext ctx, CancellationT
         }
     }
 }
+
+public class StreamingHubComplexValueTaskScenario : IScenario, IPerfTestHubReceiver
+{
+    IPerfTestHub client = default!;
+
+    public async ValueTask PrepareAsync(GrpcChannel channel)
+    {
+        this.client = await StreamingHubClient.ConnectAsync<IPerfTestHub, IPerfTestHubReceiver>(channel, this);
+    }
+
+    public async ValueTask RunAsync(PerformanceTestRunningContext ctx, CancellationToken cancellationToken)
+    {
+        while (!cancellationToken.IsCancellationRequested)
+        {
+            await client.CallMethodComplexValueTaskAsync("FooBarBaz🚀こんにちは世界", 123, 4567, 891011);
+            ctx.Increment();
+        }
+    }
+}
diff --git a/perf/BenchmarkApp/PerformanceTest.Server/PerfTestHub.cs b/perf/BenchmarkApp/PerformanceTest.Server/PerfTestHub.cs
@@ -10,11 +10,21 @@ public Task<int> CallMethodAsync(string arg1, int arg2, int arg3, int arg4)
         return Task.FromResult(0);
     }
 
+    public ValueTask<int> CallMethodValueTaskAsync(string arg1, int arg2, int arg3, int arg4)
+    {
+        return ValueTask.FromResult(0);
+    }
+
     public Task<ComplexResponse> CallMethodComplexAsync(string arg1, int arg2, int arg3, int arg4)
     {
         return Task.FromResult(ComplexResponse.Cached);
     }
 
+    public ValueTask<ComplexResponse> CallMethodComplexValueTaskAsync(string arg1, int arg2, int arg3, int arg4)
+    {
+        return ValueTask.FromResult(ComplexResponse.Cached);
+    }
+
     public Task<(int StatusCode, byte[] Data)> CallMethodLargePayloadAsync(string arg1, int arg2, int arg3, int arg4, byte[] arg5)
     {
         return Task.FromResult((123, arg5));

diff --git a/perf/BenchmarkApp/PerformanceTest.Shared/IPerfTestHub.cs b/perf/BenchmarkApp/PerformanceTest.Shared/IPerfTestHub.cs
@@ -10,7 +10,9 @@ namespace PerformanceTest.Shared
     public interface IPerfTestHub : IStreamingHub<IPerfTestHub, IPerfTestHubReceiver>
     {
         Task<int> CallMethodAsync(string arg1, int arg2, int arg3, int arg4);
+        ValueTask<int> CallMethodValueTaskAsync(string arg1, int arg2, int arg3, int arg4);
         Task<ComplexResponse> CallMethodComplexAsync(string arg1, int arg2, int arg3, int arg4);
+        ValueTask<ComplexResponse> CallMethodComplexValueTaskAsync(string arg1, int arg2, int arg3, int arg4);
         Task<(int StatusCode, byte[] Data)> CallMethodLargePayloadAsync(string arg1, int arg2, int arg3, int arg4, byte[] arg5);
     }
 

diff --git a/src/MagicOnion.Client.SourceGenerator/CodeGen/StaticStreamingHubClientGenerator.cs b/src/MagicOnion.Client.SourceGenerator/CodeGen/StaticStreamingHubClientGenerator.cs
@@ -120,24 +120,66 @@ static void EmitHelperMethods(StreamingHubClientBuildContext ctx)
         if (ctx.EnableStreamingHubDiagnosticHandler)
         {
             ctx.Writer.AppendLineWithFormat($$"""
-                            global::System.Threading.Tasks.Task<TResponse> WriteMessageWithResponseDiagnosticAsync<TRequest, TResponse>(int methodId, TRequest message, [global::System.Runtime.CompilerServices.CallerMemberName] string callerMemberName = default!)
+                            global::System.Threading.Tasks.Task<TResponse> WriteMessageWithResponseDiagnosticTaskAsync<TRequest, TResponse>(int methodId, TRequest message, [global::System.Runtime.CompilerServices.CallerMemberName] string callerMemberName = default!)
                             {
                                 if (diagnosticHandler is null)
                                 {
-                                    return base.WriteMessageWithResponseAsync<TRequest, TResponse>(methodId, message);
+                                    return base.WriteMessageWithResponseTaskAsync<TRequest, TResponse>(methodId, message);
                                 }
 
-                                return diagnosticHandler.OnMethodInvoke(this, methodId, callerMemberName, message, isFireAndForget: true, base.WriteMessageWithResponseAsync<TRequest, TResponse>);
+                                return diagnosticHandler.OnMethodInvoke(this, methodId, callerMemberName, message, isFireAndForget: true, base.WriteMessageWithResponseValueTaskOfTAsync<TRequest, TResponse>).AsTask();
                             }
 
-                            global::System.Threading.Tasks.Task<TResponse> WriteMessageFireAndForgetDiagnosticAsync<TRequest, TResponse>(int methodId, TRequest message, [global::System.Runtime.CompilerServices.CallerMemberName] string callerMemberName = default!)
+                            async global::System.Threading.Tasks.ValueTask WriteMessageWithResponseDiagnosticValueTaskAsync<TRequest, TResponse>(int methodId, TRequest message, [global::System.Runtime.CompilerServices.CallerMemberName] string callerMemberName = default!)
                             {
                                 if (diagnosticHandler is null)
                                 {
-                                    return base.WriteMessageFireAndForgetAsync<TRequest, TResponse>(methodId, message);
+                                    await base.WriteMessageWithResponseValueTaskAsync<TRequest, TResponse>(methodId, message);
+                                    return;
                                 }
 
-                                return diagnosticHandler.OnMethodInvoke(this, methodId, callerMemberName, message, isFireAndForget: true, base.WriteMessageFireAndForgetAsync<TRequest, TResponse>);
+                                await diagnosticHandler.OnMethodInvoke(this, methodId, callerMemberName, message, isFireAndForget: true, base.WriteMessageWithResponseValueTaskOfTAsync<TRequest, TResponse>);
+                            }
+
+                            global::System.Threading.Tasks.ValueTask<TResponse> WriteMessageWithResponseDiagnosticValueTaskOfTAsync<TRequest, TResponse>(int methodId, TRequest message, [global::System.Runtime.CompilerServices.CallerMemberName] string callerMemberName = default!)
+                            {
+                                if (diagnosticHandler is null)
+                                {
+                                    return base.WriteMessageWithResponseValueTaskOfTAsync<TRequest, TResponse>(methodId, message);
+                                }
+
+                                return diagnosticHandler.OnMethodInvoke(this, methodId, callerMemberName, message, isFireAndForget: true, base.WriteMessageWithResponseValueTaskOfTAsync<TRequest, TResponse>);
+                            }
+
+                            global::System.Threading.Tasks.Task<TResponse> WriteMessageFireAndForgetDiagnosticTaskAsync<TRequest, TResponse>(int methodId, TRequest message, [global::System.Runtime.CompilerServices.CallerMemberName] string callerMemberName = default!)
+                            {
+                                if (diagnosticHandler is null)
+                                {
+                                    return base.WriteMessageFireAndForgetTaskAsync<TRequest, TResponse>(methodId, message);
+                                }
+
+                                return diagnosticHandler.OnMethodInvoke(this, methodId, callerMemberName, message, isFireAndForget: true, base.WriteMessageFireAndForgetValueTaskOfTAsync<TRequest, TResponse>).AsTask();
+                            }
+
+                            async global::System.Threading.Tasks.ValueTask WriteMessageFireAndForgetDiagnosticValueTaskAsync<TRequest, TResponse>(int methodId, TRequest message, [global::System.Runtime.CompilerServices.CallerMemberName] string callerMemberName = default!)
+                            {
+                                if (diagnosticHandler is null)
+                                {
+                                    await base.WriteMessageFireAndForgetTaskAsync<TRequest, TResponse>(methodId, message);
+                                    return;
+                                }
+
+                                await diagnosticHandler.OnMethodInvoke(this, methodId, callerMemberName, message, isFireAndForget: true, base.WriteMessageFireAndForgetValueTaskOfTAsync<TRequest, TResponse>);
+                            }
+
+                            global::System.Threading.Tasks.ValueTask<TResponse> WriteMessageFireAndForgetDiagnosticValueTaskOfTAsync<TRequest, TResponse>(int methodId, TRequest message, [global::System.Runtime.CompilerServices.CallerMemberName] string callerMemberName = default!)
+                            {
+                                if (diagnosticHandler is null)
+                                {
+                                    return base.WriteMessageFireAndForgetValueTaskOfTAsync<TRequest, TResponse>(methodId, message);
+                                }
+                            
+                                return diagnosticHandler.OnMethodInvoke(this, methodId, callerMemberName, message, isFireAndForget: true, base.WriteMessageFireAndForgetValueTaskOfTAsync<TRequest, TResponse>);
                             }
             """);
             ctx.Writer.AppendLine();
@@ -217,39 +259,39 @@ static void EmitHubMethods(StreamingHubClientBuildContext ctx, bool isFireAndFor
             };
             var isReturnTypeVoid = method.MethodReturnType == MagicOnionTypeInfo.KnownTypes.System_Void;
             var writeMessageTarget = isFireAndForget ? "parent" : "this";
-            var writeMessageAsync = ctx.EnableStreamingHubDiagnosticHandler
+            var writeMessageAsyncPrefix = ctx.EnableStreamingHubDiagnosticHandler
                 ? isFireAndForget || isReturnTypeVoid
-                    ? $"{writeMessageTarget}.WriteMessageFireAndForgetDiagnosticAsync"
-                    : $"{writeMessageTarget}.WriteMessageWithResponseDiagnosticAsync"
+                    ? $"{writeMessageTarget}.WriteMessageFireAndForgetDiagnostic"
+                    : $"{writeMessageTarget}.WriteMessageWithResponseDiagnostic"
                 : isFireAndForget || isReturnTypeVoid
-                    ? $"{writeMessageTarget}.WriteMessageFireAndForgetAsync"
-                    : $"{writeMessageTarget}.WriteMessageWithResponseAsync";
+                    ? $"{writeMessageTarget}.WriteMessageFireAndForget"
+                    : $"{writeMessageTarget}.WriteMessageWithResponse";
 
             if (isFireAndForget) ctx.Writer.Append("    ");
             ctx.Writer.AppendLineWithFormat($"""
                             public {(isReturnTypeVoid ? "void" : method.MethodReturnType.FullName)} {method.MethodName}({method.Parameters.ToMethodSignaturize()})
             """);
 
             if (isFireAndForget) ctx.Writer.Append("    ");
-            if (method.MethodReturnType == MagicOnionTypeInfo.KnownTypes.System_Threading_Tasks_ValueTask)
+            if (method.MethodReturnType == MagicOnionTypeInfo.KnownTypes.System_Threading_Tasks_ValueTask || isReturnTypeVoid)
             {
-                // ValueTask
+                // ValueTask, void
                 ctx.Writer.AppendLineWithFormat($"""
-                                => new global::System.Threading.Tasks.ValueTask({writeMessageAsync}<{method.RequestType.FullName}, {method.ResponseType.FullName}>({method.HubId}{writeMessageParameters}));
+                                => {writeMessageAsyncPrefix}ValueTaskAsync<{method.RequestType.FullName}, {method.ResponseType.FullName}>({method.HubId}{writeMessageParameters});
             """);
             }
             else if (method.MethodReturnType.HasGenericArguments && method.MethodReturnType.GetGenericTypeDefinition() == MagicOnionTypeInfo.KnownTypes.System_Threading_Tasks_ValueTask)
             {
                 // ValueTask<T>
                 ctx.Writer.AppendLineWithFormat($"""
-                                => new global::System.Threading.Tasks.ValueTask<{method.ResponseType.FullName}>({writeMessageAsync}<{method.RequestType.FullName}, {method.ResponseType.FullName}>({method.HubId}{writeMessageParameters}));
+                                => {writeMessageAsyncPrefix}ValueTaskOfTAsync<{method.RequestType.FullName}, {method.ResponseType.FullName}>({method.HubId}{writeMessageParameters});
             """);
             }
             else
             {
                 // Task, Task<T>
                 ctx.Writer.AppendLineWithFormat($"""
-                                => {writeMessageAsync}<{method.RequestType.FullName}, {method.ResponseType.FullName}>({method.HubId}{writeMessageParameters});
+                                => {writeMessageAsyncPrefix}TaskAsync<{method.RequestType.FullName}, {method.ResponseType.FullName}>({method.HubId}{writeMessageParameters});
             """);
             }
         }