Skip to content

Commit 2f0c7a4

Browse files
adnanhemaniadnanhemani
authored
Generate Request IDs (if not specified); Return Request ID as a Header (apache#2602)
1 parent 3b4b995 commit 2f0c7a4

File tree

9 files changed

+402
-10
lines changed

9 files changed

+402
-10
lines changed

‎runtime/service/src/main/java/org/apache/polaris/service/config/FilterPriorities.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@
2121
import jakarta.ws.rs.Priorities;
2222

2323
public final class FilterPriorities {
24-
public static final int REALM_CONTEXT_FILTER = Priorities.AUTHENTICATION - 100;
24+
public static final int REQUEST_ID_FILTER = Priorities.AUTHENTICATION - 101;
25+
public static final int REALM_CONTEXT_FILTER = REQUEST_ID_FILTER + 1;
2526
public static final int RATE_LIMITER_FILTER = Priorities.USER;
2627
public static final int MDC_FILTER = REALM_CONTEXT_FILTER + 1;
2728
public static final int TRACING_FILTER = REALM_CONTEXT_FILTER + 2;

‎runtime/service/src/main/java/org/apache/polaris/service/events/listeners/inmemory/InMemoryBufferEventListener.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919

2020
package org.apache.polaris.service.events.listeners.inmemory;
2121

22-
import static org.apache.polaris.service.logging.LoggingMDCFilter.REQUEST_ID_KEY;
22+
import static org.apache.polaris.service.tracing.RequestIdFilter.REQUEST_ID_KEY;
2323

2424
import com.github.benmanes.caffeine.cache.Caffeine;
2525
import com.github.benmanes.caffeine.cache.LoadingCache;

‎runtime/service/src/main/java/org/apache/polaris/service/logging/LoggingMDCFilter.java‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
package org.apache.polaris.service.logging;
2020

2121
import static org.apache.polaris.service.context.RealmContextFilter.REALM_CONTEXT_KEY;
22+
import static org.apache.polaris.service.tracing.RequestIdFilter.REQUEST_ID_KEY;
2223

2324
import jakarta.annotation.Priority;
2425
import jakarta.enterprise.context.ApplicationScoped;
@@ -38,7 +39,6 @@
3839
public class LoggingMDCFilter implements ContainerRequestFilter {
3940

4041
public static final String REALM_ID_KEY = "realmId";
41-
public static final String REQUEST_ID_KEY = "requestId";
4242

4343
@Inject LoggingConfiguration loggingConfiguration;
4444

@@ -49,11 +49,7 @@ public void filter(ContainerRequestContext rc) {
4949
// Also put the MDC values in the request context for use by other filters and handlers
5050
loggingConfiguration.mdc().forEach(MDC::put);
5151
loggingConfiguration.mdc().forEach(rc::setProperty);
52-
var requestId = rc.getHeaderString(loggingConfiguration.requestIdHeaderName());
53-
if (requestId != null) {
54-
MDC.put(REQUEST_ID_KEY, requestId);
55-
rc.setProperty(REQUEST_ID_KEY, requestId);
56-
}
52+
MDC.put(REQUEST_ID_KEY, (String) rc.getProperty(REQUEST_ID_KEY));
5753
RealmContext realmContext = (RealmContext) rc.getProperty(REALM_CONTEXT_KEY);
5854
MDC.put(REALM_ID_KEY, realmContext.getRealmIdentifier());
5955
}

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
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+
package org.apache.polaris.service.tracing;
20+
21+
import jakarta.annotation.Priority;
22+
import jakarta.enterprise.context.ApplicationScoped;
23+
import jakarta.inject.Inject;
24+
import jakarta.ws.rs.container.ContainerRequestContext;
25+
import jakarta.ws.rs.container.ContainerRequestFilter;
26+
import jakarta.ws.rs.container.PreMatching;
27+
import jakarta.ws.rs.ext.Provider;
28+
import org.apache.polaris.service.config.FilterPriorities;
29+
import org.apache.polaris.service.logging.LoggingConfiguration;
30+
31+
@PreMatching
32+
@ApplicationScoped
33+
@Priority(FilterPriorities.REQUEST_ID_FILTER)
34+
@Provider
35+
public class RequestIdFilter implements ContainerRequestFilter {
36+
37+
public static final String REQUEST_ID_KEY = "requestId";
38+
39+
@Inject LoggingConfiguration loggingConfiguration;
40+
@Inject RequestIdGenerator requestIdGenerator;
41+
42+
@Override
43+
public void filter(ContainerRequestContext rc) {
44+
var requestId = rc.getHeaderString(loggingConfiguration.requestIdHeaderName());
45+
if (requestId == null) {
46+
requestId = requestIdGenerator.generateRequestId();
47+
}
48+
rc.setProperty(REQUEST_ID_KEY, requestId);
49+
}
50+
}

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

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
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 com.google.common.annotations.VisibleForTesting;
23+
import jakarta.enterprise.context.ApplicationScoped;
24+
import java.util.UUID;
25+
import java.util.concurrent.atomic.AtomicReference;
26+
27+
@ApplicationScoped
28+
public class RequestIdGenerator {
29+
static final Long COUNTER_SOFT_MAX = Long.MAX_VALUE / 2;
30+
31+
record State(String uuid, long counter) {
32+
33+
State() {
34+
this(UUID.randomUUID().toString(), 1);
35+
}
36+
37+
String requestId() {
38+
return String.format("%s_%019d", uuid, counter);
39+
}
40+
41+
State increment() {
42+
return counter >= COUNTER_SOFT_MAX ? new State() : new State(uuid, counter + 1);
43+
}
44+
}
45+
46+
final AtomicReference<State> state = new AtomicReference<>(new State());
47+
48+
public String generateRequestId() {
49+
return state.getAndUpdate(State::increment).requestId();
50+
}
51+
52+
@VisibleForTesting
53+
public void setCounter(long counter) {
54+
state.set(new State(state.get().uuid, counter));
55+
}
56+
}

‎runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
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.apache.polaris.service.tracing.RequestIdFilter.REQUEST_ID_KEY;
23+
24+
import jakarta.enterprise.context.ApplicationScoped;
25+
import jakarta.inject.Inject;
26+
import jakarta.ws.rs.container.ContainerRequestContext;
27+
import jakarta.ws.rs.container.ContainerResponseContext;
28+
import jakarta.ws.rs.container.ContainerResponseFilter;
29+
import jakarta.ws.rs.ext.Provider;
30+
import org.apache.polaris.service.logging.LoggingConfiguration;
31+
32+
@ApplicationScoped
33+
@Provider
34+
public class RequestIdResponseFilter implements ContainerResponseFilter {
35+
36+
@Inject LoggingConfiguration loggingConfiguration;
37+
38+
@Override
39+
public void filter(
40+
ContainerRequestContext requestContext, ContainerResponseContext responseContext) {
41+
responseContext
42+
.getHeaders()
43+
.add(
44+
loggingConfiguration.requestIdHeaderName(), requestContext.getProperty(REQUEST_ID_KEY));
45+
}
46+
}

‎runtime/service/src/main/java/org/apache/polaris/service/tracing/TracingFilter.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
import org.apache.polaris.core.context.RealmContext;
2929
import org.apache.polaris.service.config.FilterPriorities;
3030
import org.apache.polaris.service.context.RealmContextFilter;
31-
import org.apache.polaris.service.logging.LoggingMDCFilter;
3231
import org.eclipse.microprofile.config.inject.ConfigProperty;
3332

3433
@PreMatching
@@ -47,7 +46,7 @@ public class TracingFilter implements ContainerRequestFilter {
4746
public void filter(ContainerRequestContext rc) {
4847
if (!sdkDisabled) {
4948
Span span = Span.current();
50-
String requestId = (String) rc.getProperty(LoggingMDCFilter.REQUEST_ID_KEY);
49+
String requestId = (String) rc.getProperty(RequestIdFilter.REQUEST_ID_KEY);
5150
if (requestId != null) {
5251
span.setAttribute(REQUEST_ID_ATTRIBUTE, requestId);
5352
}

‎runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdGeneratorTest.java‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
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+
// First call should increment counter to 1
47+
assertThat(extractCounterFromRequestId(requestId)).isEqualTo(1);
48+
}
49+
50+
@Test
51+
void testGenerateRequestId_ReturnsUniqueIds() {
52+
Set<String> generatedIds = new HashSet<>();
53+
54+
// Generate multiple request IDs and verify they're all unique
55+
for (int i = 0; i < 1000; i++) {
56+
String requestId = requestIdGenerator.generateRequestId();
57+
assertThat(generatedIds).doesNotContain(requestId);
58+
generatedIds.add(requestId);
59+
}
60+
61+
assertThat(generatedIds).hasSize(1000);
62+
}
63+
64+
@Test
65+
void testCounterIncrementsSequentially() {
66+
// requestIdGenerator.setCounter(0);
67+
68+
String firstId = requestIdGenerator.generateRequestId();
69+
String secondId = requestIdGenerator.generateRequestId();
70+
String thirdId = requestIdGenerator.generateRequestId();
71+
72+
assertThat(extractCounterFromRequestId(firstId)).isEqualTo(1);
73+
assertThat(extractCounterFromRequestId(secondId)).isEqualTo(2);
74+
assertThat(extractCounterFromRequestId(thirdId)).isEqualTo(3);
75+
}
76+
77+
@Test
78+
void testCounterRotationAtSoftMax() {
79+
// Set counter close to soft max
80+
long softMax = RequestIdGenerator.COUNTER_SOFT_MAX;
81+
requestIdGenerator.setCounter(softMax);
82+
83+
String beforeRotation = requestIdGenerator.generateRequestId();
84+
String afterRotation = requestIdGenerator.generateRequestId();
85+
86+
// The UUID part should be different after rotation
87+
String beforeUuidPart = beforeRotation.substring(0, beforeRotation.lastIndexOf('_'));
88+
String afterUuidPart = afterRotation.substring(0, afterRotation.lastIndexOf('_'));
89+
assertNotEquals(beforeUuidPart, afterUuidPart);
90+
91+
assertThat(extractCounterFromRequestId(beforeRotation)).isEqualTo(softMax);
92+
// Counter reset to 1 (after increment from 0)
93+
assertThat(extractCounterFromRequestId(afterRotation)).isEqualTo(1);
94+
}
95+
96+
@Test
97+
void testSetCounterChangesNextGeneratedId() {
98+
requestIdGenerator.setCounter(100);
99+
100+
String requestId = requestIdGenerator.generateRequestId();
101+
102+
// Should increment from set value
103+
assertThat(extractCounterFromRequestId(requestId)).isEqualTo(100);
104+
}
105+
106+
private boolean isValidRequestIdFormat(String str) {
107+
try {
108+
String[] requestIdParts = str.split("_");
109+
String uuid = requestIdParts[0];
110+
String counter = requestIdParts[1];
111+
UUID.fromString(uuid);
112+
Long.parseLong(counter);
113+
return true;
114+
} catch (IllegalArgumentException e) {
115+
return false;
116+
}
117+
}
118+
119+
private long extractCounterFromRequestId(String requestId) {
120+
return Long.parseLong(requestId.split("_")[1]);
121+
}
122+
}

0 commit comments

Comments
 (0)