Skip to content

Commit ff2ba93

Browse files
FrankYang0529FrankYang0529
authored
KAFKA-13022 Optimize ClientQuotasImage#describe (#19079)
In previous implementation, the `ClientQuotasImage#describe` goes through all `ClientQuotaEntity` and checks whether entity type and entity name are matched the filter. In this PR, it goes through all `ClientQuotaEntity#entries` in the constructor and build a new data structure `Map<String, Map<String, Set<Entry<ClientQuotaEntity, ClientQuotaImage>>>>`. The first layer is entity type. There are only three entity types: user, client-id, and ip. The second layer is entity name. The last layer is a set of matched entries. With this data structure, we can speed up `describe` function. Correctness is covered by `ClientQuotasRequestTest`. Reviewers: Kuan-Po Tseng <brandboat@gmail.com>, Mickael Maison <mimaison@apache.org>, Alyssa Huang <ahuang@confluent.io>, TaiJuWu <tjwu1217@gmail.com>, Chia-Ping Tsai <chia7712@gmail.com>
1 parent 2e0596d commit ff2ba93

3 files changed

Lines changed: 505 additions & 28 deletions

File tree

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

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.kafka.jmh.metadata;
18+
19+
import org.apache.kafka.common.message.DescribeClientQuotasRequestData;
20+
import org.apache.kafka.common.message.DescribeClientQuotasResponseData;
21+
import org.apache.kafka.common.quota.ClientQuotaEntity;
22+
import org.apache.kafka.common.requests.DescribeClientQuotasRequest;
23+
import org.apache.kafka.image.ClientQuotaImage;
24+
import org.apache.kafka.image.ClientQuotasImage;
25+
import org.apache.kafka.server.config.QuotaConfig;
26+
27+
import org.openjdk.jmh.annotations.Benchmark;
28+
import org.openjdk.jmh.annotations.BenchmarkMode;
29+
import org.openjdk.jmh.annotations.Fork;
30+
import org.openjdk.jmh.annotations.Level;
31+
import org.openjdk.jmh.annotations.Measurement;
32+
import org.openjdk.jmh.annotations.Mode;
33+
import org.openjdk.jmh.annotations.OutputTimeUnit;
34+
import org.openjdk.jmh.annotations.Param;
35+
import org.openjdk.jmh.annotations.Scope;
36+
import org.openjdk.jmh.annotations.Setup;
37+
import org.openjdk.jmh.annotations.State;
38+
import org.openjdk.jmh.annotations.Warmup;
39+
40+
import java.util.HashMap;
41+
import java.util.List;
42+
import java.util.Map;
43+
import java.util.concurrent.TimeUnit;
44+
45+
@State(Scope.Benchmark)
46+
@Fork(value = 1)
47+
@Warmup(iterations = 5)
48+
@Measurement(iterations = 15)
49+
@BenchmarkMode(Mode.AverageTime)
50+
@OutputTimeUnit(TimeUnit.NANOSECONDS)
51+
public class ClientQuotasImageDescribeBenchmark {
52+
53+
@Param({"10", "100", "1000"})
54+
private int eachEntityCount;
55+
56+
private ClientQuotasImage clientQuotasImage;
57+
58+
@Setup(Level.Trial)
59+
public void setup() {
60+
clientQuotasImage = createClientQuotasImage(eachEntityCount);
61+
}
62+
63+
static ClientQuotasImage createClientQuotasImage(int eachEntityCount) {
64+
Map<ClientQuotaEntity, ClientQuotaImage> entities = new HashMap<>();
65+
ClientQuotaImage defaultImage = new ClientQuotaImage(Map.of(QuotaConfig.REQUEST_PERCENTAGE_OVERRIDE_CONFIG, 1.0));
66+
for (int i = 0; i < eachEntityCount; i++) {
67+
entities.put(new ClientQuotaEntity(Map.of(ClientQuotaEntity.USER, "user-" + i)), defaultImage);
68+
entities.put(new ClientQuotaEntity(Map.of(ClientQuotaEntity.CLIENT_ID, "client-id-" + i)), defaultImage);
69+
entities.put(new ClientQuotaEntity(Map.of(ClientQuotaEntity.IP, "ip-" + i)), defaultImage);
70+
}
71+
return new ClientQuotasImage(entities);
72+
}
73+
74+
@Benchmark
75+
public DescribeClientQuotasResponseData describeSpecified() {
76+
return clientQuotasImage.describe(new DescribeClientQuotasRequestData()
77+
.setComponents(List.of(new DescribeClientQuotasRequestData.ComponentData()
78+
.setEntityType(ClientQuotaEntity.USER)
79+
.setMatchType(DescribeClientQuotasRequest.MATCH_TYPE_SPECIFIED)
80+
.setMatch(null))));
81+
}
82+
83+
@Benchmark
84+
public DescribeClientQuotasResponseData describeDefault() {
85+
return clientQuotasImage.describe(new DescribeClientQuotasRequestData()
86+
.setComponents(List.of(new DescribeClientQuotasRequestData.ComponentData()
87+
.setEntityType(ClientQuotaEntity.USER)
88+
.setMatchType(DescribeClientQuotasRequest.MATCH_TYPE_DEFAULT)
89+
.setMatch(null))));
90+
}
91+
92+
@Benchmark
93+
public DescribeClientQuotasResponseData describeExact() {
94+
return clientQuotasImage.describe(new DescribeClientQuotasRequestData()
95+
.setComponents(List.of(new DescribeClientQuotasRequestData.ComponentData()
96+
.setEntityType(ClientQuotaEntity.USER)
97+
.setMatchType(DescribeClientQuotasRequest.MATCH_TYPE_EXACT)
98+
.setMatch("user-0"))));
99+
}
100+
}

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

