Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Look up helper class bytes when they are needed #7839

Merged
merged 2 commits into from
Feb 17, 2023
Merged

Look up helper class bytes when they are needed #7839

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9d5d729
Look up helper class bytes when they are needed
laurit Feb 16, 2023
9856863
use byte buddy abstraction
laurit Feb 17, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 38 additions & 18 deletions muzzle/src/main/java/io/opentelemetry/javaagent/tooling/HelperInjector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Supplier;
import javax.annotation.Nullable;
import net.bytebuddy.agent.builder.AgentBuilder.Transformer;
import net.bytebuddy.description.type.TypeDescription;
Expand Down Expand Up @@ -88,7 +89,7 @@ Class<?> inject(ClassLoader classLoader, String className) {
private final List<HelperResource> helperResources;
@Nullable private final ClassLoader helpersSource;
@Nullable private final Instrumentation instrumentation;
private final Map<String, byte[]> dynamicTypeMap = new LinkedHashMap<>();
private final Map<String, Supplier<byte[]>> dynamicTypeMap = new LinkedHashMap<>();

private final Cache<ClassLoader, Boolean> injectedClassLoaders = Cache.weak();
private final Cache<ClassLoader, Boolean> resourcesInjectedClassLoaders = Cache.weak();
Expand All @@ -108,7 +109,6 @@ public HelperInjector(
String requestingName,
List<String> helperClassNames,
List<HelperResource> helperResources,
// TODO can this be replaced with the context class loader?
ClassLoader helpersSource,
Instrumentation instrumentation) {
this.requestingName = requestingName;
Expand All @@ -120,7 +120,9 @@ public HelperInjector(
}

private HelperInjector(
String requestingName, Map<String, byte[]> helperMap, Instrumentation instrumentation) {
String requestingName,
Map<String, Supplier<byte[]>> helperMap,
Instrumentation instrumentation) {
this.requestingName = requestingName;

this.helperClassNames = helperMap.keySet();
Expand All @@ -135,9 +137,9 @@ public static HelperInjector forDynamicTypes(
String requestingName,
Collection<DynamicType.Unloaded<?>> helpers,
Instrumentation instrumentation) {
Map<String, byte[]> bytes = new HashMap<>(helpers.size());
Map<String, Supplier<byte[]>> bytes = new HashMap<>(helpers.size());
for (DynamicType.Unloaded<?> helper : helpers) {
bytes.put(helper.getTypeDescription().getName(), helper.getBytes());
bytes.put(helper.getTypeDescription().getName(), helper::getBytes);
}
return new HelperInjector(requestingName, bytes, instrumentation);
}
Expand All @@ -146,18 +148,27 @@ public static void setHelperInjectorListener(HelperInjectorListener listener) {
helperInjectorListener = listener;
}

private Map<String, byte[]> getHelperMap() throws IOException {
private Map<String, Supplier<byte[]>> getHelperMap() {
if (dynamicTypeMap.isEmpty()) {
Map<String, byte[]> classnameToBytes = new LinkedHashMap<>();

ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(helpersSource);
Map<String, Supplier<byte[]>> result = new LinkedHashMap<>();

for (String helperClassName : helperClassNames) {
byte[] classBytes = locator.locate(helperClassName).resolve();
classnameToBytes.put(helperClassName, classBytes);
result.put(
helperClassName,
() -> {
try (ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(helpersSource)) {
return locator.locate(helperClassName).resolve();
} catch (IOException exception) {
if (logger.isLoggable(SEVERE)) {
logger.log(
SEVERE, "Failed to read {0}", new Object[] {helperClassName}, exception);
}
throw new IllegalStateException("Failed to read " + helperClassName, exception);
}
});
}

return classnameToBytes;
return result;
} else {
return dynamicTypeMap;
}
Expand Down Expand Up @@ -248,10 +259,10 @@ private void injectHelperClasses(TypeDescription typeDescription, ClassLoader cl
new Object[] {cl, helperClassNames});
}

Map<String, byte[]> classnameToBytes = getHelperMap();
Map<String, Supplier<byte[]>> classnameToBytes = getHelperMap();
Map<String, HelperClassInjector> map =
helperInjectors.computeIfAbsent(cl, (unused) -> new ConcurrentHashMap<>());
for (Map.Entry<String, byte[]> entry : classnameToBytes.entrySet()) {
for (Map.Entry<String, Supplier<byte[]>> entry : classnameToBytes.entrySet()) {
// for boot loader we use a placeholder injector, we only need these classes to be
// in the injected classes map to later tell which of the classes are injected
HelperClassInjector injector =
Expand Down Expand Up @@ -280,9 +291,18 @@ private void injectHelperClasses(TypeDescription typeDescription, ClassLoader cl
});
}

private Map<String, Class<?>> injectBootstrapClassLoader(Map<String, byte[]> classnameToBytes)
private static Map<String, byte[]> resolve(Map<String, Supplier<byte[]>> classes) {
Map<String, byte[]> result = new LinkedHashMap<>();
for (Map.Entry<String, Supplier<byte[]>> entry : classes.entrySet()) {
result.put(entry.getKey(), entry.getValue().get());
}
return result;
}

private Map<String, Class<?>> injectBootstrapClassLoader(Map<String, Supplier<byte[]>> inject)
throws IOException {

Map<String, byte[]> classnameToBytes = resolve(inject);
if (helperInjectorListener != null) {
helperInjectorListener.onInjection(classnameToBytes);
}
Expand Down Expand Up @@ -358,16 +378,16 @@ public static Class<?> loadHelperClass(ClassLoader classLoader, String className
}

private static class HelperClassInjector {
private final byte[] bytes;
private final Supplier<byte[]> bytes;

HelperClassInjector(byte[] bytes) {
HelperClassInjector(Supplier<byte[]> bytes) {
this.bytes = bytes;
}

Class<?> inject(ClassLoader classLoader, String className) {
Map<String, Class<?>> result =
new ClassInjector.UsingReflection(classLoader)
.injectRaw(Collections.singletonMap(className, bytes));
.injectRaw(Collections.singletonMap(className, bytes.get()));
return result.get(className);
}
}
Expand Down