Added Instrumentation for ServletContainerInitializer#onStartup (#2449)

Author: tobiasstadler

apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/InitServiceNameInstrumentation.java renamed to apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/servicename/InitServiceNameInstrumentation.java

@@ -16,8 +16,10 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-package co.elastic.apm.agent.servlet;
+package co.elastic.apm.agent.servlet.servicename;
 
+import co.elastic.apm.agent.servlet.AbstractServletInstrumentation;
+import co.elastic.apm.agent.servlet.ServletServiceNameHelper;
 import co.elastic.apm.agent.servlet.adapter.JakartaServletApiAdapter;
 import co.elastic.apm.agent.servlet.adapter.JavaxServletApiAdapter;
 import net.bytebuddy.asm.Advice;

apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/servicename/ServletContainerInitializerServiceNameInstrumentation.java

@@ -0,0 +1,114 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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 co.elastic.apm.agent.servlet.servicename;
+
+import co.elastic.apm.agent.servlet.AbstractServletInstrumentation;
+import co.elastic.apm.agent.servlet.ServletServiceNameHelper;
+import co.elastic.apm.agent.servlet.adapter.JakartaServletApiAdapter;
+import co.elastic.apm.agent.servlet.adapter.JavaxServletApiAdapter;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.NamedElement;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+import javax.annotation.Nullable;
+import java.util.Set;
+
+import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+import static net.bytebuddy.matcher.ElementMatchers.nameContains;
+import static net.bytebuddy.matcher.ElementMatchers.nameEndsWith;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+/**
+ * Instruments
+ * <ul>
+ *     <li>{@link javax.servlet.ServletContainerInitializer#onStartup(java.util.Set, javax.servlet.ServletContext)} </li>
+ *     <li>{@link jakarta.servlet.ServletContainerInitializer#onStartup(java.util.Set, jakarta.servlet.ServletContext)}</li>
+ * </ul>
+ * <p>
+ * Determines the service name based on the webapp's {@code META-INF/MANIFEST.MF} file early in the startup process.
+ * As this doesn't work with runtime attachment, the service name is also determined when the first request comes in.
+ */
+public abstract class ServletContainerInitializerServiceNameInstrumentation extends AbstractServletInstrumentation {
+
+    @Override
+    public ElementMatcher<? super NamedElement> getTypeMatcherPreFilter() {
+        return nameContains("Initializer");
+    }
+
+    @Override
+    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+        return not(isInterface()).and(hasSuperType(namedOneOf(
+            "javax.servlet.ServletContainerInitializer", "jakarta.servlet.ServletContainerInitializer")));
+    }
+
+    @Override
+    public ElementMatcher<? super MethodDescription> getMethodMatcher() {
+        return named("onStartup")
+            .and(takesArguments(2))
+            .and(takesArgument(0, Set.class))
+            .and(takesArgument(1, nameEndsWith("ServletContext")));
+    }
+
+    public static class JavaxInitServiceNameInstrumentation extends ServletContainerInitializerServiceNameInstrumentation {
+
+        private static final JavaxServletApiAdapter adapter = JavaxServletApiAdapter.get();
+
+        @Override
+        public String rootClassNameThatClassloaderCanLoad() {
+            return "javax.servlet.AsyncContext";
+        }
+
+        public static class AdviceClass {
+
+            @Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
+            public static void onEnter(@Advice.Argument(1) @Nullable Object servletContext) {
+                if (servletContext instanceof javax.servlet.ServletContext) {
+                    ServletServiceNameHelper.determineServiceName(adapter, (javax.servlet.ServletContext) servletContext, tracer);
+                }
+            }
+        }
+    }
+
+    public static class JakartaInitServiceNameInstrumentation extends ServletContainerInitializerServiceNameInstrumentation {
+
+        private static final JakartaServletApiAdapter adapter = JakartaServletApiAdapter.get();
+
+        @Override
+        public String rootClassNameThatClassloaderCanLoad() {
+            return "jakarta.servlet.AsyncContext";
+        }
+
+        public static class AdviceClass {
+
+            @Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
+            public static void onEnter(@Advice.Argument(1) @Nullable Object servletContext) {
+                if (servletContext instanceof jakarta.servlet.ServletContext) {
+                    ServletServiceNameHelper.determineServiceName(adapter, (jakarta.servlet.ServletContext) servletContext, tracer);
+                }
+            }
+        }
+    }
+}

