(DO NOT MERGE) Prototype of OpenTelemetry Tracing in GAX

Google.Api.Gax.Grpc.Tests/ApiBidirectionalStreamingCallTest.cs

@@ -7,6 +7,8 @@
 
 using Google.Api.Gax.Testing;
 using System;
+using System.Diagnostics;
+using System.Linq;
 using Xunit;
 
 namespace Google.Api.Gax.Grpc.Tests
@@ -17,7 +19,7 @@ public class ApiBidirectionalStreamingCallTest
         public void FailWithRetry()
         {
             var apiCall = ApiBidirectionalStreamingCall.Create<int, int>(
-                "Method", 
+                "Method",
                 callOptions => null,
                 CallSettings.FromRetry(new RetrySettings(5, TimeSpan.Zero, TimeSpan.Zero, 1.0, e => false, RetrySettings.RandomJitter)),
                 new BidirectionalStreamingSettings(100),
@@ -52,5 +54,33 @@ public void WithLogging()
             var entry = Assert.Single(logger.ListLogEntries());
             Assert.Contains("BidiStreamingMethod", entry.Message);
         }
+
+        [Fact]
+        public void WithTracing()
+        {
+            var sourceName = "source";
+            Activity capturedActivity = null;
+
+            // We need a listener attached to an ActivitySource, otherwise activity won't be created.
+            using var listener = new ActivityListener
+            {
+                ShouldListenTo = source => source.Name == sourceName,
+                Sample = (ref ActivityCreationOptions<ActivityContext> options) => ActivitySamplingResult.AllDataAndRecorded,
+                ActivityStarted = activity => capturedActivity = activity
+            };
+            ActivitySource.AddActivityListener(listener);
+            using var source = new ActivitySource(sourceName);
+            var apiCall = ApiBidirectionalStreamingCall.Create<int, int>(
+                "BidiStreamingMethod",
+                callOptions => null,
+                null,
+                new BidirectionalStreamingSettings(100),
+                new FakeClock()).WithTracing(source);
+            apiCall.Call(null);
+            Assert.NotNull(capturedActivity);
+            Assert.Equal(ApiCallTracingExtensions.BidiStreamingCallType, capturedActivity.Tags.First(j => j.Key == ApiCallTracingExtensions.GrpcCallTypeTag).Value);
+            Assert.Equal(sourceName, capturedActivity.Source.Name);
+            Assert.Contains("BidiStreamingMethod", capturedActivity.OperationName);
+        }
     }
 }

Google.Api.Gax.Grpc.Tests/ApiCallTest.cs

@@ -9,6 +9,8 @@
 using Google.Protobuf.Reflection;
 using Grpc.Core;
 using System;
+using System.Diagnostics;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
@@ -217,6 +219,60 @@ public async Task WithLogging_Async()
             Assert.All(entries, entry => Assert.Contains("SimpleMethod", entry.Message));
         }
 
+        [Fact]
+        public void WithTracing_Sync()
+        {
+            var sourceName = "source";
+            Activity capturedActivity = null;
+
+            // We need a listener attached to an ActivitySource, otherwise activity won't be created.
+            using var listener = new ActivityListener
+            {
+                ShouldListenTo = source => source.Name == sourceName,
+                Sample = (ref ActivityCreationOptions<ActivityContext> options) => ActivitySamplingResult.AllDataAndRecorded,
+                ActivityStarted = activity => capturedActivity = activity
+            };
+            ActivitySource.AddActivityListener(listener);
+            using var source = new ActivitySource(sourceName);
+            var call = new ApiCall<SimpleRequest, SimpleResponse>(
+                "SimpleMethod",
+                (req, cs) => Task.FromResult(default(SimpleResponse)),
+                (req, cs) => null,
+                null).WithTracing(source);
+            call.Sync(new SimpleRequest(), null);           
+            Assert.NotNull(capturedActivity);
+            Assert.Contains("SimpleMethod", capturedActivity.OperationName);
+            Assert.Equal(ApiCallTracingExtensions.UnaryCallType, capturedActivity.Tags.First(j => j.Key == ApiCallTracingExtensions.GrpcCallTypeTag).Value);
+            Assert.Equal(sourceName, capturedActivity.Source.Name);
+        }
+
+        [Fact]
+        public async Task WithTracing_Async()
+        {
+            var sourceName = "source";
+            Activity capturedActivity = null;
+
+            // We need a listener attached to an ActivitySource, otherwise activity won't be created.
+            using var listener = new ActivityListener
+            {
+                ShouldListenTo = source => source.Name == sourceName,
+                Sample = (ref ActivityCreationOptions<ActivityContext> options) => ActivitySamplingResult.AllDataAndRecorded,
+                ActivityStarted = activity => capturedActivity = activity
+            };
+            ActivitySource.AddActivityListener(listener);
+            using var source = new ActivitySource(sourceName);
+            var call = new ApiCall<SimpleRequest, SimpleResponse>(
+                "SimpleMethod",
+                (req, cs) => Task.FromResult(default(SimpleResponse)),
+                (req, cs) => null,
+                null).WithTracing(source);
+            await call.Async(new SimpleRequest(), null);
+            Assert.NotNull(capturedActivity);
+            Assert.Contains("SimpleMethod", capturedActivity.OperationName);
+            Assert.Equal(ApiCallTracingExtensions.UnaryCallType, capturedActivity.Tags.First(j => j.Key == ApiCallTracingExtensions.GrpcCallTypeTag).Value);
+            Assert.Equal(sourceName, capturedActivity.Source.Name);
+        }
+
         internal class ExtractedRequestParamRequest : IMessage<ExtractedRequestParamRequest>
         {
             public string TableName { get; set; }

Google.Api.Gax.Grpc.Tests/ApiClientStreamingCallTest.cs

@@ -7,6 +7,8 @@
 
 using Google.Api.Gax.Testing;
 using System;
+using System.Diagnostics;
+using System.Linq;
 using Xunit;
 
 namespace Google.Api.Gax.Grpc.Tests
@@ -52,5 +54,33 @@ public void WithLogging()
             var entry = Assert.Single(logger.ListLogEntries());
             Assert.Contains("ClientStreamingMethod", entry.Message);
         }
