fix: instrumentation

zebrunner · Mar 8, 2024 · be48496 · be48496
1 parent 9e62492
commit be48496
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/agent/src/main/java/com/zebrunner/mcloud/grid/agent/NodeAgent.java b/agent/src/main/java/com/zebrunner/mcloud/grid/agent/NodeAgent.java
@@ -7,7 +7,6 @@
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.NameMatcher;
 import net.bytebuddy.pool.TypePool;
-import org.openqa.selenium.grid.node.local.SessionSlot;
 
 import java.lang.instrument.Instrumentation;
 import java.util.logging.Logger;
@@ -21,6 +20,8 @@
 public class NodeAgent {
     private static final Logger LOGGER = Logger.getLogger(NodeAgent.class.getName());
     private static final String RELAY_SESSION_FACTORY_CLASS = "org.openqa.selenium.grid.node.relay.RelaySessionFactory";
+    private static final String SESSION_SLOT_CLASS = "org.openqa.selenium.grid.node.local.SessionSlot";
+
     private static final String TEST_METHOD_NAME = "test";
 
     public static void premain(String args, Instrumentation instrumentation) {
@@ -29,6 +30,8 @@ public static void premain(String args, Instrumentation instrumentation) {
                     .with(new AgentBuilder.InitializationStrategy.SelfInjection.Eager())
                     .type(named(RELAY_SESSION_FACTORY_CLASS))
                     .transform((builder, type, classloader, module, protectionDomain) -> addTestMethodInterceptor(builder))
+                    .type(named(SESSION_SLOT_CLASS))
+                    .transform((builder, type, classloader, module, protectionDomain) -> addReleaseMethodInterceptor(builder))
                     .installOn(instrumentation);
         } catch (Exception e) {
             LOGGER.warning(() -> "Could not init instrumentation.");
@@ -42,6 +45,11 @@ private static DynamicType.Builder<?> addTestMethodInterceptor(DynamicType.Build
                 .intercept(to(releaseMethodInterceptor()));
     }
 
+    private static DynamicType.Builder<?> addReleaseMethodInterceptor(DynamicType.Builder<?> builder) {
+        return builder.method(isReleaseMethod())
+                .intercept(to(releaseMethodInterceptor()));
+    }
+
     public static ElementMatcher<? super MethodDescription> isTestMethod() {
         return isPublic()
                 .and(not(isStatic()))