GOBBLIN-2265: AM waits for workflow terminal state before exit

gobblin-runtime/src/main/java/org/apache/gobblin/runtime/app/ServiceBasedAppLauncher.java

@@ -109,7 +109,9 @@ public class ServiceBasedAppLauncher implements ApplicationLauncher {
   private volatile boolean hasStarted = false;
   private volatile boolean hasStopped = false;
 
-  private ServiceManager serviceManager;
+  // volatile: written once in start() and read without the monitor in awaitStopped(), so safe publication is
+  // needed (and keeps FindBugs IS2_INCONSISTENT_SYNC quiet).
+  private volatile ServiceManager serviceManager;
 
   public ServiceBasedAppLauncher(Properties properties, String appName) throws Exception {
     this.stopTime = Integer.parseInt(properties.getProperty(APP_STOP_TIME_SECONDS, DEFAULT_APP_STOP_TIME_SECONDS));
@@ -215,6 +217,23 @@ public synchronized void stop() throws ApplicationException {
     }
   }
 
+  /**
+   * Block until every managed service has reached a terminal state (the application has fully stopped). Does not
+   * itself initiate shutdown; waits for a stop triggered elsewhere. Returns {@code false} if it times out first,
+   * or {@code true} immediately if the launcher was never {@link #start() started}.
+   */
+  public boolean awaitStopped(long timeout, TimeUnit unit) throws InterruptedException {
+    if (this.serviceManager == null) {
+      return true;
+    }
+    try {
+      this.serviceManager.awaitStopped(timeout, unit);
+      return true;
+    } catch (TimeoutException te) {
+      return false;
+    }
+  }
+
   @Override
   public void close() throws IOException {
     // Do nothing

gobblin-runtime/src/test/java/org/apache/gobblin/runtime/app/ServiceBasedAppLauncherTest.java

@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.gobblin.runtime.app;
+
+import java.util.Properties;
+import java.util.concurrent.TimeUnit;
+
+import org.testng.annotations.Test;
+
+import com.google.common.util.concurrent.AbstractIdleService;
+import com.google.common.util.concurrent.Service;
+
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+
+/** Unit tests for {@link ServiceBasedAppLauncher#awaitStopped(long, java.util.concurrent.TimeUnit)}. */
+public class ServiceBasedAppLauncherTest {
+
+  /** A do-nothing service that starts and stops cleanly. */
+  private static Service noopService() {
+    return new AbstractIdleService() {
+      @Override protected void startUp() { }
+      @Override protected void shutDown() { }
+    };
+  }
+
+  private static ServiceBasedAppLauncher newLauncher(String name) throws Exception {
+    Properties props = new Properties();
+    // Bound the service-start wait so a misbehaving service can't hang the test (default is FOREVER).
+    props.setProperty(ServiceBasedAppLauncher.STARTUP_TIMEOUT_SECONDS, "30");
+    return new ServiceBasedAppLauncher(props, name);
+  }
+
+  @Test
+  public void testAwaitStoppedReturnsTrueWhenNeverStarted() throws Exception {
+    // No serviceManager exists until start(); there is nothing to wait for, so it must return immediately.
+    ServiceBasedAppLauncher launcher = newLauncher("never-started");
+    assertTrue(launcher.awaitStopped(1, TimeUnit.SECONDS));
+  }
+
+  @Test
+  public void testAwaitStoppedTimesOutWhileRunningThenReturnsAfterStop() throws Exception {
+    ServiceBasedAppLauncher launcher = newLauncher("await-stopped");
+    launcher.addService(noopService());
+    launcher.start();
+
+    // Services are RUNNING (not terminal), so a short wait must time out -> false.
+    assertFalse(launcher.awaitStopped(200, TimeUnit.MILLISECONDS),
+        "awaitStopped must time out (false) while the application is still running");
+
+    launcher.stop();
+
+    // After stop(), every service reaches a terminal state, so awaitStopped must return true.
+    assertTrue(launcher.awaitStopped(30, TimeUnit.SECONDS),
+        "awaitStopped must return true once the application has stopped");
+  }
+}

gobblin-temporal/src/main/java/org/apache/gobblin/temporal/yarn/GobblinTemporalApplicationMaster.java

@@ -21,8 +21,10 @@
 import com.google.common.util.concurrent.Service;
 import com.typesafe.config.Config;
 import com.typesafe.config.ConfigFactory;
+import com.google.common.annotations.VisibleForTesting;
 import com.typesafe.config.ConfigValueFactory;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 import lombok.Getter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -32,6 +34,7 @@
 import org.apache.gobblin.annotation.Alpha;
 import org.apache.gobblin.cluster.GobblinClusterConfigurationKeys;
 import org.apache.gobblin.cluster.GobblinClusterUtils;
+import org.apache.gobblin.configuration.ConfigurationKeys;
 import org.apache.gobblin.temporal.cluster.GobblinTemporalClusterManager;
 import org.apache.gobblin.temporal.joblauncher.GobblinTemporalJobLauncher;
 import org.apache.gobblin.util.ConfigUtils;
@@ -127,6 +130,23 @@ private YarnContainerSecurityManager buildYarnContainerSecurityManager(Config co
     return new YarnTemporalAppMasterSecurityManager(config, fs, this.eventBus, this.logCopier, this._yarnService);
   }
 
+  /** Block until the AM application has fully stopped (all services incl. {@code YarnService} terminal). */
+  @VisibleForTesting
+  boolean awaitApplicationStopped(long timeout, TimeUnit unit) throws InterruptedException {
+    return this.applicationLauncher.awaitStopped(timeout, unit);
+  }
+
+  /** The configured flow SLA ({@code gobblin.flow.sla.time}) in millis, or a generous fallback when unset. */
+  private static long flowSlaMillis(Config config) {
+    if (!config.hasPath(ConfigurationKeys.GOBBLIN_FLOW_FINISH_DEADLINE_TIME)) {
+      return TimeUnit.DAYS.toMillis(1);
+    }
+    TimeUnit unit = TimeUnit.valueOf(ConfigUtils.getString(config,
+        ConfigurationKeys.GOBBLIN_FLOW_FINISH_DEADLINE_TIME_UNIT,
+        ConfigurationKeys.DEFAULT_GOBBLIN_FLOW_FINISH_DEADLINE_TIME_UNIT));
+    return unit.toMillis(config.getLong(ConfigurationKeys.GOBBLIN_FLOW_FINISH_DEADLINE_TIME));
+  }
+
   private static Options buildOptions() {
     Options options = new Options();
     options.addOption("a", GobblinClusterConfigurationKeys.APPLICATION_NAME_OPTION_NAME, true, "Yarn application name");
@@ -162,20 +182,28 @@ public static void main(String[] args) throws Exception {
       ContainerId containerId =
           ConverterUtils.toContainerId(System.getenv().get(ApplicationConstants.Environment.CONTAINER_ID.key()));
 
+      Config config = ConfigFactory.load();
+      WorkflowExecutionStatus terminalStatus;
       try (GobblinTemporalApplicationMaster applicationMaster = new GobblinTemporalApplicationMaster(
           cmd.getOptionValue(GobblinClusterConfigurationKeys.APPLICATION_NAME_OPTION_NAME),
           cmd.getOptionValue(GobblinClusterConfigurationKeys.APPLICATION_ID_OPTION_NAME), containerId,
-          ConfigFactory.load(), new YarnConfiguration())) {
+          config, new YarnConfiguration())) {
 
         applicationMaster.start();
+
+        // start() can return while the workflow is still running (the job runs inside a service startUp(), so
+        // ServiceBasedAppLauncher proceeds once app.start.waitForServicesTimeout elapses). Exiting then would
+        // cancel/un-register the in-flight workflow and mis-report a successful long job as CANCELLED. Instead wait
+        // until the application has actually stopped (all services incl. YarnService terminal -> after un-register),
+        // bounded by the flow SLA (gobblin.flow.sla.time); the GaaS control plane cancels SLA overruns, unblocking
+        // this wait.
+        if (!applicationMaster.awaitApplicationStopped(flowSlaMillis(config), TimeUnit.MILLISECONDS)) {
+          LOGGER.warn("AM did not stop within the flow SLA; proceeding to exit");
+        }
+        terminalStatus = GobblinTemporalJobLauncher.getLastTerminalStatus();
       }
 
-      // Surface the underlying workflow outcome as the AM JVM exit code so GGW/Grid Gateway dashboards see
-      // failures end-to-end. The status was captured into a static cache by GobblinTemporalJobLauncher via
-      // handleLaunchFinalization on normal completion (this is also what the temporal YarnService uses to derive
-      // the FinalApplicationStatus reported to YARN). A null cache means the workflow never reached a terminal
-      // state under this AM (preemption/error/crash mid-run) -> non-zero exit, consistent with FinalApplicationStatus.
-      WorkflowExecutionStatus terminalStatus = GobblinTemporalJobLauncher.getLastTerminalStatus();
+      // Surface the captured workflow outcome as the AM JVM exit code (null = never reached terminal -> non-zero).
       int exitCode = GobblinTemporalJobLauncher.computeExitCode(terminalStatus);
       LOGGER.info("GobblinTemporalApplicationMaster exiting with code {} (workflow terminal status: {})",
           exitCode, terminalStatus);