diff --git a/src/DurableSDK/Commands/SetFunctionInvocationContextCommand.cs b/src/DurableSDK/Commands/SetFunctionInvocationContextCommand.cs index 943e8362..c7a05d76 100644 --- a/src/DurableSDK/Commands/SetFunctionInvocationContextCommand.cs +++ b/src/DurableSDK/Commands/SetFunctionInvocationContextCommand.cs @@ -18,6 +18,7 @@ public class SetFunctionInvocationContextCommand : PSCmdlet { internal const string ContextKey = "OrchestrationContext"; private const string DurableClientKey = "DurableClient"; + private const string InvocationIdKey = "InvocationId"; [Parameter(Mandatory = true, ParameterSetName = ContextKey)] public OrchestrationContext OrchestrationContext { get; set; } @@ -28,6 +29,12 @@ public class SetFunctionInvocationContextCommand : PSCmdlet [Parameter(Mandatory = true, ParameterSetName = DurableClientKey)] public object DurableClient { get; set; } + /// + /// The invocation id. + /// + [Parameter(Mandatory = true, ParameterSetName = InvocationIdKey)] + public string InvocationId { get; set; } + [Parameter(Mandatory = true, ParameterSetName = "Clear")] public SwitchParameter Clear { get; set; } @@ -44,11 +51,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; } diff --git a/src/Modules/Microsoft.Azure.Functions.PowerShellWorker/Microsoft.Azure.Functions.PowerShellWorker.psm1 b/src/Modules/Microsoft.Azure.Functions.PowerShellWorker/Microsoft.Azure.Functions.PowerShellWorker.psm1 index 7c1c5396..816ff957 100644 --- a/src/Modules/Microsoft.Azure.Functions.PowerShellWorker/Microsoft.Azure.Functions.PowerShellWorker.psm1 +++ b/src/Modules/Microsoft.Azure.Functions.PowerShellWorker/Microsoft.Azure.Functions.PowerShellWorker.psm1 @@ -19,6 +19,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( @@ -121,6 +131,9 @@ function Start-DurableOrchestration { $InstanceId = (New-Guid).Guid } + $invocationId = GetInvocationIdFromModulePrivateData + $headers = Get-TraceHeaders -InvocationId $invocationId + $Uri = if ($DurableClient.rpcBaseUrl) { # Fast local RPC path @@ -132,12 +145,50 @@ function Start-DurableOrchestration { } $Body = $InputObject | ConvertTo-Json -Compress - - $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(