Merge pull request #2733 from yahoo/hmusum/16-distribution-bits-only-for-prod

Enable 16 distribution bits only for prod regions

3 files changed, 84 insertions, 69 deletions

diff --git a/config-model/src/main/java/com/yahoo/vespa/model/content/cluster/ContentCluster.java b/config-model/src/main/java/com/yahoo/vespa/model/content/cluster/ContentCluster.java
index 76059f53149..1ddfbbd76e1 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/content/cluster/ContentCluster.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/content/cluster/ContentCluster.java
@@ -6,6 +6,7 @@ import com.google.common.collect.Sets;
 import com.yahoo.config.model.ConfigModelContext;
 import com.yahoo.config.model.producer.AbstractConfigProducerRoot;
 import com.yahoo.config.provision.ClusterSpec;
+import com.yahoo.config.provision.Environment;
 import com.yahoo.config.provision.RegionName;
 import com.yahoo.config.provision.Zone;
 import com.yahoo.vespa.config.content.MessagetyperouteselectorpolicyConfig;
@@ -46,6 +47,7 @@ import org.w3c.dom.Element;
 import java.util.*;
 import java.util.logging.Level;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * A content cluster.
@@ -585,19 +587,27 @@ public class ContentCluster extends AbstractConfigProducer implements StorDistri
 
     /**
      * Returns the distribution bits this cluster should use.
-     * OnHosted Vespa this is hardcoded not computed from the nodes because reducing the number of nodes is a common
-     * operation while reducing the number of distribution bits can lead to consistency problems.
+     * On Hosted Vespa this is hardcoded and not computed from the nodes because reducing the number of nodes is a common
+     * operation, while reducing the number of distribution bits can lead to consistency problems.
      * This hardcoded value should work fine from 1-200 nodes. Those who have more will need to set this value
      * in config and not remove it again if they reduce the node count.
      */
     public int distributionBits() {
-        if (zone.region().equals(RegionName.from("us-west-1"))) return 16; // TODO: Enable for all hosted zones (i.e when zone isn't default)
-        if (zone.region().equals(RegionName.from("us-central-1"))) return 16; // TODO: Enable for all hosted zones (i.e when zone isn't default)
-        if (zone.region().equals(RegionName.from("eu-west-1"))) return 16; // TODO: Enable for all hosted zones (i.e when zone isn't default)
-        if (zone.region().equals(RegionName.from("ap-northeast-1"))) return 16; // TODO: Enable for all hosted zones (i.e when zone isn't default)
-        if (zone.region().equals(RegionName.from("ap-northeast-2"))) return 16; // TODO: Enable for all hosted zones (i.e when zone isn't default)
-        if (zone.region().equals(RegionName.from("us-east-3"))) return 16; // TODO: Enable for all hosted zones (i.e when zone isn't default)
-        return DistributionBitCalculator.getDistributionBits(getNodeCountPerGroup(), getDistributionMode());
+        // TODO: Enable for all hosted zones (i.e when zone isn't default)
+        List<Zone> zonesWith16DistributionBits = Arrays.asList(createZone(Environment.prod, "us-west-1"),
+                                                               createZone(Environment.prod, "us-central-1"),
+                                                               createZone(Environment.prod, "eu-west-1"),
+                                                               createZone(Environment.prod, "ap-northeast-1"),
+                                                               createZone(Environment.prod, "ap-northeast-2"),
+                                                               createZone(Environment.prod, "us-east-3"));
+        if (zonesWith16DistributionBits.contains(zone))
+            return 16;
+        else
+            return DistributionBitCalculator.getDistributionBits(getNodeCountPerGroup(), getDistributionMode());
+    }
+
+    private Zone createZone(Environment environment, String region) {
+        return new Zone(environment, RegionName.from(region));
     }
 
     @Override
diff --git a/config-model/src/test/java/com/yahoo/vespa/model/content/ClusterTest.java b/config-model/src/test/java/com/yahoo/vespa/model/content/ClusterTest.java
index 126fcf7a583..82cbc338a8e 100644
--- a/config-model/src/test/java/com/yahoo/vespa/model/content/ClusterTest.java
+++ b/config-model/src/test/java/com/yahoo/vespa/model/content/ClusterTest.java
@@ -1,8 +1,14 @@
 // Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.model.content;
 
+import com.yahoo.config.model.deploy.DeployProperties;
+import com.yahoo.config.model.deploy.DeployState;
+import com.yahoo.config.model.test.MockRoot;
 import com.yahoo.config.model.test.TestDriver;
 import com.yahoo.config.model.test.TestRoot;
+import com.yahoo.config.provision.Environment;
+import com.yahoo.config.provision.RegionName;
+import com.yahoo.config.provision.Zone;
 import com.yahoo.vespa.config.content.core.StorDistributormanagerConfig;
 import com.yahoo.vespa.config.content.StorFilestorConfig;
 import com.yahoo.vespa.config.content.core.StorServerConfig;
