fix InstrumentationState

Author: EddeCCC

src/main/java/rocks/inspectit/gepard/agent/instrumentation/cache/process/BatchInstrumenter.java

@@ -70,6 +70,9 @@ Set<Class<?>> getNextBatch(int batchSize) {
       checkedClassesCount++;
 
       try {
+        if ("com.example.demo.GreetingController".equals(clazz.getName()))
+          System.out.println("GREET");
+
         boolean shouldInstrument = configurationResolver.shouldInstrument(clazz);
         boolean isInstrumented = instrumentationState.isInstrumented(clazz);
 

src/main/java/rocks/inspectit/gepard/agent/internal/instrumentation/InstrumentationState.java

@@ -2,52 +2,64 @@
 
 import com.github.benmanes.caffeine.cache.Cache;
 import com.github.benmanes.caffeine.cache.Caffeine;
+import java.util.Map;
 import java.util.Objects;
 import rocks.inspectit.gepard.agent.internal.instrumentation.model.ClassInstrumentationConfiguration;
 
 /** Stores the instrumentation configuration of all instrumented classes. */
 public class InstrumentationState {
 
   /** Store for every instrumented class */
-  private final Cache<Class<?>, ClassInstrumentationConfiguration> activeInstrumentations;
+  private final Cache<InstrumentedType, ClassInstrumentationConfiguration> activeInstrumentations;
 
   private InstrumentationState() {
-    this.activeInstrumentations = Caffeine.newBuilder().weakKeys().build();
+    this.activeInstrumentations = Caffeine.newBuilder().build();
   }
 
   public static InstrumentationState create() {
     return new InstrumentationState();
   }
 
   /**
-   * Checks, if the provided class type needs to be instrumented.
+   * Checks, if the provided class is already instrumented.
    *
    * @param clazz the class object
-   * @return true, if the provided type needs an instrumentation
+   * @return true, if the provided type is already instrumented.
    */
   public boolean isInstrumented(Class<?> clazz) {
-    ClassInstrumentationConfiguration activeConfig = activeInstrumentations.getIfPresent(clazz);
+    return activeInstrumentations.asMap().entrySet().stream()
+        .filter(entry -> entry.getKey().isEqualTo(clazz))
+        .map(Map.Entry::getValue)
+        .map(ClassInstrumentationConfiguration::isInstrumented)
+        .findAny()
+        .orElse(false);
+  }
+
+  public boolean isInstrumented(InstrumentedType instrumentedType) {
+    ClassInstrumentationConfiguration activeConfig =
+        activeInstrumentations.getIfPresent(instrumentedType);
 
     if (Objects.nonNull(activeConfig)) return activeConfig.isInstrumented();
     return false;
   }
 
   /**
-   * Adds the provided class type to the active instrumentations
+   * Adds the instrumented type to the active instrumentations
    *
-   * @param clazz the class object
+   * @param type the instrumented type
    */
-  public void addInstrumentation(Class<?> clazz) {
+  public void addInstrumentedType(InstrumentedType type) {
     ClassInstrumentationConfiguration newConfig = new ClassInstrumentationConfiguration(true);
-    activeInstrumentations.put(clazz, newConfig);
+    activeInstrumentations.put(type, newConfig);
   }
 
   /**
-   * Removes the provided class type from the active instrumentations
+   * Removes the type from the active instrumentations
    *
-   * @param clazz the class object
+   * @param type the uninstrumented type
    */
-  public void invalidateInstrumentation(Class<?> clazz) {
-    activeInstrumentations.invalidate(clazz);
+  public void invalidateInstrumentedType(InstrumentedType type) {
+    System.out.println("INVALIDATED");
+    activeInstrumentations.invalidate(type);
   }
 }

src/main/java/rocks/inspectit/gepard/agent/internal/instrumentation/InstrumentedType.java

@@ -0,0 +1,49 @@
+package rocks.inspectit.gepard.agent.internal.instrumentation;
+
+import net.bytebuddy.agent.builder.AgentBuilder;
+
+import java.util.Objects;
+
+/**
+ * Stores the full name as well as the class loader of a specific type.<br>
+ * Since the method {@link AgentBuilder.Transformer#transform} does not provide a {@code Class}
+ * object of the transformed type, we need to use the class name and the class loader to store our
+ * instrumented types.<br>
+ * We cannot create a {@code Class} object during transformation, since the class's byte-code is
+ * still modified and not "ready" yet. Thus, {@code Class.forName()} isn't working. However, we
+ * somehow need to remember which classes have been instrumented, so we store all necessary
+ * information to identify a specific class.
+ */
+public class InstrumentedType {
+
+  private final String typeName;
+
+  private final ClassLoader classLoader;
+
+  public InstrumentedType(final String typeName, final ClassLoader classLoader) {
+    this.typeName = typeName;
+    this.classLoader = classLoader;
+  }
+
+  /**
+   * Checks, if the provided class objects references this type
+   *
+   * @param clazz the class object
+   * @return true, if the provided class objects references this type
+   */
+  public boolean isEqualTo(Class<?> clazz) {
+    return typeName.equals(clazz.getName()) && classLoader.equals(clazz.getClassLoader());
+  }
+
+  @Override
+  public boolean equals(Object other) {
+    if (other instanceof InstrumentedType otherType)
+      return typeName.equals(otherType.typeName) && classLoader.equals(otherType.classLoader);
+    return false;
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(typeName, classLoader);
+  }
+}

src/main/java/rocks/inspectit/gepard/agent/transformation/DynamicTransformer.java

@@ -2,7 +2,6 @@
 
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 
-import java.lang.invoke.TypeDescriptor;
 import java.security.ProtectionDomain;
 import java.util.*;
 import net.bytebuddy.agent.builder.AgentBuilder;
@@ -15,6 +14,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import rocks.inspectit.gepard.agent.internal.instrumentation.InstrumentationState;
+import rocks.inspectit.gepard.agent.internal.instrumentation.InstrumentedType;
 import rocks.inspectit.gepard.agent.resolver.ConfigurationResolver;
 import rocks.inspectit.gepard.agent.transformation.advice.InspectitAdvice;
 
@@ -31,16 +31,9 @@ public class DynamicTransformer implements AgentBuilder.Transformer {
   /** The instrumentation state of the agent */
   private final InstrumentationState instrumentationState;
 
-  /**
-   * Set of instrumented types used by this transformer to prevent to call Class.forName() for every
-   * call of transform()
-   */
-  private final Set<TypeDescription> instrumentedTypes;
-
   DynamicTransformer(ConfigurationResolver resolver, InstrumentationState instrumentationState) {
     this.resolver = resolver;
     this.instrumentationState = instrumentationState;
-    this.instrumentedTypes = new HashSet<>();
   }
 
   /**
@@ -60,68 +53,22 @@ public DynamicType.Builder<?> transform(
       ClassLoader classLoader,
       JavaModule module,
       ProtectionDomain protectionDomain) {
+    InstrumentedType currentType = new InstrumentedType(typeDescription.getName(), classLoader);
     if (resolver.shouldInstrument(typeDescription)) {
-      log.info("Adding transformation to {}", typeDescription.getName()); // TODO log.debug()
+      log.debug("Adding transformation to {}", typeDescription.getName());
 
       // Currently, all methods of the type are instrumented
       ElementMatcher<? super MethodDescription> elementMatcher = isMethod();
       builder = builder.visit(Advice.to(InspectitAdvice.class).on(elementMatcher));
 
       // Mark type as instrumented
-      // TODO das macht noch Probleme im ScopeTest...
-      addInstrumentation(typeDescription, classLoader);
-    } else if (instrumentedTypes.contains(typeDescription)) {
-      log.info("Removing transformation from {}", typeDescription.getName()); // TODO log.debug()
+      instrumentationState.addInstrumentedType(currentType);
+    } else if (instrumentationState.isInstrumented(currentType)) {
+      log.debug("Removing transformation from {}", typeDescription.getName());
       // Mark type as uninstrumented or deinstrumented
-      invalidateInstrumentation(typeDescription, classLoader);
+      instrumentationState.invalidateInstrumentedType(currentType);
     }
 
     return builder;
   }
-
-  /**
-   * Marks the class as instrumented.
-   *
-   * @param typeDescription the class type
-   * @param classLoader the classloader, used for accessing the class object
-   */
-  private void addInstrumentation(TypeDescription typeDescription, ClassLoader classLoader) {
-    Class<?> instrumentedClass = toClass(typeDescription, classLoader);
-    if (Objects.nonNull(instrumentedClass)) {
-      instrumentationState.addInstrumentation(instrumentedClass);
-      instrumentedTypes.add(typeDescription);
-    }
-  }
-
-  /**
-   * Removes the class from marked instrumentations.
-   *
-   * @param typeDescription the class type description
-   * @param classLoader the classloader, used for accessing the class object
-   */
-  private void invalidateInstrumentation(TypeDescription typeDescription, ClassLoader classLoader) {
-    Class<?> deinstrumentedClass = toClass(typeDescription, classLoader);
-    if (Objects.nonNull(deinstrumentedClass)) {
-      instrumentationState.invalidateInstrumentation(deinstrumentedClass);
-      instrumentedTypes.remove(typeDescription);
-    }
-  }
-
-  /**
-   * Converts a bytebuddy {@link TypeDescription} to a {@link Class} object. The class is not
-   * initialized.
-   *
-   * @param type the type description
-   * @param classLoader the classloader, used for accessing the class object
-   * @return the class object of the provided type description
-   */
-  private Class<?> toClass(TypeDescription type, ClassLoader classLoader) {
-    Class<?> clazz = null;
-    try {
-      clazz = Class.forName(type.getName(), false, classLoader);
-    } catch (Exception e) {
-      log.error("Could not update instrumentation state for {}", type.getName(), e);
-    }
-    return clazz;
-  }
 }

src/test/java/rocks/inspectit/gepard/agent/instrumentation/cache/process/BatchInstrumenterTest.java

@@ -8,6 +8,7 @@
 import java.lang.instrument.UnmodifiableClassException;
 import java.util.HashSet;
 import java.util.Set;
+import net.bytebuddy.description.type.TypeDescription;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;

src/test/java/rocks/inspectit/gepard/agent/integrationtest/IntegrationTestBase.java

@@ -6,7 +6,6 @@
 import com.google.protobuf.util.JsonFormat;
 import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest;
 import io.opentelemetry.proto.trace.v1.Span;
-
 import java.io.IOException;
 import java.util.Collection;
 import java.util.Collections;
@@ -30,7 +29,6 @@
 import org.testcontainers.utility.MountableFile;
 import rocks.inspectit.gepard.agent.integrationtest.utils.OkHttpUtils;
 
-
 /**
  * Base class for integration tests. First starts the Observability Backend Mock and Configuration
  * Server Mock as Testcontainers. We can define a target application to test by extending this class

src/test/java/rocks/inspectit/gepard/agent/internal/instrumentation/InstrumentationStateTest.java

@@ -3,27 +3,41 @@
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import net.bytebuddy.description.type.TypeDescription;
 import org.junit.jupiter.api.Test;
 
+// TODO Write more tests
+
 class InstrumentationStateTest {
 
   private static final Class<?> TEST_CLASS = InstrumentationStateTest.class;
 
+  private static final InstrumentedType TEST_TYPE = new InstrumentedType(TEST_CLASS.getName(), TEST_CLASS.getClassLoader());
+
   @Test
-  void typeIsNotInstrumented() {
+  void classIsNotInstrumented() {
     InstrumentationState state = InstrumentationState.create();
 
     boolean isInstrumented = state.isInstrumented(TEST_CLASS);
 
     assertFalse(isInstrumented);
   }
 
+  @Test
+  void typeIsNotInstrumented() {
+    InstrumentationState state = InstrumentationState.create();
+
+    boolean isInstrumented = state.isInstrumented(TEST_TYPE);
+
+    assertFalse(isInstrumented);
+  }
+
   @Test
   void typeIsInstrumented() {
     InstrumentationState state = InstrumentationState.create();
 
-    state.addInstrumentation(TEST_CLASS);
-    boolean isInstrumented = state.isInstrumented(TEST_CLASS);
+    state.addInstrumentedType(TEST_TYPE);
+    boolean isInstrumented = state.isInstrumented(TEST_TYPE);
 
     assertTrue(isInstrumented);
   }
@@ -32,9 +46,9 @@ void typeIsInstrumented() {
   void typeIsDeinstrumented() {
     InstrumentationState state = InstrumentationState.create();
 
-    state.addInstrumentation(TEST_CLASS);
-    state.invalidateInstrumentation(TEST_CLASS);
-    boolean isInstrumented = state.isInstrumented(TEST_CLASS);
+    state.addInstrumentedType(TEST_TYPE);
+    state.invalidateInstrumentedType(TEST_TYPE);
+    boolean isInstrumented = state.isInstrumented(TEST_TYPE);
 
     assertFalse(isInstrumented);
   }