add additional data to entity work item result, when using with DTS

Author: sebastianburckhardt

eng/proto

@@ -513,6 +513,67 @@ internal static OrchestrationStatus ToCore(this P.OrchestrationStatus status)
         };
     }
 
+    /// <summary>
+    /// Converts a <see cref="P.EntityRequest" /> to a <see cref="EntityBatchRequest" />.
+    /// </summary>
+    /// <param name="entityRequest">The entity request to convert.</param>
+    /// <param name="responseTargetInstanceIds">A list containing the response target instance ids for calls.</param>
+    /// <param name="requestIds">A list containing th request ids of the operations.</param>
+    /// <returns>The converted entity batch request.</returns>
+    [return: NotNullIfNotNull(nameof(entityRequest))]
+    internal static EntityBatchRequest? ToEntityBatchRequest(
+        this P.EntityRequest? entityRequest,
+        out List<P.OrchestrationInstance?> responseTargetInstanceIds,
+        out List<string> requestIds)
+    {
+        if (entityRequest == null)
+        {
+            responseTargetInstanceIds = [];
+            requestIds = [];
+            return null;
+        }
+
+        var batchRequest = new EntityBatchRequest()
+        {
+            EntityState = entityRequest.EntityState,
+            InstanceId = entityRequest.InstanceId,
+        };
+        responseTargetInstanceIds = new(entityRequest.OperationRequests.Count);
+        requestIds = new(entityRequest.OperationRequests.Count);
+
+        foreach (P.HistoryEvent? op in entityRequest.OperationRequests)
+        {
+            if (op.EntityOperationSignaled is not null)
+            {
+                batchRequest.Operations!.Add(new OperationRequest
+                {
+                    Id = Guid.Parse(op.EntityOperationSignaled.RequestId),
+                    Operation = op.EntityOperationSignaled.Operation,
+                    Input = op.EntityOperationSignaled.Input,
+                });
+                responseTargetInstanceIds.Add(null);
+                requestIds.Add(op.EntityOperationSignaled.RequestId);
+            }
+            else if (op.EntityOperationCalled is not null)
+            {
+                batchRequest.Operations!.Add(new OperationRequest
+                {
+                    Id = Guid.Parse(op.EntityOperationCalled.RequestId),
+                    Operation = op.EntityOperationCalled.Operation,
+                    Input = op.EntityOperationCalled.Input,
+                });
+                responseTargetInstanceIds.Add(new P.OrchestrationInstance
+                {
+                    InstanceId = op.EntityOperationCalled.ParentInstanceId,
+                    ExecutionId = op.EntityOperationCalled.ParentExecutionId,
+                });
+                requestIds.Add(op.EntityOperationCalled.RequestId);
+            }
+        }
+
+        return batchRequest;
+    }
+
     /// <summary>
     /// Converts a <see cref="P.OperationRequest" /> to a <see cref="OperationRequest" />.
     /// </summary>

src/Shared/Grpc/ProtoUtils.cs

@@ -183,7 +183,25 @@ async Task ProcessWorkItemsAsync(AsyncServerStreamingCall<P.WorkItem> stream, Ca
                     {
                         this.RunBackgroundTask(
                             workItem,
-                            () => this.OnRunEntityBatchAsync(workItem.EntityRequest, workItem.CompletionToken));
+                            () => this.OnRunEntityBatchAsync(
+                                workItem.EntityRequest.ToEntityBatchRequest(),
+                                workItem.CompletionToken,
+                                responseDestinations: null,
+                                requestIds: null));
+                    }
+                    else if (workItem.RequestCase == P.WorkItem.RequestOneofCase.EntityRequestV2)
+                    {
+                        EntityBatchRequest batchRequest = workItem.EntityRequestV2.ToEntityBatchRequest(
+                            out List<P.OrchestrationInstance?> responseDestinations,
+                            out List<string> requestIds);
+
+                        this.RunBackgroundTask(
+                             workItem,
+                             () => this.OnRunEntityBatchAsync(
+                                workItem.EntityRequest.ToEntityBatchRequest(),
+                                workItem.CompletionToken,
+                                responseDestinations,
+                                requestIds));
                     }
                     else if (workItem.RequestCase == P.WorkItem.RequestOneofCase.HealthPing)
                     {
@@ -383,14 +401,17 @@ async Task OnRunActivityAsync(P.ActivityRequest request, string completionToken)
             await this.sidecar.CompleteActivityTaskAsync(response);
         }
 
-        async Task OnRunEntityBatchAsync(P.EntityBatchRequest request, string completionToken)
+        async Task OnRunEntityBatchAsync(
+            EntityBatchRequest batchRequest,
+            string completionToken,
+            List<P.OrchestrationInstance?>? responseDestinations,
+            IReadOnlyList<string>? requestIds)
         {
-            var coreEntityId = DTCore.Entities.EntityId.FromString(request.InstanceId);
+            var coreEntityId = DTCore.Entities.EntityId.FromString(batchRequest.InstanceId!);
             EntityId entityId = new(coreEntityId.Name, coreEntityId.Key);
 
             TaskName name = new(entityId.Name);
 
-            EntityBatchRequest batchRequest = request.ToEntityBatchRequest();
             EntityBatchResult? batchResult;
 
             try
@@ -441,9 +462,26 @@ async Task OnRunEntityBatchAsync(P.EntityBatchRequest request, string completion
                 };
             }
 
-            // convert the result to protobuf format and send it back
             P.EntityBatchResult response = batchResult.ToEntityBatchResult();
+
             response.CompletionToken = completionToken;
+
+            if (responseDestinations != null)
+            {
+                for (int i = 0; i < response.Results.Count; i++)
+                {
+                    response.ResponseDestinations.Add(responseDestinations[i]);
+                }
+            }
+
+            if (requestIds != null)
+            {
+                foreach (string requestId in requestIds)
+                {
+                    response.RequestIds.Add(requestId);
+                }
+            }
+
             await this.sidecar.CompleteEntityTaskAsync(response);
         }
     }

src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs

@@ -24,14 +24,14 @@ public sealed class GrpcDurableTaskWorkerOptions : DurableTaskWorkerOptions
     public CallInvoker? CallInvoker { get; set; }
 
     /// <summary>
-    /// Gets or sets a flag that indicates whether entity-related events appearing in orchestration histories should be
+    /// Gets or sets a value indicating whether entity-related events appearing in orchestration histories should be
     /// automatically converted back and forth between the old DT Core representation (JSON-encoded external events)
     /// and the new protobuf representation (explicit history events), which is used by the DTS scheduler backend.
     /// </summary>
     public bool ConvertOrchestrationEntityEvents { get; set; }
 
     /// <summary>
-    /// Gets or sets a flag that indicates whether to automatically add entity
+    /// Gets or sets a value indicating whether to automatically add entity
     /// unlock events into the history when an orchestration terminates while holding an entity lock.
     /// </summary>
     public bool InsertEntityUnlocksOnCompletion { get; set; }