diff --git a/src/AzureFunctions.PowerShell.OpenTelemetry.SDK.psd1 b/src/AzureFunctions.PowerShell.OpenTelemetry.SDK.psd1 index 0d90842..7f4c655 100644 --- a/src/AzureFunctions.PowerShell.OpenTelemetry.SDK.psd1 +++ b/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 diff --git a/src/AzureFunctions.PowerShell.OpenTelemetry.SDK/Traces/GetCurrentActivityForInvocation.cs b/src/AzureFunctions.PowerShell.OpenTelemetry.SDK/Traces/GetCurrentActivityForInvocation.cs new file mode 100644 index 0000000..984863f --- /dev/null +++ b/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 +{ + /// + /// 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. + /// + [Cmdlet(VerbsCommon.Get, "CurrentActivityForInvocation")] + [OutputType(typeof(GetActivityResponse))] + public class GetCurrentActivityForInvocation : PSCmdlet + { + /// + /// InvocationId for the current invocation + /// + [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)); + } + } +} diff --git a/src/AzureFunctions.PowerShell.OpenTelemetry.SDK/Traces/StartOpenTelemetryInvocationInternal.cs b/src/AzureFunctions.PowerShell.OpenTelemetry.SDK/Traces/StartOpenTelemetryInvocationInternal.cs index 38173bb..5ba0b0e 100644 --- a/src/AzureFunctions.PowerShell.OpenTelemetry.SDK/Traces/StartOpenTelemetryInvocationInternal.cs +++ b/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); } } diff --git a/src/OpenTelemetryEngine/ResponseObjects/GetActivityResponse.cs b/src/OpenTelemetryEngine/ResponseObjects/GetActivityResponse.cs new file mode 100644 index 0000000..57accce --- /dev/null +++ b/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; + } + } +} diff --git a/src/OpenTelemetryEngine/Traces/FunctionsActivityBuilder.cs b/src/OpenTelemetryEngine/Traces/FunctionsActivityBuilder.cs index 9e4f0ec..acbba31 100644 --- a/src/OpenTelemetryEngine/Traces/FunctionsActivityBuilder.cs +++ b/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); + } } }