+
+        [Fact]
+        public void WithTracing()
+        {
+            var sourceName = "source";
+            Activity capturedActivity = null;
+
+            // We need a listener attached to an ActivitySource, otherwise activity won't be created.
+            using var listener = new ActivityListener
+            {
+                ShouldListenTo = source => source.Name == sourceName,
+                Sample = (ref ActivityCreationOptions<ActivityContext> options) => ActivitySamplingResult.AllDataAndRecorded,
+                ActivityStarted = activity => capturedActivity = activity
+            };
+            ActivitySource.AddActivityListener(listener);
+            using var source = new ActivitySource(sourceName);
+            var apiCall = ApiClientStreamingCall.Create<int, int>(
+                "ClientStreamingMethod",
+                callOptions => null,
+                null,
+                new ClientStreamingSettings(100),
+                new FakeClock()).WithTracing(source);
+            apiCall.Call(null);
+            Assert.NotNull(capturedActivity);
+            Assert.Equal(ApiCallTracingExtensions.ClientStreamingCallType, capturedActivity.Tags.First(j => j.Key == ApiCallTracingExtensions.GrpcCallTypeTag).Value);
+            Assert.Contains("ClientStreamingMethod", capturedActivity.OperationName);
+            Assert.Equal(sourceName, capturedActivity.Source.Name);
+        }
     }
 }

Google.Api.Gax.Grpc.Tests/ApiServerStreamingCallTest.cs

@@ -8,6 +8,8 @@
 using Google.Api.Gax.Testing;
 using Grpc.Core;
 using System;
+using System.Diagnostics;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
@@ -172,5 +174,59 @@ public async Task WithLogging_Async()
             Assert.Equal(2, entries.Count);
             Assert.All(entries, entry => Assert.Contains("SimpleMethod", entry.Message));
         }