apm-agent-plugins/apm-servlet-plugin/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation

@@ -14,5 +14,7 @@ co.elastic.apm.agent.servlet.JakartaAsyncInstrumentation$JakartaStartAsyncInstru
 co.elastic.apm.agent.servlet.JakartaAsyncInstrumentation$JakartaAsyncContextInstrumentation
 co.elastic.apm.agent.servlet.JavaxRequestStreamRecordingInstrumentation
 co.elastic.apm.agent.servlet.JakartaRequestStreamRecordingInstrumentation
-co.elastic.apm.agent.servlet.InitServiceNameInstrumentation$JavaxInitServiceNameInstrumentation
-co.elastic.apm.agent.servlet.InitServiceNameInstrumentation$JakartaInitServiceNameInstrumentation
+co.elastic.apm.agent.servlet.servicename.InitServiceNameInstrumentation$JavaxInitServiceNameInstrumentation
+co.elastic.apm.agent.servlet.servicename.InitServiceNameInstrumentation$JakartaInitServiceNameInstrumentation
+co.elastic.apm.agent.servlet.servicename.ServletContainerInitializerServiceNameInstrumentation$JavaxInitServiceNameInstrumentation
+co.elastic.apm.agent.servlet.servicename.ServletContainerInitializerServiceNameInstrumentation$JakartaInitServiceNameInstrumentation

apm-agent-plugins/apm-servlet-plugin/src/test/java/co/elastic/apm/agent/servlet/InitServiceNameInstrumentationTest.java

@@ -29,18 +29,34 @@
 import javax.servlet.FilterChain;
 import javax.servlet.FilterConfig;
 import javax.servlet.Servlet;
+import javax.servlet.ServletContainerInitializer;
+import javax.servlet.ServletContext;
 import javax.servlet.ServletContextEvent;
 import javax.servlet.ServletContextListener;
 import javax.servlet.ServletException;
 import javax.servlet.ServletRequest;
 import javax.servlet.ServletResponse;
 import javax.servlet.http.HttpServlet;
 import java.io.IOException;
+import java.util.Set;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
 class InitServiceNameInstrumentationTest extends AbstractInstrumentationTest {
 
+    @Test
+    void testOnStartup() {
+        ServletContainerInitializer servletContainerInitializer = new NoopServletContainerInitializer();
+
+        CustomManifestLoader cl = new CustomManifestLoader(() -> getClass().getResourceAsStream("/TEST-MANIFEST.MF"));
+        CustomManifestLoader.withThreadContextClassLoader(cl, () -> {
+            servletContainerInitializer.onStartup(null, new MockServletContext());
+            tracer.startRootTransaction(cl).end();
+        });
+
+        assertServiceInfo();
+    }
+
     @Test
     void testContextInitialized() {
         ServletContextListener servletContextListener = new NoopServletContextListener();
@@ -87,6 +103,12 @@ private void assertServiceInfo() {
         assertThat(traceContext.getServiceVersion()).isEqualTo("1.42.0");
     }
 
+    private static class NoopServletContainerInitializer implements ServletContainerInitializer {
+        @Override
+        public void onStartup(Set<Class<?>> classes, ServletContext servletContext) {
+        }
+    }
+
     private static class NoopServletContextListener implements ServletContextListener {
         @Override
         public void contextInitialized(ServletContextEvent servletContextEvent) {