Cancellation support for LimitedConcurrencyLevelTaskScheduler, #1253 (#1275)

* Cancellation support for LimitedConcurrencyLevelTaskScheduler, #1253

* Remove CancelAsync that is not available on .NET Framework

* Throw on QueueTask after shutdown in LimitedConcurrencyLevelTaskScheduler, #1253

Match Java's ExecutorService.shutdown() RejectedExecutionException semantics
instead of silently dropping post-shutdown submissions, which would leave
awaiters hanging on a Task that never completes.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Optimize cleanup of cancellation token source

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: paulirwin

src/Lucene.Net.TestFramework/Support/Threading/LimitedConcurrencyLevelTaskScheduler.cs

@@ -48,7 +48,6 @@ to the software or derivative works that you create that run directly on a Micro
 Office or Microsoft Dynamics).
 */
 
-using J2N.Threading.Atomic;
 using System;
 using System.Collections.Generic;
 using System.Threading;
@@ -64,8 +63,6 @@ namespace Lucene.Net.Support.Threading
     /// </summary>
     internal class LimitedConcurrencyLevelTaskScheduler : TaskScheduler
     {
-        private readonly AtomicBoolean shutDown = new AtomicBoolean(false);
-
         // Indicates whether the current thread is processing work items.
         [ThreadStatic]
         private static bool _currentThreadIsProcessingItems;
@@ -76,21 +73,42 @@ internal class LimitedConcurrencyLevelTaskScheduler : TaskScheduler
         // The maximum concurrency level allowed by this scheduler.
         private readonly int _maxDegreeOfParallelism;
 
+        // A cancellation token for preventing queueing new work when shut down
+        private readonly CancellationToken _cancellationToken;
+
         // Indicates whether the scheduler is currently processing work items.
         private int _delegatesQueuedOrRunning = 0;
 
-        // Creates a new instance with the specified degree of parallelism.
-        public LimitedConcurrencyLevelTaskScheduler(int maxDegreeOfParallelism)
+        /// <summary>
+        /// Creates a new instance with the specified degree of parallelism.
+        /// </summary>
+        /// <param name="maxDegreeOfParallelism">The max degree of parallelism for tasks.</param>
+        /// <param name="cancellationToken">
+        /// A cancellation token that is used to shut down the task scheduler in an orderly manner.
+        /// If cancellation is requested, this scheduler will throw a <see cref="TaskSchedulerException"/> if a
+        /// new task is attempted to be queued.
+        /// This behaves like <c>ExecutorService.shutdown()</c> in Java (with <c>RejectedExecutionException</c>),
+        /// allowing any running tasks to finish.
+        /// </param>
+        /// <exception cref="ArgumentOutOfRangeException">if <paramref name="maxDegreeOfParallelism"/> is less than 1.</exception>
+        public LimitedConcurrencyLevelTaskScheduler(int maxDegreeOfParallelism, CancellationToken cancellationToken = default)
         {
             if (maxDegreeOfParallelism < 1) throw new ArgumentOutOfRangeException(nameof(maxDegreeOfParallelism));
             _maxDegreeOfParallelism = maxDegreeOfParallelism;
+            _cancellationToken = cancellationToken;
         }
 
         // Queues a task to the scheduler.
         protected sealed override void QueueTask(Task task)
         {
             // Don't queue any more work.
-            if (shutDown) return;
+            if (_cancellationToken.IsCancellationRequested)
+            {
+                // LUCENENET NOTE: in Java's `ExecutorService.shutdown()`, this would be RejectedExecutionException.
+                // The equivalent for a TaskScheduler is TaskSchedulerException, but the framework will wrap what
+                // we throw here in a TaskSchedulerException.
+                throw new InvalidOperationException("The task scheduler was shut down and is not accepting new tasks.");
+            }
 
             // Add the task to the list of tasks to be processed.  If there aren't enough
             // delegates currently queued or running to process tasks, schedule another.
@@ -203,17 +221,13 @@ protected sealed override IEnumerable<Task> GetScheduledTasks()
             }
         }
 
-        /// <summary>
-        /// Stops this TaskScheduler from queuing new tasks.
-        /// </summary>
-        public void Shutdown()
-        {
-            shutDown.Value = true;
-        }
-
         /// <summary>
         /// Gets a value indicating whether this TaskScheduler has been shut down.
         /// </summary>
