Support JAX-RS client typed entities - message body writer/reader

instrumentation/apache-httpclient-4.0/build.gradle.kts

@@ -39,8 +39,10 @@ afterEvaluate{
 val versions: Map<String, String> by extra
 
 dependencies {
-    implementation("org.apache.httpcomponents:httpclient:4.0")
+    api(project(":instrumentation:java-streams"))
     api("io.opentelemetry.javaagent.instrumentation:opentelemetry-javaagent-apache-httpclient-4.0:${versions["opentelemetry_java_agent"]}")
 
+    implementation("org.apache.httpcomponents:httpclient:4.0")
+
     testImplementation(project(":testing-common"))
 }

instrumentation/build.gradle.kts

@@ -38,6 +38,7 @@ dependencies{
     implementation(project(":instrumentation:okhttp:okhttp-3.0"))
     implementation(project(":instrumentation:apache-httpclient-4.0"))
     implementation(project(":instrumentation:jaxrs-client-2.0"))
+    implementation(project(":instrumentation:java-streams"))
     implementation(project(":otel-extensions"))
 }
 

instrumentation/java-streams/build.gradle.kts

@@ -0,0 +1,12 @@
+plugins {
+    `java-library`
+    id("net.bytebuddy.byte-buddy")
+}
+
+afterEvaluate{
+    io.opentelemetry.instrumentation.gradle.bytebuddy.ByteBuddyPluginConfigurator(project,
+            sourceSets.main.get(),
+            "io.opentelemetry.javaagent.tooling.muzzle.collector.MuzzleCodeGenerationPlugin",
+            project(":javaagent-tooling").configurations["instrumentationMuzzle"] + configurations.runtimeClasspath
+    ).configure()
+}

instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/instrumentation/hypertrace/apachehttpclient/v4_0/InputStreamInstrumentationModule.java → instrumentation/java-streams/src/main/java/io/opentelemetry/instrumentation/hypertrace/java/inputstream/InputStreamInstrumentationModule.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.opentelemetry.instrumentation.hypertrace.apachehttpclient.v4_0;
+package io.opentelemetry.instrumentation.hypertrace.java.inputstream;
 
 import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.AgentElementMatchers.safeHasSuperType;
 import static io.opentelemetry.javaagent.tooling.matcher.NameMatchers.namedOneOf;
@@ -36,9 +36,19 @@
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
+import org.hypertrace.agent.core.GlobalObjectRegistry;
 
 /**
- * Maybe we could add optimization to instrument the input streams only when certain classes are
+ * {@link InputStream} instrumentation. The type matcher applies to all implementations. However
+ * only streams that are in the {@link GlobalObjectRegistry#inputStreamToSpanAndBufferMap} are
+ * instrumented, otherwise the instrumentation is noop.
+ *
+ * <p>If the stream is in the {@link GlobalObjectRegistry#inputStreamToSpanAndBufferMap} then result
+ * of read methods is also passed to the buffered stream (value) from the map. The instrumentation
+ * adds buffer to span from the map when read is finished (return -1), creates new span with buffer
+ * when the original span is not recording.
+ *
+ * <p>Maybe we could add optimization to instrument the input streams only when certain classes are
  * present in classloader e.g. classes from frameworks that we instrument.
  */
 @AutoService(InstrumentationModule.class)

instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/instrumentation/hypertrace/apachehttpclient/v4_0/InputStreamUtils.java → instrumentation/java-streams/src/main/java/io/opentelemetry/instrumentation/hypertrace/java/inputstream/InputStreamUtils.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.opentelemetry.instrumentation.hypertrace.apachehttpclient.v4_0;
+package io.opentelemetry.instrumentation.hypertrace.java.inputstream;
 
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.common.AttributeKey;

instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/instrumentation/hypertrace/apachehttpclient/v4_0/OutputStreamInstrumentationModule.java → instrumentation/java-streams/src/main/java/io/opentelemetry/instrumentation/hypertrace/java/outputstream/OutputStreamInstrumentationModule.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.opentelemetry.instrumentation.hypertrace.apachehttpclient.v4_0;
+package io.opentelemetry.instrumentation.hypertrace.java.outputstream;
 
 import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.AgentElementMatchers.extendsClass;
 import static io.opentelemetry.javaagent.tooling.matcher.NameMatchers.namedOneOf;

instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/instrumentation/hypertrace/apachehttpclient/v4_0/OutputStreamUtils.java → instrumentation/java-streams/src/main/java/io/opentelemetry/instrumentation/hypertrace/java/outputstream/OutputStreamUtils.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.opentelemetry.instrumentation.hypertrace.apachehttpclient.v4_0;
+package io.opentelemetry.instrumentation.hypertrace.java.outputstream;
 
 import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
 import java.io.IOException;

instrumentation/jaxrs-client-2.0/build.gradle.kts

@@ -16,7 +16,8 @@ muzzle {
         group = "io.dropwizard"
         module = "dropwizard-client"
         versions = "[0.8.0,)"
-        assertInverse = true
+        // TODO this is set in OTEL
+//        assertInverse = true
     }
 }
 
