Enhance Migration Engine and Unit Tests

- Improved error logging in MigrationEngine to provide clearer context on exceptions.
- Added functionality to InMemoryMigrationSource to simulate shutdown behavior during body reads.
- Enhanced InMemoryMigrationTarget to allow custom exceptions for failure scenarios.
- Introduced TimerRecordingTimeProvider to facilitate testing of timer-related functionality.
- Expanded HaltThresholdTests to cover edge cases involving mixed benign and fault skips.
- Updated MigrationEngineBodyRetryTests to ensure proper handling of shutdowns during body reads.
- Added tests to verify behavior when categories are empty or already completed.
- Implemented checks for checkpoint conflicts and cancellation scenarios in MigrationEngineFailurePathTests.
- Enhanced MigrationEngineHaltTests to evaluate mixed skip scenarios and their impact on halting.
- Improved MigrationEngineOrderingTests to ensure blocked categories behave correctly.
- Added tests to MigrationEngineResumeTests to verify that restart behavior maintains original start time.
- Updated MigrationEngineRunCategoriesTests to handle cases where no categories are selected.
- Enhanced MigrationEngineSkipReasonTests to ensure accurate recording of multiple skip reasons.
- Added tests to MigrationEngineThrottleTests to verify behavior during pauses and shutdowns.

Author: warwickschroeder

src/ServiceControl.Persistence/DataMigration/MigrationEngine.cs

@@ -159,7 +159,8 @@ public async Task<MigrationCheckpoint> RunCategoryAsync(MigrationCategory catego
         }
         catch (Exception ex)
         {
-            var reason = $"{ex.GetType().Name} at cursor {checkpoint.Cursor ?? "the start"}: {ex.Message}";
+            var position = checkpoint.Cursor is null ? "at the start" : $"at cursor {checkpoint.Cursor}";
+            var reason = $"{ex.GetType().Name} {position}: {ex.Message}";
             logger.LogError(ex, "Category {CategoryId} halted at cursor {Cursor}", category.Id, checkpoint.Cursor);
             return await Settle(checkpoint with { State = MigrationCategoryState.Halted, LastError = reason }, cancellationToken);
         }

src/ServiceControl.UnitTests/Migration/Fakes/InMemoryMigrationSource.cs

@@ -24,6 +24,9 @@ public sealed class InMemoryMigrationSource : IMigrationSource
 
     public void FailBodyReads(string sourceId, int times, Exception failure) => bodyFailures[sourceId] = (times, failure);
 
+    /// <summary>Makes the body read for this row behave like a host shutting down: the token is cancelled and the read throws.</summary>
+    public (string SourceId, CancellationTokenSource Source)? StopOnBodyRead { get; set; }
+
     public int BodyReadAttempts(string sourceId) => bodyReadAttempts.GetValueOrDefault(sourceId);
 
     public Task Open(CancellationToken cancellationToken = default) => Task.CompletedTask;
