From 4407e07767cfeaa2e625253cebb4b14e1ee5288d Mon Sep 17 00:00:00 2001 From: Rafael Winterhalter Date: Fri, 1 Nov 2024 00:25:11 +0100 Subject: [PATCH] Avoid eager resolution of class files in auxiliary types. --- .../bytebuddy/agent/builder/AgentBuilder.java | 45 +++++++++++-------- 1 file changed, 26 insertions(+), 19 deletions(-) diff --git a/byte-buddy-dep/src/main/java/net/bytebuddy/agent/builder/AgentBuilder.java b/byte-buddy-dep/src/main/java/net/bytebuddy/agent/builder/AgentBuilder.java index d9d5d5b9de..731bf9ef55 100644 --- a/byte-buddy-dep/src/main/java/net/bytebuddy/agent/builder/AgentBuilder.java +++ b/byte-buddy-dep/src/main/java/net/bytebuddy/agent/builder/AgentBuilder.java @@ -3914,9 +3914,8 @@ public void register(DynamicType dynamicType, @MaybeNull ClassLoader classLoader, @MaybeNull ProtectionDomain protectionDomain, InjectionStrategy injectionStrategy) { - Map auxiliaryTypes = dynamicType.getAuxiliaryTypes(); - Map independentTypes = new LinkedHashMap(auxiliaryTypes); - for (TypeDescription auxiliaryType : auxiliaryTypes.keySet()) { + Set auxiliaryTypes = dynamicType.getAuxiliaryTypeDescriptions(), independentTypes = new LinkedHashSet(auxiliaryTypes); + for (TypeDescription auxiliaryType : auxiliaryTypes) { if (!auxiliaryType.getDeclaredAnnotations().isAnnotationPresent(AuxiliaryType.SignatureRelevant.class)) { independentTypes.remove(auxiliaryType); } @@ -3924,7 +3923,7 @@ public void register(DynamicType dynamicType, if (!independentTypes.isEmpty()) { ClassInjector classInjector = injectionStrategy.resolve(classLoader, protectionDomain); Map loadedTypeInitializers = dynamicType.getLoadedTypeInitializers(); - for (Map.Entry> entry : classInjector.inject(independentTypes).entrySet()) { + for (Map.Entry> entry : classInjector.inject(independentTypes, dynamicType).entrySet()) { loadedTypeInitializers.get(entry.getKey()).onLoad(entry.getValue()); } } @@ -4014,9 +4013,14 @@ protected static class InjectingInitializer implements LoadedTypeInitializer { private final TypeDescription instrumentedType; /** - * The auxiliary types mapped to their class file representation. + * The auxiliary types to inject. */ - private final Map rawAuxiliaryTypes; + private final Set auxiliaryTypes; + + /** + * The class file locator to use. + */ + private final ClassFileLocator classFileLocator; /** * The instrumented types and auxiliary types mapped to their loaded type initializers. @@ -4033,16 +4037,19 @@ protected static class InjectingInitializer implements LoadedTypeInitializer { * Creates a new injection initializer. * * @param instrumentedType The instrumented type. - * @param rawAuxiliaryTypes The auxiliary types mapped to their class file representation. + * @param rawAuxiliaryTypes The auxiliary types to inject. + * @param classFileLocator The class file locator to use. * @param loadedTypeInitializers The instrumented types and auxiliary types mapped to their loaded type initializers. * @param classInjector The class injector to use. */ protected InjectingInitializer(TypeDescription instrumentedType, - Map rawAuxiliaryTypes, + Set auxiliaryTypes, + ClassFileLocator classFileLocator, Map loadedTypeInitializers, ClassInjector classInjector) { this.instrumentedType = instrumentedType; - this.rawAuxiliaryTypes = rawAuxiliaryTypes; + this.auxiliaryTypes = auxiliaryTypes; + this.classFileLocator = classFileLocator; this.loadedTypeInitializers = loadedTypeInitializers; this.classInjector = classInjector; } @@ -4051,7 +4058,7 @@ protected InjectingInitializer(TypeDescription instrumentedType, * {@inheritDoc} */ public void onLoad(Class type) { - for (Map.Entry> auxiliary : classInjector.inject(rawAuxiliaryTypes).entrySet()) { + for (Map.Entry> auxiliary : classInjector.inject(auxiliaryTypes, classFileLocator).entrySet()) { loadedTypeInitializers.get(auxiliary.getKey()).onLoad(auxiliary.getValue()); } loadedTypeInitializers.get(instrumentedType).onLoad(type); @@ -4116,28 +4123,28 @@ public void register(DynamicType dynamicType, @MaybeNull ClassLoader classLoader, @MaybeNull ProtectionDomain protectionDomain, InjectionStrategy injectionStrategy) { - Map auxiliaryTypes = dynamicType.getAuxiliaryTypes(); + Set auxiliaryTypes = dynamicType.getAuxiliaryTypeDescriptions(); LoadedTypeInitializer loadedTypeInitializer; if (!auxiliaryTypes.isEmpty()) { TypeDescription instrumentedType = dynamicType.getTypeDescription(); ClassInjector classInjector = injectionStrategy.resolve(classLoader, protectionDomain); - Map independentTypes = new LinkedHashMap(auxiliaryTypes); - Map dependentTypes = new LinkedHashMap(auxiliaryTypes); - for (TypeDescription auxiliaryType : auxiliaryTypes.keySet()) { + Set independentTypes = new LinkedHashSet(auxiliaryTypes); + Set dependentTypes = new LinkedHashSet(auxiliaryTypes); + for (TypeDescription auxiliaryType : auxiliaryTypes) { (auxiliaryType.getDeclaredAnnotations().isAnnotationPresent(AuxiliaryType.SignatureRelevant.class) ? dependentTypes : independentTypes).remove(auxiliaryType); } Map loadedTypeInitializers = dynamicType.getLoadedTypeInitializers(); if (!independentTypes.isEmpty()) { - for (Map.Entry> entry : classInjector.inject(independentTypes).entrySet()) { + for (Map.Entry> entry : classInjector.inject(independentTypes, dynamicType).entrySet()) { loadedTypeInitializers.get(entry.getKey()).onLoad(entry.getValue()); } } Map lazyInitializers = new HashMap(loadedTypeInitializers); - loadedTypeInitializers.keySet().removeAll(independentTypes.keySet()); + loadedTypeInitializers.keySet().removeAll(independentTypes); loadedTypeInitializer = lazyInitializers.size() > 1 // there exist auxiliary types that need lazy loading - ? new Dispatcher.InjectingInitializer(instrumentedType, dependentTypes, lazyInitializers, classInjector) + ? new Dispatcher.InjectingInitializer(instrumentedType, dependentTypes, dynamicType, lazyInitializers, classInjector) : lazyInitializers.get(instrumentedType); } else { loadedTypeInitializer = dynamicType.getLoadedTypeInitializers().get(dynamicType.getTypeDescription()); @@ -4195,10 +4202,10 @@ public void register(DynamicType dynamicType, @MaybeNull ClassLoader classLoader, @MaybeNull ProtectionDomain protectionDomain, InjectionStrategy injectionStrategy) { - Map auxiliaryTypes = dynamicType.getAuxiliaryTypes(); + Set auxiliaryTypes = dynamicType.getAuxiliaryTypeDescriptions(); LoadedTypeInitializer loadedTypeInitializer = auxiliaryTypes.isEmpty() ? dynamicType.getLoadedTypeInitializers().get(dynamicType.getTypeDescription()) - : new Dispatcher.InjectingInitializer(dynamicType.getTypeDescription(), auxiliaryTypes, dynamicType.getLoadedTypeInitializers(), injectionStrategy.resolve(classLoader, protectionDomain)); + : new Dispatcher.InjectingInitializer(dynamicType.getTypeDescription(), auxiliaryTypes, dynamicType, dynamicType.getLoadedTypeInitializers(), injectionStrategy.resolve(classLoader, protectionDomain)); nexusAccessor.register(dynamicType.getTypeDescription().getName(), classLoader, identification, loadedTypeInitializer); } }