This is an automated email from the ASF dual-hosted git repository. xiangfu0 pushed a commit to branch xiangfu0/mse-proto-segment-list in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 04ee254c39176b62bc7d419fd50f3259c24545db Author: Xiang Fu <[email protected]> AuthorDate: Tue Sep 15 17:06:00 2026 -0700 Stop JSON-encoding MSE leaf-stage segment lists on the broker compile path On tables with many segments, the top broker CPU consumer of a multi-stage query was `WorkerMetadata.setTableSegmentsMap`: for every leaf-stage worker the planner JSON-encoded the full list of routed segment names with Jackson, on the fixed-size `multi-stage-query-compile-executor`, only for the dispatcher to copy that string into a proto `map<string, string>` custom property and for every server to JSON-parse it back (once per worker). - `WorkerMetadata` now holds the segment maps as plain objects; nothing is encoded at plan time anymore. - `QueryPlanSerDeUtils` encodes them per server at dispatch time, once per worker, in one of two wire encodings picked per request: - proto: new native `tableSegmentsMap` / `logicalTableSegmentsMap` fields of `Worker.WorkerMetadata` (`SegmentsMap` / `SegmentList` messages). - legacy JSON: the previous custom-property string, still the default. Decoding accepts both, so a new server understands every broker. - New broker config `pinot.broker.mse.proto.segment.list` (default false) and query option `protoSegmentList` to enable the proto encoding. It has to stay opt-in for one release: an older server finds no segments under the proto fields and fails the leaf stage, so it may only be enabled once the whole fleet is upgraded. - That config is live: `ProtoSegmentListPredicate` seeds it from the static broker config and then follows cluster config on the same key, registered by `BaseBrokerStarter` against the existing cluster-config change handler, and `QueryDispatcher` reads it per request. The precondition is exactly "every server already upgraded", so operators need to turn it on the moment a rolling upgrade completes, and to turn it back off at once if it misbehaves; neither should cost a broker restart. Precedence is per-query `SET protoSegmentList`, then cluster config, then static broker config, and clearing the cluster-config key falls back to the shipped `false` rather than to the static seed. - New broker config `pinot.broker.mse.query.compile.executor.threads`; the compile executor was hard-wired to half the cores and saturates (queueing every query) well before the broker does on short-query workloads. Micro-benchmark, one leaf worker, 60-char segment names (Java 25, x86): segments | legacy JSON encode | proto encode | legacy decode | proto decode 1,000 | 177 us | 132 us | 66 us | 29 us 3,000 | 624 us | 394 us | 200 us | 95 us 20,000 | 5,101 us | 2,709 us | 1,419 us | 782 us 60,000 | 15,950 us | 8,217 us | 4,602 us | 2,386 us --- .../broker/broker/helix/BaseBrokerStarter.java | 6 + .../MultiStageBrokerRequestHandler.java | 25 ++- .../MultiStageBrokerRequestHandlerTest.java | 13 ++ .../common/utils/config/QueryOptionsUtils.java | 5 + pinot-common/src/main/proto/worker.proto | 15 ++ .../common/utils/config/QueryOptionsUtilsTest.java | 8 + .../planner/physical/DispatchablePlanContext.java | 12 +- .../planner/physical/PinotDispatchPlanner.java | 2 + .../pinot/query/routing/QueryPlanSerDeUtils.java | 124 +++++++++++++-- .../apache/pinot/query/routing/WorkerMetadata.java | 66 +++----- .../query/routing/QueryPlanSerDeUtilsTest.java | 173 +++++++++++++++++++++ .../dispatch/ProtoSegmentListPredicate.java | 109 +++++++++++++ .../query/service/dispatch/QueryDispatcher.java | 28 +++- .../pinot/query/service/server/QueryServer.java | 6 +- .../dispatch/ProtoSegmentListPredicateTest.java | 146 +++++++++++++++++ .../query/service/server/QueryServerAuthzTest.java | 2 +- .../query/service/server/QueryServerTest.java | 21 ++- .../apache/pinot/spi/utils/CommonConstants.java | 26 ++++ 18 files changed, 718 insertions(+), 69 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java index 826c732d4d5..8b36ea76383 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java @@ -734,6 +734,12 @@ public abstract class BaseBrokerStarter implements ServiceStartable { _clusterConfigChangeHandler.registerClusterConfigChangeListener(ContinuousJfrStarter.INSTANCE); _clusterConfigChangeHandler.registerClusterConfigChangeListener(_serverRoutingStatsManager); + if (multiStageBrokerRequestHandler != null) { + // Registration replays the current cluster config, so a cluster-config value takes effect right away and wins + // over the static broker config seed. + _clusterConfigChangeHandler.registerClusterConfigChangeListener( + multiStageBrokerRequestHandler.getProtoSegmentListPredicate()); + } NettyInspector.registerMetrics(_brokerMetrics); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java index 08005e3be51..8ff768abe9a 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java @@ -90,6 +90,7 @@ import org.apache.pinot.query.routing.WorkerManager; import org.apache.pinot.query.runtime.MultiStageStatsTreeBuilder; import org.apache.pinot.query.runtime.plan.MultiStageQueryStats; import org.apache.pinot.query.runtime.plan.StageStatsTreeNode; +import org.apache.pinot.query.service.dispatch.ProtoSegmentListPredicate; import org.apache.pinot.query.service.dispatch.QueryDispatcher; import org.apache.pinot.spi.accounting.ThreadAccountant; import org.apache.pinot.spi.auth.TableAuthorizationResult; @@ -145,6 +146,9 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler { protected final long _extraPassiveTimeoutMs; protected final boolean _enableQueryFingerprinting; private final boolean _streamStatsDefault; + /// Live cluster-config default for the proto segment list encoding. Registered as a cluster-config + /// listener by the broker starter so operators can flip it without restarting the brokers. + private final ProtoSegmentListPredicate _protoSegmentListPredicate; @Nullable protected final String _defaultStreamingGroupByFlushThreshold; @Nullable @@ -210,11 +214,12 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler { long streamStatsDrainMs = _config.getProperty( CommonConstants.Broker.CONFIG_OF_STREAM_STATS_DRAIN_MS, CommonConstants.Broker.DEFAULT_STREAM_STATS_DRAIN_MS); + _protoSegmentListPredicate = ProtoSegmentListPredicate.create(_config); _mailboxService = new MailboxService(hostname, port, InstanceType.BROKER, config, tlsConfig); _queryDispatcher = new QueryDispatcher(_mailboxService, failureDetector, tlsConfig, isQueryCancellationEnabled(), cancelTimeout, dispatchKeepAliveTimeMs, dispatchKeepAliveTimeoutMs, dispatchKeepAliveWithoutCalls, _streamStatsDefault, - streamStatsDrainMs); + streamStatsDrainMs, _protoSegmentListPredicate); LOGGER.info("Initialized MultiStageBrokerRequestHandler on host: {}, port: {} with broker id: {}, timeout: {}ms, " + "query log max length: {}, query log max rate: {}, query cancellation enabled: {}", hostname, port, _brokerId, _brokerTimeoutMs, _queryLogger.getMaxQueryLengthToLog(), _queryLogger.getLogRateLimit(), @@ -224,8 +229,7 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler { CommonConstants.MultiStageQueryRunner.DEFAULT_OF_MULTISTAGE_EXPLAIN_INCLUDE_SEGMENT_PLAN); _queryThrottler = queryThrottler; _queryCompileExecutor = QueryThreadContext.contextAwareExecutorService( - Executors.newFixedThreadPool( - Math.max(1, Runtime.getRuntime().availableProcessors() / 2), + Executors.newFixedThreadPool(resolveQueryCompileExecutorThreads(config), new NamedThreadFactory("multi-stage-query-compile-executor"))); _defaultDisabledPlannerRules = _config.containsKey(CommonConstants.Broker.CONFIG_OF_BROKER_MSE_PLANNER_DISABLED_RULES) ? Set.copyOf( @@ -602,6 +606,15 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler { .build(); } + /// Size of the query-compile executor: the configured thread count, or half the available processors (at least 1) + /// when the config is absent or non-positive. + @VisibleForTesting + static int resolveQueryCompileExecutorThreads(PinotConfiguration config) { + int configured = config.getProperty(CommonConstants.Broker.CONFIG_OF_MSE_QUERY_COMPILE_EXECUTOR_THREADS, + CommonConstants.Broker.DEFAULT_MSE_QUERY_COMPILE_EXECUTOR_THREADS); + return configured > 0 ? configured : Math.max(1, Runtime.getRuntime().availableProcessors() / 2); + } + /// Applies broker-level defaults for MSE query options. Per-query overrides (i.e. `SET option = value` in the /// SQL text) always win because we use [Map#putIfAbsent] — a user can set the option to `0` to opt out of /// a streaming default that the cluster has enabled. @@ -1086,4 +1099,10 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler { public QueryDispatcher getQueryDispatcher() { return _queryDispatcher; } + + /// The live default for the proto segment list encoding, to be registered as a cluster-config change listener by + /// the broker starter. + public ProtoSegmentListPredicate getProtoSegmentListPredicate() { + return _protoSegmentListPredicate; + } } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java index e04423ccf31..1687d371025 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java @@ -127,6 +127,19 @@ public class MultiStageBrokerRequestHandlerTest extends QueryEnvironmentTestBase "onQueryCompletion hook must be called with the BrokerResponse from handleRequest for MSE"); } + @Test + public void testResolveQueryCompileExecutorThreads() { + int defaultThreads = Math.max(1, Runtime.getRuntime().availableProcessors() / 2); + Assert.assertEquals(MultiStageBrokerRequestHandler.resolveQueryCompileExecutorThreads(new PinotConfiguration()), + defaultThreads, "Absent config must size the executor to half the processors"); + PinotConfiguration config = new PinotConfiguration(); + config.setProperty(CommonConstants.Broker.CONFIG_OF_MSE_QUERY_COMPILE_EXECUTOR_THREADS, "12"); + Assert.assertEquals(MultiStageBrokerRequestHandler.resolveQueryCompileExecutorThreads(config), 12); + config.setProperty(CommonConstants.Broker.CONFIG_OF_MSE_QUERY_COMPILE_EXECUTOR_THREADS, "0"); + Assert.assertEquals(MultiStageBrokerRequestHandler.resolveQueryCompileExecutorThreads(config), defaultThreads, + "Non-positive config must fall back to the default sizing"); + } + @Test public void testApplyBrokerDefaultQueryOptionsInjectsStreamingGroupByFlushThreshold() throws Exception { diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java index a48c88cd146..78e74e0a0e0 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java @@ -721,6 +721,11 @@ public class QueryOptionsUtils { return option != null ? Boolean.parseBoolean(option) : defaultValue; } + public static boolean isProtoSegmentList(Map<String, String> queryOptions, boolean defaultValue) { + String option = queryOptions.get(QueryOptionKey.PROTO_SEGMENT_LIST); + return option != null ? Boolean.parseBoolean(option) : defaultValue; + } + public static boolean isMultiClusterRoutingEnabled(Map<String, String> queryOptions, boolean defaultValue) { String option = queryOptions.get(QueryOptionKey.ENABLE_MULTI_CLUSTER_ROUTING); return option != null ? Boolean.parseBoolean(option) : defaultValue; diff --git a/pinot-common/src/main/proto/worker.proto b/pinot-common/src/main/proto/worker.proto index ceca82b65ea..91f13b2d903 100644 --- a/pinot-common/src/main/proto/worker.proto +++ b/pinot-common/src/main/proto/worker.proto @@ -94,6 +94,21 @@ message WorkerMetadata { int32 workedId = 1; map<int32, bytes> mailboxInfos = 2; // Stage id to serialized MailboxInfos map<string, string> customProperty = 3; + // Leaf-stage segments to scan, keyed by table type (OFFLINE / REALTIME). Presence marks a leaf-stage worker. + // Sent instead of the JSON-encoded "tableSegmentsMap" custom property when the broker enables the proto segment + // list encoding (see the "protoSegmentList" query option); older brokers keep sending the custom property. + SegmentsMap tableSegmentsMap = 4; + // Leaf-stage segments of a logical table, keyed by physical table name (with type suffix). Same encoding rules. + SegmentsMap logicalTableSegmentsMap = 5; +} + +// Segment names to scan per table type or physical table name, see WorkerMetadata. +message SegmentsMap { + map<string, SegmentList> segments = 1; +} + +message SegmentList { + repeated string segment = 1; } message MailboxInfos { diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java index 0b8155403aa..18711317b96 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java @@ -123,6 +123,14 @@ public class QueryOptionsUtilsTest { QueryOptionsUtils.getInPredicatePruningThreshold(Map.of(IN_PREDICATE_PRUNING_THRESHOLD, "invalid")); } + @Test + public void testIsProtoSegmentList() { + assertFalse(QueryOptionsUtils.isProtoSegmentList(Map.of(), false)); + assertTrue(QueryOptionsUtils.isProtoSegmentList(Map.of(), true)); + assertTrue(QueryOptionsUtils.isProtoSegmentList(Map.of(PROTO_SEGMENT_LIST, "true"), false)); + assertFalse(QueryOptionsUtils.isProtoSegmentList(Map.of(PROTO_SEGMENT_LIST, "false"), true)); + } + @Test public void testSkipIndexesParsing() { String skipIndexesStr = "col1=inverted,range&col2=sorted"; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java index 585576dfa99..ba255e28fd4 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java @@ -198,11 +198,19 @@ public class DispatchablePlanContext { QueryServerInstance queryServerInstance = serverEntry.getValue(); serverInstanceToWorkerIdsMap.computeIfAbsent(queryServerInstance, k -> new ArrayList<>()).add(workerId); WorkerMetadata workerMetadata = new WorkerMetadata(workerId, workerIdToMailboxesMap.get(workerId)); + // A leaf-stage worker is identified by carrying a (possibly empty) segment map, so every worker of a + // scanning stage has to be present in the map. Fail loudly here instead of letting the worker decay + // into an intermediate-stage worker on the server. if (workerIdToSegmentsMap != null) { - workerMetadata.setTableSegmentsMap(workerIdToSegmentsMap.get(workerId)); + Map<String, List<String>> segmentsMap = workerIdToSegmentsMap.get(workerId); + Preconditions.checkNotNull(segmentsMap, "Missing segments map for worker id: %s", workerId); + workerMetadata.setTableSegmentsMap(segmentsMap); } if (workerIdToTableNameSegmentsMap != null) { - workerMetadata.setLogicalTableSegmentsMap(workerIdToTableNameSegmentsMap.get(workerId)); + Map<String, List<String>> tableNameSegmentsMap = workerIdToTableNameSegmentsMap.get(workerId); + Preconditions.checkNotNull(tableNameSegmentsMap, "Missing logical table segments map for worker id: %s", + workerId); + workerMetadata.setLogicalTableSegmentsMap(tableNameSegmentsMap); } workerMetadataArray[workerId] = workerMetadata; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java index a5baa46428b..37c2cf633bb 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java @@ -217,6 +217,8 @@ public class PinotDispatchPlanner { fragmentMap.put(0, reduceStage); } WorkerMetadata workerMetadata = workerMetadataList.get(0); + // Deliberately rebuilt without the leaf-stage segment maps: every TableScanNode has just been inlined into + // a ValueNode, so this stage-0 worker scans nothing and must not look like a leaf-stage worker. reduceStage.setWorkerMetadataList(List.of( new WorkerMetadata(workerMetadata.getWorkerId(), Map.of(), workerMetadata.getCustomProperties()))); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/QueryPlanSerDeUtils.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/QueryPlanSerDeUtils.java index 955f05f8d5e..361fce46038 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/QueryPlanSerDeUtils.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/QueryPlanSerDeUtils.java @@ -18,10 +18,15 @@ */ package org.apache.pinot.query.routing; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; +import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.Maps; import com.google.protobuf.ByteString; import com.google.protobuf.InvalidProtocolBufferException; +import java.io.IOException; import java.util.ArrayList; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.stream.Collectors; @@ -29,10 +34,24 @@ import org.apache.pinot.common.proto.Plan; import org.apache.pinot.common.proto.Worker; import org.apache.pinot.query.planner.plannode.PlanNode; import org.apache.pinot.query.planner.serde.PlanNodeDeserializer; +import org.apache.pinot.spi.utils.JsonUtils; /// This utility class serialize/deserialize between [Worker.StagePlan] elements to Planner elements. +/// +/// The leaf-stage segment maps of a [WorkerMetadata] have two wire encodings, picked per request by the broker: +/// +/// - **proto**: the native `tableSegmentsMap` / `logicalTableSegmentsMap` fields of [Worker.WorkerMetadata]. +/// - **legacy JSON**: a JSON string under the [WorkerMetadata#TABLE_SEGMENTS_MAP_KEY] / +/// [WorkerMetadata#LOGICAL_TABLE_SEGMENTS_MAP_KEY] custom property, which is all that servers predating the proto +/// fields understand. +/// +/// Decoding accepts both, so a server always understands every broker; the broker enables the proto encoding only +/// when every server does (see the `protoSegmentList` query option). public class QueryPlanSerDeUtils { + private static final TypeReference<Map<String, List<String>>> SEGMENTS_MAP_TYPE = new TypeReference<>() { + }; + private QueryPlanSerDeUtils() { } @@ -54,15 +73,55 @@ public class QueryPlanSerDeUtils { return new StageMetadata(protoStageMetadata.getStageId(), workerMetadataList, customProperties); } - private static WorkerMetadata fromProtoWorkerMetadata(Worker.WorkerMetadata protoWorkerMetadata) + @VisibleForTesting + static WorkerMetadata fromProtoWorkerMetadata(Worker.WorkerMetadata protoWorkerMetadata) throws InvalidProtocolBufferException { Map<Integer, ByteString> protoMailboxInfosMap = protoWorkerMetadata.getMailboxInfosMap(); Map<Integer, MailboxInfos> mailboxInfosMap = Maps.newHashMapWithExpectedSize(protoMailboxInfosMap.size()); for (Map.Entry<Integer, ByteString> entry : protoMailboxInfosMap.entrySet()) { mailboxInfosMap.put(entry.getKey(), fromProtoMailboxInfos(entry.getValue())); } - return new WorkerMetadata(protoWorkerMetadata.getWorkedId(), mailboxInfosMap, - protoWorkerMetadata.getCustomPropertyMap()); + // A broker using the legacy encoding ships the segment maps as JSON custom properties. Decode them once here and + // drop the raw strings so that the metadata never carries two copies of the same segments. + Map<String, String> customProperties = protoWorkerMetadata.getCustomPropertyMap(); + String tableSegmentsJson = customProperties.get(WorkerMetadata.TABLE_SEGMENTS_MAP_KEY); + String logicalTableSegmentsJson = customProperties.get(WorkerMetadata.LOGICAL_TABLE_SEGMENTS_MAP_KEY); + if (tableSegmentsJson != null || logicalTableSegmentsJson != null) { + customProperties = new HashMap<>(customProperties); + customProperties.remove(WorkerMetadata.TABLE_SEGMENTS_MAP_KEY); + customProperties.remove(WorkerMetadata.LOGICAL_TABLE_SEGMENTS_MAP_KEY); + } + WorkerMetadata workerMetadata = + new WorkerMetadata(protoWorkerMetadata.getWorkedId(), mailboxInfosMap, customProperties); + if (protoWorkerMetadata.hasTableSegmentsMap()) { + workerMetadata.setTableSegmentsMap(fromProtoSegmentsMap(protoWorkerMetadata.getTableSegmentsMap())); + } else if (tableSegmentsJson != null) { + workerMetadata.setTableSegmentsMap(decodeSegmentsMapJson(tableSegmentsJson)); + } + if (protoWorkerMetadata.hasLogicalTableSegmentsMap()) { + workerMetadata.setLogicalTableSegmentsMap( + fromProtoSegmentsMap(protoWorkerMetadata.getLogicalTableSegmentsMap())); + } else if (logicalTableSegmentsJson != null) { + workerMetadata.setLogicalTableSegmentsMap(decodeSegmentsMapJson(logicalTableSegmentsJson)); + } + return workerMetadata; + } + + private static Map<String, List<String>> fromProtoSegmentsMap(Worker.SegmentsMap protoSegmentsMap) { + Map<String, Worker.SegmentList> protoSegments = protoSegmentsMap.getSegmentsMap(); + Map<String, List<String>> segmentsMap = Maps.newHashMapWithExpectedSize(protoSegments.size()); + for (Map.Entry<String, Worker.SegmentList> entry : protoSegments.entrySet()) { + segmentsMap.put(entry.getKey(), new ArrayList<>(entry.getValue().getSegmentList())); + } + return segmentsMap; + } + + private static Map<String, List<String>> decodeSegmentsMapJson(String segmentsMapJson) { + try { + return JsonUtils.stringToObject(segmentsMapJson, SEGMENTS_MAP_TYPE); + } catch (IOException e) { + throw new RuntimeException("Unable to deserialize segments map: " + segmentsMapJson, e); + } } private static MailboxInfos fromProtoMailboxInfos(ByteString protoMailboxInfos) @@ -76,15 +135,60 @@ public class QueryPlanSerDeUtils { return Worker.Properties.parseFrom(protoProperties).getPropertyMap(); } - public static List<Worker.WorkerMetadata> toProtoWorkerMetadataList(List<WorkerMetadata> workerMetadataList) { - return workerMetadataList.stream().map(QueryPlanSerDeUtils::toProtoWorkerMetadata).collect(Collectors.toList()); + /// Encodes the worker metadata for the wire, with the leaf-stage segment maps as native proto fields when + /// `protoSegmentList` is set and as legacy JSON custom properties otherwise (see the class documentation). + public static List<Worker.WorkerMetadata> toProtoWorkerMetadataList(List<WorkerMetadata> workerMetadataList, + boolean protoSegmentList) { + List<Worker.WorkerMetadata> protoWorkerMetadataList = new ArrayList<>(workerMetadataList.size()); + for (WorkerMetadata workerMetadata : workerMetadataList) { + protoWorkerMetadataList.add(toProtoWorkerMetadata(workerMetadata, protoSegmentList)); + } + return protoWorkerMetadataList; + } + + private static Worker.WorkerMetadata toProtoWorkerMetadata(WorkerMetadata workerMetadata, + boolean protoSegmentList) { + Worker.WorkerMetadata.Builder builder = Worker.WorkerMetadata.newBuilder() + .setWorkedId(workerMetadata.getWorkerId()) + .putAllCustomProperty(workerMetadata.getCustomProperties()); + for (Map.Entry<Integer, MailboxInfos> entry : workerMetadata.getMailboxInfosMap().entrySet()) { + builder.putMailboxInfos(entry.getKey(), entry.getValue().toProtoBytes()); + } + Map<String, List<String>> tableSegmentsMap = workerMetadata.getTableSegmentsMap(); + if (tableSegmentsMap != null) { + if (protoSegmentList) { + builder.setTableSegmentsMap(toProtoSegmentsMap(tableSegmentsMap)); + } else { + builder.putCustomProperty(WorkerMetadata.TABLE_SEGMENTS_MAP_KEY, encodeSegmentsMapJson(tableSegmentsMap)); + } + } + Map<String, List<String>> logicalTableSegmentsMap = workerMetadata.getLogicalTableSegmentsMap(); + if (logicalTableSegmentsMap != null) { + if (protoSegmentList) { + builder.setLogicalTableSegmentsMap(toProtoSegmentsMap(logicalTableSegmentsMap)); + } else { + builder.putCustomProperty(WorkerMetadata.LOGICAL_TABLE_SEGMENTS_MAP_KEY, + encodeSegmentsMapJson(logicalTableSegmentsMap)); + } + } + return builder.build(); } - private static Worker.WorkerMetadata toProtoWorkerMetadata(WorkerMetadata workerMetadata) { - Map<Integer, ByteString> mailboxInfosMap = workerMetadata.getMailboxInfosMap().entrySet().stream() - .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().toProtoBytes())); - return Worker.WorkerMetadata.newBuilder().setWorkedId(workerMetadata.getWorkerId()) - .putAllMailboxInfos(mailboxInfosMap).putAllCustomProperty(workerMetadata.getCustomProperties()).build(); + private static Worker.SegmentsMap toProtoSegmentsMap(Map<String, List<String>> segmentsMap) { + Worker.SegmentsMap.Builder builder = Worker.SegmentsMap.newBuilder(); + for (Map.Entry<String, List<String>> entry : segmentsMap.entrySet()) { + builder.putSegments(entry.getKey(), Worker.SegmentList.newBuilder().addAllSegment(entry.getValue()).build()); + } + return builder.build(); + } + + /// JSON-encodes a segments map as `{"OFFLINE":["seg1","seg2"]}` for the legacy encoding. + private static String encodeSegmentsMapJson(Map<String, List<String>> segmentsMap) { + try { + return JsonUtils.objectToString(segmentsMap); + } catch (JsonProcessingException e) { + throw new RuntimeException("Unable to serialize segments map: " + segmentsMap, e); + } } public static Worker.MailboxInfos toProtoMailboxInfos(List<MailboxInfo> mailboxInfos) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerMetadata.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerMetadata.java index 0572a5da52d..40078c03653 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerMetadata.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerMetadata.java @@ -18,14 +18,10 @@ */ package org.apache.pinot.query.routing; -import com.fasterxml.jackson.core.JsonProcessingException; -import com.fasterxml.jackson.core.type.TypeReference; -import java.io.IOException; import java.util.HashMap; import java.util.List; import java.util.Map; import javax.annotation.Nullable; -import org.apache.pinot.spi.utils.JsonUtils; /// `WorkerMetadata` is used to send worker-level info about how to execute a stage on a particular worker. @@ -36,20 +32,28 @@ import org.apache.pinot.spi.utils.JsonUtils; /// - the mailbox info required to construct data transfer linkages. /// - the partition mechanism of the data being execute on this worker. /// +/// The segment maps are held as plain objects: they are only encoded for the wire in [QueryPlanSerDeUtils] when a +/// request is built for the server that runs the worker, so the planner never pays for encoding on the compile path. +/// /// TODO: WorkerMetadata now doesn't have info directly about how to construct the mailboxes. instead it rely on /// MailboxSendNode and MailboxReceiveNode to derive the info during runtime. this should changed to plan time soon. public class WorkerMetadata { + /// Custom-property keys under which brokers that predate the proto segment list encoding ship the segment maps as + /// JSON strings. Still written (when the proto encoding is disabled) and read by [QueryPlanSerDeUtils] so that mixed + /// broker/server versions keep working; never present in [#getCustomProperties()] of a decoded instance. public static final String TABLE_SEGMENTS_MAP_KEY = "tableSegmentsMap"; public static final String LOGICAL_TABLE_SEGMENTS_MAP_KEY = "logicalTableSegmentsMap"; private final int _workerId; private final Map<Integer, MailboxInfos> _mailboxInfosMap; private final Map<String, String> _customProperties; + @Nullable + private Map<String, List<String>> _tableSegmentsMap; + @Nullable + private Map<String, List<String>> _logicalTableSegmentsMap; public WorkerMetadata(int workerId, Map<Integer, MailboxInfos> mailboxInfosMap) { - _workerId = workerId; - _mailboxInfosMap = mailboxInfosMap; - _customProperties = new HashMap<>(); + this(workerId, mailboxInfosMap, new HashMap<>()); } public WorkerMetadata(int workerId, Map<Integer, MailboxInfos> mailboxInfosMap, @@ -71,52 +75,30 @@ public class WorkerMetadata { return _customProperties; } + /// Segments to scan keyed by table type (`OFFLINE` / `REALTIME`), or `null` for a worker that scans no physical + /// table (intermediate stage, or a logical-table leaf). @Nullable public Map<String, List<String>> getTableSegmentsMap() { - return deserializeStringSegmentListMap(TABLE_SEGMENTS_MAP_KEY); - } - - private Map<String, List<String>> deserializeStringSegmentListMap(String propertyKey) { - String tableSegmentsMapStr = _customProperties.get(propertyKey); - if (tableSegmentsMapStr != null) { - try { - return JsonUtils.stringToObject(tableSegmentsMapStr, new TypeReference<Map<String, List<String>>>() { - }); - } catch (IOException e) { - throw new RuntimeException("Unable to deserialize " + propertyKey + " : " + tableSegmentsMapStr, e); - } - } else { - return null; - } - } - - public boolean isLeafStageWorker() { - return _customProperties.containsKey(TABLE_SEGMENTS_MAP_KEY) - || _customProperties.containsKey(LOGICAL_TABLE_SEGMENTS_MAP_KEY); + return _tableSegmentsMap; } public void setTableSegmentsMap(Map<String, List<String>> tableSegmentsMap) { - String tableSegmentsMapStr; - try { - tableSegmentsMapStr = JsonUtils.objectToString(tableSegmentsMap); - } catch (JsonProcessingException e) { - throw new RuntimeException("Unable to serialize table segments map: " + tableSegmentsMap, e); - } - _customProperties.put(TABLE_SEGMENTS_MAP_KEY, tableSegmentsMapStr); + _tableSegmentsMap = tableSegmentsMap; } + /// Segments to scan keyed by physical table name (with type suffix), or `null` for a worker that scans no logical + /// table. @Nullable public Map<String, List<String>> getLogicalTableSegmentsMap() { - return deserializeStringSegmentListMap(LOGICAL_TABLE_SEGMENTS_MAP_KEY); + return _logicalTableSegmentsMap; } public void setLogicalTableSegmentsMap(Map<String, List<String>> logicalTableSegmentsMap) { - String logicalTableSegmentsMapStr; - try { - logicalTableSegmentsMapStr = JsonUtils.objectToString(logicalTableSegmentsMap); - } catch (JsonProcessingException e) { - throw new RuntimeException("Unable to serialize table segments map: " + logicalTableSegmentsMap, e); - } - _customProperties.put(LOGICAL_TABLE_SEGMENTS_MAP_KEY, logicalTableSegmentsMapStr); + _logicalTableSegmentsMap = logicalTableSegmentsMap; + } + + /// A leaf-stage worker carries a (possibly empty) segment map; an intermediate-stage worker carries none. + public boolean isLeafStageWorker() { + return _tableSegmentsMap != null || _logicalTableSegmentsMap != null; } } diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/QueryPlanSerDeUtilsTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/QueryPlanSerDeUtilsTest.java new file mode 100644 index 00000000000..8d5a22ad94c --- /dev/null +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/QueryPlanSerDeUtilsTest.java @@ -0,0 +1,173 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.routing; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.pinot.common.proto.Worker; +import org.apache.pinot.spi.utils.JsonUtils; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +/// Tests the two wire encodings of the leaf-stage segment maps in [QueryPlanSerDeUtils]: the native proto fields and +/// the legacy JSON custom properties, including a server decoding what a pre-proto broker sends. +public class QueryPlanSerDeUtilsTest { + private static final Map<String, String> CUSTOM_PROPERTIES = Map.of("foo", "bar"); + private static final Map<String, List<String>> TABLE_SEGMENTS_MAP = + Map.of("OFFLINE", List.of("seg_0", "seg_1"), "REALTIME", List.of("seg__0__0__20240101T0000Z")); + private static final Map<String, List<String>> LOGICAL_TABLE_SEGMENTS_MAP = + Map.of("t1_OFFLINE", List.of("t1_seg_0"), "t2_REALTIME", List.of("t2_seg_0", "t2_seg_1")); + + @DataProvider + public static Object[][] encodings() { + return new Object[][]{{true}, {false}}; + } + + @Test(dataProvider = "encodings") + public void testLeafWorkerRoundTrip(boolean protoSegmentList) + throws Exception { + WorkerMetadata workerMetadata = leafWorker(TABLE_SEGMENTS_MAP, null); + + Worker.WorkerMetadata proto = toProto(workerMetadata, protoSegmentList); + assertEquals(proto.hasTableSegmentsMap(), protoSegmentList); + assertFalse(proto.hasLogicalTableSegmentsMap()); + assertEquals(proto.getCustomPropertyMap().containsKey(WorkerMetadata.TABLE_SEGMENTS_MAP_KEY), !protoSegmentList); + assertFalse(proto.getCustomPropertyMap().containsKey(WorkerMetadata.LOGICAL_TABLE_SEGMENTS_MAP_KEY)); + assertEquals(proto.getCustomPropertyMap().get("foo"), "bar"); + + WorkerMetadata decoded = QueryPlanSerDeUtils.fromProtoWorkerMetadata(proto); + assertEquals(decoded.getWorkerId(), 3); + assertEquals(decoded.getTableSegmentsMap(), TABLE_SEGMENTS_MAP); + assertNull(decoded.getLogicalTableSegmentsMap()); + assertTrue(decoded.isLeafStageWorker()); + // The JSON is never surfaced as a custom property of the decoded metadata, whichever encoding was used. + assertEquals(decoded.getCustomProperties(), CUSTOM_PROPERTIES); + MailboxInfo mailboxInfo = decoded.getMailboxInfosMap().get(2).getMailboxInfos().get(0); + assertEquals(mailboxInfo.getHostname(), "localhost"); + assertEquals(mailboxInfo.getPort(), 1234); + assertEquals(mailboxInfo.getWorkerIds(), List.of(0, 1)); + } + + @Test(dataProvider = "encodings") + public void testLogicalTableLeafWorkerRoundTrip(boolean protoSegmentList) + throws Exception { + WorkerMetadata workerMetadata = leafWorker(null, LOGICAL_TABLE_SEGMENTS_MAP); + + Worker.WorkerMetadata proto = toProto(workerMetadata, protoSegmentList); + assertFalse(proto.hasTableSegmentsMap()); + assertEquals(proto.hasLogicalTableSegmentsMap(), protoSegmentList); + assertEquals(proto.getCustomPropertyMap().containsKey(WorkerMetadata.LOGICAL_TABLE_SEGMENTS_MAP_KEY), + !protoSegmentList); + + WorkerMetadata decoded = QueryPlanSerDeUtils.fromProtoWorkerMetadata(proto); + assertNull(decoded.getTableSegmentsMap()); + assertEquals(decoded.getLogicalTableSegmentsMap(), LOGICAL_TABLE_SEGMENTS_MAP); + assertTrue(decoded.isLeafStageWorker()); + assertEquals(decoded.getCustomProperties(), CUSTOM_PROPERTIES); + } + + @Test(dataProvider = "encodings") + public void testEmptySegmentListStillMarksLeafWorker(boolean protoSegmentList) + throws Exception { + // A padded worker of a partitioned table scans no segment but must still run the leaf stage. + Map<String, List<String>> emptySegments = Map.of("OFFLINE", new ArrayList<>()); + WorkerMetadata decoded = + QueryPlanSerDeUtils.fromProtoWorkerMetadata(toProto(leafWorker(emptySegments, null), protoSegmentList)); + assertEquals(decoded.getTableSegmentsMap(), emptySegments); + assertTrue(decoded.isLeafStageWorker()); + } + + @Test(dataProvider = "encodings") + public void testIntermediateWorkerRoundTrip(boolean protoSegmentList) + throws Exception { + WorkerMetadata workerMetadata = new WorkerMetadata(1, Map.of(), new HashMap<>(CUSTOM_PROPERTIES)); + + Worker.WorkerMetadata proto = toProto(workerMetadata, protoSegmentList); + assertFalse(proto.hasTableSegmentsMap()); + assertFalse(proto.hasLogicalTableSegmentsMap()); + assertEquals(proto.getCustomPropertyMap(), CUSTOM_PROPERTIES); + + WorkerMetadata decoded = QueryPlanSerDeUtils.fromProtoWorkerMetadata(proto); + assertNull(decoded.getTableSegmentsMap()); + assertNull(decoded.getLogicalTableSegmentsMap()); + assertFalse(decoded.isLeafStageWorker()); + assertEquals(decoded.getCustomProperties(), CUSTOM_PROPERTIES); + } + + /// A broker that predates the proto fields ships Jackson-encoded JSON custom properties; a new server must decode + /// exactly that. + @Test + public void testDecodesLegacyBrokerJsonCustomProperties() + throws Exception { + Worker.WorkerMetadata proto = Worker.WorkerMetadata.newBuilder().setWorkedId(7) + .putCustomProperty(WorkerMetadata.TABLE_SEGMENTS_MAP_KEY, JsonUtils.objectToString(TABLE_SEGMENTS_MAP)) + .putCustomProperty(WorkerMetadata.LOGICAL_TABLE_SEGMENTS_MAP_KEY, + JsonUtils.objectToString(LOGICAL_TABLE_SEGMENTS_MAP)) + .putCustomProperty("foo", "bar") + .build(); + + WorkerMetadata decoded = QueryPlanSerDeUtils.fromProtoWorkerMetadata(proto); + assertEquals(decoded.getWorkerId(), 7); + assertEquals(decoded.getTableSegmentsMap(), TABLE_SEGMENTS_MAP); + assertEquals(decoded.getLogicalTableSegmentsMap(), LOGICAL_TABLE_SEGMENTS_MAP); + assertTrue(decoded.isLeafStageWorker()); + assertEquals(decoded.getCustomProperties(), CUSTOM_PROPERTIES); + } + + /// The legacy encoding must stay readable by a server that predates the proto fields, which parses the custom + /// property with Jackson: pin the exact JSON shape it expects. + @Test + public void testLegacyEncodingIsTheJacksonJsonOlderServersParse() + throws Exception { + Map<String, List<String>> segmentsMap = Map.of("OFFLINE", List.of("seg_0", "seg-1.tar.gz", "s\u00ebg_2")); + Worker.WorkerMetadata proto = toProto(leafWorker(segmentsMap, null), false); + assertEquals(proto.getCustomPropertyMap().get(WorkerMetadata.TABLE_SEGMENTS_MAP_KEY), + "{\"OFFLINE\":[\"seg_0\",\"seg-1.tar.gz\",\"s\u00ebg_2\"]}"); + } + + private static WorkerMetadata leafWorker(@Nullable Map<String, List<String>> tableSegmentsMap, + @Nullable Map<String, List<String>> logicalTableSegmentsMap) { + MailboxInfos mailboxInfos = new MailboxInfos(new MailboxInfo("localhost", 1234, List.of(0, 1))); + WorkerMetadata workerMetadata = new WorkerMetadata(3, Map.of(2, mailboxInfos), new HashMap<>(CUSTOM_PROPERTIES)); + if (tableSegmentsMap != null) { + workerMetadata.setTableSegmentsMap(tableSegmentsMap); + } + if (logicalTableSegmentsMap != null) { + workerMetadata.setLogicalTableSegmentsMap(logicalTableSegmentsMap); + } + return workerMetadata; + } + + /// Serializes through the public list API and parses the bytes back, as the server does. + private static Worker.WorkerMetadata toProto(WorkerMetadata workerMetadata, boolean protoSegmentList) + throws Exception { + Worker.WorkerMetadata proto = + QueryPlanSerDeUtils.toProtoWorkerMetadataList(List.of(workerMetadata), protoSegmentList).get(0); + return Worker.WorkerMetadata.parseFrom(proto.toByteString()); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/ProtoSegmentListPredicate.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/ProtoSegmentListPredicate.java new file mode 100644 index 00000000000..7b7c20d52b1 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/ProtoSegmentListPredicate.java @@ -0,0 +1,109 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.service.dispatch; + +import java.util.Map; +import java.util.Set; +import javax.annotation.concurrent.ThreadSafe; +import org.apache.pinot.spi.config.provider.PinotClusterConfigChangeListener; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// The cluster-level default for encoding leaf-stage segment lists as native protobuf fields of the worker metadata, +/// read by [QueryDispatcher] on every request that does not carry an explicit +/// [CommonConstants.Broker.Request.QueryOptionKey#PROTO_SEGMENT_LIST] override. +/// +/// The value is seeded from the static broker configuration and can then be changed through cluster config, on the +/// same key, without restarting the brokers. That matters because the setting is only safe once every server of the +/// cluster understands the proto fields: an older server finds no segments under them, concludes the worker is not a +/// leaf-stage worker and fails its leaf stage. Operators therefore want to turn it on at the exact moment a rolling +/// upgrade completes, and to turn it back off immediately if it misbehaves, neither of which should cost a broker +/// restart. Cluster config wins over the static seed because [org.apache.pinot.common.config +/// .DefaultClusterConfigChangeHandler] replays the current cluster config to a listener as soon as it is registered; +/// clearing the key from cluster config falls back to [CommonConstants.Broker#DEFAULT_MSE_PROTO_SEGMENT_LIST], not to +/// the static seed. +/// +/// Thread-safety: `_enabled` is `volatile`, so [#isEnabled()] stays lock-free on the request path. [#onChange] is +/// `synchronized` only so that the `previous -> new` pair in its log line cannot interleave with another delivery. It +/// does *not* order deliveries: the change handler invokes listeners outside its own lock, so a delivery computed +/// from an older snapshot can still be applied after a newer one and leave a stale value until the next +/// cluster-config change. +@ThreadSafe +public class ProtoSegmentListPredicate implements PinotClusterConfigChangeListener { + private static final Logger LOGGER = LoggerFactory.getLogger(ProtoSegmentListPredicate.class); + private static final String KEY = CommonConstants.Broker.CONFIG_OF_MSE_PROTO_SEGMENT_LIST; + private static final String ENABLE_WARNING = + "Every server of the cluster must already run a version that understands the proto segment list fields. " + + "Leaf stages routed to an older server will fail. Set it back to false to revert."; + + private volatile boolean _enabled; + + public ProtoSegmentListPredicate(boolean enabled) { + _enabled = enabled; + } + + /// Seeds the value from the static broker configuration. NOTE: the Helix manager is not necessarily connected when + /// this is called, so a cluster-config override is applied later through [#onChange]. + public static ProtoSegmentListPredicate create(PinotConfiguration brokerConf) { + String rawValue = brokerConf.getProperty(KEY); + boolean enabled = rawValue == null || rawValue.isEmpty() + ? CommonConstants.Broker.DEFAULT_MSE_PROTO_SEGMENT_LIST : parseBoolean(rawValue); + LOGGER.info("Initialized {} with value: {}", KEY, enabled); + if (enabled) { + LOGGER.warn("{} is enabled in the static broker config. {}", KEY, ENABLE_WARNING); + } + return new ProtoSegmentListPredicate(enabled); + } + + public boolean isEnabled() { + return _enabled; + } + + @Override + public synchronized void onChange(Set<String> changedConfigs, Map<String, String> clusterConfigs) { + if (!changedConfigs.contains(KEY)) { + return; + } + String value = clusterConfigs.get(KEY); + boolean previous = _enabled; + _enabled = value == null || value.isEmpty() + ? CommonConstants.Broker.DEFAULT_MSE_PROTO_SEGMENT_LIST : parseBoolean(value); + if (previous == _enabled) { + return; + } + LOGGER.info("Updated {} from: {} to: {}", KEY, previous, _enabled); + if (_enabled) { + LOGGER.warn("{} was enabled live via cluster config. {}", KEY, ENABLE_WARNING); + } + } + + /// [Boolean#parseBoolean(String)] semantics (anything but `true` reads as `false`), plus a warning so that a typo + /// does not silently disable the setting. Reading the static seed as a raw string rather than through + /// [PinotConfiguration#getProperty(String, boolean)] is deliberate: that conversion is equally lenient but silent. + private static boolean parseBoolean(String value) { + String trimmed = value.trim(); + if (!trimmed.equalsIgnoreCase("true") && !trimmed.equalsIgnoreCase("false")) { + LOGGER.warn("Unrecognized boolean value '{}' for {}, reading it as false", value, KEY); + } + return Boolean.parseBoolean(trimmed); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java index 834d3b008b8..67e2cafda68 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java @@ -133,26 +133,32 @@ public class QueryDispatcher { /// Cluster-level default for stream-stats mode. Used as the fallback in [#submitAndReduce] when the query /// does not carry an explicit [QueryOptionKey#STREAM_STATS] override. private final boolean _streamStatsDefault; + /// Cluster-level default for the proto encoding of leaf-stage segment lists. Used as the fallback when the query + /// does not carry an explicit [QueryOptionKey#PROTO_SEGMENT_LIST] override. Read per request because operators can + /// change it live through cluster config; see [ProtoSegmentListPredicate]. + private final ProtoSegmentListPredicate _protoSegmentList; public QueryDispatcher(MailboxService mailboxService, FailureDetector failureDetector, @Nullable TlsConfig tlsConfig, boolean enableCancellation, Duration cancelTimeout) { this(mailboxService, failureDetector, tlsConfig, enableCancellation, cancelTimeout, - GrpcKeepAliveConfig.DISABLED, false, CommonConstants.Broker.DEFAULT_STREAM_STATS_DRAIN_MS); + GrpcKeepAliveConfig.DISABLED, false, CommonConstants.Broker.DEFAULT_STREAM_STATS_DRAIN_MS, + new ProtoSegmentListPredicate(CommonConstants.Broker.DEFAULT_MSE_PROTO_SEGMENT_LIST)); } /// Overload that accepts gRPC keep-alive settings for broker dispatch channels. A non-positive `keepAliveTimeMs` /// disables keep-alive. public QueryDispatcher(MailboxService mailboxService, FailureDetector failureDetector, @Nullable TlsConfig tlsConfig, boolean enableCancellation, Duration cancelTimeout, int keepAliveTimeMs, int keepAliveTimeoutMs, - boolean keepAliveWithoutCalls, boolean streamStatsDefault, long statsDrainMs) { + boolean keepAliveWithoutCalls, boolean streamStatsDefault, long statsDrainMs, + ProtoSegmentListPredicate protoSegmentList) { this(mailboxService, failureDetector, tlsConfig, enableCancellation, cancelTimeout, new GrpcKeepAliveConfig(keepAliveTimeMs, keepAliveTimeoutMs, keepAliveWithoutCalls), - streamStatsDefault, statsDrainMs); + streamStatsDefault, statsDrainMs, protoSegmentList); } private QueryDispatcher(MailboxService mailboxService, FailureDetector failureDetector, @Nullable TlsConfig tlsConfig, boolean enableCancellation, Duration cancelTimeout, GrpcKeepAliveConfig keepAliveConfig, - boolean streamStatsDefault, long statsDrainMs) { + boolean streamStatsDefault, long statsDrainMs, ProtoSegmentListPredicate protoSegmentList) { _cancelTimeout = cancelTimeout; _statsDrainMs = statsDrainMs; _mailboxService = mailboxService; @@ -163,6 +169,7 @@ public class QueryDispatcher { _keepAliveConfig = keepAliveConfig; _failureDetector = failureDetector; _streamStatsDefault = streamStatsDefault; + _protoSegmentList = protoSegmentList; if (enableCancellation) { _serversByQuery = new ConcurrentHashMap<>(); @@ -359,8 +366,9 @@ public class QueryDispatcher { // that stage). The streaming observer uses this to drain the session latch correctly when its stream errors // before all opchains have responded. BlockingQueue<AsyncResponse<Worker.QueryResponse>> ackQueue = new ArrayBlockingQueue<>(serversOut.size()); + boolean protoSegmentList = QueryOptionsUtils.isProtoSegmentList(queryOptions, _protoSegmentList.isEnabled()); for (QueryServerInstance server : serversOut) { - Worker.QueryRequest request = createRequest(server, stageInfos, protoRequestMetadata); + Worker.QueryRequest request = createRequest(server, stageInfos, protoRequestMetadata, protoSegmentList); int expectedForServer = 0; for (DispatchablePlanFragment stagePlan : plansWithoutRoot) { List<Integer> workerIds = stagePlan.getServerInstanceToWorkerIdMap().get(server); @@ -633,8 +641,9 @@ public class QueryDispatcher { ByteString protoRequestMetadata = QueryPlanSerDeUtils.toProtoProperties(requestMetadata); // Submit the query plan to all servers in parallel + boolean protoSegmentList = QueryOptionsUtils.isProtoSegmentList(queryOptions, _protoSegmentList.isEnabled()); BlockingQueue<AsyncResponse<E>> dispatchCallbacks = dispatch(sendRequest, serverInstancesOut, deadline, - serverInstance -> createRequest(serverInstance, stageInfos, protoRequestMetadata)); + serverInstance -> createRequest(serverInstance, stageInfos, protoRequestMetadata, protoSegmentList)); processResults(requestId, serverInstancesOut.size(), resultConsumer, deadline, dispatchCallbacks); } @@ -698,8 +707,11 @@ public class QueryDispatcher { } } + /// Builds the request for one server: the plans of the stages it takes part in, with only its own workers' + /// metadata. The leaf-stage segment lists are encoded here, once per worker, rather than at plan time. private static Worker.QueryRequest createRequest(QueryServerInstance serverInstance, - Map<DispatchablePlanFragment, StageInfo> stageInfos, ByteString protoRequestMetadata) { + Map<DispatchablePlanFragment, StageInfo> stageInfos, ByteString protoRequestMetadata, + boolean protoSegmentList) { Worker.QueryRequest.Builder requestBuilder = Worker.QueryRequest.newBuilder(); requestBuilder.setVersion(PlanVersions.V1); @@ -713,7 +725,7 @@ public class QueryDispatcher { workerMetadataList.add(stageWorkerMetadataList.get(workerId)); } List<Worker.WorkerMetadata> protoWorkerMetadataList = - QueryPlanSerDeUtils.toProtoWorkerMetadataList(workerMetadataList); + QueryPlanSerDeUtils.toProtoWorkerMetadataList(workerMetadataList, protoSegmentList); StageInfo stageInfo = entry.getValue(); Worker.StagePlan requestStagePlan = Worker.StagePlan.newBuilder() diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java index e3c44eabc10..de1674b7005 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java @@ -423,8 +423,12 @@ public class QueryServer extends PinotQueryWorkerGrpc.PinotQueryWorkerImplBase { } ByteString rootAsBytes = PlanNodeSerializer.process(explainPlan.getRootNode()).toByteString(); StageMetadata metadata = explainPlan.getStageMetadata(); + // This response travels server -> broker, so it cannot negotiate an encoding with its requester. The + // broker only reads the plan nodes of an explain response today, but the legacy encoding is the one + // every broker can decode, so use it here rather than assume the requester understands the proto + // fields. List<Worker.WorkerMetadata> protoWorkerMetadataList = - QueryPlanSerDeUtils.toProtoWorkerMetadataList(metadata.getWorkerMetadataList()); + QueryPlanSerDeUtils.toProtoWorkerMetadataList(metadata.getWorkerMetadataList(), false); builder.addStagePlan(Worker.StagePlan.newBuilder() .setRootNode(rootAsBytes) .setStageMetadata(Worker.StageMetadata.newBuilder() diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/dispatch/ProtoSegmentListPredicateTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/dispatch/ProtoSegmentListPredicateTest.java new file mode 100644 index 00000000000..8cc780dcbe3 --- /dev/null +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/dispatch/ProtoSegmentListPredicateTest.java @@ -0,0 +1,146 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.service.dispatch; + +import java.util.Map; +import java.util.Set; +import org.apache.helix.model.ClusterConfig; +import org.apache.pinot.common.config.DefaultClusterConfigChangeHandler; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + + +/// Tests for the static-config seed, the live cluster-config update path and the precedence between the two that +/// [ProtoSegmentListPredicate] documents. +public class ProtoSegmentListPredicateTest { + private static final String KEY = CommonConstants.Broker.CONFIG_OF_MSE_PROTO_SEGMENT_LIST; + + @Test + public void testCreateUsesShippedDefaultWhenUnset() { + assertFalse(ProtoSegmentListPredicate.create(new PinotConfiguration()).isEnabled()); + } + + @Test + public void testCreateReadsStaticBrokerConfig() { + assertTrue(ProtoSegmentListPredicate.create(configWith("true")).isEnabled()); + assertFalse(ProtoSegmentListPredicate.create(configWith("false")).isEnabled()); + } + + /// An empty static value means "not set", exactly as an empty cluster-config value does, so the two paths cannot + /// disagree about what a blank entry means. + @Test + public void testCreateTreatsEmptyStaticValueAsUnset() { + assertFalse(ProtoSegmentListPredicate.create(configWith("")).isEnabled()); + } + + /// A typo must never enable the setting, on either path: the encoding is only safe on a fully upgraded cluster. + @Test + public void testUnrecognizedValueReadsAsDisabled() { + assertFalse(ProtoSegmentListPredicate.create(configWith("ture")).isEnabled()); + + ProtoSegmentListPredicate predicate = new ProtoSegmentListPredicate(true); + predicate.onChange(Set.of(KEY), Map.of(KEY, "yes")); + assertFalse(predicate.isEnabled()); + } + + @Test + public void testOnChangeEnablesAndDisables() { + ProtoSegmentListPredicate predicate = new ProtoSegmentListPredicate(false); + predicate.onChange(Set.of(KEY), Map.of(KEY, "true")); + assertTrue(predicate.isEnabled()); + predicate.onChange(Set.of(KEY), Map.of(KEY, "false")); + assertFalse(predicate.isEnabled()); + } + + @Test + public void testOnChangeTrimsAndIsCaseInsensitive() { + ProtoSegmentListPredicate predicate = new ProtoSegmentListPredicate(false); + predicate.onChange(Set.of(KEY), Map.of(KEY, " TRUE ")); + assertTrue(predicate.isEnabled()); + } + + @Test + public void testOnChangeIgnoresChangeThatDoesNotTouchTheKey() { + ProtoSegmentListPredicate predicate = new ProtoSegmentListPredicate(true); + // The key is present in the config snapshot but not in the changed set, so the value is kept. + predicate.onChange(Set.of("some.other.key"), Map.of(KEY, "false", "some.other.key", "x")); + assertTrue(predicate.isEnabled()); + } + + /// Clearing the key from cluster config falls back to the shipped default rather than to the static seed, which is + /// the safe direction for a setting that is only valid on a fully upgraded cluster. + @Test + public void testOnChangeResetsToDefaultWhenValueRemovedOrEmpty() { + ProtoSegmentListPredicate predicate = new ProtoSegmentListPredicate(true); + predicate.onChange(Set.of(KEY), Map.of()); + assertFalse(predicate.isEnabled()); + + predicate = new ProtoSegmentListPredicate(true); + predicate.onChange(Set.of(KEY), Map.of(KEY, "")); + assertFalse(predicate.isEnabled()); + } + + /// The documented precedence "cluster config beats the static seed" holds only because the change handler replays + /// the current snapshot to a listener as it is registered. Pinned here against the real handler rather than left to + /// the javadoc, since that replay is what lets an operator flip the encoding without restarting the brokers. + @Test + public void testRegistrationReplayLetsClusterConfigWinOverStaticSeed() { + DefaultClusterConfigChangeHandler handler = new DefaultClusterConfigChangeHandler(); + handler.onClusterConfigChange(clusterConfig(Map.of(KEY, "true")), null); + + ProtoSegmentListPredicate predicate = ProtoSegmentListPredicate.create(configWith("false")); + assertFalse(predicate.isEnabled(), "static seed applies before registration"); + + assertTrue(handler.registerClusterConfigChangeListener(predicate)); + assertTrue(predicate.isEnabled(), "cluster config must win over the static seed"); + + // Registration really did wire the listener up, so a later change still reaches it: this is the live disable + // path an operator relies on to revert without a broker restart. + handler.onClusterConfigChange(clusterConfig(Map.of(KEY, "false")), null); + assertFalse(predicate.isEnabled()); + } + + /// The other half of the same contract: a replayed snapshot that does not carry the key must leave the static seed + /// alone. + @Test + public void testRegistrationReplayPreservesStaticSeed() { + DefaultClusterConfigChangeHandler handler = new DefaultClusterConfigChangeHandler(); + handler.onClusterConfigChange(clusterConfig(Map.of("some.other.key", "x")), null); + + ProtoSegmentListPredicate predicate = ProtoSegmentListPredicate.create(configWith("true")); + assertTrue(handler.registerClusterConfigChangeListener(predicate)); + assertTrue(predicate.isEnabled(), "an unrelated cluster config must not clear the seed"); + } + + private static ClusterConfig clusterConfig(Map<String, String> configs) { + ClusterConfig clusterConfig = new ClusterConfig("testCluster"); + configs.forEach((key, value) -> clusterConfig.getRecord().setSimpleField(key, value)); + return clusterConfig; + } + + private static PinotConfiguration configWith(String value) { + PinotConfiguration config = new PinotConfiguration(); + config.setProperty(KEY, value); + return config; + } +} diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerAuthzTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerAuthzTest.java index 101e0d91a0a..3271a548f9b 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerAuthzTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerAuthzTest.java @@ -155,7 +155,7 @@ public class QueryServerAuthzTest { DispatchablePlanFragment stagePlan = queryPlan.getQueryStageMap().get(stageId); Plan.PlanNode rootNode = PlanNodeSerializer.process(stagePlan.getPlanFragment().getFragmentRoot()); List<Worker.WorkerMetadata> workerMetadataList = - QueryPlanSerDeUtils.toProtoWorkerMetadataList(stagePlan.getWorkerMetadataList()); + QueryPlanSerDeUtils.toProtoWorkerMetadataList(stagePlan.getWorkerMetadataList(), false); ByteString customProperty = QueryPlanSerDeUtils.toProtoProperties(stagePlan.getCustomProperties()); // this particular test set requires the request to have a single QueryServerInstance to dispatch to diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerTest.java index fd8020942f6..c8c177fd51c 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerTest.java @@ -235,13 +235,26 @@ public class QueryServerTest extends QueryTestSet { @Test(dataProvider = "testSql") public void testWorkerAcceptsWorkerRequestCorrect(String sql) throws Exception { + testWorkerAcceptsWorkerRequestCorrect(sql, false); + } + + /// Same as [#testWorkerAcceptsWorkerRequestCorrect(String)] with the leaf-stage segment lists shipped as native + /// proto fields instead of the legacy JSON custom property. + @Test(dataProvider = "testSql") + public void testWorkerAcceptsProtoSegmentListRequestCorrect(String sql) + throws Exception { + testWorkerAcceptsWorkerRequestCorrect(sql, true); + } + + private void testWorkerAcceptsWorkerRequestCorrect(String sql, boolean protoSegmentList) + throws Exception { DispatchableSubPlan queryPlan = _queryEnvironment.planQuery(sql); Set<DispatchablePlanFragment> stagePlans = queryPlan.getQueryStagesWithoutRoot(); // Ignore reduce stage (stage 0) for (DispatchablePlanFragment stagePlan : stagePlans) { int stageId = stagePlan.getPlanFragment().getFragmentId(); // only get one worker request out. - Worker.QueryRequest queryRequest = getQueryRequest(queryPlan, stageId); + Worker.QueryRequest queryRequest = getQueryRequest(queryPlan, stageId, protoSegmentList); Map<String, String> requestMetadata = QueryPlanSerDeUtils.fromProtoProperties(queryRequest.getMetadata()); // submit the request for testing. @@ -321,10 +334,14 @@ public class QueryServerTest extends QueryTestSet { } private Worker.QueryRequest getQueryRequest(DispatchableSubPlan queryPlan, int stageId) { + return getQueryRequest(queryPlan, stageId, false); + } + + private Worker.QueryRequest getQueryRequest(DispatchableSubPlan queryPlan, int stageId, boolean protoSegmentList) { DispatchablePlanFragment stagePlan = queryPlan.getQueryStageMap().get(stageId); Plan.PlanNode rootNode = PlanNodeSerializer.process(stagePlan.getPlanFragment().getFragmentRoot()); List<Worker.WorkerMetadata> workerMetadataList = - QueryPlanSerDeUtils.toProtoWorkerMetadataList(stagePlan.getWorkerMetadataList()); + QueryPlanSerDeUtils.toProtoWorkerMetadataList(stagePlan.getWorkerMetadataList(), protoSegmentList); ByteString customProperty = QueryPlanSerDeUtils.toProtoProperties(stagePlan.getCustomProperties()); // this particular test set requires the request to have a single QueryServerInstance to dispatch to diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index bb29493d497..0e915aa7fc7 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -636,6 +636,24 @@ public class CommonConstants { public static final String CONFIG_OF_STREAM_STATS_DRAIN_MS = "pinot.broker.mse.stream.stats.drain.ms"; public static final long DEFAULT_STREAM_STATS_DRAIN_MS = 50L; + /// Cluster-level default for shipping the leaf-stage segment lists of a multi-stage query as native protobuf + /// fields of the worker metadata instead of a JSON string custom property. The proto encoding skips a JSON encode + /// per leaf-stage worker on the broker and a JSON parse per worker on the server. Individual queries may override + /// this default via the [Request.QueryOptionKey#PROTO_SEGMENT_LIST] query option. Requires all servers to + /// understand the proto encoding; enabling it on a mixed-version cluster with older servers fails their leaf + /// stages, so only enable it once the whole fleet has been upgraded. Brokers also watch this key in cluster + /// config, so it can be turned on right after a rolling upgrade, and back off again, without restarting them. + public static final String CONFIG_OF_MSE_PROTO_SEGMENT_LIST = "pinot.broker.mse.proto.segment.list"; + public static final boolean DEFAULT_MSE_PROTO_SEGMENT_LIST = false; + + /// Number of threads of the executor that parses, compiles and plans multi-stage queries. Planning is CPU bound + /// and this executor has a fixed size, so it can saturate (queueing every query behind it) well before the broker + /// runs out of CPU on a workload of many short multi-stage queries. A non-positive value (the default) sizes it to + /// half the available processors, with a minimum of 1. + public static final String CONFIG_OF_MSE_QUERY_COMPILE_EXECUTOR_THREADS = + "pinot.broker.mse.query.compile.executor.threads"; + public static final int DEFAULT_MSE_QUERY_COMPILE_EXECUTOR_THREADS = -1; + public static final String CONFIG_OF_USE_FIXED_REPLICA = "pinot.broker.use.fixed.replica"; public static final boolean DEFAULT_USE_FIXED_REPLICA = false; @@ -899,6 +917,14 @@ public class CommonConstants { /// other transport error) during dispatch, the broker cancels the query and surfaces the error to the /// client. public static final String STREAM_STATS = "streamStats"; + /// When set to true, the broker ships the leaf-stage segment lists as native protobuf fields of the worker + /// metadata instead of the legacy JSON string custom property, skipping a JSON encode per leaf-stage worker on + /// the broker and a JSON parse per worker on every server. + /// + /// **Mixed-version note.** All servers must understand the proto encoding when this option is enabled: an + /// older server finds no segments to scan and fails the leaf stage. Enable it (per query, or cluster-wide via + /// `pinot.broker.mse.proto.segment.list`) only once the whole fleet has been upgraded. + public static final String PROTO_SEGMENT_LIST = "protoSegmentList"; /// If set, changes the explain behavior in multi-stage engine. /// /// `true` means to ask servers for the physical plan while false means to just use logical plan. --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
