TEZ-4349. DAGClient gets stuck with invalid cached DAGStatus

tez-api/src/main/java/org/apache/tez/common/CachedEntity.java

@@ -0,0 +1,104 @@
+/**
+ * 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.tez.common;
+
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.apache.hadoop.yarn.util.Clock;
+import org.apache.hadoop.yarn.util.MonotonicClock;
+
+/**
+ * A thread safe implementation used as a container for cacheable entries with Expiration times.
+ * It supports custom {@link Clock} to control the elapsed time calculation.
+ * @param <T> the data object type.
+ */
+public class CachedEntity<T> {
+  private final AtomicReference<T> entryDataRef;
+  private final Clock cacheClock;
+  private final long expiryDurationMS;
+  private volatile long entryTimeStamp;
+
+  public CachedEntity(TimeUnit expiryTimeUnit, long expiryLength, Clock clock) {
+    entryDataRef = new AtomicReference<>(null);
+    cacheClock = clock;
+    expiryDurationMS = TimeUnit.MILLISECONDS.convert(expiryLength, expiryTimeUnit);
+    entryTimeStamp = 0;
+  }
+
+  public CachedEntity(TimeUnit expiryTimeUnit, long expiryLength) {
+    this(expiryTimeUnit, expiryLength, new MonotonicClock());
+  }
+
+  /**
+   *
+   * @return true if expiration timestamp is 0, or the elapsed time since last update is
+   *         greater than {@link #expiryDurationMS}
+   */
+  public boolean isExpired() {
+    return (entryTimeStamp == 0)
+        || ((cacheClock.getTime() - entryTimeStamp) > expiryDurationMS);
+  }
+
+  /**
+   * If the entry has expired, it reset the cache reference through {@link #clearExpiredEntry()}.
+   * @return cached data if the timestamp is valid. Null, if the timestamp has expired.
+   */
+  public T getValue() {
+    if (isExpired()) { // quick check for expiration
+      if (clearExpiredEntry()) { // remove reference to the expired entry
+        return null;
+      }
+    }
+    return entryDataRef.get();
+  }
+
+  /**
+   * Safely sets the cached data.
+   * @param newEntry
+   */
+  public void setValue(T newEntry) {
+    T currentEntry = entryDataRef.get();
+    while (!entryDataRef.compareAndSet(currentEntry, newEntry)) {
+      currentEntry = entryDataRef.get();
+    }
+    entryTimeStamp = cacheClock.getTime();
+  }
+
+  /**
+   * Enforces the expiration of the cached entry.
+   */
+  public void enforceExpiration() {
+    entryTimeStamp = 0;
+  }
+
+  /**
+   * Safely deletes the reference to the data if it was not null.
+   * @return true if the reference is set to Null. False indicates that another thread
+   *         updated the cache.
+   */
+  private boolean clearExpiredEntry() {
+    T currentEntry = entryDataRef.get();
+    if (currentEntry == null) {
+      return true;
+    }
+    // the current value is not null: try to reset it.
+    // if the CAS is successful, then we won't override a recent update to the cache.
+    return (entryDataRef.compareAndSet(currentEntry, null));
+  }
+}

tez-api/src/main/java/org/apache/tez/dag/api/TezConfiguration.java

