history event streaming test added

Author: YunchuWang

test/Grpc.IntegrationTests/GrpcSidecar/Grpc/TaskHubGrpcServer.cs

@@ -31,6 +31,9 @@ public class TaskHubGrpcServer : P.TaskHubSidecarService.TaskHubSidecarServiceBa
     readonly TaskHubDispatcherHost dispatcherHost;
     readonly IsConnectedSignal isConnectedSignal = new();
     readonly SemaphoreSlim sendWorkItemLock = new(initialCount: 1);
+    readonly ConcurrentDictionary<string, List<P.HistoryEvent>> streamingPastEvents = new(StringComparer.OrdinalIgnoreCase);
+
+    volatile bool supportsHistoryStreaming;
 
     // Initialized when a client connects to this service to receive work-item commands.
     IServerStreamWriter<P.WorkItem>? workerToClientStream;
@@ -479,6 +482,8 @@ static P.GetInstanceResponse CreateGetInstanceResponse(OrchestrationState state,
 
     public override async Task GetWorkItems(P.GetWorkItemsRequest request, IServerStreamWriter<P.WorkItem> responseStream, ServerCallContext context)
     {
+        // Record whether the client supports history streaming
+        this.supportsHistoryStreaming = request.Capabilities.Contains(P.WorkerCapability.HistoryStreaming);
         // Use a lock to mitigate the race condition where we signal the dispatch host to start but haven't
         // yet saved a reference to the client response stream.
         lock (this.isConnectedSignal)
@@ -521,6 +526,35 @@ public override async Task GetWorkItems(P.GetWorkItemsRequest request, IServerSt
         }
     }
 
+    public override async Task StreamInstanceHistory(P.StreamInstanceHistoryRequest request, IServerStreamWriter<P.HistoryChunk> responseStream, ServerCallContext context)
+    {
+        if (this.streamingPastEvents.TryGetValue(request.InstanceId, out List<P.HistoryEvent>? pastEvents))
+        {
+            const int MaxChunkBytes = 256 * 1024; // 256KB per chunk to simulate chunked streaming
+            int currentSize = 0;
+            P.HistoryChunk chunk = new();
+
+            foreach (P.HistoryEvent e in pastEvents)
+            {
+                int eventSize = e.CalculateSize();
+                if (currentSize > 0 && currentSize + eventSize > MaxChunkBytes)
+                {
+                    await responseStream.WriteAsync(chunk);
+                    chunk = new P.HistoryChunk();
+                    currentSize = 0;
+                }
+
+                chunk.Events.Add(e);
+                currentSize += eventSize;
+            }
+
+            if (chunk.Events.Count > 0)
+            {
+                await responseStream.WriteAsync(chunk);
+            }
+        }
+    }
+
     /// <summary>
     /// Invoked by the <see cref="TaskHubDispatcherHost"/> when a work item is available, proxies the call to execute an orchestrator over a gRPC channel.
     /// </summary>
@@ -547,16 +581,37 @@ async Task<GrpcOrchestratorExecutionResult> ITaskExecutor.ExecuteOrchestrator(
 
         try
         {
+            var orkRequest = new P.OrchestratorRequest
+            {
+                InstanceId = instance.InstanceId,
+                ExecutionId = instance.ExecutionId,
+                NewEvents = { newEvents.Select(ProtobufUtils.ToHistoryEventProto) },
+                OrchestrationTraceContext = orchestrationTraceContext,
+            };
+
+            // Decide whether to stream based on total size of past events (> 1MiB)
+            List<P.HistoryEvent> protoPastEvents = pastEvents.Select(ProtobufUtils.ToHistoryEventProto).ToList();
+            int totalBytes = 0;
+            foreach (P.HistoryEvent ev in protoPastEvents)
+            {
+                totalBytes += ev.CalculateSize();
+            }
+
+            if (this.supportsHistoryStreaming && totalBytes > (1024))
+            {
+                orkRequest.RequiresHistoryStreaming = true;
+                // Store past events to serve via StreamInstanceHistory
+                this.streamingPastEvents[instance.InstanceId] = protoPastEvents;
+            }
+            else
+            {
+                // Embed full history in the work item
+                orkRequest.PastEvents.AddRange(protoPastEvents);
+            }
+
             await this.SendWorkItemToClientAsync(new P.WorkItem
             {
-                OrchestratorRequest = new P.OrchestratorRequest
-                {
-                    InstanceId = instance.InstanceId,
-                    ExecutionId = instance.ExecutionId,
-                    NewEvents = { newEvents.Select(ProtobufUtils.ToHistoryEventProto) },
-                    OrchestrationTraceContext = orchestrationTraceContext,
-                    PastEvents = { pastEvents.Select(ProtobufUtils.ToHistoryEventProto) },
-                }
+                OrchestratorRequest = orkRequest,
             });
         }
         catch

test/Grpc.IntegrationTests/LargePayloadTests.cs

@@ -88,6 +88,64 @@ public async Task LargeOrchestrationInputAndOutputAndCustomStatus()
         Assert.Contains(JsonSerializer.Serialize(largeInput + largeInput), fakeStore.uploadedPayloads);
     }
 
+    // Validates history streaming path resolves externalized inputs/outputs in HistoryChunk.
+    [Fact]
+    public async Task HistoryStreaming_ResolvesPayloads()
+    {
+        // Make payloads large enough so that past events history exceeds 1 MiB to trigger streaming
+        string largeInput = new string('H', 2 * 1024 * 1024);   // 2 MiB
+        string largeOutput = new string('O', 2 * 1024 * 1024);  // 2 MiB
+        TaskName orch = nameof(HistoryStreaming_ResolvesPayloads);
+
+        InMemoryPayloadStore store = new InMemoryPayloadStore();
+
+        await using HostTestLifetime server = await this.StartWorkerAsync(
+            worker =>
+            {
+                worker.AddTasks(tasks => tasks.AddOrchestratorFunc<string, string>(
+                    orch,
+                    async (ctx, input) =>
+                    {
+                        // Emit several events so that the serialized history size grows
+                        for (int i = 0; i < 50; i++)
+                        {
+                            await ctx.CreateTimer(TimeSpan.FromMilliseconds(10), CancellationToken.None);
+                        }
+                        return largeOutput;
+                    }));
+
+                worker.UseExternalizedPayloads(opts =>
+                {
+                    opts.ExternalizeThresholdBytes = 1024;
+                    opts.ContainerName = "test";
+                    opts.ConnectionString = "UseDevelopmentStorage=true";
+                });
+                worker.Services.AddSingleton<IPayloadStore>(store);
+            },
+            client =>
+            {
+                // Enable client to resolve outputs on query
+                client.UseExternalizedPayloads(opts =>
+                {
+                    opts.ExternalizeThresholdBytes = 1024;
+                    opts.ContainerName = "test";
+                    opts.ConnectionString = "UseDevelopmentStorage=true";
+                });
+                client.Services.AddSingleton<IPayloadStore>(store);
+            });
+
+        // Start orchestration with large input to exercise history input resolution
+        string instanceId = await server.Client.ScheduleNewOrchestrationInstanceAsync(orch, largeInput);
+        OrchestrationMetadata completed = await server.Client.WaitForInstanceCompletionAsync(
+            instanceId, getInputsAndOutputs: true, this.TimeoutToken);
+
+        Assert.Equal(OrchestrationRuntimeStatus.Completed, completed.RuntimeStatus);
+        Assert.Equal(largeInput, completed.ReadInputAs<string>());
+        Assert.Equal(largeOutput, completed.ReadOutputAs<string>());
+        Assert.True(store.UploadCount >= 2);
+        Assert.True(store.DownloadCount >= 2);
+    }
+
     // Validates client externalizes large suspend and resume reasons.
     [Fact]
     public async Task SuspendAndResume_Reason_IsExternalizedByClient()