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,64 @@
+/*
+ * 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 net.bytebuddy.pool.TypePool;
+
+/**
+ * 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<EnhancementState> enhancementState = new ThreadLocal<>();
+
+	@Override
+	public TypePool.Resolution find(final String name) {
+		final EnhancementState enhancementState = getEnhancementState();
+		if ( enhancementState != null && enhancementState.getClassName().equals( name ) ) {
+			return enhancementState.getTypePoolResolution();
+		}
+		return super.find( name );
+	}
+
+	EnhancementState getEnhancementState() {
+		return enhancementState.get();
+	}
+
+	void setEnhancementState(EnhancementState state) {
+		enhancementState.set( state );
+	}
+
+	void removeEnhancementState() {
+		enhancementState.remove();
+	}
+
+	static final class EnhancementState {
+		private final String className;
+		private final ClassFileLocator.Resolution classFileResolution;
+		private TypePool.Resolution typePoolResolution;
+
+		public EnhancementState(String className, ClassFileLocator.Resolution classFileResolution) {
+			this.className = className;
+			this.classFileResolution = classFileResolution;
+		}
+
+		public String getClassName() {
+			return className;
+		}
+
+		public ClassFileLocator.Resolution getClassFileResolution() {
+			return classFileResolution;
+		}
+
+		public TypePool.Resolution getTypePoolResolution() {
+			return typePoolResolution;
+		}
+
+		public void setTypePoolResolution(TypePool.Resolution typePoolResolution) {
+			this.typePoolResolution = typePoolResolution;
+		}
+	}
+}

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

@@ -0,0 +1,38 @@
+/*
+ * 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;
+
+/**
+ * A delegating ClassFileLocator that allows overriding the resolution for the class that is currently being enhanced.
+ */
+final class EnhancerClassFileLocator implements ClassFileLocator {
+
+	private final EnhancerCacheProvider cacheProvider;
+	private final ClassFileLocator delegate;
+
+	public EnhancerClassFileLocator(EnhancerCacheProvider cacheProvider, ClassFileLocator delegate) {
+		this.cacheProvider = cacheProvider;
+		this.delegate = delegate;
+	}
+
+	@Override
+	public Resolution locate(final String name) throws IOException {
+		final EnhancerCacheProvider.EnhancementState enhancementState = cacheProvider.getEnhancementState();
+		if ( enhancementState != null && enhancementState.getClassName().equals( name ) ) {
+			return enhancementState.getClassFileResolution();
+		}
+		return delegate.locate( name );
+	}
+
+	@Override
+	public void close() throws IOException {
+		delegate.close();
+	}
+
+}

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

@@ -23,9 +23,8 @@ public interface EnhancerClassLocator extends TypePool {
 	 * This can optionally be used to remove an explicit mapping when it's no longer
 	 * essential to retain it.
 	 * The underlying implementation might ignore the operation.
-	 * @param className
 	 */
-	void deregisterClassNameAndBytes(String className);
+	void deregisterClassNameAndBytes();
 
 	/**
 	 * @return the underlying {@link ClassFileLocator}

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

@@ -140,7 +140,7 @@ public byte[] enhance(String className, byte[] originalBytes) throws Enhancement
 			throw new EnhancementException( "Failed to enhance class " + className, e );
 		}
 		finally {
-			typePool.deregisterClassNameAndBytes( safeClassName );
+			typePool.deregisterClassNameAndBytes();
 		}
 	}
 
@@ -158,7 +158,7 @@ public void discoverTypes(String className, byte[] originalBytes) {
 			throw new EnhancementException( "Failed to discover types for class " + className, e );
 		}
 		finally {
-			typePool.deregisterClassNameAndBytes( className );
+			typePool.deregisterClassNameAndBytes();
 		}
 	}
 

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

@@ -4,23 +4,21 @@
  */
 package org.hibernate.bytecode.enhance.internal.bytebuddy;
 
-import java.util.Objects;
-import java.util.concurrent.ConcurrentHashMap;
-
 import net.bytebuddy.dynamic.ClassFileLocator;
 import net.bytebuddy.pool.TypePool;
 
+import java.util.Objects;
+
 /**
  * A TypePool suitable for loading user's classes,
  * potentially in parallel operations.
  */
 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 +60,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,44 +70,35 @@ 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( cacheProvider, 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 ) ) );
+		final EnhancerCacheProvider.EnhancementState currentEnhancementState = poolCache.getEnhancementState();
+		if ( currentEnhancementState != null ) {
+			throw new IllegalStateException( "Re-entrant enhancement is not supported: " + className );
+		}
+		final EnhancerCacheProvider.EnhancementState state = new EnhancerCacheProvider.EnhancementState(
+				className,
+				new ClassFileLocator.Resolution.Explicit( Objects.requireNonNull( bytes ) )
+		);
+		// Set the state first because the ClassFileLocator needs this in the doDescribe() call below
+		poolCache.setEnhancementState( state );
+		state.setTypePoolResolution( doDescribe( className ) );
 	}
 
 	@Override
-	public void deregisterClassNameAndBytes(final String className) {
-		locator.remove( className );
+	public void deregisterClassNameAndBytes() {
+		poolCache.removeEnhancementState();
 	}
 
 	@Override
 	public ClassFileLocator asClassFileLocator() {
 		return locator;
 	}
-
 }

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