@@ -14,15 +20,20 @@ import com.yahoo.vespa.model.VespaModel;
 import com.yahoo.vespa.model.container.ContainerCluster;
 import com.yahoo.vespa.model.content.cluster.ContentCluster;
 import com.yahoo.vespa.model.content.engines.ProtonEngine;
+import com.yahoo.vespa.model.content.utils.ContentClusterBuilder;
+import com.yahoo.vespa.model.content.utils.ContentClusterUtils;
+import com.yahoo.vespa.model.content.utils.SearchDefinitionBuilder;
 import com.yahoo.vespa.model.test.utils.ApplicationPackageUtils;
 import com.yahoo.vespa.model.test.utils.VespaModelCreatorWithMockPkg;
 import org.junit.Test;
 
 import java.util.List;
 
+import static junit.framework.TestCase.assertEquals;
 import static org.hamcrest.CoreMatchers.is;
 import static org.junit.Assert.*;
 
+// TODO Rename to ContentClusterTest
 public class ClusterTest extends ContentBaseTest {
 
     private final static String HOSTS = "<admin version='2.0'><adminserver hostalias='mockhost' /></admin>";
@@ -412,20 +423,8 @@ public class ClusterTest extends ContentBaseTest {
             "</content>"
         );
 
-        {
-            FleetcontrollerConfig.Builder builder = new FleetcontrollerConfig.Builder();
-            cluster.getConfig(builder);
-            cluster.getClusterControllerConfig().getConfig(builder);
-            FleetcontrollerConfig config = new FleetcontrollerConfig(builder);
-            assertEquals(8, config.ideal_distribution_bits());
-        }
+        assertDistributionBitsInConfig(cluster, 8);
 
-        {
-            StorDistributormanagerConfig.Builder builder = new StorDistributormanagerConfig.Builder();
-            cluster.getConfig(builder);
-            StorDistributormanagerConfig config = new StorDistributormanagerConfig(builder);
-            assertEquals(8, config.minsplitcount());
-        }
         cluster = parse(
             "<content version=\"1.0\" id=\"storage\">\n" +
             "  <documents/>" +
@@ -438,20 +437,7 @@ public class ClusterTest extends ContentBaseTest {
             "</content>"
         );
 
-        {
-            FleetcontrollerConfig.Builder builder = new FleetcontrollerConfig.Builder();
-            cluster.getConfig(builder);
-            cluster.getClusterControllerConfig().getConfig(builder);
-            FleetcontrollerConfig config = new FleetcontrollerConfig(builder);
-            assertEquals(8, config.ideal_distribution_bits());
-        }
-
-        {
-            StorDistributormanagerConfig.Builder builder = new StorDistributormanagerConfig.Builder();
-            cluster.getConfig(builder);
-            StorDistributormanagerConfig config = new StorDistributormanagerConfig(builder);
-            assertEquals(8, config.minsplitcount());
-        }
+        assertDistributionBitsInConfig(cluster, 8);
     }
 
     @Test
@@ -469,20 +455,8 @@ public class ClusterTest extends ContentBaseTest {
             "</content>"
         );
 
-        {
-            FleetcontrollerConfig.Builder builder = new FleetcontrollerConfig.Builder();
-            cluster.getConfig(builder);
-            cluster.getClusterControllerConfig().getConfig(builder);
-            FleetcontrollerConfig config = new FleetcontrollerConfig(builder);
-            assertEquals(8, config.ideal_distribution_bits());
-        }
+        assertDistributionBitsInConfig(cluster, 8);
 
-        {
-            StorDistributormanagerConfig.Builder builder = new StorDistributormanagerConfig.Builder();
-            cluster.getConfig(builder);
-            StorDistributormanagerConfig config = new StorDistributormanagerConfig(builder);
-            assertEquals(8, config.minsplitcount());
-        }
         cluster = parse(
             "<content version=\"1.0\" id=\"storage\">\n" +
             "  <documents/>" +
@@ -498,23 +472,20 @@ public class ClusterTest extends ContentBaseTest {
             "</content>"
         );
 
-        {
-            FleetcontrollerConfig.Builder builder = new FleetcontrollerConfig.Builder();
-            cluster.getConfig(builder);
-            cluster.getClusterControllerConfig().getConfig(builder);
-            FleetcontrollerConfig config = new FleetcontrollerConfig(builder);
-            assertEquals(8, config.ideal_distribution_bits());
-        }
-
-        {
-            StorDistributormanagerConfig.Builder builder = new StorDistributormanagerConfig.Builder();
-            cluster.getConfig(builder);
-            StorDistributormanagerConfig config = new StorDistributormanagerConfig(builder);
-            assertEquals(8, config.minsplitcount());
-        }
+        assertDistributionBitsInConfig(cluster, 8);
     }
 
     @Test
+    public void testZoneDependentDistributionBits() throws Exception {
+        String xml = new ContentClusterBuilder().docTypes("test").getXml();
+
+        ContentCluster prodWith16Bits = createWithZone(xml, new Zone(Environment.prod, RegionName.from("us-east-3")));
+        assertDistributionBitsInConfig(prodWith16Bits, 16);
+
+        ContentCluster stagingNot16Bits = createWithZone(xml, new Zone(Environment.staging, RegionName.from("us-east-3")));
+        assertDistributionBitsInConfig(stagingNot16Bits, 8);
+    }
+    @Test
     public void testGenerateSearchNodes()
     {
         ContentCluster cluster = parse(
@@ -818,4 +789,31 @@ public class ClusterTest extends ContentBaseTest {
         assertThat(cluster.getSearch().getSearchNodes().size(), is(1));
         assertTrue(cluster.getSearch().getSearchNodes().get(0).getPreShutdownCommand().isPresent());
     }
+
+    private ContentCluster createWithZone(String clusterXml, Zone zone) throws Exception {
+        DeployState.Builder deployStateBuilder = new DeployState.Builder().properties(new DeployProperties.Builder()
+                                                                                              .hostedVespa(true)
+                                                                                              .zone(zone)
+                                                                                              .build());
+        List<String> searchDefinitions = SearchDefinitionBuilder.createSearchDefinitions("test");
+        MockRoot root = ContentClusterUtils.createMockRoot(searchDefinitions, deployStateBuilder);
+        ContentCluster cluster = ContentClusterUtils.createCluster(clusterXml, root);
+        root.freezeModelTopology();
+        cluster.validate();
+        return cluster;
+    }
+
+    private void assertDistributionBitsInConfig(ContentCluster cluster, int distributionBits) {
+        FleetcontrollerConfig.Builder builder = new FleetcontrollerConfig.Builder();
+        cluster.getConfig(builder);
+        cluster.getClusterControllerConfig().getConfig(builder);
+        FleetcontrollerConfig config = new FleetcontrollerConfig(builder);
+        assertEquals(distributionBits, config.ideal_distribution_bits());
+
+        StorDistributormanagerConfig.Builder sdBuilder = new StorDistributormanagerConfig.Builder();
+        cluster.getConfig(sdBuilder);
+        StorDistributormanagerConfig storDistributormanagerConfig = new StorDistributormanagerConfig(sdBuilder);
+        assertEquals(distributionBits, storDistributormanagerConfig.minsplitcount());
+    }
+
 }
diff --git a/config-model/src/test/java/com/yahoo/vespa/model/content/utils/ContentClusterUtils.java b/config-model/src/test/java/com/yahoo/vespa/model/content/utils/ContentClusterUtils.java
index 1ac892f551a..aebe1b76644 100644
--- a/config-model/src/test/java/com/yahoo/vespa/model/content/utils/ContentClusterUtils.java
+++ b/config-model/src/test/java/com/yahoo/vespa/model/content/utils/ContentClusterUtils.java
@@ -32,14 +32,17 @@ public class ContentClusterUtils {
     }
 
     private static MockRoot createMockRoot(HostProvisioner provisioner, List<String> searchDefinitions) {
-        ApplicationPackage applicationPackage = new MockApplicationPackage.Builder().withSearchDefinitions(searchDefinitions).build();
-        DeployState deployState = new DeployState.Builder()
-                .applicationPackage(applicationPackage)
-                .modelHostProvisioner(provisioner)
-                .build();
-        return new MockRoot("", deployState);
+        return createMockRoot(provisioner, searchDefinitions, new DeployState.Builder());
+    }
 
+    private static MockRoot createMockRoot(HostProvisioner provisioner, List<String> searchDefinitions, DeployState.Builder deployStateBuilder) {
+        ApplicationPackage applicationPackage = new MockApplicationPackage.Builder().withSearchDefinitions(searchDefinitions).build();
+        deployStateBuilder.applicationPackage(applicationPackage)
+                          .modelHostProvisioner(provisioner)
+                          .build();
+        return new MockRoot("", deployStateBuilder.build());
     }
+
     public static MockRoot createMockRoot(String[] hosts, List<String> searchDefinitions) throws Exception {
         return createMockRoot(new InMemoryProvisioner(true, hosts), searchDefinitions);
     }
@@ -48,6 +51,10 @@ public class ContentClusterUtils {
         return createMockRoot(new SingleNodeProvisioner(), searchDefinitions);
     }
 
+    public static MockRoot createMockRoot(List<String> searchDefinitions, DeployState.Builder deployStateBuilder) {
+        return createMockRoot(new SingleNodeProvisioner(), searchDefinitions, deployStateBuilder);
+    }
+
     public static ContentCluster createCluster(String clusterXml, MockRoot root) throws Exception {
         Document doc = XML.getDocument(clusterXml);
         Admin admin = new Admin(root, new Yamas("vespa", 60), new Metrics(), Collections.emptyMap(), false);