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

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 {
+        $PrivateData['InvocationId']
+    }
+}
+
 function Get-DurableStatus {
     [CmdletBinding()]
     param(
@@ -114,6 +124,8 @@ function Start-DurableOrchestration {
         }
 
     $Body = $InputObject | ConvertTo-Json -Compress -Depth 100
+
+    $InvocationId = GetInvocationIdFromModulePrivateData
 
     $null = Invoke-RestMethod -Uri $Uri -Method 'POST' -ContentType 'application/json' -Body $Body
 

src/DurableEngine/Utilities/IPowerShellServices.cs

@@ -13,6 +13,8 @@ public interface IPowerShellServices
     {
         void SetDurableClient(object durableClient);
 
+        void SetInvocationId(string invocationId);
+
         void SetOrchestrationContext(OrchestrationContext orchestrationContext);
 
         void ClearOrchestrationContext();

src/DurableEngine/Utilities/PowerShellServices.cs

@@ -45,6 +45,15 @@ public void SetDurableClient(object durableClient)
             _hasSetOrchestrationContext = true;
         }
 
+        public void SetInvocationId(string invocationId)
+        {
+            _pwsh.AddCommand(SetFunctionInvocationContextCommand)
+                .AddParameter("InvocationId", invocationId)
+                .InvokeAndClearCommands();
+
+            _hasSetOrchestrationContext = true;
+        }
+
         public void SetOrchestrationContext(OrchestrationContext orchestrationContext)
         {
             _pwsh.AddCommand(SetFunctionInvocationContextCommand)

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;
             }