36
36
import org .apache .lucene .util .BitSetIterator ;
37
37
import org .apache .lucene .util .Bits ;
38
38
import org .apache .lucene .util .SparseFixedBitSet ;
39
- import org .elasticsearch .common .util .concurrent .ConcurrentCollections ;
40
39
import org .elasticsearch .core .Releasable ;
41
40
import org .elasticsearch .lucene .util .CombinedBitSet ;
42
41
import org .elasticsearch .search .dfs .AggregatedDfs ;
53
52
import java .util .List ;
54
53
import java .util .Objects ;
55
54
import java .util .PriorityQueue ;
56
- import java .util .Set ;
57
55
import java .util .concurrent .Callable ;
58
56
import java .util .concurrent .Executor ;
59
57
@@ -80,7 +78,6 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable {
80
78
// don't create slices with less than this number of docs
81
79
private final int minimumDocsPerSlice ;
82
80
83
- private final Set <Thread > timeoutOverwrites = ConcurrentCollections .newConcurrentSet ();
84
81
private volatile boolean timeExceeded = false ;
85
82
86
83
/** constructor for non-concurrent search */
@@ -356,6 +353,8 @@ private <C extends Collector, T> T search(Weight weight, CollectorManager<C, T>
356
353
}
357
354
}
358
355
356
+ private static final ThreadLocal <Boolean > timeoutOverwrites = ThreadLocal .withInitial (() -> false );
357
+
359
358
/**
360
359
* Similar to the lucene implementation, with the following changes made:
361
360
* 1) postCollection is performed after each segment is collected. This is needed for aggregations, performed by search threads
@@ -379,12 +378,12 @@ public void search(List<LeafReaderContext> leaves, Weight weight, Collector coll
379
378
try {
380
379
// Search phase has finished, no longer need to check for timeout
381
380
// otherwise the aggregation post-collection phase might get cancelled.
382
- boolean added = timeoutOverwrites .add ( Thread . currentThread ()) ;
383
- assert added ;
381
+ assert timeoutOverwrites .get () == false ;
382
+ timeoutOverwrites . set ( true ) ;
384
383
doAggregationPostCollection (collector );
385
384
} finally {
386
- boolean removed = timeoutOverwrites .remove ( Thread . currentThread () );
387
- assert removed ;
385
+ assert timeoutOverwrites .get ( );
386
+ timeoutOverwrites . set ( false ) ;
388
387
}
389
388
}
390
389
}
@@ -402,7 +401,7 @@ public boolean timeExceeded() {
402
401
}
403
402
404
403
public void throwTimeExceededException () {
405
- if (timeoutOverwrites .contains ( Thread . currentThread () ) == false ) {
404
+ if (timeoutOverwrites .get ( ) == false ) {
406
405
throw new TimeExceededException ();
407
406
}
408
407
}
0 commit comments