Remove dead module awaiter cancellation

src/ModularPipelines/Distributed/Master/DistributedModuleExecutor.cs

@@ -182,7 +182,7 @@ internal static void CompleteCancelledModules(
         IModuleResultRegistrar resultRegistrar,
         CancellationToken cancellationToken)
     {
-        var cancelledModules = scheduler.CancelPendingModules(cancelModuleResultAwaiters: false);
+        var cancelledModules = scheduler.CancelPendingModules();
         resultRegistrar.RegisterTerminatedResultsForCancelledModules(
             cancelledModules,
             new OperationCanceledException(cancellationToken));

src/ModularPipelines/Distributed/Worker/WorkerModuleScheduler.cs

@@ -32,7 +32,7 @@ public void MarkModuleCompleted(Type moduleType, bool success, Exception? except
 
     public ModuleState? GetModuleState(Type moduleType) => null;
 
-    public IReadOnlyList<IModule> CancelPendingModules(bool cancelModuleResultAwaiters = true)
+    public IReadOnlyList<IModule> CancelPendingModules()
     {
         return [];
     }

src/ModularPipelines/Engine/GeneratedModuleMetadata.cs

@@ -208,8 +208,6 @@ IModuleResult CreateFailure(
 
     ILogger GetOutputLogger(IServiceProvider serviceProvider);
 
-    void CancelCompletionSource(IModule module);
-
     void SetCompletionSource(IModule module, IModuleResult result);
 
     Task<IModuleResult> ExecuteAsync(
@@ -257,11 +255,6 @@ public ILogger GetOutputLogger(IServiceProvider serviceProvider)
         return serviceProvider.GetRequiredService<ILogger<TModule>>();
     }
 
-    public void CancelCompletionSource(IModule module)
-    {
-        ((Module<TResult>) module).CompletionSource.TrySetCanceled();
-    }
-
     public void SetCompletionSource(IModule module, IModuleResult result)
     {
         ((Module<TResult>) module).CompletionSource.TrySetResult((ModuleResult<TResult>) result);

src/ModularPipelines/Engine/IModuleScheduler.cs

@@ -49,13 +49,11 @@ internal interface IModuleScheduler : IDisposable
     ModuleState? GetModuleState(Type moduleType);
 
     /// <summary>
-    /// Cancels all modules that are queued or pending (not yet executing)
-    /// This is used when the pipeline is cancelled to ensure TaskCompletionSources are properly completed.
+    /// Cancels all modules that are queued or pending (not yet executing).
+    /// AlwaysRun modules are excluded and are allowed to complete.
+    /// This cancels only the scheduler's internal completion sources. Call
+    /// <c>RegisterTerminatedResultsForCancelledModules</c> for the returned modules to complete their public result tasks.
     /// </summary>
-    /// <param name="cancelModuleResultAwaiters">
-    /// Whether to cancel typed module result awaiters immediately. Set to <see langword="false"/>
-    /// when terminated results will be registered after scheduler cancellation.
-    /// </param>
     /// <returns>The modules transitioned to the completed state by cancellation.</returns>
-    IReadOnlyList<IModule> CancelPendingModules(bool cancelModuleResultAwaiters = true);
+    IReadOnlyList<IModule> CancelPendingModules();
 }

src/ModularPipelines/Engine/IModuleStateTracker.cs

@@ -34,15 +34,12 @@ internal interface IModuleStateTracker
 
     /// <summary>
     /// Cancels all modules that are queued or pending (not yet executing).
-    /// This is used when the pipeline is cancelled to ensure TaskCompletionSources are properly completed.
+    /// This cancels only the scheduler's internal completion sources. Call
+    /// <c>RegisterTerminatedResultsForCancelledModules</c> for the returned modules to complete their public result tasks.
     /// Note: AlwaysRun modules are not cancelled as they should be allowed to complete.
     /// </summary>
-    /// <param name="cancelModuleResultAwaiters">
-    /// Whether to cancel typed module result awaiters immediately. Set to <see langword="false"/>
-    /// when terminated results will be registered after scheduler cancellation.
-    /// </param>
     /// <returns>The modules transitioned to the completed state by cancellation.</returns>
-    IReadOnlyList<IModule> CancelPendingModules(bool cancelModuleResultAwaiters = true);
+    IReadOnlyList<IModule> CancelPendingModules();
 
     /// <summary>
     /// Gets the state for a specific module.

src/ModularPipelines/Engine/ModuleExecutor.cs

@@ -189,7 +189,7 @@ private async Task<IEnumerable<IModule>> ExecuteWithSchedulerAsync(
         catch (Exception exception)
         {
             var cancelledModules =
-                scheduler.CancelPendingModules(cancelModuleResultAwaiters: false);
+                scheduler.CancelPendingModules();
             _resultRegistrar.RegisterTerminatedResultsForCancelledModules(
                 cancelledModules,
                 exception);
@@ -230,7 +230,7 @@ private async Task<IEnumerable<IModule>> ExecuteWithSchedulerAsync(
     private void RegisterCancellationCallback(CancellationTokenSource cancellationTokenSource, IModuleScheduler scheduler)
     {
         cancellationTokenSource.Token.Register(
-            () => scheduler.CancelPendingModules(cancelModuleResultAwaiters: false));
+            () => scheduler.CancelPendingModules());
     }
 
     private async Task<Exception?> ExecuteWorkerPoolAsync(
@@ -277,7 +277,7 @@ await Parallel.ForEachAsync(
                         if (isFirstFailure)
                         {
                             var cancelledModules =
-                                scheduler.CancelPendingModules(cancelModuleResultAwaiters: false);
+                                scheduler.CancelPendingModules();
                             _resultRegistrar.RegisterTerminatedResultsForCancelledModules(
                                 cancelledModules,
                                 ex);

src/ModularPipelines/Engine/ModuleScheduler.cs

@@ -216,21 +216,19 @@ public void MarkModuleCompleted(Type moduleType, bool success, Exception? except
     }
 
     /// <summary>
-    /// Cancels all modules that are queued or pending (not yet executing)
-    /// This is used when the pipeline is cancelled to ensure TaskCompletionSources are properly completed
+    /// Cancels all modules that are queued or pending (not yet executing).
+    /// This cancels only the scheduler's internal completion sources. Call
+    /// <c>RegisterTerminatedResultsForCancelledModules</c> for the returned modules to complete their public result tasks.
     /// Note: AlwaysRun modules are not cancelled as they should be allowed to complete.
     /// </summary>
-    /// <param name="cancelModuleResultAwaiters">
-    /// Whether to cancel typed module result awaiters immediately.
-    /// </param>
-    public IReadOnlyList<IModule> CancelPendingModules(bool cancelModuleResultAwaiters = true)
+    public IReadOnlyList<IModule> CancelPendingModules()
     {
         if (IsDisposed)
         {
             return [];
         }
 
-        return _stateTracker.CancelPendingModules(cancelModuleResultAwaiters);
+        return _stateTracker.CancelPendingModules();
     }
 
     public void Dispose()

src/ModularPipelines/Engine/ModuleStateTracker.cs

@@ -257,7 +257,7 @@ public void MarkModuleCompleted(Type moduleType, bool success, Exception? except
     }
 
     /// <inheritdoc />
-    public IReadOnlyList<IModule> CancelPendingModules(bool cancelModuleResultAwaiters = true)
+    public IReadOnlyList<IModule> CancelPendingModules()
     {
         List<(ModuleState State, ModuleExecutionState OriginalState)> cancelledModules;
 
@@ -288,10 +288,6 @@ public IReadOnlyList<IModule> CancelPendingModules(bool cancelModuleResultAwaite
         foreach (var (moduleState, _) in cancelledModules)
         {
             moduleState.CompletionSource.TrySetCanceled();
-            if (cancelModuleResultAwaiters)
-            {
-                ModuleCompletionSourceCanceller.Cancel(moduleState.Module, moduleState.ModuleType);
-            }
         }
 
         // Logging outside lock

test/ModularPipelines.Distributed.UnitTests/Master/DistributedModuleExecutorTests.cs

@@ -192,7 +192,7 @@ private static Mock<IModuleScheduler> CreateMockScheduler(params ModuleState[] m
                 return tcs.Task;
             });
         scheduler.Setup(s => s.MarkModuleStarted(It.IsAny<Type>())).Returns(true);
-        scheduler.Setup(s => s.CancelPendingModules(It.IsAny<bool>())).Returns([]);
+        scheduler.Setup(s => s.CancelPendingModules()).Returns([]);
 
         return scheduler;
     }
@@ -898,7 +898,7 @@ public async Task Cancellation_After_Start_Claim_Completes_Module_Result(
         var module = new DistributedModule();
         var moduleState = new ModuleState(module, typeof(DistributedModule));
         var scheduler = CreateMockScheduler(moduleState);
-        scheduler.Setup(s => s.CancelPendingModules(false))
+        scheduler.Setup(s => s.CancelPendingModules())
             .Returns([]);
         scheduler.Setup(s => s.MarkModuleStarted(typeof(DistributedModule)))
             .Returns(() =>
@@ -954,7 +954,7 @@ public async Task Publish_Failure_Completes_Module_Result(
         var module = new DistributedModule();
         var moduleState = new ModuleState(module, typeof(DistributedModule));
         var scheduler = CreateMockScheduler(moduleState);
-        scheduler.Setup(s => s.CancelPendingModules(false))
+        scheduler.Setup(s => s.CancelPendingModules())
             .Returns([]);
 
         var publishException = new InvalidOperationException("Broker unavailable");

test/ModularPipelines.UnitTests/Distributed/DistributedModuleExecutorTests.cs

@@ -30,7 +30,7 @@ public void CompleteCancelledModules_RegistersTerminatedResults()
         var scheduler = new Mock<IModuleScheduler>();
         var resultRegistrar = new Mock<IModuleResultRegistrar>();
         scheduler
-            .Setup(x => x.CancelPendingModules(false))
+            .Setup(x => x.CancelPendingModules())
             .Returns(cancelledModules);
 
         DistributedModuleExecutor.CompleteCancelledModules(

test/ModularPipelines.UnitTests/Engine/GeneratedModuleMetadataTests.cs

@@ -99,21 +99,6 @@ public async Task Cancelled_Result_Registration_Defers_AlwaysRun_Completion()
         }
     }
 
-    [Test]
-    public async Task Generated_Runtime_Cancels_Typed_Completion_Source()
-    {
-        var module = new GeneratedMetadataDependencyModule();
-
-        var found = GeneratedModuleMetadata.TryGetRuntime(module.GetType(), out var runtime);
-        runtime.CancelCompletionSource(module);
-
-        using (Assert.Multiple())
-        {
-            await Assert.That(found).IsTrue();
-            await Assert.That(module.CompletionSource.Task.IsCanceled).IsTrue();
-        }
-    }
-
     [Test]
     public async Task Generated_Runtime_Resolves_Unbuffered_Output_Logger()
     {

test/ModularPipelines.UnitTests/Engine/ModuleExecutorLoggingTests.cs

@@ -103,7 +103,7 @@ public async Task SuccessfulCompletion_DoesNotLogCancellation()
 
         var logOutput = logs.ToString();
         await Assert.That(logOutput).DoesNotContain("Cancellation triggered");
-        scheduler.Verify(x => x.CancelPendingModules(false), Times.Once);
+        scheduler.Verify(x => x.CancelPendingModules(), Times.Once);
     }
 
     [Test]
@@ -119,7 +119,7 @@ public async Task SchedulerAndAlwaysRunFaults_AreAggregatedAfterRegisteringTermi
         scheduler.SetupGet(x => x.ReadyModules).Returns(readyModules.Reader);
         scheduler.Setup(x => x.RunSchedulerAsync(It.IsAny<CancellationToken>()))
             .Returns(Task.CompletedTask);
-        scheduler.Setup(x => x.CancelPendingModules(false))
+        scheduler.Setup(x => x.CancelPendingModules())
             .Returns(cancelledModules);
         var schedulerFactory = new Mock<IModuleSchedulerFactory>();
         schedulerFactory.Setup(x => x.Create()).Returns(scheduler.Object);
@@ -668,7 +668,7 @@ public async Task CancelPendingModules_WithPendingModule_LogsCancellation()
     }
 
     [Test]
-    public async Task CancelPendingModules_CompletesPendingModuleAwaitable()
+    public async Task CancelPendingModules_LeavesModuleResultForRegistrarCompletion()
     {
         var module = new LaterModule();
         var state = new ModuleState(module, module.GetType());
@@ -680,8 +680,8 @@ public async Task CancelPendingModules_CompletesPendingModuleAwaitable()
 
         using (Assert.Multiple())
         {
-            await Assert.That(module.CompletionSource.Task.IsCanceled).IsTrue();
-            await Assert.That(((IInternalModule) module).ResultTask.IsCompleted).IsTrue();
+            await Assert.That(state.CompletionSource.Task.IsCanceled).IsTrue();
+            await Assert.That(((IInternalModule) module).ResultTask.IsCompleted).IsFalse();
         }
     }
 

test/ModularPipelines.UnitTests/Engine/TaskCompletionSourceTests.cs

@@ -292,7 +292,7 @@ private static async Task<WeakReference> CreateAbandonedSchedulerFaultAsync(Exce
         scheduler.SetupGet(x => x.ReadyModules).Returns(readyModules.Reader);
         scheduler.Setup(x => x.RunSchedulerAsync(It.IsAny<CancellationToken>()))
             .Returns(schedulerTaskSource.Task);
-        scheduler.Setup(x => x.CancelPendingModules(false)).Returns([]);
+        scheduler.Setup(x => x.CancelPendingModules()).Returns([]);
         var schedulerFactory = new Mock<IModuleSchedulerFactory>();
         schedulerFactory.Setup(x => x.Create()).Returns(scheduler.Object);
         var alwaysRunHandler = new Mock<IAlwaysRunHandler>();