Add StopAsync and keep a dispatched run from starting after Stop [minor]

CLAUDE.md

@@ -90,7 +90,7 @@ Import-Module ./scripts/PSBuild.psm1
 - Exception handling: Exceptions from actions are captured and can be rethrown via `RethrowExceptions()`
 - The library prevents overlapping executions - if an action is still running when the next interval arrives, it waits until the current action completes
 - `LastRunTimestamp` (a `Stopwatch` timestamp) tracks when the action last executed (based on `IntervalType`), and is what scheduling measures from. `LastRunTime` records the same moment as wall-clock time for reporting only: the wall clock can be set back, which would hold the action off for as long as it moved
-- Polling can be stopped/restarted via `Stop()` and `Restart()`/`RestartAsync()`
+- Polling can be stopped/restarted via `Stop()`/`StopAsync()` and `Restart()`/`RestartAsync()`. A dispatched run re-checks `StopGeneration` under the lock before it begins, so `Stop()` skips one that has not started; `StopAsync()` also waits for the polling loop and any run in progress
 
 ## Testing
 

IntervalAction.Test/IntervalActionTests.cs

@@ -2,6 +2,10 @@
 
 [assembly: DoNotParallelize]
 
+// Most tests here exercise the synchronous Stop() on purpose: what it does while a run is still in
+// flight is the behavior under test, and StopAsync() would wait that run out. StopAsync() has tests of its own.
+#pragma warning disable CA1849 // Call async methods when in an async method
+
 namespace ktsu.IntervalAction.Test;
 
 using Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -46,7 +50,7 @@
 			{
 				Interlocked.Increment(ref executions);
 				// Simulate a long running task.
 				Thread.Sleep(500);
 			},
 			IntervalType = IntervalType.FromLastStart
 		};
@@ -295,27 +299,82 @@
 		{
 			PollingInterval = TimeSpan.FromMilliseconds(10),
 			ActionInterval = TimeSpan.Zero,
-			Action = () => counter++,
+			Action = () => Interlocked.Increment(ref counter),
 			IntervalType = IntervalType.FromLastStart
 		};
 
 		IntervalAction intervalAction = IntervalAction.Start(options);
 		// Allow the polling loop to execute a few times.
 		await Task.Delay(30).ConfigureAwait(false);
 
-		// Act
-		intervalAction.Stop();
-		// Await the polling task to ensure the loop has exited.
-		await intervalAction.PollingTask.ConfigureAwait(false);
-		int counterAfterStop = counter;
+		// Act: StopAsync returns only once the loop has exited and the last run has finished, so
+		// nothing is left that could still change the counter
+		await intervalAction.StopAsync().ConfigureAwait(false);
+		int counterAfterStop = Volatile.Read(ref counter);
 
 		// Wait additional time to verify no further actions are executed.
 		await Task.Delay(30).ConfigureAwait(false);
-		Assert.AreEqual(counterAfterStop, counter);
+		Assert.AreEqual(counterAfterStop, Volatile.Read(ref counter));
+		Assert.IsTrue(intervalAction.PollingTask.IsCompleted, "The polling loop should have exited.");
+		Assert.IsTrue(intervalAction.ActionTask?.IsCompleted ?? true, "No run should still be in flight.");
 
 		intervalAction.RethrowExceptions();
 	}
 
