diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java
index d54d94f47d7c..5cb3cb56d9e9 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java
@@ -264,11 +264,15 @@ public long add(WindowedValue<T> data) throws IOException {
         }
       }
 
-      Windmill.KeyedMessageBundle.Builder keyedOutput = productionMap.get(key);
-      if (keyedOutput == null) {
-        keyedOutput = Windmill.KeyedMessageBundle.newBuilder().setKey(key);
-        productionMap.put(key, keyedOutput);
-      }
+      Windmill.KeyedMessageBundle.Builder keyedOutput =
+          productionMap.computeIfAbsent(
+              key,
+              (k) -> {
+                Windmill.KeyedMessageBundle.Builder builder =
+                    Windmill.KeyedMessageBundle.newBuilder();
+                builder.setKey(k);
+                return builder;
+              });
 
       try {
         messageBuilder