Receiver and SenderContext

micrometer-core/build.gradle

@@ -64,6 +64,7 @@ dependencies {
     optionalApi 'org.jetbrains.kotlinx:kotlinx-coroutines-core'
 
     testImplementation 'io.projectreactor:reactor-test'
+    testImplementation project(":micrometer-observation-test")
 
     // dependency injection tests
     testImplementation 'javax.inject:javax.inject'

micrometer-core/src/main/java/io/micrometer/core/instrument/binder/okhttp3/DefaultOkHttpObservationConvention.java

@@ -75,13 +75,13 @@ public DefaultOkHttpObservationConvention(String metricName) {
 
     @Override
     public KeyValues getLowCardinalityKeyValues(OkHttpContext context) {
-        OkHttpMetricsEventListener.CallState state = context.getState();
+        OkHttpObservationInterceptor.CallState state = context.getState();
         Request request = state.request;
         boolean requestAvailable = request != null;
         Function<Request, String> urlMapper = context.getUrlMapper();
-        Iterable<Tag> extraTags = context.getExtraTags();
-        Iterable<BiFunction<Request, Response, Tag>> contextSpecificTags = context.getContextSpecificTags();
-        Iterable<Tag> unknownRequestTags = context.getUnknownRequestTags();
+        Iterable<KeyValue> extraTags = context.getExtraTags();
+        Iterable<BiFunction<Request, Response, KeyValue>> contextSpecificTags = context.getContextSpecificTags();
+        Iterable<KeyValue> unknownRequestTags = context.getUnknownRequestTags();
         boolean includeHostTag = context.isIncludeHostTag();
         // TODO: Tags to key values and back - maybe we can improve this?
         KeyValues keyValues = KeyValues.of(
@@ -90,20 +90,19 @@ public KeyValues getLowCardinalityKeyValues(OkHttpContext context) {
                 OkHttpDocumentedObservation.OkHttpLegacyLowCardinalityTags.URI.of(getUriTag(urlMapper, state, request)),
                 OkHttpDocumentedObservation.OkHttpLegacyLowCardinalityTags.STATUS
                         .of(getStatusMessage(state.response, state.exception)))
-                .and(tagsToKeyValues(stream(extraTags.spliterator(), false)))
+                .and(extraTags)
                 .and(stream(contextSpecificTags.spliterator(), false)
                         .map(contextTag -> contextTag.apply(request, state.response))
                         .map(tag -> KeyValue.of(tag.getKey(), tag.getValue())).collect(toList()))
-                .and(getRequestTags(request, tagsToKeyValues(stream(unknownRequestTags.spliterator(), false))))
-                .and(generateTagsForRoute(request));
+                .and(getRequestTags(request, unknownRequestTags)).and(generateTagsForRoute(request));
         if (includeHostTag) {
             keyValues = KeyValues.of(keyValues).and(OkHttpDocumentedObservation.OkHttpLegacyLowCardinalityTags.HOST
                     .of(requestAvailable ? request.url().host() : TAG_VALUE_UNKNOWN));
         }
         return keyValues;
     }
 
-    private String getUriTag(Function<Request, String> urlMapper, OkHttpMetricsEventListener.CallState state,
+    private String getUriTag(Function<Request, String> urlMapper, OkHttpObservationInterceptor.CallState state,
             @Nullable Request request) {
         if (request == null) {
             return TAG_VALUE_UNKNOWN;
@@ -133,11 +132,18 @@ private Iterable<KeyValue> getRequestTags(@Nullable Request request, Iterable<Ke
             if (requestTag != null) {
                 return tagsToKeyValues(requestTag.stream());
             }
+            KeyValues keyValues = request.tag(KeyValues.class);
+            if (keyValues != null) {
+                return keyValues;
+            }
         }
         Object requestTag = request.tag();
         if (requestTag instanceof Tags) {
             return tagsToKeyValues(((Tags) requestTag).stream());
         }
+        else if (requestTag instanceof KeyValues) {
+            return (Iterable<KeyValue>) requestTag;
+        }
         return KeyValues.empty();
     }
 

micrometer-core/src/main/java/io/micrometer/core/instrument/binder/okhttp3/OkHttpContext.java

@@ -15,62 +15,72 @@
  */
 package io.micrometer.core.instrument.binder.okhttp3;
 
-import io.micrometer.core.instrument.Tag;
-import io.micrometer.observation.transport.http.context.HttpClientContext;
+import io.micrometer.common.KeyValue;
+import io.micrometer.observation.transport.SenderContext;
+import io.micrometer.observation.transport.Kind;
+import io.micrometer.observation.transport.RequestReplySenderContext;
 import okhttp3.Request;
 import okhttp3.Response;
 
 import java.util.function.BiFunction;
 import java.util.function.Function;
 
 /**
- * A {@link HttpClientContext} for OkHttp3.
+ * A {@link SenderContext} for OkHttp3.
  *
  * @author Marcin Grzejszczak
  * @since 1.10.0
  */
-public class OkHttpContext extends HttpClientContext {
-
-    private final OkHttpMetricsEventListener.CallState state;
+public class OkHttpContext extends RequestReplySenderContext<Request.Builder, Response> {
 
     private final Function<Request, String> urlMapper;
 
-    private final Iterable<Tag> extraTags;
+    private final Iterable<KeyValue> extraTags;
 
-    private final Iterable<BiFunction<Request, Response, Tag>> contextSpecificTags;
+    private final Iterable<BiFunction<Request, Response, KeyValue>> contextSpecificTags;
 
-    private final Iterable<Tag> unknownRequestTags;
+    private final Iterable<KeyValue> unknownRequestTags;
 
     private final boolean includeHostTag;
 
-    public OkHttpContext(OkHttpMetricsEventListener.CallState state, Function<Request, String> urlMapper,
-            Iterable<Tag> extraTags, Iterable<BiFunction<Request, Response, Tag>> contextSpecificTags,
-            Iterable<Tag> unknownRequestTags, boolean includeHostTag) {
-        this.state = state;
+    private OkHttpObservationInterceptor.CallState state;
+
+    public OkHttpContext(Function<Request, String> urlMapper, Iterable<KeyValue> extraTags,
+            Iterable<BiFunction<Request, Response, KeyValue>> contextSpecificTags,
+            Iterable<KeyValue> unknownRequestTags, boolean includeHostTag) {
+        super((carrier, key, value) -> {
+            if (carrier != null) {
+                carrier.header(key, value);
+            }
+        }, Kind.CLIENT);
         this.urlMapper = urlMapper;
         this.extraTags = extraTags;
         this.contextSpecificTags = contextSpecificTags;
         this.unknownRequestTags = unknownRequestTags;
         this.includeHostTag = includeHostTag;
     }
 
-    public OkHttpMetricsEventListener.CallState getState() {
+    public void setState(OkHttpObservationInterceptor.CallState state) {
+        this.state = state;
+    }
+
+    public OkHttpObservationInterceptor.CallState getState() {
         return state;
     }
 
     public Function<Request, String> getUrlMapper() {
         return urlMapper;
     }
 
-    public Iterable<Tag> getExtraTags() {
+    public Iterable<KeyValue> getExtraTags() {
         return extraTags;
     }
 
-    public Iterable<BiFunction<Request, Response, Tag>> getContextSpecificTags() {
+    public Iterable<BiFunction<Request, Response, KeyValue>> getContextSpecificTags() {
         return contextSpecificTags;
     }
 
-    public Iterable<Tag> getUnknownRequestTags() {
+    public Iterable<KeyValue> getUnknownRequestTags() {
         return unknownRequestTags;
     }