-        public bool IsShutdown => shutDown;
+        /// <remarks>
+        /// This simply returns whether the cancellation token provided to the constructor
+        /// has requested cancellation.
+        /// </remarks>
+        public bool IsShutdown => _cancellationToken.IsCancellationRequested;
     }
 }

src/Lucene.Net.TestFramework/Util/LuceneTestCase.cs

@@ -2084,14 +2084,16 @@ public static IndexSearcher NewSearcher(IndexReader r, bool maybeWrap, bool wrap
             {
                 int threads = 0;
                 LimitedConcurrencyLevelTaskScheduler ex;
+                CancellationTokenSource cts = null;
                 if (random.NextBoolean())
                 {
                     ex = null;
                 }
                 else
                 {
+                    cts = new CancellationTokenSource(); // LUCENENET NOTE: this is cleaned up in ReaderClosedListenerAnonymousClass
                     threads = TestUtil.NextInt32(random, 1, 8);
-                    ex = new LimitedConcurrencyLevelTaskScheduler(threads);
+                    ex = new LimitedConcurrencyLevelTaskScheduler(threads, cts.Token);
                     //ex = new ThreadPoolExecutor(threads, threads, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<IThreadRunnable>(), new NamedThreadFactory("LuceneTestCase"));
                     // uncomment to intensify LUCENE-3840
                     // ex.prestartAllCoreThreads();
@@ -2102,7 +2104,7 @@ public static IndexSearcher NewSearcher(IndexReader r, bool maybeWrap, bool wrap
                     {
                         Console.WriteLine("NOTE: newSearcher using ExecutorService with " + threads + " threads");
                     }
-                    r.AddReaderDisposedListener(new ReaderClosedListenerAnonymousClass(ex));
+                    r.AddReaderDisposedListener(new ReaderClosedListenerAnonymousClass(cts));
                 }
                 IndexSearcher ret;
                 if (wrapWithAssertions)
@@ -3262,20 +3264,25 @@ public static double RandomGaussian() // LUCENENET: CA1822: Mark members as stat
             return Random.NextGaussian();
         }
 
+#nullable enable
         private sealed class ReaderClosedListenerAnonymousClass : IReaderDisposedListener
         {
-            private readonly LimitedConcurrencyLevelTaskScheduler ex;
+            private readonly CancellationTokenSource? cts; // LUCENENET-specific: cancellation support, can be null
 
-            public ReaderClosedListenerAnonymousClass(LimitedConcurrencyLevelTaskScheduler ex)
+            public ReaderClosedListenerAnonymousClass(CancellationTokenSource? cts)
             {
-                this.ex = ex;
+                this.cts = cts;
             }
 
             public void OnDispose(IndexReader reader)
             {
-                ex?.Shutdown();
-                //TestUtil.ShutdownExecutorService(ex);
+                if (cts != null)
+                {
+                    cts.Cancel();
+                    cts.Dispose();
+                }
             }
         }
+#nullable restore
     }
 }

src/Lucene.Net.Tests.Suggest/Spell/TestSpellChecker.cs

@@ -472,11 +472,19 @@ public async Task TestConcurrentAccess()
             int numThreads = 5 + Random.nextInt(5);
             var tasks = new ConcurrentBag<Task>();
             SpellCheckWorker[] workers = new SpellCheckWorker[numThreads];
-            var executor = new LimitedConcurrencyLevelTaskScheduler(numThreads); // LUCENENET NOTE: Not sure why in Java they decided to pass the max concurrent threads as all of the threads, but this demonstrates how to use a custom TaskScheduler in .NET.
+
+            // LUCENENET NOTE: Not sure why in Java they decided to pass the max concurrent threads as all of the threads, but this demonstrates how to use a custom TaskScheduler in .NET.
+            // LUCENENET NOTE: This cancellation token/source is intentionally separate from the one below, because it is solely used as an equivalent
+            // to ExecutorService.shutdown(), which just stops queueing new tasks.
+            using var executorShutdown = new CancellationTokenSource();
+            var executor = new LimitedConcurrencyLevelTaskScheduler(numThreads, executorShutdown.Token);
+
             using var shutdown = new CancellationTokenSource();
             var cancellationToken = shutdown.Token;
+
             var stop = new AtomicBoolean(false);
             var taskFactory = new TaskFactory(executor);
+
             for (int i = 0; i < numThreads; i++)
             {
                 SpellCheckWorker spellCheckWorker = new SpellCheckWorker(this, r, stop, cancellationToken, taskNum: i);
@@ -495,7 +503,12 @@ public async Task TestConcurrentAccess()
             }
 
             stop.Value = true;
