xiangfu0 commented on code in PR #19568: URL: https://github.com/apache/pinot/pull/19568#discussion_r4067139557
########## 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) { Review Comment: Adopted — thanks, this is clearly the better shape, and to answer the question directly: no, it was not considered. The first version was built around an operator flipping a live flag, and auto-detection never came up. `pinot.broker.mse.proto.segment.list` is now a `NEVER` / `SAFE` / `ALWAYS` mode with `SAFE` as the default, and the cluster-config listener, its registration comment and `testRegistrationBeforeFirstDeliveryStillLetsClusterConfigWin` are all gone. The per-query option stays as the override. The runbook in the description has shrunk to "nothing to do, here is the log line". I took both of your weaknesses as written: exact equality against `PinotVersion.VERSION`, and a missing or `UNKNOWN` version reads as outdated. Beyond that I diverged from `SendStatsPredicate` in four places, each on purpose: 1. **Only servers are checked.** Brokers never decode the fields, so a broker on another version is no reason to stay on JSON. 2. **An unreadable instance config counts as outdated.** `SendStatsPredicate` treats it as non-problematic, which is fine when the downside is missing stats. Here the downside is a failed query, so it fails closed. 3. **It starts disabled.** `Safe` starts with `_sendStats = true`. This one stays on the legacy encoding until the first delivery proves every server is current. 4. **Multi-cluster queries always use the legacy encoding in `SAFE`.** `enableMultiClusterRouting` sends leaf stages to `RemoteClusterBrokerRoutingManager` servers, and a watcher on the local cluster cannot see their versions. The old runbook had the same hole, since "every server in the cluster" silently excluded remote clusters. `ALWAYS` still forces proto. On registration I did not go through `_instanceConfigChangeHandlers`, for two reasons: - `ClusterChangeHandler` lives in `pinot-broker`, while the predicate sits next to the dispatcher in `pinot-query-runtime`. - `processClusterChange(changeType)` carries no changed path. Every instance-config change would re-read every instance config, on top of the routing manager's own full read. Instead it is a Helix `InstanceConfigChangeListener` on the spectator manager, with `@PreFetch(enabled = false)`, exactly like `SendStatsPredicate` on the server. A data change reads only the changed instance, and a child change or INIT does the full read. `BaseBrokerStarter` gains one line, `watchInstanceConfigs(_spectatorHelixManager)`, which is a no-op outside `SAFE`. A registration failure is logged rather than thrown, so it leaves `SAFE` on the legacy encoding instead of failing broker startup. One operational gotcha worth writing down: an instance config left behind by a decommissioned old server keeps `SAFE` on JSON until it is removed. The transition log names the servers holding it back (capped at 10), so this is greppable. ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java: ########## @@ -633,8 +641,9 @@ private <E> void execute(long requestId, Set<DispatchablePlanFragment> stagePlan 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)); Review Comment: I worked through the timing and don't think build-first helps here, so I've left the loop as it is. Happy to be corrected if I'm missing something about `send`. `send` is an async gRPC call, so treat its cost `s_i` as small next to the encode `e_i`. The start time of server `k` (1-indexed in iteration order) under each ordering: | | server `k` starts at | last server starts at | |---|---|---| | master: encodes done in the planner, right before dispatch | `Σ_all e + Σ_{i≤k} s_i` | `Σ e + Σ s` | | this PR: interleaved | `Σ_{i≤k} (e_i + s_i)` | `Σ e + Σ s` | | build-first | `Σ_all e + Σ_{i≤k} s_i` | `Σ e + Σ s` | So the last server, which gates the query, starts at the same moment in all three. The encodes were already on the critical path on master; they just ran inside `callAsync` instead of inside the loop. What interleaving changes is that every other server starts *earlier* than on master. Build-first gives that back and reproduces master's schedule exactly. The skew you're seeing is early servers starting early, not the last one starting late. Leaf stages that start sooner also feed their receivers sooner, so it is never worse. Actually shortening the critical path means building the per-server requests in parallel. That is a real idea but a separate change, and as you say it needs a wide fan-out against a large table to justify it. With `SAFE` now the default, homogeneous clusters also get the proto encode, which is roughly half the JSON cost. `submitWithStream` is the same analysis. I replaced the benchmark caveat in the description with this reasoning, so it no longer implies a win only on the opt-in path. ########## pinot-query-planner/src/main/java/org/apache/pinot/query/routing/QueryPlanSerDeUtils.java: ########## @@ -54,15 +73,55 @@ private static StageMetadata fromProtoStageMetadata(Worker.StageMetadata protoSt 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()) { Review Comment: Agreed on the direction. On the WARN, though, I think the premise doesn't hold for this case, so I haven't added it. An old server that receives the proto encoding does not silently scan nothing. `isLeafStageWorker()` is false, so the stage goes through `PlanNodeToOpChain`, where `visitTableScan` returns an `ErrorOperator` with `QUERY_EXECUTION`: "Plan node of type TableScanNode is not supported in OpChain execution." That has been the behavior since #16257 (1.5.0); before that it threw `UnsupportedOperationException("Plan node of type TableScanNode is not supported!")`. Either way the query fails with an error, and the error names the node type. A WARN would also live in new server code, and the dangerous scenario is precisely the one where the server runs old code. It cannot gain a log line retroactively. On a new server, a `TableScanNode` stage with no segment map can only come from a planner bug, and that fails through the same `ErrorOperator`. The likelier form of that bug, a worker missing from its stage's segment map, is now caught at plan time with the worker id by the new check in `DispatchablePlanContext`. The part that matters, making the misconfiguration unreachable, is what `SAFE` now does. The one remaining way in is `ALWAYS` or `SET protoSegmentList = true` against an old server, and that fails loudly as above. I corrected the description, which said the same thing ("fails the leaf stage") but now names the operator. ########## pinot-query-planner/src/main/java/org/apache/pinot/query/routing/QueryPlanSerDeUtils.java: ########## @@ -54,15 +73,55 @@ private static StageMetadata fromProtoStageMetadata(Worker.StageMetadata protoSt 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); Review Comment: Good catch, fixed. I went with always-unmodifiable rather than always-copy. The proto map view already is unmodifiable, and on master the server always received that view, so this restores the invariant without adding a copy to the proto path. The legacy path strips the JSON keys from its copy and wraps it in `Collections.unmodifiableMap`. `QueryPlanSerDeUtilsTest#testDecodedCustomPropertiesAreUnmodifiable` pins it for leaf and intermediate workers under both encodings. ########## 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<>()); Review Comment: Right, that test only covered a non-empty map holding an empty list. I added `testZeroEntrySegmentsMapStillMarksLeafWorker` for `Map.of()` under both encodings. It goes through real bytes (`toProto` round-trips via `parseFrom`), asserts `hasTableSegmentsMap()` / `hasLogicalTableSegmentsMap()` on the parsed message in the proto case, and checks that the decoded worker is still a leaf. The legacy path encodes to `{}` and decodes to an empty, non-null map, so it lands in the same place. I left the existing test as it is: `Map.of("OFFLINE", List.of())` is exactly what `PlanFragmentAndMailboxAssignment` builds for dimension-table leaves, so it is worth keeping pinned too. ########## pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerTest.java: ########## @@ -235,13 +235,26 @@ public void onCompleted() { @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) Review Comment: Added both, and making `SAFE` the default gave the integration test something more useful to assert than result equality alone: - `MultiStageEngineIntegrationTest#testProtoSegmentListEncodingIsTransparent` first asserts that the broker's predicate is `SAFE` and turns itself on from the versions the real cluster publishes, then that it stays off for multi-cluster queries. That is the half no unit test can reach. It then runs three queries with `SET protoSegmentList = false` and `= true` and compares the rows: a single-leaf count, a filtered group-by, and a join with two leaf stages. - `BaseLogicalTableIntegrationTest#testProtoSegmentListPreservesLogicalTableResults` does the same comparison for logical tables (`logicalTableSegmentsMap`). Because it sits in the base class, it runs in all eight offline / realtime / hybrid subclasses. A side effect worth knowing about: every instance of a test cluster reports the same version, so with `SAFE` as the default every existing MSE integration test now runs the proto encoding. The explicit `false` is what keeps the legacy path covered. ########## pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java: ########## @@ -198,11 +198,19 @@ public Map<Integer, DispatchablePlanFragment> constructDispatchablePlanFragmentM 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); Review Comment: I verified your read of the populating sites: `WorkerManager` at 447 / 666 / 954 / 1344 / 1431 and `PlanFragmentAndMailboxAssignment` all fill the segment map in the same per-worker loop as the server assignment. 1431 even enforces it through `checkLeafWorkerAssignment`. That reasoning is now in the comment at the check, so it outlives this thread, and the description calls the change out under its own heading with the `"null"`-string history. I have kept it in this PR rather than splitting it, for two reasons: - Pinot squash-merges, so a separate commit inside the PR would not survive to bisect. It would have to be a separate PR. - A separate PR would land the guard against master's representation, where a missing entry means something different. It is this PR's presence-means-leaf semantics that make the guard necessary: without it, a missing entry would quietly demote the worker to an intermediate one and fail on the server with a `TableScanNode` error, far from the cause. If you still prefer the split I'll do it, but I think it reads better as part of the change that needs it. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
