3636import org .apache .lucene .util .BitSetIterator ;
3737import org .apache .lucene .util .Bits ;
3838import org .apache .lucene .util .SparseFixedBitSet ;
39- import org .elasticsearch .common .util .concurrent .ConcurrentCollections ;
4039import org .elasticsearch .core .Releasable ;
4140import org .elasticsearch .lucene .util .CombinedBitSet ;
4241import org .elasticsearch .search .dfs .AggregatedDfs ;
5352import java .util .List ;
5453import java .util .Objects ;
5554import java .util .PriorityQueue ;
56- import java .util .Set ;
5755import java .util .concurrent .Callable ;
5856import java .util .concurrent .Executor ;
5957import java .util .stream .Collectors ;
@@ -82,7 +80,6 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable {
8280 // don't create slices with less than this number of docs
8381 private final int minimumDocsPerSlice ;
8482
85- private final Set <Thread > timeoutOverwrites = ConcurrentCollections .newConcurrentSet ();
8683 private volatile boolean timeExceeded = false ;
8784
8885 /** constructor for non-concurrent search */
@@ -374,6 +371,8 @@ private <C extends Collector, T> T search(Weight weight, CollectorManager<C, T>
374371 }
375372 }
376373
374+ private static final ThreadLocal <Boolean > timeoutOverwrites = ThreadLocal .withInitial (() -> false );
375+
377376 /**
378377 * Similar to the lucene implementation, with the following changes made:
379378 * 1) postCollection is performed after each segment is collected. This is needed for aggregations, performed by search threads
@@ -397,12 +396,12 @@ public void search(LeafReaderContextPartition[] leaves, Weight weight, Collector
397396 try {
398397 // Search phase has finished, no longer need to check for timeout
399398 // otherwise the aggregation post-collection phase might get cancelled.
400- boolean added = timeoutOverwrites .add ( Thread . currentThread ()) ;
401- assert added ;
399+ assert timeoutOverwrites .get () == false ;
400+ timeoutOverwrites . set ( true ) ;
402401 doAggregationPostCollection (collector );
403402 } finally {
404- boolean removed = timeoutOverwrites .remove ( Thread . currentThread () );
405- assert removed ;
403+ assert timeoutOverwrites .get ( );
404+ timeoutOverwrites . set ( false ) ;
406405 }
407406 }
408407 }
@@ -420,7 +419,7 @@ public boolean timeExceeded() {
420419 }
421420
422421 public void throwTimeExceededException () {
423- if (timeoutOverwrites .contains ( Thread . currentThread () ) == false ) {
422+ if (timeoutOverwrites .get ( ) == false ) {
424423 throw new TimeExceededException ();
425424 }
426425 }
0 commit comments