-            executor.Shutdown(); // Stop allowing tasks to queue
+
+            // LUCENENET NOTE: This is technically pointless, since by this point we've already queued all the tasks.
+            // Leaving this here for compatibility and possible future test changes that might use it properly.
+            // ReSharper disable once MethodHasAsyncOverload - not available in .NET Framework
+            executorShutdown.Cancel(); // Stop queueing new tasks
+
             try
             {
                 // wait for 60 seconds - usually this is very fast but coverage runs could take quite long
@@ -526,7 +539,7 @@ public async Task TestConcurrentAccess()
             AssertSearchersClosed();
         }
 
-        private void AssertLastSearcherOpen(int numSearchers)
+        private static void AssertLastSearcherOpen(int numSearchers)
         {
             assertEquals(numSearchers, searchers.Count);
             IndexSearcher[] searcherArray = searchers.ToArray();
@@ -545,7 +558,7 @@ private void AssertLastSearcherOpen(int numSearchers)
             }
         }
 
-        private void AssertSearchersClosed()
+        private static void AssertSearchersClosed()
         {
             foreach (IndexSearcher searcher in searchers)
             {

src/Lucene.Net.Tests/Search/TestIndexSearcher.cs

@@ -134,7 +134,7 @@ public virtual void TestHugeN()
                 }
             }
 
-            // LUCENENET: .NET doesn't have a way to shut down the TaskScheduler explicitly
+            // LUCENENET: shutdown not needed here since all searches above run synchronously
             //TestUtil.ShutdownExecutorService(service);
         }
 
@@ -361,7 +361,7 @@ public virtual void TestCancellation_MultiThreaded_CancelDuringSearch_Throws()
                 Assume.That(r.Leaves.Count >= 2, "Test requires a multi-segment index");
 
                 using CancellationTokenSource cts = new CancellationTokenSource();
-                TaskScheduler service = new LimitedConcurrencyLevelTaskScheduler(4);
+                TaskScheduler service = new LimitedConcurrencyLevelTaskScheduler(4); // LUCENENET NOTE: intentionally NOT passing cts.Token here since that parameter is for shutdown only, and that's not what we're testing
                 IndexSearcher searcher = new CancelAfterFirstLeafSearcher(r, service, cts);
 
                 Exception ex = Assert.Catch(() => searcher.Search(new MatchAllDocsQuery(), 10, cts.Token));

src/Lucene.Net.Tests/Support/Threading/JSR166TestCase.cs

@@ -351,7 +351,8 @@ public void joinPool(TaskScheduler exec)
         {
             try
             {
-                exec.Shutdown();
+                // LUCENENET NOTE: no need to call a shutdown cancellation token here,
+                // since any tasks would already be queued by the time we get here
                 assertTrue(exec.AwaitTermination(TimeSpan.FromMilliseconds(LONG_DELAY_MS)));
             }
             // catch (SecurityException ok) // LUCENENET - not needed
@@ -448,6 +449,9 @@ public TaskState(TaskScheduler scheduler)
 
             public void NewTask(Action action)
             {
+                if (_factory.Scheduler is LimitedConcurrencyLevelTaskScheduler { IsShutdown: true })
+                    return;
+
                 var task = _factory.StartNew(action);
                 _tasks.Add(task);
             }
@@ -517,14 +521,6 @@ public static int GetTaskCount(this TaskScheduler scheduler)
             return 0;
         }
 
-        public static void Shutdown(this TaskScheduler scheduler)
-        {
-            if (scheduler is LimitedConcurrencyLevelTaskScheduler lcl)
-            {
-                lcl.Shutdown();
-            }
-        }
-
         public static bool IsTerminated(this TaskScheduler scheduler)
         {
             if (scheduler is LimitedConcurrencyLevelTaskScheduler lcl

src/Lucene.Net.Tests/Support/Threading/TestLimitedConcurrencyLevelTaskScheduler.cs

@@ -1,6 +1,7 @@
 // Based on tests from Apache Harmony:
 // https://github.com/apache/harmony/blob/02970cb7227a335edd2c8457ebdde0195a735733/classlib/modules/concurrent/src/test/java/ThreadPoolExecutorTest.java
 
+using Lucene.Net.Attributes;
 using NUnit.Framework;
 using System;
 using System.Threading;
@@ -135,7 +136,6 @@ public void TestGetCompletedTaskCount()
             AssumeTrue($"Expected 1, but got {p2.GetCompletedTaskCount()} - this may be a timing issue.", p2.GetCompletedTaskCount() == 1);
 
             // LUCENENET NOTE: not catching SecurityException because that's not relevant here
-            p2.Shutdown();
             joinPool(p2);
         }
 
@@ -191,18 +191,64 @@ public void TestGetTaskCount()
         }
 
         /// <summary>
