diff --git a/src/DurableTask.ServiceBus/ServiceBusOrchestrationService.cs b/src/DurableTask.ServiceBus/ServiceBusOrchestrationService.cs index ae451e03..f759ff59 100644 --- a/src/DurableTask.ServiceBus/ServiceBusOrchestrationService.cs +++ b/src/DurableTask.ServiceBus/ServiceBusOrchestrationService.cs @@ -62,7 +62,7 @@ public class ServiceBusOrchestrationService : IOrchestrationService, IOrchestrat // as every fetched message also creates a tracking message which counts towards this limit. const int MaxMessageCount = 80; const int SessionStreamWarningSizeInBytes = 150 * 1024; - const int StatusPollingIntervalInSeconds = 2; + static readonly TimeSpan StatusPollingInterval = TimeSpan.FromSeconds(2); const int DuplicateDetectionWindowInHours = 4; /// @@ -1222,9 +1222,12 @@ public async Task ForceTerminateTaskOrchestrationAsync(string instanceId, string /// /// Wait for an orchestration to reach any terminal state within the given timeout /// - /// The execution id of the orchestration + /// The execution id of the orchestration to wait for. When null, empty, or + /// whitespace, the current generation of the instance is followed. Otherwise only that execution is + /// tracked, except for a execution, + /// where the current generation of the instance is followed instead. /// Instance to wait for - /// Max timeout to wait + /// Max timeout to wait. Only positive values, , or are allowed. /// Task cancellation token public async Task WaitForOrchestrationAsync( string instanceId, @@ -1236,20 +1239,68 @@ public async Task WaitForOrchestrationAsync( if (string.IsNullOrWhiteSpace(instanceId)) { - throw new ArgumentException("instanceId"); + throw new ArgumentException("The instance id cannot be null, empty, or whitespace.", nameof(instanceId)); } - double timeoutSeconds = timeout.TotalSeconds; + bool isInfiniteTimeSpan = timeout == Timeout.InfiniteTimeSpan; + if (timeout < TimeSpan.Zero && !isInfiniteTimeSpan) + { + throw new ArgumentException($"The parameter {nameof(timeout)} cannot be negative unless it is Timeout.InfiniteTimeSpan." + + $" The value for {nameof(timeout)} was '{timeout}'." + + $" Please provide a positive timeout value, TimeSpan.Zero or Timeout.InfiniteTimeSpan.", + nameof(timeout)); + } + + bool pinnedToExecution = !string.IsNullOrWhiteSpace(executionId); - while (!cancellationToken.IsCancellationRequested && timeoutSeconds > 0) + while (!cancellationToken.IsCancellationRequested) { - OrchestrationState state = (await GetOrchestrationStateAsync(instanceId, false))?.FirstOrDefault(); + OrchestrationState state = pinnedToExecution + ? await GetOrchestrationStateAsync(instanceId, executionId) + : (await GetOrchestrationStateAsync(instanceId, false))?.FirstOrDefault(); + + // A pinned execution that continued-as-new is only a tombstone: the live orchestration + // moved on to a new execution id, so stop pinning and follow the current generation. + // The tombstone still counts as this iteration's status check, so fall through to the + // timeout accounting below rather than re-querying immediately. + if (pinnedToExecution && state?.OrchestrationStatus == OrchestrationStatus.ContinuedAsNew) + { + pinnedToExecution = false; + } + + // ContinuedAsNew is never a final state: a new generation always follows it. The built-in + // AzureTableInstanceStore hides these rows from the non-pinned lookup, but that is not + // guaranteed by IOrchestrationServiceInstanceStore, so keep polling if one surfaces. if (state == null || (state.OrchestrationStatus == OrchestrationStatus.Running) - || (state.OrchestrationStatus == OrchestrationStatus.Pending)) + || (state.OrchestrationStatus == OrchestrationStatus.Pending) + || (state.OrchestrationStatus == OrchestrationStatus.Suspended) + || (state.OrchestrationStatus == OrchestrationStatus.ContinuedAsNew)) { - await Task.Delay(StatusPollingIntervalInSeconds * 1000, cancellationToken); - timeoutSeconds -= StatusPollingIntervalInSeconds; + TimeSpan delay = StatusPollingInterval; + + if (!isInfiniteTimeSpan) + { + // The timeout condition is checked after the status check so that a user-provided + // timeout of `TimeSpan.Zero` still checks the status of the orchestration once + // before returning. + if (timeout <= TimeSpan.Zero) + { + break; + } + + // Only the time actually spent waiting is charged against the budget, and the + // last delay is clamped to what remains, so the full timeout window is polled + // before giving up. + if (timeout < delay) + { + delay = timeout; + } + + timeout -= delay; + } + + await Task.Delay(delay, cancellationToken); } else { diff --git a/src/DurableTask.ServiceBus/Tracking/AzureTableClient.cs b/src/DurableTask.ServiceBus/Tracking/AzureTableClient.cs index c0e28a8d..e4baf99b 100644 --- a/src/DurableTask.ServiceBus/Tracking/AzureTableClient.cs +++ b/src/DurableTask.ServiceBus/Tracking/AzureTableClient.cs @@ -124,7 +124,7 @@ internal async Task DeleteJumpStartTableIfExistsAsync() } } - public Task> QueryOrchestrationStatesAsync( + public virtual Task> QueryOrchestrationStatesAsync( OrchestrationStateQuery stateQuery) { var query = CreateQueryInternal(stateQuery, false); @@ -138,7 +138,7 @@ public async Task> QueryOrchestrationSt return await QueryTableSegmentAsync(this.historyTableClient, query, continuationToken, count); } - public Task> QueryJumpStartOrchestrationsAsync(OrchestrationStateQuery stateQuery) + public virtual Task> QueryJumpStartOrchestrationsAsync(OrchestrationStateQuery stateQuery) { // TODO: Enable segmented query for paging purpose var query = CreateQueryInternal(stateQuery, true); diff --git a/src/DurableTask.ServiceBus/Tracking/AzureTableInstanceStore.cs b/src/DurableTask.ServiceBus/Tracking/AzureTableInstanceStore.cs index 818ae3cf..d8c98a09 100644 --- a/src/DurableTask.ServiceBus/Tracking/AzureTableInstanceStore.cs +++ b/src/DurableTask.ServiceBus/Tracking/AzureTableInstanceStore.cs @@ -90,6 +90,20 @@ public AzureTableInstanceStore(string hubName, Uri endpoint, TokenCredential cre DateTimeUtils.SetMinDateTimeForStorageEmulator(); } + /// + /// Creates a new AzureTableInstanceStore over the supplied table client. Used by tests to + /// substitute the table client; production code uses the connection string or credential + /// constructors above. + /// + /// The table client to issue queries through + internal AzureTableInstanceStore(AzureTableClient tableClient) + { + this.tableClient = tableClient ?? throw new ArgumentNullException(nameof(tableClient)); + + // Workaround an issue with Storage that throws exceptions for any date < 1600 so DateTime.Min cannot be used + DateTimeUtils.SetMinDateTimeForStorageEmulator(); + } + /// /// Runs initialization to prepare the storage for use /// @@ -217,17 +231,28 @@ public async Task> GetOrchestratio /// The matching orchestration state or null if not found public async Task GetOrchestrationStateAsync(string instanceId, string executionId) { + // Retries here mirror the instance-id overload above so that both lookups survive the same + // transient table failures. Without them a single failure that escapes the storage SDK's own + // retries would fault every caller that polls by execution id. AzureTableOrchestrationStateEntity result = - (await this.tableClient.QueryOrchestrationStatesAsync( - new OrchestrationStateQuery().AddInstanceFilter(instanceId, executionId)).ConfigureAwait(false)).FirstOrDefault(); + (await Utils.ExecuteWithRetries(() => this.tableClient.QueryOrchestrationStatesAsync( + new OrchestrationStateQuery().AddInstanceFilter(instanceId, executionId)), + string.Empty, + "GetOrchestrationStateAsync-stateEntity", + MaxRetriesTableStore, + IntervalBetweenRetriesSecs).ConfigureAwait(false)).FirstOrDefault(); // ReSharper disable once ConvertIfStatementToNullCoalescingExpression if (result == null) { // Query from JumpStart table - result = (await this.tableClient.QueryJumpStartOrchestrationsAsync( + result = (await Utils.ExecuteWithRetries(() => this.tableClient.QueryJumpStartOrchestrationsAsync( new OrchestrationStateQuery() - .AddInstanceFilter(instanceId, executionId)) + .AddInstanceFilter(instanceId, executionId)), + string.Empty, + "GetOrchestrationStateAsync-jumpStartEntity", + MaxRetriesTableStore, + IntervalBetweenRetriesSecs) .ConfigureAwait(false)) .FirstOrDefault(); } diff --git a/test/DurableTask.ServiceBus.Tests/AzureTableInstanceStoreRetryTests.cs b/test/DurableTask.ServiceBus.Tests/AzureTableInstanceStoreRetryTests.cs new file mode 100644 index 00000000..78d6ef7b --- /dev/null +++ b/test/DurableTask.ServiceBus.Tests/AzureTableInstanceStoreRetryTests.cs @@ -0,0 +1,151 @@ +// ---------------------------------------------------------------------------------- +// Copyright Microsoft Corporation +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// http://www.apache.org/licenses/LICENSE-2.0 +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ---------------------------------------------------------------------------------- + +namespace DurableTask.ServiceBus.Tests +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Threading.Tasks; + using global::Azure; + using DurableTask.Core; + using DurableTask.Core.Tracking; + using DurableTask.ServiceBus.Tracking; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + /// + /// Verifies that the execution-id lookup on retries transient + /// table failures, matching the instance-id overload. Polling by execution id would otherwise fault + /// on the first transient error that escapes the storage SDK's own retries. + /// + [TestClass] + public class AzureTableInstanceStoreRetryTests + { + const string InstanceId = "instance-1"; + const string ExecutionId = "generation-1"; + const string HubName = "testhub"; + + // Parsed only; constructing a TableServiceClient performs no I/O. + const string FakeStorageConnectionString = "UseDevelopmentStorage=true"; + + static readonly DateTime BaseTime = new DateTime(2024, 1, 1, 0, 0, 0, DateTimeKind.Utc); + + static AzureTableOrchestrationStateEntity CreateEntity() + { + return new AzureTableOrchestrationStateEntity(new OrchestrationState + { + OrchestrationInstance = new OrchestrationInstance + { + InstanceId = InstanceId, + ExecutionId = ExecutionId + }, + OrchestrationStatus = OrchestrationStatus.Completed, + CreatedTime = BaseTime, + LastUpdatedTime = BaseTime, + Output = "final output" + }); + } + + /// + /// A transient failure on the state table must be retried rather than surfaced to the caller. + /// + [TestMethod] + public async Task GetOrchestrationStateAsync_ByExecutionId_RetriesTransientStateTableFailure() + { + var tableClient = new FailOnceTableClient { FailuresBeforeSuccess = 1 }; + var store = new AzureTableInstanceStore(tableClient); + + OrchestrationStateInstanceEntity state = await store.GetOrchestrationStateAsync(InstanceId, ExecutionId); + + Assert.IsNotNull(state, "A single transient failure must not fault the execution-id lookup."); + Assert.AreEqual(ExecutionId, state.State.OrchestrationInstance.ExecutionId); + Assert.AreEqual("final output", state.State.Output); + Assert.AreEqual(2, tableClient.StateQueryCount, "Expected one failed attempt followed by a successful retry."); + } + + /// + /// The JumpStart fallback is a separate query and needs the same protection: without it a + /// transient failure there faults the lookup even though the state table responded fine. + /// + [TestMethod] + public async Task GetOrchestrationStateAsync_ByExecutionId_RetriesTransientJumpStartTableFailure() + { + var tableClient = new FailOnceTableClient + { + // Force the JumpStart fallback by leaving the state table empty. + StateTableIsEmpty = true, + JumpStartFailuresBeforeSuccess = 1 + }; + + var store = new AzureTableInstanceStore(tableClient); + + OrchestrationStateInstanceEntity state = await store.GetOrchestrationStateAsync(InstanceId, ExecutionId); + + Assert.IsNotNull(state, "A single transient failure must not fault the JumpStart fallback."); + Assert.AreEqual(ExecutionId, state.State.OrchestrationInstance.ExecutionId); + Assert.AreEqual(2, tableClient.JumpStartQueryCount, "Expected one failed attempt followed by a successful retry."); + } + + /// + /// Table client that throws a retryable storage error a fixed number of times before returning + /// a row, so a lookup that lacks retries fails and one that retries succeeds. + /// + sealed class FailOnceTableClient : AzureTableClient + { + public FailOnceTableClient() + : base(HubName, FakeStorageConnectionString) + { + } + + public int FailuresBeforeSuccess { get; set; } + + public int JumpStartFailuresBeforeSuccess { get; set; } + + public bool StateTableIsEmpty { get; set; } + + public int StateQueryCount { get; private set; } + + public int JumpStartQueryCount { get; private set; } + + public override Task> QueryOrchestrationStatesAsync( + OrchestrationStateQuery stateQuery) + { + this.StateQueryCount++; + + if (this.StateQueryCount <= this.FailuresBeforeSuccess) + { + throw new RequestFailedException(503, "The server is busy."); + } + + IEnumerable result = this.StateTableIsEmpty + ? Enumerable.Empty() + : new[] { CreateEntity() }; + + return Task.FromResult(result); + } + + public override Task> QueryJumpStartOrchestrationsAsync( + OrchestrationStateQuery stateQuery) + { + this.JumpStartQueryCount++; + + if (this.JumpStartQueryCount <= this.JumpStartFailuresBeforeSuccess) + { + throw new RequestFailedException(503, "The server is busy."); + } + + return Task.FromResult>(new[] { CreateEntity() }); + } + } + } +} diff --git a/test/DurableTask.ServiceBus.Tests/WaitForOrchestrationTests.cs b/test/DurableTask.ServiceBus.Tests/WaitForOrchestrationTests.cs new file mode 100644 index 00000000..faee4592 --- /dev/null +++ b/test/DurableTask.ServiceBus.Tests/WaitForOrchestrationTests.cs @@ -0,0 +1,649 @@ +// ---------------------------------------------------------------------------------- +// Copyright Microsoft Corporation +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// http://www.apache.org/licenses/LICENSE-2.0 +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ---------------------------------------------------------------------------------- + +namespace DurableTask.ServiceBus.Tests +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Threading; + using System.Threading.Tasks; + using DurableTask.Core; + using DurableTask.Core.Tracking; + using DurableTask.ServiceBus.Settings; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + /// + /// Unit tests for . + /// These use an in-memory instance store so they do not require a live Service Bus namespace. + /// + [TestClass] + public class WaitForOrchestrationTests + { + const string InstanceId = "instance-1"; + const string FakeConnectionString = + "Endpoint=sb://test.servicebus.windows.net/;SharedAccessKeyName=key;SharedAccessKey=dGVzdGtleQ=="; + + static readonly DateTime BaseTime = new DateTime(2024, 1, 1, 0, 0, 0, DateTimeKind.Utc); + + static ServiceBusOrchestrationService CreateService(FakeInstanceStore instanceStore) + { + return new ServiceBusOrchestrationService( + FakeConnectionString, + "testhub", + instanceStore, + null, + new ServiceBusOrchestrationServiceSettings()); + } + + static OrchestrationState CreateState( + string executionId, + OrchestrationStatus status, + DateTime createdTime, + string output = null) + { + return new OrchestrationState + { + OrchestrationInstance = new OrchestrationInstance + { + InstanceId = InstanceId, + ExecutionId = executionId + }, + OrchestrationStatus = status, + CreatedTime = createdTime, + LastUpdatedTime = createdTime, + Output = output + }; + } + + /// + /// The core of the fix: while waiting on a specific execution, state written by an earlier run + /// of the same instance id must never be returned, even though it is the only readable state. + /// + [TestMethod] + public async Task WaitForOrchestration_PinnedExecution_IgnoresStateFromPreviousRun() + { + var store = new FakeInstanceStore(); + + // A previous run of the same instance id that already completed. + store.States.Add(CreateState("previous-run", OrchestrationStatus.Completed, BaseTime, "stale output")); + + // The new run's state only becomes readable after the first poll. + store.OnQuery = s => + { + if (s.QueryCount == 2) + { + s.States.Add(CreateState("current-run", OrchestrationStatus.Completed, BaseTime.AddMinutes(5), "fresh output")); + } + }; + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + "current-run", + TimeSpan.FromSeconds(30), + CancellationToken.None); + + Assert.IsNotNull(state); + Assert.AreEqual("current-run", state.OrchestrationInstance.ExecutionId, "Returned state from the wrong run."); + Assert.AreEqual("fresh output", state.Output); + } + + /// + /// If the new run never becomes readable, the wait must time out rather than return the + /// previous run's result. + /// + [TestMethod] + public async Task WaitForOrchestration_PinnedExecution_TimesOutRatherThanReturnPreviousRun() + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("previous-run", OrchestrationStatus.Completed, BaseTime, "stale output")); + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + "current-run", + TimeSpan.FromSeconds(2), + CancellationToken.None); + + Assert.IsNull(state, "State from a previous run of the same instance id must not be returned."); + } + + /// + /// A ContinuedAsNew row is a tombstone that is never updated again, so the wait must follow the + /// new generation instead of polling the pinned execution forever. + /// + [TestMethod] + public async Task WaitForOrchestration_PinnedExecution_ContinuedAsNew_FollowsNextGeneration() + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.ContinuedAsNew, BaseTime, "next input")); + store.States.Add(CreateState("generation-2", OrchestrationStatus.Completed, BaseTime.AddMinutes(1), "final output")); + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + "generation-1", + TimeSpan.FromSeconds(30), + CancellationToken.None); + + Assert.IsNotNull(state); + Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus); + Assert.AreEqual("generation-2", state.OrchestrationInstance.ExecutionId); + Assert.AreEqual("final output", state.Output); + } + + /// + /// Un-pinning after a ContinuedAsNew tombstone must not re-query within the same iteration: + /// the tombstone was itself this iteration's status check, so a zero timeout still performs + /// exactly one lookup rather than reaching the next generation for free. + /// + [TestMethod] + public async Task WaitForOrchestration_Timeout_Zero_ContinuedAsNew_ChecksOnceAndDoesNotFollowNextGeneration() + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.ContinuedAsNew, BaseTime, "next input")); + store.States.Add(CreateState("generation-2", OrchestrationStatus.Completed, BaseTime.AddMinutes(1), "final output")); + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + "generation-1", + TimeSpan.Zero, + CancellationToken.None); + + Assert.IsNull(state, "A zero timeout must not reach the next generation after un-pinning."); + Assert.AreEqual(1, store.QueryCount, "A zero timeout must perform exactly one lookup."); + } + + /// + /// Suspended is a pause, not a terminal state: an orchestration that is never resumed has no + /// result, so the wait must time out rather than report it as finished. + /// + [TestMethod] + public async Task WaitForOrchestration_Suspended_IsNotReturnedAsTerminal() + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.Suspended, BaseTime)); + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + "generation-1", + TimeSpan.FromSeconds(2), + CancellationToken.None); + + Assert.IsNull(state, "A suspended orchestration has not completed and must not be returned."); + } + + /// + /// Callers that do not supply an execution id keep the legacy behavior of following the + /// current generation of the instance. + /// + [DataTestMethod] + [DataRow(null)] + [DataRow("")] + [DataRow(" ")] + public async Task WaitForOrchestration_WithoutExecutionId_UsesCurrentGeneration(string executionId) + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.ContinuedAsNew, BaseTime, "next input")); + store.States.Add(CreateState("generation-2", OrchestrationStatus.Completed, BaseTime.AddMinutes(1), "final output")); + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + executionId, + TimeSpan.FromSeconds(30), + CancellationToken.None); + + Assert.IsNotNull(state); + Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus); + Assert.AreEqual("generation-2", state.OrchestrationInstance.ExecutionId); + Assert.AreEqual(0, store.PinnedQueryCount, "An empty execution id must not be queried as an exact execution."); + } + + /// + /// Hiding ContinuedAsNew rows from the current generation lookup is an implementation detail of + /// AzureTableInstanceStore, not a guarantee of IOrchestrationServiceInstanceStore. A store that + /// surfaces them must not cause a tombstone to be reported as the final state. + /// + [TestMethod] + public async Task WaitForOrchestration_StoreWithoutContinuedAsNewFilter_DoesNotReturnTombstone() + { + var store = new FakeInstanceStore { FilterContinuedAsNew = false }; + store.States.Add(CreateState("generation-1", OrchestrationStatus.ContinuedAsNew, BaseTime, "next input")); + + store.OnQuery = s => + { + if (s.QueryCount == 2) + { + s.States.Add(CreateState("generation-2", OrchestrationStatus.Completed, BaseTime.AddMinutes(1), "final output")); + } + }; + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + null, + TimeSpan.FromSeconds(30), + CancellationToken.None); + + Assert.IsNotNull(state); + Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus); + Assert.AreEqual("generation-2", state.OrchestrationInstance.ExecutionId); + } + + [TestMethod] + public async Task WaitForOrchestration_PinnedExecution_ReturnsFailedState() + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.Failed, BaseTime, "boom")); + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + "generation-1", + TimeSpan.FromSeconds(30), + CancellationToken.None); + + Assert.IsNotNull(state); + Assert.AreEqual(OrchestrationStatus.Failed, state.OrchestrationStatus); + Assert.AreEqual(1, store.PinnedQueryCount, "The pinned lookup should have been used."); + Assert.AreEqual(0, store.LatestQueryCount, "The current generation lookup should not have been used."); + } + + /// + /// A suspended orchestration that is resumed and then runs to completion must return the + /// final state, not stop at the intermediate Suspended or Running rows. + /// + [TestMethod] + public async Task WaitForOrchestration_Suspended_ResumedAndCompleted_ReturnsFinalState() + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.Suspended, BaseTime)); + + // Suspended -> Running (resumed) -> Completed, one transition per poll. + store.OnQuery = s => + { + OrchestrationStatus? next = s.QueryCount == 2 ? OrchestrationStatus.Running + : s.QueryCount == 3 ? OrchestrationStatus.Completed + : (OrchestrationStatus?)null; + + if (next != null) + { + s.States.Clear(); + s.States.Add(CreateState( + "generation-1", + next.Value, + BaseTime, + next == OrchestrationStatus.Completed ? "final output" : null)); + } + }; + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + "generation-1", + TimeSpan.FromSeconds(30), + CancellationToken.None); + + Assert.IsNotNull(state, "The resumed orchestration completed and must be returned."); + Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus); + Assert.AreEqual("final output", state.Output); + Assert.AreEqual(3, store.QueryCount, "The wait should have polled through Suspended and Running."); + } + + /// + /// Negative timeouts are a caller bug. Timeout.InfiniteTimeSpan is itself negative (-1ms), so it + /// must be excluded from this check; values adjacent to it must still be rejected. + /// + [DataTestMethod] + [DataRow(-2)] + [DataRow(-1000)] + [DataRow(-60000)] + public async Task WaitForOrchestration_Timeout_Negative_Throws(int timeoutMilliseconds) + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.Completed, BaseTime, "final output")); + + ServiceBusOrchestrationService service = CreateService(store); + + await Assert.ThrowsExceptionAsync( + () => service.WaitForOrchestrationAsync( + InstanceId, + "generation-1", + TimeSpan.FromMilliseconds(timeoutMilliseconds), + CancellationToken.None), + $"A timeout of {timeoutMilliseconds}ms should be rejected."); + + Assert.AreEqual(0, store.QueryCount, "The timeout should be validated before any lookup."); + } + + /// + /// A zero timeout means "check once and return", so a state that is already terminal is returned. + /// + [TestMethod] + public async Task WaitForOrchestration_Timeout_Zero_ChecksOnceAndReturnsTerminalState() + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.Completed, BaseTime, "final output")); + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + "generation-1", + TimeSpan.Zero, + CancellationToken.None); + + Assert.IsNotNull(state, "A zero timeout must still perform one status check."); + Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus); + Assert.AreEqual(1, store.QueryCount); + } + + /// + /// A zero timeout must not wait: if the orchestration is not yet terminal it returns null after + /// a single check rather than polling. + /// + [TestMethod] + public async Task WaitForOrchestration_Timeout_Zero_DoesNotPollWhenNotComplete() + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.Running, BaseTime)); + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + "generation-1", + TimeSpan.Zero, + CancellationToken.None); + + Assert.IsNull(state); + Assert.AreEqual(1, store.QueryCount, "A zero timeout must check exactly once and not poll."); + } + + /// + /// A positive timeout polls until it elapses, then gives up and returns null. + /// + [TestMethod] + public async Task WaitForOrchestration_Timeout_Positive_PollsThenReturnsNull() + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.Running, BaseTime)); + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + "generation-1", + TimeSpan.FromSeconds(4), + CancellationToken.None); + + Assert.IsNull(state, "The orchestration never completed, so the wait must time out."); + Assert.IsTrue(store.QueryCount > 1, $"A 4 second timeout should poll more than once, polled {store.QueryCount} time(s)."); + } + + /// + /// The whole timeout window must be polled. Charging a full polling interval against the budget + /// before the delay is awaited would abandon the wait about halfway through and miss an + /// orchestration that completes late in the window. + /// + [TestMethod] + public async Task WaitForOrchestration_Timeout_Positive_PollsForTheFullTimeout() + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.Running, BaseTime)); + + // With a 2 second polling interval a 4 second timeout allows checks at roughly t=0, t=2 + // and t=4, so state that only becomes terminal on the third check is still observed. + store.OnQuery = s => + { + if (s.QueryCount == 3) + { + s.States.Clear(); + s.States.Add(CreateState("generation-1", OrchestrationStatus.Completed, BaseTime, "final output")); + } + }; + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + "generation-1", + TimeSpan.FromSeconds(4), + CancellationToken.None); + + Assert.IsNotNull(state, "The orchestration completed within the timeout and must be returned."); + Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus); + Assert.AreEqual("final output", state.Output); + } + + /// + /// A timeout shorter than the polling interval must wait out its window rather than returning on + /// the first check, but must not overshoot it by delaying for a whole polling interval. + /// + [TestMethod] + public async Task WaitForOrchestration_Timeout_ShorterThanPollingInterval_WaitsWithoutOvershooting() + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.Running, BaseTime)); + + ServiceBusOrchestrationService service = CreateService(store); + + // The timeout is kept far below the 2 second polling interval so the three possible + // behaviours are widely separated in time and the assertions below do not depend on + // precise scheduling: clamping the delay takes about 200ms, delaying for a whole + // interval takes about 2s, and charging the budget before the delay returns instantly. + TimeSpan timeout = TimeSpan.FromMilliseconds(200); + + var stopwatch = System.Diagnostics.Stopwatch.StartNew(); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + "generation-1", + timeout, + CancellationToken.None); + + stopwatch.Stop(); + + Assert.IsNull(state, "The orchestration never completed, so the wait must time out."); + Assert.AreEqual(2, store.QueryCount, "The wait must check once, wait out its window, then check again."); + + // Task.Delay never returns early, so this only fails if the wait did not delay at all. + Assert.IsTrue( + stopwatch.Elapsed >= TimeSpan.FromMilliseconds(150), + $"The wait must use its window instead of returning immediately, took {stopwatch.ElapsedMilliseconds}ms."); + + // Well clear of the ~200ms a clamped delay needs, and well below the 2s a full polling + // interval would take. + Assert.IsTrue( + stopwatch.Elapsed < TimeSpan.FromMilliseconds(1500), + $"The wait must clamp its delay to the remaining timeout instead of waiting a whole polling interval, took {stopwatch.ElapsedMilliseconds}ms."); + } + + /// + /// Timeout.InfiniteTimeSpan is negative, so a naive remaining-time check would treat it as already + /// elapsed and return null on the first poll instead of waiting. + /// + [TestMethod] + public async Task WaitForOrchestration_Timeout_Infinite_WaitsForCompletion() + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.Running, BaseTime)); + + store.OnQuery = s => + { + if (s.QueryCount == 2) + { + s.States.Clear(); + s.States.Add(CreateState("generation-1", OrchestrationStatus.Completed, BaseTime, "final output")); + } + }; + + ServiceBusOrchestrationService service = CreateService(store); + + OrchestrationState state = await service.WaitForOrchestrationAsync( + InstanceId, + "generation-1", + Timeout.InfiniteTimeSpan, + CancellationToken.None); + + Assert.IsNotNull(state, "An infinite timeout must keep waiting instead of giving up immediately."); + Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus); + Assert.AreEqual("final output", state.Output); + } + + /// + /// An infinite wait must still observe cancellation, otherwise it can never be stopped. + /// + [TestMethod] + public async Task WaitForOrchestration_Timeout_Infinite_HonorsCancellation() + { + var store = new FakeInstanceStore(); + store.States.Add(CreateState("generation-1", OrchestrationStatus.Running, BaseTime)); + + using (var cts = new CancellationTokenSource(TimeSpan.FromSeconds(3))) + { + ServiceBusOrchestrationService service = CreateService(store); + + Task waitTask = service.WaitForOrchestrationAsync( + InstanceId, + "generation-1", + Timeout.InfiniteTimeSpan, + cts.Token); + + // An infinite wait that ignores its token would never complete, so bound the await to + // fail the test instead of hanging the run. + Task finished = await Task.WhenAny(waitTask, Task.Delay(TimeSpan.FromSeconds(30))); + + Assert.AreSame(waitTask, finished, "An infinite wait must stop once its token is cancelled."); + + try + { + Assert.IsNull(await waitTask, "A cancelled wait must not return a state."); + } + catch (OperationCanceledException) + { + // Also acceptable: the polling delay observes the token directly. + } + + // Whichever path ends the wait, it must have kept polling until cancellation rather + // than treating the negative Timeout.InfiniteTimeSpan as an elapsed budget and + // bailing out after the first lookup. + Assert.IsTrue( + store.QueryCount > 1, + $"An infinite wait must poll until cancelled, polled {store.QueryCount} time(s)."); + } + } + + /// + /// In-memory instance store that mimics the query semantics of AzureTableInstanceStore. + /// + sealed class FakeInstanceStore : IOrchestrationServiceInstanceStore + { + public List States { get; } = new List(); + + /// + /// Mimics AzureTableInstanceStore, which excludes ContinuedAsNew rows from the + /// current generation lookup. + /// + public bool FilterContinuedAsNew { get; set; } = true; + + /// + /// Invoked before every state lookup so a test can make state readable after N polls. + /// + public Action OnQuery { get; set; } + + public int PinnedQueryCount { get; private set; } + + public int LatestQueryCount { get; private set; } + + public int QueryCount => this.PinnedQueryCount + this.LatestQueryCount; + + public int MaxHistoryEntryLength => 1024; + + public Task> GetOrchestrationStateAsync(string instanceId, bool allInstances) + { + this.LatestQueryCount++; + this.OnQuery?.Invoke(this); + + IEnumerable matches = this.States + .Where(s => s.OrchestrationInstance.InstanceId == instanceId); + + if (!allInstances && this.FilterContinuedAsNew) + { + matches = matches.Where(s => s.OrchestrationStatus != OrchestrationStatus.ContinuedAsNew); + } + + if (allInstances) + { + return Task.FromResult(matches.Select(Wrap)); + } + + OrchestrationState latest = matches.OrderByDescending(s => s.LastUpdatedTime).FirstOrDefault(); + + return Task.FromResult(latest == null + ? Enumerable.Empty() + : new[] { Wrap(latest) }.AsEnumerable()); + } + + public Task GetOrchestrationStateAsync(string instanceId, string executionId) + { + this.PinnedQueryCount++; + this.OnQuery?.Invoke(this); + + OrchestrationState match = this.States.FirstOrDefault( + s => s.OrchestrationInstance.InstanceId == instanceId && + s.OrchestrationInstance.ExecutionId == executionId); + + return Task.FromResult(match == null ? null : Wrap(match)); + } + + static OrchestrationStateInstanceEntity Wrap(OrchestrationState state) + { + return new OrchestrationStateInstanceEntity { State = state }; + } + + public Task InitializeStoreAsync(bool recreate) => throw new NotImplementedException(); + + public Task DeleteStoreAsync() => throw new NotImplementedException(); + + public Task WriteEntitiesAsync(IEnumerable entities) => throw new NotImplementedException(); + + public Task> GetEntitiesAsync(string instanceId, string executionId) => throw new NotImplementedException(); + + public Task DeleteEntitiesAsync(IEnumerable entities) => throw new NotImplementedException(); + + public Task> GetOrchestrationHistoryEventsAsync(string instanceId, string executionId) => throw new NotImplementedException(); + + public Task PurgeOrchestrationHistoryEventsAsync(DateTime thresholdDateTimeUtc, OrchestrationStateTimeRangeFilterType timeRangeFilterType) => throw new NotImplementedException(); + + public Task WriteJumpStartEntitiesAsync(IEnumerable entities) => throw new NotImplementedException(); + + public Task DeleteJumpStartEntitiesAsync(IEnumerable entities) => throw new NotImplementedException(); + + public Task> GetJumpStartEntitiesAsync(int top) => throw new NotImplementedException(); + } + } +}