pr feedback

palantir · Nov 13, 2024 · 443fc26 · 443fc26
1 parent 52aeec6
commit 443fc26
Showing 1 changed file with 5 additions and 10 deletions.
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
@@ -26,11 +26,11 @@ static final class Key<T> {
         private final Class<T> valueClass;
         private final Supplier<T> factory;
 
-        T cast(final Object value) {
+        private T cast(final Object value) {
             return valueClass.cast(value);
         }
 
-        Supplier<T> getFactory() {
+        private Supplier<T> getFactory() {
             return factory;
         }
 
@@ -44,13 +44,8 @@ Supplier<T> getFactory() {
     private final Map<Key<?>, Object> state = new HashMap<>();
 
     <T> T getState(Key<T> 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"));
     }
 }