Task + ValueTask + TraceFix

coronabytes · Oct 4, 2024 · 1c95761 · 1c95761
1 parent 7a3d656
commit 1c95761
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 6 deletions.
diff --git a/Core.ServiceMesh.SourceGen.Tests/UnitTest1.cs b/Core.ServiceMesh.SourceGen.Tests/UnitTest1.cs
@@ -21,9 +21,12 @@ namespace SampleApp;
                      public interface ISomeService
                      {
                          ValueTask A(string a);
+                         Task Aa(string a);
                          ValueTask<string> B(string b);
+                         Task<string> Bb(string b);
                          ValueTask<T> C<T>(T a, T b) where T : INumber<T>;
                          IAsyncEnumerable<string> D(string d);
+                         IAsyncEnumerable<T> E<T>(T e) where T : INumber<T>;
                      }
 
                      [ServiceMesh("someservice")]
@@ -32,11 +35,19 @@ public class SomeService : ISomeService
                          public async ValueTask A(string a)
                          {
 
+                         }
+                         public async Task Aa(string a)
+                         {
+
                          }
                          public async ValueTask<string> B(string b)
                          {
                             return b + " " + b;
                          }
+                         public async Task<string> Bb(string b)
+                         {
+                            return b + " " + b;
+                         }
                          public async ValueTask<T> C<T>(T a, T b) where T : INumber<T>
                          {
                             return a + b;
@@ -47,8 +58,15 @@ public async IAsyncEnumerable<string> D(string d)
                             yield return "b";
                             yield return "c";
                          }
+                         public async IAsyncEnumerable<T> E<T>(T e) where T : INumber<T>
+                         {
+                            yield return e;
+                            yield return e;
+                            yield return e;
+                         }
                      }
                      """;
+
         return TestHelper.VerifySourceGen(source);
     }
 }

diff --git a/Core.ServiceMesh.SourceGen/ServiceMeshGenerator.cs b/Core.ServiceMesh.SourceGen/ServiceMeshGenerator.cs
@@ -2,8 +2,10 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Linq;
+using System.Reflection.Metadata;
 using System.Text;
 using System.Threading;
+using System.Threading.Tasks;
 using Core.ServiceMesh.SourceGen.Core;
 using Core.ServiceMesh.SourceGen.Model;
 using Microsoft.CodeAnalysis;
@@ -105,7 +107,7 @@ private static void BuildRemoteProxy(SourceProductionContext context, ServiceDes
         builder.AppendLine("[System.ComponentModel.EditorBrowsable(System.ComponentModel.EditorBrowsableState.Never)]");
         builder.AppendLine(
             $"public sealed class {service.ClassName}RemoteProxy(IServiceMesh mesh) : {service.ClassName}");
-        builder.AppendLine("{");
+        builder.Append("{");
 
         foreach (var m in service.Methods)
         {
@@ -114,7 +116,7 @@ private static void BuildRemoteProxy(SourceProductionContext context, ServiceDes
             var genericsEx = $"[{string.Join(", ", m.Generics.Select(x => $"typeof({x})"))}]";
             var invoke = $"await mesh.RequestAsync(subject, {parameterEx}, {genericsEx});";
 
-            if (m.Return.StartsWith("ValueTask<"))
+            if (m.Return.StartsWith("ValueTask<") || m.Return.StartsWith("Task<"))
                 invoke =
                     $"return await mesh.RequestAsync<{m.ReturnArguments[0]}>(subject, {parameterEx}, {genericsEx});";
 
@@ -157,18 +159,29 @@ private static void BuildTraceProxy(SourceProductionContext context, ServiceDesc
         builder.AppendLine("[System.ComponentModel.EditorBrowsable(System.ComponentModel.EditorBrowsableState.Never)]");
         builder.AppendLine(
             $"public sealed class {service.ClassName}TraceProxy({service.ClassName} svc) : I{service.ClassName}");
-        builder.AppendLine("{");
+        builder.Append("{");
 
         foreach (var m in service.Methods)
         {
             builder.AppendLine();
 
-            // TODO: use await for correct trace?
+            var invoke = $"await svc.{m.Name}({string.Join(", ", m.ParameterNames)});";
+
+            if (m.Return.StartsWith("ValueTask<") || m.Return.StartsWith("Task<"))
+                invoke =
+                    $"return await svc.{m.Name}({string.Join(", ", m.ParameterNames)});";
+
+            if (m.Return.StartsWith("IAsyncEnumerable<"))
+                invoke = $"""
+                          await foreach (var msg in svc.{m.Name}({string.Join(", ", m.ParameterNames)}))
+                                      yield return msg;
+                          """;
+
             var code = $$"""
-                             public {{m.Return}} {{m.Name}}{{(m.Generics.Any() ? "<" + string.Join(",", m.Generics.Select(x => x)) + ">" : string.Empty)}}({{string.Join(", ", m.Parameters.Select(x => x))}}) {{(m.Constraints.Any() ? string.Join(", ", m.Constraints) : string.Empty)}}
+                             public async {{m.Return}} {{m.Name}}{{(m.Generics.Any() ? "<" + string.Join(",", m.Generics.Select(x => x)) + ">" : string.Empty)}}({{string.Join(", ", m.Parameters.Select(x => x))}}) {{(m.Constraints.Any() ? string.Join(", ", m.Constraints) : string.Empty)}}
                              {
                                 using var activity = ServiceMeshActivity.Source.StartActivity("REQ {{service.ServiceName}}.{{m.Name}}", ActivityKind.Internal, Activity.Current?.Context ?? default);
-                                return svc.{{m.Name}}({{string.Join(", ", m.ParameterNames)}}); 
+                                {{invoke}}
                              }
                          """;