From 7eacbfab3c2a784f8995a844690504d7cb125bfb Mon Sep 17 00:00:00 2001
From: Sanne Grinovero <sanne@hibernate.org>
Date: Thu, 4 May 2023 21:41:21 +0300
Subject: [PATCH] HHH-16545 PersistenceUtil.MetadataCache needs to be
 threadsafe

---
 .../internal/util/PersistenceUtilHelper.java  | 47 +++++++++++++------
 1 file changed, 33 insertions(+), 14 deletions(-)

diff --git a/hibernate-core/src/main/java/org/hibernate/jpa/internal/util/PersistenceUtilHelper.java b/hibernate-core/src/main/java/org/hibernate/jpa/internal/util/PersistenceUtilHelper.java
index 509d19ec97..5490ec9950 100644
--- a/hibernate-core/src/main/java/org/hibernate/jpa/internal/util/PersistenceUtilHelper.java
+++ b/hibernate-core/src/main/java/org/hibernate/jpa/internal/util/PersistenceUtilHelper.java
@@ -6,6 +6,7 @@
  */
 package org.hibernate.jpa.internal.util;
 
+import java.io.ObjectStreamException;
 import java.io.Serializable;
 import java.lang.reflect.Field;
 import java.lang.reflect.InvocationTargetException;
@@ -16,7 +17,7 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.WeakHashMap;
+
 import jakarta.persistence.spi.LoadState;
 
 import org.hibernate.HibernateException;
@@ -411,24 +412,42 @@ public final class PersistenceUtilHelper {
 	}
 
 	/**
-	 * Cache hierarchy and member resolution in a weak hash map
+	 * Cache hierarchy and member resolution, taking care to not leak
+	 * references to Class instances.
 	 */
-	//TODO not really thread-safe
-	public static class MetadataCache implements Serializable {
-		private transient Map<Class<?>, ClassMetadataCache> classCache = new WeakHashMap<>();
+	public static final class MetadataCache implements Serializable {
 
+		private final ClassValue<ClassMetadataCache> metadataCacheClassValue;
 
-		private void readObject(java.io.ObjectInputStream stream) {
-			classCache = new WeakHashMap<>();
+		public MetadataCache() {
+			this( new MetadataClassValue() );
 		}
 
-		ClassMetadataCache getClassMetadata(Class<?> clazz) {
-			ClassMetadataCache classMetadataCache = classCache.get( clazz );
-			if ( classMetadataCache == null ) {
-				classMetadataCache = new ClassMetadataCache( clazz );
-				classCache.put( clazz, classMetadataCache );
-			}
-			return classMetadataCache;
+		//To help with serialization: no need to serialize the actual metadataCacheClassValue field
+		private MetadataCache(ClassValue<ClassMetadataCache> metadataCacheClassValue) {
+			this.metadataCacheClassValue = metadataCacheClassValue;
+		}
+
+		Object writeReplace() throws ObjectStreamException {
+			//Writing a different instance which doesn't include the cache
+			return new MetadataCache(null);
+		}
+
+		private Object readResolve() throws ObjectStreamException {
+			//Ensure we do instantiate a new cache instance on deserialization
+			return new MetadataCache();
+		}
+
+		ClassMetadataCache getClassMetadata(final Class<?> clazz) {
+			return metadataCacheClassValue.get( clazz );
+		}
+
+	}
+
+	private static final class MetadataClassValue extends ClassValue<ClassMetadataCache> {
+		@Override
+		protected ClassMetadataCache computeValue(final Class type) {
+			return new ClassMetadataCache( type );
 		}
 	}