@@ -31,6 +32,7 @@ afterEvaluate{
 val versions: Map<String, String> by extra
 
 dependencies {
+    api(project(":instrumentation:java-streams"))
     api("io.opentelemetry.javaagent.instrumentation:opentelemetry-javaagent-jaxrs-client-2.0-common:${versions["opentelemetry_java_agent"]}")
 
     compileOnly("javax.ws.rs:javax.ws.rs-api:2.0.1")

instrumentation/jaxrs-client-2.0/src/main/java/io/opentelemetry/instrumentation/hypertrace/jaxrs/v2_0/JaxrsClientBodyCaptureFilter.java

@@ -26,11 +26,8 @@
 import javax.ws.rs.client.ClientRequestFilter;
 import javax.ws.rs.client.ClientResponseContext;
 import javax.ws.rs.client.ClientResponseFilter;
-import javax.ws.rs.core.Form;
-import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.MultivaluedMap;
 import org.hypertrace.agent.config.Config.AgentConfig;
-import org.hypertrace.agent.core.ContentTypeUtils;
 import org.hypertrace.agent.core.HypertraceConfig;
 import org.hypertrace.agent.core.HypertraceSemanticAttributes;
 import org.slf4j.Logger;
@@ -57,28 +54,8 @@ public void filter(ClientRequestContext requestContext) {
             HypertraceSemanticAttributes::httpRequestHeader,
             requestContext.getStringHeaders());
       }
-      if (requestContext.hasEntity()
-          && agentConfig.getDataCapture().getHttpBody().getRequest().getValue()) {
-        MediaType mediaType = requestContext.getMediaType();
-        if (mediaType == null || !ContentTypeUtils.shouldCapture(mediaType.toString())) {
-          return;
-        }
-
-        Object entity = requestContext.getEntity();
-        if (entity != null) {
-          if (entity instanceof Form) {
-            Form form = (Form) entity;
-            String content = getUrlEncodedContent(form);
-            currentSpan.setAttribute(HypertraceSemanticAttributes.HTTP_REQUEST_BODY, content);
-          } else {
-            currentSpan.setAttribute(
-                HypertraceSemanticAttributes.HTTP_REQUEST_BODY, entity.toString());
-          }
-        }
-      }
-      requestContext.getEntity();
     } catch (Exception ex) {
-      log.error("Exception while getting request entity or headers", ex);
+      log.error("Exception while getting request headers", ex);
     }
   }
 
