diff --git a/runtime/service/src/main/java/org/apache/polaris/service/tracing/DefaultRequestIdGenerator.java b/runtime/service/src/main/java/org/apache/polaris/service/tracing/DefaultRequestIdGenerator.java
new file mode 100644
index 0000000000..47f17291fa
--- /dev/null
+++ b/runtime/service/src/main/java/org/apache/polaris/service/tracing/DefaultRequestIdGenerator.java
@@ -0,0 +1,69 @@
+/*
+ * 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 io.smallrye.mutiny.Uni;
+import jakarta.annotation.Nonnull;
+import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.ws.rs.container.ContainerRequestContext;
+import java.util.UUID;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * Default implementation of {@link RequestIdGenerator}, striking a balance between randomness and
+ * performance.
+ *
+ *
The IDs generated by this generator are of the form: {@code UUID_COUNTER}. The UUID part is
+ * randomly generated at startup, and the counter is incremented for each request.
+ *
+ *
In the unlikely event that the counter overflows, a new UUID is generated and the counter is
+ * reset to 1.
+ */
+@ApplicationScoped
+public class DefaultRequestIdGenerator implements RequestIdGenerator {
+
+ record RequestId(UUID uuid, long counter) {
+
+ RequestId() {
+ this(UUID.randomUUID(), 1);
+ }
+
+ @Override
+ @Nonnull
+ public String toString() {
+ return String.format("%s_%019d", uuid(), counter());
+ }
+
+ RequestId increment() {
+ return counter == Long.MAX_VALUE ? new RequestId() : new RequestId(uuid, counter + 1);
+ }
+ }
+
+ final AtomicReference state = new AtomicReference<>(new RequestId());
+
+ @Override
+ public Uni generateRequestId(ContainerRequestContext requestContext) {
+ return Uni.createFrom().item(nextRequestId().toString());
+ }
+
+ RequestId nextRequestId() {
+ return state.getAndUpdate(RequestId::increment);
+ }
+}
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 973313b9de..460a732331 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
@@ -18,33 +18,60 @@
*/
package org.apache.polaris.service.tracing;
-import jakarta.annotation.Priority;
-import jakarta.enterprise.context.ApplicationScoped;
+import io.smallrye.mutiny.Uni;
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 jakarta.ws.rs.container.ContainerResponseContext;
+import jakarta.ws.rs.core.MediaType;
+import jakarta.ws.rs.core.Response;
+import org.apache.iceberg.rest.responses.ErrorResponse;
import org.apache.polaris.service.config.FilterPriorities;
import org.apache.polaris.service.logging.LoggingConfiguration;
+import org.jboss.resteasy.reactive.server.ServerRequestFilter;
+import org.jboss.resteasy.reactive.server.ServerResponseFilter;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
-@PreMatching
-@ApplicationScoped
-@Priority(FilterPriorities.REQUEST_ID_FILTER)
-@Provider
-public class RequestIdFilter implements ContainerRequestFilter {
+public class RequestIdFilter {
public static final String REQUEST_ID_KEY = "requestId";
+ private static final Logger LOGGER = LoggerFactory.getLogger(RequestIdFilter.class);
+
@Inject LoggingConfiguration loggingConfiguration;
@Inject RequestIdGenerator requestIdGenerator;
- @Override
- public void filter(ContainerRequestContext rc) {
+ @ServerRequestFilter(preMatching = true, priority = FilterPriorities.REQUEST_ID_FILTER)
+ public Uni assignRequestId(ContainerRequestContext rc) {
var requestId = rc.getHeaderString(loggingConfiguration.requestIdHeaderName());
- if (requestId == null) {
- requestId = requestIdGenerator.generateRequestId();
+ return (requestId != null
+ ? Uni.createFrom().item(requestId)
+ : requestIdGenerator.generateRequestId(rc))
+ .onItem()
+ .invoke(id -> rc.setProperty(REQUEST_ID_KEY, id))
+ .onItemOrFailure()
+ .transform((id, error) -> error == null ? null : errorResponse(error));
+ }
+
+ @ServerResponseFilter
+ public void addResponseHeader(
+ ContainerRequestContext request, ContainerResponseContext response) {
+ String requestId = (String) request.getProperty(REQUEST_ID_KEY);
+ if (requestId != null) { // can be null if request ID generation fails
+ response.getHeaders().add(loggingConfiguration.requestIdHeaderName(), requestId);
}
- rc.setProperty(REQUEST_ID_KEY, requestId);
+ }
+
+ private static Response errorResponse(Throwable error) {
+ LOGGER.error("Failed to generate request ID", error);
+ return Response.status(Response.Status.INTERNAL_SERVER_ERROR)
+ .type(MediaType.APPLICATION_JSON_TYPE)
+ .entity(
+ ErrorResponse.builder()
+ .responseCode(Response.Status.INTERNAL_SERVER_ERROR.getStatusCode())
+ .withMessage("Request ID generation failed")
+ .withType("RequestIdGenerationError")
+ .build())
+ .build();
}
}
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 9ddd0b04c8..4264f0e381 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
@@ -19,38 +19,20 @@
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.AtomicReference;
+import io.smallrye.mutiny.Uni;
+import jakarta.ws.rs.container.ContainerRequestContext;
-@ApplicationScoped
-public class RequestIdGenerator {
- static final Long COUNTER_SOFT_MAX = Long.MAX_VALUE / 2;
-
- 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);
- }
- }
-
- final AtomicReference state = new AtomicReference<>(new State());
-
- public String generateRequestId() {
- return state.getAndUpdate(State::increment).requestId();
- }
-
- @VisibleForTesting
- public void setCounter(long counter) {
- state.set(new State(state.get().uuid, counter));
- }
+/**
+ * A generator for request IDs.
+ *
+ * @see RequestIdFilter
+ */
+public interface RequestIdGenerator {
+
+ /**
+ * Generates a new request ID. IDs must be fast to generate and unique.
+ *
+ * @param requestContext The JAX-RS request context
+ */
+ Uni generateRequestId(ContainerRequestContext requestContext);
}
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
deleted file mode 100644
index 728960182d..0000000000
--- a/runtime/service/src/main/java/org/apache/polaris/service/tracing/RequestIdResponseFilter.java
+++ /dev/null
@@ -1,46 +0,0 @@
-/*
- * 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.enterprise.context.ApplicationScoped;
-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;
-
-@ApplicationScoped
-@Provider
-public class RequestIdResponseFilter implements ContainerResponseFilter {
-
- @Inject LoggingConfiguration loggingConfiguration;
-
- @Override
- public void filter(
- ContainerRequestContext requestContext, ContainerResponseContext responseContext) {
- responseContext
- .getHeaders()
- .add(
- loggingConfiguration.requestIdHeaderName(), 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 8b6859ca36..d794317d47 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
@@ -47,9 +47,7 @@ public void filter(ContainerRequestContext rc) {
if (!sdkDisabled) {
Span span = Span.current();
String requestId = (String) rc.getProperty(RequestIdFilter.REQUEST_ID_KEY);
- if (requestId != null) {
- span.setAttribute(REQUEST_ID_ATTRIBUTE, requestId);
- }
+ span.setAttribute(REQUEST_ID_ATTRIBUTE, requestId);
RealmContext realmContext =
(RealmContext) rc.getProperty(RealmContextFilter.REALM_CONTEXT_KEY);
span.setAttribute(REALM_ID_ATTRIBUTE, realmContext.getRealmIdentifier());
diff --git a/runtime/service/src/test/java/org/apache/polaris/service/tracing/DefaultRequestIdGeneratorTest.java b/runtime/service/src/test/java/org/apache/polaris/service/tracing/DefaultRequestIdGeneratorTest.java
new file mode 100644
index 0000000000..13dc54baed
--- /dev/null
+++ b/runtime/service/src/test/java/org/apache/polaris/service/tracing/DefaultRequestIdGeneratorTest.java
@@ -0,0 +1,83 @@
+/*
+ * 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 java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentSkipListSet;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import org.apache.polaris.service.tracing.DefaultRequestIdGenerator.RequestId;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class DefaultRequestIdGeneratorTest {
+
+ private DefaultRequestIdGenerator requestIdGenerator;
+
+ @BeforeEach
+ void setUp() {
+ requestIdGenerator = new DefaultRequestIdGenerator();
+ }
+
+ @Test
+ void testGeneratesUniqueIds() {
+ Set generatedIds = new ConcurrentSkipListSet<>();
+ try (ExecutorService executor = Executors.newFixedThreadPool(10)) {
+ for (int i = 0; i < 1000; i++) {
+ executor.execute(() -> generatedIds.add(requestIdGenerator.nextRequestId().toString()));
+ }
+ }
+ assertThat(generatedIds).hasSize(1000);
+ }
+
+ @Test
+ void testCounterIncrementsSequentially() {
+ assertThat(requestIdGenerator.nextRequestId().counter()).isEqualTo(1L);
+ assertThat(requestIdGenerator.nextRequestId().counter()).isEqualTo(2L);
+ assertThat(requestIdGenerator.nextRequestId().counter()).isEqualTo(3L);
+ }
+
+ @Test
+ void testCounterRotationAtMax() {
+ requestIdGenerator.state.set(new RequestId(UUID.randomUUID(), Long.MAX_VALUE));
+
+ var beforeRotation = requestIdGenerator.nextRequestId();
+ var afterRotation = requestIdGenerator.nextRequestId();
+
+ // The UUID should be different after rotation
+ assertThat(beforeRotation.uuid()).isNotEqualTo(afterRotation.uuid());
+
+ // The counter should be reset to 1 after rotation
+ assertThat(beforeRotation.counter()).isEqualTo(Long.MAX_VALUE);
+ assertThat(afterRotation.counter()).isEqualTo(1L);
+ }
+
+ @Test
+ void testRequestIdToString() {
+ var uuid = UUID.randomUUID();
+ assertThat(new RequestId(uuid, 1L).toString()).isEqualTo(uuid + "_0000000000000000001");
+ assertThat(new RequestId(uuid, 12345L).toString()).isEqualTo(uuid + "_0000000000000012345");
+ assertThat(new RequestId(uuid, Long.MAX_VALUE).toString())
+ .isEqualTo(uuid + "_9223372036854775807");
+ }
+}
diff --git a/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdFilterTest.java b/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdFilterTest.java
new file mode 100644
index 0000000000..ddcf5894ba
--- /dev/null
+++ b/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdFilterTest.java
@@ -0,0 +1,104 @@
+/*
+ * 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 io.restassured.RestAssured.given;
+import static org.hamcrest.CoreMatchers.anything;
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.is;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import io.quarkus.test.common.http.TestHTTPEndpoint;
+import io.quarkus.test.junit.QuarkusTest;
+import io.quarkus.test.junit.mockito.InjectSpy;
+import io.restassured.http.ContentType;
+import io.restassured.specification.RequestSpecification;
+import io.smallrye.mutiny.Uni;
+import org.apache.polaris.service.catalog.api.IcebergRestOAuth2Api;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+@QuarkusTest
+@TestHTTPEndpoint(IcebergRestOAuth2Api.class)
+@SuppressWarnings("UastIncorrectHttpHeaderInspection")
+public class RequestIdFilterTest {
+
+ @InjectSpy RequestIdGenerator requestIdGenerator;
+
+ @BeforeEach
+ void resetMocks() {
+ Mockito.reset(requestIdGenerator);
+ }
+
+ @Test
+ void testSuccessWithGeneratedRequestId() {
+ givenTokenRequest()
+ .when()
+ .post()
+ .then()
+ .statusCode(200)
+ .body(containsString("access_token"))
+ .header("Polaris-Request-Id", anything());
+ verify(requestIdGenerator, times(1)).generateRequestId(any());
+ }
+
+ @Test
+ void testSuccessWithCustomRequestId() {
+ givenTokenRequest()
+ .header("Polaris-Request-Id", "custom-request-id")
+ .when()
+ .post()
+ .then()
+ .statusCode(200)
+ .body(containsString("access_token"))
+ .header("Polaris-Request-Id", "custom-request-id");
+ verify(requestIdGenerator, never()).generateRequestId(any());
+ }
+
+ @Test
+ void testError() {
+ doReturn(Uni.createFrom().failure(new RuntimeException("test error")))
+ .when(requestIdGenerator)
+ .generateRequestId(any());
+ givenTokenRequest()
+ .when()
+ .post()
+ .then()
+ .statusCode(500)
+ .body("error.message", is("Request ID generation failed"))
+ .body("error.type", is("RequestIdGenerationError"))
+ .body("error.code", is(500));
+ verify(requestIdGenerator, times(1)).generateRequestId(any());
+ }
+
+ private static RequestSpecification givenTokenRequest() {
+ return given()
+ .contentType(ContentType.URLENC)
+ .formParam("grant_type", "client_credentials")
+ .formParam("scope", "PRINCIPAL_ROLE:ALL")
+ .formParam("client_id", "test-admin")
+ .formParam("client_secret", "test-secret");
+ }
+}
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
deleted file mode 100644
index 08edfe4810..0000000000
--- a/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdGeneratorTest.java
+++ /dev/null
@@ -1,122 +0,0 @@
-/*
- * 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);
- // First call should increment counter to 1
- assertThat(extractCounterFromRequestId(requestId)).isEqualTo(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(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);
-
- 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);
-
- assertThat(extractCounterFromRequestId(beforeRotation)).isEqualTo(softMax);
- // Counter reset to 1 (after increment from 0)
- assertThat(extractCounterFromRequestId(afterRotation)).isEqualTo(1);
- }
-
- @Test
- void testSetCounterChangesNextGeneratedId() {
- requestIdGenerator.setCounter(100);
-
- String requestId = requestIdGenerator.generateRequestId();
-
- // Should increment from set value
- assertThat(extractCounterFromRequestId(requestId)).isEqualTo(100);
- }
-
- 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;
- }
- }
-
- private long extractCounterFromRequestId(String requestId) {
- return Long.parseLong(requestId.split("_")[1]);
- }
-}
diff --git a/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdHeaderTest.java b/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdHeaderTest.java
index ba64fddccc..8a653476df 100644
--- a/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdHeaderTest.java
+++ b/runtime/service/src/test/java/org/apache/polaris/service/tracing/RequestIdHeaderTest.java
@@ -27,7 +27,6 @@
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;
@@ -91,14 +90,13 @@ private Response request(Map headers) {
@Test
public void testRequestIdHeaderSpecified() {
String requestId = "pre-requested-request-id";
- 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));
+ Map headers = Map.of(REALM_HEADER, REALM, REQUEST_ID_HEADER, requestId);
+ assertThat(sendRequest(headers)).isEqualTo(requestId);
+ assertThat(sendRequest(headers)).isEqualTo(requestId);
String newRequestId = "new-pre-requested-request-id";
- headers.put(REQUEST_ID_HEADER, newRequestId);
- assertThat(sendRequest(headers)).matches(s -> s.equals(newRequestId));
+ headers = Map.of(REALM_HEADER, REALM, REQUEST_ID_HEADER, newRequestId);
+ assertThat(sendRequest(headers)).isEqualTo(newRequestId);
}
@Test
@@ -106,10 +104,9 @@ public void testRequestIdHeaderNotSpecified() {
Map headers = Map.of(REALM_HEADER, REALM);
Set requestIds = new HashSet<>();
for (int i = 0; i < 10; i++) {
- String requestId = sendRequest(headers);
- assertThat(requestIds).doesNotContain(requestId);
- requestIds.add(requestId);
+ requestIds.add(sendRequest(headers));
}
+ assertThat(requestIds).hasSize(10);
}
private String sendRequest(Map headers) {