Merge pull request #746 from Cysharp/feature/StreamingHubDiagnosticHa…

…ndler [Preview] Introduce IStreamingHubDiagnosticHandler
Cysharp · Mar 21, 2024 · da8112b · da8112b
2 parents 015e5a0 + fe66d55
commit da8112b
Show file tree

Hide file tree

Showing 110 changed files with 1,467 additions and 19 deletions.
diff --git a/Directory.Packages.props b/Directory.Packages.props
@@ -37,7 +37,7 @@
     <PackageVersion Include="Microsoft.AspNetCore.Mvc.Testing" Version="7.0.0" />
     <PackageVersion Include="Microsoft.Extensions.Diagnostics.Testing" Version="8.0.0" />
     <!-- from https://pkgs.dev.azure.com/dnceng/public/_packaging/dotnet-tools/nuget/v3/index.json -->
-    <PackageVersion Include="Microsoft.CodeAnalysis.CSharp.SourceGenerators.Testing.XUnit" Version="1.1.2-beta1.23163.2" />
+    <PackageVersion Include="Microsoft.CodeAnalysis.CSharp.SourceGenerators.Testing.XUnit" Version="1.1.2-beta1.24163.6" />
     <PackageVersion Include="Microsoft.NET.Test.Sdk" Version="17.3.2" />
     <PackageVersion Include="NSubstitute" Version="5.0.0" />
     <PackageVersion Include="Testcontainers" Version="2.3.0" />

diff --git a/src/MagicOnion.Client.SourceGenerator/CodeGen/MagicOnionInitializerGenerator.cs b/src/MagicOnion.Client.SourceGenerator/CodeGen/MagicOnionInitializerGenerator.cs
@@ -59,6 +59,17 @@ partial class {{generationContext.InitializerPartialTypeName}}
                     internal static void Register() => TryRegisterProviderFactory();
             """);
         }
+
+        if (generationContext.Options.EnableStreamingHubDiagnosticHandler)
+        {
+            writer.AppendLineWithFormat($$"""
+                    /// <summary>
+                    /// Gets or sets a diagnostic handler for the StreamingHub.
+                    /// </summary>
+                    public static global::MagicOnion.Client.IStreamingHubDiagnosticHandler StreamingHubDiagnosticHandler { get; set; }
+            """);
+        }
+
         writer.AppendLineWithFormat($$"""
 
                     /// <summary>