@@ -1975,6 +1975,18 @@ static Set<String> getPropertySet() {
       TEZ_PREFIX + "test.minicluster.app.wait.on.shutdown.secs";
   public static final long TEZ_TEST_MINI_CLUSTER_APP_WAIT_ON_SHUTDOWN_SECS_DEFAULT = 30;
 
+  /**
+   * Long value
+   * Status Cache timeout window in minutes for the DAGClient.
+   */
+  @Private
+  @ConfigurationScope(Scope.CLIENT)
+  @ConfigurationProperty(type="long")
+  public static final String TEZ_CLIENT_DAG_STATUS_CACHE_TIMEOUT_MINUTES = TEZ_PREFIX
+      + "client.dag.status.cache.timeout-minutes";
+  // Default timeout is 5 minutes.
+  public static final long TEZ_CLIENT_DAG_STATUS_CACHE_TIMEOUT_MINUTES_DEFAULT = 5;
+
   /**
    * Long value
    * Time to wait (in milliseconds) for yarn app's diagnotics is available

tez-api/src/main/java/org/apache/tez/dag/api/client/DAGClientImpl.java

@@ -27,9 +27,11 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 
 import com.google.common.annotations.VisibleForTesting;
 import org.apache.hadoop.security.UserGroupInformation;
+import org.apache.tez.common.CachedEntity;
 import org.apache.tez.common.Preconditions;
 
 import org.apache.hadoop.yarn.exceptions.ApplicationNotFoundException;
@@ -58,13 +60,15 @@ public class DAGClientImpl extends DAGClient {
   private final String dagId;
   private final TezConfiguration conf;
   private final FrameworkClient frameworkClient;
-
+  /**
+   * Container to cache the last {@link DAGStatus}.
+   */
+  private final CachedEntity<DAGStatus> cachedDAGStatusRef;
   @VisibleForTesting
   protected DAGClientInternal realClient;
-  private boolean dagCompleted = false;
+  private volatile boolean dagCompleted = false;
   @VisibleForTesting
   protected boolean isATSEnabled = false;
-  private DAGStatus cachedDagStatus = null;
   Map<String, VertexStatus> cachedVertexStatus = new HashMap<String, VertexStatus>();
 
   private static final long SLEEP_FOR_COMPLETION = 500;
@@ -110,6 +114,28 @@ public DAGClientImpl(ApplicationId appId, String dagId, TezConfiguration conf,
     this.diagnoticsWaitTimeout = conf.getLong(
         TezConfiguration.TEZ_CLIENT_DIAGNOSTICS_WAIT_TIMEOUT_MS,
         TezConfiguration.TEZ_CLIENT_DIAGNOSTICS_WAIT_TIMEOUT_MS_DEFAULT);
+    cachedDAGStatusRef = initCacheDAGRefFromConf(conf);
+  }
+
+  /**
+   * Constructs a new {@link CachedEntity} for {@link DAGStatus}
+   * @param conf TEZ configuration parameters.
+   * @return a caching entry to hold the {@link DAGStatus}.
+   */
+  protected CachedEntity initCacheDAGRefFromConf(TezConfiguration conf) {
+    long clientDAGStatusCacheTimeOut = conf.getLong(
+        TezConfiguration.TEZ_CLIENT_DAG_STATUS_CACHE_TIMEOUT_MINUTES,
+        TezConfiguration.TEZ_CLIENT_DAG_STATUS_CACHE_TIMEOUT_MINUTES_DEFAULT);
+    if (clientDAGStatusCacheTimeOut <= 0) {
+      LOG.error("DAG Status cache timeout interval should be positive. Enforcing default value.");
+      clientDAGStatusCacheTimeOut =
+          TezConfiguration.TEZ_CLIENT_DAG_STATUS_CACHE_TIMEOUT_MINUTES_DEFAULT;
+    }
+    return new CachedEntity<>(TimeUnit.MINUTES, clientDAGStatusCacheTimeOut);
+  }
+
+  protected CachedEntity getCachedDAGStatusRef() {
+    return cachedDAGStatusRef;
   }
 
   @Override
@@ -133,13 +159,11 @@ public DAGStatus getDAGStatus(@Nullable Set<StatusGetOpts> statusOptions,
     }
 
     long startTime = System.currentTimeMillis();
-    boolean refreshStatus;
-    DAGStatus dagStatus;
-    if(cachedDagStatus != null) {
-      dagStatus = cachedDagStatus;
-      refreshStatus = true;
-    } else {
-      // For the first lookup only. After this cachedDagStatus should be populated.
+
+    DAGStatus dagStatus = cachedDAGStatusRef.getValue();
+    boolean refreshStatus = true;
+    if (dagStatus == null) {
+      // the first lookup only or when the cachedDAG has expired
       dagStatus = getDAGStatus(statusOptions);
       refreshStatus = false;
     }
@@ -221,13 +245,14 @@ protected DAGStatus getDAGStatusInternal(@Nullable Set<StatusGetOpts> statusOpti
       final DAGStatus dagStatus = getDAGStatusViaAM(statusOptions, timeout);
 
       if (!dagCompleted) {
-        if (dagStatus != null) {
-          cachedDagStatus = dagStatus;
+        if (dagStatus != null) { // update the cached DAGStatus
+          cachedDAGStatusRef.setValue(dagStatus);
           return dagStatus;
         }
-        if (cachedDagStatus != null) {
+        DAGStatus cachedDAG = cachedDAGStatusRef.getValue();
+        if (cachedDAG != null) {
           // could not get from AM (not reachable/ was killed). return cached status.
-          return cachedDagStatus;
+          return cachedDAG;
         }
       }
 
@@ -253,8 +278,11 @@ protected DAGStatus getDAGStatusInternal(@Nullable Set<StatusGetOpts> statusOpti
 
     // dag completed and Timeline service is either not enabled or does not have completion status
     // return cached status if completion info is present.
-    if (dagCompleted && cachedDagStatus != null && cachedDagStatus.isCompleted()) {
-      return cachedDagStatus;
+    if (dagCompleted) {
+      DAGStatus cachedDag = cachedDAGStatusRef.getValue();
+      if (cachedDag != null && cachedDag.isCompleted()) {
+        return cachedDag;
+      }
     }
 
     // everything else fails rely on RM.