Jackie-Jiang commented on code in PR #20106:
URL: https://github.com/apache/pinot/pull/20106#discussion_r4519237531
##########
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:
Fixed in a609309182 (included in current head 33ecaff224). The helper now
uses `MetricValueUtils.getPartitionGaugeValue(...)` exactly as suggested, and
the manual registry lookup/cast is removed. The readiness, offset,
document-count, and final salary assertions are unchanged. Compilation and
formatting/style/license checks passed; no tests were run locally for this
update, and CI is running. Your passing run and gate-disabled negative control
provide useful confirmation that the regression exercises the intended behavior.
--
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]
