Issue:3933 - Initial Checkin for Multinode setup for local gradle run

Signed-off-by: Pranit Kumar <[email protected]>
opensearch-project · Jul 21, 2022 · 1f54bd2 · 1f54bd2
1 parent 6d7a152
commit 1f54bd2
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 6 deletions.
diff --git a/buildSrc/src/main/java/org/opensearch/gradle/testclusters/OpenSearchCluster.java b/buildSrc/src/main/java/org/opensearch/gradle/testclusters/OpenSearchCluster.java
@@ -84,6 +84,8 @@ public class OpenSearchCluster implements TestClusterConfiguration, Named {
     private final ArchiveOperations archiveOperations;
     private int nodeIndex = 0;
 
+    private int zoneCount = 1;
+
     public OpenSearchCluster(
         String clusterName,
         Project project,
@@ -104,13 +106,20 @@ public OpenSearchCluster(
         this.bwcJdk = bwcJdk;
 
         // Always add the first node
-        addNode(clusterName + "-0");
+        addNode(clusterName + "-0", "zone-1");
         // configure the cluster name eagerly so all nodes know about it
         this.nodes.all((node) -> node.defaultConfig.put("cluster.name", safeName(clusterName)));
 
         addWaitForClusterHealth();
     }
 
+    public void setNumberOfZones(int zoneCount) {
+        if (zoneCount < 1) {
+            throw new IllegalArgumentException("Number of zones should be >= 1 but was " + zoneCount + " for " + this);
+        }
+        this.zoneCount = zoneCount;
+    }
+
     public void setNumberOfNodes(int numberOfNodes) {
         checkFrozen();
 
@@ -124,12 +133,20 @@ public void setNumberOfNodes(int numberOfNodes) {
             );
         }
 
+        if(numberOfNodes < zoneCount) {
+            throw new IllegalArgumentException("Number of nodes should be >= zoneCount but was " + numberOfNodes + " for " + this.zoneCount);
+        }
+
+        int currentZone = 1;
         for (int i = nodes.size(); i < numberOfNodes; i++) {
-            addNode(clusterName + "-" + i);
+            //addNode(clusterName + "-" + i);
+            currentZone = (currentZone >= zoneCount) ? 1 : (currentZone + 1);
+            String zoneName = "zone-" + currentZone;
+            addNode(clusterName + "-" + i, zoneName);
         }
     }
 
-    private void addNode(String nodeName) {
+    private void addNode(String nodeName, String zoneName) {
         OpenSearchNode newNode = new OpenSearchNode(
             path,
             nodeName,
@@ -138,7 +155,8 @@ private void addNode(String nodeName) {
             fileSystemOperations,
             archiveOperations,
             workingDirBase,
-            bwcJdk
+            bwcJdk,
+            zoneName
         );
         // configure the cluster name eagerly
         newNode.defaultConfig.put("cluster.name", safeName(clusterName));

diff --git a/buildSrc/src/main/java/org/opensearch/gradle/testclusters/OpenSearchNode.java b/buildSrc/src/main/java/org/opensearch/gradle/testclusters/OpenSearchNode.java
@@ -175,6 +175,8 @@ public class OpenSearchNode implements TestClusterConfiguration {
     private final Config legacyESConfig;
     private Config currentConfig;
 
+    private String zone;
+
     OpenSearchNode(
         String path,
         String name,
@@ -183,7 +185,8 @@ public class OpenSearchNode implements TestClusterConfiguration {
         FileSystemOperations fileSystemOperations,
         ArchiveOperations archiveOperations,
         File workingDirBase,
-        Jdk bwcJdk
+        Jdk bwcJdk,
+        String zone
     ) {
         this.path = path;
         this.name = name;
@@ -205,6 +208,7 @@ public class OpenSearchNode implements TestClusterConfiguration {
         opensearchConfig = Config.getOpenSearchConfig(workingDir);
         legacyESConfig = Config.getLegacyESConfig(workingDir);
         currentConfig = opensearchConfig;
+        this.zone = zone;
     }
 
     /*
@@ -1239,6 +1243,7 @@ private void createConfiguration() {
         baseConfig.put("path.logs", confPathLogs.toAbsolutePath().toString());
         baseConfig.put("path.shared_data", workingDir.resolve("sharedData").toString());
         baseConfig.put("node.attr.testattr", "test");
+        baseConfig.put("node.attr.availabilityzone", zone);
         baseConfig.put("node.portsfile", "true");
         baseConfig.put("http.port", httpPort);
         if (getVersion().onOrAfter(Version.fromString("6.7.0"))) {

diff --git a/gradle/run.gradle b/gradle/run.gradle
@@ -37,8 +37,31 @@ testClusters {
   }
 }
 
+testClusters {
+  runMultiNodeTask {
+    testDistribution = 'archive'
+  }
+}
+
+final String nodes = findProperty('nodeCount') ?: "3"
+final String zones = findProperty('zoneCount') ?: "1"
+final int nodeCount = Integer.valueOf(nodes)
+final int zoneCount = Integer.valueOf(zones)
+
+testClusters {
+  runMultiNodeTask {
+    testDistribution = 'archive'
+    numberOfZones = zoneCount
+    numberOfNodes = nodeCount
+  }
+}
+
 tasks.register("run", RunTask) {
-  useCluster testClusters.runTask;
+  if (project.hasProperty('nodeCount')) {
+    useCluster testClusters.runMultiNodeTask;
+  } else {
+    useCluster testClusters.runTask;
+  }
   description = 'Runs opensearch in the foreground'
   group = 'Verification'