xiangfu0 commented on code in PR #20106:
URL: https://github.com/apache/pinot/pull/20106#discussion_r4500658003
##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java:
##########
@@ -225,6 +247,210 @@ public void testRebalanceWithBatching()
"Failed to drop servers");
}
+ @Test
+ public void testRebalanceWaitsForOlderSegmentBeforeIncrement()
+ throws Exception {
+ String tableName = "partialUpsertReadiness";
+ String tableNameWithType =
TableNameBuilder.REALTIME.tableNameWithType(tableName);
+ String topic = getKafkaTopic() + "-readiness";
+ String tenant = "readinessTenant";
+ List<BaseServerStarter> servers = new ArrayList<>();
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ Future<RebalanceResult> rebalance = null;
+ DelayedTableUpsertMetadataManager delayedManager = null;
+ boolean tableAdded = false;
+ boolean schemaAdded = false;
+ boolean topicCreated = false;
+ try {
+ BaseServerStarter source = startOneServer(NUM_SERVERS);
+ servers.add(source);
+ _resourceManager.updateInstanceTags(source.getInstanceId(), tenant +
"_REALTIME", false);
+ PinotConfiguration destinationConf = getServerConf(NUM_SERVERS + 1);
+
destinationConf.setProperty(CommonConstants.Server.INSTANCE_DATA_MANAGER_CONFIG_PREFIX
+ "."
+ + CommonConstants.Server.Upsert.CONFIG_PREFIX + "."
+ + CommonConstants.Server.Upsert.DEFAULT_METADATA_MANAGER_CLASS,
+ DelayedTableUpsertMetadataManager.class.getName());
+ BaseServerStarter destination = startOneServer(destinationConf);
+ servers.add(destination);
+ _resourceManager.updateInstanceTags(destination.getInstanceId(), tenant
+ "_REALTIME", false);
+
+ createKafkaTopic(topic);
+ topicCreated = true;
+ Schema schema = new Schema.SchemaBuilder().setSchemaName(tableName)
+ .addSingleValueDimension(PRIMARY_KEY_COL, DataType.INT)
+ .addMetric("salary", DataType.INT)
+ .addDateTime(getTimeColumnName(), DataType.LONG,
"1:MILLISECONDS:EPOCH", "1:MILLISECONDS")
+ .setPrimaryKeyColumns(List.of(PRIMARY_KEY_COL)).build();
+ addSchema(schema);
+ schemaAdded = true;
+ UpsertConfig upsert = new UpsertConfig(UpsertConfig.Mode.PARTIAL);
+ upsert.setPartialUpsertStrategies(Map.of("salary",
UpsertConfig.Strategy.INCREMENT));
+ TableConfig tableConfig = createCSVUpsertTableConfig(tableName, topic, 2,
+ getCSVDecoderProperties(",", PRIMARY_KEY_COL + ",salary," +
getTimeColumnName()), upsert, PRIMARY_KEY_COL);
+ tableConfig.setTenantConfig(new TenantConfig(getBrokerTenant(), tenant,
null));
+ addTableConfig(tableConfig);
+ tableAdded = true;
+ waitForClusterStateSettled(tableNameWithType, 2, 1, 60_000L);
+
+ Map<String, Map<String, String>> initialAssignment =
+ HelixHelper.getTableIdealState(_helixManager,
tableNameWithType).getRecord().getMapFields();
+ assertEquals(initialAssignment.values().stream()
+ .filter(instances ->
instances.containsKey(source.getInstanceId())).count(), 1L);
+ assertEquals(initialAssignment.values().stream()
+ .filter(instances ->
instances.containsKey(destination.getInstanceId())).count(), 1L);
+ int partition = initialAssignment.entrySet().stream()
+ .filter(entry ->
entry.getValue().containsKey(source.getInstanceId()))
+ .mapToInt(entry -> new
LLCSegmentName(entry.getKey()).getPartitionGroupId()).findFirst().orElseThrow();
+
+ // Warm both servers on this same table, so the old table-wide readiness
latch has already passed.
+ // S1 owns the base keys; S2 contains disjoint keys, leaving the target
key's history only in S1.
+ MurmurPartitionFunction partitionFunction = new
MurmurPartitionFunction(2, null);
+ int[][] keys = new int[2][];
+ for (int p = 0; p < 2; p++) {
+ int partitionId = p;
+ keys[p] = IntStream.range(0, 100)
+ .filter(key ->
partitionFunction.getPartition(Integer.toString(key)) ==
partitionId).limit(2).toArray();
+ assertEquals(keys[p].length, 2);
+ }
+ long timestamp = System.currentTimeMillis();
+ pushCsvIntoKafka(List.of(keys[0][0] + ",100," + timestamp, keys[1][0] +
",100," + timestamp), topic, 0);
+ waitForSalarySum(tableName, 200);
+ getOrCreateAdminClient().getTableClient().forceCommit(tableName);
+ waitForClusterStateSettled(tableNameWithType, 4, 1, 60_000L);
+ pushCsvIntoKafka(List.of(keys[0][1] + ",1," + (timestamp + 1),
keys[1][1] + ",1," + (timestamp + 1)), topic, 0);
+ waitForSalarySum(tableName, 202);
+ getOrCreateAdminClient().getTableClient().forceCommit(tableName);
+ waitForClusterStateSettled(tableNameWithType, 6, 1, 60_000L);
+
+ Map<String, Map<String, String>> assignment =
+ HelixHelper.getTableIdealState(_helixManager,
tableNameWithType).getRecord().getMapFields();
+ String baseSegment = findReadinessSegment(assignment, partition, 0);
+ String predecessor = findReadinessSegment(assignment, partition, 1);
+ String consuming = findReadinessSegment(assignment, partition, 2);
+ RealtimeTableDataManager destinationTable = (RealtimeTableDataManager)
destination.getServerInstance()
+ .getInstanceDataManager().getTableDataManager(tableNameWithType);
+ delayedManager = (DelayedTableUpsertMetadataManager)
destinationTable.getTableUpsertMetadataManager();
+ delayedManager.delaySegment(baseSegment);
+
+ // Move the source's partition onto the already-warm destination. Run
the synchronous rebalancer separately
+ // so this thread can hold S1 until S2 has loaded and S3 has reached its
readiness wait.
+ _resourceManager.updateInstanceTags(source.getInstanceId(), "", false);
+ RebalanceConfig config = new RebalanceConfig();
+ config.setDryRun(false);
+ config.setReassignInstances(true);
+ config.setIncludeConsuming(true);
+ config.setForceCommit(false);
+ config.setMinAvailableReplicas(0);
+ rebalance = executor.submit(() ->
_tableRebalancer.rebalance(tableConfig, config, null));
+ assertTrue(delayedManager.awaitDelayedLoad(60, TimeUnit.SECONDS), "S1
never reached the destination");
+ TestUtils.waitForCondition(aVoid -> {
+ SegmentDataManager manager =
destinationTable.getSegmentDataManager(predecessor);
+ return manager instanceof ImmutableSegmentDataManager
+ && ((ImmutableSegmentDataManager)
manager).isPartitionMetadataInitialized()
+ && destinationTable.getSegmentDataManager(consuming) instanceof
RealtimeSegmentDataManager
+ && isWaitingForReadiness(tableNameWithType, partition);
+ }, 100L, 60_000L, "S3 did not wait for S1 after its immediate
predecessor S2 was loaded");
+ RealtimeSegmentDataManager consumer =
+ (RealtimeSegmentDataManager)
destinationTable.getSegmentDataManager(consuming);
+ String initialOffset = consumer.getCurrentOffset().toString();
+ assertEquals(consumer.getSegment().getNumDocsIndexed(), 0);
+
+ pushCsvIntoKafka(List.of(keys[partition][0] + ",7," + (timestamp + 2),
+ keys[partition][0] + ",11," + (timestamp + 3)), topic, 0);
+ assertTrue(isWaitingForReadiness(tableNameWithType, partition));
+ assertEquals(consumer.getCurrentOffset().toString(), initialOffset);
+ assertEquals(consumer.getSegment().getNumDocsIndexed(), 0);
+ assertFalse(rebalance.isDone(), "Rebalance completed while S1 metadata
was still blocked");
+
+ delayedManager.releaseDelayedLoad();
+ assertEquals(rebalance.get(60, TimeUnit.SECONDS).getStatus(),
RebalanceResult.Status.DONE);
+ waitForClusterStateSettled(tableNameWithType, 6, 1, 60_000L);
+ waitForSalarySum(tableName, 220);
+ assertEquals(getPinotConnection().execute("SELECT salary FROM " +
tableName + " WHERE " + PRIMARY_KEY_COL
+ + " = " + keys[partition][0]).getResultSet(0).getInt(0), 118);
+ } finally {
+ if (delayedManager != null) {
+ delayedManager.releaseDelayedLoad();
+ }
+ if (rebalance != null) {
+ rebalance.cancel(true);
+ }
+ executor.shutdownNow();
+ try {
+ assertTrue(executor.awaitTermination(60, TimeUnit.SECONDS), "Rebalance
executor did not stop");
+ } finally {
+ try {
+ if (tableAdded) {
+ dropRealtimeTable(tableName);
+ waitForEVToDisappear(tableNameWithType);
+ }
+ } finally {
+ try {
+ stopReadinessServers(servers);
+ } finally {
+ SegmentBuildTimeLeaseExtender.initExecutor();
+ try {
+ if (schemaAdded) {
+ deleteSchema(tableName);
+ }
+ } finally {
+ if (topicCreated) {
+ deleteKafkaTopic(topic);
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
+ private void stopReadinessServers(List<BaseServerStarter> servers) {
+ AssertionError failure = new AssertionError("Failed to clean up readiness
test servers");
+ for (BaseServerStarter server : servers) {
+ try {
+ server.stop();
+ } catch (Exception | AssertionError e) {
+ failure.addSuppressed(e);
+ }
+ }
+ for (BaseServerStarter server : servers) {
+ try {
+ TestUtils.waitForCondition(aVoid ->
_resourceManager.dropInstance(server.getInstanceId()).isSuccessful(),
+ 60_000L, "Failed to drop readiness test server");
+ } catch (Exception | AssertionError e) {
+ failure.addSuppressed(e);
+ }
+ }
+ if (failure.getSuppressed().length > 0) {
+ throw failure;
+ }
+ }
+
+ private static String findReadinessSegment(Map<String, Map<String, String>>
assignment, int partition, int sequence) {
+ return assignment.keySet().stream().map(LLCSegmentName::new)
+ .filter(name -> name.getPartitionGroupId() == partition &&
name.getSequenceNumber() == sequence)
+ .map(LLCSegmentName::getSegmentName).findFirst().orElseThrow();
+ }
+
+ private void waitForSalarySum(String tableName, int expected) {
+ TestUtils.waitForCondition(aVoid -> {
+ try {
+ return getPinotConnection().execute("SELECT SUM(salary) FROM " +
tableName)
+ .getResultSet(0).getDouble(0) == expected;
+ } catch (Exception e) {
+ return false;
+ }
+ }, 100L, 60_000L, "Unexpected salary total for " + tableName + ": expected
" + expected);
+ }
+
+ private static boolean isWaitingForReadiness(String tableNameWithType, int
partition) {
+ String metric = CommonConstants.Server.DEFAULT_METRICS_PREFIX
+ + ServerGauge.CONSUMERS_WAITING_FOR_READINESS.getGaugeName() + "." +
tableNameWithType + "." + partition;
+ PinotMetric gauge = ServerMetrics.get().getMetricsRegistry().allMetrics()
+ .get(PinotMetricUtils.makePinotMetricName(ServerMetrics.class,
metric));
+ return gauge instanceof PinotGauge && ((Number) ((PinotGauge<?>)
gauge).getValue()).longValue() > 0;
Review Comment:
**MAJOR: this check is always false, so the new test can never pass.**
`YammerMetricsRegistry.allMetrics()` wraps every entry as a `YammerMetric`,
which is not a `PinotGauge`, so `instanceof PinotGauge` never matches. Locally,
597 polls saw S2 initialized, S3 registered and the gauge present, but this
returned `false` each time. Reading it with
`MetricValueUtils.getPartitionGaugeValue(ServerMetrics.get(),
tableNameWithType, partition, ServerGauge.CONSUMERS_WAITING_FOR_READINESS) > 0`
makes the test pass. With the history gate disabled, the fixed test fails as it
should.
--
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]
