Merge branch 'apache:master' into ETL-18960

Author: agam-99

gobblin-api/src/main/java/org/apache/gobblin/service/ServiceConfigKeys.java

@@ -159,6 +159,8 @@ public class ServiceConfigKeys {
   public static final String NUM_DAG_PROC_THREADS_KEY = GOBBLIN_SERVICE_DAG_PROCESSING_ENGINE_PREFIX + "numThreads";
   public static final String DAG_PROC_ENGINE_NON_RETRYABLE_EXCEPTIONS_KEY = GOBBLIN_SERVICE_DAG_PROCESSING_ENGINE_PREFIX + "nonRetryableExceptions";
   public static final Integer DEFAULT_NUM_DAG_PROC_THREADS = 3;
+  public static final String NUM_SPEC_CATALOG_LISTENER_THREADS_KEY = GOBBLIN_SERVICE_PREFIX + "specCatalogListener.numThreads";
+  public static final int DEFAULT_NUM_SPEC_CATALOG_LISTENER_THREADS = 3;
   public static final long DEFAULT_FLOW_FINISH_DEADLINE_MILLIS = TimeUnit.HOURS.toMillis(24);
 
   public static final String ERROR_PATTERN_STORE_CLASS = ServiceConfigKeys.GOBBLIN_SERVICE_PREFIX + "errorPatternStore.class";

gobblin-cluster/src/main/java/org/apache/gobblin/cluster/GobblinClusterConfigurationKeys.java

@@ -154,6 +154,8 @@ public class GobblinClusterConfigurationKeys {
 
   public static final String HELIX_JOB_STOP_TIMEOUT_SECONDS = GOBBLIN_CLUSTER_PREFIX + "helix.job.stopTimeoutSeconds";
   public static final long DEFAULT_HELIX_JOB_STOP_TIMEOUT_SECONDS = 10L;
+  public static final String TASK_STATE_MODEL_FACTORY_SHUTDOWN_TIMEOUT_SECONDS = GOBBLIN_CLUSTER_PREFIX + "taskStateModelFactory.shutdownTimeoutSeconds";
+  public static final long DEFAULT_TASK_STATE_MODEL_FACTORY_SHUTDOWN_TIMEOUT_SECONDS = 300L;
   public static final String TASK_RUNNER_SUITE_BUILDER = GOBBLIN_CLUSTER_PREFIX + "taskRunnerSuite.builder";
 
   public static final String HELIX_JOB_NAME_KEY = GOBBLIN_CLUSTER_PREFIX + "helixJobName";

gobblin-cluster/src/main/java/org/apache/gobblin/cluster/GobblinTaskRunner.java

@@ -177,6 +177,7 @@ public class GobblinTaskRunner implements StandardMetricsBridge {
   protected final String applicationId;
   private final boolean isMetricReportingFailureFatal;
   private final boolean isEventReportingFailureFatal;
+  private final long taskStateModelFactoryShutdownTimeoutSeconds;
 
   public GobblinTaskRunner(String applicationName,
       String helixInstanceName,
@@ -213,6 +214,10 @@ public GobblinTaskRunner(String applicationName,
         ConfigurationKeys.GOBBLIN_TASK_EVENT_REPORTING_FAILURE_FATAL,
         ConfigurationKeys.DEFAULT_GOBBLIN_TASK_EVENT_REPORTING_FAILURE_FATAL);
 
+    this.taskStateModelFactoryShutdownTimeoutSeconds = ConfigUtils.getLong(this.clusterConfig,
+        GobblinClusterConfigurationKeys.TASK_STATE_MODEL_FACTORY_SHUTDOWN_TIMEOUT_SECONDS,
+        GobblinClusterConfigurationKeys.DEFAULT_TASK_STATE_MODEL_FACTORY_SHUTDOWN_TIMEOUT_SECONDS);
+
     logger.info("Configured GobblinTaskRunner work dir to: {}", this.appWorkPath.toString());
 
     this.isContainerExitOnHealthCheckFailureEnabled = ConfigUtils.getBoolean(config, GobblinClusterConfigurationKeys.CONTAINER_EXIT_ON_HEALTH_CHECK_FAILURE_ENABLED,
@@ -415,20 +420,43 @@ public synchronized void stop() {
 
     logger.info("Stopping the Gobblin Task runner");
 
-    // Stop metric reporting
-    if (this.containerMetrics.isPresent()) {
-      this.containerMetrics.get().stopMetricsReporting();
-    }
-
     try {
-      stopServices();
+      try {
+        stopServices();
+      } finally {
+        logger.info("All services are stopped.");
+        shutdownTaskStateModelFactory();
+        disconnectHelixManager();
+      }
     } finally {
-      logger.info("All services are stopped.");
-      this.taskStateModelFactory.shutdown();
-      disconnectHelixManager();
+      // Stop metric reporting only after tasks have completed and emitted their final GTEs.
+      // Placed in an outer finally to guarantee execution even if shutdownTaskStateModelFactory()
+      // or disconnectHelixManager() throws an unexpected runtime exception.
+      if (this.containerMetrics.isPresent()) {
+        this.containerMetrics.get().stopMetricsReporting();
+      }
+      this.isStopped = true;
     }
+  }
 
-    this.isStopped = true;
+  private void shutdownTaskStateModelFactory() {
+    this.taskStateModelFactory.shutdown();
+    long timeoutMs = TimeUnit.SECONDS.toMillis(this.taskStateModelFactoryShutdownTimeoutSeconds);
+    long startTime = System.currentTimeMillis();
+    while (!this.taskStateModelFactory.isTerminated()
+        && System.currentTimeMillis() - startTime < timeoutMs) {
+      try {
+        Thread.sleep(1000);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        logger.warn("Interrupted while waiting for task state model factory to terminate");
+        break;
+      }
+    }
+    if (!this.taskStateModelFactory.isTerminated()) {
+      logger.warn("Task state model factory did not terminate within {} seconds",
+          this.taskStateModelFactoryShutdownTimeoutSeconds);
+    }
   }
 
   private void stopServices() {