Merge 08d07e7 into 84197e2

opentracing-contrib · Apr 23, 2019 · 1058c5c · 1058c5c
2 parents 84197e2 + 08d07e7
commit 1058c5c
Show file tree

Hide file tree

Showing 2 changed files with 85 additions and 5 deletions.
diff --git a/...ain/java/io/opentracing/contrib/elasticsearch/common/TracingHttpClientConfigCallback.java b/...ain/java/io/opentracing/contrib/elasticsearch/common/TracingHttpClientConfigCallback.java
@@ -27,37 +27,68 @@
 import org.apache.http.HttpResponseInterceptor;
 import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
 import org.elasticsearch.client.RestClientBuilder;
+import org.elasticsearch.client.RestClientBuilder.HttpClientConfigCallback;
 
 
 public class TracingHttpClientConfigCallback implements RestClientBuilder.HttpClientConfigCallback {
 
   private final Tracer tracer;
   private final Function<HttpRequest, String> spanNameProvider;
+  private final HttpClientConfigCallback callback;
 
   public TracingHttpClientConfigCallback(Tracer tracer,
-      Function<HttpRequest, String> spanNameProvider) {
+      Function<HttpRequest, String> spanNameProvider,
+      HttpClientConfigCallback callback) {
     this.tracer = tracer;
     this.spanNameProvider = spanNameProvider;
+    this.callback = callback;
+  }
+
+  public TracingHttpClientConfigCallback(Tracer tracer,
+      Function<HttpRequest, String> spanNameProvider) {
+    this(tracer, spanNameProvider, null);
   }
 
   /**
    * Default span name provider (ClientSpanNameProvider.REQUEST_METHOD_NAME) is used
    */
   public TracingHttpClientConfigCallback(Tracer tracer) {
-    this(tracer, ClientSpanNameProvider.REQUEST_METHOD_NAME);
+    this(tracer, ClientSpanNameProvider.REQUEST_METHOD_NAME, null);
   }
 
   /**
-   * GlobalTracer is used to get tracer Default span name provider (ClientSpanNameProvider.REQUEST_METHOD_NAME)
+   * Default span name provider (ClientSpanNameProvider.REQUEST_METHOD_NAME) is used
+   */
+  public TracingHttpClientConfigCallback(Tracer tracer, HttpClientConfigCallback callback) {
+    this(tracer, ClientSpanNameProvider.REQUEST_METHOD_NAME, callback);
+  }
+
+  /**
+   * GlobalTracer is used to get tracer. Default span name provider (ClientSpanNameProvider.REQUEST_METHOD_NAME)
    * is used
    */
   public TracingHttpClientConfigCallback() {
-    this(GlobalTracer.get(), ClientSpanNameProvider.REQUEST_METHOD_NAME);
+    this(GlobalTracer.get(), ClientSpanNameProvider.REQUEST_METHOD_NAME, null);
+  }
+
+  /**
+   * GlobalTracer is used to get tracer. Default span name provider (ClientSpanNameProvider.REQUEST_METHOD_NAME)
+   * is used
+   */
+  public TracingHttpClientConfigCallback(HttpClientConfigCallback callback) {
+    this(GlobalTracer.get(), ClientSpanNameProvider.REQUEST_METHOD_NAME, callback);
   }
 
   @Override
   public HttpAsyncClientBuilder customizeHttpClient(
-      final HttpAsyncClientBuilder httpClientBuilder) {
+      final HttpAsyncClientBuilder httpAsyncClientBuilder) {
+
+    HttpAsyncClientBuilder httpClientBuilder;
+    if (callback != null) {
+      httpClientBuilder = callback.customizeHttpClient(httpAsyncClientBuilder);
+    } else {
+      httpClientBuilder = httpAsyncClientBuilder;
+    }
 
     httpClientBuilder.addInterceptorFirst((HttpRequestInterceptor) (request, context) -> {
       SpanBuilder spanBuilder = tracer.buildSpan(spanNameProvider.apply(request))

diff --git a/...lasticsearch6-client/src/test/java/io/opentracing/contrib/elasticsearch6/TracingTest.java b/...lasticsearch6-client/src/test/java/io/opentracing/contrib/elasticsearch6/TracingTest.java
@@ -42,6 +42,7 @@
 import org.elasticsearch.client.Response;
 import org.elasticsearch.client.ResponseListener;
 import org.elasticsearch.client.RestClient;
+import org.elasticsearch.client.RestClientBuilder.HttpClientConfigCallback;
 import org.elasticsearch.client.transport.TransportClient;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.transport.TransportAddress;
@@ -140,6 +141,54 @@ public void onFailure(Exception exception) {
     assertNull(mockTracer.activeSpan());
   }
 
+  @Test
+  public void restClientWithCallback() throws Exception {
+    RestClient restClient = RestClient.builder(
+        new HttpHost("localhost", HTTP_PORT, "http"))
+        .setHttpClientConfigCallback(new TracingHttpClientConfigCallback(mockTracer,
+            (HttpClientConfigCallback) httpClientBuilder -> httpClientBuilder))
+        .build();
+
+    HttpEntity entity = new NStringEntity(
+        "{\n" +
+            "    \"user\" : \"kimchy\",\n" +
+            "    \"post_date\" : \"2009-11-15T14:12:12\",\n" +
+            "    \"message\" : \"trying out Elasticsearch\"\n" +
+            "}", ContentType.APPLICATION_JSON);
+
+    Request request = new Request("PUT", "/twitter/tweet/1");
+    request.setEntity(entity);
+
+    Response indexResponse = restClient.performRequest(request);
+
+    assertNotNull(indexResponse);
+
+    Request request2 = new Request("PUT", "/twitter/tweet/2");
+    request2.setEntity(entity);
+
+    final CountDownLatch latch = new CountDownLatch(1);
+    restClient
+        .performRequestAsync(request2, new ResponseListener() {
+          @Override
+          public void onSuccess(Response response) {
+            latch.countDown();
+          }
+
+          @Override
+          public void onFailure(Exception exception) {
+            latch.countDown();
+          }
+        });
+
+    latch.await(30, TimeUnit.SECONDS);
+    restClient.close();
+
+    List<MockSpan> finishedSpans = mockTracer.finishedSpans();
+    assertEquals(2, finishedSpans.size());
+    checkSpans(finishedSpans, "PUT");
+    assertNull(mockTracer.activeSpan());
+  }
+
   @Test
   public void transportClient() throws Exception {