@@ -126,12 +137,24 @@ static StreamingHubClientFactoryCache()
             """);
         foreach (var hubInfo in serviceCollection.Hubs)
         {
-            writer.AppendLineWithFormat($$"""
+            if (generationContext.Options.EnableStreamingHubDiagnosticHandler)
+            {
+                writer.AppendLineWithFormat($$"""
+                                if (typeof(TStreamingHub) == typeof({{hubInfo.ServiceType.FullName}}) && typeof(TReceiver) == typeof({{hubInfo.Receiver.ReceiverType.FullName}}))
+                                {
+                                    factory = ((global::MagicOnion.Client.StreamingHubClientFactoryDelegate<{{hubInfo.ServiceType.FullName}}, {{hubInfo.Receiver.ReceiverType.FullName}}>)((a, _, b, c, d, e) => new MagicOnionGeneratedClient.{{hubInfo.GetClientFullName()}}(a, b, c, d, e, StreamingHubDiagnosticHandler)));
+                                }
+            """);
+            }
+            else
+            {
+                writer.AppendLineWithFormat($$"""
                                 if (typeof(TStreamingHub) == typeof({{hubInfo.ServiceType.FullName}}) && typeof(TReceiver) == typeof({{hubInfo.Receiver.ReceiverType.FullName}}))
                                 {
                                     factory = ((global::MagicOnion.Client.StreamingHubClientFactoryDelegate<{{hubInfo.ServiceType.FullName}}, {{hubInfo.Receiver.ReceiverType.FullName}}>)((a, _, b, c, d, e) => new MagicOnionGeneratedClient.{{hubInfo.GetClientFullName()}}(a, b, c, d, e)));
                                 }
             """);
+            }
         }
 
         writer.AppendLineWithFormat($$$"""

diff --git a/src/MagicOnion.Client.SourceGenerator/CodeGen/StaticStreamingHubClientGenerator.cs b/src/MagicOnion.Client.SourceGenerator/CodeGen/StaticStreamingHubClientGenerator.cs
@@ -7,15 +7,18 @@ public class StaticStreamingHubClientGenerator
 {
     class StreamingHubClientBuildContext
     {
-        public StreamingHubClientBuildContext(MagicOnionStreamingHubInfo hub, StringBuilder writer)
+        public StreamingHubClientBuildContext(MagicOnionStreamingHubInfo hub, StringBuilder writer, bool enableStreamingHubDiagnosticHandler)
         {
             Hub = hub;
             Writer = writer;
+            EnableStreamingHubDiagnosticHandler = enableStreamingHubDiagnosticHandler;
         }
 
         public MagicOnionStreamingHubInfo Hub { get; }
 
         public StringBuilder Writer { get; }
+
+        public bool EnableStreamingHubDiagnosticHandler { get; }
     }
 
     public static string Build(GenerationContext generationContext, MagicOnionStreamingHubInfo hubInfo)
@@ -25,7 +28,7 @@ public static string Build(GenerationContext generationContext, MagicOnionStream
 
         EmitHeader(generationContext, writer);
 
-        var buildContext = new StreamingHubClientBuildContext(hubInfo, writer);
+        var buildContext = new StreamingHubClientBuildContext(hubInfo, writer, generationContext.Options.EnableStreamingHubDiagnosticHandler);
 
         EmitPreamble(generationContext, buildContext);
         EmitHubClientClass(generationContext, buildContext);
@@ -79,6 +82,17 @@ static void EmitPostscript(GenerationContext generationContext, StreamingHubClie
         }
     }
 
+    static void EmitProperties(StreamingHubClientBuildContext ctx)
+    {
+        if (ctx.EnableStreamingHubDiagnosticHandler)
+        {
+            ctx.Writer.AppendLineWithFormat($"""
+                            readonly global::MagicOnion.Client.IStreamingHubDiagnosticHandler diagnosticHandler;
+            """);
+            ctx.Writer.AppendLine();
+        }
+    }
+
     static void EmitHubClientClass(GenerationContext generationContext, StreamingHubClientBuildContext ctx)
     {
         ctx.Writer.AppendLineWithFormat($$"""
@@ -88,6 +102,7 @@ public class {{ctx.Hub.GetClientFullName()}} : global::MagicOnion.Client.Streami
             """);
         EmitProperties(ctx);
         EmitConstructor(ctx);
+        EmitHelperMethods(ctx);
         EmitHubMethods(ctx, isFireAndForget: false);
         EmitFireAndForget(ctx);
         EmitOnBroadcastEvent(ctx);
@@ -98,18 +113,54 @@ public class {{ctx.Hub.GetClientFullName()}} : global::MagicOnion.Client.Streami
         // }
     }
 
-    static void EmitProperties(StreamingHubClientBuildContext ctx)
+    static void EmitHelperMethods(StreamingHubClientBuildContext ctx)
     {
+        if (ctx.EnableStreamingHubDiagnosticHandler)
+        {
+            ctx.Writer.AppendLineWithFormat($$"""
+                            async global::System.Threading.Tasks.Task<TResponse> WriteMessageWithResponseDiagnosticAsync<TRequest, TResponse>(int methodId, TRequest message, [global::System.Runtime.CompilerServices.CallerMemberName] string callerMemberName = default!)
+                            {
+                                var requestId = global::System.Guid.NewGuid();
+                                diagnosticHandler?.OnRequestBegin(this, requestId, callerMemberName, message, isFireAndForget: false);
+                                var result = await base.WriteMessageWithResponseAsync<TRequest, TResponse>(methodId, message).ConfigureAwait(false);
+                                diagnosticHandler?.OnRequestEnd(this, requestId, callerMemberName, result);
+                                return result;
+                            }
+
+                            async global::System.Threading.Tasks.Task<TResponse> WriteMessageFireAndForgetDiagnosticAsync<TRequest, TResponse>(int methodId, TRequest message, [global::System.Runtime.CompilerServices.CallerMemberName] string callerMemberName = default!)
+                            {
+                                var requestId = global::System.Guid.NewGuid();
+                                diagnosticHandler?.OnRequestBegin(this, requestId, callerMemberName, message, isFireAndForget: true);
+                                var result = await base.WriteMessageFireAndForgetAsync<TRequest, TResponse>(methodId, message).ConfigureAwait(false);
+                                diagnosticHandler?.OnRequestEnd(this, requestId, callerMemberName, result);
+                                return result;
+                            }
+            """);
+            ctx.Writer.AppendLine();
+        }
     }
-    
+
     static void EmitConstructor(StreamingHubClientBuildContext ctx)
     {
-        ctx.Writer.AppendLineWithFormat($$"""
+        if (ctx.EnableStreamingHubDiagnosticHandler)
+        {
+            ctx.Writer.AppendLineWithFormat($$"""
+                            public {{ctx.Hub.GetClientFullName()}}(global::Grpc.Core.CallInvoker callInvoker, global::System.String host, global::Grpc.Core.CallOptions options, global::MagicOnion.Serialization.IMagicOnionSerializerProvider serializerProvider, global::MagicOnion.Client.IMagicOnionClientLogger logger, global::MagicOnion.Client.IStreamingHubDiagnosticHandler diagnosticHandler)
+                                : base("{{ctx.Hub.ServiceType.Name}}", callInvoker, host, options, serializerProvider, logger)
+                            {
+                                this.diagnosticHandler = diagnosticHandler;
+                            }
+            """);
+        }
+        else
+        {
+            ctx.Writer.AppendLineWithFormat($$"""
                             public {{ctx.Hub.GetClientFullName()}}(global::Grpc.Core.CallInvoker callInvoker, global::System.String host, global::Grpc.Core.CallOptions options, global::MagicOnion.Serialization.IMagicOnionSerializerProvider serializerProvider, global::MagicOnion.Client.IMagicOnionClientLogger logger)
                                 : base("{{ctx.Hub.ServiceType.Name}}", callInvoker, host, options, serializerProvider, logger)
                             {
                             }
             """);
+        }
         ctx.Writer.AppendLine();
     }
 
