#27 - ENH: Add HttpClientRequest setAttribute() getAttribute() to pass custom attributes between RequestIntercept methods

Author: rbygrave

client/src/main/java/io/avaje/http/client/DHttpClientRequest.java

@@ -52,6 +52,7 @@ class DHttpClientRequest implements HttpClientRequest, HttpClientResponse {
   private boolean suppressLogging;
   private long startAsyncNanos;
   private String label;
+  private Map<String, Object> customAttributes;
 
   DHttpClientRequest(DHttpClientContext context, Duration requestTimeout) {
     this.context = context;
@@ -83,6 +84,21 @@ public String label() {
     return label;
   }
 
+  @Override
+  public HttpClientRequest setAttribute(String key, Object value) {
+    if (customAttributes == null) {
+      customAttributes = new HashMap<>();
+    }
+    customAttributes.put(key, value);
+    return this;
+  }
+
+  @SuppressWarnings("unchecked")
+  @Override
+  public <E> E getAttribute(String key) {
+    return customAttributes == null ? null : (E) customAttributes.get(key);
+  }
+
   @Override
   public HttpClientRequest requestTimeout(Duration requestTimeout) {
     this.requestTimeout = requestTimeout;

client/src/main/java/io/avaje/http/client/HttpClientRequest.java

@@ -61,6 +61,26 @@ public interface HttpClientRequest {
    */
   String label();
 
+  /**
+   * Used to pass custom attribute between {@link RequestIntercept} methods.
+   * <p>
+   * Allows us to pass something between {@code beforeRequest} and {@code afterResponse}
+   * methods of a {@link RequestIntercept} or between multiple {@link RequestIntercept}.
+   *
+   * @param key   The unique key used to store the attribute
+   * @param value The attribute to store
+   */
+  HttpClientRequest setAttribute(String key, Object value);
+
+  /**
+   * Return a custom attribute typically set by a {@link RequestIntercept#beforeRequest(HttpClientRequest)}.
+   *
+   * @param key The key for the custom attribute
+   * @param <E> The inferred type of the attribute
+   * @return The custom attribute
+   */
+  <E> E getAttribute(String key);
+
   /**
    * Set the request timeout to use for this request. When not set the default
    * request timeout will be used.

client/src/test/java/io/avaje/http/client/RetryTest.java

@@ -33,20 +33,35 @@ void retryTest() {
     assertThat(res.body()).isEqualTo("All good at 3rd attempt");
 
     assertThat(myIntercept.responseTimeMicros).isGreaterThan(1);
+    assertThat(myIntercept.customAttributeTimeMillis).isGreaterThan(1);
+
     assertThat(myIntercept.counter).isEqualTo(1);
     assertThat(myIntercept.label).isEqualTo("http_client_hello_retry");
   }
 
   static class MyIntercept implements RequestIntercept {
     int counter;
     long responseTimeMicros;
+    long customAttributeTimeMillis;
     String label;
 
+    @Override
+    public void beforeRequest(HttpClientRequest request) {
+      final String label = request.setAttribute("MY_START_TIME", System.currentTimeMillis()).label();
+      assertThat(label).isEqualTo("http_client_hello_retry");
+    }
+
     /**
      * Not called for the retry attempts. Only called on the final success or error response.
      */
     @Override
     public void afterResponse(HttpResponse<?> response, HttpClientRequest request) {
+
+      final String does_not_exist = request.getAttribute("DOES_NOT_EXIST");
+      assertThat(does_not_exist).isNull();
+
+      long start = request.getAttribute("MY_START_TIME");
+      customAttributeTimeMillis = System.currentTimeMillis() - start;
       counter++;
       responseTimeMicros = request.responseTimeMicros();
       label = request.label();