From ecfddec515e2ffa255cd57da641914107201d4ae Mon Sep 17 00:00:00 2001 From: mittalprince Date: Tue, 18 Aug 2026 13:01:58 +0530 Subject: [PATCH] Refresh assigned task's datastream list when a dedup datastream joins an existing task When a new datastream dedupes into the task(s) of an already-established datastream (same effective TASK_PREFIX), the assignment (task names) does not change, so the normal assignment diff/rebalance path never touches the affected task's ZK node and the task's cached datastreams list on every instance goes stale. Narrow the fix to only broadcast a datastream update (which triggers every instance to refresh its locally-assigned tasks' datastreams lists) when a newly-READY datastream actually joins an existing task group, rather than on every datastream creation. This avoids extra onAssignmentChange() churn for the common case of a brand-new datastream getting a brand-new task. Also updates TestCoordinator to reflect that a task's representative datastream can now correctly be any datastream in its dedup group. Addresses review feedback from kanishkjaiswal2015. --- .../datastream/server/TestCoordinator.java | 87 +++++++++++++++++-- .../datastream/server/Coordinator.java | 35 ++++++++ 2 files changed, 115 insertions(+), 7 deletions(-) diff --git a/datastream-server-restli/src/test/java/com/linkedin/datastream/server/TestCoordinator.java b/datastream-server-restli/src/test/java/com/linkedin/datastream/server/TestCoordinator.java index ba61f735e..4b7c50f27 100644 --- a/datastream-server-restli/src/test/java/com/linkedin/datastream/server/TestCoordinator.java +++ b/datastream-server-restli/src/test/java/com/linkedin/datastream/server/TestCoordinator.java @@ -482,9 +482,11 @@ public void testCoordinationWithBroadcastStrategy() throws Exception { ds1.setStatus(DatastreamStatus.PAUSED); DatastreamTestUtils.updateDatastreams(zkClient, testCluster, ds1); - // Check that the datastream are running. - assertConnectorAssignment(connector1, WAIT_TIMEOUT_MS, "datastream1", "datastream2"); - assertConnectorAssignment(connector2, WAIT_TIMEOUT_MS, "datastream1", "datastream2"); + // Check that the datastream are running. Datastream3 is deduped into datastream1's task, and with the task's + // datastreams list now correctly refreshed to reflect the dedup, datastream3 is an equally valid + // representative of that task alongside datastream1. + assertConnectorAssignment(connector1, WAIT_TIMEOUT_MS, 2, "datastream1", "datastream2", "datastream3"); + assertConnectorAssignment(connector2, WAIT_TIMEOUT_MS, 2, "datastream1", "datastream2", "datastream3"); // Verify that No Tasks are parked (because DS1 and DS3 are in the same group, and DS3 is not paused) Assert.assertEquals(zkClient.getChildren(pausedPath).size(), 0); @@ -3209,6 +3211,63 @@ public void testDatastreamDedupeMetadataCopy() throws Exception { zkClient.close(); } + /** + * Test that when a "dedup" datastream is created that reuses the task(s) of an already-assigned + * datastream (same taskPrefix/partitions, so the assignment/task name does not change), the already-assigned + * task's cached {@link DatastreamTask#getDatastreams()} list is refreshed to include the new datastream, instead + * of going stale. Without the fix, the task object survives the dedupe unchanged (same object, per + * {@link DatastreamTaskImpl#equals}), but its datastreams list would still only contain the original stream. + */ + @Test + public void testTaskDatastreamsRefreshedAfterDedupeReusesExistingTask() throws Exception { + String testCluster = "testTaskDatastreamsRefreshedAfterDedupeReusesExistingTask"; + String connectorType = "connectorType"; + + TestHookConnector connector1 = new TestHookConnector("connector1", connectorType); + + Coordinator coordinator1 = createCoordinator(_zkConnectionString, testCluster); + coordinator1.addConnector(connectorType, connector1, new BroadcastStrategy(Optional.empty()), false, + new SourceBasedDeduper(), null); + coordinator1.start(); + + ZkClient zkClient = new ZkClient(_zkConnectionString); + + Datastream[] datastreams = DatastreamTestUtils.createDatastreams(connectorType, "stream1", "stream2"); + datastreams[0].getMetadata().put(DatastreamMetadataConstants.TASK_PREFIX, "MyPrefix"); + datastreams[1].getMetadata().put(DatastreamMetadataConstants.TASK_PREFIX, "MyPrefix"); + + // Store and wait for stream1 to be assigned; this creates the task that stream2 will later dedupe into + DatastreamTestUtils.storeDatastreams(zkClient, testCluster, datastreams[0]); + Assert.assertTrue(PollUtils.poll(() -> connector1.getTasks().size() == 1, 50, WAIT_TIMEOUT_MS)); + DatastreamTask task = connector1.getTasks().get(0); + Assert.assertTrue(PollUtils.poll(() -> task.getDatastreams().size() == 1, 50, WAIT_TIMEOUT_MS)); + + // Dedupe stream2 against stream1's destination, then store it + datastreams[1].setSource(datastreams[0].getSource()); + datastreams[1].removeDestination(); + coordinator1.initializeDatastream(datastreams[1]); + DatastreamTestUtils.storeDatastreams(zkClient, testCluster, datastreams[1]); + + // Wait for stream2 to become READY + Assert.assertTrue(PollUtils.poll(() -> { + Datastream ds = DatastreamTestUtils.getDatastream(zkClient, testCluster, "stream2"); + return ds.getStatus() == DatastreamStatus.READY; + }, 50, WAIT_TIMEOUT_MS)); + + // No new task should have been created for the dedup stream: it must reuse the existing task + Assert.assertEquals(connector1.getTasks().size(), 1); + Assert.assertEquals(connector1.getTasks().get(0), task); + + // The reused task's datastreams list must be refreshed to include the newly deduped stream + Assert.assertTrue(PollUtils.poll(() -> task.getDatastreams().size() == 2, 50, WAIT_TIMEOUT_MS), + "Task's datastreams list was not refreshed after dedupe; still contains: " + task.getDatastreams()); + Assert.assertTrue(task.getDatastreams().stream().anyMatch(ds -> ds.getName().equals("stream2"))); + + coordinator1.stop(); + coordinator1.getDatastreamCache().getZkclient().close(); + zkClient.close(); + } + private static class TestCoordinatorWithSpyZkAdapter extends Coordinator { TestCoordinatorWithSpyZkAdapter(CachedDatastreamReader testDatastreamCache, Properties testConfig) throws DatastreamException { @@ -4509,11 +4568,20 @@ private String getNumThroughputViolatingTopicsMetric(String key) { // tasks with the specified names. private void assertConnectorAssignment(TestHookConnector connector, long timeoutMs, String... datastreamNames) throws InterruptedException { + assertConnectorAssignment(connector, timeoutMs, datastreamNames.length, datastreamNames); + } + + // helper method: assert that within a timeout value, the connector is assigned exactly expectedTaskCount tasks, + // each represented (via getDatastreams().get(0)) by one of the given acceptable datastream names. Useful when a + // task's representative datastream name may vary across runs (e.g. once its datastreams list is refreshed to + // reflect a dedup, the task may equally be represented by any datastream in its group). + private void assertConnectorAssignment(TestHookConnector connector, long timeoutMs, int expectedTaskCount, + String... datastreamNames) throws InterruptedException { final long interval = timeoutMs < 100 ? timeoutMs : 100; - boolean result = - PollUtils.poll(() -> validateAssignment(connector.getTasks(), datastreamNames), interval, timeoutMs); + boolean result = PollUtils.poll(() -> validateAssignment(connector.getTasks(), expectedTaskCount, datastreamNames), + interval, timeoutMs); LOG.info( String.format("assertConnectorAssignment. Connector: %s, Connector Tasks: %s, ASSERT: %s", connector.getName(), @@ -4534,8 +4602,13 @@ private boolean waitTillAssignmentIsComplete(int totalTasks, long timeoutMs, Tes } private boolean validateAssignment(List assignment, String... datastreamNames) { - if (assignment.size() != datastreamNames.length) { - LOG.warn("Expected size: " + datastreamNames.length + ", Actual size: " + assignment.size()); + return validateAssignment(assignment, datastreamNames.length, datastreamNames); + } + + private boolean validateAssignment(List assignment, int expectedTaskCount, + String... datastreamNames) { + if (assignment.size() != expectedTaskCount) { + LOG.warn("Expected size: " + expectedTaskCount + ", Actual size: " + assignment.size()); return false; } diff --git a/datastream-server/src/main/java/com/linkedin/datastream/server/Coordinator.java b/datastream-server/src/main/java/com/linkedin/datastream/server/Coordinator.java index b2a09b056..e3efccdc4 100644 --- a/datastream-server/src/main/java/com/linkedin/datastream/server/Coordinator.java +++ b/datastream-server/src/main/java/com/linkedin/datastream/server/Coordinator.java @@ -1362,6 +1362,16 @@ private void handleDatastreamAddOrDelete() { return; } + // Snapshot of task prefixes that already have an established (non-INITIALIZING, non-deleting/expired) + // datastream *before* this batch is processed. Used below to detect a "dedup" datastream: one that is about + // to become READY and reuses the task(s) of one of these already-established datastreams (same taskPrefix), + // as opposed to a brand-new datastream that will get a brand-new task. + Set establishedTaskPrefixes = activeStreams.stream() + .filter(ds -> ds.getStatus() != DatastreamStatus.INITIALIZING) + .map(Coordinator::getEffectiveTaskPrefix) + .collect(Collectors.toSet()); + + boolean anyDatastreamJoinedExistingTask = false; for (Datastream ds : allStreams) { if (ds.getStatus() == DatastreamStatus.INITIALIZING) { try { @@ -1379,6 +1389,9 @@ private void handleDatastreamAddOrDelete() { shouldRetry = true; } else { recordStreamProvisioningTime(ds); + if (establishedTaskPrefixes.contains(getEffectiveTaskPrefix(ds))) { + anyDatastreamJoinedExistingTask = true; + } } } catch (Exception e) { _log.warn("Failed to update the destination of new datastream {}", ds, e); @@ -1392,6 +1405,19 @@ private void handleDatastreamAddOrDelete() { } } + if (anyDatastreamJoinedExistingTask) { + // A newly READY datastream reused the task(s) of an existing, already-assigned datastream (same + // taskPrefix/partitions, e.g. destination dedup). In that case, the assignment (task names) does not + // change, so the normal assignment diff/rebalance path never touches the affected task's ZK node and the + // task's cached _datastreams list on every instance goes stale (see getDatastreams() javadoc). Explicitly + // broadcast so every live instance re-reads datastream groups and refreshes the datastreams list of its + // already-assigned tasks via onDatastreamUpdate(), the same way an explicit datastream update/ + // pause-partitions call already does today. We only do this when a task is actually reused (as opposed to + // every datastream creation) to avoid an extra onAssignmentChange() dispatch for the common case of a + // brand-new datastream getting a brand-new task. + broadcastDatastreamUpdate(); + } + if (shouldRetry) { _metrics.updateKeyedMeter(CoordinatorMetrics.KeyedMeter.HANDLE_DATASTREAM_ADD_OR_DELETE_NUM_RETRIES, 1); @@ -1412,6 +1438,15 @@ private void handleDatastreamAddOrDelete() { _log.info("END: Coordinator::handleDatastreamAddOrDelete."); } + /** + * Get the effective task prefix of a datastream: the explicit {@code system.taskPrefix} metadata if present + * (e.g. copied from another datastream during destination dedup), otherwise the datastream's own name. + */ + private static String getEffectiveTaskPrefix(Datastream datastream) { + String taskPrefix = Objects.requireNonNull(datastream.getMetadata()).get(DatastreamMetadataConstants.TASK_PREFIX); + return taskPrefix != null ? taskPrefix : DatastreamTaskImpl.getTaskPrefix(datastream); + } + /** * Records the time from {@code system.creation.ms} to the {@code INITIALIZING -> READY} * transition for a newly created datastream. Updates the {@code streamProvisioningTimeMs} histogram and the