Use Collections.unmodifiableMap

Author: SiyaoIsHiding

core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java

@@ -19,10 +19,11 @@
 
 import com.datastax.oss.driver.api.core.cql.Statement;
 import com.datastax.oss.driver.api.core.session.Request;
-import com.datastax.oss.protocol.internal.util.collection.NullAllowingImmutableMap;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 
 /**
@@ -68,26 +69,14 @@ default String getCustomPayloadKey() {
   default Statement<?> getDecoratedStatement(
       @NonNull Statement<?> statement, @NonNull String requestId) {
 
-    // local copy
-    Map<String, ByteBuffer> existing = statement.getCustomPayload();
+    Map<String, ByteBuffer> existing = new HashMap<>(statement.getCustomPayload());
     String key = getCustomPayloadKey();
 
-    NullAllowingImmutableMap.Builder<String, ByteBuffer> builder =
-        NullAllowingImmutableMap.builder();
+    // Add or overwrite
+    existing.put(key, ByteBuffer.wrap(requestId.getBytes(StandardCharsets.UTF_8)));
 
-    // Copy all existing entries except the one we’re overwriting
-    if (existing != null && !existing.isEmpty()) {
-      for (Map.Entry<String, ByteBuffer> e : existing.entrySet()) {
-        if (!e.getKey().equals(key)) {
-          builder.put(e.getKey(), e.getValue());
-        }
-      }
-    }
+    Map<String, ByteBuffer> unmodifiableMap = Collections.unmodifiableMap(existing);
 
-    // Put (or overwrite) the target key
-    builder.put(key, ByteBuffer.wrap(requestId.getBytes(StandardCharsets.UTF_8)));
-
-    Map<String, ByteBuffer> newPayload = builder.build();
-    return statement.setCustomPayload(newPayload);
+    return statement.setCustomPayload(unmodifiableMap);
   }
 }

integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java

@@ -17,12 +17,14 @@
  */
 package com.datastax.oss.driver.core.tracker;
 
+import static com.datastax.oss.driver.Assertions.assertThatStage;
 import static org.assertj.core.api.Assertions.assertThat;
 
 import com.datastax.oss.driver.api.core.CqlSession;
 import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
 import com.datastax.oss.driver.api.core.config.DriverConfigLoader;
 import com.datastax.oss.driver.api.core.cql.ResultSet;
+import com.datastax.oss.driver.api.core.cql.SimpleStatement;
 import com.datastax.oss.driver.api.core.cql.Statement;
 import com.datastax.oss.driver.api.core.session.Request;
 import com.datastax.oss.driver.api.core.tracker.RequestIdGenerator;
@@ -122,4 +124,23 @@ public void should_not_write_id_to_custom_payload_when_key_is_not_set() {
       assertThat(rs.getExecutionInfo().getRequest().getCustomPayload().get("trace_key")).isNull();
     }
   }
+
+  @Test
+  public void should_succeed_with_null_value_in_custom_payload() {
+    DriverConfigLoader loader =
+        SessionUtils.configLoaderBuilder()
+            .withString(
+                DefaultDriverOption.REQUEST_ID_GENERATOR_CLASS, "W3CContextRequestIdGenerator")
+            .build();
+    try (CqlSession session = SessionUtils.newSession(ccmRule, loader)) {
+      String query = "SELECT * FROM system.local";
+      Map<String, ByteBuffer> customPayload =
+          new NullAllowingImmutableMap.Builder<String, ByteBuffer>(1).put("my_key", null).build();
+      SimpleStatement statement =
+          SimpleStatement.newInstance(query).setCustomPayload(customPayload);
+      assertThatStage(session.executeAsync(statement)).isSuccess();
+      ResultSet rs = session.execute(query);
+      assertThat(rs.getExecutionInfo().getRequest().getCustomPayload().get("trace_key")).isNull();
+    }
+  }
 }