Issue:3933 -Multi-node setup for local gradle run

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,21 @@ public OpenSearchCluster(
         this.bwcJdk = bwcJdk;
 
         // Always add the first node
-        addNode(clusterName + "-0");
+        String zone = hasZoneProperty() ? "zone-1" : "";
+        addNode(clusterName + "-0", zone);
         // 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 +134,31 @@ public void setNumberOfNodes(int numberOfNodes) {
             );
         }
 
-        for (int i = nodes.size(); i < numberOfNodes; i++) {
-            addNode(clusterName + "-" + i);
+        if (numberOfNodes < zoneCount) {
+            throw new IllegalArgumentException(
+                "Number of nodes should be >= zoneCount but was " + numberOfNodes + " for " + this.zoneCount
+            );
         }
+
+        if (hasZoneProperty()) {
+            int currentZone;
+            for (int i = nodes.size(); i < numberOfNodes; i++) {
+                currentZone = i % zoneCount + 1;
+                String zoneName = "zone-" + currentZone;
+                addNode(clusterName + "-" + i, zoneName);
+            }
+        } else {
+            for (int i = nodes.size(); i < numberOfNodes; i++) {
+                addNode(clusterName + "-" + i, "");
+            }
+        }
+    }
+
+    private boolean hasZoneProperty() {
+        return this.project.findProperty("numZones") != null;
     }
 
-    private void addNode(String nodeName) {
+    private void addNode(String nodeName, String zoneName) {
         OpenSearchNode newNode = new OpenSearchNode(
             path,
             nodeName,
@@ -138,7 +167,8 @@ private void addNode(String nodeName) {
             fileSystemOperations,
             archiveOperations,
             workingDirBase,
-            bwcJdk
+            bwcJdk,
+            zoneName
         );
         // configure the cluster name eagerly
         newNode.defaultConfig.put("cluster.name", safeName(clusterName));

buildSrc/src/main/java/org/opensearch/gradle/testclusters/OpenSearchNode.java

@@ -32,6 +32,7 @@
 package org.opensearch.gradle.testclusters;
 
 import org.apache.commons.io.FileUtils;
+import org.apache.commons.lang.StringUtils;
 import org.opensearch.gradle.Architecture;
 import org.opensearch.gradle.DistributionDownloadPlugin;
 import org.opensearch.gradle.OpenSearchDistribution;
@@ -175,6 +176,8 @@ public class OpenSearchNode implements TestClusterConfiguration {
     private final Config legacyESConfig;
     private Config currentConfig;
 
+    private String zone;
+
     OpenSearchNode(
         String path,
         String name,
@@ -183,7 +186,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 +209,7 @@ public class OpenSearchNode implements TestClusterConfiguration {
         opensearchConfig = Config.getOpenSearchConfig(workingDir);
         legacyESConfig = Config.getLegacyESConfig(workingDir);
         currentConfig = opensearchConfig;
+        this.zone = zone;
     }
 
     /*
@@ -1239,6 +1244,10 @@ 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");
+        if (StringUtils.isNotBlank(zone)) {
+            baseConfig.put("cluster.routing.allocation.awareness.attributes", "zone");
+            baseConfig.put("node.attr.zone", zone);
+        }
         baseConfig.put("node.portsfile", "true");
         baseConfig.put("http.port", httpPort);
         if (getVersion().onOrAfter(Version.fromString("6.7.0"))) {

gradle/run.gradle

@@ -31,9 +31,14 @@ import org.opensearch.gradle.testclusters.RunTask
 
 apply plugin: 'opensearch.testclusters'
 
+def numNodes = findProperty('numNodes') as Integer ?: 1
+def numZones = findProperty('numZones') as Integer ?: 1
+
 testClusters {
   runTask {
     testDistribution = 'archive'
+    if (numZones > 1) numberOfZones = numZones
+    if (numNodes > 1) numberOfNodes = numNodes
   }
 }