Pass cancellation token further down in multithreaded code, #922

Author: paulirwin

src/Lucene.Net/Search/IndexSearcher.cs

@@ -539,7 +539,7 @@ protected virtual TopDocs Search(Weight weight, ScoreDoc? after, int nDocs, Canc
 
                 for (int i = 0; i < m_leafSlices.Length; i++) // search each sub
                 {
-                    runner.Submit(new SearcherCallableNoSort(@lock, this, m_leafSlices[i], weight, after, nDocs, hq).Call);
+                    runner.Submit(new SearcherCallableNoSort(@lock, this, m_leafSlices[i], weight, after, nDocs, hq, cancellationToken).Call);
                 }
 
                 int totalHits = 0;
@@ -649,7 +649,7 @@ protected virtual TopFieldDocs Search(Weight weight, FieldDoc? after, int nDocs,
 
                 for (int i = 0; i < m_leafSlices.Length; i++) // search each leaf slice
                 {
-                    runner.Submit(new SearcherCallableWithSort(@lock, this, m_leafSlices[i], weight, after, nDocs, topCollector, sort, doDocScores, doMaxScore).Call);
+                    runner.Submit(new SearcherCallableWithSort(@lock, this, m_leafSlices[i], weight, after, nDocs, topCollector, sort, doDocScores, doMaxScore, cancellationToken).Call);
                 }
 
                 int totalHits = 0;
@@ -756,6 +756,8 @@ protected virtual void Search(IList<AtomicReaderContext> leaves, Weight weight,
                     }
                 }
             }
+
+            cancellationToken.ThrowIfCancellationRequested(); // LUCENENET specific - cancellation support
         }
 
         /// <summary>
@@ -859,21 +861,23 @@ public virtual Weight CreateNormalizedWeight(Query query)
             private readonly int nDocs;
             private readonly HitQueue hq;
             private readonly LeafSlice slice;
+            private readonly CancellationToken cancellationToken;
 
-            public SearcherCallableNoSort(ReentrantLock @lock, IndexSearcher searcher, LeafSlice slice, Weight weight, ScoreDoc? after, int nDocs, HitQueue hq)
+            public SearcherCallableNoSort(ReentrantLock @lock, IndexSearcher searcher, LeafSlice slice, Weight weight, ScoreDoc? after, int nDocs, HitQueue hq, CancellationToken cancellationToken)
             {
                 this.@lock = @lock;
                 this.searcher = searcher;
                 this.weight = weight;
                 this.after = after;
                 this.nDocs = nDocs;
                 this.hq = hq;
+                this.cancellationToken = cancellationToken;
                 this.slice = slice;
             }
 
             public TopDocs Call()
             {
-                TopDocs docs = searcher.Search(slice.Leaves, weight, after, nDocs);
+                TopDocs docs = searcher.Search(slice.Leaves, weight, after, nDocs, cancellationToken);
                 ScoreDoc[] scoreDocs = docs.ScoreDocs;
                 //it would be so nice if we had a thread-safe insert
                 @lock.Lock();
@@ -911,8 +915,9 @@ public TopDocs Call()
             private readonly FieldDoc? after;
             private readonly bool doDocScores;
             private readonly bool doMaxScore;
+            private readonly CancellationToken cancellationToken;
 
-            public SearcherCallableWithSort(ReentrantLock @lock, IndexSearcher searcher, LeafSlice slice, Weight weight, FieldDoc? after, int nDocs, TopFieldCollector hq, Sort sort, bool doDocScores, bool doMaxScore)
+            public SearcherCallableWithSort(ReentrantLock @lock, IndexSearcher searcher, LeafSlice slice, Weight weight, FieldDoc? after, int nDocs, TopFieldCollector hq, Sort sort, bool doDocScores, bool doMaxScore, CancellationToken cancellationToken)
             {
                 this.@lock = @lock;
                 this.searcher = searcher;
@@ -924,14 +929,15 @@ public SearcherCallableWithSort(ReentrantLock @lock, IndexSearcher searcher, Lea
                 this.after = after;
                 this.doDocScores = doDocScores;
                 this.doMaxScore = doMaxScore;
+                this.cancellationToken = cancellationToken;
             }
 
             private readonly FakeScorer fakeScorer = new FakeScorer();
 
             public TopFieldDocs Call()
             {
                 if (Debugging.AssertsEnabled) Debugging.Assert(slice.Leaves.Length == 1);
-                TopFieldDocs docs = searcher.Search(slice.Leaves, weight, after, nDocs, sort, true, doDocScores || sort.NeedsScores, doMaxScore);
+                TopFieldDocs docs = searcher.Search(slice.Leaves, weight, after, nDocs, sort, true, doDocScores || sort.NeedsScores, doMaxScore, cancellationToken);
                 @lock.Lock();
                 try
                 {
@@ -991,7 +997,7 @@ public void Dispose()
 
             public void Submit(Func<T> task)
             {
-                this.service.Submit(task);
+                this.service.Submit(task, cancellationToken);
                 ++numTasks;
             }
 

src/Lucene.Net/Support/Threading/TaskSchedulerCompletionService.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Lucene.Net.Support.Threading
@@ -31,9 +32,9 @@ public TaskSchedulerCompletionService(TaskScheduler scheduler)
             this.factory = new TaskFactory<T>(scheduler ?? TaskScheduler.Default);
         }
 
-        public Task<T> Submit(Func<T> task)
+        public Task<T> Submit(Func<T> task, CancellationToken cancellationToken = default)
         {
-            var t = factory.StartNew(task);
+            var t = factory.StartNew(task, cancellationToken);
             taskQueue.Enqueue(t);
             return t;
         }
@@ -43,4 +44,4 @@ public Task<T> Take()
             return taskQueue.Dequeue();
         }
     }
-}
+}