+
+        [Fact]
+        public void WithTracing_Sync()
+        {
+            var sourceName = "source";
+            Activity capturedActivity = null;
+
+            // We need a listener attached to an ActivitySource, otherwise activity won't be created.
+            using var listener = new ActivityListener
+            {
+                ShouldListenTo = source => source.Name == sourceName,
+                Sample = (ref ActivityCreationOptions<ActivityContext> options) => ActivitySamplingResult.AllDataAndRecorded,
+                ActivityStarted = activity => capturedActivity = activity
+            };
+            ActivitySource.AddActivityListener(listener);
+            using var source = new ActivitySource(sourceName);
+            var call = new ApiServerStreamingCall<SimpleRequest, SimpleResponse>(
+                "SimpleMethod",
+                (req, cs) => Task.FromResult(default(AsyncServerStreamingCall<SimpleResponse>)),
+                (req, cs) => null,
+                null).WithTracing(source);
+            call.Call(new SimpleRequest(), null);
+            Assert.NotNull(capturedActivity);
+            Assert.Contains("SimpleMethod", capturedActivity.OperationName);
+            Assert.Equal(ApiCallTracingExtensions.ServerStreamingCallType, capturedActivity.Tags.First(j => j.Key == ApiCallTracingExtensions.GrpcCallTypeTag).Value);
+            Assert.Equal(sourceName, capturedActivity.Source.Name);
+        }
+
+        [Fact]
+        public async Task WithTracing_Async()
+        {
+            var sourceName = "source";
+            Activity capturedActivity = null;
+
+            // We need a listener attached to an ActivitySource, otherwise activity won't be created.
+            using var listener = new ActivityListener
+            {
+                ShouldListenTo = source => source.Name == sourceName,
+                Sample = (ref ActivityCreationOptions<ActivityContext> options) => ActivitySamplingResult.AllDataAndRecorded,
+                ActivityStarted = activity => capturedActivity = activity
+            };
+            ActivitySource.AddActivityListener(listener);
+            using var source = new ActivitySource(sourceName);
+            var call = new ApiServerStreamingCall<SimpleRequest, SimpleResponse>(
+                "SimpleMethod",
+                (req, cs) => Task.FromResult(default(AsyncServerStreamingCall<SimpleResponse>)),
+                (req, cs) => null,
+                null).WithTracing(source);
+            await call.CallAsync(new SimpleRequest(), null);
+            Assert.NotNull(capturedActivity);
+            Assert.Contains("SimpleMethod", capturedActivity.OperationName);
+            Assert.Equal(ApiCallTracingExtensions.ServerStreamingCallType, capturedActivity.Tags.FirstOrDefault(j => j.Key == ApiCallTracingExtensions.GrpcCallTypeTag).Value);
+            Assert.Equal(sourceName, capturedActivity.Source.Name);
+        }
     }
 }

Google.Api.Gax.Grpc/ApiBidirectionalStreamingCall.cs

@@ -8,6 +8,7 @@
 using Grpc.Core;
 using Microsoft.Extensions.Logging;
 using System;
+using System.Diagnostics;
 
 namespace Google.Api.Gax.Grpc
 {
@@ -82,6 +83,10 @@ internal ApiBidirectionalStreamingCall<TRequest, TResponse> WithLogging(ILogger
             logger is null
                 ? this
                 : new ApiBidirectionalStreamingCall<TRequest, TResponse>(_methodName, _call.WithLogging(logger, _methodName), BaseCallSettings, StreamingSettings);
-
+
+        internal ApiBidirectionalStreamingCall<TRequest, TResponse> WithTracing(ActivitySource source) =>
+            source is null
+                ? this
+                : new ApiBidirectionalStreamingCall<TRequest, TResponse>(_methodName, _call.WithTracing(source, _methodName), BaseCallSettings, StreamingSettings);
     }
 }

Google.Api.Gax.Grpc/ApiCall.cs

@@ -9,6 +9,7 @@
 using Grpc.Core;
 using Microsoft.Extensions.Logging;
 using System;
+using System.Diagnostics;
 using System.Threading.Tasks;
 
 namespace Google.Api.Gax.Grpc
