Skip to content

Commit 6f90986

Browse files
committed
review comments from @adutra
1 parent 5a75b0d commit 6f90986

File tree

4 files changed

+147
-43
lines changed

4 files changed

+147
-43
lines changed

runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,11 @@
3333
@Priority(FilterPriorities.REQUEST_ID_FILTER)
3434
@Provider
3535
public class RequestIdFilter implements ContainerRequestFilter {
36-
@Inject RequestIdGenerator requestIdGenerator;
3736

3837
public static final String REQUEST_ID_KEY = "requestId";
3938

4039
@Inject LoggingConfiguration loggingConfiguration;
40+
@Inject RequestIdGenerator requestIdGenerator;
4141

4242
@Override
4343
public void filter(ContainerRequestContext rc) {

runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -22,21 +22,28 @@
2222
import com.google.common.annotations.VisibleForTesting;
2323
import jakarta.enterprise.context.ApplicationScoped;
2424
import java.util.UUID;
25+
import java.util.concurrent.atomic.AtomicBoolean;
2526
import java.util.concurrent.atomic.AtomicLong;
2627

2728
@ApplicationScoped
2829
public class RequestIdGenerator {
29-
private static String BASE_DEFAULT_UUID = UUID.randomUUID().toString();
30+
private volatile String baseDefaultUuid = UUID.randomUUID().toString();
3031
static final AtomicLong COUNTER = new AtomicLong();
3132
static final Long COUNTER_SOFT_MAX = Long.MAX_VALUE / 2;
3233

34+
private final AtomicBoolean resetInProgress = new AtomicBoolean(false);
35+
3336
public String generateRequestId() {
34-
String requestId = BASE_DEFAULT_UUID + "_" + COUNTER.incrementAndGet();
35-
if (COUNTER.get() >= COUNTER_SOFT_MAX) {
36-
// If we get anywhere close to danger of overflowing Long (which, theoretically,
37-
// would be extremely unlikely) rotate the UUID and start again.
38-
BASE_DEFAULT_UUID = UUID.randomUUID().toString();
39-
COUNTER.set(0);
37+
long currentCounter = COUNTER.incrementAndGet();
38+
String requestId = baseDefaultUuid + "_" + currentCounter;
39+
if (currentCounter >= COUNTER_SOFT_MAX) {
40+
if (resetInProgress.compareAndSet(false, true)) {
41+
// If we get anywhere close to danger of overflowing Long (which, theoretically,
42+
// would be extremely unlikely) rotate the UUID and start again.
43+
baseDefaultUuid = UUID.randomUUID().toString();
44+
COUNTER.set(0);
45+
resetInProgress.set(false);
46+
}
4047
}
4148
return requestId;
4249
}
Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.polaris.service.tracing;
21+
22+
import static org.assertj.core.api.Assertions.assertThat;
23+
import static org.junit.jupiter.api.Assertions.assertNotEquals;
24+
25+
import java.util.HashSet;
26+
import java.util.Set;
27+
import java.util.UUID;
28+
import org.junit.jupiter.api.BeforeEach;
29+
import org.junit.jupiter.api.Test;
30+
31+
public class RequestIdGeneratorTest {
32+
33+
private RequestIdGenerator requestIdGenerator;
34+
35+
@BeforeEach
36+
void setUp() {
37+
requestIdGenerator = new RequestIdGenerator();
38+
}
39+
40+
@Test
41+
void testGenerateRequestId_ReturnsValidFormat() {
42+
String requestId = requestIdGenerator.generateRequestId();
43+
44+
assertThat(requestId).isNotNull();
45+
assertThat(requestId).matches(this::isValidRequestIdFormat);
46+
assertThat(requestId).contains("_1"); // First call should increment counter to 1
47+
}
48+
49+
@Test
50+
void testGenerateRequestId_ReturnsUniqueIds() {
51+
Set<String> generatedIds = new HashSet<>();
52+
53+
// Generate multiple request IDs and verify they're all unique
54+
for (int i = 0; i < 1000; i++) {
55+
String requestId = requestIdGenerator.generateRequestId();
56+
assertThat(generatedIds).doesNotContain(requestId);
57+
generatedIds.add(requestId);
58+
}
59+
60+
assertThat(generatedIds).hasSize(1000);
61+
}
62+
63+
@Test
64+
void testCounterIncrementsSequentially() {
65+
requestIdGenerator.setCounter(0);
66+
67+
String firstId = requestIdGenerator.generateRequestId();
68+
String secondId = requestIdGenerator.generateRequestId();
69+
String thirdId = requestIdGenerator.generateRequestId();
70+
71+
assertThat(firstId).endsWith("_1");
72+
assertThat(secondId).endsWith("_2");
73+
assertThat(thirdId).endsWith("_3");
74+
}
75+
76+
@Test
77+
void testCounterRotationAtSoftMax() {
78+
// Set counter close to soft max
79+
long softMax = RequestIdGenerator.COUNTER_SOFT_MAX;
80+
requestIdGenerator.setCounter(softMax - 1);
81+
82+
String beforeRotation = requestIdGenerator.generateRequestId();
83+
String afterRotation = requestIdGenerator.generateRequestId();
84+
85+
// The UUID part should be different after rotation
86+
String beforeUuidPart = beforeRotation.substring(0, beforeRotation.lastIndexOf('_'));
87+
String afterUuidPart = afterRotation.substring(0, afterRotation.lastIndexOf('_'));
88+
assertNotEquals(beforeUuidPart, afterUuidPart);
89+
90+
// After rotation, counter should be reset and UUID should be different
91+
assertThat(beforeRotation).endsWith("_" + softMax);
92+
assertThat(afterRotation).endsWith("_1"); // Counter reset to 1 (after increment from 0)
93+
}
94+
95+
@Test
96+
void testSetCounterChangesNextGeneratedId() {
97+
requestIdGenerator.setCounter(100);
98+
99+
String requestId = requestIdGenerator.generateRequestId();
100+
101+
assertThat(requestId).endsWith("_101"); // Should increment from set value
102+
}
103+
104+
private boolean isValidRequestIdFormat(String str) {
105+
try {
106+
String[] requestIdParts = str.split("_");
107+
String uuid = requestIdParts[0];
108+
String counter = requestIdParts[1];
109+
UUID.fromString(uuid);
110+
Long.parseLong(counter);
111+
return true;
112+
} catch (IllegalArgumentException e) {
113+
return false;
114+
}
115+
}
116+
}

runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java renamed to runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdHeaderTest.java

Lines changed: 16 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -16,24 +16,24 @@
1616
* specific language governing permissions and limitations
1717
* under the License.
1818
*/
19-
package org.apache.polaris.service.admin;
19+
package org.apache.polaris.service.tracing;
2020

2121
import static org.assertj.core.api.Assertions.assertThat;
2222

2323
import io.quarkus.test.junit.QuarkusTest;
2424
import io.quarkus.test.junit.QuarkusTestProfile;
2525
import io.quarkus.test.junit.TestProfile;
26-
import jakarta.inject.Inject;
2726
import jakarta.ws.rs.client.Entity;
2827
import jakarta.ws.rs.core.MultivaluedHashMap;
2928
import jakarta.ws.rs.core.Response;
3029
import java.net.URI;
30+
import java.util.HashMap;
31+
import java.util.HashSet;
3132
import java.util.Map;
3233
import java.util.Objects;
33-
import java.util.UUID;
34+
import java.util.Set;
3435
import org.apache.polaris.service.it.env.PolarisApiEndpoints;
3536
import org.apache.polaris.service.it.env.PolarisClient;
36-
import org.apache.polaris.service.tracing.RequestIdGenerator;
3737
import org.junit.jupiter.api.Test;
3838

3939
@QuarkusTest
@@ -52,8 +52,6 @@ public Map<String, String> getConfigOverrides() {
5252
}
5353
}
5454

55-
@Inject RequestIdGenerator requestIdGenerator;
56-
5755
private static final String REQUEST_ID_HEADER = "x-test-request-id-random";
5856
private static final String REALM_HEADER = "realm";
5957
private static final String REALM = "realm1";
@@ -93,41 +91,24 @@ private Response request(Map<String, String> headers) {
9391
@Test
9492
public void testRequestIdHeaderSpecified() {
9593
String requestId = "pre-requested-request-id";
96-
Map<String, String> headers = Map.of(REALM_HEADER, REALM, REQUEST_ID_HEADER, requestId);
94+
HashMap<String, String> headers =
95+
new HashMap<>(Map.of(REALM_HEADER, REALM, REQUEST_ID_HEADER, requestId));
96+
assertThat(sendRequest(headers)).matches(s -> s.equals(requestId));
9797
assertThat(sendRequest(headers)).matches(s -> s.equals(requestId));
98-
}
9998

100-
@Test
101-
public void testRequestIdHeaderNotSpecified() {
102-
Map<String, String> headers = Map.of(REALM_HEADER, REALM);
103-
assertThat(sendRequest(headers)).matches(this::isValidDefaultUUID);
99+
String newRequestId = "new-pre-requested-request-id";
100+
headers.put(REQUEST_ID_HEADER, newRequestId);
101+
assertThat(sendRequest(headers)).matches(s -> s.equals(newRequestId));
104102
}
105103

106104
@Test
107-
public void testRequestIdHeaderNotSpecifiedAndCounterExhausted() {
108-
requestIdGenerator.setCounter(Long.MAX_VALUE / 2 + 1);
105+
public void testRequestIdHeaderNotSpecified() {
109106
Map<String, String> headers = Map.of(REALM_HEADER, REALM);
110-
String requestId = sendRequest(headers);
111-
assertThat(requestId).matches(this::isValidDefaultUUID);
112-
String currentRequestIdBase = requestId.split("_")[0];
113-
String uuidBase = currentRequestIdBase;
114-
115-
requestId = sendRequest(headers);
116-
assertThat(requestId).matches(this::isValidDefaultUUID);
117-
currentRequestIdBase = requestId.split("_")[0];
118-
assertThat(currentRequestIdBase).isNotEqualTo(uuidBase);
119-
}
120-
121-
private boolean isValidDefaultUUID(String str) {
122-
try {
123-
String[] requestIdParts = str.split("_");
124-
String uuid = requestIdParts[0];
125-
String counter = requestIdParts[1];
126-
UUID.fromString(uuid);
127-
Long.parseLong(counter);
128-
return true;
129-
} catch (IllegalArgumentException e) {
130-
return false;
107+
Set<String> requestIds = new HashSet<>();
108+
for (int i = 0; i < 10; i++) {
109+
String requestId = sendRequest(headers);
110+
assertThat(requestIds).doesNotContain(requestId);
111+
requestIds.add(requestId);
131112
}
132113
}
133114

0 commit comments

Comments
 (0)