get context flowing across dotnet CLI calls

Author: baronfel

src/Cli/dotnet/Activities.cs

@@ -7,4 +7,7 @@ namespace Microsoft.DotNet.Cli;
 public static class Activities
 {
     public static ActivitySource s_source = new("dotnet-cli", Product.Version);
+
+    public const string DOTNET_CLI_TRACEPARENT = nameof(DOTNET_CLI_TRACEPARENT);
+    public const string DOTNET_CLI_TRACESTATE = nameof(DOTNET_CLI_TRACESTATE);
 }

src/Cli/dotnet/CommandFactory/CommandResolution/MuxerCommandResolver.cs

@@ -1,24 +1,54 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-#nullable disable
-
+using System.Diagnostics;
 using Microsoft.DotNet.Cli.Utils;
 using Microsoft.DotNet.Cli.Utils.Extensions;
+using OpenTelemetry;
+using OpenTelemetry.Context.Propagation;
 
 namespace Microsoft.DotNet.Cli.CommandFactory.CommandResolution;
 
 public class MuxerCommandResolver : ICommandResolver
 {
-    public CommandSpec Resolve(CommandResolverArguments commandResolverArguments)
+    public CommandSpec? Resolve(CommandResolverArguments commandResolverArguments)
     {
         if (commandResolverArguments.CommandName == Muxer.MuxerName)
         {
             var muxer = new Muxer();
             var escapedArgs = ArgumentEscaper.EscapeAndConcatenateArgArrayForProcessStart(
                 commandResolverArguments.CommandArguments.OrEmptyIfNull());
-            return new CommandSpec(muxer.MuxerPath, escapedArgs);
+            var env = MakeActivityContextEnvironment();
+            return new CommandSpec(muxer.MuxerPath, escapedArgs, env);
         }
         return null;
     }
+
+    private Dictionary<string, string>? MakeActivityContextEnvironment()
+    {
+        var currentActivity = Activity.Current;
+        var currentBaggage = Baggage.Current;
+        if (currentActivity == null)
+        {
+            return null;
+        }
+        var contextToInject = currentActivity.Context;
+        var propagationContext = new PropagationContext(contextToInject, currentBaggage);
+        var envDict = new Dictionary<string, string>();
+        Propagators.DefaultTextMapPropagator.Inject(propagationContext, envDict, WriteTraceStateIntoEnv);
+        return envDict;
+    }
+
+    private void WriteTraceStateIntoEnv(Dictionary<string, string> dictionary, string key, string value)
+    {
+        switch (key)
+        {
+            case "traceparent":
+                dictionary[Activities.DOTNET_CLI_TRACEPARENT] = value;
+                break;
+            case "tracestate":
+                dictionary[Activities.DOTNET_CLI_TRACESTATE] = value;
+                break;
+        }
+    }
 }

src/Cli/dotnet/Program.cs

@@ -17,6 +17,7 @@
 using Microsoft.Extensions.EnvironmentAbstractions;
 using NuGet.Frameworks;
 using OpenTelemetry;
+using OpenTelemetry.Context.Propagation;
 using OpenTelemetry.Metrics;
 using OpenTelemetry.Resources;
 using OpenTelemetry.Trace;
@@ -54,12 +55,16 @@ public class Program
         .AddOtlpExporter()
         .Build();
 
+    private static ActivityContext s_parentActivityContext = default;
+    private static ActivityKind s_activityKind = ActivityKind.Internal;
+
 
     public static int Main(string[] args)
     {
         // capture the time to we can compute muxer/host startup overhead
         DateTime mainTimeStamp = DateTime.Now;
-        using var _mainActivity = Activities.s_source.StartActivity("main");
+        (s_parentActivityContext, s_activityKind) = DeriveParentActivityContextFromEnv();
+        using var _mainActivity = Activities.s_source.StartActivity("main", s_activityKind, s_parentActivityContext);
         using AutomaticEncodingRestorer _encodingRestorer = new();
 
         // Setting output encoding is not available on those platforms
@@ -117,9 +122,44 @@ public static int Main(string[] args)
         }
     }
 
+    /// <summary>
+    /// uses the OpenTelemetrySDK's Propagation API to derive the parent activity context and kind
+    /// from the DOTNET_CLI_TRACEPARENT and DOTNET_CLI_TRACESTATE environment variables.
+    /// </summary>
+    /// <returns></returns>
+    private static (ActivityContext parentActivityContext, ActivityKind kind) DeriveParentActivityContextFromEnv()
+    {
+        var traceParent = Env.GetEnvironmentVariable(Activities.DOTNET_CLI_TRACEPARENT);
+        var traceState = Env.GetEnvironmentVariable(Activities.DOTNET_CLI_TRACESTATE);
+        static IEnumerable<string> GetValueFromCarrier(Dictionary<string, IEnumerable<string>> carrier, string key)
+        {
+            return carrier.TryGetValue(key, out var value) ? value : Enumerable.Empty<string>();
+        }
+
+        if (string.IsNullOrEmpty(traceParent))
+        {
+            return (default, ActivityKind.Internal);
+        }
+        var carriermap = new Dictionary<string, IEnumerable<string>>
+        {
+            { "traceparent", [traceParent] },
+            { "tracestate", [traceState] }
+        };
+
+        // Use the OpenTelemetry Propagator to extract the parent activity context and kind. For some reason this isn't set by the OTel SDK like docs say it should be.
+        Sdk.SetDefaultTextMapPropagator(new CompositeTextMapPropagator([
+            new TraceContextPropagator(),
+            new BaggagePropagator()
+        ]));
+        var parentActivityContext = Propagators.DefaultTextMapPropagator.Extract(default, carriermap, GetValueFromCarrier);
+        var kind = parentActivityContext.ActivityContext.IsRemote ? ActivityKind.Server : ActivityKind.Internal;
+
+        return (parentActivityContext.ActivityContext, kind);
+    }
+
     private static void TrackHostStartup(DateTime mainTimeStamp)
     {
-        var hostStartupActivity = Activities.s_source.CreateActivity("host-startup", ActivityKind.Server);
+        var hostStartupActivity = Activities.s_source.CreateActivity("host-startup", s_activityKind, s_parentActivityContext);
         hostStartupActivity?.SetStartTime(Process.GetCurrentProcess().StartTime);
         hostStartupActivity?.SetEndTime(mainTimeStamp);
         hostStartupActivity?.SetStatus(ActivityStatusCode.Ok);
@@ -138,23 +178,51 @@ private static void SetupMSBuildEnvironmentInvariants()
         }
     }
 