-        /// <see cref="LimitedConcurrencyLevelTaskScheduler.IsShutdown"/> is false before shutdown, true after
+        /// Tests that a canceled token does not queue new tasks
         /// </summary>
         [Test]
-        public void TestIsShutdown()
+        public void TestCancellation()
         {
-            var p1 = new LimitedConcurrencyLevelTaskScheduler(1);
-            assertFalse(p1.IsShutdown);
-            p1.Shutdown(); // LUCENENET NOTE: not catching SecurityException because that's not relevant here
-            assertTrue(p1.IsShutdown);
+            using var cts = new CancellationTokenSource();
+            TaskScheduler p1 = new LimitedConcurrencyLevelTaskScheduler(1, cts.Token);
+
+            assertEquals(0, p1.GetTaskCount());
+
+            try
+            {
+                p1.Execute(MediumRunnable); // queue first task, not yet canceled
+                // NOTE: no need to sleep here, we're just checking if tasks got queued
+                assertEquals(1, p1.GetTaskCount());
+
+                cts.Cancel();
+                p1.Execute(MediumRunnable); // queue second task (or try to)
+                assertEquals(1, p1.GetTaskCount()); // cancellation should prevent queueing new tasks
+            }
+            catch (Exception)
+            {
+                unexpectedException();
+            }
+
+            assertTrue(cts.Token.IsCancellationRequested);
             joinPool(p1);
         }
 
+        /// <summary>
+        /// Submitting a task after shutdown via <see cref="TaskFactory.StartNew(Action)"/>
+        /// surfaces an <see cref="InvalidOperationException"/> (wrapped by the TPL in a
+        /// <see cref="TaskSchedulerException"/>), matching Java's
+        /// <c>RejectedExecutionException</c> semantics for <c>ExecutorService</c>
+        /// after <c>shutdown()</c>. Silently dropping the task would let awaiters
+        /// hang on a Task that will never complete.
+        /// </summary>
+        /// <remarks>
+        /// LUCENENET specific - this exercises the direct <see cref="TaskScheduler"/>
+        /// submission path used outside of <see cref="JSR166TestCaseExtensions"/>
+        /// (e.g., <see cref="Lucene.Net.Search.IndexSearcher"/>'s parallel slice path).
+        /// </remarks>
+        [Test, LuceneNetSpecific]
+        public void TestQueueTaskAfterShutdownThrows()
+        {
+            using var cts = new CancellationTokenSource();
+            var scheduler = new LimitedConcurrencyLevelTaskScheduler(1, cts.Token);
+            var factory = new TaskFactory(scheduler);
+
+            cts.Cancel();
+
+            // NOTE: we intentionally do not want to pass cts.Token to StartNew, to test QueueTask.
+            var ex = Assert.Throws<TaskSchedulerException>(() => factory.StartNew(() => { }, CancellationToken.None));
+            assertTrue(
+                $"Expected InvalidOperationException inner, got {ex!.InnerException?.GetType().FullName}",
+                ex.InnerException is InvalidOperationException);
+        }
+
         /// <summary>
         /// isTerminated is false before termination, true after
         /// </summary>
@@ -215,7 +261,8 @@ public void TestIsShutdown()
         [Test]
         public void TestIsTerminated()
         {
-            TaskScheduler p1 = new LimitedConcurrencyLevelTaskScheduler(1);
+            using var cts = new CancellationTokenSource();
+            TaskScheduler p1 = new LimitedConcurrencyLevelTaskScheduler(1, cts.Token);
             assertFalse(p1.IsTerminated());
 
             try
@@ -224,7 +271,9 @@ public void TestIsTerminated()
             }
             finally
             {
-                p1.Shutdown(); // LUCENENET NOTE: not catching SecurityException because that's not relevant here
+                // LUCENENET NOTE: not catching SecurityException because that's not relevant here
+                // LUCENENET NOTE: canceling here is of questionable utility, since no more tasks are queued
+                cts.Cancel();
             }
 
             try