Cancellation support for LimitedConcurrencyLevelTaskScheduler, #1253

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,34 @@ 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, <see cref="QueueTask"/> will not queue any more tasks.
+        /// This behaves like <c>ExecutorService.shutdown()</c> in Java, 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) return;
 
             // 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 +213,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,15 @@ public static IndexSearcher NewSearcher(IndexReader r, bool maybeWrap, bool wrap
             {
                 int threads = 0;
                 LimitedConcurrencyLevelTaskScheduler ex;
+                var cts = new CancellationTokenSource();
                 if (random.NextBoolean())
                 {
                     ex = null;
                 }
                 else
                 {
                     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 +2103,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)
@@ -3264,17 +3265,17 @@ public static double RandomGaussian() // LUCENENET: CA1822: Mark members as stat
 
         private sealed class ReaderClosedListenerAnonymousClass : IReaderDisposedListener
         {
-            private readonly LimitedConcurrencyLevelTaskScheduler ex;
+            private readonly CancellationTokenSource cts; // LUCENENET-specific: cancellation support
 
-            public ReaderClosedListenerAnonymousClass(LimitedConcurrencyLevelTaskScheduler ex)
+            public ReaderClosedListenerAnonymousClass(CancellationTokenSource cts)
             {
-                this.ex = ex;
+                this.cts = cts;
             }
 
             public void OnDispose(IndexReader reader)
             {
-                ex?.Shutdown();
-                //TestUtil.ShutdownExecutorService(ex);
+                cts.Cancel();
+                cts.Dispose();
             }
         }
     }

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,11 @@ 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.
+            await executorShutdown.CancelAsync(); // Stop queueing new tasks
+
             try
             {
                 // wait for 60 seconds - usually this is very fast but coverage runs could take quite long
@@ -526,7 +538,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 +557,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

@@ -135,7 +135,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,15 +190,32 @@ 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);
         }
 
@@ -215,7 +231,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 +241,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