From 49a182be27a0b41d9e4fc2d9d0fbbc16e8a13022 Mon Sep 17 00:00:00 2001 From: adnanhemani Date: Thu, 18 Sep 2025 00:36:46 -0700 Subject: [PATCH 01/10] Generate Request IDs (if not specified); Return Request ID as a Header --- .../service/config/FilterPriorities.java | 1 + .../service/logging/LoggingMDCFilter.java | 6 - .../service/tracing/RequestIdFilter.java | 52 ++++++++ .../tracing/RequestIdResponseFilter.java | 38 ++++++ .../service/tracing/TracingFilter.java | 3 +- .../service/admin/RequestIdHeaderTest.java | 123 ++++++++++++++++++ 6 files changed, 215 insertions(+), 8 deletions(-) create mode 100644 runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java create mode 100644 runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java create mode 100644 runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java diff --git a/runtime/service/src/main/java/org/apache/polaris/service/config/FilterPriorities.java b/runtime/service/src/main/java/org/apache/polaris/service/config/FilterPriorities.java index eed2d43530..06fd801684 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/config/FilterPriorities.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/config/FilterPriorities.java @@ -21,6 +21,7 @@ import jakarta.ws.rs.Priorities; public final class FilterPriorities { + public static final int REQUEST_ID_FILTER = Priorities.AUTHENTICATION - 101; public static final int REALM_CONTEXT_FILTER = Priorities.AUTHENTICATION - 100; public static final int RATE_LIMITER_FILTER = Priorities.USER; public static final int MDC_FILTER = REALM_CONTEXT_FILTER + 1; diff --git a/runtime/service/src/main/java/org/apache/polaris/service/logging/LoggingMDCFilter.java b/runtime/service/src/main/java/org/apache/polaris/service/logging/LoggingMDCFilter.java index c3b078ea16..8274a4c551 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/logging/LoggingMDCFilter.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/logging/LoggingMDCFilter.java @@ -38,7 +38,6 @@ public class LoggingMDCFilter implements ContainerRequestFilter { public static final String REALM_ID_KEY = "realmId"; - public static final String REQUEST_ID_KEY = "requestId"; @Inject LoggingConfiguration loggingConfiguration; @@ -49,11 +48,6 @@ public void filter(ContainerRequestContext rc) { // Also put the MDC values in the request context for use by other filters and handlers loggingConfiguration.mdc().forEach(MDC::put); loggingConfiguration.mdc().forEach(rc::setProperty); - var requestId = rc.getHeaderString(loggingConfiguration.requestIdHeaderName()); - if (requestId != null) { - MDC.put(REQUEST_ID_KEY, requestId); - rc.setProperty(REQUEST_ID_KEY, requestId); - } RealmContext realmContext = (RealmContext) rc.getProperty(REALM_CONTEXT_KEY); MDC.put(REALM_ID_KEY, realmContext.getRealmIdentifier()); } diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java new file mode 100644 index 0000000000..86ba719cc4 --- /dev/null +++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java @@ -0,0 +1,52 @@ +/* + * 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.polaris.service.tracing; + +import jakarta.annotation.Priority; +import jakarta.enterprise.context.ApplicationScoped; +import jakarta.inject.Inject; +import jakarta.ws.rs.container.ContainerRequestContext; +import jakarta.ws.rs.container.ContainerRequestFilter; +import jakarta.ws.rs.container.PreMatching; +import jakarta.ws.rs.ext.Provider; +import java.util.UUID; +import org.apache.polaris.service.config.FilterPriorities; +import org.apache.polaris.service.logging.LoggingConfiguration; +import org.slf4j.MDC; + +@PreMatching +@ApplicationScoped +@Priority(FilterPriorities.REQUEST_ID_FILTER) +@Provider +public class RequestIdFilter implements ContainerRequestFilter { + + public static final String REQUEST_ID_KEY = "requestId"; + + @Inject LoggingConfiguration loggingConfiguration; + + @Override + public void filter(ContainerRequestContext rc) { + var requestId = rc.getHeaderString(loggingConfiguration.requestIdHeaderName()); + if (requestId == null) { + requestId = UUID.randomUUID().toString(); + } + MDC.put(REQUEST_ID_KEY, requestId); + rc.setProperty(REQUEST_ID_KEY, requestId); + } +} diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java new file mode 100644 index 0000000000..085c224391 --- /dev/null +++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java @@ -0,0 +1,38 @@ +/* + * 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.polaris.service.tracing; + +import static org.apache.polaris.service.tracing.RequestIdFilter.REQUEST_ID_KEY; + +import jakarta.ws.rs.container.ContainerRequestContext; +import jakarta.ws.rs.container.ContainerResponseContext; +import jakarta.ws.rs.container.ContainerResponseFilter; +import jakarta.ws.rs.ext.Provider; + +@Provider +public class RequestIdResponseFilter implements ContainerResponseFilter { + private static final String REQUEST_ID_HEADER = "X-Request-Id"; + + @Override + public void filter( + ContainerRequestContext requestContext, ContainerResponseContext responseContext) { + responseContext.getHeaders().add(REQUEST_ID_HEADER, requestContext.getProperty(REQUEST_ID_KEY)); + } +} diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/TracingFilter.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/TracingFilter.java index b4c0562351..8b6859ca36 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/tracing/TracingFilter.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/TracingFilter.java @@ -28,7 +28,6 @@ import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.service.config.FilterPriorities; import org.apache.polaris.service.context.RealmContextFilter; -import org.apache.polaris.service.logging.LoggingMDCFilter; import org.eclipse.microprofile.config.inject.ConfigProperty; @PreMatching @@ -47,7 +46,7 @@ public class TracingFilter implements ContainerRequestFilter { public void filter(ContainerRequestContext rc) { if (!sdkDisabled) { Span span = Span.current(); - String requestId = (String) rc.getProperty(LoggingMDCFilter.REQUEST_ID_KEY); + String requestId = (String) rc.getProperty(RequestIdFilter.REQUEST_ID_KEY); if (requestId != null) { span.setAttribute(REQUEST_ID_ATTRIBUTE, requestId); } diff --git a/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java b/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java new file mode 100644 index 0000000000..e60a4dcc5e --- /dev/null +++ b/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java @@ -0,0 +1,123 @@ +/* + * 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.polaris.service.admin; + +import static org.assertj.core.api.Assertions.assertThat; + +import io.quarkus.test.junit.QuarkusTest; +import io.quarkus.test.junit.QuarkusTestProfile; +import io.quarkus.test.junit.TestProfile; +import jakarta.ws.rs.client.Entity; +import jakarta.ws.rs.core.MultivaluedHashMap; +import jakarta.ws.rs.core.Response; +import java.net.URI; +import java.util.Map; +import java.util.Objects; +import java.util.UUID; +import org.apache.polaris.service.it.env.PolarisApiEndpoints; +import org.apache.polaris.service.it.env.PolarisClient; +import org.junit.jupiter.api.Test; + +@QuarkusTest +@TestProfile(RequestIdHeaderTest.Profile.class) +public class RequestIdHeaderTest { + public static class Profile implements QuarkusTestProfile { + @Override + public Map getConfigOverrides() { + return Map.of( + "polaris.log.request-id-header-name", + REQUEST_ID_HEADER, + "polaris.bootstrap.credentials", + String.format("%s,%s,%s", REALM, CLIENT_ID, CLIENT_SECRET), + "polaris.realm-context.header-name", + REALM_HEADER, + "polaris.realm-context.realms", + REALM); + } + } + + private static final String REQUEST_ID_HEADER = "x-test-request-id-random"; + private static final String RESPONSE_REQUEST_ID_HEADER = "X-Request-Id"; + private static final String REALM_HEADER = "realm"; + private static final String REALM = "realm1"; + private static final String CLIENT_ID = "client1"; + private static final String CLIENT_SECRET = "secret1"; + + private static final URI baseUri = + URI.create( + "http://localhost:" + + Objects.requireNonNull( + Integer.getInteger("quarkus.http.test-port"), + "System property not set correctly: quarkus.http.test-port")); + + private Response request(Map headers) { + try (PolarisClient client = + PolarisClient.polarisClient(new PolarisApiEndpoints(baseUri, REALM, headers))) { + return client + .catalogApiPlain() + .request("v1/oauth/tokens") + .post( + Entity.form( + new MultivaluedHashMap<>( + Map.of( + "grant_type", + "client_credentials", + "scope", + "PRINCIPAL_ROLE:ALL", + "client_id", + CLIENT_ID, + "client_secret", + CLIENT_SECRET)))); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + @Test + public void testRequestIdHeaderSpecified() { + String requestId = "pre-requested-request-id"; + Map headers = Map.of(REALM_HEADER, REALM, REQUEST_ID_HEADER, requestId); + try (Response response = request(headers)) { + assertThat(response.getHeaders()).containsKey(RESPONSE_REQUEST_ID_HEADER); + assertThat(response.getHeaders().get(RESPONSE_REQUEST_ID_HEADER)).hasSize(1); + assertThat(response.getHeaders().get(RESPONSE_REQUEST_ID_HEADER)) + .allMatch(s -> s.equals(requestId)); + } + } + + @Test + public void testRequestIdHeaderNotSpecified() { + Map headers = Map.of(REALM_HEADER, REALM); + try (Response response = request(headers)) { + assertThat(response.getHeaders()).containsKey(RESPONSE_REQUEST_ID_HEADER); + assertThat(response.getHeaders().get(RESPONSE_REQUEST_ID_HEADER)).hasSize(1); + assertThat(response.getHeaders().get(RESPONSE_REQUEST_ID_HEADER)) + .allMatch(s -> isValidUUID(s.toString())); + } + } + + private boolean isValidUUID(String str) { + try { + UUID.fromString(str); + return true; + } catch (IllegalArgumentException e) { + return false; + } + } +} From f44d70ac7853b1ac78d0b02503c9dae439d31668 Mon Sep 17 00:00:00 2001 From: adnanhemani Date: Thu, 18 Sep 2025 18:16:03 -0700 Subject: [PATCH 02/10] review comments from @adutra --- .../polaris/service/logging/LoggingMDCFilter.java | 2 ++ .../polaris/service/tracing/RequestIdFilter.java | 2 -- .../service/tracing/RequestIdResponseFilter.java | 10 ++++++++-- 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/runtime/service/src/main/java/org/apache/polaris/service/logging/LoggingMDCFilter.java b/runtime/service/src/main/java/org/apache/polaris/service/logging/LoggingMDCFilter.java index 8274a4c551..fb0112cbb6 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/logging/LoggingMDCFilter.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/logging/LoggingMDCFilter.java @@ -19,6 +19,7 @@ package org.apache.polaris.service.logging; import static org.apache.polaris.service.context.RealmContextFilter.REALM_CONTEXT_KEY; +import static org.apache.polaris.service.tracing.RequestIdFilter.REQUEST_ID_KEY; import jakarta.annotation.Priority; import jakarta.enterprise.context.ApplicationScoped; @@ -48,6 +49,7 @@ public void filter(ContainerRequestContext rc) { // Also put the MDC values in the request context for use by other filters and handlers loggingConfiguration.mdc().forEach(MDC::put); loggingConfiguration.mdc().forEach(rc::setProperty); + MDC.put(REQUEST_ID_KEY, (String) rc.getProperty(REQUEST_ID_KEY)); RealmContext realmContext = (RealmContext) rc.getProperty(REALM_CONTEXT_KEY); MDC.put(REALM_ID_KEY, realmContext.getRealmIdentifier()); } diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java index 86ba719cc4..53d3bcf96b 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java @@ -28,7 +28,6 @@ import java.util.UUID; import org.apache.polaris.service.config.FilterPriorities; import org.apache.polaris.service.logging.LoggingConfiguration; -import org.slf4j.MDC; @PreMatching @ApplicationScoped @@ -46,7 +45,6 @@ public void filter(ContainerRequestContext rc) { if (requestId == null) { requestId = UUID.randomUUID().toString(); } - MDC.put(REQUEST_ID_KEY, requestId); rc.setProperty(REQUEST_ID_KEY, requestId); } } diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java index 085c224391..6d252a9c27 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java @@ -21,18 +21,24 @@ import static org.apache.polaris.service.tracing.RequestIdFilter.REQUEST_ID_KEY; +import jakarta.inject.Inject; import jakarta.ws.rs.container.ContainerRequestContext; import jakarta.ws.rs.container.ContainerResponseContext; import jakarta.ws.rs.container.ContainerResponseFilter; import jakarta.ws.rs.ext.Provider; +import org.apache.polaris.service.logging.LoggingConfiguration; @Provider public class RequestIdResponseFilter implements ContainerResponseFilter { - private static final String REQUEST_ID_HEADER = "X-Request-Id"; + + @Inject LoggingConfiguration loggingConfiguration; @Override public void filter( ContainerRequestContext requestContext, ContainerResponseContext responseContext) { - responseContext.getHeaders().add(REQUEST_ID_HEADER, requestContext.getProperty(REQUEST_ID_KEY)); + responseContext + .getHeaders() + .add( + loggingConfiguration.requestIdHeaderName(), requestContext.getProperty(REQUEST_ID_KEY)); } } From 1c10523ea091a9fc4cf8d55de9ddc731f2baba33 Mon Sep 17 00:00:00 2001 From: adnanhemani Date: Thu, 18 Sep 2025 19:03:23 -0700 Subject: [PATCH 03/10] pass tests --- .../service/tracing/RequestIdResponseFilter.java | 2 ++ .../polaris/service/admin/RequestIdHeaderTest.java | 14 ++++++-------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java index 6d252a9c27..728960182d 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java @@ -21,6 +21,7 @@ import static org.apache.polaris.service.tracing.RequestIdFilter.REQUEST_ID_KEY; +import jakarta.enterprise.context.ApplicationScoped; import jakarta.inject.Inject; import jakarta.ws.rs.container.ContainerRequestContext; import jakarta.ws.rs.container.ContainerResponseContext; @@ -28,6 +29,7 @@ import jakarta.ws.rs.ext.Provider; import org.apache.polaris.service.logging.LoggingConfiguration; +@ApplicationScoped @Provider public class RequestIdResponseFilter implements ContainerResponseFilter { diff --git a/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java b/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java index e60a4dcc5e..335d3b9dcb 100644 --- a/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java +++ b/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java @@ -53,7 +53,6 @@ public Map getConfigOverrides() { } private static final String REQUEST_ID_HEADER = "x-test-request-id-random"; - private static final String RESPONSE_REQUEST_ID_HEADER = "X-Request-Id"; private static final String REALM_HEADER = "realm"; private static final String REALM = "realm1"; private static final String CLIENT_ID = "client1"; @@ -94,10 +93,9 @@ public void testRequestIdHeaderSpecified() { String requestId = "pre-requested-request-id"; Map headers = Map.of(REALM_HEADER, REALM, REQUEST_ID_HEADER, requestId); try (Response response = request(headers)) { - assertThat(response.getHeaders()).containsKey(RESPONSE_REQUEST_ID_HEADER); - assertThat(response.getHeaders().get(RESPONSE_REQUEST_ID_HEADER)).hasSize(1); - assertThat(response.getHeaders().get(RESPONSE_REQUEST_ID_HEADER)) - .allMatch(s -> s.equals(requestId)); + assertThat(response.getHeaders()).containsKey(REQUEST_ID_HEADER); + assertThat(response.getHeaders().get(REQUEST_ID_HEADER)).hasSize(1); + assertThat(response.getHeaders().get(REQUEST_ID_HEADER)).allMatch(s -> s.equals(requestId)); } } @@ -105,9 +103,9 @@ public void testRequestIdHeaderSpecified() { public void testRequestIdHeaderNotSpecified() { Map headers = Map.of(REALM_HEADER, REALM); try (Response response = request(headers)) { - assertThat(response.getHeaders()).containsKey(RESPONSE_REQUEST_ID_HEADER); - assertThat(response.getHeaders().get(RESPONSE_REQUEST_ID_HEADER)).hasSize(1); - assertThat(response.getHeaders().get(RESPONSE_REQUEST_ID_HEADER)) + assertThat(response.getHeaders()).containsKey(REQUEST_ID_HEADER); + assertThat(response.getHeaders().get(REQUEST_ID_HEADER)).hasSize(1); + assertThat(response.getHeaders().get(REQUEST_ID_HEADER)) .allMatch(s -> isValidUUID(s.toString())); } } From 6632c212043951679ff3cd28c9b1b90ef49bb92e Mon Sep 17 00:00:00 2001 From: adnanhemani Date: Fri, 19 Sep 2025 14:41:13 -0700 Subject: [PATCH 04/10] review from @adutra --- .../org/apache/polaris/service/admin/RequestIdHeaderTest.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java b/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java index 335d3b9dcb..e6e9a36624 100644 --- a/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java +++ b/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java @@ -43,8 +43,6 @@ public Map getConfigOverrides() { return Map.of( "polaris.log.request-id-header-name", REQUEST_ID_HEADER, - "polaris.bootstrap.credentials", - String.format("%s,%s,%s", REALM, CLIENT_ID, CLIENT_SECRET), "polaris.realm-context.header-name", REALM_HEADER, "polaris.realm-context.realms", From 0bd7a21bf4f95c59b687f39dd583737fd48c68a0 Mon Sep 17 00:00:00 2001 From: adnanhemani Date: Fri, 19 Sep 2025 18:01:09 -0700 Subject: [PATCH 05/10] merge from main --- .../events/listeners/inmemory/InMemoryEventListener.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runtime/service/src/main/java/org/apache/polaris/service/events/listeners/inmemory/InMemoryEventListener.java b/runtime/service/src/main/java/org/apache/polaris/service/events/listeners/inmemory/InMemoryEventListener.java index 8af66b1941..003889d3b8 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/events/listeners/inmemory/InMemoryEventListener.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/events/listeners/inmemory/InMemoryEventListener.java @@ -19,7 +19,7 @@ package org.apache.polaris.service.events.listeners.inmemory; -import static org.apache.polaris.service.logging.LoggingMDCFilter.REQUEST_ID_KEY; +import static org.apache.polaris.service.tracing.RequestIdFilter.REQUEST_ID_KEY; import com.github.benmanes.caffeine.cache.Caffeine; import com.github.benmanes.caffeine.cache.LoadingCache; From ea02c79070d0813d6043554bd143230b8103b9b9 Mon Sep 17 00:00:00 2001 From: adnanhemani Date: Wed, 24 Sep 2025 22:55:43 -0700 Subject: [PATCH 06/10] Add RequestIdGenerator --- .../service/tracing/RequestIdFilter.java | 4 +- .../service/tracing/RequestIdGenerator.java | 48 +++++++++++++++++++ .../service/admin/RequestIdHeaderTest.java | 48 ++++++++++++++----- 3 files changed, 85 insertions(+), 15 deletions(-) create mode 100644 runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java index 53d3bcf96b..ae1a60178d 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java @@ -25,7 +25,6 @@ import jakarta.ws.rs.container.ContainerRequestFilter; import jakarta.ws.rs.container.PreMatching; import jakarta.ws.rs.ext.Provider; -import java.util.UUID; import org.apache.polaris.service.config.FilterPriorities; import org.apache.polaris.service.logging.LoggingConfiguration; @@ -34,6 +33,7 @@ @Priority(FilterPriorities.REQUEST_ID_FILTER) @Provider public class RequestIdFilter implements ContainerRequestFilter { + @Inject RequestIdGenerator requestIdGenerator; public static final String REQUEST_ID_KEY = "requestId"; @@ -43,7 +43,7 @@ public class RequestIdFilter implements ContainerRequestFilter { public void filter(ContainerRequestContext rc) { var requestId = rc.getHeaderString(loggingConfiguration.requestIdHeaderName()); if (requestId == null) { - requestId = UUID.randomUUID().toString(); + requestId = requestIdGenerator.generateRequestId(); } rc.setProperty(REQUEST_ID_KEY, requestId); } diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java new file mode 100644 index 0000000000..276a8636f3 --- /dev/null +++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java @@ -0,0 +1,48 @@ +/* + * 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.polaris.service.tracing; + +import com.google.common.annotations.VisibleForTesting; +import jakarta.enterprise.context.ApplicationScoped; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicLong; + +@ApplicationScoped +public class RequestIdGenerator { + private static String BASE_DEFAULT_UUID = UUID.randomUUID().toString(); + static final AtomicLong COUNTER = new AtomicLong(); + static final Long COUNTER_SOFT_MAX = Long.MAX_VALUE / 2; + + public String generateRequestId() { + String requestId = BASE_DEFAULT_UUID + "_" + COUNTER.incrementAndGet(); + if (COUNTER.get() >= COUNTER_SOFT_MAX) { + // If we get anywhere close to danger of overflowing Long (which, theoretically, + // would be extremely unlikely) rotate the UUID and start again. + BASE_DEFAULT_UUID = UUID.randomUUID().toString(); + COUNTER.set(0); + } + return requestId; + } + + @VisibleForTesting + public void setCounter(long counter) { + COUNTER.set(counter); + } +} diff --git a/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java b/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java index e6e9a36624..6d6307894d 100644 --- a/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java +++ b/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java @@ -23,6 +23,7 @@ import io.quarkus.test.junit.QuarkusTest; import io.quarkus.test.junit.QuarkusTestProfile; import io.quarkus.test.junit.TestProfile; +import jakarta.inject.Inject; import jakarta.ws.rs.client.Entity; import jakarta.ws.rs.core.MultivaluedHashMap; import jakarta.ws.rs.core.Response; @@ -32,6 +33,7 @@ import java.util.UUID; import org.apache.polaris.service.it.env.PolarisApiEndpoints; import org.apache.polaris.service.it.env.PolarisClient; +import org.apache.polaris.service.tracing.RequestIdGenerator; import org.junit.jupiter.api.Test; @QuarkusTest @@ -50,6 +52,8 @@ public Map getConfigOverrides() { } } + @Inject RequestIdGenerator requestIdGenerator; + private static final String REQUEST_ID_HEADER = "x-test-request-id-random"; private static final String REALM_HEADER = "realm"; private static final String REALM = "realm1"; @@ -90,30 +94,48 @@ private Response request(Map headers) { public void testRequestIdHeaderSpecified() { String requestId = "pre-requested-request-id"; Map headers = Map.of(REALM_HEADER, REALM, REQUEST_ID_HEADER, requestId); - try (Response response = request(headers)) { - assertThat(response.getHeaders()).containsKey(REQUEST_ID_HEADER); - assertThat(response.getHeaders().get(REQUEST_ID_HEADER)).hasSize(1); - assertThat(response.getHeaders().get(REQUEST_ID_HEADER)).allMatch(s -> s.equals(requestId)); - } + assertThat(sendRequest(headers)).matches(s -> s.equals(requestId)); } @Test public void testRequestIdHeaderNotSpecified() { Map headers = Map.of(REALM_HEADER, REALM); - try (Response response = request(headers)) { - assertThat(response.getHeaders()).containsKey(REQUEST_ID_HEADER); - assertThat(response.getHeaders().get(REQUEST_ID_HEADER)).hasSize(1); - assertThat(response.getHeaders().get(REQUEST_ID_HEADER)) - .allMatch(s -> isValidUUID(s.toString())); - } + assertThat(sendRequest(headers)).matches(this::isValidDefaultUUID); + } + + @Test + public void testRequestIdHeaderNotSpecifiedAndCounterExhausted() { + requestIdGenerator.setCounter(Long.MAX_VALUE / 2 + 1); + Map headers = Map.of(REALM_HEADER, REALM); + String requestId = sendRequest(headers); + assertThat(requestId).matches(this::isValidDefaultUUID); + String currentRequestIdBase = requestId.split("_")[0]; + String uuidBase = currentRequestIdBase; + + requestId = sendRequest(headers); + assertThat(requestId).matches(this::isValidDefaultUUID); + currentRequestIdBase = requestId.split("_")[0]; + assertThat(currentRequestIdBase).isNotEqualTo(uuidBase); } - private boolean isValidUUID(String str) { + private boolean isValidDefaultUUID(String str) { try { - UUID.fromString(str); + String[] requestIdParts = str.split("_"); + String uuid = requestIdParts[0]; + String counter = requestIdParts[1]; + UUID.fromString(uuid); + Long.parseLong(counter); return true; } catch (IllegalArgumentException e) { return false; } } + + private String sendRequest(Map headers) { + try (Response response = request(headers)) { + assertThat(response.getHeaders()).containsKey(REQUEST_ID_HEADER); + assertThat(response.getHeaders().get(REQUEST_ID_HEADER)).hasSize(1); + return response.getHeaders().get(REQUEST_ID_HEADER).getFirst().toString(); + } + } } From 6f90986c1a3112754e34813fe480726a5c3823c9 Mon Sep 17 00:00:00 2001 From: adnanhemani Date: Thu, 25 Sep 2025 15:49:42 -0700 Subject: [PATCH 07/10] review comments from @adutra --- .../service/tracing/RequestIdFilter.java | 2 +- .../service/tracing/RequestIdGenerator.java | 21 ++-- .../tracing/RequestIdGeneratorTest.java | 116 ++++++++++++++++++ .../RequestIdHeaderTest.java | 51 +++----- 4 files changed, 147 insertions(+), 43 deletions(-) create mode 100644 runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdGeneratorTest.java rename runtime/service/src/test/java/org/apache/polaris/service/{admin => tracing}/RequestIdHeaderTest.java (74%) diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java index ae1a60178d..973313b9de 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdFilter.java @@ -33,11 +33,11 @@ @Priority(FilterPriorities.REQUEST_ID_FILTER) @Provider public class RequestIdFilter implements ContainerRequestFilter { - @Inject RequestIdGenerator requestIdGenerator; public static final String REQUEST_ID_KEY = "requestId"; @Inject LoggingConfiguration loggingConfiguration; + @Inject RequestIdGenerator requestIdGenerator; @Override public void filter(ContainerRequestContext rc) { diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java index 276a8636f3..ede69d0f24 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java @@ -22,21 +22,28 @@ import com.google.common.annotations.VisibleForTesting; import jakarta.enterprise.context.ApplicationScoped; import java.util.UUID; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; @ApplicationScoped public class RequestIdGenerator { - private static String BASE_DEFAULT_UUID = UUID.randomUUID().toString(); + private volatile String baseDefaultUuid = UUID.randomUUID().toString(); static final AtomicLong COUNTER = new AtomicLong(); static final Long COUNTER_SOFT_MAX = Long.MAX_VALUE / 2; + private final AtomicBoolean resetInProgress = new AtomicBoolean(false); + public String generateRequestId() { - String requestId = BASE_DEFAULT_UUID + "_" + COUNTER.incrementAndGet(); - if (COUNTER.get() >= COUNTER_SOFT_MAX) { - // If we get anywhere close to danger of overflowing Long (which, theoretically, - // would be extremely unlikely) rotate the UUID and start again. - BASE_DEFAULT_UUID = UUID.randomUUID().toString(); - COUNTER.set(0); + long currentCounter = COUNTER.incrementAndGet(); + String requestId = baseDefaultUuid + "_" + currentCounter; + if (currentCounter >= COUNTER_SOFT_MAX) { + if (resetInProgress.compareAndSet(false, true)) { + // If we get anywhere close to danger of overflowing Long (which, theoretically, + // would be extremely unlikely) rotate the UUID and start again. + baseDefaultUuid = UUID.randomUUID().toString(); + COUNTER.set(0); + resetInProgress.set(false); + } } return requestId; } diff --git a/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdGeneratorTest.java b/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdGeneratorTest.java new file mode 100644 index 0000000000..8bb1c3099e --- /dev/null +++ b/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdGeneratorTest.java @@ -0,0 +1,116 @@ +/* + * 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.polaris.service.tracing; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.assertNotEquals; + +import java.util.HashSet; +import java.util.Set; +import java.util.UUID; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class RequestIdGeneratorTest { + + private RequestIdGenerator requestIdGenerator; + + @BeforeEach + void setUp() { + requestIdGenerator = new RequestIdGenerator(); + } + + @Test + void testGenerateRequestId_ReturnsValidFormat() { + String requestId = requestIdGenerator.generateRequestId(); + + assertThat(requestId).isNotNull(); + assertThat(requestId).matches(this::isValidRequestIdFormat); + assertThat(requestId).contains("_1"); // First call should increment counter to 1 + } + + @Test + void testGenerateRequestId_ReturnsUniqueIds() { + Set generatedIds = new HashSet<>(); + + // Generate multiple request IDs and verify they're all unique + for (int i = 0; i < 1000; i++) { + String requestId = requestIdGenerator.generateRequestId(); + assertThat(generatedIds).doesNotContain(requestId); + generatedIds.add(requestId); + } + + assertThat(generatedIds).hasSize(1000); + } + + @Test + void testCounterIncrementsSequentially() { + requestIdGenerator.setCounter(0); + + String firstId = requestIdGenerator.generateRequestId(); + String secondId = requestIdGenerator.generateRequestId(); + String thirdId = requestIdGenerator.generateRequestId(); + + assertThat(firstId).endsWith("_1"); + assertThat(secondId).endsWith("_2"); + assertThat(thirdId).endsWith("_3"); + } + + @Test + void testCounterRotationAtSoftMax() { + // Set counter close to soft max + long softMax = RequestIdGenerator.COUNTER_SOFT_MAX; + requestIdGenerator.setCounter(softMax - 1); + + String beforeRotation = requestIdGenerator.generateRequestId(); + String afterRotation = requestIdGenerator.generateRequestId(); + + // The UUID part should be different after rotation + String beforeUuidPart = beforeRotation.substring(0, beforeRotation.lastIndexOf('_')); + String afterUuidPart = afterRotation.substring(0, afterRotation.lastIndexOf('_')); + assertNotEquals(beforeUuidPart, afterUuidPart); + + // After rotation, counter should be reset and UUID should be different + assertThat(beforeRotation).endsWith("_" + softMax); + assertThat(afterRotation).endsWith("_1"); // Counter reset to 1 (after increment from 0) + } + + @Test + void testSetCounterChangesNextGeneratedId() { + requestIdGenerator.setCounter(100); + + String requestId = requestIdGenerator.generateRequestId(); + + assertThat(requestId).endsWith("_101"); // Should increment from set value + } + + private boolean isValidRequestIdFormat(String str) { + try { + String[] requestIdParts = str.split("_"); + String uuid = requestIdParts[0]; + String counter = requestIdParts[1]; + UUID.fromString(uuid); + Long.parseLong(counter); + return true; + } catch (IllegalArgumentException e) { + return false; + } + } +} diff --git a/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java b/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdHeaderTest.java similarity index 74% rename from runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java rename to runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdHeaderTest.java index 6d6307894d..ba64fddccc 100644 --- a/runtime/service/src/test/java/org/apache/polaris/service/admin/RequestIdHeaderTest.java +++ b/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdHeaderTest.java @@ -16,24 +16,24 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.polaris.service.admin; +package org.apache.polaris.service.tracing; import static org.assertj.core.api.Assertions.assertThat; import io.quarkus.test.junit.QuarkusTest; import io.quarkus.test.junit.QuarkusTestProfile; import io.quarkus.test.junit.TestProfile; -import jakarta.inject.Inject; import jakarta.ws.rs.client.Entity; import jakarta.ws.rs.core.MultivaluedHashMap; import jakarta.ws.rs.core.Response; import java.net.URI; +import java.util.HashMap; +import java.util.HashSet; import java.util.Map; import java.util.Objects; -import java.util.UUID; +import java.util.Set; import org.apache.polaris.service.it.env.PolarisApiEndpoints; import org.apache.polaris.service.it.env.PolarisClient; -import org.apache.polaris.service.tracing.RequestIdGenerator; import org.junit.jupiter.api.Test; @QuarkusTest @@ -52,8 +52,6 @@ public Map getConfigOverrides() { } } - @Inject RequestIdGenerator requestIdGenerator; - private static final String REQUEST_ID_HEADER = "x-test-request-id-random"; private static final String REALM_HEADER = "realm"; private static final String REALM = "realm1"; @@ -93,41 +91,24 @@ private Response request(Map headers) { @Test public void testRequestIdHeaderSpecified() { String requestId = "pre-requested-request-id"; - Map headers = Map.of(REALM_HEADER, REALM, REQUEST_ID_HEADER, requestId); + HashMap headers = + new HashMap<>(Map.of(REALM_HEADER, REALM, REQUEST_ID_HEADER, requestId)); + assertThat(sendRequest(headers)).matches(s -> s.equals(requestId)); assertThat(sendRequest(headers)).matches(s -> s.equals(requestId)); - } - @Test - public void testRequestIdHeaderNotSpecified() { - Map headers = Map.of(REALM_HEADER, REALM); - assertThat(sendRequest(headers)).matches(this::isValidDefaultUUID); + String newRequestId = "new-pre-requested-request-id"; + headers.put(REQUEST_ID_HEADER, newRequestId); + assertThat(sendRequest(headers)).matches(s -> s.equals(newRequestId)); } @Test - public void testRequestIdHeaderNotSpecifiedAndCounterExhausted() { - requestIdGenerator.setCounter(Long.MAX_VALUE / 2 + 1); + public void testRequestIdHeaderNotSpecified() { Map headers = Map.of(REALM_HEADER, REALM); - String requestId = sendRequest(headers); - assertThat(requestId).matches(this::isValidDefaultUUID); - String currentRequestIdBase = requestId.split("_")[0]; - String uuidBase = currentRequestIdBase; - - requestId = sendRequest(headers); - assertThat(requestId).matches(this::isValidDefaultUUID); - currentRequestIdBase = requestId.split("_")[0]; - assertThat(currentRequestIdBase).isNotEqualTo(uuidBase); - } - - private boolean isValidDefaultUUID(String str) { - try { - String[] requestIdParts = str.split("_"); - String uuid = requestIdParts[0]; - String counter = requestIdParts[1]; - UUID.fromString(uuid); - Long.parseLong(counter); - return true; - } catch (IllegalArgumentException e) { - return false; + Set requestIds = new HashSet<>(); + for (int i = 0; i < 10; i++) { + String requestId = sendRequest(headers); + assertThat(requestIds).doesNotContain(requestId); + requestIds.add(requestId); } } From 1eee9ff4c3a2039c437b701141fee437157eb448 Mon Sep 17 00:00:00 2001 From: adnanhemani Date: Thu, 25 Sep 2025 17:46:53 -0700 Subject: [PATCH 08/10] review comments from @flyrain --- .../polaris/service/config/FilterPriorities.java | 2 +- .../service/tracing/RequestIdGenerator.java | 14 ++++++-------- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/runtime/service/src/main/java/org/apache/polaris/service/config/FilterPriorities.java b/runtime/service/src/main/java/org/apache/polaris/service/config/FilterPriorities.java index 06fd801684..5ce1664092 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/config/FilterPriorities.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/config/FilterPriorities.java @@ -22,7 +22,7 @@ public final class FilterPriorities { public static final int REQUEST_ID_FILTER = Priorities.AUTHENTICATION - 101; - public static final int REALM_CONTEXT_FILTER = Priorities.AUTHENTICATION - 100; + public static final int REALM_CONTEXT_FILTER = REQUEST_ID_FILTER + 1; public static final int RATE_LIMITER_FILTER = Priorities.USER; public static final int MDC_FILTER = REALM_CONTEXT_FILTER + 1; public static final int TRACING_FILTER = REALM_CONTEXT_FILTER + 2; diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java index ede69d0f24..803f84a718 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java @@ -36,14 +36,12 @@ public class RequestIdGenerator { public String generateRequestId() { long currentCounter = COUNTER.incrementAndGet(); String requestId = baseDefaultUuid + "_" + currentCounter; - if (currentCounter >= COUNTER_SOFT_MAX) { - if (resetInProgress.compareAndSet(false, true)) { - // If we get anywhere close to danger of overflowing Long (which, theoretically, - // would be extremely unlikely) rotate the UUID and start again. - baseDefaultUuid = UUID.randomUUID().toString(); - COUNTER.set(0); - resetInProgress.set(false); - } + if (currentCounter >= COUNTER_SOFT_MAX && resetInProgress.compareAndSet(false, true)) { + // If we get anywhere close to danger of overflowing Long (which, theoretically, + // would be extremely unlikely) rotate the UUID and start again. + baseDefaultUuid = UUID.randomUUID().toString(); + COUNTER.set(0); + resetInProgress.set(false); } return requestId; } From f44b8907aa58469a80067e7b903454e1bb5d148b Mon Sep 17 00:00:00 2001 From: adnanhemani Date: Thu, 25 Sep 2025 18:29:46 -0700 Subject: [PATCH 09/10] added sync block --- .../service/tracing/RequestIdGenerator.java | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java index 803f84a718..044a265dc7 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java @@ -22,7 +22,6 @@ import com.google.common.annotations.VisibleForTesting; import jakarta.enterprise.context.ApplicationScoped; import java.util.UUID; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; @ApplicationScoped @@ -31,17 +30,18 @@ public class RequestIdGenerator { static final AtomicLong COUNTER = new AtomicLong(); static final Long COUNTER_SOFT_MAX = Long.MAX_VALUE / 2; - private final AtomicBoolean resetInProgress = new AtomicBoolean(false); - public String generateRequestId() { long currentCounter = COUNTER.incrementAndGet(); String requestId = baseDefaultUuid + "_" + currentCounter; - if (currentCounter >= COUNTER_SOFT_MAX && resetInProgress.compareAndSet(false, true)) { - // If we get anywhere close to danger of overflowing Long (which, theoretically, - // would be extremely unlikely) rotate the UUID and start again. - baseDefaultUuid = UUID.randomUUID().toString(); - COUNTER.set(0); - resetInProgress.set(false); + if (currentCounter >= COUNTER_SOFT_MAX) { + synchronized (this) { + if (COUNTER.get() >= COUNTER_SOFT_MAX) { + // If we get anywhere close to danger of overflowing Long (which, theoretically, + // would be extremely unlikely) rotate the UUID and start again. + baseDefaultUuid = UUID.randomUUID().toString(); + COUNTER.set(0); + } + } } return requestId; } From 588a26d6ad53dd49a6fc08002bd5cc6f13eb49a6 Mon Sep 17 00:00:00 2001 From: adnanhemani Date: Mon, 29 Sep 2025 15:16:15 -0700 Subject: [PATCH 10/10] change generator logic, per @adutra --- .../service/tracing/RequestIdGenerator.java | 37 ++++++++++--------- .../tracing/RequestIdGeneratorTest.java | 26 ++++++++----- 2 files changed, 36 insertions(+), 27 deletions(-) diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java index 044a265dc7..9ddd0b04c8 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdGenerator.java @@ -22,32 +22,35 @@ import com.google.common.annotations.VisibleForTesting; import jakarta.enterprise.context.ApplicationScoped; import java.util.UUID; -import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; @ApplicationScoped public class RequestIdGenerator { - private volatile String baseDefaultUuid = UUID.randomUUID().toString(); - static final AtomicLong COUNTER = new AtomicLong(); static final Long COUNTER_SOFT_MAX = Long.MAX_VALUE / 2; - public String generateRequestId() { - long currentCounter = COUNTER.incrementAndGet(); - String requestId = baseDefaultUuid + "_" + currentCounter; - if (currentCounter >= COUNTER_SOFT_MAX) { - synchronized (this) { - if (COUNTER.get() >= COUNTER_SOFT_MAX) { - // If we get anywhere close to danger of overflowing Long (which, theoretically, - // would be extremely unlikely) rotate the UUID and start again. - baseDefaultUuid = UUID.randomUUID().toString(); - COUNTER.set(0); - } - } + record State(String uuid, long counter) { + + State() { + this(UUID.randomUUID().toString(), 1); + } + + String requestId() { + return String.format("%s_%019d", uuid, counter); + } + + State increment() { + return counter >= COUNTER_SOFT_MAX ? new State() : new State(uuid, counter + 1); } - return requestId; + } + + final AtomicReference state = new AtomicReference<>(new State()); + + public String generateRequestId() { + return state.getAndUpdate(State::increment).requestId(); } @VisibleForTesting public void setCounter(long counter) { - COUNTER.set(counter); + state.set(new State(state.get().uuid, counter)); } } diff --git a/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdGeneratorTest.java b/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdGeneratorTest.java index 8bb1c3099e..08edfe4810 100644 --- a/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdGeneratorTest.java +++ b/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdGeneratorTest.java @@ -43,7 +43,8 @@ void testGenerateRequestId_ReturnsValidFormat() { assertThat(requestId).isNotNull(); assertThat(requestId).matches(this::isValidRequestIdFormat); - assertThat(requestId).contains("_1"); // First call should increment counter to 1 + // First call should increment counter to 1 + assertThat(extractCounterFromRequestId(requestId)).isEqualTo(1); } @Test @@ -62,22 +63,22 @@ void testGenerateRequestId_ReturnsUniqueIds() { @Test void testCounterIncrementsSequentially() { - requestIdGenerator.setCounter(0); + // requestIdGenerator.setCounter(0); String firstId = requestIdGenerator.generateRequestId(); String secondId = requestIdGenerator.generateRequestId(); String thirdId = requestIdGenerator.generateRequestId(); - assertThat(firstId).endsWith("_1"); - assertThat(secondId).endsWith("_2"); - assertThat(thirdId).endsWith("_3"); + assertThat(extractCounterFromRequestId(firstId)).isEqualTo(1); + assertThat(extractCounterFromRequestId(secondId)).isEqualTo(2); + assertThat(extractCounterFromRequestId(thirdId)).isEqualTo(3); } @Test void testCounterRotationAtSoftMax() { // Set counter close to soft max long softMax = RequestIdGenerator.COUNTER_SOFT_MAX; - requestIdGenerator.setCounter(softMax - 1); + requestIdGenerator.setCounter(softMax); String beforeRotation = requestIdGenerator.generateRequestId(); String afterRotation = requestIdGenerator.generateRequestId(); @@ -87,9 +88,9 @@ void testCounterRotationAtSoftMax() { String afterUuidPart = afterRotation.substring(0, afterRotation.lastIndexOf('_')); assertNotEquals(beforeUuidPart, afterUuidPart); - // After rotation, counter should be reset and UUID should be different - assertThat(beforeRotation).endsWith("_" + softMax); - assertThat(afterRotation).endsWith("_1"); // Counter reset to 1 (after increment from 0) + assertThat(extractCounterFromRequestId(beforeRotation)).isEqualTo(softMax); + // Counter reset to 1 (after increment from 0) + assertThat(extractCounterFromRequestId(afterRotation)).isEqualTo(1); } @Test @@ -98,7 +99,8 @@ void testSetCounterChangesNextGeneratedId() { String requestId = requestIdGenerator.generateRequestId(); - assertThat(requestId).endsWith("_101"); // Should increment from set value + // Should increment from set value + assertThat(extractCounterFromRequestId(requestId)).isEqualTo(100); } private boolean isValidRequestIdFormat(String str) { @@ -113,4 +115,8 @@ private boolean isValidRequestIdFormat(String str) { return false; } } + + private long extractCounterFromRequestId(String requestId) { + return Long.parseLong(requestId.split("_")[1]); + } }