diff --git a/src/AzureFunctions.PowerShell.Durable.SDK.psm1 b/src/AzureFunctions.PowerShell.Durable.SDK.psm1 index 4806a49..8c6dc51 100644 --- a/src/AzureFunctions.PowerShell.Durable.SDK.psm1 +++ b/src/AzureFunctions.PowerShell.Durable.SDK.psm1 @@ -22,6 +22,16 @@ function GetDurableClientFromModulePrivateData { } } +function GetInvocationIdFromModulePrivateData { + $PrivateData = $PSCmdlet.MyInvocation.MyCommand.Module.PrivateData + if ($null -eq $PrivateData -or $null -eq $PrivateData['InvocationId']) { + return $null + } + else { + return $PrivateData['InvocationId'] + } +} + function Get-DurableStatus { [CmdletBinding()] param( @@ -103,6 +113,9 @@ function Start-DurableOrchestration { $InstanceId = (New-Guid).Guid } + $invocationId = GetInvocationIdFromModulePrivateData + $headers = Get-TraceHeaders -InvocationId $invocationId + $Uri = if ($DurableClient.rpcBaseUrl) { # Fast local RPC path @@ -115,11 +128,49 @@ function Start-DurableOrchestration { $Body = $InputObject | ConvertTo-Json -Compress -Depth 100 - $null = Invoke-RestMethod -Uri $Uri -Method 'POST' -ContentType 'application/json' -Body $Body + $null = Invoke-RestMethod -Uri $Uri -Method 'POST' -ContentType 'application/json' -Body $Body -Headers $headers return $instanceId } +function Get-TraceHeaders { + param( + [string] $InvocationId + ) + + if ($null -eq $InvocationId -or $InvocationId -eq "") { + return @{} # Return an empty headers object + } + + # Check if Get-CurrentActivityForInvocation is available + if (-not (Get-Command -Name Get-CurrentActivityForInvocation -ErrorAction SilentlyContinue)) { + Write-Warning "Get-CurrentActivityForInvocation is not available. Skipping call." + return @{} # Return an empty headers object + } + + $activityResponse = Get-CurrentActivityForInvocation -InvocationId $invocationId + $activity = $activityResponse.activity + + $traceId = $activity.TraceId + $spanId = $activity.SpanId + $traceFlags = $activity.TraceFlags + $traceState = $activity.TraceStateString + + $flag = "00" + if ($null -ne $traceFlags -and $traceFlags -eq "Recorded") { + $flag = "01" + } + + $traceparent = "00-$traceId-$spanId-$flag" + + $headers = @{ + "traceparent" = $traceparent + "tracestate" = $traceState + } + + return $headers +} + function Stop-DurableOrchestration { [CmdletBinding()] param( diff --git a/src/DurableSDK/Commands/Internals/SetFunctionInvocationContextCommand.cs b/src/DurableSDK/Commands/Internals/SetFunctionInvocationContextCommand.cs index 2e3bc7f..a4381ea 100644 --- a/src/DurableSDK/Commands/Internals/SetFunctionInvocationContextCommand.cs +++ b/src/DurableSDK/Commands/Internals/SetFunctionInvocationContextCommand.cs @@ -23,6 +23,7 @@ public class SetFunctionInvocationContextCommand : PSCmdlet { private const string ContextKey = "OrchestrationContext"; private const string DurableClientKey = "DurableClient"; + private const string InvocationIdKey = "InvocationId"; /// /// The orchestration context. @@ -36,6 +37,12 @@ public class SetFunctionInvocationContextCommand : PSCmdlet [Parameter(Mandatory = true, ParameterSetName = DurableClientKey)] public object DurableClient { get; set; } + /// + /// The function invocation ID. + /// + [Parameter(Mandatory = true, ParameterSetName = InvocationIdKey)] + public string InvocationId { get; set; } + /// /// Whether or not to clear the privateData of this module. /// @@ -70,11 +77,16 @@ protected override void EndProcessing() privateData[DurableClientKey] = DurableClient; break; + case InvocationIdKey: + privateData[InvocationIdKey] = InvocationId; + break; + default: if (Clear.IsPresent) { privateData.Remove(ContextKey); privateData.Remove(DurableClientKey); + privateData.Remove(InvocationIdKey); } break; }