Add invocation id to module private data and add a Get-CurrentActivityForInvocation cmdlet (#25)

Author: bachuv

src/AzureFunctions.PowerShell.OpenTelemetry.SDK.psd1

@@ -70,7 +70,8 @@
         'Stop-OpenTelemetryInvocationInternal',
         'New-FunctionsOpenTelemetryTracerBuilder',
         'Start-FunctionsOpenTelemetrySpan',
-        'Stop-FunctionsOpenTelemetrySpan'
+        'Stop-FunctionsOpenTelemetrySpan',
+        'Get-CurrentActivityForInvocation'
     )
 
     # Variables to export from this module

src/AzureFunctions.PowerShell.OpenTelemetry.SDK/Traces/GetCurrentActivityForInvocation.cs

@@ -0,0 +1,33 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System.Management.Automation;
+using OpenTelemetryEngine.Traces;
+using OpenTelemetryEngine.ResponseObjects;
+
+namespace AzureFunctions.PowerShell.OpenTelemetry.SDK
+{
+    /// <summary>
+    /// This cmdlet is used to start an OpenTelemetry invocation.
+    /// Currently, that just means starting a (hidden) span which is a copy of the host span. 
+    /// This allows all new spans created by the user or by dependent modules to link back to the host span using TraceParent. 
+    /// </summary>
+    [Cmdlet(VerbsCommon.Get, "CurrentActivityForInvocation")]
+    [OutputType(typeof(GetActivityResponse))]
+    public class GetCurrentActivityForInvocation : PSCmdlet
+    {
+        /// <summary>
+        /// InvocationId for the current invocation
+        /// </summary>
+        [Parameter(Mandatory = true)]
+        public string InvocationId { get; set; } = string.Empty;
+
+        // This method will be called for each input received from the pipeline to this cmdlet; if no input is received, this method is not called
+        protected override void ProcessRecord()
+        {
+            WriteObject(FunctionsActivityBuilder.GetActivityForInvocation(InvocationId));
+        }
+    }
+}

src/AzureFunctions.PowerShell.OpenTelemetry.SDK/Traces/StartOpenTelemetryInvocationInternal.cs

@@ -6,6 +6,7 @@
 using System.Management.Automation;
 using OpenTelemetryEngine.Traces;
 using OpenTelemetryEngine.ResponseObjects;
+using System.Management.Automation.Runspaces;
 
 namespace AzureFunctions.PowerShell.OpenTelemetry.SDK
 {
@@ -42,6 +43,15 @@ public class StartOpenTelemetryInvocationInternal : PSCmdlet
         protected override void ProcessRecord()
         {
             var newActivity = FunctionsActivityBuilder.StartInternalActivity(InvocationId, TraceParent, TraceState);
+
+            try
+            {
+                string script = "param($id) Set-FunctionInvocationContext -InvocationId $id";
+                _ = this.InvokeCommand.InvokeScript(script, false, PipelineResultTypes.Output, null, new object[] { InvocationId });
+            }
+            catch (ParameterBindingException) { }
+            catch (CommandNotFoundException) { }
+
             WriteObject(newActivity);
         }
     }

src/OpenTelemetryEngine/ResponseObjects/GetActivityResponse.cs

@@ -0,0 +1,14 @@
+using System.Diagnostics;
+
+namespace OpenTelemetryEngine.ResponseObjects
+{
+    public class GetActivityResponse
+    {
+        public Activity? activity;
+
+        public GetActivityResponse(Activity? activity)
+        {
+            this.activity = activity;
+        }
+    }
+}

src/OpenTelemetryEngine/Traces/FunctionsActivityBuilder.cs

@@ -94,6 +94,15 @@ public static void StopActivity(FunctionsActivityResponse? response)
         {
             response?.activity?.Stop();
         }
+
+        public static GetActivityResponse GetActivityForInvocation(string invocationId)
+        {
+            if (internalActivitiesByInvocationId.ContainsKey(invocationId))
+            {
+                return new GetActivityResponse(internalActivitiesByInvocationId[invocationId]);
+            }
+            return new GetActivityResponse(Activity.Current);
+        }
     }
 }