perf(spanner): stop re-parsing the request id on every RPC (#14353)

RequestIdTargetTracker keys its cache on the logical request key, but it
accepted the x-goog-spanner-request-id header value and derived that key
by calling XGoogSpannerRequestId.of(String), which runs a six-group
regex plus four Long.parseLong calls. Every caller already holds the
parsed XGoogSpannerRequestId in the gRPC CallOptions.

HeaderInterceptor paid this twice per RPC (a get on the first response
and a remove on close). KeyAwareChannel added a third by rendering the
object it already held back into a header string with getHeaderValue(),
only for the tracker to parse it again.

Location-aware routing is the only producer of tracked targets and is
disabled unless the instance type is OMNI, so for most clients the cache
stays empty for the lifetime of the process and all of that work is
wasted. The tracker now accepts the parsed request id (or the logical
key, where the caller already has one) and short circuits on a volatile
flag that is set the first time a target is recorded.

This also fixes a latent bug in KeyAwareChannel.onClose, which passed an
already-normalized key into remove(String). The regex never matched it,
so every RPC on that path constructed and threw an
IllegalStateException, including its concatenated message and stack
trace fill-in, before the catch block fell back to returning the raw
string.

Measured per RPC (JMH, 2 forks x 8 x 1s, -prof gc):

```
  before                     581.2 ns   1376 B
  after, default config        2.0 ns      0 B
  after, location-aware on    42.5 ns     56 B
```

Author: olavloite

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java

@@ -17,6 +17,7 @@
 
 import static com.google.api.gax.grpc.GrpcCallContext.TRACER_KEY;
 import static com.google.cloud.spanner.BuiltInMetricsConstant.UNDEFINED_PROJECT_ID;
+import static com.google.cloud.spanner.XGoogSpannerRequestId.REQUEST_ID_CALL_OPTIONS_KEY;
 import static com.google.cloud.spanner.spi.v1.SpannerRpcViews.DATABASE_ID;
 import static com.google.cloud.spanner.spi.v1.SpannerRpcViews.INSTANCE_ID;
 import static com.google.cloud.spanner.spi.v1.SpannerRpcViews.METHOD;
@@ -111,6 +112,9 @@ public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
     ApiTracer tracer = callOptions.getOption(TRACER_KEY);
     CompositeTracer compositeTracer =
         tracer instanceof CompositeTracer ? (CompositeTracer) tracer : null;
+    // The same request id that RequestIdInterceptor writes to the request header. Reading it from
+    // the call options avoids having to parse the header value back into an object on every RPC.
+    XGoogSpannerRequestId parsedRequestId = callOptions.getOption(REQUEST_ID_CALL_OPTIONS_KEY);
     return new SimpleForwardingClientCall<ReqT, RespT>(next.newCall(method, callOptions)) {
       @Override
       public void start(Listener<RespT> responseListener, Metadata headers) {
@@ -133,7 +137,8 @@ public void start(Listener<RespT> responseListener, Metadata headers) {
                 @Override
                 public void onHeaders(Metadata metadata) {
                   try {
-                    recordFirstResponseLatency(requestId, startedAtNanos, firstResponseRecorded);
+                    recordFirstResponseLatency(
+                        parsedRequestId, startedAtNanos, firstResponseRecorded);
                     String serverTiming = metadata.get(SERVER_TIMING_HEADER_KEY);
                     try {
                       // Get gfe and afe Latency value
@@ -170,7 +175,8 @@ public void onClose(Status status, Metadata trailers) {
                           LEVEL, "Unable to get built-in metric attributes {0}", e.getMessage());
                     }
                     if (status.isOk()) {
-                      recordFirstResponseLatency(requestId, startedAtNanos, firstResponseRecorded);
+                      recordFirstResponseLatency(
+                          parsedRequestId, startedAtNanos, firstResponseRecorded);
                     }
                     recordBuiltInMetrics(
                         compositeTracer,
@@ -183,7 +189,7 @@ public void onClose(Status status, Metadata trailers) {
                   } catch (Throwable throwable) {
                     LOGGER.log(Level.WARNING, "Error recording metrics in onClose", throwable);
                   } finally {
-                    RequestIdTargetTracker.remove(requestId);
+                    RequestIdTargetTracker.remove(parsedRequestId);
                     super.onClose(status, trailers);
                   }
                 }
@@ -263,7 +269,9 @@ private void recordBuiltInMetrics(
   }
 
   private void recordFirstResponseLatency(
-      String requestId, long startedAtNanos, AtomicBoolean firstResponseRecorded) {
+      @Nullable XGoogSpannerRequestId requestId,
+      long startedAtNanos,
+      AtomicBoolean firstResponseRecorded) {
     if (!firstResponseRecorded.compareAndSet(false, true)) {
       return;
     }

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java

@@ -913,15 +913,12 @@ public void sendMessage(RequestT message) {
         selectedPreferLeader = preferLeader;
         this.channelFinder = finder;
         selectedEndpoint.incrementActiveRequests();
-        XGoogSpannerRequestId requestId = callOptions.getOption(REQUEST_ID_CALL_OPTIONS_KEY);
-        if (requestId != null) {
-          RequestIdTargetTracker.record(
-              requestId.getHeaderValue(),
-              selectedDatabaseScope,
-              selectedTargetEndpoint,
-              operationUid,
-              selectedPreferLeader);
-        }
+        RequestIdTargetTracker.record(
+            logicalRequestKey,
+            selectedDatabaseScope,
+            selectedTargetEndpoint,
+            operationUid,
+            selectedPreferLeader);
 
         // Record real traffic for idle eviction tracking.
         parentChannel.onRequestRouted(endpoint);
@@ -1105,11 +1102,7 @@ boolean maybeRerouteUnimplementedGenericCall(io.grpc.Status status) {
         selectedPreferLeader = false;
         channelFinder = null;
         selectedEndpoint.incrementActiveRequests();
-        XGoogSpannerRequestId requestId = callOptions.getOption(REQUEST_ID_CALL_OPTIONS_KEY);
-        if (requestId != null) {
-          RequestIdTargetTracker.record(
-              requestId.getHeaderValue(), null, selectedTargetEndpoint, 0L, false);
-        }
+        RequestIdTargetTracker.record(logicalRequestKey, null, selectedTargetEndpoint, 0L, false);
         parentChannel.onRequestRouted(defaultEndpoint);
         recordRouteSelectionTrace(methodDescriptor, defaultEndpoint.getAddress(), true, false);
 
@@ -1326,7 +1319,7 @@ public void onClose(io.grpc.Status status, Metadata trailers) {
       if (call.selectedEndpoint != null) {
         call.selectedEndpoint.decrementActiveRequests();
       }
-      RequestIdTargetTracker.remove(call.logicalRequestKey);
+      RequestIdTargetTracker.removeLogicalKey(call.logicalRequestKey);
       call.maybeClearAffinity();
       super.onClose(status, trailers);
     }

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdTargetTracker.java

@@ -23,6 +23,16 @@
 import java.util.concurrent.TimeUnit;
 import javax.annotation.Nullable;
 
+/**
+ * Associates an in-flight request with the endpoint that location-aware routing selected for it, so
+ * that {@link HeaderInterceptor} can attribute the observed latency back to that endpoint.
+ *
+ * <p>Entries are keyed by {@link XGoogSpannerRequestId#getLogicalRequestKey()}, which is stable
+ * across all attempts of the same logical RPC. Callers pass that key directly rather than the
+ * {@code x-goog-spanner-request-id} header value: every call site either holds the parsed {@link
+ * XGoogSpannerRequestId} (it is carried in the gRPC {@code CallOptions}) or has already derived the
+ * key from it, so re-parsing the header string here would be pure overhead on the per-RPC path.
+ */
 final class RequestIdTargetTracker {
   @VisibleForTesting static final long MAX_TRACKED_TARGETS = 1_000_000L;
 
@@ -37,54 +47,78 @@ final class RequestIdTargetTracker {
           .expireAfterWrite(10, TimeUnit.MINUTES)
           .build();
 
+  /**
+   * Whether a routing target has ever been recorded. Location-aware routing is the only producer of
+   * entries and is disabled unless the instance type is {@code OMNI}, so for most clients this
+   * cache stays empty for the lifetime of the process. Checking this flag first lets the per-RPC
+   * lookup and removal return immediately instead of probing the cache.
+   *
+   * <p>The flag only ever transitions from {@code false} to {@code true}, and every RPC reads it
+   * from whichever thread it happens to run on. It is therefore written only when it is not already
+   * set, so that recording a target does not repeatedly invalidate the cache line that all those
+   * readers share. Visibility of the entries themselves does not depend on this flag; {@link Cache}
+   * provides its own guarantees.
+   */
+  private static volatile boolean tracking;
+
   private RequestIdTargetTracker() {}
 
   static void record(
-      String requestId,
+      @Nullable String logicalRequestKey,
       @Nullable String databaseScope,
-      String targetEndpoint,
+      @Nullable String targetEndpoint,
       long operationUid,
       boolean preferLeader) {
-    String trackingKey = normalizeRequestKey(requestId);
-    if (trackingKey == null || targetEndpoint == null || targetEndpoint.isEmpty()) {
+    if (logicalRequestKey == null
+        || logicalRequestKey.isEmpty()
+        || targetEndpoint == null
+        || targetEndpoint.isEmpty()) {
       return;
     }
     TARGETS.put(
-        trackingKey, new RoutingTarget(databaseScope, targetEndpoint, operationUid, preferLeader));
+        logicalRequestKey,
+        new RoutingTarget(databaseScope, targetEndpoint, operationUid, preferLeader));
+    if (!tracking) {
+      tracking = true;
+    }
   }
 
   @Nullable
-  static RoutingTarget get(String requestId) {
-    String trackingKey = normalizeRequestKey(requestId);
-    if (trackingKey == null) {
-      return null;
-    }
-    return TARGETS.getIfPresent(trackingKey);
+  static RoutingTarget get(@Nullable XGoogSpannerRequestId requestId) {
+    String logicalRequestKey = trackingKey(requestId);
+    return logicalRequestKey == null ? null : TARGETS.getIfPresent(logicalRequestKey);
+  }
+
+  static void remove(@Nullable XGoogSpannerRequestId requestId) {
+    removeLogicalKey(trackingKey(requestId));
   }
 
-  static void remove(String requestId) {
-    String trackingKey = normalizeRequestKey(requestId);
-    if (trackingKey == null) {
+  static void removeLogicalKey(@Nullable String logicalRequestKey) {
+    if (!tracking || logicalRequestKey == null || logicalRequestKey.isEmpty()) {
       return;
     }
-    TARGETS.invalidate(trackingKey);
+    TARGETS.invalidate(logicalRequestKey);
+  }
+
+  /**
+   * Returns the cache key for {@code requestId}, or {@code null} if the key cannot be derived or
+   * nothing is being tracked. Deriving the key allocates a string, so it is skipped entirely while
+   * the cache is known to be empty.
+   */
+  @Nullable
+  private static String trackingKey(@Nullable XGoogSpannerRequestId requestId) {
+    return tracking && requestId != null ? requestId.getLogicalRequestKey() : null;
   }
 
   @VisibleForTesting
-  static void clear() {
-    TARGETS.invalidateAll();
+  static boolean isTracking() {
+    return tracking;
   }
 
   @VisibleForTesting
-  static String normalizeRequestKey(String requestId) {
-    if (requestId == null || requestId.isEmpty()) {
-      return null;
-    }
-    try {
-      return XGoogSpannerRequestId.of(requestId).getLogicalRequestKey();
-    } catch (IllegalStateException e) {
-      return requestId;
-    }
+  static void clear() {
+    TARGETS.invalidateAll();
+    tracking = false;
   }
 
   static final class RoutingTarget {

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/HeaderInterceptorTest.java

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner.spi.v1;
 
 import static com.google.api.gax.grpc.GrpcCallContext.TRACER_KEY;
+import static com.google.cloud.spanner.XGoogSpannerRequestId.REQUEST_ID_CALL_OPTIONS_KEY;
 import static com.google.cloud.spanner.XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
@@ -25,6 +26,7 @@
 
 import com.google.cloud.spanner.CompositeTracer;
 import com.google.cloud.spanner.SpannerRpcMetrics;
+import com.google.cloud.spanner.XGoogSpannerRequestId;
 import com.google.common.collect.ImmutableList;
 import io.grpc.CallOptions;
 import io.grpc.Channel;
@@ -170,18 +172,22 @@ public void recordServerTimingHeaderMetrics(
           }
         };
 
-    CallOptions callOptions = CallOptions.DEFAULT.withOption(TRACER_KEY, throwingTracer);
+    XGoogSpannerRequestId requestId = XGoogSpannerRequestId.of(1, 1, 1, 1);
+    CallOptions callOptions =
+        CallOptions.DEFAULT
+            .withOption(TRACER_KEY, throwingTracer)
+            .withOption(REQUEST_ID_CALL_OPTIONS_KEY, requestId);
     MethodDescriptor<String, String> methodDescriptor = createMethodDescriptor();
     FakeChannel channel = new FakeChannel();
 
-    String requestId = "1.0000000000000001.1.1.1.1";
-    RequestIdTargetTracker.record(requestId, "test-database", "endpoint-1", 100L, false);
+    RequestIdTargetTracker.record(
+        requestId.getLogicalRequestKey(), "test-database", "endpoint-1", 100L, false);
     assertNotNull(RequestIdTargetTracker.get(requestId));
 
     ClientCall<String, String> call =
         interceptor.interceptCall(methodDescriptor, callOptions, channel);
     CapturingListener<String> responseListener = new CapturingListener<>();
-    call.start(responseListener, createDefaultHeaders(requestId));
+    call.start(responseListener, createDefaultHeaders(requestId.getHeaderValue()));
 
     // Deliver onClose - even though metric recording throws, onClose must propagate downstream
     channel.lastListener.onClose(Status.OK, new Metadata());