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

Author: andystaples

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; }
 
+        /// <summary>
+        /// The invocation id.
+        /// </summary>
+        [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;
             }

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(