From 171838ac918caaa799a185133d2e2b5aa508568f Mon Sep 17 00:00:00 2001 From: Yash Mayya Date: Sat, 25 Jan 2025 18:44:48 +0530 Subject: [PATCH] Allow overriding cluster configs in quickstarts --- .../tools/MultistageEngineQuickStart.java | 7 ++++++ .../apache/pinot/tools/QuickStartBase.java | 4 ++++ .../org/apache/pinot/tools/Quickstart.java | 2 +- .../tools/admin/command/QuickstartRunner.java | 23 +++++++++++++++++-- 4 files changed, 33 insertions(+), 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 index e249c37c0b70..447544aa8723 100644 --- 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 getConfigOverrides() { return configOverrides; } + @Override + protected Map 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 index b8296e40937f..388192d62399 100644 --- 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 getConfigOverrides() { } } + protected Map 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 index b5894ca968d9..0ebb0053b7fd 100644 --- 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 index 77cf329e0def..01b37ca3a51c 100644 --- 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 _configOverrides; + private final Map _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 tableRequests, int numContr int numServers, int numMinions, File tempDir, Map configOverrides) throws Exception { this(tableRequests, numControllers, numBrokers, numServers, numMinions, tempDir, true, null, configOverrides, null, - true); + true, Map.of()); } public QuickstartRunner(List tableRequests, int numControllers, int numBrokers, int numServers, int numMinions, File tempDir, boolean enableIsolation, AuthProvider authProvider, - Map configOverrides, String zkExternalAddress, boolean deleteExistingData) + Map configOverrides, String zkExternalAddress, boolean deleteExistingData, + Map clusterConfigOverrides) throws Exception { _tableRequests = tableRequests; _numControllers = numControllers; @@ -98,6 +105,7 @@ public QuickstartRunner(List 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.ADMINISTRATOR, _zkExternalAddress != null ? _zkExternalAddress : ZK_ADDRESS); + helixManager.connect(); + HelixConfigScope scope = + new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).forCluster(CLUSTER_NAME).build(); + for (Map.Entry entry : _clusterConfigOverrides.entrySet()) { + helixManager.getConfigAccessor().set(scope, entry.getKey(), entry.getValue()); + } } private void startBrokers()