diff --git a/utils/src/main/java/software/amazon/awssdk/utils/AttributeMap.java b/utils/src/main/java/software/amazon/awssdk/utils/AttributeMap.java index 12eee34f0842..7c9f29a2d053 100644 --- a/utils/src/main/java/software/amazon/awssdk/utils/AttributeMap.java +++ b/utils/src/main/java/software/amazon/awssdk/utils/AttributeMap.java @@ -327,17 +327,17 @@ private void internalPut(Key key, Value value) { */ private Value internalComputeIfAbsent(Key key, Supplier> value) { checkCopyOnUpdate(); - return attributes.compute(key, (k, v) -> { - if (v == null || resolveValue(v) == null) { - Value newValue = value.get(); - Validate.notNull(newValue, "Supplied value must not be null."); - if (v != null) { - dependencyGraph.valueUpdated(v, newValue); - } - return newValue; + Value currentValue = attributes.get(key); + if (currentValue == null || resolveValue(currentValue) == null) { + Value newValue = value.get(); + Validate.notNull(newValue, "Supplied value must not be null."); + if (currentValue != null) { + dependencyGraph.valueUpdated(currentValue, newValue); } - return v; - }); + attributes.put(key, newValue); + return newValue; + } + return currentValue; } private void checkCopyOnUpdate() { diff --git a/utils/src/test/java/software/amazon/awssdk/utils/AttributeMapTest.java b/utils/src/test/java/software/amazon/awssdk/utils/AttributeMapTest.java index fc5a1c2a9d9b..2392996ddf2f 100644 --- a/utils/src/test/java/software/amazon/awssdk/utils/AttributeMapTest.java +++ b/utils/src/test/java/software/amazon/awssdk/utils/AttributeMapTest.java @@ -30,7 +30,6 @@ import java.util.concurrent.TimeUnit; import org.junit.Test; import org.mockito.Mockito; -import org.w3c.dom.Attr; public class AttributeMapTest { @@ -158,6 +157,20 @@ public void lazyAttributes_notReResolvedAfterToBuilderBuild() { verify(lazyRead, Mockito.times(1)).run(); } + @Test + public void lazyAttributes_canUpdateTheMap_andBeUpdatedWithPutLazyIfAbsent() { + AttributeMap.Builder map = AttributeMap.builder(); + map.putLazyIfAbsent(STRING_KEY, c -> { + // Force a modification to the underlying map. We wouldn't usually do this so explicitly, but + // this can happen internally within AttributeMap when resolving uncached lazy values, + // so it needs to be possible. + map.put(INTEGER_KEY, 0); + return "string"; + }); + map.putLazyIfAbsent(STRING_KEY, c -> "string"); // Force the value to be resolved within the putLazyIfAbsent + assertThat(map.get(STRING_KEY)).isEqualTo("string"); + } + @Test public void changesInBuilder_doNotAffectBuiltMap() { AttributeMap.Builder builder = mapBuilderWithLazyString();