36
36
import org .opensearch .action .OriginalIndices ;
37
37
import org .opensearch .action .support .IndicesOptions ;
38
38
import org .opensearch .cluster .ClusterState ;
39
+ import org .opensearch .cluster .node .DiscoveryNode ;
39
40
import org .opensearch .cluster .routing .GroupShardsIterator ;
40
41
import org .opensearch .common .UUIDs ;
41
42
import org .opensearch .common .collect .Tuple ;
42
43
import org .opensearch .common .settings .ClusterSettings ;
43
44
import org .opensearch .common .settings .Settings ;
45
+ import org .opensearch .common .unit .TimeValue ;
44
46
import org .opensearch .common .util .concurrent .AtomicArray ;
45
47
import org .opensearch .common .util .concurrent .OpenSearchExecutors ;
46
48
import org .opensearch .common .util .set .Sets ;
55
57
import org .opensearch .index .shard .ShardNotFoundException ;
56
58
import org .opensearch .search .SearchPhaseResult ;
57
59
import org .opensearch .search .SearchShardTarget ;
60
+ import org .opensearch .search .builder .SearchSourceBuilder ;
58
61
import org .opensearch .search .internal .AliasFilter ;
59
62
import org .opensearch .search .internal .InternalSearchResponse ;
60
63
import org .opensearch .search .internal .ShardSearchContextId ;
65
68
import org .opensearch .test .OpenSearchTestCase ;
66
69
import org .opensearch .threadpool .TestThreadPool ;
67
70
import org .opensearch .threadpool .ThreadPool ;
71
+ import org .opensearch .transport .ReceiveTimeoutTransportException ;
68
72
import org .opensearch .transport .Transport ;
69
73
import org .junit .After ;
70
74
import org .junit .Before ;
89
93
import java .util .function .BiFunction ;
90
94
import java .util .stream .IntStream ;
91
95
96
+ import org .mockito .Mockito ;
97
+
98
+ import static org .opensearch .action .search .SearchTransportService .QUERY_ACTION_NAME ;
92
99
import static org .opensearch .tasks .TaskResourceTrackingService .TASK_RESOURCE_USAGE ;
93
100
import static org .hamcrest .Matchers .equalTo ;
94
101
import static org .hamcrest .Matchers .greaterThanOrEqualTo ;
@@ -138,6 +145,7 @@ private AbstractSearchAsyncAction<SearchPhaseResult> createAction(
138
145
false ,
139
146
expected ,
140
147
resourceUsage ,
148
+ false ,
141
149
new SearchShardIterator (null , null , Collections .emptyList (), null )
142
150
);
143
151
}
@@ -151,6 +159,7 @@ private AbstractSearchAsyncAction<SearchPhaseResult> createAction(
151
159
final boolean catchExceptionWhenExecutePhaseOnShard ,
152
160
final AtomicLong expected ,
153
161
final TaskResourceUsage resourceUsage ,
162
+ final boolean blockTheFirstQueryPhase ,
154
163
final SearchShardIterator ... shards
155
164
) {
156
165
@@ -179,7 +188,7 @@ private AbstractSearchAsyncAction<SearchPhaseResult> createAction(
179
188
.setNodeId (randomAlphaOfLengthBetween (1 , 5 ))
180
189
.build ();
181
190
threadPool .getThreadContext ().addResponseHeader (TASK_RESOURCE_USAGE , taskResourceInfo .toString ());
182
-
191
+ AtomicBoolean firstShard = new AtomicBoolean ( true );
183
192
return new AbstractSearchAsyncAction <SearchPhaseResult >(
184
193
"test" ,
185
194
logger ,
@@ -207,7 +216,13 @@ private AbstractSearchAsyncAction<SearchPhaseResult> createAction(
207
216
) {
208
217
@ Override
209
218
protected SearchPhase getNextPhase (final SearchPhaseResults <SearchPhaseResult > results , SearchPhaseContext context ) {
210
- return null ;
219
+ return new SearchPhase ("test" ) {
220
+ @ Override
221
+ public void run () {
222
+ listener .onResponse (new SearchResponse (null , null , 0 , 0 , 0 , 0 , null , null ));
223
+ assertingListener .onPhaseEnd (context , null );
224
+ }
225
+ };
211
226
}
212
227
213
228
@ Override
@@ -218,6 +233,17 @@ protected void executePhaseOnShard(
218
233
) {
219
234
if (failExecutePhaseOnShard ) {
220
235
listener .onFailure (new ShardNotFoundException (shardIt .shardId ()));
236
+ } else if (blockTheFirstQueryPhase && firstShard .compareAndSet (true , false )) {
237
+ // Sleep and throw ReceiveTimeoutTransportException to simulate node blocked
238
+ try {
239
+ Thread .sleep (request .source ().timeout ().millis ());
240
+ } catch (InterruptedException e ) {}
241
+ ;
242
+ DiscoveryNode node = Mockito .mock (DiscoveryNode .class );
243
+ Mockito .when (node .getName ()).thenReturn ("test_nodes" );
244
+ listener .onFailure (
245
+ new ReceiveTimeoutTransportException (node , QUERY_ACTION_NAME , "request_id [171] timed out after [413ms]" )
246
+ );
221
247
} else {
222
248
if (catchExceptionWhenExecutePhaseOnShard ) {
223
249
try {
@@ -227,6 +253,7 @@ protected void executePhaseOnShard(
227
253
}
228
254
} else {
229
255
listener .onResponse (new QuerySearchResult ());
256
+
230
257
}
231
258
}
232
259
}
@@ -587,6 +614,7 @@ public void onFailure(Exception e) {
587
614
false ,
588
615
new AtomicLong (),
589
616
new TaskResourceUsage (randomLong (), randomLong ()),
617
+ false ,
590
618
shards
591
619
);
592
620
action .run ();
@@ -635,6 +663,7 @@ public void onFailure(Exception e) {
635
663
false ,
636
664
new AtomicLong (),
637
665
new TaskResourceUsage (randomLong (), randomLong ()),
666
+ false ,
638
667
shards
639
668
);
640
669
action .run ();
@@ -688,6 +717,7 @@ public void onFailure(Exception e) {
688
717
catchExceptionWhenExecutePhaseOnShard ,
689
718
new AtomicLong (),
690
719
new TaskResourceUsage (randomLong (), randomLong ()),
720
+ false ,
691
721
shards
692
722
);
693
723
action .run ();
@@ -791,6 +821,41 @@ public void testOnPhaseListenersWithDfsType() throws InterruptedException {
791
821
assertEquals (0 , testListener .getPhaseCurrent (searchDfsQueryThenFetchAsyncAction .getSearchPhaseName ()));
792
822
}
793
823
824
+ public void testExecutePhaseOnShardBlockAndRetrunPartialResult () {
825
+ // on shard is blocked in query phase
826
+ final Index index = new Index ("test" , UUID .randomUUID ().toString ());
827
+
828
+ final SearchShardIterator [] shards = IntStream .range (0 , 2 + randomInt (4 ))
829
+ .mapToObj (i -> new SearchShardIterator (null , new ShardId (index , i ), List .of ("n1" ), null , null , null ))
830
+ .toArray (SearchShardIterator []::new );
831
+
832
+ SearchRequest searchRequest = new SearchRequest ().allowPartialSearchResults (true );
833
+ searchRequest .source (new SearchSourceBuilder ());
834
+ long timeoutMills = 500 ;
835
+ searchRequest .source ().timeout (new TimeValue (timeoutMills , TimeUnit .MILLISECONDS ));
836
+ searchRequest .setMaxConcurrentShardRequests (shards .length );
837
+ final AtomicBoolean successed = new AtomicBoolean (false );
838
+ long current = System .currentTimeMillis ();
839
+
840
+ final ArraySearchPhaseResults <SearchPhaseResult > queryResult = new ArraySearchPhaseResults <>(shards .length );
841
+ AbstractSearchAsyncAction <SearchPhaseResult > action = createAction (searchRequest , queryResult , new ActionListener <>() {
842
+ @ Override
843
+ public void onResponse (SearchResponse response ) {
844
+ successed .set (true );
845
+ }
846
+
847
+ @ Override
848
+ public void onFailure (Exception e ) {
849
+ successed .set (false );
850
+ }
851
+ }, false , false , false , new AtomicLong (), new TaskResourceUsage (randomLong (), randomLong ()), true , shards );
852
+ action .run ();
853
+ long s = System .currentTimeMillis () - current ;
854
+ assertTrue (s > timeoutMills );
855
+ assertTrue (successed .get ());
856
+
857
+ }
858
+
794
859
private SearchDfsQueryThenFetchAsyncAction createSearchDfsQueryThenFetchAsyncAction (
795
860
List <SearchRequestOperationsListener > searchRequestOperationsListeners
796
861
) {
0 commit comments