Close temporal metrics scope on job completion

Author: abhishekmjain

gobblin-temporal/src/main/java/org/apache/gobblin/temporal/cluster/GobblinTemporalTaskRunner.java

@@ -126,6 +126,7 @@ public class GobblinTemporalTaskRunner implements StandardMetricsBridge {
   private final boolean isMetricReportingFailureFatal;
   private final boolean isEventReportingFailureFatal;
   private final List<TemporalWorker> workers;
+  private final WorkflowServiceStubs workflowServiceStubs;
 
   public GobblinTemporalTaskRunner(String applicationName,
       String applicationId,
@@ -163,6 +164,9 @@ public GobblinTemporalTaskRunner(String applicationName,
         ConfigurationKeys.DEFAULT_GOBBLIN_TASK_EVENT_REPORTING_FAILURE_FATAL);
     this.workers = new ArrayList<>();
 
+    String connectionUri = clusterConfig.getString(GobblinTemporalConfigurationKeys.TEMPORAL_CONNECTION_STRING);
+    this.workflowServiceStubs = TemporalWorkflowClientFactory.createServiceInstance(connectionUri);
+
     logger.info("GobblinTaskRunner({}): applicationName {}, applicationId {}, taskRunnerId {}, config {}, appWorkDir {}",
         this.isTaskDriver ? "taskDriver" : "worker",
         applicationName,
@@ -241,12 +245,9 @@ public void start()
   private TemporalWorker initiateWorker() throws Exception {
     logger.info("Starting Temporal Worker");
 
-    String connectionUri = clusterConfig.getString(GobblinTemporalConfigurationKeys.TEMPORAL_CONNECTION_STRING);
-    WorkflowServiceStubs service = TemporalWorkflowClientFactory.createServiceInstance(connectionUri);
-
     String namespace = ConfigUtils.getString(clusterConfig, GobblinTemporalConfigurationKeys.GOBBLIN_TEMPORAL_NAMESPACE,
             GobblinTemporalConfigurationKeys.DEFAULT_GOBBLIN_TEMPORAL_NAMESPACE);
-    WorkflowClient client = TemporalWorkflowClientFactory.createClientInstance(service, namespace);
+    WorkflowClient client = TemporalWorkflowClientFactory.createClientInstance(workflowServiceStubs, namespace);
 
     String workerClassName = ConfigUtils.getString(clusterConfig,
         GobblinTemporalConfigurationKeys.WORKER_CLASS, GobblinTemporalConfigurationKeys.DEFAULT_WORKER_CLASS);
@@ -292,6 +293,12 @@ public synchronized void stop() {
       this.containerMetrics.get().stopMetricsReporting();
     }
 
+    try {
+      this.workflowServiceStubs.getOptions().getMetricsScope().close();
+    } catch (Exception e) {
+      logger.error("Exception occurred while closing MetricsScope", e);
+    }
+
     workers.forEach(TemporalWorker::shutdown);
 
     logger.info("All services are stopped.");

gobblin-temporal/src/main/java/org/apache/gobblin/temporal/joblauncher/GobblinTemporalJobLauncher.java

@@ -17,6 +17,7 @@
 
 package org.apache.gobblin.temporal.joblauncher;
 
+import java.io.IOException;
 import java.util.List;
 import java.util.Properties;
 import java.util.concurrent.ConcurrentHashMap;
@@ -188,4 +189,15 @@ protected void removeTasksFromCurrentJob(List<String> workUnitIdsToRemove) {
   protected void addTasksToCurrentJob(List<WorkUnit> workUnitsToAdd) {
     log.warn("NOT IMPLEMENTED: Temporal addTasksToCurrentJob");
   }
+
+  @Override
+  public void close() throws IOException {
+    try {
+      this.workflowServiceStubs.getOptions().getMetricsScope().close();
+    } catch (Exception e) {
+      log.error("Exception occurred while closing MetricsScope ", e);
+    } finally {
+      super.close();
+    }
+  }
 }

gobblin-temporal/src/main/java/org/apache/gobblin/temporal/joblauncher/GobblinTemporalJobScheduler.java

@@ -210,15 +210,16 @@ public void handleNewJobConfigArrival(NewJobConfigArrivalEvent newJobArrival) {
         LOGGER.info("No job schedule found, so running job " + jobUri);
         GobblinTemporalJobLauncherListener listener = new GobblinTemporalJobLauncherListener(this.launcherMetrics);
         JobLauncher launcher = buildJobLauncher(newJobArrival.getJobConfig());
-        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
+        try {
+          launcher.launchJob(listener);
+        } finally {
           try {
             launcher.cancelJob(listener);
-          } catch (JobException e) {
-            LOGGER.error("Failed to cancel the job during shutdown", e);
-            throw new RuntimeException(e);
           }
-        }));
-        launcher.launchJob(listener);
+          finally {
+            launcher.close();
+          }
+        }
       }
     } catch (Exception je) {
       LOGGER.error("Failed to schedule or run job " + jobUri, je);