@@ -70,6 +73,13 @@ public async IAsyncEnumerable<MigrationBatch> Read(
         await Task.Yield();
 
         var attempt = bodyReadAttempts[sourceId] = BodyReadAttempts(sourceId) + 1;
+
+        if (StopOnBodyRead is { } stop && stop.SourceId == sourceId)
+        {
+            await stop.Source.CancelAsync();
+            throw new OperationCanceledException(stop.Source.Token);
+        }
+
         if (bodyFailures.TryGetValue(sourceId, out var failures) && attempt <= failures.Times)
         {
             throw failures.Failure;

src/ServiceControl.UnitTests/Migration/Fakes/InMemoryMigrationTarget.cs

@@ -16,6 +16,12 @@ public sealed class InMemoryMigrationTarget(IMigrationCheckpointStore checkpoint
     public int DefaultBatchSize { get; set; } = 3;
     public string NoBatchSizeFor { get; set; }
     public int? FailOnCallNumber { get; set; }
+
+    /// <summary>What FailOnCallNumber throws, when the default simulated failure is the wrong shape for the test.</summary>
+    public Exception FailWith { get; set; }
+
+    /// <summary>Cancels the token on this call and then writes normally, so the stop surfaces from the source's next batch.</summary>
+    public (int CallNumber, CancellationTokenSource Source)? CancelOnCall { get; set; }
     public (int CallNumber, CancellationTokenSource Source)? StopOnCall { get; set; }
     int callCount;
 
@@ -45,7 +51,12 @@ public async Task<MigrationWriteResult> Write(
 
         if (FailOnCallNumber == callCount)
         {
-            throw new InvalidOperationException($"Simulated failure on write {callCount}");
+            throw FailWith ?? new InvalidOperationException($"Simulated failure on write {callCount}");
+        }
+
+        if (CancelOnCall is { } cancel && cancel.CallNumber == callCount)
+        {
+            await cancel.Source.CancelAsync();
         }
 
         var keys = writtenKeysByCategory.TryGetValue(category.Id, out var existingKeys) ? existingKeys : writtenKeysByCategory[category.Id] = [];

src/ServiceControl.UnitTests/Migration/Fakes/TimerRecordingTimeProvider.cs

@@ -0,0 +1,35 @@
+#nullable enable
+namespace ServiceControl.UnitTests.Migration.Fakes;
+
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using Microsoft.Extensions.Time.Testing;
+
+// Signals each timer the engine creates, so a test only advances the clock once a pause is waiting on it.
+public sealed class TimerRecordingTimeProvider : TimeProvider
+{
+    readonly FakeTimeProvider clock = new();
+
+    public SemaphoreSlim TimerCreated { get; } = new(0);
+
+    public List<TimeSpan> DueTimes { get; } = [];
+
+    public void Advance(TimeSpan delta) => clock.Advance(delta);
+
+    public override DateTimeOffset GetUtcNow() => clock.GetUtcNow();
+
+    public override long GetTimestamp() => clock.GetTimestamp();
+
+    public override long TimestampFrequency => clock.TimestampFrequency;
+
+    public override TimeZoneInfo LocalTimeZone => clock.LocalTimeZone;
+
+    public override ITimer CreateTimer(TimerCallback callback, object? state, TimeSpan dueTime, TimeSpan period)
+    {
+        var timer = clock.CreateTimer(callback, state, dueTime, period);
+        DueTimes.Add(dueTime);
+        TimerCreated.Release();
+        return timer;
+    }
+}

src/ServiceControl.UnitTests/Migration/HaltThresholdTests.cs

@@ -36,12 +36,40 @@ public void Does_not_halt_a_large_category_with_only_a_small_proportion_skipped(
     [Test]
     public void Exactly_at_both_boundaries_does_not_halt_because_both_must_be_exceeded()
     {
-        // 100 of 2,000 is exactly 5% and exactly the floor. "Exceed" means strictly past, not "at".
+        // Exactly the floor, which settles it before the proportion is ever worked out. "Exceed" means strictly past, not "at".
         var exceeded = HaltThreshold.Exceeded(skippedCount: 100, totalCount: 2_000, percentThreshold: 5, minimumFloor: 100);
 
         Assert.That(exceeded, Is.False);
     }
 
+    [Test]
+    public void Exactly_on_the_proportion_does_not_halt_once_the_floor_is_behind_it()
+    {
+        // 101 of 2,020 is exactly 5% with the floor already passed, so this is the only shape that
+        // reaches the proportion comparison and depends on it being strictly greater.
+        var exceeded = HaltThreshold.Exceeded(skippedCount: 101, totalCount: 2_020, percentThreshold: 5, minimumFloor: 100);
+
+        Assert.That(exceeded, Is.False);
+    }
+
+    [Test]
+    public void A_hair_past_the_proportion_halts_once_the_floor_is_behind_it()
+    {
+        // 101 of 2,000 is 5.05%: the same skip count as above, one row's worth over the line.
+        var exceeded = HaltThreshold.Exceeded(skippedCount: 101, totalCount: 2_000, percentThreshold: 5, minimumFloor: 100);
+
+        Assert.That(exceeded, Is.True);
+    }
+
+    [Test]
+    public void A_skip_count_with_nothing_processed_never_halts_and_never_divides_by_zero()
+    {
+        // Past the floor with a zero total, which is the only input that reaches the division guard.
+        var exceeded = HaltThreshold.Exceeded(skippedCount: 101, totalCount: 0, percentThreshold: 5, minimumFloor: 100);
+
+        Assert.That(exceeded, Is.False);
+    }
+
     [Test]
     public void No_rows_processed_yet_never_halts()
     {

src/ServiceControl.UnitTests/Migration/MigrationEngineBodyRetryTests.cs

@@ -4,6 +4,7 @@ namespace ServiceControl.UnitTests.Migration;
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging.Abstractions;
 using Microsoft.Extensions.Time.Testing;
@@ -43,13 +44,25 @@ public async Task A_body_that_fails_until_the_last_attempt_then_succeeds_is_writ
         }
     }
 
-    [Test]
-    public async Task A_body_read_that_fails_as_a_defect_halts_the_category_on_the_first_attempt_without_skipping_the_message()
+    // Every type the engine treats as a defect. Each fails the same way on every attempt, so retrying one
+    // would only turn a code fault into skipped messages.
+    static readonly Exception[] Defects =
+    [
+        new NotSupportedException("this source cannot read bodies"),
+        new NotImplementedException("not written yet"),
+        new InvalidOperationException("the session is closed"),
+        new ArgumentException("the id is not a document id"),
+        new NullReferenceException("no attachment"),
+        new InvalidCastException("not an attachment")
+    ];
+
+    [TestCaseSource(nameof(Defects))]
+    public async Task A_body_read_that_fails_as_a_defect_halts_the_category_on_the_first_attempt_without_skipping_the_message(Exception defect)
     {
         var category = MigrationCategoryRegistry.Find("UnresolvedAndRetryIssuedFailedMessages")!;
         var source = new InMemoryMigrationSource();
         source.Seed(category.Id, Row("msg-1"));
-        source.FailBodyReads("msg-1", MigrationEngine.MaxBodyReadAttempts, new NotSupportedException("this source cannot read bodies"));
+        source.FailBodyReads("msg-1", MigrationEngine.MaxBodyReadAttempts, defect);
         var checkpointStore = new InMemoryMigrationCheckpointStore();
         var target = new InMemoryMigrationTarget(checkpointStore);
         var options = new MigrationEngineOptions(TimeSpan.Zero, 5, 100, []) { BodyRetryBackoff = TimeSpan.Zero };
@@ -60,13 +73,78 @@ public async Task A_body_read_that_fails_as_a_defect_halts_the_category_on_the_f
         using (Assert.EnterMultipleScope())
         {
             Assert.That(checkpoint.State, Is.EqualTo(MigrationCategoryState.Halted));
-            Assert.That(checkpoint.LastError, Does.Contain(nameof(NotSupportedException)));
+            Assert.That(checkpoint.LastError, Does.Contain(defect.GetType().Name));
             Assert.That(source.BodyReadAttempts("msg-1"), Is.EqualTo(1));
             Assert.That(checkpoint.SkippedCount, Is.Zero);
             Assert.That(checkpoint.SkipReasons, Is.Null);
         }
     }
 
+    [Test]
+    public async Task A_shutdown_during_a_body_read_stops_the_run_instead_of_skipping_the_message()
+    {
+        // Retrying a shutdown to the attempt limit and then recording the message as permanently
+        // unreadable is the one path here that silently loses a customer's failed message.
+        var category = MigrationCategoryRegistry.Find("UnresolvedAndRetryIssuedFailedMessages")!;
+        var source = new InMemoryMigrationSource();
+        source.Seed(category.Id, Row("msg-1"));
+        using var stopping = new CancellationTokenSource();
+        source.StopOnBodyRead = ("msg-1", stopping);
+        var checkpointStore = new InMemoryMigrationCheckpointStore();
+        var target = new InMemoryMigrationTarget(checkpointStore);
+        var options = new MigrationEngineOptions(TimeSpan.Zero, 5, 100, []) { BodyRetryBackoff = TimeSpan.Zero };
+        var engine = new MigrationEngine(source, target, checkpointStore, new FakeTimeProvider(), options, NullLogger<MigrationEngine>.Instance);
+
+        Assert.ThrowsAsync<OperationCanceledException>(() => engine.RunCategoryAsync(category, stopping.Token));
+
+        var persisted = await checkpointStore.Read(category.Id);
+        using (Assert.EnterMultipleScope())
+        {
+            Assert.That(source.BodyReadAttempts("msg-1"), Is.EqualTo(1), "a shutdown is not a transient body failure, so it is not retried");
+            Assert.That(persisted!.State, Is.EqualTo(MigrationCategoryState.InProgress));
+            Assert.That(persisted.SkippedCount, Is.Zero, "the message is still there to copy on the next run");
+            Assert.That(target.WrittenRows(category.Id), Is.Empty);
+        }
+    }
+
+    [Test]
+    public async Task The_configured_backoff_is_waited_between_body_read_attempts()
+    {
+        // Without the wait, three attempts against a body store that is briefly down all fail inside a
+        // millisecond and the message is skipped for an outage it would have survived.
+        var category = MigrationCategoryRegistry.Find("UnresolvedAndRetryIssuedFailedMessages")!;
+        var source = new InMemoryMigrationSource();
+        source.Seed(category.Id, Row("msg-1"));
+        var body = new MigrationBody(new byte[] { 1 }, "text/plain");
+        source.SetBody("msg-1", body);
+        source.FailBodyReads("msg-1", MigrationEngine.MaxBodyReadAttempts - 1, new TimeoutException("body store unreachable"));
+        var checkpointStore = new InMemoryMigrationCheckpointStore();
+        var target = new InMemoryMigrationTarget(checkpointStore);
+        var clock = new TimerRecordingTimeProvider();
+        var backoff = TimeSpan.FromMilliseconds(200);
+        var options = new MigrationEngineOptions(TimeSpan.Zero, 5, 100, []) { BodyRetryBackoff = backoff };
+        var engine = new MigrationEngine(source, target, checkpointStore, clock, options, NullLogger<MigrationEngine>.Instance);
+
+        var runTask = engine.RunCategoryAsync(category);
+
+        // Two failures, so a wait after each before the attempt that succeeds.
+        for (var waitNumber = 1; waitNumber <= MigrationEngine.MaxBodyReadAttempts - 1; waitNumber++)
+        {
+            Assert.That(await clock.TimerCreated.WaitAsync(TimeSpan.FromSeconds(5)), Is.True, $"backoff {waitNumber} never started");
+            Assert.That(runTask.IsCompleted, Is.False, $"backoff {waitNumber} should still be pending");
+            clock.Advance(backoff);
+        }
+
+        var checkpoint = await runTask.WaitAsync(TimeSpan.FromSeconds(5));
+
+        using (Assert.EnterMultipleScope())
+        {
+            Assert.That(checkpoint.State, Is.EqualTo(MigrationCategoryState.Complete));
+            Assert.That(clock.DueTimes, Is.EqualTo(new[] { backoff, backoff }), "no wait after the final attempt, which has nothing left to retry");
+            Assert.That(target.WrittenRows(category.Id).Single().Body, Is.EqualTo(body));
+        }
+    }
+
     [Test]
     public async Task The_skip_warning_for_an_unreadable_body_carries_the_last_attempts_exception()
     {

src/ServiceControl.UnitTests/Migration/MigrationEngineCopyTests.cs

@@ -40,6 +40,72 @@ public async Task Copies_every_row_and_finishes_Complete_when_nothing_was_skippe
         }
     }
 
+    [Test]
+    public async Task A_category_with_no_rows_finishes_Complete_without_a_cursor()
+    {
+        // The ordinary state of several required categories on a small instance: nothing to copy is a
+        // finished category, not a category that never ran.
+        var category = MigrationCategoryRegistry.Find("MessageRedirects")!;
+        var source = new InMemoryMigrationSource();
+        var checkpointStore = new InMemoryMigrationCheckpointStore();
+        var target = new InMemoryMigrationTarget(checkpointStore);
+        var options = new MigrationEngineOptions(TimeSpan.Zero, 5, 100, []);
+        var engine = new MigrationEngine(source, target, checkpointStore, new FakeTimeProvider(), options, NullLogger<MigrationEngine>.Instance);
+
+        var checkpoint = await engine.RunCategoryAsync(category);
+
+        using (Assert.EnterMultipleScope())
+        {
+            Assert.That(checkpoint.State, Is.EqualTo(MigrationCategoryState.Complete));
+            Assert.That((checkpoint.CopiedCount, checkpoint.SkippedCount), Is.EqualTo((0L, 0L)));
+            Assert.That(checkpoint.Cursor, Is.Null);
+            Assert.That(checkpoint.SettledAt, Is.Not.Null);
+        }
+    }
+
+    [Test]
+    public async Task The_moment_a_category_settles_comes_from_the_injected_clock()
+    {
+        var category = MigrationCategoryRegistry.Find("KnownEndpoints")!;
+        var source = new InMemoryMigrationSource();
+        source.Seed(category.Id, Row("a"));
+        var checkpointStore = new InMemoryMigrationCheckpointStore();
+        var target = new InMemoryMigrationTarget(checkpointStore);
+        var settledAt = new DateTimeOffset(2026, 3, 4, 5, 6, 7, TimeSpan.Zero);
+        var clock = new FakeTimeProvider(settledAt);
+        var options = new MigrationEngineOptions(TimeSpan.Zero, 5, 100, []);
+        var engine = new MigrationEngine(source, target, checkpointStore, clock, options, NullLogger<MigrationEngine>.Instance);
+
+        var checkpoint = await engine.RunCategoryAsync(category);
+
+        Assert.That(checkpoint.SettledAt, Is.EqualTo(settledAt.UtcDateTime), "a wall-clock read here would drift from every other time the migration reports");
+    }
+
+    [Test]
+    public async Task A_category_already_CompleteWithErrors_is_left_alone_on_a_second_run()
+    {
+        // Finished with a few skips is finished. Re-reading it would copy the whole category again and
+        // count its skips a second time.
+        var category = MigrationCategoryRegistry.Find("KnownEndpoints")!;
+        var source = new InMemoryMigrationSource();
+        source.Seed(category.Id, Row("a"), Row("b"));
+        var checkpointStore = new InMemoryMigrationCheckpointStore();
+        var finishedWithSkips = new MigrationCheckpoint(category.Id, MigrationCategoryState.CompleteWithErrors, "b", 1, 1, 2,
+            new Dictionary<MigrationSkipReason, long> { [MigrationSkipReason.BodyUnreadable] = 1 }, DateTime.UtcNow, DateTime.UtcNow, DateTime.UtcNow, null);
+        await checkpointStore.Upsert(finishedWithSkips);
+        var target = new InMemoryMigrationTarget(checkpointStore);
+        var options = new MigrationEngineOptions(TimeSpan.Zero, 5, 100, []);
+        var engine = new MigrationEngine(source, target, checkpointStore, new FakeTimeProvider(), options, NullLogger<MigrationEngine>.Instance);
+
+        var checkpoint = await engine.RunCategoryAsync(category);
+
+        using (Assert.EnterMultipleScope())
+        {
+            Assert.That(checkpoint, Is.EqualTo(finishedWithSkips with { Version = 1 }), "the row is read back untouched, at the version the seeding save left it");
+            Assert.That(target.WrittenRows(category.Id), Is.Empty);
+        }
+    }
+
     [Test]
     public async Task A_category_already_Complete_is_left_alone_on_a_second_run()
     {