From 443fc2653ce1b6388943b7e9248b47c0a5cb9bea Mon Sep 17 00:00:00 2001
From: Brian Laub <blaub@palantir.com>
Date: Wed, 13 Nov 2024 10:11:05 -0500
Subject: [PATCH] pr feedback

---
 .../com/palantir/dialogue/core/ChannelState.java  | 15 +++++----------
 1 file changed, 5 insertions(+), 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
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<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"));
     }
 }