Maintenance: add default retry rule for Armeria client (#375)

Signed-off-by: qchea <[email protected]>
opensearch-project · Oct 6, 2021 · b3d1ff9 · b3d1ff9
1 parent ba0d9f1
commit b3d1ff9
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/...core/src/integrationTest/java/com/amazon/dataprepper/integration/EndToEndRawSpanTest.java b/...core/src/integrationTest/java/com/amazon/dataprepper/integration/EndToEndRawSpanTest.java
@@ -18,6 +18,8 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.protobuf.ByteString;
 import com.linecorp.armeria.client.Clients;
+import com.linecorp.armeria.client.retry.RetryRule;
+import com.linecorp.armeria.client.retry.RetryingClient;
 import com.linecorp.armeria.internal.shaded.bouncycastle.util.encoders.Hex;
 import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest;
 import io.opentelemetry.proto.collector.trace.v1.TraceServiceGrpc;
@@ -152,8 +154,10 @@ private void refreshIndices(final RestHighLevelClient restHighLevelClient) throw
     }
 
     private void sendExportTraceServiceRequestToSource(final int port, final ExportTraceServiceRequest request) {
-        Clients.newClient(String.format("gproto+http://127.0.0.1:%d/", port),
-                TraceServiceGrpc.TraceServiceBlockingStub.class).export(request);
+        TraceServiceGrpc.TraceServiceBlockingStub client = Clients.builder(String.format("gproto+http://127.0.0.1:%d/", port))
+                .decorator(RetryingClient.newDecorator(RetryRule.failsafe()))
+                .build(TraceServiceGrpc.TraceServiceBlockingStub.class);
+        client.export(request);
     }
 
     private List<Map<String, Object>> getSourcesFromSearchHits(final SearchHits searchHits) {

diff --git a/...e/src/integrationTest/java/com/amazon/dataprepper/integration/EndToEndServiceMapTest.java b/...e/src/integrationTest/java/com/amazon/dataprepper/integration/EndToEndServiceMapTest.java
@@ -27,6 +27,8 @@
 import java.util.stream.Stream;
 
 import com.linecorp.armeria.client.Clients;
+import com.linecorp.armeria.client.retry.RetryRule;
+import com.linecorp.armeria.client.retry.RetryingClient;
 import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest;
 import io.opentelemetry.proto.collector.trace.v1.TraceServiceGrpc;
 import io.opentelemetry.proto.common.v1.AnyValue;
@@ -132,8 +134,10 @@ private void refreshIndices(final RestHighLevelClient restHighLevelClient) throw
     }
 
     private void sendExportTraceServiceRequestToSource(final int port, final ExportTraceServiceRequest request) {
-        Clients.newClient(String.format("gproto+http://127.0.0.1:%d/", port),
-                TraceServiceGrpc.TraceServiceBlockingStub.class).export(request);
+        TraceServiceGrpc.TraceServiceBlockingStub client = Clients.builder(String.format("gproto+http://127.0.0.1:%d/", port))
+                .decorator(RetryingClient.newDecorator(RetryRule.failsafe()))
+                .build(TraceServiceGrpc.TraceServiceBlockingStub.class);
+        client.export(request);
     }
 
     private List<Map<String, Object>> getSourcesFromIndex(final RestHighLevelClient restHighLevelClient, final String index) throws IOException {