Change approaches - move cache into the marshaller instead of SDKFields

Author: alextwoods

core/protocols/aws-json-protocol/src/main/java/software/amazon/awssdk/protocols/json/internal/marshall/JsonProtocolMarshaller.java

@@ -30,6 +30,7 @@
 import java.util.EnumMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.core.SdkField;
@@ -65,6 +66,12 @@ public class JsonProtocolMarshaller implements ProtocolMarshaller<SdkHttpFullReq
 
     private static final JsonMarshallerRegistry MARSHALLER_REGISTRY = createMarshallerRegistry();
 
+    // Caches the resolved marshaller for non-PAYLOAD fields, keyed by SdkField identity.
+    // SdkField instances are static final per generated model class, so identity-based lookup is correct.
+    // ConcurrentHashMap is used for thread safety; the one-time put per SdkField is negligible.
+    private static final ConcurrentHashMap<SdkField<?>, JsonMarshaller<Object>> MARSHALLER_CACHE =
+        new ConcurrentHashMap<>();
+
     private final URI endpoint;
     private final StructuredJsonGenerator jsonGenerator;
     private final SdkHttpFullRequest.Builder request;
@@ -412,11 +419,8 @@ private void marshallFieldViaRegistry(SdkField<?> field, Object val) {
                                .marshall(val, marshallerContext, field.locationName(), (SdkField<Object>) field);
             return;
         }
-        JsonMarshaller<Object> marshaller = field.cachedMarshaller(MARSHALLER_REGISTRY);
-        if (marshaller == null) {
-            marshaller = MARSHALLER_REGISTRY.getMarshaller(field.location(), field.marshallingType(), val);
-            field.cacheMarshaller(MARSHALLER_REGISTRY, marshaller);
-        }
+        JsonMarshaller<Object> marshaller = MARSHALLER_CACHE.computeIfAbsent(field,
+            f -> MARSHALLER_REGISTRY.getMarshaller(f.location(), f.marshallingType(), val));
         marshaller.marshall(val, marshallerContext, field.locationName(), (SdkField<Object>) field);
     }
 

core/protocols/aws-json-protocol/src/test/java/software/amazon/awssdk/protocols/json/internal/marshall/CachedNonPayloadMarshallingTest.java

@@ -77,36 +77,21 @@ void header_string_secondCall_usesCachedMarshaller() {
         // Use the SAME SdkField instance for both calls so the cache is shared
         SdkField<String> field = headerField("x-custom-header", obj -> "headerValue");
 
-        // First call — populates the cache
+        // First call — populates the internal marshaller cache
         SdkPojo pojo1 = new SimplePojo(field);
         SdkHttpFullRequest result1 = createMarshaller().marshall(pojo1);
 
-        // After first marshalling, the cache should be populated on the SdkField.
-        // We can't access the exact registry key, but we can verify the field has
-        // a non-null cached marshaller by checking that a second marshalling produces
-        // identical output.
-
         // Second call — should use cached marshaller
         SdkPojo pojo2 = new SimplePojo(field);
         SdkHttpFullRequest result2 = createMarshaller().marshall(pojo2);
 
-        // Both calls produce identical header output
+        // Both calls produce identical header output, confirming the cached path works
         assertThat(result1.firstMatchingHeader("x-custom-header"))
             .isPresent()
             .hasValue("headerValue");
         assertThat(result2.firstMatchingHeader("x-custom-header"))
             .isPresent()
             .hasValue("headerValue");
-
-        // Verify the cache was populated: the field should have a non-null cached
-        // marshaller for at least one registry key. Since we can't access the private
-        // MARSHALLER_REGISTRY, we verify indirectly: the field's cachedMarshaller
-        // with a dummy key returns null (different key), but the fact that both calls
-        // succeeded with identical output confirms the cached path works.
-        Object cachedWithDifferentKey = field.cachedMarshaller(new Object());
-        assertThat(cachedWithDifferentKey)
-            .as("Different registry key should return null")
-            .isNull();
     }
 
     // ---- QUERY_PARAM tests ----

core/sdk-core/src/main/java/software/amazon/awssdk/core/SdkField.java

@@ -50,15 +50,6 @@ public final class SdkField<TypeT> {
     private final Map<TraitType, Trait> l1Traits;
     private final Map<Class<? extends Trait>, Trait> l2Traits;
 
-    // Single-slot marshaller cache. Two volatile fields are used instead of an AtomicReference to an immutable
-    // holder to avoid per-SdkField object allocation. The read in cachedMarshaller() is not atomic across both
-    // fields: between reading the key and reading the marshaller, another thread could overwrite both. This is
-    // safe because (1) in practice there is only one registry per protocol, so all threads converge to the same
-    // marshaller, and (2) the worst case with multiple registries is a benign cache miss or a single call using
-    // a marshaller from a different registry, which self-corrects on the next call.
-    private volatile Object cachedMarshaller;
-    private volatile Object cachedMarshallerRegistryKey;
-
     private SdkField(Builder<TypeT> builder) {
         this.memberName = builder.memberName;
         this.marshallingType = builder.marshallingType;
@@ -262,33 +253,6 @@ public boolean containsTrait(Class<? extends Trait> clzz, TraitType type) {
         return getTrait(clzz, type) != null;
     }
 
-    /**
-     * Returns the cached marshaller for the given registry key, or null if not cached.
-     * Uses reference identity ({@code ==}) for the registry key comparison.
-     *
-     * @param registryKey The registry key to match against the cached key.
-     * @param <T> The type of the cached marshaller.
-     * @return The cached marshaller if the registry key matches, or null.
-     */
-    @SuppressWarnings("unchecked")
-    public <T> T cachedMarshaller(Object registryKey) {
-        if (cachedMarshallerRegistryKey == registryKey) {
-            return (T) cachedMarshaller;
-        }
-        return null;
-    }
-
-    /**
-     * Caches the resolved marshaller for the given registry key.
-     *
-     * @param registryKey The registry key to associate with the cached marshaller.
-     * @param marshaller The marshaller instance to cache.
-     */
-    public void cacheMarshaller(Object registryKey, Object marshaller) {
-        this.cachedMarshaller = marshaller;
-        this.cachedMarshallerRegistryKey = registryKey;
-    }
-
     /**
      * Retrieves the current value of 'this' field from the given POJO. Uses the getter passed into the {@link Builder}.
      *