Add Invocation Id to Set-FunctionInvocationContext and propagate trace information in HTTP requests to correlate client and orchestration functions (#99)

Author: bachuv

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(

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";
 
         /// <summary>
         /// The orchestration context.
@@ -36,6 +37,12 @@ public class SetFunctionInvocationContextCommand : PSCmdlet
         [Parameter(Mandatory = true, ParameterSetName = DurableClientKey)]
         public object DurableClient { get; set; }
 
+        /// <summary>
+        /// The function invocation ID.
+        /// </summary>
+        [Parameter(Mandatory = true, ParameterSetName = InvocationIdKey)]
+        public string InvocationId { get; set; }
+
         /// <summary>
         /// Whether or not to clear the privateData of this module.
         /// </summary>
@@ -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;
             }