From dd6f26e9db9d494162af76339acb6a85fd635f95 Mon Sep 17 00:00:00 2001
From: Bryce Anderson <bl_anderson@apple.com>
Date: Wed, 19 Feb 2025 17:32:18 -0700
Subject: [PATCH] loadbalancer: fix flaky test (#3195)

Motivation:

We have a flaky test. This is because we're able to get a connection
but only transiently. This happens because the DefaultHost is willing
to add a connection to an expired host without connections, which is a
state that is about to be closed.

Modifications:

Don't add the connection to the host if the host is is the transient
expired state with zero connections.

Fixes #2677
---
 .../java/io/servicetalk/loadbalancer/DefaultHost.java    | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/servicetalk-loadbalancer/src/main/java/io/servicetalk/loadbalancer/DefaultHost.java b/servicetalk-loadbalancer/src/main/java/io/servicetalk/loadbalancer/DefaultHost.java
index 39cdfda5ec..5f44d6b7b7 100644
--- a/servicetalk-loadbalancer/src/main/java/io/servicetalk/loadbalancer/DefaultHost.java
+++ b/servicetalk-loadbalancer/src/main/java/io/servicetalk/loadbalancer/DefaultHost.java
@@ -167,11 +167,11 @@ public boolean markExpired() {
             } else if (oldState.state == State.CLOSED) {
                 return true;
             }
-            Object nextState = oldState.connections.isEmpty() ? State.CLOSED : State.EXPIRED;
+            assert oldState.state == State.ACTIVE || oldState.state == State.UNHEALTHY;
             if (connStateUpdater.compareAndSet(this, oldState, oldState.toExpired())) {
                 cancelIfHealthCheck(oldState);
                 hostObserver.onHostMarkedExpired(oldState.connections.size());
-                if (nextState == State.CLOSED) {
+                if (oldState.connections.isEmpty()) {
                     // Trigger the callback to remove the host from usedHosts array.
                     this.closeAsync().subscribe();
                     return true;
@@ -326,7 +326,10 @@ private boolean addConnection(final C connection) {
         int addAttempt = 0;
         for (;;) {
             final ConnState previous = connStateUpdater.get(this);
-            if (previous.state == State.CLOSED) {
+            if (previous.state == State.CLOSED ||
+                    // an expired state with no connections is only a transient state and the host is going
+                    // to be closing, so reject the connection.
+                    previous.state == State.EXPIRED && previous.connections.isEmpty()) {
                 return false;
             }
             ++addAttempt;