@@ -160,6 +211,13 @@ static void EmitHubMethods(StreamingHubClientBuildContext ctx, bool isFireAndFor
                 // new DynamicArgumentTuple(arg1, arg2, ...)
                 _ => $", {method.Parameters.ToNewDynamicArgumentTuple()}",
             };
+            var writeMessageAsync = ctx.EnableStreamingHubDiagnosticHandler
+                ? isFireAndForget
+                    ? "parent.WriteMessageFireAndForgetDiagnosticAsync"
+                    : "this.WriteMessageWithResponseDiagnosticAsync"
+                : isFireAndForget
+                    ? "parent.WriteMessageFireAndForgetAsync"
+                    : "base.WriteMessageWithResponseAsync";
 
             if (isFireAndForget) ctx.Writer.Append("    ");
             ctx.Writer.AppendLineWithFormat($"""
@@ -171,21 +229,21 @@ static void EmitHubMethods(StreamingHubClientBuildContext ctx, bool isFireAndFor
             {
                 // ValueTask
                 ctx.Writer.AppendLineWithFormat($"""
-                                => new global::System.Threading.Tasks.ValueTask({(isFireAndForget ? "parent.WriteMessageFireAndForgetAsync" : "base.WriteMessageWithResponseAsync")}<{method.RequestType.FullName}, {method.ResponseType.FullName}>({method.HubId}{writeMessageParameters}));
+                                => new global::System.Threading.Tasks.ValueTask({writeMessageAsync}<{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}>({(isFireAndForget ? "parent.WriteMessageFireAndForgetAsync" : "base.WriteMessageWithResponseAsync")}<{method.RequestType.FullName}, {method.ResponseType.FullName}>({method.HubId}{writeMessageParameters}));
+                                => new global::System.Threading.Tasks.ValueTask<{method.ResponseType.FullName}>({writeMessageAsync}<{method.RequestType.FullName}, {method.ResponseType.FullName}>({method.HubId}{writeMessageParameters}));
             """);
             }
             else
             {
                 // Task, Task<T>
                 ctx.Writer.AppendLineWithFormat($"""
-                                => {(isFireAndForget ? "parent.WriteMessageFireAndForgetAsync" : "base.WriteMessageWithResponseAsync")}<{method.RequestType.FullName}, {method.ResponseType.FullName}>({method.HubId}{writeMessageParameters});
+                                => {writeMessageAsync}<{method.RequestType.FullName}, {method.ResponseType.FullName}>({method.HubId}{writeMessageParameters});
             """);
             }
         }
@@ -210,14 +268,29 @@ protected override void OnBroadcastEvent(global::System.Int32 methodId, global::
                 _ => string.Join(", ", Enumerable.Range(1, method.Parameters.Count).Select(x => $"value.Item{x}"))
             };
 
-            ctx.Writer.AppendLineWithFormat($$"""
+            if (ctx.EnableStreamingHubDiagnosticHandler)
+            {
+                ctx.Writer.AppendLineWithFormat($$"""
                                     case {{method.HubId}}: // {{method.MethodReturnType.ToDisplayName()}} {{method.MethodName}}({{method.Parameters.ToMethodSignaturize()}})
                                         {
                                             var value = base.Deserialize<{{method.RequestType.FullName}}>(data);
+                                            diagnosticHandler?.OnBroadcastEvent(this, "{{method.MethodName}}", value);
                                             receiver.{{method.MethodName}}({{methodArgs}});
                                         }
                                         break;
             """);
+            }
+            else
+            {
+                ctx.Writer.AppendLineWithFormat($$"""
+                                    case {{method.HubId}}: // {{method.MethodReturnType.ToDisplayName()}} {{method.MethodName}}({{method.Parameters.ToMethodSignaturize()}})
+                                        {
+                                            var value = base.Deserialize<{{method.RequestType.FullName}}>(data);
+                                            receiver.{{method.MethodName}}({{methodArgs}});
+                                        }
+                                        break;
+            """);
+            }
         }
         ctx.Writer.AppendLine("""
                                 }

diff --git a/src/MagicOnion.Client.SourceGenerator/GenerationContext.cs b/src/MagicOnion.Client.SourceGenerator/GenerationContext.cs
@@ -29,7 +29,19 @@ public enum SerializerType
     MemoryPack = 1,
 }
 
-public record GenerationOptions(SerializerType Serializer, bool DisableAutoRegistration, string MessagePackFormatterNamespace)
+public record GenerationOptions(
+    SerializerType Serializer,
+    bool DisableAutoRegistration,
+    string MessagePackFormatterNamespace,
+    bool EnableStreamingHubDiagnosticHandler,
+    string GenerateFileHintNamePrefix
+)
 {
-    public static GenerationOptions Default { get; } = new (SerializerType.MessagePack, false, "MessagePack.Formatters");
+    public static GenerationOptions Default { get; } = new (
+        SerializerType.MessagePack,
+        DisableAutoRegistration: false,
+        MessagePackFormatterNamespace: "MessagePack.Formatters",
+        EnableStreamingHubDiagnosticHandler: false,
+        GenerateFileHintNamePrefix: string.Empty
+    );
 }
diff --git a/src/MagicOnion.Client.SourceGenerator/MagicOnionClientSourceGenerator.Emitter.cs b/src/MagicOnion.Client.SourceGenerator/MagicOnionClientSourceGenerator.Emitter.cs
@@ -44,6 +44,13 @@ internal class {{MagicOnionClientGenerationAttributeName}} : global::System.Attr
                         /// Gets or set the namespace of pre-generated MessagePackFormatters. The default value is <c>MessagePack.Formatters</c>.
                         /// </summary>
                         public string MessagePackFormatterNamespace { get; set; } = "MessagePack.Formatters";
+
+                        /// <summary>
+                        /// Gets or set whether to enable the StreamingHandler diagnostic handler. This is for debugging purpose. The default value is <see langword="false" />.
+                        /// </summary>
+                        public bool EnableStreamingHubDiagnosticHandler { get; set; } = false;
+
+                        public string GenerateFileHintNamePrefix { get; set; } = string.Empty;
 
                         public global::System.Type[] TypesContainedInTargetAssembly { get; }
 
@@ -76,7 +83,7 @@ public static void Emit(GenerationContext context, ImmutableArray<INamedTypeSymb
 
             foreach (var (path, source) in generated)
             {
-                context.SourceProductionContext.AddSource(path, source);
+                context.SourceProductionContext.AddSource(context.Options.GenerateFileHintNamePrefix + path, source);
             }
         }
 

diff --git a/src/MagicOnion.Client.SourceGenerator/MagicOnionClientSourceGenerator.Parser.cs b/src/MagicOnion.Client.SourceGenerator/MagicOnionClientSourceGenerator.Parser.cs
@@ -93,6 +93,12 @@ static GenerationOptions ParseClientGenerationOptions(AttributeData attr)
                 case nameof(GenerationOptions.MessagePackFormatterNamespace):
                     options = options with { MessagePackFormatterNamespace = (string)namedArg.Value.Value! };
                     break;
+                case nameof(GenerationOptions.EnableStreamingHubDiagnosticHandler):
+                    options = options with { EnableStreamingHubDiagnosticHandler = (bool)namedArg.Value.Value! };
+                    break;
+                case nameof(GenerationOptions.GenerateFileHintNamePrefix):
+                    options = options with { GenerateFileHintNamePrefix = (string)namedArg.Value.Value! };
+                    break;
             }
         }
 

diff --git a/...lient.Unity/Assets/Scripts/MagicOnion/MagicOnion.Client/IStreamingHubDiagnosticHandler.cs b/...lient.Unity/Assets/Scripts/MagicOnion/MagicOnion.Client/IStreamingHubDiagnosticHandler.cs
@@ -0,0 +1,43 @@
+using System;
+
+namespace MagicOnion.Client
+{
+    /// <summary>
+    /// [Preview] The interface of the handler for StreamingHub diagnostics. This API may change in the future.
+    /// </summary>
+    public interface IStreamingHubDiagnosticHandler
+    {
+        /// <summary>
+        /// The callback method at the beginning of a Hub method request. This API may change in the future.
+        /// </summary>
+        /// <typeparam name="THub"></typeparam>
+        /// <typeparam name="TRequest"></typeparam>
+        /// <param name="hubInstance"></param>
+        /// <param name="requestId"></param>
+        /// <param name="methodName"></param>
+        /// <param name="request"></param>
+        /// <param name="isFireAndForget"></param>
+        void OnRequestBegin<THub, TRequest>(THub hubInstance, Guid requestId, string methodName, TRequest request, bool isFireAndForget);
+
+        /// <summary>
+        /// [Preview] The callback method at the end of a Hub method request. This API may change in the future.
+        /// </summary>
+        /// <typeparam name="THub"></typeparam>
+        /// <typeparam name="TResponse"></typeparam>
+        /// <param name="hubInstance"></param>
+        /// <param name="requestId"></param>
+        /// <param name="methodName"></param>
+        /// <param name="response"></param>
+        void OnRequestEnd<THub, TResponse>(THub hubInstance, Guid requestId, string methodName, TResponse response);
+
+        /// <summary>
+        /// [Preview] The callback method when a method of HubReceiver is invoked. This API may change in the future.
+        /// </summary>
+        /// <typeparam name="THub"></typeparam>
+        /// <typeparam name="T"></typeparam>
+        /// <param name="hubInstance"></param>
+        /// <param name="methodName"></param>
+        /// <param name="value"></param>
+        void OnBroadcastEvent<THub, T>(THub hubInstance, string methodName, T value);
+    }
+}
diff --git a/....Unity/Assets/Scripts/MagicOnion/MagicOnion.Client/IStreamingHubDiagnosticHandler.cs.meta b/....Unity/Assets/Scripts/MagicOnion/MagicOnion.Client/IStreamingHubDiagnosticHandler.cs.meta