updatestest

Author: nytian

eng/targets/Release.props

@@ -17,7 +17,7 @@
   </PropertyGroup>
 
   <PropertyGroup>
-    <VersionPrefix>1.12.0</VersionPrefix>
+    <VersionPrefix>1.13.0</VersionPrefix>
     <VersionSuffix></VersionSuffix>
   </PropertyGroup>
 

test/Grpc.IntegrationTests/GrpcDurableTaskClientIntegrationTests.cs

@@ -229,6 +229,7 @@ public async Task PurgeInstances_WithFilter_EndToEnd()
     [InlineData(true)]
     public async Task RestartAsync_EndToEnd(bool restartWithNewInstanceId)
     {
+        using var cts = new CancellationTokenSource(TimeSpan.FromMinutes(1)); // Reduced timeout to 1 minute
         await using HostTestLifetime server = await this.StartAsync();
 
         // Start an initial orchestration with shouldThrow = false to ensure it completes successfully
@@ -238,8 +239,8 @@ public async Task RestartAsync_EndToEnd(bool restartWithNewInstanceId)
         // Wait for it to start and then complete
         await server.Client.WaitForInstanceStartAsync(originalInstanceId, default);
         await server.Client.RaiseEventAsync(originalInstanceId, "event", default);
-        await server.Client.WaitForInstanceCompletionAsync(originalInstanceId, default);
-
+        await server.Client.WaitForInstanceCompletionAsync(originalInstanceId, cts.Token);
+        
         // Verify the original orchestration completed
         OrchestrationMetadata? originalMetadata = await server.Client.GetInstanceAsync(originalInstanceId, true);
         originalMetadata.Should().NotBeNull();
@@ -258,32 +259,29 @@ public async Task RestartAsync_EndToEnd(bool restartWithNewInstanceId)
             restartedInstanceId.Should().Be(originalInstanceId);
         }
 
-        // Wait for the restarted orchestration to start
-        await server.Client.WaitForInstanceStartAsync(restartedInstanceId, default);
-
-        // Verify the restarted orchestration has the same input
-        OrchestrationMetadata? restartedMetadata = await server.Client.GetInstanceAsync(restartedInstanceId, true);
-        restartedMetadata.Should().NotBeNull();
-        restartedMetadata!.Name.Should().Be(OrchestrationName);
-        restartedMetadata.SerializedInput.Should().Be("false");
-
         // Complete the restarted orchestration
-        await server.Client.RaiseEventAsync(restartedInstanceId, "event", default);
-        await server.Client.WaitForInstanceCompletionAsync(restartedInstanceId, default);
+        await server.Client.RaiseEventAsync(restartedInstanceId, "event");
+        
+        // Wait for completion (with shorter timeout)
+        using var completionCts = new CancellationTokenSource(TimeSpan.FromSeconds(30));
+        await server.Client.WaitForInstanceCompletionAsync(restartedInstanceId, completionCts.Token);
 
         // Verify the restarted orchestration completed
-        restartedMetadata = await server.Client.GetInstanceAsync(restartedInstanceId, true);
+        var restartedMetadata = await server.Client.GetInstanceAsync(restartedInstanceId, true);
         restartedMetadata.Should().NotBeNull();
+        restartedMetadata!.Name.Should().Be(OrchestrationName);
+        restartedMetadata.SerializedInput.Should().Be("false");
         restartedMetadata!.RuntimeStatus.Should().Be(OrchestrationRuntimeStatus.Completed);
     }
 
     [Fact]
     public async Task RestartAsync_InstanceNotFound_ThrowsArgumentException()
     {
+        using var cts = new CancellationTokenSource(TimeSpan.FromMinutes(1)); // 1-minute timeout
         await using HostTestLifetime server = await this.StartAsync();
 
         // Try to restart a non-existent orchestration
-        Func<Task> restartAction = () => server.Client.RestartAsync("non-existent-instance-id");
+        Func<Task> restartAction = () => server.Client.RestartAsync("non-existent-instance-id", cancellation: cts.Token);
 
         // Should throw ArgumentException
         await restartAction.Should().ThrowAsync<ArgumentException>()

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

@@ -339,12 +339,20 @@ static P.GetInstanceResponse CreateGetInstanceResponse(OrchestrationState state,
         {
             // Get the original orchestration state
             IList<OrchestrationState> states = await this.client.GetOrchestrationStateAsync(request.InstanceId, false);
+            
             if (states == null || states.Count == 0)
             {
                 throw new RpcException(new Status(StatusCode.NotFound, $"An orchestration with the instanceId {request.InstanceId} was not found."));
             }
 
             OrchestrationState state = states[0];
+            
+            // Check if the state is null
+            if (state == null)
+            {
+                throw new RpcException(new Status(StatusCode.NotFound, $"An orchestration with the instanceId {request.InstanceId} was not found."));
+            }
+            
             string newInstanceId = request.RestartWithNewInstanceId ? Guid.NewGuid().ToString("N") : request.InstanceId;
 
             // Create a new orchestration instance

test/Grpc.IntegrationTests/GrpcSidecar/InMemoryOrchestrationService.cs

@@ -444,8 +444,21 @@ public void AddMessage(TaskMessage message)
             SerializedInstanceState state = this.store.GetOrAdd(instanceId, id => new SerializedInstanceState(id, executionId));
             lock (state)
             {
+                bool isRestart = state.ExecutionId != null && state.ExecutionId != executionId;
+                
                 if (message.Event is ExecutionStartedEvent startEvent)
                 {
+                    // For restart scenarios, clear the history and reset the state
+                    if (isRestart && state.IsCompleted)
+                    {
+                        state.HistoryEventsJson.Clear();
+                        state.ExecutionId = executionId; // Always update the execution ID
+                        state.IsCompleted = false;
+                    }
+                    
+                    // Add the ExecutionStartedEvent to history so the orchestration worker has proper context
+                    state.HistoryEventsJson.Add(startEvent);
+                    
                     OrchestrationState newStatusRecord = new()
                     {
                         OrchestrationInstance = startEvent.OrchestrationInstance,
@@ -460,7 +473,6 @@ public void AddMessage(TaskMessage message)
                     };
 
                     state.StatusRecordJson = JsonValue.Create(newStatusRecord);
-                    state.HistoryEventsJson.Clear();
                     state.IsCompleted = false;
                 }
                 else if (state.IsCompleted)
@@ -661,13 +673,11 @@ public async ValueTask<SerializedInstanceState> TakeNextAsync(CancellationToken
 
             public void Schedule(SerializedInstanceState state)
             {
-                // TODO: There is a race condition here. If another thread is calling TakeNextAsync
-                //       and removed the queue item before updating the dictionary, then we'll fail
-                //       to update the readyToRunQueue and the orchestration will get stuck.
-                if (this.readyInstances.TryAdd(state.InstanceId, state))
-                {
-                    this.readyToRunQueue.Writer.TryWrite(state);
-                }
+                // Always update the dictionary entry (for restart scenarios)
+                this.readyInstances[state.InstanceId] = state;
+                
+                // Always try to write to the queue
+                this.readyToRunQueue.Writer.TryWrite(state);
             }
         }