Reuse servlet31 advice in sparkweb

instrumentation/servlet/servlet-3.1/src/main/java/io/opentelemetry/instrumentation/hypertrace/servlet/v3_1/Servlet31Advice.java

@@ -0,0 +1,140 @@
+/*
+ * Copyright The Hypertrace Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.opentelemetry.instrumentation.hypertrace.servlet.v3_1;
+
+import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.Servlet3HttpServerTracer.TRACER;
+
+import io.opentelemetry.instrumentation.hypertrace.servlet.common.ServletSpanDecorator;
+import io.opentelemetry.trace.Span;
+import java.util.Enumeration;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import net.bytebuddy.asm.Advice;
+import org.hypertrace.agent.blocking.BlockingProvider;
+import org.hypertrace.agent.blocking.BlockingResult;
+import org.hypertrace.agent.core.DynamicConfig;
+import org.hypertrace.agent.core.HypertraceSemanticAttributes;
+
+public class Servlet31Advice {
+
+  // request attribute key injected at first filerChain.doFilter
+  private static final String ALREADY_LOADED = "__org.hypertrace.agent.on_start_executed";
+
+  @Advice.OnMethodEnter(suppress = Throwable.class, skipOn = BlockingResult.class)
+  public static Object start(
+      @Advice.Argument(value = 0, readOnly = false) ServletRequest request,
+      @Advice.Argument(value = 1, readOnly = false) ServletResponse response,
+      @Advice.Local("rootStart") Boolean rootStart) {
+    if (!(request instanceof HttpServletRequest) || !(response instanceof HttpServletResponse)) {
+      return null;
+    }
+
+    if (!DynamicConfig.isEnabled(InstrumentationName.INSTRUMENTATION_NAME)) {
+      return null;
+    }
+
+    // TODO run on every doFilter and check if user removed wrapper
+    // TODO what if user unwraps request and reads the body?
+
+    // run the instrumentation only for the root FilterChain.doFilter()
+    if (request.getAttribute(ALREADY_LOADED) != null) {
+      return null;
+    }
+    request.setAttribute(ALREADY_LOADED, true);
+
+    HttpServletRequest httpRequest = (HttpServletRequest) request;
+    HttpServletResponse httpResponse = (HttpServletResponse) response;
+    Span currentSpan = TRACER.getCurrentSpan();
+
+    rootStart = true;
+    response = new BufferingHttpServletResponse(httpResponse);
+    request = new BufferingHttpServletRequest(httpRequest, (HttpServletResponse) response);
+
+    ServletSpanDecorator.addSessionId(currentSpan, httpRequest);
+
+    // set request headers
+    Enumeration<String> headerNames = httpRequest.getHeaderNames();
+    Map<String, String> headers = new HashMap<>();
+    while (headerNames.hasMoreElements()) {
+      String headerName = headerNames.nextElement();
+      String headerValue = httpRequest.getHeader(headerName);
+      currentSpan.setAttribute(
+          HypertraceSemanticAttributes.httpRequestHeader(headerName), headerValue);
+      headers.put(headerName, headerValue);
+    }
+    BlockingResult blockingResult = BlockingProvider.getBlockingEvaluator().evaluate(headers);
+    currentSpan.setAttribute(
+        HypertraceSemanticAttributes.OPA_RESULT, blockingResult.blockExecution());
+    if (blockingResult.blockExecution()) {
+      httpResponse.setStatus(403);
+      currentSpan.setAttribute(HypertraceSemanticAttributes.OPA_REASON, blockingResult.getReason());
+      return blockingResult;
+    }
+    return null;
+  }
+
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+  public static void stopSpan(
+      @Advice.Argument(0) ServletRequest request,
+      @Advice.Argument(1) ServletResponse response,
+      @Advice.Local("rootStart") Boolean rootStart) {
+    if (rootStart != null) {
+      if (!(request instanceof BufferingHttpServletRequest)
+          || !(response instanceof BufferingHttpServletResponse)) {
+        return;
+      }
+
+      request.removeAttribute(ALREADY_LOADED);
+      Span currentSpan = TRACER.getCurrentSpan();
+
+      AtomicBoolean responseHandled = new AtomicBoolean(false);
+      if (request.isAsyncStarted()) {
+        try {
+          request
+              .getAsyncContext()
+              .addListener(new BodyCaptureAsyncListener(responseHandled, currentSpan));
+        } catch (IllegalStateException e) {
+          // org.eclipse.jetty.server.Request may throw an exception here if request became
+          // finished after check above. We just ignore that exception and move on.
+        }
+      }
+
+      if (!request.isAsyncStarted() && responseHandled.compareAndSet(false, true)) {
+        BufferingHttpServletResponse bufferingResponse = (BufferingHttpServletResponse) response;
+        BufferingHttpServletRequest bufferingRequest = (BufferingHttpServletRequest) request;
+
+        // set response headers
+        for (String headerName : bufferingResponse.getHeaderNames()) {
+          String headerValue = bufferingResponse.getHeader(headerName);
+          currentSpan.setAttribute(
+              HypertraceSemanticAttributes.httpResponseHeader(headerName), headerValue);
+        }
+        // Bodies are captured at the end after all user processing.
+        currentSpan.setAttribute(
+            HypertraceSemanticAttributes.HTTP_REQUEST_BODY,
+            bufferingRequest.getBufferedBodyAsString());
+        currentSpan.setAttribute(
+            HypertraceSemanticAttributes.HTTP_RESPONSE_BODY, bufferingResponse.getBufferAsString());
+      }
+    }
+  }
+}

instrumentation/servlet/servlet-3.1/src/main/java/io/opentelemetry/instrumentation/hypertrace/servlet/v3_1/Servlet31BodyInstrumentation.java

@@ -16,7 +16,6 @@
 
 package io.opentelemetry.instrumentation.hypertrace.servlet.v3_1;
 
-import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.Servlet3HttpServerTracer.TRACER;
 import static io.opentelemetry.javaagent.tooling.ClassLoaderMatcher.hasClassesNamed;
 import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.AgentElementMatchers.safeHasSuperType;
 import static io.opentelemetry.javaagent.tooling.matcher.NameMatchers.namedOneOf;
@@ -26,25 +25,11 @@
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import com.google.auto.service.AutoService;
-import io.opentelemetry.instrumentation.hypertrace.servlet.common.ServletSpanDecorator;
 import io.opentelemetry.javaagent.tooling.Instrumenter;
-import io.opentelemetry.trace.Span;
-import java.util.Enumeration;
-import java.util.HashMap;
 import java.util.Map;
-import java.util.concurrent.atomic.AtomicBoolean;
-import javax.servlet.ServletRequest;
-import javax.servlet.ServletResponse;
-import javax.servlet.http.HttpServletRequest;
-import javax.servlet.http.HttpServletResponse;
-import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
-import org.hypertrace.agent.blocking.BlockingProvider;
-import org.hypertrace.agent.blocking.BlockingResult;
-import org.hypertrace.agent.core.DynamicConfig;
-import org.hypertrace.agent.core.HypertraceSemanticAttributes;
 
 /**
  * TODO https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1395 is resolved
@@ -105,6 +90,7 @@ public String[] helperClassNames() {
       packageName + ".BufferingHttpServletRequest",
       packageName + ".BufferingHttpServletRequest$ServletInputStreamWrapper",
       packageName + ".BodyCaptureAsyncListener",
+      packageName + ".Servlet31Advice",
     };
   }
 
@@ -115,112 +101,6 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
             .and(takesArgument(0, named("javax.servlet.ServletRequest")))
             .and(takesArgument(1, named("javax.servlet.ServletResponse")))
             .and(isPublic()),
-        FilterAdvice.class.getName());
-  }
-
-  public static class FilterAdvice {
-    // request attribute key injected at first filerChain.doFilter
-    private static final String ALREADY_LOADED = "__org.hypertrace.agent.on_start_executed";
-
-    @Advice.OnMethodEnter(suppress = Throwable.class, skipOn = BlockingResult.class)
-    public static Object start(
-        @Advice.Argument(value = 0, readOnly = false) ServletRequest request,
-        @Advice.Argument(value = 1, readOnly = false) ServletResponse response,
-        @Advice.Local("rootStart") Boolean rootStart) {
-      if (!(request instanceof HttpServletRequest) || !(response instanceof HttpServletResponse)) {
-        return null;
-      }
-
-      if (!DynamicConfig.isEnabled(InstrumentationName.INSTRUMENTATION_NAME)) {
-        return null;
-      }
-
-      // TODO run on every doFilter and check if user removed wrapper
-      // TODO what if user unwraps request and reads the body?
-
-      // run the instrumentation only for the root FilterChain.doFilter()
-      if (request.getAttribute(ALREADY_LOADED) != null) {
-        return null;
-      }
-      request.setAttribute(ALREADY_LOADED, true);
-
-      HttpServletRequest httpRequest = (HttpServletRequest) request;
-      HttpServletResponse httpResponse = (HttpServletResponse) response;
-      Span currentSpan = TRACER.getCurrentSpan();
-
-      rootStart = true;
-      response = new BufferingHttpServletResponse(httpResponse);
-      request = new BufferingHttpServletRequest(httpRequest, (HttpServletResponse) response);
-
-      ServletSpanDecorator.addSessionId(currentSpan, httpRequest);
-
-      // set request headers
-      Enumeration<String> headerNames = httpRequest.getHeaderNames();
-      Map<String, String> headers = new HashMap<>();
-      while (headerNames.hasMoreElements()) {
-        String headerName = headerNames.nextElement();
-        String headerValue = httpRequest.getHeader(headerName);
-        currentSpan.setAttribute(
-            HypertraceSemanticAttributes.httpRequestHeader(headerName), headerValue);
-        headers.put(headerName, headerValue);
-      }
-      BlockingResult blockingResult = BlockingProvider.getBlockingEvaluator().evaluate(headers);
-      currentSpan.setAttribute(
-          HypertraceSemanticAttributes.OPA_RESULT, blockingResult.blockExecution());
-      if (blockingResult.blockExecution()) {
-        httpResponse.setStatus(403);
-        currentSpan.setAttribute(
-            HypertraceSemanticAttributes.OPA_REASON, blockingResult.getReason());
-        return blockingResult;
-      }
-      return null;
-    }
-
-    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void stopSpan(
-        @Advice.Argument(0) ServletRequest request,
-        @Advice.Argument(1) ServletResponse response,
-        @Advice.Local("rootStart") Boolean rootStart) {
-      if (rootStart != null) {
-        if (!(request instanceof BufferingHttpServletRequest)
-            || !(response instanceof BufferingHttpServletResponse)) {
-          return;
-        }
-
-        request.removeAttribute(ALREADY_LOADED);
-        Span currentSpan = TRACER.getCurrentSpan();
-
-        AtomicBoolean responseHandled = new AtomicBoolean(false);
-        if (request.isAsyncStarted()) {
-          try {
-            request
-                .getAsyncContext()
-                .addListener(new BodyCaptureAsyncListener(responseHandled, currentSpan));
-          } catch (IllegalStateException e) {
-            // org.eclipse.jetty.server.Request may throw an exception here if request became
-            // finished after check above. We just ignore that exception and move on.
-          }
-        }
-
-        if (!request.isAsyncStarted() && responseHandled.compareAndSet(false, true)) {
-          BufferingHttpServletResponse bufferingResponse = (BufferingHttpServletResponse) response;
-          BufferingHttpServletRequest bufferingRequest = (BufferingHttpServletRequest) request;
-
-          // set response headers
-          for (String headerName : bufferingResponse.getHeaderNames()) {
-            String headerValue = bufferingResponse.getHeader(headerName);
-            currentSpan.setAttribute(
-                HypertraceSemanticAttributes.httpResponseHeader(headerName), headerValue);
-          }
-          // Bodies are captured at the end after all user processing.
-          currentSpan.setAttribute(
-              HypertraceSemanticAttributes.HTTP_REQUEST_BODY,
-              bufferingRequest.getBufferedBodyAsString());
-          currentSpan.setAttribute(
-              HypertraceSemanticAttributes.HTTP_RESPONSE_BODY,
-              bufferingResponse.getBufferAsString());
-        }
-      }
-    }
+        Servlet31Advice.class.getName());
   }
 }