diff --git a/dialogue-core/src/main/java/com/palantir/dialogue/core/ChannelState.java b/dialogue-core/src/main/java/com/palantir/dialogue/core/ChannelState.java index 5880f1591..a7c8caf46 100644 --- a/dialogue-core/src/main/java/com/palantir/dialogue/core/ChannelState.java +++ b/dialogue-core/src/main/java/com/palantir/dialogue/core/ChannelState.java @@ -26,11 +26,11 @@ static final class Key { private final Class valueClass; private final Supplier factory; - T cast(final Object value) { + private T cast(final Object value) { return valueClass.cast(value); } - Supplier getFactory() { + private Supplier getFactory() { return factory; } @@ -44,13 +44,8 @@ Supplier getFactory() { private final Map, Object> state = new HashMap<>(); T getState(Key key) { - if (state.containsKey(key)) { - return key.cast(state.get(key)); - } else { - T value = key.getFactory().get(); - Preconditions.checkNotNull(value, "state factory cannot produce a null value"); - state.put(key, value); - return value; - } + return key.cast(Preconditions.checkNotNull( + state.computeIfAbsent(key, keyValue -> keyValue.getFactory().get()), + "state factory cannot produce a null value")); } }