HHH-19230 Avoid class loader leak in enhancement and improve bytebuddy type caching efficiency

Author: beikov

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/EnhancerCacheProvider.java

@@ -0,0 +1,41 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.bytecode.enhance.internal.bytebuddy;
+
+import net.bytebuddy.pool.TypePool;
+
+import java.util.AbstractMap;
+import java.util.Map;
+
+/**
+ * A simple cache provider that allows overriding the resolution for the class that is currently being enhanced.
+ */
+final class EnhancerCacheProvider extends TypePool.CacheProvider.Simple {
+
+	private final ThreadLocal<Map.Entry<String, TypePool.Resolution>> enhancedTypeResolution = new ThreadLocal<>();
+
+	@Override
+	public TypePool.Resolution find(final String name) {
+		final Map.Entry<String, TypePool.Resolution> entry = enhancedTypeResolution.get();
+		if ( entry != null && entry.getKey().equals( name ) ) {
+			return entry.getValue();
+		}
+		return super.find( name );
+	}
+
+	@Override
+	public void clear() {
+		super.clear();
+		enhancedTypeResolution.remove();
+	}
+
+	public void registerEnhancedType(final String className, TypePool.Resolution resolution) {
+		enhancedTypeResolution.set( new AbstractMap.SimpleEntry<>( className, resolution ) );
+	}
+
+	public void deregisterEnhancedType(final String name) {
+		enhancedTypeResolution.remove();
+	}
+}

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/EnhancerClassFileLocator.java

@@ -0,0 +1,52 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.bytecode.enhance.internal.bytebuddy;
+
+import net.bytebuddy.dynamic.ClassFileLocator;
+
+import java.io.IOException;
+import java.util.AbstractMap;
+import java.util.Map;
+
+/**
+ * A delegating ClassFileLocator that allows overriding the resolution for the class that is currently being enhanced.
+ */
+final class EnhancerClassFileLocator implements ClassFileLocator {
+
+	private final ThreadLocal<Map.Entry<String, ClassFileLocator.Resolution>> enhancedTypeResolution = new ThreadLocal<>();
+	private final ClassFileLocator delegate;
+
+	public EnhancerClassFileLocator(ClassFileLocator delegate) {
+		this.delegate = delegate;
+	}
+
+	@Override
+	public Resolution locate(final String name) throws IOException {
+		final Map.Entry<String, ClassFileLocator.Resolution> entry = enhancedTypeResolution.get();
+		if ( entry != null && entry.getKey().equals( name ) ) {
+			return entry.getValue();
+		}
+		return delegate.locate( name );
+	}
+
+	@Override
+	public void close() throws IOException {
+		try {
+			delegate.close();
+		}
+		finally {
+			enhancedTypeResolution.remove();
+		}
+	}
+
+	public void registerEnhancedType(final String className, ClassFileLocator.Resolution resolution) {
+		enhancedTypeResolution.set( new AbstractMap.SimpleEntry<>( className, resolution ) );
+	}
+
+	public void deregisterEnhancedType(final String name) {
+		enhancedTypeResolution.remove();
+	}
+
+}

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/ModelTypePool.java

@@ -7,6 +7,7 @@
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 
+import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.ClassFileLocator;
 import net.bytebuddy.pool.TypePool;
 
@@ -16,11 +17,10 @@
  */
 public class ModelTypePool extends TypePool.Default implements EnhancerClassLocator {
 
-	private final ConcurrentHashMap<String, Resolution> resolutions = new ConcurrentHashMap<>();
-	private final OverridingClassFileLocator locator;
-	private final SafeCacheProvider poolCache;
+	private final EnhancerClassFileLocator locator;
+	private final EnhancerCacheProvider poolCache;
 
-	private ModelTypePool(SafeCacheProvider cacheProvider, OverridingClassFileLocator classFileLocator, CoreTypePool parent) {
+	private ModelTypePool(EnhancerCacheProvider cacheProvider, EnhancerClassFileLocator classFileLocator, CoreTypePool parent) {
 		super( cacheProvider, classFileLocator, ReaderMode.FAST, parent );
 		this.poolCache = cacheProvider;
 		this.locator = classFileLocator;
@@ -62,7 +62,7 @@ public static EnhancerClassLocator buildModelTypePool(ClassFileLocator classFile
 	 * @return
 	 */
 	public static EnhancerClassLocator buildModelTypePool(ClassFileLocator classFileLocator, CoreTypePool coreTypePool) {
-		return buildModelTypePool( classFileLocator, coreTypePool, new SafeCacheProvider() );
+		return buildModelTypePool( classFileLocator, coreTypePool, new EnhancerCacheProvider() );
 	}
 
 	/**
@@ -72,39 +72,23 @@ public static EnhancerClassLocator buildModelTypePool(ClassFileLocator classFile
 	 * @param cacheProvider
 	 * @return
 	 */
-	public static EnhancerClassLocator buildModelTypePool(ClassFileLocator classFileLocator, CoreTypePool coreTypePool, SafeCacheProvider cacheProvider) {
+	static EnhancerClassLocator buildModelTypePool(ClassFileLocator classFileLocator, CoreTypePool coreTypePool, EnhancerCacheProvider cacheProvider) {
 		Objects.requireNonNull( classFileLocator );
 		Objects.requireNonNull( coreTypePool );
 		Objects.requireNonNull( cacheProvider );
-		return new ModelTypePool( cacheProvider, new OverridingClassFileLocator( classFileLocator ), coreTypePool );
-	}
-
-	@Override
-	protected Resolution doDescribe(final String name) {
-		final Resolution resolution = resolutions.get( name );
-		if ( resolution != null ) {
-			return resolution;
-		}
-		else {
-			return resolutions.computeIfAbsent( name, super::doDescribe );
-		}
+		return new ModelTypePool( cacheProvider, new EnhancerClassFileLocator( classFileLocator ), coreTypePool );
 	}
 
 	@Override
 	public void registerClassNameAndBytes(final String className, final byte[] bytes) {
-		//Very important: ensure the registered override is actually effective in case this class
-		//was already resolved in the recent past; this could have happened for example as a side effect
-		//of symbol resolution during enhancement of a different class, or very simply when attempting
-		//to re-enhanced the same class - which happens frequently in WildFly because of the class transformers
-		//being triggered concurrently by multiple parallel deployments.
-		resolutions.remove( className );
-		poolCache.remove( className );
-		locator.put( className, new ClassFileLocator.Resolution.Explicit( Objects.requireNonNull( bytes ) ) );
+		locator.registerEnhancedType( className, new ClassFileLocator.Resolution.Explicit( Objects.requireNonNull( bytes ) ) );
+		poolCache.registerEnhancedType( className, doDescribe( className ) );
 	}
 
 	@Override
 	public void deregisterClassNameAndBytes(final String className) {
-		locator.remove( className );
+		poolCache.deregisterEnhancedType( className );
+		locator.deregisterEnhancedType( className );
 	}
 
 	@Override

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/OverridingClassFileLocator.java

@@ -60,9 +60,6 @@ public byte[] transform(
 		catch (final Exception e) {
 			throw new TransformerException( "Error performing enhancement of " + className, e );
 		}
-		finally {
-			bytecodeProvider.resetCaches();
-		}
 	}
 
 	@Override