Handle edges when cached StoragePolicy cluster state has no distributors

* Avoid division by zero in `getRandomTargetSpec()`
* Ensure policy context is set when `NoDistributorsAvailableException` is thrown

2 files changed, 38 insertions, 4 deletions

diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/StoragePolicy.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/StoragePolicy.java
index adf9d5ee912..46d0cd3fc37 100644
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/StoragePolicy.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/StoragePolicy.java
@@ -104,7 +104,9 @@ public class StoragePolicy extends SlobrokPolicy {
         String getRandomTargetSpec(RoutingContext context) {
             Targets targets = validTargets.get();
             // Try to use list of random targets, if at least X % of the nodes are up
-            while (100 * targets.list.size() / targets.total >= requiredUpPercentageToSendToKnownGoodNodes) {
+            while ((targets.total != 0) &&
+                   (100 * targets.list.size() / targets.total >= requiredUpPercentageToSendToKnownGoodNodes))
+            {
                 int randIndex = randomizer.nextInt(targets.list.size());
                 String targetSpec = getTargetSpec(targets.list.get(randIndex), context);
                 if (targetSpec != null) {
@@ -408,9 +410,10 @@ public class StoragePolicy extends SlobrokPolicy {
                     log.log(LogLevel.DEBUG, "No distributors available; clearing cluster state");
                     safeCachedClusterState.set(null);
                     sendRandomReason = "No distributors available. Sending to random distributor.";
+                    context.setContext(createRandomDistributorTargetContext());
                 }
             } else {
-                context.setContext(new MessageContext(null));
+                context.setContext(createRandomDistributorTargetContext());
                 sendRandomReason = "No cluster state cached. Sending to random distributor.";
             }
             if (context.shouldTrace(1)) {
@@ -419,6 +422,10 @@ public class StoragePolicy extends SlobrokPolicy {
             return hostFetcher.getRandomTargetSpec(context);
         }
 
+        private static MessageContext createRandomDistributorTargetContext() {
+            return new MessageContext(null);
+        }
+
         private static Optional<ClusterState> clusterStateFromReply(final WrongDistributionReply reply) {
             try {
                 return Optional.of(new ClusterState(reply.getSystemState()));
diff --git a/documentapi/src/test/java/com/yahoo/documentapi/messagebus/protocol/test/storagepolicy/BasicTests.java b/documentapi/src/test/java/com/yahoo/documentapi/messagebus/protocol/test/storagepolicy/BasicTests.java
index 14d8818bfec..8a6df061430 100644
--- a/documentapi/src/test/java/com/yahoo/documentapi/messagebus/protocol/test/storagepolicy/BasicTests.java
+++ b/documentapi/src/test/java/com/yahoo/documentapi/messagebus/protocol/test/storagepolicy/BasicTests.java
@@ -3,6 +3,8 @@ package com.yahoo.documentapi.messagebus.protocol.test.storagepolicy;
 
 import com.yahoo.collections.Pair;
 import com.yahoo.documentapi.messagebus.protocol.WrongDistributionReply;
+import com.yahoo.messagebus.Error;
+import com.yahoo.messagebus.ErrorCode;
 import com.yahoo.messagebus.Reply;
 import com.yahoo.messagebus.routing.RoutingNode;
 import org.junit.Test;
@@ -19,11 +21,11 @@ public class BasicTests extends StoragePolicyTestEnvironment {
     @Test
     public void testNormalUsage() {
         setClusterNodes(new int[]{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 });
-            // First we want a wrong distribution reply, so make sure we don't try correct node on random
+        // First we want a wrong distribution reply, so make sure we don't try correct node on random
         policyFactory.avoidPickingAtRandom(bucketOneNodePreference[0]);
         RoutingNode target = select();
         replyWrongDistribution(target, "foo", 5, "version:1 bits:16 distributor:10 storage:10");
-            // Then send to correct node and verify that
+        // Then send to correct node and verify that
         sendToCorrectNode("foo", bucketOneNodePreference[0]);
     }
 
@@ -76,4 +78,29 @@ public class BasicTests extends StoragePolicyTestEnvironment {
         assertEquals(Arrays.toString(bucketOneNodePreference), Arrays.toString(result));
     }
 
+    private void letPolicyCacheStateWithNoAvailableDistributors() {
+        setClusterNodes(new int[]{ 0 });
+        String clusterState = " bits:16 storage:1"; // No distributors online
+        var target = select();
+        // Update to cluster state with no distributors online
+        replyWrongDistribution(target, "foo", null, clusterState);
+    }
+
+    @Test
+    public void random_node_target_used_if_no_distributors_available_in_cached_state() {
+        letPolicyCacheStateWithNoAvailableDistributors();
+        select(); // Will trigger failure if no target is set, i.e. if an exception is thrown
+    }
+
+    @Test
+    public void policy_context_set_on_no_distributors_exception_during_select() {
+        letPolicyCacheStateWithNoAvailableDistributors();
+        // Re-select with no distributors. Should still get a target selection context.
+        var target = select();
+        // Trigger a reply merge with error, which will try to inspect the context.
+        // If an exception is thrown internally from a missing context, the reply will not be
+        // set correctly and this will fail.
+        replyError(target, new Error(ErrorCode.FATAL_ERROR, "oh no!"));
+    }
+
 }