-    internal static int ProcessArgs(string[] args)
+    private static string GetCommandName(ParseResult r)
+    {
+        // walk the parent command tree to find the top-level command name and get the full command name for this parseresult
+        List<string> parentNames = [r.CommandResult.Command.Name];
+        var current = r.CommandResult.Parent;
+        while (current is CommandResult parentCommandResult)
+        {
+            parentNames.Add(parentCommandResult.Command.Name);
+            current = parentCommandResult.Parent;
+        }
+        parentNames.Reverse();
+        return string.Join(' ', parentNames);
+    }
+    private static void SetDisplayName(Activity activity, ParseResult parseResult)
+    {
+        if (activity == null)
+        {
+            return;
+        }
+        var name = GetCommandName(parseResult);
+        // Set the display name to the full command name
+        activity.DisplayName = name;
+
+        // Set the command name as an attribute for better filtering in telemetry
+        activity.SetTag("command.name", name);
+    }
+
+    internal static int ProcessArgs(string[] args, ITelemetry telemetryClient = null)
     {
         ParseResult parseResult;
-        using (var _parseActivity = Activities.s_source.StartActivity("parse"))
+        using (var _parseActivity = Activities.s_source.StartActivity("parse", s_activityKind, s_parentActivityContext))
         {
             // If we get C# file path as the first argument, parse as `dotnet run file.cs`.
             parseResult = args is [{ } filePath, ..] && VirtualProjectBuildingCommand.IsValidEntryPointPath(filePath)
                 ? Parser.Instance.Parse(["run", .. args])
                 : Parser.Instance.Parse(args);
 
+            SetDisplayName(_parseActivity, parseResult);
             // Avoid create temp directory with root permission and later prevent access in non sudo
             // This method need to be run very early before temp folder get created
             // https://github.com/dotnet/sdk/issues/20195
             SudoEnvironmentDirectoryOverride.OverrideEnvironmentVariableToTmp(parseResult);
         }
 
-        using (var _firstTimeUseActivity = Activities.s_source.StartActivity("first-time-use"))
+        using (var _firstTimeUseActivity = Activities.s_source.StartActivity("first-time-use", s_activityKind, s_parentActivityContext))
         {
             IFirstTimeUseNoticeSentinel firstTimeUseNoticeSentinel = new FirstTimeUseNoticeSentinel();
 
@@ -221,17 +289,7 @@ internal static int ProcessArgs(string[] args)
         int exitCode;
         if (parseResult.CanBeInvoked())
         {
-            using var _invocationActivity = Activities.s_source.StartActivity("invocation");
-            // walk the parent command tree to find the top-level command name and get the full command name for this parseresult
-            List<string> parentNames = [parseResult.CommandResult.Command.Name];
-            var current = parseResult.CommandResult.Parent;
-            while (current is CommandResult parentCommandResult)
-            {
-                parentNames.Add(parentCommandResult.Command.Name);
-                current = parentCommandResult.Parent;
-            }
-            parentNames.Reverse();
-            _invocationActivity?.DisplayName = string.Join(' ', parentNames);
+            using var _invocationActivity = Activities.s_source.StartActivity("invocation", s_activityKind, s_parentActivityContext);
             try
             {
                 exitCode = parseResult.Invoke();
@@ -246,14 +304,14 @@ internal static int ProcessArgs(string[] args)
         {
             try
             {
-                var _lookupExternalCommandActivity = Activities.s_source.StartActivity("lookup-external-command");
+                var _lookupExternalCommandActivity = Activities.s_source.StartActivity("lookup-external-command", s_activityKind, s_parentActivityContext);
                 var resolvedCommand = CommandFactoryUsingResolver.Create(
                         "dotnet-" + parseResult.GetValue(Parser.DotnetSubCommand),
                         args.GetSubArguments(),
                         FrameworkConstants.CommonFrameworks.NetStandardApp15);
                 _lookupExternalCommandActivity?.Dispose();
 
-                var _executionActivity = Activities.s_source.StartActivity("execute-extensible-command");
+                var _executionActivity = Activities.s_source.StartActivity("execute-extensible-command", s_activityKind, s_parentActivityContext);
                 var result = resolvedCommand.Execute();
                 _executionActivity?.Dispose();