KAFKA-13022: Optimize ClientQuotasImage#describe

jmh-benchmarks/src/main/java/org/apache/kafka/jmh/metadata/ClientQuotasImageDescribeBenchmark.java

@@ -0,0 +1,99 @@
+/*
+ * 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.metadata;
+
+import org.apache.kafka.common.message.DescribeClientQuotasRequestData;
+import org.apache.kafka.common.quota.ClientQuotaEntity;
+import org.apache.kafka.common.requests.DescribeClientQuotasRequest;
+import org.apache.kafka.image.ClientQuotaImage;
+import org.apache.kafka.image.ClientQuotasImage;
+import org.apache.kafka.server.config.QuotaConfig;
+
+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.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+@State(Scope.Benchmark)
+@Fork(value = 1)
+@Warmup(iterations = 5)
+@Measurement(iterations = 15)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+public class ClientQuotasImageDescribeBenchmark {
+
+    @Param({"10", "100", "1000"})
+    private int eachEntityCount;
+
+    private ClientQuotasImage clientQuotasImage;
+
+    @Setup(Level.Trial)
+    public void setup() {
+        clientQuotasImage = createClientQuotasImage(eachEntityCount);
+    }
+
+    static ClientQuotasImage createClientQuotasImage(int eachEntityCount) {
+        Map<ClientQuotaEntity, ClientQuotaImage> entities = new HashMap<>();
+        ClientQuotaImage defaultImage = new ClientQuotaImage(Map.of(QuotaConfig.REQUEST_PERCENTAGE_OVERRIDE_CONFIG, 1.0));
+        for (int i = 0; i < eachEntityCount; i++) {
+            entities.put(new ClientQuotaEntity(Map.of(ClientQuotaEntity.USER, "user-" + i)), defaultImage);
+            entities.put(new ClientQuotaEntity(Map.of(ClientQuotaEntity.CLIENT_ID, "client-id-" + i)), defaultImage);
+            entities.put(new ClientQuotaEntity(Map.of(ClientQuotaEntity.IP, "ip-" + i)), defaultImage);
+        }
+        return new ClientQuotasImage(entities);
+    }
+
+    @Benchmark
+    public void describeSpecified() {
+        clientQuotasImage.describe(new DescribeClientQuotasRequestData()
+            .setComponents(List.of(new DescribeClientQuotasRequestData.ComponentData()
+                .setEntityType(ClientQuotaEntity.USER)
+                .setMatchType(DescribeClientQuotasRequest.MATCH_TYPE_SPECIFIED)
+                .setMatch(null))));
+    }
+
+    @Benchmark
+    public void describeDefault() {
+        clientQuotasImage.describe(new DescribeClientQuotasRequestData()
+            .setComponents(List.of(new DescribeClientQuotasRequestData.ComponentData()
+                .setEntityType(ClientQuotaEntity.USER)
+                .setMatchType(DescribeClientQuotasRequest.MATCH_TYPE_DEFAULT)
+                .setMatch(null))));
+    }
+
+    @Benchmark
+    public void describeExact() {
+        clientQuotasImage.describe(new DescribeClientQuotasRequestData()
+            .setComponents(List.of(new DescribeClientQuotasRequestData.ComponentData()
+                .setEntityType(ClientQuotaEntity.USER)
+                .setMatchType(DescribeClientQuotasRequest.MATCH_TYPE_EXACT)
+                .setMatch("user-0"))));
+    }
+}

metadata/src/main/java/org/apache/kafka/image/ClientQuotasImage.java

@@ -54,8 +54,26 @@ public final class ClientQuotasImage {
 
     private final Map<ClientQuotaEntity, ClientQuotaImage> entities;
 
+    // Map from entity type to entity name to set of entries. The entity type could be "user", "client-id", and "ip".
+    // {
+    //   "user": { "user1": [entry1, entry2], "user2": [entry3] },
+    //   "client-id": { "client-id1": [entry4], "client-id2": [entry5] },
+    //   "ip": { "ip1": [entry6], "ip2": [entry7] }
+    // }
+    private final Map<String, Map<String, Set<Entry<ClientQuotaEntity, ClientQuotaImage>>>> entitiesByType;
+
     public ClientQuotasImage(Map<ClientQuotaEntity, ClientQuotaImage> entities) {
         this.entities = Collections.unmodifiableMap(entities);
+        Map<String, Map<String, Set<Entry<ClientQuotaEntity, ClientQuotaImage>>>> entitiesByType = new HashMap<>();
+        for (Entry<ClientQuotaEntity, ClientQuotaImage> entry : entities.entrySet()) {
+            ClientQuotaEntity entity = entry.getKey();
+            for (Entry<String, String> entityEntry : entity.entries().entrySet()) {
+                entitiesByType.putIfAbsent(entityEntry.getKey(), new HashMap<>());
+                entitiesByType.get(entityEntry.getKey()).putIfAbsent(entityEntry.getValue(), new HashSet<>());
+                entitiesByType.get(entityEntry.getKey()).get(entityEntry.getValue()).add(entry);
+            }
+        }
+        this.entitiesByType = Collections.unmodifiableMap(entitiesByType);
     }
 
     public boolean isEmpty() {
@@ -126,40 +144,45 @@ public DescribeClientQuotasResponseData describe(DescribeClientQuotasRequestData
                     "user or clientId filter component.");
             }
         }
-        // TODO: this is O(N). We should add indexing here to speed it up. See KAFKA-13022.
-        for (Entry<ClientQuotaEntity, ClientQuotaImage> entry : entities.entrySet()) {
-            ClientQuotaEntity entity = entry.getKey();
-            ClientQuotaImage quotaImage = entry.getValue();
-            if (matches(entity, exactMatch, typeMatch, request.strict())) {
-                response.entries().add(toDescribeEntry(entity, quotaImage));
-            }
-        }
-        return response;
-    }
 
-    private static boolean matches(ClientQuotaEntity entity,
-                                   Map<String, String> exactMatch,
-                                   Set<String> typeMatch,
-                                   boolean strict) {
-        if (strict) {
-            if (entity.entries().size() != exactMatch.size() + typeMatch.size()) {
-                return false;
+        Set<ClientQuotaEntity> addedEntities = new HashSet<>();
+        for (Entry<String, String> exactMatchEntry : exactMatch.entrySet()) {
+            if (entitiesByType.containsKey(exactMatchEntry.getKey()) &&
+                entitiesByType.get(exactMatchEntry.getKey()).containsKey(exactMatchEntry.getValue())) {
+                for (Entry<ClientQuotaEntity, ClientQuotaImage> entry : entitiesByType.get(exactMatchEntry.getKey()).get(exactMatchEntry.getValue())) {
+                    if (request.strict() && !entry.getKey().entries().equals(exactMatch)) {
+                        continue;
+                    }
+                    if (!addedEntities.contains(entry.getKey())) {
+                        addedEntities.add(entry.getKey());
+                        response.entries().add(toDescribeEntry(entry.getKey(), entry.getValue()));
+                    }
+                }
             }
         }
-        for (Entry<String, String> entry : exactMatch.entrySet()) {
-            if (!entity.entries().containsKey(entry.getKey())) {
-                return false;
-            }
-            if (!Objects.equals(entity.entries().get(entry.getKey()), entry.getValue())) {
-                return false;
+
+        for (String type : typeMatch) {
+            if (entitiesByType.containsKey(type)) {
+                for (Set<Entry<ClientQuotaEntity, ClientQuotaImage>> entrySet : entitiesByType.get(type).values()) {
+                    for (Entry<ClientQuotaEntity, ClientQuotaImage> entry : entrySet) {
+                        if (request.strict() && entry.getKey().entries().size() != typeMatch.size()) {
+                            continue;
+                        }
+                        if (!addedEntities.contains(entry.getKey())) {
+                            addedEntities.add(entry.getKey());
+                            response.entries().add(toDescribeEntry(entry.getKey(), entry.getValue()));
+                        }
+                    }
+                }
             }
         }
-        for (String type : typeMatch) {
-            if (!entity.entries().containsKey(type)) {
-                return false;
+
+        if (!request.strict() && exactMatch.isEmpty() && typeMatch.isEmpty()) {
+            for (Entry<ClientQuotaEntity, ClientQuotaImage> entry : entities.entrySet()) {
+                response.entries().add(toDescribeEntry(entry.getKey(), entry.getValue()));
             }
         }
-        return true;
+        return response;
     }
 
     private static EntryData toDescribeEntry(ClientQuotaEntity entity,