KAFKA-20804: Reduce lock contention in ProducerMetadata#add

clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerMetadata.java

@@ -29,19 +29,19 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
 import java.util.OptionalInt;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 
 public class ProducerMetadata extends Metadata {
     // If a topic hasn't been accessed for this many milliseconds, it is removed from the cache.
     private final long metadataIdleMs;
 
     /* Topics with expiry time */
-    private final Map<String, Long> topics = new HashMap<>();
+    private final Map<String, Long> topics = new ConcurrentHashMap<>();
     private final Set<String> newTopics = new HashSet<>();
     private final Logger log;
     private final Time time;
@@ -70,11 +70,21 @@ public synchronized MetadataRequest.Builder newMetadataRequestBuilderForNewTopic
         return new MetadataRequest.Builder(new ArrayList<>(newTopics), true);
     }
 
-    public synchronized void add(String topic, long nowMs) {
+    public void add(String topic, long nowMs) {
         Objects.requireNonNull(topic, "topic cannot be null");
-        if (topics.put(topic, nowMs + metadataIdleMs) == null) {
-            newTopics.add(topic);
-            requestUpdateForNewTopics();
+        long expiryTime = nowMs + metadataIdleMs;
+        // replace() only writes if the topic is still present, so there's no window in which a topic
+        // concurrently evicted by retainTopic() could get silently re-inserted without newTopics bookkeeping.
+        if (topics.replace(topic, expiryTime) != null) {
+            return;
+        }
+        synchronized (this) {
+            // New (or concurrently-evicted) topic: topics.put() and newTopics.add() must happen atomically
+            // with retainTopic(), hence the shared lock.
+            if (topics.put(topic, expiryTime) == null) {
+                newTopics.add(topic);
+                requestUpdateForNewTopics();
+            }
         }
     }
 
@@ -124,15 +134,20 @@ public synchronized boolean retainTopic(String topic, boolean isInternal, long n
         Long expireMs = topics.get(topic);
         if (expireMs == null) {
             return false;
-        } else if (newTopics.contains(topic)) {
+        }
+        if (newTopics.contains(topic)) {
             return true;
-        } else if (expireMs <= nowMs) {
+        }
+        if (expireMs > nowMs) {
+            return true;
+        }
+        // Only remove if the expiry we read is still current: the lock-free refresh path in add() can
+        // race with this check, and a plain remove(topic) would drop a concurrently-refreshed entry.
+        if (topics.remove(topic, expireMs)) {
             log.debug("Removing unused topic {} from the metadata list, expiryMs {} now {}", topic, expireMs, nowMs);
-            topics.remove(topic);
             return false;
-        } else {
-            return true;
         }
+        return true;
     }
 
     /**

clients/src/test/java/org/apache/kafka/clients/producer/internals/ProducerMetadataTest.java

@@ -304,6 +304,34 @@ public void testRequestUpdateForTopic() {
         assertFalse(metadata.updateRequested());
     }
 
+    @Test
+    public void testRetainTopic() {
+        final String unknownTopic = "unknown-topic";
+        final String newTopic = "new-topic";
+        final String activeTopic = "active-topic";
+
+        // A topic that was never added isn't retained.
+        assertFalse(metadata.retainTopic(unknownTopic, false, 0));
+
+        // A "new" topic (added but not yet resolved by an update) is retained even past its
+        // nominal expiry, since it hasn't had a chance to be fetched yet.
+        metadata.add(newTopic, 0);
+        assertTrue(metadata.newTopics().contains(newTopic));
+        assertTrue(metadata.retainTopic(newTopic, false, METADATA_IDLE_MS * 10));
+
+        // Once resolved (no longer "new"), a topic that hasn't reached its idle expiry is retained.
+        metadata.add(activeTopic, 0);
+        metadata.updateWithCurrentRequestVersion(responseWithTopics(Set.of(newTopic, activeTopic)), true, 0);
+        assertFalse(metadata.newTopics().contains(activeTopic));
+        assertTrue(metadata.retainTopic(activeTopic, false, METADATA_IDLE_MS - 1));
+
+        // Past its idle expiry, the topic is evicted: retainTopic() returns false and actually
+        // removes the topic, so a subsequent call also returns false.
+        assertFalse(metadata.retainTopic(activeTopic, false, METADATA_IDLE_MS));
+        assertFalse(metadata.containsTopic(activeTopic));
+        assertFalse(metadata.retainTopic(activeTopic, false, METADATA_IDLE_MS));
+    }
+
     private MetadataResponse responseWithCurrentTopics() {
         return responseWithTopics(metadata.topics());
     }

jmh-benchmarks/src/main/java/org/apache/kafka/jmh/producer/ProducerMetadataAddBenchmark.java

@@ -0,0 +1,88 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kafka.jmh.producer;
+
+import org.apache.kafka.clients.producer.internals.ProducerMetadata;
+import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.internals.LogContext;
+
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.annotations.Warmup;
+
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Benchmarks {@link ProducerMetadata#add} refreshing an already-tracked topic's expiry. Every
+ * {@code send()} call does this refresh, so with many application threads sharing one producer,
+ * it's on the hot path and contends across threads for a topic that's virtually always already
+ * tracked. This benchmark simulates that: many threads concurrently refreshing a shared pool of
+ * already-registered topics.
+ *
+ * To measure the effect of a change to {@code add()}, run this benchmark before and after the
+ * change and compare the throughput, rather than hand-mirroring the old implementation here
+ * where it could drift out of sync with reality.
+ */
+@State(Scope.Benchmark)
+@Fork(value = 1)
+@Warmup(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
+@Measurement(iterations = 10, time = 1, timeUnit = TimeUnit.SECONDS)
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@Threads(8)
+public class ProducerMetadataAddBenchmark {
+
+    private static final int TOPIC_COUNT = 200;
+    private static final long METADATA_IDLE_MS = TimeUnit.MINUTES.toMillis(5);
+
+    private String[] topics;
+    private ProducerMetadata metadata;
+
+    @Setup(Level.Trial)
+    public void setup() {
+        topics = new String[TOPIC_COUNT];
+        for (int i = 0; i < TOPIC_COUNT; i++) {
+            topics[i] = "topic-" + i;
+        }
+
+        metadata = new ProducerMetadata(100L, 1000L, TimeUnit.MINUTES.toMillis(5), METADATA_IDLE_MS,
+            new LogContext(), new ClusterResourceListeners(), Time.SYSTEM);
+
+        long nowMs = Time.SYSTEM.milliseconds();
+        for (String topic : topics) {
+            metadata.add(topic, nowMs);
+        }
+    }
+
+    @Benchmark
+    public void addExistingTopic() {
+        String topic = topics[ThreadLocalRandom.current().nextInt(TOPIC_COUNT)];
+        metadata.add(topic, Time.SYSTEM.milliseconds());
+    }
+}