[vpj] Push Job Timeout (#1786)

There should be a timeout on the entire push job as a whole. If exceeded, the push job should be cancelled. 😭

Author: KaiSernLim

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

@@ -98,6 +98,7 @@
 import com.linkedin.venice.exceptions.ErrorType;
 import com.linkedin.venice.exceptions.VeniceException;
 import com.linkedin.venice.exceptions.VeniceResourceAccessException;
+import com.linkedin.venice.exceptions.VeniceTimeoutException;
 import com.linkedin.venice.hadoop.exceptions.VeniceInvalidInputException;
 import com.linkedin.venice.hadoop.input.kafka.KafkaInputDictTrainer;
 import com.linkedin.venice.hadoop.mapreduce.datawriter.jobs.DataWriterMRJob;
@@ -164,6 +165,8 @@
 import java.util.Optional;
 import java.util.Properties;
 import java.util.Set;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import org.apache.avro.Schema;
@@ -245,6 +248,7 @@ public class VenicePushJob implements AutoCloseable {
   private final PushJobHeartbeatSenderFactory pushJobHeartbeatSenderFactory;
   private PushJobHeartbeatSender pushJobHeartbeatSender = null;
   private boolean pushJobStatusUploadDisabledHasBeenLogged = false;
+  private final ScheduledExecutorService timeoutExecutor;
 
   /**
    * @param jobId  id of the job
@@ -253,6 +257,7 @@ public class VenicePushJob implements AutoCloseable {
   public VenicePushJob(String jobId, Properties vanillaProps) {
     this.jobId = jobId;
     this.props = getVenicePropsFromVanillaProps(Objects.requireNonNull(vanillaProps, "VPJ props cannot be null"));
+    this.timeoutExecutor = Executors.newSingleThreadScheduledExecutor();
     LOGGER.info("Constructing {}: {}", VenicePushJob.class.getSimpleName(), props.toString(true));
     this.sslProperties = Lazy.of(() -> {
       try {
@@ -657,11 +662,7 @@ DataWriterComputeJob getDataWriterComputeJob() {
    */
   public void run() {
     try {
-      Optional<SSLFactory> sslFactory = VPJSSLUtils.createSSLFactory(
-          pushJobSetting.enableSSL,
-          props.getString(SSL_FACTORY_CLASS_NAME, DEFAULT_SSL_FACTORY_CLASS_NAME),
-          this.sslProperties);
-      initControllerClient(pushJobSetting.storeName, sslFactory);
+      initControllerClient(pushJobSetting.storeName);
       pushJobSetting.clusterName = controllerClient.getClusterName();
       LOGGER.info(
           "The store {} is discovered in Venice cluster {}",
@@ -672,6 +673,7 @@ public void run() {
         initKIFRepushDetails();
       }
 
+      setupJobTimeoutMonitor();
       initPushJobDetails();
       logGreeting();
       sendPushJobDetailsToController();
@@ -887,6 +889,20 @@ public void run() {
     }
   }
 
+  /**
+   * Timeout on the entire push job that kills the job if it runs longer than the store's configured bootstrap timeout.
+   */
+  private void setupJobTimeoutMonitor() {
+    long bootstrapToOnlineTimeoutInHours =
+        getStoreResponse(pushJobSetting.storeName).getStore().getBootstrapToOnlineTimeoutInHours();
+    timeoutExecutor.schedule(() -> {
+      cancel();
+      throw new VeniceTimeoutException(
+          "Failing push-job for store " + pushJobSetting.storeName + " which is still running after "
+              + bootstrapToOnlineTimeoutInHours + " hours.");
+    }, bootstrapToOnlineTimeoutInHours, TimeUnit.HOURS);
+  }
+
   private void buildHDFSSchemaDir() throws IOException {
     // Build the full path for HDFSRmdSchemaSource:
     // RMD schemas: <job_temp_dir>/rmd_schemas
@@ -1199,9 +1215,12 @@ protected InputDataInfoProvider getInputDataInfoProvider() {
    *    2. A mock controller client is provided
    *
    * @param storeName
-   * @param sslFactory
    */
-  private void initControllerClient(String storeName, Optional<SSLFactory> sslFactory) {
+  private void initControllerClient(String storeName) {
+    Optional<SSLFactory> sslFactory = VPJSSLUtils.createSSLFactory(
+        pushJobSetting.enableSSL,
+        props.getString(SSL_FACTORY_CLASS_NAME, DEFAULT_SSL_FACTORY_CLASS_NAME),
+        this.sslProperties);
     final String controllerD2ZkHost;
     if (pushJobSetting.multiRegion) {
       // In multi region mode, push jobs will communicate with parent controller
@@ -2312,7 +2331,6 @@ void pollStatusUntilComplete(
      * no more than {@link DEFAULT_JOB_STATUS_IN_UNKNOWN_STATE_TIMEOUT_MS}.
      */
     long unknownStateStartTimeMs = 0;
-    long pollStartTimeMs = System.currentTimeMillis();
 
     String topicToMonitor = getTopicToMonitor(pushJobSetting);
 
@@ -2382,14 +2400,6 @@ void pollStatusUntilComplete(
         }
         return;
       }
-      long bootstrapToOnlineTimeoutInHours =
-          VenicePushJob.this.pushJobSetting.storeResponse.getStore().getBootstrapToOnlineTimeoutInHours();
-      long durationMs = LatencyUtils.getElapsedTimeFromMsToMs(pollStartTimeMs);
-      if (durationMs > TimeUnit.HOURS.toMillis(bootstrapToOnlineTimeoutInHours)) {
-        throw new VeniceException(
-            "Failing push-job for store " + VenicePushJob.this.pushJobSetting.storeResponse.getName()
-                + " which is still running after " + TimeUnit.MILLISECONDS.toHours(durationMs) + " hours.");
-      }
       if (!overallStatus.equals(ExecutionStatus.UNKNOWN)) {
         unknownStateStartTimeMs = 0;
       } else if (unknownStateStartTimeMs == 0) {
@@ -2557,9 +2567,8 @@ private String pushJobPropertiesToString(
   }
 
   /**
-   * A cancel method for graceful cancellation of the running Job to be invoked as a result of user actions.
-   *
-   * @throws Exception
+   * A cancel method for graceful cancellation of the running Job to be invoked as a result of user actions or due to
+   * the job exceeding bootstrapToOnlineTimeoutInHours.
    */
   public void cancel() {
     killJob(pushJobSetting, controllerClient);
@@ -2572,7 +2581,7 @@ public void cancel() {
     sendPushJobDetailsToController();
   }
 
-  private void killJob(PushJobSetting pushJobSetting, ControllerClient controllerClient) {
+  void killJob(PushJobSetting pushJobSetting, ControllerClient controllerClient) {
     // Attempting to kill job. There's a race condition, but meh. Better kill when you know it's running
     killDataWriterJob();
     if (!pushJobSetting.isIncrementalPush) {
@@ -2597,7 +2606,7 @@ private void killJob(PushJobSetting pushJobSetting, ControllerClient controllerC
     }
   }
 
-  private void killDataWriterJob() {
+  void killDataWriterJob() {
     if (dataWriterComputeJob == null) {
       LOGGER.warn("No op to kill a null data writer job");
       return;
@@ -2671,6 +2680,7 @@ private static Path getLatestPath(Path path, FileSystem fs) throws IOException {
 
   @Override
   public void close() {
+    timeoutExecutor.shutdownNow();
     closeVeniceWriter();
     Utils.closeQuietlyWithErrorLogged(dataWriterComputeJob);
     Utils.closeQuietlyWithErrorLogged(controllerClient);
@@ -2684,7 +2694,6 @@ public void close() {
   }
 
   public static void main(String[] args) {
-
     if (args.length != 1) {
       Utils.exit("USAGE: java -jar venice-push-job-all.jar <VPJ_config_file_path>");
     }

clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/mapreduce/datawriter/jobs/DataWriterMRJob.java

@@ -328,7 +328,7 @@ public DataWriterTaskTracker getTaskTracker() {
   }
 
   @Override
-  protected void runComputeJob() {
+  public void runComputeJob() {
     LOGGER.info("Triggering MR job for data writer");
     try {
       runningJob = JobUtils.runJobWithConfig(jobConf, jobClientWrapper);
@@ -346,6 +346,7 @@ public PushJobSetting getPushJobSetting() {
   public void kill() {
     if (runningJob == null) {
       LOGGER.warn("No op to kill a null running job");
+      super.kill();
       return;
     }
     try {

clients/venice-push-job/src/main/java/com/linkedin/venice/jobs/DataWriterComputeJob.java

@@ -2,6 +2,7 @@
 
 import static com.linkedin.venice.ConfigKeys.PASS_THROUGH_CONFIG_PREFIXES_LIST_KEY;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.linkedin.venice.ConfigKeys;
 import com.linkedin.venice.exceptions.VeniceException;
 import com.linkedin.venice.hadoop.PushJobSetting;
@@ -88,7 +89,8 @@ public static void populateWithPassThroughConfigs(VeniceProperties props, Config
 
   public abstract DataWriterTaskTracker getTaskTracker();
 
-  protected void validateJob() {
+  @VisibleForTesting
+  public void validateJob() {
     DataWriterTaskTracker dataWriterTaskTracker = getTaskTracker();
     if (dataWriterTaskTracker == null) {
       throw new VeniceException("DataWriterTaskTracker is not set. Unable to validate the job status.");
@@ -160,7 +162,8 @@ protected void validateJob() {
     }
   }
 
-  protected abstract void runComputeJob();
+  @VisibleForTesting
+  public abstract void runComputeJob();
 
   @Override
   public void configure(VeniceProperties properties) {

clients/venice-push-job/src/main/java/com/linkedin/venice/spark/datawriter/jobs/AbstractDataWriterSparkJob.java

@@ -323,7 +323,7 @@ public PushJobSetting getPushJobSetting() {
   }
 
   @Override
-  protected void runComputeJob() {
+  public void runComputeJob() {
     // Load data from input path
     Dataset<Row> dataFrame = getInputDataFrame();
     validateDataFrame(dataFrame);