Allow overriding cluster configs in quickstarts

apache · Jan 25, 2025 · d5848d8 · d5848d8
1 parent c239952
commit d5848d8
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 3 deletions.
diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/MultistageEngineQuickStart.java b/pinot-tools/src/main/java/org/apache/pinot/tools/MultistageEngineQuickStart.java
@@ -297,6 +297,13 @@ protected Map<String, Object> getConfigOverrides() {
     return configOverrides;
   }
 
+  @Override
+  protected Map<String, String> getClusterConfigOverrides() {
+    return Map.of(
+        CommonConstants.Helix.CONFIG_OF_MULTI_STAGE_ENGINE_MAX_SERVER_QUERY_THREADS, "50"
+    );
+  }
+
   @Override
   protected int getNumQuickstartRunnerServers() {
     return 3;

diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/QuickStartBase.java b/pinot-tools/src/main/java/org/apache/pinot/tools/QuickStartBase.java
@@ -291,6 +291,10 @@ protected Map<String, Object> getConfigOverrides() {
     }
   }
 
+  protected Map<String, String> getClusterConfigOverrides() {
+    return Map.of();
+  }
+
   protected String[] getDefaultBatchTableDirectories() {
     return DEFAULT_OFFLINE_TABLE_DIRECTORIES;
   }

diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/Quickstart.java b/pinot-tools/src/main/java/org/apache/pinot/tools/Quickstart.java
@@ -67,7 +67,7 @@ public void execute()
 
     QuickstartRunner runner =
         new QuickstartRunner(quickstartTableRequests, 1, 1, getNumQuickstartRunnerServers(), 1, quickstartRunnerDir,
-            true, getAuthProvider(), getConfigOverrides(), _zkExternalAddress, true);
+            true, getAuthProvider(), getConfigOverrides(), _zkExternalAddress, true, getClusterConfigOverrides());
 
     printStatus(Color.CYAN, "***** Starting Zookeeper, controller, broker, server and minion *****");
     runner.startAll();

diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/QuickstartRunner.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/QuickstartRunner.java
@@ -29,6 +29,11 @@
 import java.util.Map;
 import java.util.Random;
 import org.apache.commons.io.FileUtils;
+import org.apache.helix.HelixManager;
+import org.apache.helix.HelixManagerFactory;
+import org.apache.helix.InstanceType;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
 import org.apache.pinot.spi.auth.AuthProvider;
 import org.apache.pinot.spi.config.tenant.TenantRole;
 import org.apache.pinot.spi.env.PinotConfiguration;
@@ -69,6 +74,7 @@ public class QuickstartRunner {
   private final boolean _enableTenantIsolation;
   private final AuthProvider _authProvider;
   private final Map<String, Object> _configOverrides;
+  private final Map<String, String> _clusterConfigOverrides;
   private final boolean _deleteExistingData;
 
   // If this field is non-null, an embedded Zookeeper instance will not be launched
@@ -82,12 +88,13 @@ public QuickstartRunner(List<QuickstartTableRequest> tableRequests, int numContr
       int numServers, int numMinions, File tempDir, Map<String, Object> configOverrides)
       throws Exception {
     this(tableRequests, numControllers, numBrokers, numServers, numMinions, tempDir, true, null, configOverrides, null,
-        true);
+        true, Map.of());
   }
 
   public QuickstartRunner(List<QuickstartTableRequest> tableRequests, int numControllers, int numBrokers,
       int numServers, int numMinions, File tempDir, boolean enableIsolation, AuthProvider authProvider,
-      Map<String, Object> configOverrides, String zkExternalAddress, boolean deleteExistingData)
+      Map<String, Object> configOverrides, String zkExternalAddress, boolean deleteExistingData,
+      Map<String, String> clusterConfigOverrides)
       throws Exception {
     _tableRequests = tableRequests;
     _numControllers = numControllers;
@@ -98,6 +105,7 @@ public QuickstartRunner(List<QuickstartTableRequest> tableRequests, int numContr
     _enableTenantIsolation = enableIsolation;
     _authProvider = authProvider;
     _configOverrides = new HashMap<>(configOverrides);
+    _clusterConfigOverrides = clusterConfigOverrides;
     if (numMinions > 0) {
       // configure the controller to schedule tasks when minion is enabled
       _configOverrides.put("controller.task.scheduler.enabled", true);
@@ -127,11 +135,22 @@ private void startControllers()
           .setTenantIsolation(_enableTenantIsolation)
           .setDataDir(new File(_tempDir, DEFAULT_CONTROLLER_DIR + i).getAbsolutePath())
           .setConfigOverrides(_configOverrides);
+
       if (!controllerStarter.execute()) {
         throw new RuntimeException("Failed to start Controller");
       }
       _controllerPorts.add(DEFAULT_CONTROLLER_PORT + i);
     }
+
+    HelixManager helixManager =
+        HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "localhost_" + _controllerPorts.get(0),
+            InstanceType.CONTROLLER, _zkExternalAddress != null ? _zkExternalAddress : ZK_ADDRESS);
+    helixManager.connect();
+    HelixConfigScope scope =
+        new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).forCluster(CLUSTER_NAME).build();
+    for (Map.Entry<String, String> entry : _clusterConfigOverrides.entrySet()) {
+      helixManager.getConfigAccessor().set(scope, entry.getKey(), entry.getValue());
+    }
   }
 
   private void startBrokers()