Record data-loader execution time within Apollo Tracing results (#3691)

Shane32 · web-flow · commit 8360265f4848 · 2023-11-05T14:34:06.000-05:00
* Record data-loader execution time within apollo tracing results

* Update

* Update src/GraphQL.DataLoader.Tests/ApolloTracingTests.cs

* Update src/GraphQL.DataLoader.Tests/ApolloTracingTests.cs
diff --git a/src/GraphQL.DataLoader.Tests/ApolloTracingTests.cs b/src/GraphQL.DataLoader.Tests/ApolloTracingTests.cs
@@ -0,0 +1,49 @@
+using GraphQL.Instrumentation;
+using GraphQL.Types;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace GraphQL.DataLoader.Tests;
+
+public class ApolloTracingTests
+{
+    // verifies that the time spent executing the data loader is included with the
+    // resolver duration, as the resolver has not 'completed' until the data loader completes
+    //
+    // note that (a) for batch data loaders, the time spent in the data loader is counted for all resolvers
+    // waiting on that data loader, and (b) any delay caused by the execution engine is counted within the
+    // resolver execution time
+    //
+    // perhaps it would be more ideal to separately record the data loader execution time within the apollo
+    // tracing output (e.g. 'extensions.tracing.execution.dataloaders'), but this does not currently happen,
+    // and the execution engine is not aware of which queued data loaders run as a single batch operation or not,
+    // making implementation difficult
+    [Fact]
+    public async Task DataLoaderTimeRecorded()
+    {
+        var query = new ObjectGraphType() { Name = "Query" };
+        query.Field<StringGraphType>("test")
+            .Resolve(_ => new SimpleDataLoader<string>(async _ =>
+            {
+                await Task.Delay(2000);
+                return "Ok";
+            }));
+
+        var serviceCollection = new ServiceCollection();
+        serviceCollection.AddGraphQL(b => b
+            .AddSchema(provider => new Schema(provider) { Query = query })
+            .AddSystemTextJson()
+            .UseApolloTracing());
+        var services = serviceCollection.BuildServiceProvider();
+        var serializer = services.GetRequiredService<IGraphQLTextSerializer>();
+        var executer = services.GetRequiredService<IDocumentExecuter>();
+        var ret = await executer.ExecuteAsync(new()
+        {
+            Query = "{test}",
+            RequestServices = services,
+        });
+        var apolloTrace = ret.Extensions["tracing"].ShouldBeOfType<ApolloTrace>();
+        var resolverData = apolloTrace.Execution.Resolvers.Single();
+        resolverData.Path.ShouldBe(new object[] { "test" });
+        resolverData.Duration.ShouldBeGreaterThanOrEqualTo(2000000000L);
+    }
+}
diff --git a/src/GraphQL.DataLoader.Tests/GraphQL.DataLoader.Tests.csproj b/src/GraphQL.DataLoader.Tests/GraphQL.DataLoader.Tests.csproj
@@ -27,6 +27,7 @@
     <ProjectReference Include="..\GraphQL.NewtonsoftJson\GraphQL.NewtonsoftJson.csproj" />
     <ProjectReference Include="..\GraphQL.SystemTextJson\GraphQL.SystemTextJson.csproj" />
     <ProjectReference Include="..\GraphQL.DataLoader\GraphQL.DataLoader.csproj" />
+    <ProjectReference Include="..\GraphQL.MicrosoftDI\GraphQL.MicrosoftDI.csproj" />
   </ItemGroup>
 
 </Project>
diff --git a/src/GraphQL/Instrumentation/InstrumentFieldsMiddleware.cs b/src/GraphQL/Instrumentation/InstrumentFieldsMiddleware.cs
@@ -1,3 +1,5 @@
+using GraphQL.DataLoader;
+
 namespace GraphQL.Instrumentation
 {
     /// <summary>
@@ -25,8 +27,41 @@ public class InstrumentFieldsMiddleware : IFieldMiddleware
                 { "path", context.ResponsePath },
             };
 
-            using (context.Metrics.Subject("field", name, metadata))
-                return await next(context).ConfigureAwait(false);
+            var marker = context.Metrics.Subject("field", name, metadata);
+            var disposeMarker = true;
+            try
+            {
+                var ret = await next(context).ConfigureAwait(false);
+                if (ret is IDataLoaderResult dataLoaderResult)
+                {
+                    disposeMarker = false;
+                    return new CompleteDataLoaderResult(dataLoaderResult, marker);
+                }
+                return ret;
+            }
+            finally
+            {
+                if (disposeMarker)
+                    marker.Dispose();
+            }
+        }
+
+        private class CompleteDataLoaderResult : IDataLoaderResult
+        {
+            private readonly IDataLoaderResult _baseDataLoaderResult;
+            private readonly Metrics.Marker _marker;
+
+            public CompleteDataLoaderResult(IDataLoaderResult baseDataLoaderResult, Metrics.Marker marker)
+            {
+                _baseDataLoaderResult = baseDataLoaderResult;
+                _marker = marker;
+            }
+
+            public async Task<object?> GetResultAsync(CancellationToken cancellationToken = default)
+            {
+                using (_marker)
+                    return await _baseDataLoaderResult.GetResultAsync(cancellationToken).ConfigureAwait(false);
+            }
         }
     }
 }
