From 1f410799de496b6d86f381bfc7fd4b37f03f3eb3 Mon Sep 17 00:00:00 2001
From: Sanne Grinovero <sanne@hibernate.org>
Date: Thu, 16 May 2024 17:38:14 +0100
Subject: [PATCH] Potential workaround for concurrent enhancement race
 conditions

---
 .../deployment/HibernateEntityEnhancer.java   | 45 ++++++++++++++++---
 1 file changed, 38 insertions(+), 7 deletions(-)

diff --git a/extensions/hibernate-orm/deployment/src/main/java/io/quarkus/hibernate/orm/deployment/HibernateEntityEnhancer.java b/extensions/hibernate-orm/deployment/src/main/java/io/quarkus/hibernate/orm/deployment/HibernateEntityEnhancer.java
index 70a00a31ac531..7a352605ab39b 100644
--- a/extensions/hibernate-orm/deployment/src/main/java/io/quarkus/hibernate/orm/deployment/HibernateEntityEnhancer.java
+++ b/extensions/hibernate-orm/deployment/src/main/java/io/quarkus/hibernate/orm/deployment/HibernateEntityEnhancer.java
@@ -17,6 +17,8 @@
 import io.quarkus.hibernate.orm.deployment.integration.QuarkusClassFileLocator;
 import io.quarkus.hibernate.orm.deployment.integration.QuarkusEnhancementContext;
 import net.bytebuddy.ClassFileVersion;
+import net.bytebuddy.dynamic.ClassFileLocator;
+import net.bytebuddy.pool.TypePool;
 
 /**
  * Used to transform bytecode by registering to
@@ -94,16 +96,11 @@ private static class EnhancerHolder {
         private volatile Enhancer actualEnhancer;
 
         public Enhancer getEnhancer() {
-            //Lazily initialized for multiple reasons:
-            //1)it's expensive: try to skip it if we can; this is actually not unlikely to happen as the transformation is cacheable.
-            //2)We want the Advice loaders of the Hibernate ORM implementation to be initialized within the scope in which we
-            //have the transformation classloader installed in the thread's context.
+            //Lazily initialized as it's expensive and might not be necessary: these transformations are cacheable.
             if (actualEnhancer == null) {
                 synchronized (this) {
                     if (actualEnhancer == null) {
-                        EnhancerClassLocator modelPool = ModelTypePool.buildModelTypePool(QuarkusClassFileLocator.INSTANCE,
-                                CORE_POOL);
-                        actualEnhancer = PROVIDER.getEnhancer(QuarkusEnhancementContext.INSTANCE, modelPool);
+                        actualEnhancer = PROVIDER.getEnhancer(QuarkusEnhancementContext.INSTANCE, new ThreadsafeLocator());
                     }
                 }
             }
@@ -111,4 +108,38 @@ public Enhancer getEnhancer() {
         }
     }
 
+    private static final class ThreadsafeLocator implements EnhancerClassLocator {
+
+        final CacheProvider sharedCache = TypePool.CacheProvider.Simple.withObjectType();
+
+        final ThreadLocal<EnhancerClassLocator> localLocator = ThreadLocal
+                .withInitial(() -> ModelTypePool.buildModelTypePool(QuarkusClassFileLocator.INSTANCE,
+                        CORE_POOL, sharedCache));
+
+        @Override
+        public void registerClassNameAndBytes(String s, byte[] bytes) {
+            localLocator.get().registerClassNameAndBytes(s, bytes);
+        }
+
+        @Override
+        public void deregisterClassNameAndBytes(String s) {
+            localLocator.get().deregisterClassNameAndBytes(s);
+        }
+
+        @Override
+        public ClassFileLocator asClassFileLocator() {
+            return localLocator.get().asClassFileLocator();
+        }
+
+        @Override
+        public Resolution describe(String s) {
+            return localLocator.get().describe(s);
+        }
+
+        @Override
+        public void clear() {
+            //not useful
+        }
+    }
+
 }