Skip to content

Commit

Permalink
refactor: add method hook state
Browse files Browse the repository at this point in the history
  • Loading branch information
EddeCCC committed Oct 1, 2024
1 parent eec2ccd commit 4f8f091
Show file tree
Hide file tree
Showing 7 changed files with 241 additions and 109 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ public class MethodHook implements IMethodHook {

@Override
public void onEnter(Object[] instrumentedMethodArgs, Object thiz) {
// Using our log4j does not work here...
// Using our log4j here will not be visible in the target application...
String message =
String.format(
"inspectIT: Enter MethodHook with %d args in %s",
Expand All @@ -25,7 +25,7 @@ public void onEnter(Object[] instrumentedMethodArgs, Object thiz) {
@Override
public void onExit(
Object[] instrumentedMethodArgs, Object thiz, Object returnValue, Throwable thrown) {
// Using our log4j does not work here...
// Using our log4j here will not be visible in the target application...
String exceptionMessage = Objects.nonNull(thrown) ? thrown.getMessage() : "no exception";
String message =
String.format(
Expand Down
Original file line number Diff line number Diff line change
@@ -1,10 +1,7 @@
/* (C) 2024 */
package rocks.inspectit.gepard.agent.instrumentation.hook;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
Expand All @@ -21,7 +18,8 @@
import rocks.inspectit.gepard.bootstrap.instrumentation.noop.NoopMethodHook;

/**
* TODO implement, add java doc
* Implements {@link IHookManager} and serves as singleton, which returns the {@link IMethodHook}
* for every instrumented method. The hooks are stored within the {@link MethodHookState}.
*
* <p>Note: In inspectIT Ocelot we don't implement {@link IHookManager} directly, but instead pass
* over {@link #getHook} as lambda to {@link Instances}. This should avoid issues with spring
Expand All @@ -32,10 +30,10 @@ public class MethodHookManager implements IHookManager {
private static final Logger log = LoggerFactory.getLogger(MethodHookManager.class);

/** Stores classes and all of their hooked methods. Will be kept up-to-date during runtime. */
private final Cache<Class<?>, HookedMethods> hooks;
private final MethodHookState hookState;

private MethodHookManager() {
hooks = Caffeine.newBuilder().weakKeys().build();
hookState = new MethodHookState();
}

/**
Expand All @@ -56,96 +54,47 @@ public static MethodHookManager create() {

@Override
public IMethodHook getHook(Class<?> clazz, String methodSignature) {
HookedMethods hookedMethods = hooks.getIfPresent(clazz);
HookedMethods hookedMethods = hookState.getIfPresent(clazz);
if (Objects.nonNull(hookedMethods)) return hookedMethods.getActiveHook(methodSignature);
return NoopMethodHook.INSTANCE;
}

/**
* @param clazz the instrumented class
* @param configuration the instrumentation configuration for the class
*/
public void updateHooksFor(Class<?> clazz, ClassInstrumentationConfiguration configuration) {
log.debug("Updating hooks for {}", clazz.getName());
ElementMatcher.Junction<MethodDescription> methodMatcher = configuration.methodMatcher();
TypeDescription type = TypeDescription.ForLoadedType.of(clazz);
Set<MethodDescription.InDefinedShape> matchedMethods =
type.getDeclaredMethods().stream()
.filter(methodMatcher::matches)
.collect(Collectors.toSet());
String className = clazz.getName();
log.debug("Updating hooks for {}", className);
Set<MethodDescription.InDefinedShape> instrumentedMethods =
getInstrumentedMethods(clazz, configuration);

ClassHookConfiguration classConfiguration = new ClassHookConfiguration();
if (!matchedMethods.isEmpty()) matchedMethods.forEach(classConfiguration::putHookConfiguration);

removeObsoleteHooks(clazz, matchedMethods);
updateHooks(clazz, classConfiguration);
}

private void removeObsoleteHooks(
Class<?> clazz, Set<MethodDescription.InDefinedShape> matchedMethods) {
Set<String> matchedSignatures =
matchedMethods.stream().map(this::getSignature).collect(Collectors.toSet());
instrumentedMethods.forEach(classConfiguration::putHookConfiguration);

HookedMethods hookedMethods = hooks.get(clazz, c -> new HookedMethods());
Set<String> methodSignatures = hookedMethods.getMethodSignatures();
int removeCounter = hookState.removeObsoleteHooks(clazz, instrumentedMethods);
log.debug("Removed {} obsolete method hooks for {}", removeCounter, className);

AtomicInteger operationCounter = new AtomicInteger(0);
methodSignatures.stream()
.filter(signature -> !matchedSignatures.contains(signature))
.forEach(
signature -> {
removeHook(clazz, signature);
operationCounter.addAndGet(1);
});
log.debug("Removed {} obsolete method hooks for {}", operationCounter.get(), clazz.getName());
}

private void updateHooks(Class<?> clazz, ClassHookConfiguration classConfiguration) {
AtomicInteger operationCounter = new AtomicInteger(0);
classConfiguration
.asMap()
.forEach(
(method, active) -> {
// Currently always true, later we should compare the current with the new config
if (active) {
String signature = getSignature(method);
Optional<MethodHook> maybeHook = getCurrentHook(clazz, signature);
if (maybeHook.isEmpty()) {
setHook(clazz, signature, new MethodHook());
operationCounter.addAndGet(1);
}
}
});
int updateCounter = hookState.updateHooks(clazz, classConfiguration);
// TODO This should be DEBUG, but should also be visible in tests
log.info("Updated {} method hooks for {}", operationCounter.get(), clazz.getName());
}

// TODO Alternativen suchen + Methode auslagern
// Wichtig, dass die Signaturen mit denen von ByteBuddy passen!
private String getSignature(MethodDescription methodDescription) {
String methodName = methodDescription.getName();
String parameters =
methodDescription.getParameters().asTypeList().stream()
.map(type -> type.asErasure().getTypeName())
.collect(Collectors.joining(","));
return methodName + "(" + parameters + ")";
}

private void setHook(Class<?> declaringClass, String methodSignature, MethodHook newHook) {
hooks
.asMap()
.computeIfAbsent(declaringClass, (v) -> new HookedMethods())
.putMethod(methodSignature, newHook);
}

private void removeHook(Class<?> declaringClass, String methodSignature) {
HookedMethods hookedMethods = hooks.getIfPresent(declaringClass);
if (Objects.nonNull(hookedMethods)) {
hookedMethods.removeMethod(methodSignature);
if (hookedMethods.noActiveHooks()) hooks.invalidate(declaringClass);
}
log.info("Updated {} method hooks for {}", updateCounter, className);
}

private Optional<MethodHook> getCurrentHook(Class<?> clazz, String methodSignature) {
HookedMethods hookedMethods = hooks.getIfPresent(clazz);
return Optional.ofNullable(hookedMethods)
.map(methods -> methods.getActiveHook(methodSignature));
/**
* Returns all methods of the provided class, which are included in the instrumentation
* configuration.
*
* @param clazz the class containing the methods
* @param configuration the instrumentation configuration for the class
* @return the set of all instrumented methods of the class
*/
private Set<MethodDescription.InDefinedShape> getInstrumentedMethods(
Class<?> clazz, ClassInstrumentationConfiguration configuration) {
ElementMatcher.Junction<MethodDescription> methodMatcher = configuration.methodMatcher();
TypeDescription type = TypeDescription.ForLoadedType.of(clazz);
return type.getDeclaredMethods().stream()
.filter(methodMatcher::matches)
.collect(Collectors.toSet());
}

/** Remove at shutdown */
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,149 @@
/* (C) 2024 */
package rocks.inspectit.gepard.agent.instrumentation.hook;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import net.bytebuddy.description.method.MethodDescription;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import rocks.inspectit.gepard.agent.instrumentation.hook.configuration.ClassHookConfiguration;
import rocks.inspectit.gepard.agent.instrumentation.hook.configuration.HookedMethods;

/** Stores the method hook configurations of all instrumented classes. */
public class MethodHookState {
private static final Logger log = LoggerFactory.getLogger(MethodHookState.class);

private final Cache<Class<?>, HookedMethods> hooks;

public MethodHookState() {
this.hooks = Caffeine.newBuilder().weakKeys().build();
}

/**
* @param clazz the class containing the methods
* @return the hooked methods for the provided class
*/
public HookedMethods getIfPresent(Class<?> clazz) {
return hooks.getIfPresent(clazz);
}

/**
* Removes all methods hooks of the provided class, which are no longer necessary. This is done by
* comparing the current instrumented methods with the already hooked methods. If a hooked method
* is not part of the instrumented methods, it will be removed.
*
* @param clazz the class containing the methods
* @param instrumentedMethods the methods, which should be hooked
* @return the amount of hooks removed
*/
public int removeObsoleteHooks(
Class<?> clazz, Set<MethodDescription.InDefinedShape> instrumentedMethods) {
Set<String> matchedSignatures =
instrumentedMethods.stream().map(this::getSignature).collect(Collectors.toSet());

HookedMethods hookedMethods = hooks.get(clazz, c -> new HookedMethods());
Set<String> methodSignatures = hookedMethods.getMethodSignatures();

AtomicInteger operationCounter = new AtomicInteger(0);
methodSignatures.stream()
.filter(signature -> !matchedSignatures.contains(signature))
.forEach(
signature -> {
removeHook(clazz, signature);
operationCounter.addAndGet(1);
});
return operationCounter.get();
}

/**
* Updates the method hooks by overwriting the previous hooks. Currently, there is no difference
* between method hooks, thus we don't need to compare them. This might change in the future.
*
* @param clazz the class containing the methods
* @param classConfiguration the hook configuration for the provided class
* @return the amount of updated hooks
*/
public int updateHooks(Class<?> clazz, ClassHookConfiguration classConfiguration) {
AtomicInteger operationCounter = new AtomicInteger(0);
classConfiguration
.asMap()
.forEach(
(method, active) -> {
// Currently always true, later we should compare the current with the new config
if (active) {
String signature = getSignature(method);
Optional<MethodHook> maybeHook = getCurrentHook(clazz, signature);
if (maybeHook.isEmpty()) {
setHook(clazz, signature, new MethodHook());
operationCounter.addAndGet(1);
}
}
});
return operationCounter.get();
}

/**
* Creates the signature for the provided method. <br>
* The signature of this method would be: {@code
* getSignature(net.bytebuddy.description.method.MethodDescription)} <br>
* It should match with the method signatures used by ByteBuddy. Multiple method parameter will be
* concatenated with "," and without spaces.
*
* @param methodDescription the method description
* @return the signature of the provided method
*/
private String getSignature(MethodDescription methodDescription) {
String methodName = methodDescription.getName();
String parameters =
methodDescription.getParameters().asTypeList().stream()
.map(type -> type.asErasure().getTypeName())
.collect(Collectors.joining(","));
return methodName + "(" + parameters + ")";
}

/**
* Overwrite the hook for a specific method of the provided class.
*
* @param declaringClass the class containing the method
* @param methodSignature the method signature to be hooked
* @param newHook the hook for the method
*/
private void setHook(Class<?> declaringClass, String methodSignature, MethodHook newHook) {
hooks
.asMap()
.computeIfAbsent(declaringClass, (v) -> new HookedMethods())
.putMethod(methodSignature, newHook);
}

/**
* Removes the hook for the specific method of the provided class
*
* @param declaringClass the class containing the method
* @param methodSignature the method, whose hook should be removed
*/
private void removeHook(Class<?> declaringClass, String methodSignature) {
HookedMethods hookedMethods = hooks.getIfPresent(declaringClass);
if (Objects.nonNull(hookedMethods)) {
hookedMethods.removeMethod(methodSignature);
if (hookedMethods.noActiveHooks()) hooks.invalidate(declaringClass);
}
}

/**
* Returns the hook for the specific method of the provided class.
*
* @param clazz the class containing the method
* @param methodSignature the method, which might be hooked
* @return the hook of the method, if existing
*/
private Optional<MethodHook> getCurrentHook(Class<?> clazz, String methodSignature) {
HookedMethods hookedMethods = hooks.getIfPresent(clazz);
return Optional.ofNullable(hookedMethods)
.map(methods -> methods.getActiveHook(methodSignature));
}
}

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,10 @@
import java.util.Set;
import net.bytebuddy.description.method.MethodDescription;

/**
* Stores the hook configuration of all methods for a specific class. Currently, there is no complex
* hook configuration, thus we only use booleans.
*/
public class ClassHookConfiguration {

/** Set of methods and their hook configuration. Currently, just true. */
Expand All @@ -15,14 +19,25 @@ public ClassHookConfiguration() {
this.hookConfigurations = new HashMap<>();
}

/**
* @return the configuration as map
*/
public Map<MethodDescription, Boolean> asMap() {
return hookConfigurations;
}

/**
* @return the set of all methods within this configuration
*/
public Set<MethodDescription> getMethods() {
return hookConfigurations.keySet();
}

/**
* Stores a new hook configuration.
*
* @param method the method, which should be put into the configurations
*/
public void putHookConfiguration(MethodDescription method) {
hookConfigurations.put(method, true);
}
Expand Down
Loading

0 comments on commit 4f8f091

Please sign in to comment.