[SPARK-59928] Fix kubernetes.client.http.response.latency.nanos always recording near zero

Author: ykisana

docs/configuration.md

@@ -263,6 +263,12 @@ has a retry left for it, even if it does not retry it, e.g. on a request timeout
 a client which never retries, i.e. a Kubernetes event write or the recording of the Kueue
 `PodsReady` condition, is not counted there when it gets no response.
 
+`kubernetes.client.http.response.latency.nanos` ends when the response headers arrive, so it does
+not include reading the body. A request re-sent after a `401` is recorded twice, and the two
+values add up to the time from sending it until the re-sent response arrives, including the token
+refresh. WebSocket upgrades, e.g. of a watch, are not recorded, and a rejected one can be counted
+more than once in `kubernetes.client.http.response`.
+
 ### Latency for State Transition
 
 Spark Operator also measures the latency between each state transition for apps, in the format of

docs/migration_guide.md

@@ -116,6 +116,12 @@ affect building the project, CI, tests or examples.
   `NullPointerException` for most properties. The operator logs each invalid value once as a
   warning, instead of as an error with a stack trace whenever it reads the value. Use `true` or
   `false`, which behave the same in both versions ([SPARK-XXXXX](https://issues.apache.org/jira/browse/SPARK-XXXXX)).
+- Since 1.1.0, `kubernetes.client.http.response.latency.nanos` records the latency of HTTP
+  responses until their headers arrive. 1.0 recorded nearly zero regardless of the latency. A
+  request re-sent after a `401` is recorded twice, and the two values add up to the time from
+  sending it until the re-sent response arrives, including the token refresh. It no longer records
+  a WebSocket upgrade response, e.g. of a watch, so its count can be
+  lower than that of `kubernetes.client.http.response` ([SPARK-59928](https://issues.apache.org/jira/browse/SPARK-59928)).
 
 ### SparkApplication
 

spark-operator/src/main/java/org/apache/spark/k8s/operator/metrics/source/KubernetesMetricsInterceptor.java

@@ -29,6 +29,7 @@
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
 
 import com.codahale.metrics.Histogram;
 import com.codahale.metrics.Meter;
@@ -90,6 +91,25 @@ public KubernetesMetricsInterceptor() {
     }
   }
 
+  /**
+   * Called before a request is sent, to allow wrapping the consumer of the response body.
+   *
+   * <p>Wraps the consumer to carry the request start time, which {@link #after(HttpRequest,
+   * HttpResponse, AsyncBody.Consumer)} finds with {@link AsyncBody.Consumer#unwrap(Class)}. fabric8
+   * passes to it the outermost consumer after all interceptors have wrapped it, so an interceptor
+   * which wraps the consumer after this one must delegate {@code unwrap(...)}, or no latency is
+   * recorded.
+   *
+   * @param consumer the consumer of the response body
+   * @param request the request about to be sent
+   * @return the consumer carrying the request start time
+   */
+  @Override
+  public AsyncBody.Consumer<List<ByteBuffer>> consumer(
+      AsyncBody.Consumer<List<ByteBuffer>> consumer, HttpRequest request) {
+    return new TimedConsumer(consumer, this, new AtomicLong(System.nanoTime()));
+  }
+
   /**
    * Called before a request to allow for the manipulation of the request
    *
@@ -102,21 +122,35 @@ public void before(BasicBuilder builder, HttpRequest request, RequestTags tags)
   }
 
   /**
-   * Called after a non-WebSocket HTTP response is received. The body might or might not be already
-   * consumed.
+   * Called after an HTTP response is received. The body might or might not be already consumed.
    *
    * <p>Should be used to analyze response codes and headers, original response shouldn't be
    * altered.
    *
+   * <p>fabric8 also calls this for a WebSocket upgrade response with a {@code null} consumer, and
+   * again for an HTTP request re-sent after {@link #afterFailure(BasicBuilder, HttpResponse,
+   * RequestTags)}. The two values of a request re-sent this way add up to the time from sending it
+   * until the re-sent response arrives.
+   *
    * @param request the original request sent to the server.
    * @param response the response received from the server.
+   * @param consumer the consumer of the response body, or {@code null} for a WebSocket upgrade.
    */
   @Override
+  @SuppressWarnings("PMD.CompareObjectsWithEquals")
   public void after(
       HttpRequest request,
       HttpResponse<?> response,
       AsyncBody.Consumer<List<ByteBuffer>> consumer) {
-    updateResponseMetrics(response, System.nanoTime());
+    TimedConsumer timedConsumer = consumer == null ? null : consumer.unwrap(TimedConsumer.class);
+    while (timedConsumer != null && timedConsumer.owner() != this) {
+      timedConsumer = timedConsumer.delegate().unwrap(TimedConsumer.class);
+    }
+    if (timedConsumer != null) {
+      long now = System.nanoTime();
+      responseLatency.update(now - timedConsumer.startTimeNanos().getAndSet(now));
+    }
+    updateResponseMetrics(response);
   }
 
   /**
@@ -139,7 +173,8 @@ public CompletableFuture<Boolean> afterFailure(
   /**
    * Called after a connection attempt fails.
    *
-   * <p>This method will be invoked on each failed connection attempt.
+   * <p>fabric8 invokes this method only for a failed connection attempt which has a retry left, not
+   * for the last attempt.
    *
    * @param request the HTTP request.
    * @param failure the Java exception that caused the failure.
@@ -183,11 +218,9 @@ private void updateRequestMetrics(HttpRequest request) {
         });
   }
 
-  private void updateResponseMetrics(HttpResponse response, long startTimeNanos) {
+  private void updateResponseMetrics(HttpResponse response) {
     Objects.requireNonNull(response);
-    final long latency = System.nanoTime() - startTimeNanos;
     responseRateMeter.mark();
-    responseLatency.update(latency);
     getMeterByResponseCode(response.code()).mark();
     if (KUBERNETES_CLIENT_METRICS_GROUP_BY_RESPONSE_CODE_GROUP_ENABLED.getValue()) {
       responseCodeGroupMeters.get(response.code() / 100 - 1).mark();
@@ -229,4 +262,26 @@ public Optional<Pair<String, String>> parseNamespaceScopedResource(String path)
       return Optional.empty();
     }
   }
+
+  /**
+   * Consumer of a response body which carries the start time of its request, restarted at each
+   * response so that the values of a re-sent request add up to its latency. It also carries the
+   * interceptor which created it, so that each interceptor on a client uses its own start time.
+   */
+  private record TimedConsumer(
+      AsyncBody.Consumer<List<ByteBuffer>> delegate,
+      KubernetesMetricsInterceptor owner,
+      AtomicLong startTimeNanos)
+      implements AsyncBody.Consumer<List<ByteBuffer>> {
+    @Override
+    public void consume(List<ByteBuffer> value, AsyncBody asyncBody) throws Exception {
+      delegate.consume(value, asyncBody);
+    }
+
+    @Override
+    public <U> U unwrap(Class<U> target) {
+      U self = AsyncBody.Consumer.super.unwrap(target);
+      return self != null ? self : delegate.unwrap(target);
+    }
+  }
 }

spark-operator/src/test/java/org/apache/spark/k8s/operator/metrics/source/KubernetesMetricsInterceptorTest.java

@@ -19,19 +19,29 @@
 
 package org.apache.spark.k8s.operator.metrics.source;
 
+import static java.net.HttpURLConnection.HTTP_UNAUTHORIZED;
+import static org.awaitility.Awaitility.await;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
+import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
+import com.codahale.metrics.Histogram;
 import com.codahale.metrics.Meter;
 import com.codahale.metrics.Metric;
+import com.codahale.metrics.Snapshot;
 import io.fabric8.kubernetes.api.model.ConfigMap;
 import io.fabric8.kubernetes.api.model.ObjectMeta;
+import io.fabric8.kubernetes.api.model.StatusBuilder;
+import io.fabric8.kubernetes.client.Config;
+import io.fabric8.kubernetes.client.ConfigBuilder;
 import io.fabric8.kubernetes.client.KubernetesClient;
+import io.fabric8.kubernetes.client.http.AsyncBody;
 import io.fabric8.kubernetes.client.http.Interceptor;
+import io.fabric8.kubernetes.client.informers.SharedIndexInformer;
 import io.fabric8.kubernetes.client.server.mock.EnableKubernetesMockClient;
 import io.fabric8.kubernetes.client.server.mock.KubernetesMockServer;
 import org.junit.jupiter.api.AfterEach;
@@ -50,6 +60,9 @@
 @TestMethodOrder(MethodOrderer.OrderAnnotation.class)
 class KubernetesMetricsInterceptorTest {
 
+  private static final String CONFIG_MAP_PATH =
+      "/api/v1/namespaces/spark-system/configmaps/spark-job-operator-configuration";
+
   private KubernetesMockServer mockServer;
   private KubernetesClient kubernetesClient;
 
@@ -125,6 +138,137 @@ void testWhenKubernetesServerNotWorking() {
     }
   }
 
+  @Test
+  @Order(3)
+  void testResponseLatency() {
+    KubernetesMetricsInterceptor metricsInterceptor = new KubernetesMetricsInterceptor();
+    List<Interceptor> interceptors = List.of(metricsInterceptor);
+    try (KubernetesClient client =
+        KubernetesClientFactory.buildKubernetesClient(
+            interceptors, kubernetesClient.getConfiguration())) {
+      ConfigMap configMap = createConfigMap();
+      mockServer
+          .expect()
+          .get()
+          .delay(50)
+          .withPath(CONFIG_MAP_PATH)
+          .andReturn(200, configMap)
+          .once();
+      client.resource(configMap).get();
+
+      Histogram latency = latency(metricsInterceptor);
+      Assertions.assertEquals(1, latency.getCount());
+      Snapshot snapshot = latency.getSnapshot();
+      Assertions.assertTrue(snapshot.getMax() >= 50_000_000L, () -> "max: " + snapshot.getMax());
+      Assertions.assertTrue(
+          snapshot.getMax() < 30_000_000_000L, () -> "max: " + snapshot.getMax());
+    }
+  }
+
+  @Test
+  @Order(4)
+  void testResponseLatencyOfResentRequest() {
+    KubernetesMetricsInterceptor metricsInterceptor = new KubernetesMetricsInterceptor();
+    List<Interceptor> interceptors = List.of(metricsInterceptor);
+    Config config =
+        new ConfigBuilder(kubernetesClient.getConfiguration())
+            .withOauthTokenProvider(() -> "token")
+            .build();
+    try (KubernetesClient client =
+        KubernetesClientFactory.buildKubernetesClient(interceptors, config)) {
+      ConfigMap configMap = createConfigMap();
+      mockServer
+          .expect()
+          .get()
+          .delay(200)
+          .withPath(CONFIG_MAP_PATH)
+          .andReturn(HTTP_UNAUTHORIZED, new StatusBuilder().withCode(HTTP_UNAUTHORIZED).build())
+          .once();
+      mockServer.expect().get().withPath(CONFIG_MAP_PATH).andReturn(200, configMap).once();
+      client.resource(configMap).get();
+
+      Assertions.assertEquals(2, meterCount(metricsInterceptor, "http.response"));
+      Assertions.assertEquals(1, meterCount(metricsInterceptor, "http.response.401"));
+      Assertions.assertEquals(0, meterCount(metricsInterceptor, "failed"));
+      Histogram latency = latency(metricsInterceptor);
+      Assertions.assertEquals(2, latency.getCount());
+      Snapshot snapshot = latency.getSnapshot();
+      Assertions.assertTrue(snapshot.getMax() >= 200_000_000L, () -> "max: " + snapshot.getMax());
+      Assertions.assertTrue(snapshot.getMin() < 200_000_000L, () -> "min: " + snapshot.getMin());
+    }
+  }
+
+  @Test
+  @Order(5)
+  void testResponseLatencyOfWebSocketUpgrade() {
+    KubernetesMetricsInterceptor metricsInterceptor = new KubernetesMetricsInterceptor();
+    List<Interceptor> interceptors = List.of(metricsInterceptor);
+    try (KubernetesClient client =
+            KubernetesClientFactory.buildKubernetesClient(
+                interceptors, kubernetesClient.getConfiguration());
+        SharedIndexInformer<ConfigMap> informer =
+            client.configMaps().inNamespace("spark-system").inform()) {
+      await()
+          .untilAsserted(
+              () -> {
+                Assertions.assertEquals(2, meterCount(metricsInterceptor, "http.response"));
+                Assertions.assertEquals(1, meterCount(metricsInterceptor, "http.response.101"));
+                Assertions.assertEquals(1, latency(metricsInterceptor).getCount());
+              });
+      Assertions.assertTrue(informer.isWatching());
+    }
+  }
+
+  @Test
+  @Order(6)
+  void testResponseLatencyOfTwoInterceptors() {
+    KubernetesMetricsInterceptor first = new KubernetesMetricsInterceptor();
+    KubernetesMetricsInterceptor second = new KubernetesMetricsInterceptor() {};
+    List<Interceptor> interceptors = List.of(first, second);
+    try (KubernetesClient client =
+        KubernetesClientFactory.buildKubernetesClient(
+            interceptors, kubernetesClient.getConfiguration())) {
+      ConfigMap configMap = createConfigMap();
+      mockServer
+          .expect()
+          .get()
+          .delay(50)
+          .withPath(CONFIG_MAP_PATH)
+          .andReturn(200, configMap)
+          .once();
+      client.resource(configMap).get();
+
+      for (KubernetesMetricsInterceptor interceptor : List.of(first, second)) {
+        Histogram latency = latency(interceptor);
+        Assertions.assertEquals(1, latency.getCount());
+        Snapshot snapshot = latency.getSnapshot();
+        Assertions.assertTrue(snapshot.getMin() >= 50_000_000L, () -> "min: " + snapshot.getMin());
+      }
+    }
+  }
+
+  @Test
+  @Order(7)
+  void testConsumerUnwrapsDelegate() {
+    AsyncBody.Consumer<List<ByteBuffer>> delegate = (value, asyncBody) -> {};
+    AsyncBody.Consumer<List<ByteBuffer>> consumer =
+        new KubernetesMetricsInterceptor().consumer(delegate, null);
+    Assertions.assertSame(delegate, consumer.unwrap(delegate.getClass()));
+  }
+
+  private static long meterCount(KubernetesMetricsInterceptor interceptor, String name) {
+    Meter meter = interceptor.metricRegistry().getMeters().get(name);
+    Assertions.assertNotNull(meter, name);
+    return meter.getCount();
+  }
+
+  private static Histogram latency(KubernetesMetricsInterceptor interceptor) {
+    Histogram histogram =
+        interceptor.metricRegistry().getHistograms().get("http.response.latency.nanos");
+    Assertions.assertNotNull(histogram, "http.response.latency.nanos");
+    return histogram;
+  }
+
   private static SparkApplication createSparkApplication() {
     ObjectMeta meta = new ObjectMeta();
     meta.setName("sample-spark-application");