Fix classloading in ConstructorCache

Use the SDK standard ClassLoaderHelper#getClass to try and load the
given class from a ClassLoader.

Previously, we were only looking at the ClassLoader returned by
ClassLoaderHelper#contextClassLoader which either returns the Thread's
context ClassLoader, or the system ClassLoader if that isn't present.
This is an issue because it's possible that neither of these
ClassLoaders are correct; we should also be looking at the ClassLoader
that loaded the calling class.

As part of this change, the internal Map used in ConstructorCache has
been changed. Rather than doing a two-level mapping from String -> (Map
of ClassLoader to Class), we have a single level mapping from String ->
Class. Class is still referenced via a WeakReference so it is still free
to be GC'd at any point.

Author: dagnir

core/checksums/src/main/java/software/amazon/awssdk/checksums/internal/ConstructorCache.java

@@ -17,23 +17,25 @@
 
 import java.lang.ref.WeakReference;
 import java.lang.reflect.Constructor;
-import java.util.Collections;
 import java.util.Map;
 import java.util.Optional;
-import java.util.WeakHashMap;
 import java.util.concurrent.ConcurrentHashMap;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.utils.ClassLoaderHelper;
 import software.amazon.awssdk.utils.Logger;
 
 /**
- * A cache that stores classes and their constructors by class name and class loader.
+ * A cache that stores classes and exposes a method to retrieve its zero-argument constructor.
  * <p>
- * This cache uses weak references to both class loaders and classes, allowing them to be garbage collected
- * when no longer needed. It provides methods to retrieve the zero-argument constructor for a class,
- * based on the current thread's context class loader or the system class loader.
+ * This cache stores weak references to loaded classes, allowing them to be garbage collected at any point.
+ * <p>
+ * Classes are loaded by first attempting to load it via the thread context {@code ClassLoader} (or system {@code ClassLoader} if
+ * the calling thread does not have one). If that fails, will attempt using the {@code ClassLoader} that loaded this class, and
+ * finally, attempt using the class that loaded {@link ClassLoaderHelper}.
  * <p>
  * If a class or its zero-argument constructor cannot be found, an empty result is returned.
+ *
+ * @see ClassLoaderHelper#loadClass(String, Class[])
  */
 @SdkInternalApi
 public final class ConstructorCache {
@@ -43,36 +45,36 @@ public final class ConstructorCache {
      * Cache storing classes by class name and class loader.
      * Uses weak references to allow garbage collection when not needed.
      */
-    private final Map<String, Map<ClassLoader, Optional<WeakReference<Class<?>>>>> classesByClassName =
+    private final Map<String, Optional<WeakReference<Class<?>>>> classesByClassName =
         new ConcurrentHashMap<>();
 
     /**
      * Retrieve the class for the given class name from the context or system class loader.
      * Returns an empty result if the class is not found.
      */
     private Optional<Class<?>> getClass(String className) {
-        Map<ClassLoader, Optional<WeakReference<Class<?>>>> classesByClassLoader =
-            classesByClassName.computeIfAbsent(className, k -> Collections.synchronizedMap(new WeakHashMap<>()));
-
-        ClassLoader classLoader = ClassLoaderHelper.contextClassLoader();
-        Optional<WeakReference<Class<?>>> classRef = classesByClassLoader.computeIfAbsent(classLoader, k -> {
+        Optional<WeakReference<Class<?>>> classRef = classesByClassName.computeIfAbsent(className, k -> {
             try {
-                Class<?> clazz = classLoader.loadClass(className);
+                Class<?> clazz = ClassLoaderHelper.loadClass(k, false);
                 return Optional.of(new WeakReference<>(clazz));
             } catch (ClassNotFoundException e) {
                 return Optional.empty();
             }
         });
 
-        // if the WeakReference to the class has been garbage collected, remove it from the cache and try again
+        // Were we able to find this class?
         if (classRef.isPresent()) {
             Class<?> clazz = classRef.get().get();
+            // Class hasn't been GC'd
             if (clazz != null) {
                 return Optional.of(clazz);
             }
-            classesByClassLoader.remove(classLoader);
+            // if the WeakReference to the class has been garbage collected, it has been unloaded.
+            // Remove it from the cache and try a fresh load
+            classesByClassName.remove(className);
             return getClass(className);
         }
+
         return Optional.empty();
     }