42
42
import org .apache .hc .core5 .http .ContentType ;
43
43
import org .apache .hc .core5 .http .HttpEntity ;
44
44
import org .apache .hc .core5 .http .HttpHost ;
45
+ import org .apache .hc .core5 .http .ProtocolVersion ;
45
46
import org .apache .hc .core5 .http .io .entity .InputStreamEntity ;
46
47
import org .apache .hc .core5 .http .io .entity .StringEntity ;
47
48
import org .apache .hc .core5 .http .message .BasicClassicHttpResponse ;
49
+ import org .apache .hc .core5 .http .message .RequestLine ;
50
+ import org .apache .hc .core5 .http .message .StatusLine ;
48
51
import org .apache .hc .core5 .http .nio .AsyncPushConsumer ;
49
52
import org .apache .hc .core5 .http .nio .AsyncRequestProducer ;
50
53
import org .apache .hc .core5 .http .nio .AsyncResponseConsumer ;
57
60
import org .opensearch .Version ;
58
61
import org .opensearch .action .bulk .BackoffPolicy ;
59
62
import org .opensearch .action .search .SearchRequest ;
63
+ import org .opensearch .client .ResponseException ;
60
64
import org .opensearch .client .RestClient ;
61
65
import org .opensearch .client .http .HttpUriRequestProducer ;
62
66
import org .opensearch .client .nio .HeapBufferedAsyncResponseConsumer ;
83
87
import java .io .IOException ;
84
88
import java .io .InputStreamReader ;
85
89
import java .io .UncheckedIOException ;
90
+ import java .net .ConnectException ;
86
91
import java .net .URL ;
87
92
import java .nio .charset .StandardCharsets ;
88
93
import java .util .Queue ;
89
94
import java .util .concurrent .ExecutorService ;
90
95
import java .util .concurrent .Future ;
91
96
import java .util .concurrent .LinkedBlockingQueue ;
92
97
import java .util .concurrent .atomic .AtomicBoolean ;
98
+ import java .util .concurrent .atomic .AtomicInteger ;
93
99
import java .util .concurrent .atomic .AtomicReference ;
94
100
import java .util .function .Consumer ;
95
101
import java .util .stream .Stream ;
96
102
103
+ import org .mockito .Mockito ;
104
+
97
105
import static org .opensearch .common .unit .TimeValue .timeValueMillis ;
98
106
import static org .opensearch .common .unit .TimeValue .timeValueMinutes ;
99
107
import static org .hamcrest .Matchers .empty ;
@@ -515,7 +523,7 @@ public void testInvalidJsonThinksRemoteIsNotES() throws IOException {
515
523
Exception e = expectThrows (RuntimeException .class , () -> sourceWithMockedRemoteCall ("some_text.txt" ).start ());
516
524
assertEquals (
517
525
"Error parsing the response, remote is likely not an OpenSearch instance" ,
518
- e .getCause ().getCause ().getCause ().getMessage ()
526
+ e .getCause ().getCause ().getCause ().getCause (). getMessage ()
519
527
);
520
528
}
521
529
@@ -524,7 +532,7 @@ public void testUnexpectedJsonThinksRemoteIsNotES() throws IOException {
524
532
Exception e = expectThrows (RuntimeException .class , () -> sourceWithMockedRemoteCall ("main/2_3_3.json" ).start ());
525
533
assertEquals (
526
534
"Error parsing the response, remote is likely not an OpenSearch instance" ,
527
- e .getCause ().getCause ().getCause ().getMessage ()
535
+ e .getCause ().getCause ().getCause ().getCause (). getMessage ()
528
536
);
529
537
}
530
538
@@ -702,4 +710,105 @@ private static ClassicHttpRequest getRequest(AsyncRequestProducer requestProduce
702
710
assertThat (requestProducer , instanceOf (HttpUriRequestProducer .class ));
703
711
return ((HttpUriRequestProducer ) requestProducer ).getRequest ();
704
712
}
713
+
714
+ RemoteScrollableHitSource createRemoteSourceWithFailure (
715
+ boolean shouldMockRemoteVersion ,
716
+ Exception failure ,
717
+ AtomicInteger invocationCount
718
+ ) {
719
+ CloseableHttpAsyncClient httpClient = new CloseableHttpAsyncClient () {
720
+
721
+ @ Override
722
+ public void close () throws IOException {}
723
+
724
+ @ Override
725
+ public void close (CloseMode closeMode ) {}
726
+
727
+ @ Override
728
+ public void start () {}
729
+
730
+ @ Override
731
+ public void register (String hostname , String uriPattern , Supplier <AsyncPushConsumer > supplier ) {}
732
+
733
+ @ Override
734
+ public void initiateShutdown () {}
735
+
736
+ @ Override
737
+ public IOReactorStatus getStatus () {
738
+ return null ;
739
+ }
740
+
741
+ @ Override
742
+ protected <T > Future <T > doExecute (
743
+ HttpHost target ,
744
+ AsyncRequestProducer requestProducer ,
745
+ AsyncResponseConsumer <T > responseConsumer ,
746
+ HandlerFactory <AsyncPushConsumer > pushHandlerFactory ,
747
+ HttpContext context ,
748
+ FutureCallback <T > callback
749
+ ) {
750
+ invocationCount .getAndIncrement ();
751
+ callback .failed (failure );
752
+ return null ;
753
+ }
754
+
755
+ @ Override
756
+ public void awaitShutdown (org .apache .hc .core5 .util .TimeValue waitTime ) throws InterruptedException {}
757
+ };
758
+ return sourceWithMockedClient (shouldMockRemoteVersion , httpClient );
759
+ }
760
+
761
+ void verifyRetries (boolean shouldMockRemoteVersion , Exception failureResponse , boolean expectedToRetry ) {
762
+ retriesAllowed = 5 ;
763
+ AtomicInteger invocations = new AtomicInteger ();
764
+ invocations .set (0 );
765
+ RemoteScrollableHitSource source = createRemoteSourceWithFailure (shouldMockRemoteVersion , failureResponse , invocations );
766
+
767
+ Throwable e = expectThrows (RuntimeException .class , source ::start );
768
+ int expectedInvocations = 0 ;
769
+ if (shouldMockRemoteVersion ) {
770
+ expectedInvocations += 1 ; // first search
771
+ if (expectedToRetry ) expectedInvocations += retriesAllowed ;
772
+ } else {
773
+ expectedInvocations = 1 ; // the first should fail and not trigger any retry.
774
+ }
775
+
776
+ assertEquals (expectedInvocations , invocations .get ());
777
+
778
+ // Unwrap the some artifacts from the test
779
+ while (e .getMessage ().equals ("failed" )) {
780
+ e = e .getCause ();
781
+ }
782
+ // There is an additional wrapper for ResponseException.
783
+ if (failureResponse instanceof ResponseException ) {
784
+ e = e .getCause ();
785
+ }
786
+
787
+ assertSame (failureResponse , e );
788
+ }
789
+
790
+ ResponseException withResponseCode (int statusCode , String errorMsg ) throws IOException {
791
+ org .opensearch .client .Response mockResponse = Mockito .mock (org .opensearch .client .Response .class );
792
+ Mockito .when (mockResponse .getEntity ()).thenReturn (new StringEntity (errorMsg , ContentType .TEXT_PLAIN ));
793
+ Mockito .when (mockResponse .getStatusLine ()).thenReturn (new StatusLine (new BasicClassicHttpResponse (statusCode , errorMsg )));
794
+ Mockito .when (mockResponse .getRequestLine ()).thenReturn (new RequestLine ("GET" , "/" , new ProtocolVersion ("https" , 1 , 1 )));
795
+ return new ResponseException (mockResponse );
796
+ }
797
+
798
+ public void testRetryOnCallFailure () throws Exception {
799
+ // First call succeeds. Search calls failing with 5xxs and 429s should be retried but not 400s.
800
+ verifyRetries (true , withResponseCode (500 , "Internal Server Error" ), true );
801
+ verifyRetries (true , withResponseCode (429 , "Too many requests" ), true );
802
+ verifyRetries (true , withResponseCode (400 , "Client Error" ), false );
803
+
804
+ // First call succeeds. Search call failed with exceptions other than ResponseException
805
+ verifyRetries (true , new ConnectException ("blah" ), true ); // should retry connect exceptions.
806
+ verifyRetries (true , new RuntimeException ("foobar" ), false );
807
+
808
+ // First call(remote version lookup) failed and no retries expected
809
+ verifyRetries (false , withResponseCode (500 , "Internal Server Error" ), false );
810
+ verifyRetries (false , withResponseCode (429 , "Too many requests" ), false );
811
+ verifyRetries (false , withResponseCode (400 , "Client Error" ), false );
812
+ verifyRetries (false , new ConnectException ("blah" ), false );
813
+ }
705
814
}
0 commit comments