+	[TestMethod]
+	public async Task NoActionRunsOrStartsAfterStopAsync()
+	{
+		// Arrange: polling and running as fast as the class allows, so a stop usually lands while a run
+		// has been dispatched to the thread pool but has not begun, or is in progress
+		const int trials = 500;
+		int startedAfterStop = 0;
+		int runningAtStop = 0;
+
+		for (int trial = 0; trial < trials; trial++)
+		{
+			int running = 0;
+			int stopped = 0;
+			int startedLate = 0;
+			IntervalActionOptions options = new()
+			{
+				PollingInterval = TimeSpan.FromMilliseconds(1),
+				ActionInterval = TimeSpan.Zero,
+				Action = () =>
+				{
+					Interlocked.Increment(ref running);
+					if (Volatile.Read(ref stopped) != 0)
+					{
+						Interlocked.Increment(ref startedLate);
+					}
+
+					Thread.SpinWait(200);
+					Interlocked.Decrement(ref running);
+				},
+				IntervalType = IntervalType.FromLastStart
+			};
+
+			IntervalAction intervalAction = IntervalAction.Start(options);
+			await Task.Delay(trial % 3).ConfigureAwait(false);
+
+			// Act
+			await intervalAction.StopAsync().ConfigureAwait(false);
+			Volatile.Write(ref stopped, 1);
+			runningAtStop += Volatile.Read(ref running);
+
+			// Give a run that was dispatched but not yet begun the chance to begin, if one could
+			await Task.Delay(1).ConfigureAwait(false);
+			await Task.Yield();
+			startedAfterStop += Volatile.Read(ref startedLate);
+
+			intervalAction.RethrowExceptions();
+		}
+
+		// Assert
+		string outcome = $"Of {trials} trials, a run was still executing when StopAsync returned in {runningAtStop}, and a run began after it returned in {startedAfterStop}.";
+		Assert.AreEqual(0, runningAtStop, outcome);
+		Assert.AreEqual(0, startedAfterStop, outcome);
+	}
+
 	[TestMethod]
 	public async Task RethrowExceptionsThrowsException()
 	{

IntervalAction/IntervalAction.cs

@@ -221,8 +221,65 @@ public static IntervalAction Start(IntervalActionOptions intervalActionOptions)
 	/// <remarks>
 	/// A restart that is still pending when this is called is cancelled rather than left to resume
 	/// polling once it finishes. A restart requested after this call is unaffected.
+	/// A run that was dispatched but has not begun by the time this is called is skipped. A run that
+	/// had already begun is not waited for: use <see cref="StopAsync"/> to wait until it has finished.
 	/// </remarks>
-	public void Stop()
+	public void Stop() => StopCore(out _);
+
+	/// <summary>
+	/// Stops the polling of the action, then waits for the polling loop and any run of the action
+	/// that is still in flight to finish.
+	/// </summary>
+	/// <returns>
+	/// A task that completes once no run of the action is executing and none can start, until the
+	/// next <see cref="Restart"/> or <see cref="RestartAsync"/>.
+	/// </returns>
+	/// <remarks>
+	/// <para>
+	/// Use this rather than <see cref="Stop"/> before disposing anything the action uses. A run that
+	/// <see cref="Stop"/> skips never starts, and a run that had already started is awaited.
+	/// </para>
+	/// <para>
+	/// The returned task does not fault when the action did. Call <see cref="RethrowExceptions"/>
+	/// afterwards to observe a failure from the final run.
+	/// </para>
+	/// <para>
+	/// Do not await this from inside the action: it waits for the action to finish, so it would never complete.
+	/// </para>
+	/// </remarks>
+	public async Task StopAsync()
+	{
+		// A restart already pending when the stop lands is awaited too: it may be between deciding to
+		// start a loop and publishing it as PollingTask. One requested after the stop is not waited on.
+		StopCore(out Task pendingRestart);
+		await WaitAndDiscardOutcomeAsync(pendingRestart).ConfigureAwait(false);
+
+		Task pollingTask;
+		lock (Lock)
+		{
+			pollingTask = PollingTask;
+		}
+
+		await WaitAndDiscardOutcomeAsync(pollingTask).ConfigureAwait(false);
+
+		// The loop has exited, so it can no longer claim a new run: what is in ActionTask now is the last one
+		Task? actionTask;
+		lock (Lock)
+		{
+			actionTask = ActionTask;
+		}
+
+		if (actionTask is not null)
+		{
+			await WaitAndDiscardOutcomeAsync(actionTask).ConfigureAwait(false);
+		}
+	}
+
+	/// <summary>
+	/// Stops polling and cancels the running asynchronous action, if any.
+	/// </summary>
+	/// <param name="pendingRestart">The restart that was pending, or last completed, when the stop took effect.</param>
+	private void StopCore(out Task pendingRestart)
 	{
 		CancellationTokenSource? actionCancellation;
 
@@ -231,6 +288,7 @@ public void Stop()
 			StopGeneration++;
 			EndPolling();
 			actionCancellation = ActionCancellation;
+			pendingRestart = RestartGate;
 		}
 
 		// Cancelled outside the lock, because cancelling runs the action's own token callbacks
@@ -337,7 +395,9 @@ private async Task RestartCoreAsync(Task previousRestart, long stopGeneration)
 
 				while (shouldPoll)
 				{
-					TryRun();
+					// Passing the generation this loop was started under lets TryRun refuse to claim a run
+					// if Stop() lands between the ShouldPoll read above and the claim
+					TryRun(stopGeneration);
 
 					try
 					{
@@ -375,6 +435,27 @@ private async Task RestartCoreAsync(Task previousRestart, long stopGeneration)
 	/// </summary>
 	/// <returns><c>true</c> if the action was started; otherwise, <c>false</c>.</returns>
 	internal bool TryRun()
+	{
+		long stopGeneration;
+
+		lock (Lock)
+		{
+			stopGeneration = StopGeneration;
+		}
+
+		return TryRun(stopGeneration);
+	}
+
+	/// <summary>
+	/// Attempts to run the action if the specified interval has passed since the last execution, and
+	/// no <see cref="Stop"/> has happened since <paramref name="stopGeneration"/> was read.
+	/// </summary>
+	/// <param name="stopGeneration">
+	/// The value of <see cref="StopGeneration"/> the caller is running under. The run is neither
+	/// claimed, nor, once claimed, begun, if a <see cref="Stop"/> has happened since.
+	/// </param>
+	/// <returns><c>true</c> if the action was dispatched; otherwise, <c>false</c>.</returns>
+	private bool TryRun(long stopGeneration)
 	{
 		Ensure.NotNull(Action);
 
@@ -393,21 +474,26 @@ internal bool TryRun()
 				ActionTask = null;
 			}
 
+			if (StopGeneration != stopGeneration)
+			{
+				return false;
+			}
+
 			if (ActionInterval >= TimeSpan.Zero && ActionTask is null && HasIntervalElapsed())
 			{
 				if (AsyncAction is { } asyncAction)
 				{
 					CancellationTokenSource cancellation = new();
 					ActionCancellation = cancellation;
-					ActionTask = Task.Run(() => RunAsyncAction(asyncAction, cancellation), CancellationToken.None);
+					ActionTask = Task.Run(() => RunAsyncAction(asyncAction, cancellation, stopGeneration), CancellationToken.None);
 					return true;
 				}
 
 				ActionTask = Task.Run(() =>
 				{
-					if (IntervalType == IntervalType.FromLastStart)
+					if (!TryBeginRun(stopGeneration))
 					{
-						RecordRun();
+						return;
 					}
 
 					Action();
@@ -425,19 +511,51 @@ internal bool TryRun()
 		}
 	}
 
+	/// <summary>
+	/// Decides, once a dispatched run is finally executing, whether it may still call the action.
+	/// </summary>
+	/// <param name="stopGeneration">The value of <see cref="StopGeneration"/> the run was claimed under.</param>
+	/// <returns><see langword="false"/> if <see cref="Stop"/> has been called since the run was claimed.</returns>
+	/// <remarks>
+	/// The run is claimed under <see cref="Lock"/>, but the thread pool starts it later. Without this
+	/// check, a <see cref="Stop"/> in between would return, the polling loop would exit, and the action
+	/// would still begin afterwards, possibly against resources the caller had already disposed.
+	/// </remarks>
+	private bool TryBeginRun(long stopGeneration)
+	{
+		lock (Lock)
+		{
+			if (StopGeneration != stopGeneration)
+			{
+				return false;
+			}
+
+			// Recorded under the same lock as the check, rather than through RecordRun, so a run that
+			// begins is always one that a stop had not yet reached
+			if (IntervalType == IntervalType.FromLastStart)
+			{
+				LastRunTime = DateTimeOffset.Now;
+				LastRunTimestamp = Stopwatch.GetTimestamp();
+			}
+
+			return true;
+		}
+	}
+
 	/// <summary>
 	/// Runs <see cref="AsyncAction"/> once, so that the returned task completes when its work does.
 	/// </summary>
 	/// <param name="asyncAction">The action to run.</param>
 	/// <param name="cancellation">The source whose token the action receives, and which this run disposes.</param>
+	/// <param name="stopGeneration">The value of <see cref="StopGeneration"/> the run was claimed under.</param>
 	/// <returns>A task that completes, or faults, when the action does.</returns>
-	private async Task RunAsyncAction(Func<CancellationToken, Task> asyncAction, CancellationTokenSource cancellation)
+	private async Task RunAsyncAction(Func<CancellationToken, Task> asyncAction, CancellationTokenSource cancellation, long stopGeneration)
 	{
 		try
 		{
-			if (IntervalType == IntervalType.FromLastStart)
+			if (!TryBeginRun(stopGeneration))
 			{
-				RecordRun();
+				return;
 			}
 
 			try

README.md

@@ -116,7 +116,11 @@ var intervalAction = IntervalAction.Start(new()
 
 // For cleanup when done
 await Task.Delay(60000); // Run for 1 minute
-intervalAction.Stop();
+
+// StopAsync cancels the token, then waits for the run in progress to finish, so nothing the
+// action uses is disposed while it is still running. Stop() returns without waiting.
+await intervalAction.StopAsync();
+intervalAction.RethrowExceptions(); // Reports a failure from the final run, if there was one
 ```
 
 ### Dynamic Interval Adjustment
@@ -161,9 +165,9 @@ The main class for scheduling recurring actions.
 | Name | Parameters | Return Type | Description |
 |------|------------|-------------|-------------|
 | `Start` | `IntervalActionOptions options` | `IntervalAction` | Static factory method to create and start an interval action |
-| `Stop` | | `void` | Stops the scheduled action from running |
+| `Stop` | | `void` | Stops the scheduled action from running; a run already in progress is not waited for |
+| `StopAsync` | | `Task` | Stops the scheduled action, then completes once the polling loop has exited and any run in progress has finished. Use it before disposing anything the action uses |
 | `Restart` | | `void` | Restarts a previously stopped action |
-| `Dispose` | | `void` | Cleans up resources and stops the action |
 
 ### `IntervalActionOptions` Class