From 885f7f818def4e80c490f91d8658b934ae884539 Mon Sep 17 00:00:00 2001
From: jrhee17 <guins_j@guins.org>
Date: Mon, 5 Aug 2024 10:30:10 +0900
Subject: [PATCH] address nullaway failure

---
 .../xds/client/endpoint/DefaultLoadBalancer.java    |  3 ++-
 .../endpoint/LocalityRoutingStateFactory.java       | 13 ++++++-------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/xds/src/main/java/com/linecorp/armeria/xds/client/endpoint/DefaultLoadBalancer.java b/xds/src/main/java/com/linecorp/armeria/xds/client/endpoint/DefaultLoadBalancer.java
index e2bd1b89857..14708c44fe1 100644
--- a/xds/src/main/java/com/linecorp/armeria/xds/client/endpoint/DefaultLoadBalancer.java
+++ b/xds/src/main/java/com/linecorp/armeria/xds/client/endpoint/DefaultLoadBalancer.java
@@ -134,7 +134,8 @@ HostsSource hostSourceToUse(DefaultLbState lbState, XdsRandom random,
         if (priority != 0 ||
             localityRoutingState == null ||
             localityRoutingState.state() == State.NO_LOCALITY_ROUTING ||
-            prioritySet.routingEnabled() <= random.nextInt(100, ROUTING_ENABLED)) {
+            prioritySet.routingEnabled() <= random.nextInt(100, ROUTING_ENABLED) ||
+            localityRoutingState.localHostSet() == null) {
             return new HostsSource(priority, sourceType(hostAvailability));
         }
 
diff --git a/xds/src/main/java/com/linecorp/armeria/xds/client/endpoint/LocalityRoutingStateFactory.java b/xds/src/main/java/com/linecorp/armeria/xds/client/endpoint/LocalityRoutingStateFactory.java
index 9f3ce14b1ca..50d6ddd3573 100644
--- a/xds/src/main/java/com/linecorp/armeria/xds/client/endpoint/LocalityRoutingStateFactory.java
+++ b/xds/src/main/java/com/linecorp/armeria/xds/client/endpoint/LocalityRoutingStateFactory.java
@@ -47,7 +47,8 @@ LocalityRoutingState create(PrioritySet upstreamPrioritySet, PrioritySet localPr
         // Only priority 0 is supported
         final HostSet upstreamHostSet = upstreamPrioritySet.hostSets().get(0);
         final HostSet localHostSet = localPrioritySet.hostSets().get(0);
-        if (earlyExitNonLocalityRouting(upstreamHostSet, localHostSet, localLocality,
+        if (upstreamHostSet == null || localHostSet == null ||
+            earlyExitNonLocalityRouting(upstreamHostSet, localHostSet, localLocality,
                                         upstreamPrioritySet.minClusterSize())) {
             return new LocalityRoutingState(State.NO_LOCALITY_ROUTING, localHostSet, localLocality);
         }
@@ -121,10 +122,11 @@ static final class LocalityRoutingState {
         private final State state;
         private final long localPercentageToRoute;
         private final List<ResidualCapacity> residualCapacities;
+        @Nullable
         private final HostSet localHostSet;
         private final Locality localLocality;
 
-        private LocalityRoutingState(State state, HostSet localHostSet, Locality localLocality) {
+        private LocalityRoutingState(State state, @Nullable HostSet localHostSet, Locality localLocality) {
             assert state == State.NO_LOCALITY_ROUTING || state == State.LOCALITY_DIRECT;
             this.state = state;
             localPercentageToRoute = 0;
@@ -147,6 +149,7 @@ State state() {
             return state;
         }
 
+        @Nullable
         HostSet localHostSet() {
             return localHostSet;
         }
@@ -238,12 +241,8 @@ private static Map<Locality, LocalityPercentages> calculateLocalityPercentages(
         return localityPercentagesBuilder.build();
     }
 
-    private static boolean earlyExitNonLocalityRouting(@Nullable HostSet upstreamHostSet,
-                                                       @Nullable HostSet localHostSet,
+    private static boolean earlyExitNonLocalityRouting(HostSet upstreamHostSet, HostSet localHostSet,
                                                        Locality localLocality, long minClusterSize) {
-        if (upstreamHostSet == null || localHostSet == null) {
-            return true;
-        }
         final Map<Locality, EndpointGroup> upstreamHealthyHostsPerLocality =
                 upstreamHostSet.healthyEndpointGroupPerLocality();
         final Map<Locality, EndpointGroup> localHealthyHostsPerLocality =