Add Stripe client telemetry to request headers

build.gradle

@@ -47,6 +47,7 @@ dependencies {
     testCompile group: 'junit', name: 'junit', version:'4.12'
     testCompile group: 'org.mockito', name: 'mockito-core', version:'2.22.0'
     testRuntime group: 'org.slf4j', name: 'slf4j-api', version: '1.7.25'
+    testImplementation 'com.squareup.okhttp3:mockwebserver:3.12.1'
 }
 
 jar {

src/main/java/com/stripe/Stripe.java

@@ -19,6 +19,7 @@ public abstract class Stripe {
   public static volatile String apiVersion;
   public static volatile String clientId;
   public static volatile String partnerId;
+  public static volatile boolean enableTelemetry = false;
 
   // Note that URLConnection reserves the value of 0 to mean "infinite
   // timeout", so we use -1 here to represent an unset value which should

src/main/java/com/stripe/net/ClientTelemetryPayload.java

@@ -0,0 +1,8 @@
+package com.stripe.net;
+
+import com.google.gson.annotations.SerializedName;
+
+public class ClientTelemetryPayload {
+  @SerializedName("last_request_metrics")
+  public RequestMetrics lastRequestMetrics;
+}

src/main/java/com/stripe/net/LiveStripeResponseGetter.java

@@ -43,13 +43,15 @@
 import java.util.ListIterator;
 import java.util.Map;
 import java.util.Scanner;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import javax.net.ssl.HttpsURLConnection;
 import javax.net.ssl.SSLSocketFactory;
 
 import lombok.Cleanup;
 
 public class LiveStripeResponseGetter implements StripeResponseGetter {
   private static final String DNS_CACHE_TTL_PROPERTY_NAME = "networkaddress.cache.ttl";
+  private static final int MAX_REQUEST_METRICS_BUFFER_SIZE = 100;
 
   private static final class Parameter {
     public final String key;
@@ -149,6 +151,13 @@ static Map<String, String> getHeaders(RequestOptions options) {
     if (options.getStripeAccount() != null) {
       headers.put("Stripe-Account", options.getStripeAccount());
     }
+
+    RequestMetrics lastRequestMetrics = prevRequestMetrics.poll();
+    if (Stripe.enableTelemetry && lastRequestMetrics != null) {
+      headers.put("X-Stripe-Client-Telemetry",
+          ApiResource.GSON.toJson(lastRequestMetrics.payload()));
+    }
+
     return headers;
   }
 
@@ -385,6 +394,9 @@ private static class StripeError {
     String charge;
   }
 
+  private static ConcurrentLinkedQueue<RequestMetrics> prevRequestMetrics =
+      new ConcurrentLinkedQueue<RequestMetrics>();
+
   // represents OAuth API errors returned as JSON
   // handleOAuthError uses this class to raise the appropriate OAuthException
   private static class StripeOAuthError {
@@ -523,8 +535,12 @@ private static <T> T staticRequest(
       ApiResource.RequestMethod method, String url, Map<String, Object> params,
       Class<T> clazz, ApiResource.RequestType type, RequestOptions options)
       throws StripeException {
+    long requestStart = System.currentTimeMillis();
+
     StripeResponse response = rawRequest(method, url, params, type, options);
 
+    long requestDurationMS = System.currentTimeMillis() - requestStart;
+
     int responseCode = response.code();
     String responseBody = response.body();
     String requestId = response.requestId();
@@ -544,6 +560,11 @@ private static <T> T staticRequest(
       StripeObject obj = (StripeObject)resource;
       obj.setLastResponse(response);
     }
+
+    if (Stripe.enableTelemetry && prevRequestMetrics.size() < MAX_REQUEST_METRICS_BUFFER_SIZE) {
+      prevRequestMetrics.add(new RequestMetrics(requestId, requestDurationMS));
+    }
+
     return resource;
   }
 

src/main/java/com/stripe/net/RequestMetrics.java

@@ -0,0 +1,25 @@
+package com.stripe.net;
+
+import com.google.gson.annotations.SerializedName;
+
+public class RequestMetrics {
+  @SerializedName("request_id")
+  public String requestId;
+
+  @SerializedName("request_duration_ms")
+  public long requestDurationMs;
+
+  public RequestMetrics(String requestId, long requestDurationMS) {
+    this.requestId = requestId;
+    this.requestDurationMs = requestDurationMS;
+  }
+
+  /**
+   * Constructs the JSON payload to be sent in the X-Stripe-Client-Telemetry header.
+   */
+  public ClientTelemetryPayload payload() {
+    ClientTelemetryPayload p = new ClientTelemetryPayload();
+    p.lastRequestMetrics = this;
+    return p;
+  }
+}

src/test/java/com/stripe/functional/TelemetryTest.java

@@ -0,0 +1,152 @@
+package com.stripe.functional;
+
+import static junit.framework.TestCase.assertEquals;
+import static junit.framework.TestCase.assertNull;
+import static junit.framework.TestCase.assertTrue;
+
+import com.stripe.BaseStripeTest;
+import com.stripe.Stripe;
+import com.stripe.exception.StripeException;
+import com.stripe.model.Balance;
+import com.stripe.net.ApiResource;
+import com.stripe.net.ClientTelemetryPayload;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+
+import okhttp3.mockwebserver.MockResponse;
+import okhttp3.mockwebserver.MockWebServer;
+import okhttp3.mockwebserver.RecordedRequest;
+import org.junit.Test;
+
+
+public class TelemetryTest extends BaseStripeTest {
+  @Test
+  public void testTelemetryEnabled() throws StripeException, IOException, InterruptedException {
+    MockWebServer server = new MockWebServer();
+    server.enqueue(new MockResponse().setBody("{}").addHeader("Request-Id", "req_1")
+        .setBodyDelay(30, TimeUnit.MILLISECONDS));
+    server.enqueue(new MockResponse().setBody("{}").addHeader("Request-Id", "req_2")
+        .setBodyDelay(120, TimeUnit.MILLISECONDS));
+    server.enqueue(new MockResponse().setBody("{}").addHeader("Request-Id", "req_3"));
+    server.start();
+
+    Stripe.overrideApiBase(server.url("").toString());
+    Stripe.enableTelemetry = true;
+
+    Balance b1 = Balance.retrieve();
+    RecordedRequest request1 = server.takeRequest();
+    assertNull(request1.getHeader("X-Stripe-Client-Telemetry"));
+
+    Balance b2 = Balance.retrieve();
+    RecordedRequest request2 = server.takeRequest();
+    String telemetry1 = request2.getHeader("X-Stripe-Client-Telemetry");
+    ClientTelemetryPayload payload1 = ApiResource.GSON.fromJson(
+        telemetry1, ClientTelemetryPayload.class);
+    assertEquals(payload1.lastRequestMetrics.requestId, "req_1");
+    assertTrue(payload1.lastRequestMetrics.requestDurationMs > 30);
+
+    Balance b3 = Balance.retrieve();
+    RecordedRequest request3 = server.takeRequest();
+    String telemetry2 = request3.getHeader("X-Stripe-Client-Telemetry");
+    ClientTelemetryPayload payload2 = ApiResource.GSON.fromJson(
+        telemetry2, ClientTelemetryPayload.class);
+    assertEquals(payload2.lastRequestMetrics.requestId, "req_2");
+    assertTrue(payload2.lastRequestMetrics.requestDurationMs > 120);
+
+    server.shutdown();
+  }
+
+  @Test
+  public void testTelemetryDisabled() throws StripeException, IOException, InterruptedException {
+    MockWebServer server = new MockWebServer();
+    server.enqueue(new MockResponse().setBody("{}").addHeader("Request-Id", "req_1"));
+    server.enqueue(new MockResponse().setBody("{}").addHeader("Request-Id", "req_2"));
+    server.enqueue(new MockResponse().setBody("{}").addHeader("Request-Id", "req_3"));
+    server.start();
+
+    Stripe.overrideApiBase(server.url("").toString());
+    Stripe.enableTelemetry = false;
+
+    Balance b1 = Balance.retrieve();
+    RecordedRequest request1 = server.takeRequest();
+    assertNull(request1.getHeader("X-Stripe-Client-Telemetry"));
+
+    Balance b2 = Balance.retrieve();
+    RecordedRequest request2 = server.takeRequest();
+    assertNull(request2.getHeader("X-Stripe-Client-Telemetry"));
+
+    server.shutdown();
+  }
+
+  @Test
+  public void testTelemetryWorksWithConcurrentRequests() throws IOException, InterruptedException {
+    MockWebServer server = new MockWebServer();
+
+    for (int i = 0; i < 20; i++) {
+      server.enqueue(new MockResponse().setBody("{}")
+          .addHeader("Request-Id", "req_" + i)
+          .setBodyDelay(50, TimeUnit.MILLISECONDS));
+    }
+    server.start();
+
+    Stripe.overrideApiBase(server.url("").toString());
+    Stripe.enableTelemetry = true;
+
+    Runnable work = new Runnable() {
+      @Override
+      public void run() {
+        try {
+          Balance.retrieve();
+        } catch (StripeException e) {
+          assertNull(e);
+        }
+      }
+    };
+
+    // the first 10 requests will not contain telemetry
+    ArrayList<Thread> threads = new ArrayList<>();
+    for (int i = 0; i < 10; i++) {
+      Thread t = new Thread(work);
+      threads.add(t);
+      t.start();
+    }
+    for (int i = 0; i < 10; i++) {
+      threads.get(i).join();
+    }
+    threads.clear();
+
+    // the following 10 requests will contain telemetry
+    for (int i = 0; i < 10; i++) {
+      Thread t = new Thread(work);
+      threads.add(t);
+      t.start();
+    }
+    for (int i = 0; i < 10; i++) {
+      threads.get(i).join();
+    }
+
+    Set<String> seenRequestIds = new HashSet<>();
+
+    for (int i = 0; i < 10; i++) {
+      RecordedRequest request = server.takeRequest();
+      assertNull(request.getHeader("X-Stripe-Client-Telemetry"));
+    }
+
+    for (int i = 0; i < 10; i++) {
+      RecordedRequest request = server.takeRequest();
+      String telemetry = request.getHeader("X-Stripe-Client-Telemetry");
+      ClientTelemetryPayload payload = ApiResource.GSON.fromJson(
+          telemetry, ClientTelemetryPayload.class);
+      seenRequestIds.add(payload.lastRequestMetrics.requestId);
+    }
+
+    // check that each telemetry payload corresponds to a unique request id
+    assertEquals(10, seenRequestIds.size());
+
+    server.shutdown();
+  }
+}