Implement a mechanism for adding ad-hoc listeners to OKHTTP requests (#6660)

nikita-tkachenko-datadog · web-flow · commit da71297afd11 · 2024-02-14T21:45:09.000+01:00
diff --git a/communication/build.gradle b/communication/build.gradle
@@ -20,6 +20,7 @@ dependencies {
   testImplementation deps.bytebuddy
   testImplementation group: 'org.msgpack', name: 'msgpack-core', version: '0.8.20'
   testImplementation group: 'org.msgpack', name: 'jackson-dataformat-msgpack', version: '0.8.20'
+  testImplementation group: 'com.squareup.okhttp3', name: 'mockwebserver', version: versions.okhttp_legacy
 }
 
 ext {
@@ -33,6 +34,7 @@ ext {
     'datadog.communication.http.OkHttpUtils',
     'datadog.communication.http.OkHttpUtils.1',
     'datadog.communication.http.OkHttpUtils.ByteBufferRequestBody',
+    'datadog.communication.http.OkHttpUtils.CustomListener',
     'datadog.communication.http.OkHttpUtils.GZipByteBufferRequestBody',
     'datadog.communication.http.OkHttpUtils.GZipRequestBodyDecorator',
     'datadog.communication.http.OkHttpUtils.JsonRequestBody',
diff --git a/communication/src/main/java/datadog/communication/http/OkHttpUtils.java b/communication/src/main/java/datadog/communication/http/OkHttpUtils.java
@@ -23,6 +23,7 @@
 import okhttp3.ConnectionSpec;
 import okhttp3.Credentials;
 import okhttp3.Dispatcher;
+import okhttp3.EventListener;
 import okhttp3.HttpUrl;
 import okhttp3.MediaType;
 import okhttp3.OkHttpClient;
@@ -100,6 +101,8 @@ public static OkHttpClient buildHttpClient(
         timeoutMillis);
   }
 
+  public abstract static class CustomListener extends EventListener {}
+
   private static OkHttpClient buildHttpClient(
       final String unixDomainSocketPath,
       final String namedPipe,
@@ -115,6 +118,12 @@ private static OkHttpClient buildHttpClient(
     final OkHttpClient.Builder builder = new OkHttpClient.Builder();
 
     builder
+        .eventListenerFactory(
+            call -> {
+              Request request = call.request();
+              CustomListener listener = request.tag(CustomListener.class);
+              return listener != null ? listener : EventListener.NONE;
+            })
         .connectTimeout(timeoutMillis, MILLISECONDS)
         .writeTimeout(timeoutMillis, MILLISECONDS)
         .readTimeout(timeoutMillis, MILLISECONDS)
diff --git a/communication/src/test/groovy/datadog/communication/http/OkHttpUtilsTest.groovy b/communication/src/test/groovy/datadog/communication/http/OkHttpUtilsTest.groovy
@@ -0,0 +1,46 @@
+package datadog.communication.http
+
+import okhttp3.Call
+import okhttp3.Request
+import okhttp3.mockwebserver.MockResponse
+import okhttp3.mockwebserver.MockWebServer
+import spock.lang.Shared
+import spock.lang.Specification
+
+class OkHttpUtilsTest extends Specification {
+
+  @Shared
+  MockWebServer server = new MockWebServer()
+
+  def cleanupSpec() {
+    server.shutdown()
+  }
+
+  def "It is possible to register a custom listener for HTTP requests"() {
+    setup:
+    def url = server.url("/")
+    def client = OkHttpUtils.buildHttpClient(url, 1000)
+    def listener = new TestListener()
+
+    server.enqueue(new MockResponse())
+
+    when:
+    def request = new Request.Builder()
+      .url(url)
+      .get()
+      .tag(OkHttpUtils.CustomListener, listener)
+      .build()
+    client.newCall(request).execute()
+
+    then:
+    listener.notified
+  }
+
+  private static final class TestListener extends OkHttpUtils.CustomListener {
+    private boolean notified
+
+    void callStart(Call call) {
+      notified = true
+    }
+  }
+}
