loadbalancer: Add ConnectTracker and make HealthIndicator extend it (#2818)

Motivation:

We can't rely on the request path to pass in connection
failures because if we can't get a connection we can never
add the observers. That means we need the `Host` implementation
to pass it in directly.

Modifications:

Add two new methods to `HealthIndicator` that provide the current
time and and a callback for submitting connection failures.

Result:

We should be able to track connection failures now and factor
those into the `HealthTracker` implementations.

Author: bryce-anderson

servicetalk-http-netty/src/main/java/io/servicetalk/http/netty/DefaultHttpLoadBalancerFactory.java

@@ -368,23 +368,23 @@ public Single<StreamingHttpResponse> request(final StreamingHttpRequest request)
 
             return Single.defer(() -> {
                 final RequestTracker theTracker = new AtMostOnceDeliveryRequestTracker(tracker);
-                final long startTime = theTracker.beforeStart();
+                final long startTime = theTracker.beforeRequestStart();
 
                 return delegate.request(request)
                         .liftSync(new BeforeFinallyHttpOperator(new TerminalSignalConsumer() {
                             @Override
                             public void onComplete() {
-                                theTracker.onSuccess(startTime);
+                                theTracker.onRequestSuccess(startTime);
                             }
 
                             @Override
                             public void onError(final Throwable throwable) {
-                                theTracker.onError(startTime, errorClassFunction.apply(throwable));
+                                theTracker.onRequestError(startTime, errorClassFunction.apply(throwable));
                             }
 
                             @Override
                             public void cancel() {
-                                theTracker.onError(startTime, ErrorClass.CANCELLED);
+                                theTracker.onRequestError(startTime, ErrorClass.CANCELLED);
                             }
                         }, /*discardEventsAfterCancel*/ true))
 
@@ -397,7 +397,7 @@ public void cancel() {
                             final ErrorClass eClass = peerResponseErrorClassifier.apply(response);
                             if (eClass != null) {
                                 // The onError is triggered before the body is actually consumed.
-                                theTracker.onError(startTime, eClass);
+                                theTracker.onRequestError(startTime, eClass);
                             }
                             return response;
                         })
@@ -453,21 +453,21 @@ private AtMostOnceDeliveryRequestTracker(final RequestTracker original) {
             }
 
             @Override
-            public long beforeStart() {
-                return original.beforeStart();
+            public long beforeRequestStart() {
+                return original.beforeRequestStart();
             }
 
             @Override
-            public void onSuccess(final long beforeStartTimeNs) {
+            public void onRequestSuccess(final long beforeStartTimeNs) {
                 if (doneUpdater.compareAndSet(this, 0, 1)) {
-                    original.onSuccess(beforeStartTimeNs);
+                    original.onRequestSuccess(beforeStartTimeNs);
                 }
             }
 
             @Override
-            public void onError(final long beforeStartTimeNs, final ErrorClass errorClass) {
+            public void onRequestError(final long beforeStartTimeNs, final ErrorClass errorClass) {
                 if (doneUpdater.compareAndSet(this, 0, 1)) {
-                    original.onError(beforeStartTimeNs, errorClass);
+                    original.onRequestError(beforeStartTimeNs, errorClass);
                 }
             }
         }

servicetalk-loadbalancer-experimental/src/main/java/io/servicetalk/loadbalancer/ConnectTracker.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright © 2024 Apple Inc. and the ServiceTalk project authors
+ *
+ * Licensed 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 io.servicetalk.loadbalancer;
+
+/**
+ * An interface for tracking connection establishment measurements.
+ * This has an intended usage similar to the {@link RequestTracker} but with a focus on connection establishment
+ * metrics.
+ */
+interface ConnectTracker {
+
+    /**
+     * Get the current time in nanoseconds.
+     * Note: this must not be a stateful API. Eg, it does not necessarily have a correlation with any other method call
+     * and such shouldn't be used as a method of counting in the same way that {@link RequestTracker} is used.
+     * @return the current time in nanoseconds.
+     */
+    long beforeConnectStart();
+
+    /**
+     * Callback to notify the parent {@link HealthChecker} that an attempt to connect to this host has succeeded.
+     * @param beforeConnectStart the time that the connection attempt was initiated.
+     */
+    void onConnectSuccess(long beforeConnectStart);
+
+    /**
+     * Callback to notify the parent {@link HealthChecker} that an attempt to connect to this host has failed.
+     * @param beforeConnectStart the time that the connection attempt was initiated.
+     */
+    void onConnectError(long beforeConnectStart);
+}

servicetalk-loadbalancer-experimental/src/main/java/io/servicetalk/loadbalancer/DefaultHost.java

@@ -17,15 +17,18 @@
 
 import io.servicetalk.client.api.ConnectionFactory;
 import io.servicetalk.client.api.ConnectionLimitReachedException;
+import io.servicetalk.client.api.DelegatingConnectionFactory;
 import io.servicetalk.client.api.LoadBalancedConnection;
 import io.servicetalk.concurrent.api.AsyncContext;
 import io.servicetalk.concurrent.api.Completable;
 import io.servicetalk.concurrent.api.ListenableAsyncCloseable;
 import io.servicetalk.concurrent.api.Single;
+import io.servicetalk.concurrent.api.TerminalSignalConsumer;
 import io.servicetalk.concurrent.internal.DefaultContextMap;
 import io.servicetalk.concurrent.internal.DelayedCancellable;
 import io.servicetalk.context.api.ContextMap;
 import io.servicetalk.loadbalancer.LoadBalancerObserver.HostObserver;
+import io.servicetalk.transport.api.TransportObserver;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -112,10 +115,12 @@ private enum State {
         this.lbDescription = requireNonNull(lbDescription, "lbDescription");
         this.address = requireNonNull(address, "address");
         this.linearSearchSpace = linearSearchSpace;
-        this.connectionFactory = requireNonNull(connectionFactory, "connectionFactory");
+        this.healthIndicator = healthIndicator;
+        requireNonNull(connectionFactory, "connectionFactory");
+        this.connectionFactory = healthIndicator == null ? connectionFactory :
+                new InstrumentedConnectionFactory<>(connectionFactory, healthIndicator);
         this.healthCheckConfig = healthCheckConfig;
         this.hostObserver = requireNonNull(hostObserver, "hostObserver");
-        this.healthIndicator = healthIndicator;
         this.closeable = toAsyncCloseable(this::doClose);
         hostObserver.onHostCreated(address);
     }
@@ -235,7 +240,7 @@ public Single<C> newConnection(
             Single<? extends C> establishConnection = connectionFactory.newConnection(address, actualContext, null);
             if (healthCheckConfig != null) {
                 // Schedule health check before returning
-                establishConnection = establishConnection.beforeOnError(this::markUnhealthy);
+                establishConnection = establishConnection.beforeOnError(this::onConnectionError);
             }
             return establishConnection
                     .flatMap(newCnx -> {
@@ -302,7 +307,7 @@ private void markHealthy(final HealthCheck originalHealthCheckState) {
         }
     }
 
-    private void markUnhealthy(final Throwable cause) {
+    private void onConnectionError(Throwable cause) {
         assert healthCheckConfig != null;
         for (;;) {
             ConnState previous = connStateUpdater.get(this);
@@ -646,4 +651,56 @@ public String toString() {
                     '}';
         }
     }
+
+    private static final class InstrumentedConnectionFactory<Addr, C extends LoadBalancedConnection>
+            extends DelegatingConnectionFactory<Addr, C> {
+
+        private final ConnectTracker connectTracker;
+
+        InstrumentedConnectionFactory(final ConnectionFactory<Addr, C> delegate, ConnectTracker connectTracker) {
+            super(delegate);
+            this.connectTracker = connectTracker;
+        }
+
+        @Override
+        public Single<C> newConnection(Addr addr, @Nullable ContextMap context, @Nullable TransportObserver observer) {
+            return Single.defer(() -> {
+                final long connectStartTime = connectTracker.beforeConnectStart();
+                return delegate().newConnection(addr, context, observer)
+                        .beforeFinally(new ConnectSignalConsumer<>(connectStartTime, connectTracker))
+                        .shareContextOnSubscribe();
+            });
+        }
+    }
+
+    private static class ConnectSignalConsumer<C extends LoadBalancedConnection> implements TerminalSignalConsumer {
+
+        private final ConnectTracker connectTracker;
+        private final long connectStartTime;
+
+        ConnectSignalConsumer(final long connectStartTime, final ConnectTracker connectTracker) {
+            this.connectStartTime = connectStartTime;
+            this.connectTracker = connectTracker;
+        }
+
+        @Override
+        public void onComplete() {
+            connectTracker.onConnectSuccess(connectStartTime);
+        }
+
+        @Override
+        public void cancel() {
+            // We assume cancellation is the result of some sort of timeout.
+            doOnError();
+        }
+
+        @Override
+        public void onError(Throwable t) {
+            doOnError();
+        }
+
+        private void doOnError() {
+            connectTracker.onConnectError(connectStartTime);
+        }
+    }
 }

servicetalk-loadbalancer-experimental/src/main/java/io/servicetalk/loadbalancer/DefaultRequestTracker.java

@@ -81,19 +81,19 @@ abstract class DefaultRequestTracker implements RequestTracker, ScoreSupplier {
     protected abstract long currentTimeNanos();
 
     @Override
-    public final long beforeStart() {
+    public final long beforeRequestStart() {
         pendingUpdater.incrementAndGet(this);
         return currentTimeNanos();
     }
 
     @Override
-    public void onSuccess(final long startTimeNanos) {
+    public void onRequestSuccess(final long startTimeNanos) {
         pendingUpdater.decrementAndGet(this);
         calculateAndStore((ewma, currentLatency) -> currentLatency, startTimeNanos);
     }
 
     @Override
-    public void onError(final long startTimeNanos, ErrorClass errorClass) {
+    public void onRequestError(final long startTimeNanos, ErrorClass errorClass) {
         pendingUpdater.decrementAndGet(this);
         calculateAndStore(errorClass == ErrorClass.CANCELLED ? this:: cancelPenalty : this::errorPenalty,
                 startTimeNanos);

servicetalk-loadbalancer-experimental/src/main/java/io/servicetalk/loadbalancer/HealthIndicator.java

@@ -26,7 +26,7 @@
  * health check system can give the host information about it's perceived health and the host can give the
  * health check system information about request results.
  */
-interface HealthIndicator extends RequestTracker, ScoreSupplier, Cancellable {
+interface HealthIndicator extends RequestTracker, ConnectTracker, ScoreSupplier, Cancellable {
 
     /**
      * Whether the host is considered healthy by the HealthIndicator.

servicetalk-loadbalancer-experimental/src/main/java/io/servicetalk/loadbalancer/RequestTracker.java

@@ -21,13 +21,13 @@
  * A tracker of latency of an action over time.
  * <p>
  * The usage of the RequestTracker is intended to follow the simple workflow:
- * - At initiation of an action for which a request is must call {@link RequestTracker#beforeStart()} and save the
- *   timestamp much like would be done when using a stamped lock.
- * - Once the request event is complete only one of the {@link RequestTracker#onSuccess(long)} or
- *   {@link RequestTracker#onError(long, ErrorClass)} methods must be called and called exactly once.
- * In other words, every call to {@link RequestTracker#beforeStart()} must be followed by exactly one call to either of
- * the completion methods {@link RequestTracker#onSuccess(long)} or
- * {@link RequestTracker#onError(long, ErrorClass)}. Failure to do so can cause state corruption in the
+ * - At initiation of an action for which a request is must call {@link RequestTracker#beforeRequestStart()} and save
+ *   the timestamp much like would be done when using a stamped lock.
+ * - Once the request event is complete only one of the {@link RequestTracker#onRequestSuccess(long)} or
+ *   {@link RequestTracker#onRequestError(long, ErrorClass)} methods must be called and called exactly once.
+ * In other words, every call to {@link RequestTracker#beforeRequestStart()} must be followed by exactly one call to
+ * either of the completion methods {@link RequestTracker#onRequestSuccess(long)} or
+ * {@link RequestTracker#onRequestError(long, ErrorClass)}. Failure to do so can cause state corruption in the
  * {@link RequestTracker} implementations which may track not just latency but also the outstanding requests.
  */
 public interface RequestTracker {
@@ -40,20 +40,20 @@ public interface RequestTracker {
      *
      * @return Current time in nanoseconds.
      */
-    long beforeStart();
+    long beforeRequestStart();
 
     /**
      * Records a successful completion of the action for which latency is to be tracked.
      *
-     * @param beforeStartTimeNs return value from {@link #beforeStart()}.
+     * @param beforeStartTimeNs return value from {@link #beforeRequestStart()}.
      */
-    void onSuccess(long beforeStartTimeNs);
+    void onRequestSuccess(long beforeStartTimeNs);
 
     /**
      * Records a failed completion of the action for which latency is to be tracked.
      *
-     * @param beforeStartTimeNs return value from {@link #beforeStart()}.
+     * @param beforeStartTimeNs return value from {@link #beforeRequestStart()}.
      * @param errorClass the class of error that triggered this method.
      */
-    void onError(long beforeStartTimeNs, ErrorClass errorClass);
+    void onRequestError(long beforeStartTimeNs, ErrorClass errorClass);
 }

servicetalk-loadbalancer-experimental/src/main/java/io/servicetalk/loadbalancer/XdsHealthIndicator.java

@@ -115,20 +115,41 @@ public final boolean isHealthy() {
     }
 
     @Override
-    public final void onSuccess(final long beforeStartTimeNs) {
-        super.onSuccess(beforeStartTimeNs);
+    public final void onRequestSuccess(final long beforeStartTimeNs) {
+        super.onRequestSuccess(beforeStartTimeNs);
         successes.incrementAndGet();
         consecutive5xx.set(0);
         LOGGER.trace("Observed success for address {}", address);
     }
 
     @Override
-    public final void onError(final long beforeStartTimeNs, ErrorClass errorClass) {
-        super.onError(beforeStartTimeNs, errorClass);
+    public final void onRequestError(final long beforeStartTimeNs, ErrorClass errorClass) {
+        super.onRequestError(beforeStartTimeNs, errorClass);
         // For now, don't consider cancellation to be an error or a success.
-        if (errorClass == ErrorClass.CANCELLED) {
-            return;
+        if (errorClass != ErrorClass.CANCELLED) {
+            doOnError();
         }
+    }
+
+    @Override
+    public long beforeConnectStart() {
+        return currentTimeNanos();
+    }
+
+    @Override
+    public void onConnectError(long beforeConnectStart) {
+        // This assumes that the connect request was intended to be used for a request dispatch which
+        // will have now failed. This is not strictly true: a connection can be acquired and simply not
+        // used, but in practice it's a very good assumption.
+        doOnError();
+    }
+
+    @Override
+    public void onConnectSuccess(long beforeConnectStart) {
+        // noop: the request path will now determine if the request was a success or failure.
+    }
+
+    private void doOnError() {
         failures.incrementAndGet();
         final int consecutiveFailures = consecutive5xx.incrementAndGet();
         final OutlierDetectorConfig localConfig = currentConfig();

servicetalk-loadbalancer-experimental/src/test/java/io/servicetalk/loadbalancer/DefaultHostTest.java

@@ -29,11 +29,14 @@
 import java.util.function.Predicate;
 import javax.annotation.Nullable;
 
+import static io.servicetalk.concurrent.api.Single.failed;
 import static io.servicetalk.concurrent.api.Single.succeeded;
+import static io.servicetalk.concurrent.internal.DeliberateException.DELIBERATE_EXCEPTION;
 import static io.servicetalk.loadbalancer.HealthCheckConfig.DEFAULT_HEALTH_CHECK_FAILED_CONNECTIONS_THRESHOLD;
 import static io.servicetalk.loadbalancer.UnhealthyHostConnectionFactory.UNHEALTHY_HOST_EXCEPTION;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.is;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -233,4 +236,17 @@ void forwardsHealthIndicatorScore() {
         assertThat(host.score(), is(10));
         verify(healthIndicator, times(1)).score();
     }
+
+    @Test
+    void connectFailuresAreForwardedToHealthIndicator() {
+        connectionFactory = new TestConnectionFactory(address -> failed(DELIBERATE_EXCEPTION));
+        HealthIndicator healthIndicator = mock(HealthIndicator.class);
+        buildHost(healthIndicator);
+        verify(mockHostObserver, times(1)).onHostCreated("address");
+        Throwable underlying = assertThrows(ExecutionException.class, () ->
+                host.newConnection(cxn -> true, false, null).toFuture().get()).getCause();
+        assertEquals(DELIBERATE_EXCEPTION, underlying);
+        verify(healthIndicator, times(1)).beforeConnectStart();
+        verify(healthIndicator, times(1)).onConnectError(0L);
+    }
 }