TEZ-4250: Optimise TaskImpl::getCounters.

Change-Id: I7db4db3a4d2cb70f2e1d96fbaa2f853524fbce8a
apache · Jun 21, 2023 · c539e82 · c539e82
1 parent d20f334
commit c539e82
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 6 deletions.
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/TaskAttempt.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/TaskAttempt.java
@@ -21,6 +21,7 @@
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
 
+import com.google.common.annotations.VisibleForTesting;
 import org.apache.hadoop.yarn.api.records.Container;
 import org.apache.hadoop.yarn.api.records.ContainerId;
 import org.apache.hadoop.yarn.api.records.NodeId;
@@ -63,13 +64,20 @@ public void setLocalityCounter(DAGCounter localityCounter) {
         }
       }
     }
+
+    @VisibleForTesting
+    public void setCounters(TezCounters counters) {
+      this.counters = counters;
+    }
   }
 
   Task getTask();
   TaskAttemptReport getReport();
   List<String> getDiagnostics();
   TaskAttemptTerminationCause getTerminationCause();
   TezCounters getCounters();
+  @VisibleForTesting
+  void setCounters(TezCounters counters);
   float getProgress();
   TaskAttemptState getState();
   TaskAttemptState getStateNoLock();

diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskAttemptImpl.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskAttemptImpl.java
@@ -654,6 +654,17 @@ public TezCounters getCounters() {
       readLock.unlock();
     }
   }
+
+  @VisibleForTesting
+  @Override
+  public void setCounters(TezCounters counters) {
+    writeLock.lock();
+    try {
+      reportedStatus.setCounters(counters);
+    } finally {
+      writeLock.unlock();
+    }
+  }
 
   TaskStatistics getStatistics() {
     return this.statistics;

diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskImpl.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskImpl.java
@@ -465,15 +465,19 @@ public TaskReport getReport() {
 
   @Override
   public TezCounters getCounters() {
-    TezCounters counters = new TezCounters();
-    counters.incrAllCounters(this.counters);
+    TezCounters counters = null;
+    if (getVertex().isSpeculationEnabled()) {
+      counters = new TezCounters();
+      counters.incrAllCounters(this.counters);
+    }
     readLock.lock();
     try {
       TaskAttempt bestAttempt = selectBestAttempt();
-      if (bestAttempt != null) {
+      if (bestAttempt != null && counters != null) {
         counters.incrAllCounters(bestAttempt.getCounters());
+        return counters;
       }
-      return counters;
+      return (bestAttempt != null) ? bestAttempt.getCounters() : TaskAttemptImpl.EMPTY_COUNTERS;
     } finally {
       readLock.unlock();
     }
@@ -1522,10 +1526,9 @@ public void transition(TaskImpl task, TaskEvent event) {
   void setCounters(TezCounters counters) {
     try {
       writeLock.lock();
-      this.counters = counters;
+      selectBestAttempt().setCounters(counters);
     } finally {
       writeLock.unlock();
     }
   }
-
 }