From 585cc559215b1fb4c0e8d91e57d0d9b8f13545d8 Mon Sep 17 00:00:00 2001
From: Mateusz Rzeszutek <mrzeszutek@splunk.com>
Date: Fri, 27 Aug 2021 14:34:47 +0200
Subject: [PATCH] Add request param to EndTimeExtractor (#3947)

---
 .../instrumentation/api/instrumenter/EndTimeExtractor.java  | 4 ++--
 .../instrumentation/api/instrumenter/Instrumenter.java      | 4 ++--
 .../api/instrumenter/InstrumenterBuilder.java               | 5 +++--
 .../instrumentation/api/instrumenter/InstrumenterTest.java  | 6 +++---
 .../javaagent/instrumentation/jms/JmsSingletons.java        | 3 ++-
 5 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/EndTimeExtractor.java b/instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/EndTimeExtractor.java
index 234b00e00c..79f0d8b64f 100644
--- a/instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/EndTimeExtractor.java
+++ b/instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/EndTimeExtractor.java
@@ -14,8 +14,8 @@ import org.checkerframework.checker.nullness.qual.Nullable;
  * must be comparable.
  */
 @FunctionalInterface
-public interface EndTimeExtractor<RESPONSE> {
+public interface EndTimeExtractor<REQUEST, RESPONSE> {
 
   /** Returns the timestamp marking the end of the response processing. */
-  Instant extract(@Nullable RESPONSE response);
+  Instant extract(REQUEST request, @Nullable RESPONSE response);
 }
diff --git a/instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/Instrumenter.java b/instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/Instrumenter.java
index bb37660010..105c934b3f 100644
--- a/instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/Instrumenter.java
+++ b/instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/Instrumenter.java
@@ -72,7 +72,7 @@ public class Instrumenter<REQUEST, RESPONSE> {
   private final List<? extends RequestListener> requestListeners;
   private final ErrorCauseExtractor errorCauseExtractor;
   @Nullable private final StartTimeExtractor<REQUEST> startTimeExtractor;
-  @Nullable private final EndTimeExtractor<RESPONSE> endTimeExtractor;
+  @Nullable private final EndTimeExtractor<REQUEST, RESPONSE> endTimeExtractor;
   private final boolean disabled;
   private final SpanSuppressionStrategy spanSuppressionStrategy;
 
@@ -187,7 +187,7 @@ public class Instrumenter<REQUEST, RESPONSE> {
     }
 
     if (endTimeExtractor != null) {
-      span.end(endTimeExtractor.extract(response));
+      span.end(endTimeExtractor.extract(request, response));
     } else {
       span.end();
     }
diff --git a/instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterBuilder.java b/instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterBuilder.java
index 93535b0723..fc43ef608f 100644
--- a/instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterBuilder.java
+++ b/instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterBuilder.java
@@ -54,7 +54,7 @@ public final class InstrumenterBuilder<REQUEST, RESPONSE> {
       SpanStatusExtractor.getDefault();
   ErrorCauseExtractor errorCauseExtractor = ErrorCauseExtractor.jdk();
   @Nullable StartTimeExtractor<REQUEST> startTimeExtractor = null;
-  @Nullable EndTimeExtractor<RESPONSE> endTimeExtractor = null;
+  @Nullable EndTimeExtractor<REQUEST, RESPONSE> endTimeExtractor = null;
   boolean disabled = false;
 
   private boolean enableSpanSuppressionByType = ENABLE_SPAN_SUPPRESSION_BY_TYPE;
@@ -130,7 +130,8 @@ public final class InstrumenterBuilder<REQUEST, RESPONSE> {
    * determining start and end timestamps to the OpenTelemetry SDK.
    */
   public InstrumenterBuilder<REQUEST, RESPONSE> setTimeExtractors(
-      StartTimeExtractor<REQUEST> startTimeExtractor, EndTimeExtractor<RESPONSE> endTimeExtractor) {
+      StartTimeExtractor<REQUEST> startTimeExtractor,
+      EndTimeExtractor<REQUEST, RESPONSE> endTimeExtractor) {
     this.startTimeExtractor = requireNonNull(startTimeExtractor);
     this.endTimeExtractor = requireNonNull(endTimeExtractor);
     return this;
diff --git a/instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterTest.java b/instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterTest.java
index d250ee502e..359364c89d 100644
--- a/instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterTest.java
+++ b/instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterTest.java
@@ -487,7 +487,7 @@ class InstrumenterTest {
     Instrumenter<Instant, Instant> instrumenter =
         Instrumenter.<Instant, Instant>newBuilder(
                 otelTesting.getOpenTelemetry(), "test", request -> "test span")
-            .setTimeExtractors(request -> request, response -> response)
+            .setTimeExtractors(request -> request, (request, response) -> response)
             .newInstrumenter();
 
     Instant startTime = Instant.ofEpochSecond(100);
@@ -912,12 +912,12 @@ class InstrumenterTest {
   }
 
   private static Instrumenter<Map<String, String>, Map<String, String>> getInstrumenterWithType(
-      boolean enableInstrumenation, AttributesExtractor... attributeExtractors) {
+      boolean enableInstrumentation, AttributesExtractor... attributeExtractors) {
     InstrumenterBuilder<Map<String, String>, Map<String, String>> builder =
         Instrumenter.<Map<String, String>, Map<String, String>>newBuilder(
                 otelTesting.getOpenTelemetry(), "test", unused -> "span")
             .addAttributesExtractors(attributeExtractors)
-            .enableInstrumentationTypeSuppression(enableInstrumenation);
+            .enableInstrumentationTypeSuppression(enableInstrumentation);
 
     return builder.newClientInstrumenter(Map::put);
   }
diff --git a/instrumentation/jms-1.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jms/JmsSingletons.java b/instrumentation/jms-1.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jms/JmsSingletons.java
index e31326bdd9..f4caa86818 100644
--- a/instrumentation/jms-1.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jms/JmsSingletons.java
+++ b/instrumentation/jms-1.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jms/JmsSingletons.java
@@ -36,7 +36,8 @@ public final class JmsSingletons {
         Instrumenter.<MessageWithDestination, Void>newBuilder(
                 otel, INSTRUMENTATION_NAME, spanNameExtractor)
             .addAttributesExtractor(attributesExtractor)
-            .setTimeExtractors(MessageWithDestination::getStartTime, response -> Instant.now())
+            .setTimeExtractors(
+                MessageWithDestination::getStartTime, (request, response) -> Instant.now())
             .newInstrumenter(SpanKindExtractor.alwaysConsumer());
     LISTENER_INSTRUMENTER =
         Instrumenter.<MessageWithDestination, Void>newBuilder(