Lines changed: 99 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -48,11 +48,50 @@
4848
* <p>
4949
* This class is thread-safe.
5050
*/
51-
public record ClientQuotasImage(Map<ClientQuotaEntity, ClientQuotaImage> entities) {
51+
public final class ClientQuotasImage {
5252
public static final ClientQuotasImage EMPTY = new ClientQuotasImage(Map.of());
5353

54+
private final Map<ClientQuotaEntity, ClientQuotaImage> entities;
55+
56+
// Map from entity type to entity name to set of entries. The entity type could be "user", "client-id", and "ip".
57+
// {
58+
// "user": { "user1": {entity1: image1}, "user2": {entity2: image2} },
59+
// "client-id": { "client-id1": {entity3: image3}, "client-id2": {entity4: image4} },
60+
// "ip": { "ip1": {entity5: image5}, "ip2": {entity6: image6} }
61+
// }
62+
private final Map<String, Map<String, Map<ClientQuotaEntity, ClientQuotaImage>>> entitiesByTypeAndName;
63+
64+
// Map from entity type to set of entries. The entity type could be "user", "client-id", and "ip".
65+
// {
66+
// "user": { entity1: image1, entity2: image2 },
67+
// "client-id": { entity3: image3, entity4: image4 },
68+
// "ip": { entity5: image5, entity6: image6 }
69+
// }
70+
private final Map<String, Map<ClientQuotaEntity, ClientQuotaImage>> entitiesByType;
71+
5472
public ClientQuotasImage(Map<ClientQuotaEntity, ClientQuotaImage> entities) {
5573
this.entities = Collections.unmodifiableMap(entities);
74+
var entitiesByTypeAndName = new HashMap<String, Map<String, Map<ClientQuotaEntity, ClientQuotaImage>>>();
75+
var entitiesByType = new HashMap<String, Map<ClientQuotaEntity, ClientQuotaImage>>();
76+
for (var entry : entities.entrySet()) {
77+
ClientQuotaEntity entity = entry.getKey();
78+
for (var entityEntry : entity.entries().entrySet()) {
79+
entitiesByTypeAndName
80+
.computeIfAbsent(entityEntry.getKey(), k -> new HashMap<>())
81+
.computeIfAbsent(entityEntry.getValue(), k -> new HashMap<>())
82+
.put(entity, entry.getValue());
83+
84+
entitiesByType
85+
.computeIfAbsent(entityEntry.getKey(), k -> new HashMap<>())
86+
.put(entity, entry.getValue());
87+
}
88+
}
89+
this.entitiesByTypeAndName = Collections.unmodifiableMap(entitiesByTypeAndName);
90+
this.entitiesByType = Collections.unmodifiableMap(entitiesByType);
91+
}
92+
93+
public Map<ClientQuotaEntity, ClientQuotaImage> entities() {
94+
return entities;
5695
}
5796

5897
public boolean isEmpty() {
@@ -68,7 +107,6 @@ public void write(ImageWriter writer) {
68107
}
69108

70109
public DescribeClientQuotasResponseData describe(DescribeClientQuotasRequestData request) {
71-
DescribeClientQuotasResponseData response = new DescribeClientQuotasResponseData();
72110
Map<String, String> exactMatch = new HashMap<>();
73111
Set<String> typeMatch = new HashSet<>();
74112
for (DescribeClientQuotasRequestData.ComponentData component : request.components()) {
@@ -118,40 +156,61 @@ public DescribeClientQuotasResponseData describe(DescribeClientQuotasRequestData
118156
"user or clientId filter component.");
119157
}
120158
}
121-
// TODO: this is O(N). We should add indexing here to speed it up. See KAFKA-13022.
122-
for (Entry<ClientQuotaEntity, ClientQuotaImage> entry : entities.entrySet()) {
123-
ClientQuotaEntity entity = entry.getKey();
124-
ClientQuotaImage quotaImage = entry.getValue();
125-
if (matches(entity, exactMatch, typeMatch, request.strict())) {
126-
response.entries().add(toDescribeEntry(entity, quotaImage));
127-
}
128-
}
129-
return response;
159+
160+
return matches(exactMatch, typeMatch, request.strict());
130161
}
131162

132-
private static boolean matches(ClientQuotaEntity entity,
133-
Map<String, String> exactMatch,
134-
Set<String> typeMatch,
135-
boolean strict) {
136-
if (strict) {
137-
if (entity.entries().size() != exactMatch.size() + typeMatch.size()) {
138-
return false;
163+
private DescribeClientQuotasResponseData matches(
164+
Map<String, String> exactMatch,
165+
Set<String> typeMatch,
166+
boolean strict
167+
) {
168+
DescribeClientQuotasResponseData response = new DescribeClientQuotasResponseData();
169+
Map<ClientQuotaEntity, ClientQuotaImage> candidates = null;
170+
// Case 1: exact match exists. Filter candidates based on exact match first and then type match
171+
if (!exactMatch.isEmpty()) {
172+
for (Entry<String, String> exactMatchEntry : exactMatch.entrySet()) {
173+
String entityType = exactMatchEntry.getKey();
174+
String entityName = exactMatchEntry.getValue();
175+
var nameMap = entitiesByTypeAndName.get(entityType);
176+
var matches = Map.<ClientQuotaEntity, ClientQuotaImage>of();
177+
if (nameMap != null) matches = nameMap.getOrDefault(entityName, Map.of());
178+
if (candidates == null) {
179+
candidates = new HashMap<>(matches);
180+
} else {
181+
candidates.keySet().retainAll(matches.keySet());
182+
}
139183
}
140-
}
141-
for (Entry<String, String> entry : exactMatch.entrySet()) {
142-
if (!entity.entries().containsKey(entry.getKey())) {
143-
return false;
184+
185+
for (String type : typeMatch) {
186+
candidates.keySet().retainAll(entitiesByType.getOrDefault(type, Map.of()).keySet());
187+
}
188+
} else if (!typeMatch.isEmpty()) {
189+
// Case 2: no exact match, only type match exists
190+
for (String type : typeMatch) {
191+
Map<ClientQuotaEntity, ClientQuotaImage> matches = entitiesByType.getOrDefault(type, Map.of());
192+
if (candidates == null) {
193+
candidates = new HashMap<>(matches);
194+
} else {
195+
candidates.keySet().retainAll(matches.keySet());
196+
}
144197
}
145-
if (!Objects.equals(entity.entries().get(entry.getKey()), entry.getValue())) {
146-
return false;
198+
} else if (!strict) {
199+
// Case 3: no exact match, no type match, no strict, return all entries
200+
for (Entry<ClientQuotaEntity, ClientQuotaImage> entry : entities.entrySet()) {
201+
response.entries().add(toDescribeEntry(entry.getKey(), entry.getValue()));
147202
}
203+
return response;
148204
}
149-
for (String type : typeMatch) {
150-
if (!entity.entries().containsKey(type)) {
151-
return false;
205+
206+
if (candidates != null) {
207+
for (Entry<ClientQuotaEntity, ClientQuotaImage> entry : candidates.entrySet()) {
208+
if (!strict || entry.getKey().entries().size() == exactMatch.size() + typeMatch.size()) {
209+
response.entries().add(toDescribeEntry(entry.getKey(), entry.getValue()));
210+
}
152211
}
153212
}
154-
return true;
213+
return response;
155214
}
156215

157216
private static EntryData toDescribeEntry(ClientQuotaEntity entity,
@@ -166,6 +225,18 @@ private static EntryData toDescribeEntry(ClientQuotaEntity entity,
166225
return data;
167226
}
168227

228+
@Override
229+
public boolean equals(Object o) {
230+
if (this == o) return true;
231+
if (!(o instanceof ClientQuotasImage other)) return false;
232+
return entities.equals(other.entities);
233+
}
234+
235+
@Override
236+
public int hashCode() {
237+
return Objects.hash(entities);
238+
}
239+
169240
@Override
170241
public String toString() {
171242
return new ClientQuotasImageNode(this).stringify();

0 commit comments

Comments
 (0)