@@ -100,26 +77,8 @@ public void filter(ClientRequestContext requestContext, ClientResponseContext re
             responseContext.getHeaders());
       }
     } catch (Exception ex) {
-      log.error("Exception while getting response entity or headers", ex);
-    }
-  }
-
-  private static String getUrlEncodedContent(Form form) {
-    MultivaluedMap<String, String> formMap = form.asMap();
-    StringBuilder sb = new StringBuilder();
-    if (formMap != null) {
-      for (Map.Entry<String, List<String>> entry : formMap.entrySet()) {
-        if (sb.length() > 0) {
-          sb.append("&");
-        }
-        for (String value : entry.getValue()) {
-          sb.append(entry.getKey());
-          sb.append("=");
-          sb.append(value);
-        }
-      }
+      log.error("Exception while getting response headers", ex);
     }
-    return sb.toString();
   }
 
   private static void captureHeaders(

instrumentation/jaxrs-client-2.0/src/main/java/io/opentelemetry/instrumentation/hypertrace/jaxrs/v2_0/JaxrsClientEntityInterceptor.java

@@ -16,63 +16,116 @@
 
 package io.opentelemetry.instrumentation.hypertrace.jaxrs.v2_0;
 
-import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.trace.Span;
-import io.opentelemetry.api.trace.Tracer;
-import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.instrumentation.jaxrsclient.v2_0.ClientTracingFilter;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
 import javax.ws.rs.WebApplicationException;
+import javax.ws.rs.core.HttpHeaders;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.ext.ReaderInterceptor;
 import javax.ws.rs.ext.ReaderInterceptorContext;
+import javax.ws.rs.ext.WriterInterceptor;
+import javax.ws.rs.ext.WriterInterceptorContext;
 import org.hypertrace.agent.config.Config.AgentConfig;
+import org.hypertrace.agent.core.ContentEncodingUtils;
+import org.hypertrace.agent.core.ContentLengthUtils;
 import org.hypertrace.agent.core.ContentTypeUtils;
+import org.hypertrace.agent.core.GlobalObjectRegistry;
+import org.hypertrace.agent.core.GlobalObjectRegistry.SpanAndBuffer;
 import org.hypertrace.agent.core.HypertraceConfig;
 import org.hypertrace.agent.core.HypertraceSemanticAttributes;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-public class JaxrsClientEntityInterceptor implements ReaderInterceptor {
+public class JaxrsClientEntityInterceptor implements ReaderInterceptor, WriterInterceptor {
 
   private static final Logger log = LoggerFactory.getLogger(JaxrsClientEntityInterceptor.class);
 
-  private static final Tracer TRACER =
-      OpenTelemetry.getGlobalTracer("org.hypertrace.java.jaxrs.client");
-
+  /** Writing response body to input stream */
   @Override
   public Object aroundReadFrom(ReaderInterceptorContext context)
       throws IOException, WebApplicationException {
 
-    Object entity = context.proceed();
+    MediaType mediaType = context.getMediaType();
+    AgentConfig agentConfig = HypertraceConfig.get();
+    if (mediaType == null
+        || !ContentTypeUtils.shouldCapture(mediaType.toString())
+        || !agentConfig.getDataCapture().getHttpBody().getResponse().getValue()) {
+      return context.proceed();
+    }
 
     Object spanObj = context.getProperty(ClientTracingFilter.SPAN_PROPERTY_NAME);
     if (!(spanObj instanceof Span)) {
       log.error(
           "Span object is not present in the context properties, response object will not be captured");
-      return entity;
+      return context.proceed();
     }
     Span currentSpan = (Span) spanObj;
 
-    MediaType mediaType = context.getMediaType();
-    AgentConfig agentConfig = HypertraceConfig.get();
-    if (mediaType == null
-        || !ContentTypeUtils.shouldCapture(mediaType.toString())
-        || !agentConfig.getDataCapture().getHttpBody().getResponse().getValue()) {
-      return entity;
+    // TODO as optimization the type could be checked here and if it is a primitive type e.g. String
+    // it could be read directly.
+    //    context.getType();
+
+    InputStream entityStream = context.getInputStream();
+    Object entity = null;
+    try {
+      String encodingStr = context.getHeaders().getFirst(HttpHeaders.CONTENT_ENCODING);
+      String contentLengthStr = context.getHeaders().getFirst(HttpHeaders.CONTENT_LENGTH);
+      int contentLength = ContentLengthUtils.parseLength(contentLengthStr);
+
+      ByteArrayOutputStream buffer = new ByteArrayOutputStream(contentLength);
+      GlobalObjectRegistry.inputStreamToSpanAndBufferMap.put(
+          entityStream,
+          new SpanAndBuffer(
+              currentSpan,
+              buffer,
+              HypertraceSemanticAttributes.HTTP_RESPONSE_BODY,
+              ContentEncodingUtils.toCharset(encodingStr)));
+      entity = context.proceed();
+    } catch (Exception ex) {
+      log.error("Exception while capturing response body", ex);
     }
+    return entity;
+  }
+
+  /** Writing request body to output stream */
+  @Override
+  public void aroundWriteTo(WriterInterceptorContext context)
+      throws IOException, WebApplicationException {
+
+    Object spanObj = context.getProperty(ClientTracingFilter.SPAN_PROPERTY_NAME);
+    if (!(spanObj instanceof Span)) {
+      log.error(
+          "Span object is not present in the context properties, request body will not be captured");
+      context.proceed();
+      return;
+    }
+    Span currentSpan = (Span) spanObj;
 
-    if (currentSpan.isRecording()) {
-      currentSpan.setAttribute(HypertraceSemanticAttributes.HTTP_RESPONSE_BODY, entity.toString());
-    } else {
-      TRACER
-          .spanBuilder(HypertraceSemanticAttributes.ADDITIONAL_DATA_SPAN_NAME)
-          .setParent(Context.root().with(currentSpan))
-          .setAttribute(HypertraceSemanticAttributes.HTTP_RESPONSE_BODY, entity.toString())
-          .startSpan()
-          .end();
+    AgentConfig agentConfig = HypertraceConfig.get();
+    if (agentConfig.getDataCapture().getHttpBody().getRequest().getValue()) {
+      MediaType mediaType = context.getMediaType();
+      if (mediaType == null || !ContentTypeUtils.shouldCapture(mediaType.toString())) {
+        context.proceed();
+        return;
+      }
     }
 
-    return entity;
+    // TODO length is not known
+    ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+    OutputStream entityStream = context.getOutputStream();
+    try {
+      GlobalObjectRegistry.outputStreamToBufferMap.put(entityStream, buffer);
+      context.proceed();
+    } catch (Exception ex) {
+      log.error("Failed to capture request body", ex);
+    } finally {
+      GlobalObjectRegistry.outputStreamToBufferMap.remove(entityStream);
+      // TODO encoding is not known
+      currentSpan.setAttribute(HypertraceSemanticAttributes.HTTP_REQUEST_BODY, buffer.toString());
+    }
   }
 }