@@ -158,6 +159,7 @@ internal ApiCall<TRequest, TResponse> WithMergedBaseCallSettings(CallSettings ca
 
         // Note: it seems unfortunate that we can't reuse whatever logger is already configured here, but we don't
         // have access to it. This is the logger for logging "I'm about to back off / retry" rather than the actual calls.
+        // TODO: Check whether ActivitySource is needed or since it is called later we don't need it.
         internal ApiCall<TRequest, TResponse> WithRetry(IClock clock, IScheduler scheduler, ILogger retryLogger) =>
             new ApiCall<TRequest, TResponse>(
                 _methodName,
@@ -173,6 +175,14 @@ logger is null
                     _syncCall.WithLogging(logger, _methodName),
                     BaseCallSettings);
 
+        internal ApiCall<TRequest, TResponse> WithTracing(ActivitySource source) =>
+           source is null
+               ? this
+               : new ApiCall<TRequest, TResponse>(
+                   _methodName, _asyncCall.WithTracing(source, _methodName),
+                   _syncCall.WithTracing(source, _methodName),
+                   BaseCallSettings);
+
         /// <summary>
         /// Constructs a new <see cref="ApiCall{TRequest, TResponse}"/> that applies an overlay to
         /// the underlying <see cref="CallSettings"/>. If a value exists in both the original and

Google.Api.Gax.Grpc/ApiCallTracingExtensions.cs

@@ -0,0 +1,106 @@
+/*
+ * Copyright 2023 Google Inc. All Rights Reserved.
+ * Use of this source code is governed by a BSD-style
+ * license that can be found in the LICENSE file or at
+ * https://developers.google.com/open-source/licenses/bsd
+ */
+
+using Grpc.Core;
+using System;
+using System.Diagnostics;
+using System.Threading.Tasks;
+
+namespace Google.Api.Gax.Grpc;
+
+internal static class ApiCallTracingExtensions
+{
+    internal const string GrpcCallTypeTag = "grpc.call.type";
+    internal const string UnaryCallType = "unary";
+    internal const string ServerStreamingCallType = "server_streaming";
+    internal const string ClientStreamingCallType = "client_streaming";
+    internal const string BidiStreamingCallType = "bidi_streaming";
+
+    // By design, the code is mostly duplicated between the async and sync calls.
+
+    #region Unary calls
+    // Async call
+    internal static Func<TRequest, CallSettings, Task<TResponse>> WithTracing<TRequest, TResponse>(
+        this Func<TRequest, CallSettings, Task<TResponse>> fn, ActivitySource activitySource, string methodName) =>
+        async (request, callSettings) =>
+        {
+            using var activity = activitySource?.StartActivity($"{methodName}/{fn.Method.Name}");
+            activity?.SetTag(GrpcCallTypeTag, UnaryCallType);
+            try
+            {
+                return await fn(request, callSettings).ConfigureAwait(false);
+            }
+            catch (Exception ex)
+            {
+                activity?.SetException(ex);
+                throw;
+            }
+        };
+
+    // Sync call
+    internal static Func<TRequest, CallSettings, TResponse> WithTracing<TRequest, TResponse>(
+        this Func<TRequest, CallSettings, TResponse> fn, ActivitySource activitySource, string methodName) =>
+        (request, callSettings) =>
+        {
+            using var activity = activitySource?.StartActivity($"{methodName}/{fn.Method.Name}");
+            activity?.SetTag(GrpcCallTypeTag, UnaryCallType);
+            try
+            {
+                return fn(request, callSettings);
+            }
+            catch (Exception ex)
+            {
+                activity?.SetException(ex);
+                throw;
+            }
+        };
+    #endregion
+
+    #region Server streaming
+    // Async call
+    internal static Func<TRequest, CallSettings, Task<AsyncServerStreamingCall<TResponse>>> WithTracing<TRequest, TResponse>(
+        this Func<TRequest, CallSettings, Task<AsyncServerStreamingCall<TResponse>>> fn, ActivitySource activitySource, string methodName) =>
+        async (request, callSettings) =>
+        {
+            using var activity = activitySource?.StartActivity($"{methodName}/{fn.Method.Name}");
+            activity?.SetTag(GrpcCallTypeTag, ServerStreamingCallType);
+            return await fn(request, callSettings).ConfigureAwait(false);
+        };
+
+    // Sync call
+    internal static Func<TRequest, CallSettings, AsyncServerStreamingCall<TResponse>> WithTracing<TRequest, TResponse>(
+        this Func<TRequest, CallSettings, AsyncServerStreamingCall<TResponse>> fn, ActivitySource activitySource, string methodName) =>
+        (request, callSettings) =>
+        {
+            using var activity = activitySource?.StartActivity($"{methodName}/{fn.Method.Name}");
+            activity?.SetTag(GrpcCallTypeTag, ServerStreamingCallType);
+            return fn(request, callSettings);
+        };
+    #endregion
+
+    #region Client streaming
+    internal static Func<CallSettings, AsyncClientStreamingCall<TRequest, TResponse>> WithTracing<TRequest, TResponse>(
+        this Func<CallSettings, AsyncClientStreamingCall<TRequest, TResponse>> fn, ActivitySource activitySource, string methodName) =>
+        callSettings =>
+        {
+            using var activity = activitySource?.StartActivity($"{methodName}/{fn.Method.Name}");
+            activity?.SetTag(GrpcCallTypeTag, ClientStreamingCallType);
+            return fn(callSettings);
+        };
+    #endregion
+
+    #region Bidi streaming
+    internal static Func<CallSettings, AsyncDuplexStreamingCall<TRequest, TResponse>> WithTracing<TRequest, TResponse>(
+        this Func<CallSettings, AsyncDuplexStreamingCall<TRequest, TResponse>> fn, ActivitySource activitySource, string methodName) =>
+        callSettings =>
+        {
+            using var activity = activitySource?.StartActivity($"{methodName}/{fn.Method.Name}");
+            activity?.SetTag(GrpcCallTypeTag, BidiStreamingCallType);
+            return fn(callSettings);
+        };
+    #endregion
+}