[vpj] Add kill detection during data writing phase

When a push job is killed by the controller (e.g., a user push supersedes
a repush), the VPJ now detects this during the data writing phase instead
of waiting until pollStatusUntilComplete(). A periodic kill-check monitor
queries the controller for the push status and kills the data writer job
immediately if the push has been terminated, avoiding hours of wasted
resource usage on large stores.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: pthirun

clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/VenicePushJob.java

@@ -180,6 +180,7 @@
 import java.util.Set;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import org.apache.avro.Schema;
@@ -261,6 +262,8 @@ public class VenicePushJob implements AutoCloseable {
   private final PushJobHeartbeatSenderFactory pushJobHeartbeatSenderFactory;
   private PushJobHeartbeatSender pushJobHeartbeatSender = null;
   private volatile boolean pushJobStatusUploadDisabledHasBeenLogged = false;
+  private ScheduledFuture<?> pushJobKillCheckScheduledFuture;
+  private volatile boolean pushJobKilledByController = false;
   private final ScheduledExecutorService timeoutExecutor;
   private static final int VERSION_SWAP_BUFFER_TIME_MINUTES = 20;
 
@@ -827,7 +830,13 @@ public void run() {
         LOGGER.info("Incremental Push Version: {}", pushJobSetting.incrementalPushVersion);
         getVeniceWriter(pushJobSetting)
             .broadcastStartOfIncrementalPush(pushJobSetting.incrementalPushVersion, new HashMap<>());
-        runJobAndUpdateStatus();
+        startPushJobKillCheckMonitor();
+        try {
+          runJobAndUpdateStatus();
+        } finally {
+          stopPushJobKillCheckMonitor();
+        }
+        throwIfPushJobKilledByController();
         getVeniceWriter(pushJobSetting)
             .broadcastEndOfIncrementalPush(pushJobSetting.incrementalPushVersion, Collections.emptyMap());
       } else {
@@ -849,7 +858,13 @@ public void run() {
            * {@link createNewStoreVersion(PushJobSetting, long, ControllerClient, String, VeniceProperties)}
            */
         }
-        runJobAndUpdateStatus();
+        startPushJobKillCheckMonitor();
+        try {
+          runJobAndUpdateStatus();
+        } finally {
+          stopPushJobKillCheckMonitor();
+        }
+        throwIfPushJobKilledByController();
 
         if (!pushJobSetting.suppressEndOfPushMessage) {
           if (pushJobSetting.sendControlMessagesDirectly) {
@@ -994,6 +1009,59 @@ private void setupJobTimeoutMonitor() {
     }, timeoutMs, TimeUnit.MILLISECONDS);
   }
 
+  /**
+   * Schedules a periodic task that checks whether the push job has been killed by the controller.
+   * This runs during the data writing phase to detect early kills (e.g., when a user push supersedes
+   * a repush) and abort the data writer job promptly instead of wasting resources.
+   */
+  void startPushJobKillCheckMonitor() {
+    String topicToMonitor = getTopicToMonitor(pushJobSetting);
+    long intervalMs = pushJobSetting.pollJobStatusIntervalMs;
+    LOGGER.info("Starting push job kill check monitor for topic: {} with interval: {} ms", topicToMonitor, intervalMs);
+    pushJobKillCheckScheduledFuture = timeoutExecutor.scheduleAtFixedRate(() -> {
+      try {
+        JobStatusQueryResponse response = ControllerClient.retryableRequest(
+            controllerClient,
+            pushJobSetting.controllerStatusPollRetries,
+            client -> client.queryOverallJobStatus(topicToMonitor, Optional.empty(), null, false));
+        if (response.isError()) {
+          LOGGER.warn(
+              "Kill check monitor could not query job status for topic: {}. Error: {}",
+              topicToMonitor,
+              response.getError());
+          return;
+        }
+        ExecutionStatus status = getExecutionStatusFromControllerResponse(response);
+        if (status.isTerminal() && status.isError()) {
+          LOGGER.error(
+              "Kill check monitor detected that push job for topic: {} has been killed. Status: {}",
+              topicToMonitor,
+              status);
+          pushJobKilledByController = true;
+          killDataWriterJob();
+        }
+      } catch (Exception e) {
+        LOGGER.warn("Kill check monitor encountered an error while checking job status", e);
+      }
+    }, intervalMs, intervalMs, TimeUnit.MILLISECONDS);
+  }
+
+  void stopPushJobKillCheckMonitor() {
+    if (pushJobKillCheckScheduledFuture != null) {
+      pushJobKillCheckScheduledFuture.cancel(false);
+      pushJobKillCheckScheduledFuture = null;
+      LOGGER.info("Stopped push job kill check monitor");
+    }
+  }
+
+  private void throwIfPushJobKilledByController() {
+    if (pushJobKilledByController) {
+      throw new VeniceException(
+          "Push job for store " + pushJobSetting.storeName + " (topic: " + pushJobSetting.topic
+              + ") was killed by the controller during the data writing phase.");
+    }
+  }
+
   private void buildHDFSSchemaDir() throws IOException {
     // Build the full path for HDFSRmdSchemaSource:
     // RMD schemas: <job_temp_dir>/rmd_schemas

clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobTest.java

@@ -1839,4 +1839,78 @@ public void testTargetRegionPushWithDeferredSwapVersionStatusChecks(
       }
     }
   }
+
+  /**
+   * Test that VPJ detects a killed push during the data writing phase and kills the data writer job.
+   * This simulates the scenario where a repush is superseded by a user push, and the controller kills
+   * the repush version while data is still being written.
+   */
+  @Test(dataProvider = "DataWriterJobClasses")
+  public void testPushJobKilledDuringDataWriting(Class<? extends DataWriterComputeJob> dataWriterJobClass)
+      throws Exception {
+    Properties props = getVpjRequiredProperties();
+    props.put(KEY_FIELD_PROP, "id");
+    props.put(VALUE_FIELD_PROP, "name");
+    props.put(DATA_WRITER_COMPUTE_JOB_CLASS, dataWriterJobClass.getCanonicalName());
+    ControllerClient client = getClient();
+
+    // Simulate controller returning ERROR status (push was killed)
+    JobStatusQueryResponse killResponse = mock(JobStatusQueryResponse.class);
+    doReturn("ERROR").when(killResponse).getStatus();
+    doReturn(false).when(killResponse).isError();
+    doReturn(killResponse).when(client).queryOverallJobStatus(anyString(), any(), any(), anyBoolean());
+
+    try (VenicePushJob pushJob = getSpyVenicePushJob(props, client)) {
+      PushJobSetting pushJobSetting = pushJob.getPushJobSetting();
+      pushJobSetting.pollJobStatusIntervalMs = 10; // Poll quickly for the test
+
+      CountDownLatch dataWriterRunningLatch = new CountDownLatch(1);
+      CountDownLatch dataWriterKilledLatch = new CountDownLatch(1);
+
+      // Stall the data writer job until it gets killed by the kill-check monitor
+      doCallRealMethod().when(pushJob).runJobAndUpdateStatus();
+      doCallRealMethod().when(pushJob).startPushJobKillCheckMonitor();
+      doCallRealMethod().when(pushJob).stopPushJobKillCheckMonitor();
+      doCallRealMethod().when(pushJob).killDataWriterJob();
+
+      DataWriterComputeJob dataWriterJob = spy(pushJob.getDataWriterComputeJob());
+      pushJob.setDataWriterComputeJob(dataWriterJob);
+      doNothing().when(dataWriterJob).configure(any(), any());
+      doNothing().when(dataWriterJob).validateJob();
+
+      Answer<Void> stallDataWriterJob = invocation -> {
+        dataWriterRunningLatch.countDown();
+        if (!dataWriterKilledLatch.await(10, TimeUnit.SECONDS)) {
+          fail("Timed out waiting for the data writer job to be killed by kill-check monitor");
+        }
+        throw new VeniceException("Data writer job was killed");
+      };
+      doAnswer(stallDataWriterJob).when(dataWriterJob).runComputeJob();
+
+      // When dataWriterJob.kill() is called, release the stalled data writer
+      doAnswer(invocation -> {
+        invocation.callRealMethod();
+        dataWriterKilledLatch.countDown();
+        return null;
+      }).when(dataWriterJob).kill();
+
+      skipVPJValidation(pushJob);
+      // Override skipVPJValidation's stub on runJobAndUpdateStatus
+      doCallRealMethod().when(pushJob).runJobAndUpdateStatus();
+
+      try {
+        pushJob.run();
+        fail("Expected VeniceException due to push job being killed during data writing");
+      } catch (VeniceException e) {
+        assertTrue(
+            e.getMessage().contains("killed by the controller during the data writing phase")
+                || e.getMessage().contains("Data writer job was killed"),
+            "Unexpected error message: " + e.getMessage());
+      }
+
+      assertEquals(dataWriterRunningLatch.getCount(), 0, "Data writer job should have started");
+      assertEquals(dataWriterKilledLatch.getCount(), 0, "Data writer job should have been killed");
+      verify(dataWriterJob, times(1)).kill();
+    }
+  }
 }