Open Telemetry prototype

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

@@ -8,7 +8,10 @@
 using Google.Protobuf;
 using Google.Protobuf.Reflection;
 using Grpc.Core;
+using Microsoft.Extensions.Logging;
 using System;
+using System.Diagnostics;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
@@ -217,6 +220,49 @@ public async Task WithLogging_Async()
             Assert.All(entries, entry => Assert.Contains("SimpleMethod", entry.Message));
         }
 
+        [Fact]
+        public void WithTracing_Sync()
+        {
+            using var helper = new ActivityHelper();
+            var call = new ApiCall<SimpleRequest, SimpleResponse>(
+                "SimpleMethod",
+                (req, cs) => Task.FromResult(default(SimpleResponse)),
+                (req, cs) => null,
+                null).WithTracing(helper.Source);
+
+            call.Sync(new SimpleRequest(), null);
+
+            var activity = helper.CapturedActivity;
+            Assert.NotNull(activity);
+            Assert.Contains("SimpleMethod", activity.OperationName);
+            Assert.Equal(ActivityStatusCode.Ok, activity.Status);
+        }
+
+        private class ActivityHelper : IDisposable
+        {
+            public ActivitySource Source { get; }
+            public Activity CapturedActivity { get; private set; }
+            private readonly ActivityListener _listener;
+
+            internal ActivityHelper([CallerMemberName] string name = null)
+            {
+                Source = new ActivitySource(name);
+                _listener = new ActivityListener
+                {
+                    ShouldListenTo = candidate => candidate == Source,
+                    Sample = (ref ActivityCreationOptions<ActivityContext> options) => ActivitySamplingResult.AllDataAndRecorded,
+                    ActivityStarted = activity => CapturedActivity = activity
+                };
+                ActivitySource.AddActivityListener(_listener);
+            }
+
+            public void Dispose()
+            {
+                _listener.Dispose();
+                Source.Dispose();
+            }
+        }
+
         internal class ExtractedRequestParamRequest : IMessage<ExtractedRequestParamRequest>
         {
             public string TableName { get; set; }

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
 {
@@ -83,5 +84,9 @@ logger is null
                 ? this
                 : new ApiBidirectionalStreamingCall<TRequest, TResponse>(_methodName, _call.WithLogging(logger, _methodName), BaseCallSettings, StreamingSettings);
 
+        internal ApiBidirectionalStreamingCall<TRequest, TResponse> WithTracing(ActivitySource activitySource) =>
+            activitySource is null
+                ? this
+                : new ApiBidirectionalStreamingCall<TRequest, TResponse>(_methodName, _call.WithTracing(activitySource, _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
@@ -173,6 +174,14 @@ logger is null
                     _syncCall.WithLogging(logger, _methodName),
                     BaseCallSettings);
 
+        internal ApiCall<TRequest, TResponse> WithTracing(ActivitySource activitySource) =>
+            activitySource is null
+                ? this
+                : new ApiCall<TRequest, TResponse>(
+                    _methodName, _asyncCall.WithTracing(activitySource, _methodName),
+                    _syncCall.WithTracing(activitySource, _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,138 @@
+/*  
+ * Copyright 2024 Google LLC 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;
+
+/// <summary>
+/// Extension methods to provide tracing via <see cref="ActivitySource"/>.
+/// </summary>
+internal static class ApiCallTracingExtensions
+{
+    internal const string AttributeExceptionEventName = "exception";
+    internal const string AttributeExceptionType = "exception.type";
+    internal const string AttributeExceptionMessage = "exception.message";
+    internal const string AttributeExceptionStacktrace = "exception.stacktrace";
+
+    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";
+
+    // Unary async
+    internal static Func<TRequest, CallSettings, Task<TResponse>> WithTracing<TRequest, TResponse>(
+        this Func<TRequest, CallSettings, Task<TResponse>> fn, ActivitySource activitySource, string methodName)
+    {
+        GaxPreconditions.CheckNotNull(activitySource, nameof(activitySource));
+        var activityName = FormatActivityName(fn, methodName);
+        return async (request, callSettings) =>
+        {
+            using var activity = activitySource.StartActivity(activityName, ActivityKind.Client);
+            activity?.SetTag(GrpcCallTypeTag, UnaryCallType);
+            // TODO: Add a tag with the name of the client, in case a custom source has been provided?
+            try
+            {
+                var response = await fn(request, callSettings).ConfigureAwait(false);
+                activity?.SetStatus(ActivityStatusCode.Ok);
+                return response;
+            }
+            catch (Exception ex) when (SetActivityException(activity, ex))
+            {
+                // We'll never actually get here, because SetActivityException always returns false.
+                // Alternative: catch without an exception filter, make SetActivityException return void,
+                // and call ExceptionDispatchInfo.Capture(e).Throw();.
+                throw;
+            }
+        };
+    }
+
+    // Unary sync
+    internal static Func<TRequest, CallSettings, TResponse> WithTracing<TRequest, TResponse>(
+        this Func<TRequest, CallSettings, TResponse> fn, ActivitySource activitySource, string methodName)
+    {
+        GaxPreconditions.CheckNotNull(activitySource, nameof(activitySource));
+        var activityName = FormatActivityName(fn, methodName);
+        return (request, callSettings) =>
+        {
+            using var activity = activitySource.StartActivity(activityName, ActivityKind.Client);
+            activity?.SetTag(GrpcCallTypeTag, UnaryCallType);
+            // TODO: Add a tag with the name of the client, in case a custom source has been provided?
+            try
+            {
+                var response = fn(request, callSettings);
+                activity?.SetStatus(ActivityStatusCode.Ok);
+                return response;
+            }
+            catch (Exception ex) when (SetActivityException(activity, ex))
+            {
+                // We'll never actually get here, because SetActivityException always returns false.
+                // Alternative: catch without an exception filter, make SetActivityException return void,
+                // and call ExceptionDispatchInfo.Capture(e).Throw();.
+                // (As this is a sync method we may be okay just to use a throw statement. But we need to
+                // validate that we're not losing any info.)
+                throw;
+            }
+        };
+    }
+
+    // Server-streaming async
+    internal static Func<TRequest, CallSettings, Task<AsyncServerStreamingCall<TResponse>>> WithTracing<TRequest, TResponse>(
+        this Func<TRequest, CallSettings, Task<AsyncServerStreamingCall<TResponse>>> fn, ActivitySource activitySource, string methodName) =>
+        fn;
+
+    // Server-streaming sync
+    internal static Func<TRequest, CallSettings, AsyncServerStreamingCall<TResponse>> WithTracing<TRequest, TResponse>(
+        this Func<TRequest, CallSettings, AsyncServerStreamingCall<TResponse>> fn, ActivitySource activitySource, string methodName) =>
+        fn;
+
+    // Client-streaming
+    internal static Func<CallSettings, AsyncClientStreamingCall<TRequest, TResponse>> WithTracing<TRequest, TResponse>(
+        this Func<CallSettings, AsyncClientStreamingCall<TRequest, TResponse>> fn, ActivitySource activitySource, string methodName) =>
+        fn;
+
+    // Bidi-streaming
+    internal static Func<CallSettings, AsyncDuplexStreamingCall<TRequest, TResponse>> WithTracing<TRequest, TResponse>(
+        this Func<CallSettings, AsyncDuplexStreamingCall<TRequest, TResponse>> fn, ActivitySource activitySource, string methodName) =>
+        fn;
+
+    // This is still very much up in the air, and may even require changes to the parameters, so that we get more information
+    // (e.g. the full RPC name, the client name etc).
+    private static string FormatActivityName(Delegate fn, string methodName) => $"{fn.Method.Name}/{methodName}";
+
+    // TODO: See if there's a standard way of doing this. It seems odd to have to do it ourselves.
+    /// <summary>
+    /// Sets an exception within an activity. We may wish to expose this publicly for integration purposes.
+    /// This always returns false, so that it can be used as an exception filter.
+    /// </summary>
+    private static bool SetActivityException(Activity activity, Exception ex)
+    {
+        if (ex is null || activity is null)
+        {
+            return false;
+        }
+
+        var tagsCollection = new ActivityTagsCollection
+        {
+            { AttributeExceptionType, ex.GetType().FullName },
+            { AttributeExceptionStacktrace, ex.ToString() },
+        };
+
+        if (!string.IsNullOrWhiteSpace(ex.Message))
+        {
+            tagsCollection[AttributeExceptionMessage] = ex.Message;
+        }
+
+        activity.SetStatus(ActivityStatusCode.Error, ex.Message);
+        activity.AddEvent(new ActivityEvent(AttributeExceptionEventName, default, tagsCollection));
+        return false;
+    }
+}

Google.Api.Gax.Grpc/ApiClientStreamingCall.cs

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

Google.Api.Gax.Grpc/ApiServerStreamingCall.cs

@@ -8,6 +8,7 @@
 using Grpc.Core;
 using Microsoft.Extensions.Logging;
 using System;
+using System.Diagnostics;
 using System.Threading.Tasks;
 
 namespace Google.Api.Gax.Grpc
@@ -144,5 +145,14 @@ logger is null
                     _asyncCall.WithLogging(logger, _methodName),
                     _syncCall.WithLogging(logger, _methodName),
                     BaseCallSettings);
+
+        internal ApiServerStreamingCall<TRequest, TResponse> WithTracing(ActivitySource activitySource) =>
+            activitySource is null
+                ? this
+                : new ApiServerStreamingCall<TRequest, TResponse>(
+                    _methodName,
+                    _asyncCall.WithTracing(activitySource, _methodName),
+                    _syncCall.WithTracing(activitySource, _methodName),
+                    BaseCallSettings);
     }
 }

Google.Api.Gax.Grpc/ClientHelper.cs

@@ -9,6 +9,7 @@
 using Grpc.Core;
 using Microsoft.Extensions.Logging;
 using System;
+using System.Diagnostics;
 
 namespace Google.Api.Gax.Grpc
 {
@@ -28,6 +29,8 @@ public class ClientHelper
         /// </summary>
         private readonly CallSettings _versionCallSettings;
 
+        private readonly ActivitySource _activitySource;
+
         /// <summary>
         /// Constructs a helper from the given settings.
         /// Behavior is undefined if settings are changed after construction.
@@ -61,6 +64,7 @@ public ClientHelper(Options options)
             {
                 _versionCallSettings = _versionCallSettings.WithHeader(ApiVersionHeaderName, options.ApiVersion);
             }
+            _activitySource = settings.ActivitySource ?? options.ActivitySource;
         }
 
         /// <summary>
@@ -104,6 +108,7 @@ public ApiCall<TRequest, TResponse> BuildApiCall<TRequest, TResponse>(
             // These operations are applied in reverse order.
             // I.e. Version header is added first, then retry is performed.
             return ApiCall.Create(methodName, asyncGrpcCall, syncGrpcCall, baseCallSettings, Clock)
+                .WithTracing(_activitySource)
                 .WithLogging(Logger)
                 .WithRetry(Clock, Scheduler, Logger)
                 .WithMergedBaseCallSettings(_versionCallSettings);
@@ -128,6 +133,7 @@ public ApiServerStreamingCall<TRequest, TResponse> BuildApiCall<TRequest, TRespo
             // These operations are applied in reverse order.
             // I.e. Version header is added first, then retry is performed.
             return ApiServerStreamingCall.Create(methodName, grpcCall, baseCallSettings, Clock)
+                .WithTracing(_activitySource)
                 .WithLogging(Logger)
                 .WithMergedBaseCallSettings(_versionCallSettings);
         }
@@ -152,6 +158,7 @@ public ApiBidirectionalStreamingCall<TRequest, TResponse> BuildApiCall<TRequest,
         {
             CallSettings baseCallSettings = _clientCallSettings.MergedWith(perMethodCallSettings);
             return ApiBidirectionalStreamingCall.Create(methodName, grpcCall, baseCallSettings, streamingSettings, Clock)
+                .WithTracing(_activitySource)
                 .WithLogging(Logger)
                 .WithMergedBaseCallSettings(_versionCallSettings);
         }
@@ -176,6 +183,7 @@ public ApiClientStreamingCall<TRequest, TResponse> BuildApiCall<TRequest, TRespo
         {
             CallSettings baseCallSettings = _clientCallSettings.MergedWith(perMethodCallSettings);
             return ApiClientStreamingCall.Create(methodName, grpcCall, baseCallSettings, streamingSettings, Clock)
+                .WithTracing(_activitySource)
                 .WithLogging(Logger)
                 .WithMergedBaseCallSettings(_versionCallSettings);
         }
@@ -204,6 +212,12 @@ public sealed class Options
             /// The API version to send in the x-goog-api-version header, if any. This may be null.
             /// </summary>
             public string ApiVersion { get; set; }
+
+            /// <summary>
+            /// The activity source to use for tracing, if any. This may be null. This is ignored
+            /// if <see cref="Settings"/> specifies an activity source.
+            /// </summary>
+            public ActivitySource ActivitySource { get; set; }
         }
     }
 }

Google.Api.Gax.Grpc/ServiceSettingsBase.cs

@@ -7,6 +7,7 @@
 
 using Grpc.Core;
 using Grpc.Core.Interceptors;
+using System.Diagnostics;
 
 namespace Google.Api.Gax.Grpc
 {
@@ -42,6 +43,7 @@ protected ServiceSettingsBase(ServiceSettingsBase existing)
             Scheduler = existing.Scheduler;
             VersionHeaderBuilder = existing.VersionHeaderBuilder.Clone();
             Interceptor = existing.Interceptor;
+            ActivitySource = existing.ActivitySource;
         }
 
         /// <summary>
@@ -84,5 +86,11 @@ protected ServiceSettingsBase(ServiceSettingsBase existing)
         /// on or after that date are aware of this property.
         /// </summary>
         public Interceptor Interceptor { get; set; }
+
+        /// <summary>
+        /// An optional <see cref="ActivitySource"/>, which can override the default activity source used for tracing
+        /// calls from this client.
+        /// </summary>
+        public ActivitySource ActivitySource { get; set; }
     }
 }