Extract interface for RequestIdGenerator

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.
+ *
+ * <p>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.
+ *
+ * <p>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<RequestId> state = new AtomicReference<>(new RequestId());
+
+  @Override
+  public Uni<String> generateRequestId(ContainerRequestContext requestContext) {
+    return Uni.createFrom().item(nextRequestId().toString());
+  }
+
+  RequestId nextRequestId() {
+    return state.getAndUpdate(RequestId::increment);
+  }
+}

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<Response> 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();
   }
 }

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> 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<String> generateRequestId(ContainerRequestContext requestContext);
 }

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());

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<String> 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");
+  }
+}