From b736ca310845b58add81148ae8cc96c071b59f80 Mon Sep 17 00:00:00 2001
From: Nikolay Martynov <mar.kolya@gmail.com>
Date: Thu, 16 May 2019 17:24:55 -0400
Subject: [PATCH] Handle null callback in HttpAsyncApache instrumentation

---
 .../ApacheHttpAsyncClientInstrumentation.java | 31 +++++++++---
 ...acheHttpAsyncClientNullCallbackTest.groovy | 50 +++++++++++++++++++
 2 files changed, 75 insertions(+), 6 deletions(-)
 create mode 100644 dd-java-agent/instrumentation/apache-httpasyncclient-4/src/test/groovy/ApacheHttpAsyncClientNullCallbackTest.groovy

diff --git a/dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpAsyncClientInstrumentation.java b/dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpAsyncClientInstrumentation.java
index 78682b5247..a612c34c38 100644
--- a/dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpAsyncClientInstrumentation.java
+++ b/dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpAsyncClientInstrumentation.java
@@ -180,6 +180,7 @@ public class ApacheHttpAsyncClientInstrumentation extends Instrumenter.Default {
       }
       this.clientSpan = clientSpan;
       this.context = context;
+      // Note: this can be null in real life, so we have to handle this carefully
       this.delegate = delegate;
     }
 
@@ -190,11 +191,11 @@ public class ApacheHttpAsyncClientInstrumentation extends Instrumenter.Default {
       clientSpan.finish(); // Finish span before calling delegate
 
       if (parentContinuation == null) {
-        delegate.completed(result);
+        completeDelegate(result);
       } else {
         try (final TraceScope scope = parentContinuation.activate()) {
           scope.setAsyncPropagation(true);
-          delegate.completed(result);
+          completeDelegate(result);
         }
       }
     }
@@ -207,11 +208,11 @@ public class ApacheHttpAsyncClientInstrumentation extends Instrumenter.Default {
       clientSpan.finish(); // Finish span before calling delegate
 
       if (parentContinuation == null) {
-        delegate.failed(ex);
+        failDelegate(ex);
       } else {
         try (final TraceScope scope = parentContinuation.activate()) {
           scope.setAsyncPropagation(true);
-          delegate.failed(ex);
+          failDelegate(ex);
         }
       }
     }
@@ -223,14 +224,32 @@ public class ApacheHttpAsyncClientInstrumentation extends Instrumenter.Default {
       clientSpan.finish(); // Finish span before calling delegate
 
       if (parentContinuation == null) {
-        delegate.cancelled();
+        cancelDelegate();
       } else {
         try (final TraceScope scope = parentContinuation.activate()) {
           scope.setAsyncPropagation(true);
-          delegate.cancelled();
+          cancelDelegate();
         }
       }
     }
+
+    private void completeDelegate(final T result) {
+      if (delegate != null) {
+        delegate.completed(result);
+      }
+    }
+
+    private void failDelegate(final Exception ex) {
+      if (delegate != null) {
+        delegate.failed(ex);
+      }
+    }
+
+    private void cancelDelegate() {
+      if (delegate != null) {
+        delegate.cancelled();
+      }
+    }
   }
 
   public static class HttpHeadersInjectAdapter implements TextMap {
diff --git a/dd-java-agent/instrumentation/apache-httpasyncclient-4/src/test/groovy/ApacheHttpAsyncClientNullCallbackTest.groovy b/dd-java-agent/instrumentation/apache-httpasyncclient-4/src/test/groovy/ApacheHttpAsyncClientNullCallbackTest.groovy
new file mode 100644
index 0000000000..2015e4956e
--- /dev/null
+++ b/dd-java-agent/instrumentation/apache-httpasyncclient-4/src/test/groovy/ApacheHttpAsyncClientNullCallbackTest.groovy
@@ -0,0 +1,50 @@
+import datadog.trace.agent.test.base.HttpClientTest
+import datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator
+import org.apache.http.client.methods.HttpGet
+import org.apache.http.impl.nio.client.HttpAsyncClients
+import org.apache.http.message.BasicHeader
+import spock.lang.AutoCleanup
+import spock.lang.Shared
+
+import java.util.concurrent.Future
+
+class ApacheHttpAsyncClientNullCallbackTest extends HttpClientTest<ApacheHttpAsyncClientDecorator> {
+
+  @AutoCleanup
+  @Shared
+  def client = HttpAsyncClients.createDefault()
+
+  def setupSpec() {
+    client.start()
+  }
+
+  @Override
+  int doRequest(String method, URI uri, Map<String, String> headers, Closure callback) {
+    assert method == "GET"
+
+    HttpGet request = new HttpGet(uri)
+    headers.entrySet().each {
+      request.addHeader(new BasicHeader(it.key, it.value))
+    }
+
+    // The point here is to test case when callback is null - fire-and-forget style
+    // So to make sure request is done we start request, wait for future to finish
+    // and then call callback if present.
+    Future future = client.execute(request, null)
+    future.get()
+    if (callback != null) {
+      callback()
+    }
+    return 200
+  }
+
+  @Override
+  ApacheHttpAsyncClientDecorator decorator() {
+    return ApacheHttpAsyncClientDecorator.DECORATE
+  }
+
+  @Override
+  Integer statusOnRedirectError() {
+    return 302
+  }
+}