TEZ-4227 Introduce convenient methods in TezID subclasses

Change-Id: I6cabfa75e9b6b62e41ba8c2cc5e3d2d1a8a49102

Author: ghanko

tez-common/src/main/java/org/apache/tez/dag/records/DAGIDHolder.java

@@ -0,0 +1,11 @@
+package org.apache.tez.dag.records;
+
+import org.apache.hadoop.yarn.api.records.ApplicationId;
+
+public interface DAGIDHolder {
+    TezDAGID getDAGId();
+
+    default ApplicationId getApplicationId() {
+        return getDAGId().getApplicationId();
+    }
+}

tez-common/src/main/java/org/apache/tez/dag/records/TaskAttemptIDHolder.java

@@ -0,0 +1,10 @@
+package org.apache.tez.dag.records;
+
+public interface TaskAttemptIDHolder extends TaskIDHolder {
+    TezTaskAttemptID getTaskAttemptID();
+
+    @Override
+    default TezTaskID getTaskID() {
+        return getTaskAttemptID().getTaskID();
+    }
+}

tez-common/src/main/java/org/apache/tez/dag/records/TaskIDHolder.java

@@ -0,0 +1,10 @@
+package org.apache.tez.dag.records;
+
+public interface TaskIDHolder extends VertexIDHolder {
+    TezTaskID getTaskID();
+
+    @Override
+    default TezVertexID getVertexID() {
+        return getTaskID().getVertexID();
+    }
+}

tez-common/src/main/java/org/apache/tez/dag/records/TezTaskAttemptID.java

@@ -45,7 +45,7 @@
  */
 @InterfaceAudience.Public
 @InterfaceStability.Stable
-public class TezTaskAttemptID extends TezID {
+public class TezTaskAttemptID extends TezID implements TaskIDHolder {
   public static final String ATTEMPT = "attempt";
   private TezTaskID taskId;
 
@@ -73,6 +73,7 @@ private TezTaskAttemptID(TezTaskID taskId, int id) {
   }
 
   /** Returns the {@link TezTaskID} object that this task attempt belongs to */
+  @Override
   public TezTaskID getTaskID() {
     return taskId;
   }

tez-common/src/main/java/org/apache/tez/dag/records/TezTaskID.java

@@ -41,7 +41,7 @@
  */
 @InterfaceAudience.Public
 @InterfaceStability.Stable
-public class TezTaskID extends TezID {
+public class TezTaskID extends TezID implements VertexIDHolder {
   public static final String TASK = "task";
   private final int serializingHash;
 
@@ -79,6 +79,7 @@ public int getSerializingHash() {
   }
 
   /** Returns the {@link TezVertexID} object that this task belongs to */
+  @Override
   public TezVertexID getVertexID() {
     return vertexId;
   }

tez-common/src/main/java/org/apache/tez/dag/records/TezVertexID.java

@@ -44,7 +44,7 @@
  */
 @InterfaceAudience.Public
 @InterfaceStability.Stable
-public class TezVertexID extends TezID {
+public class TezVertexID extends TezID implements DAGIDHolder {
   public static final String VERTEX = "vertex";
   static final ThreadLocal<FastNumberFormat> tezVertexIdFormat = new ThreadLocal<FastNumberFormat>() {
 
@@ -79,6 +79,7 @@ private TezVertexID(TezDAGID dagId, int id) {
   }
 
   /** Returns the {@link TezDAGID} object that this tip belongs to */
+  @Override
   public TezDAGID getDAGId() {
     return dagId;
   }
@@ -158,5 +159,4 @@ public static TezVertexID fromString(String vertexIdStr) {
     }
     return null;
   }
-
 }

tez-common/src/main/java/org/apache/tez/dag/records/VertexIDHolder.java

@@ -0,0 +1,10 @@
+package org.apache.tez.dag.records;
+
+public interface VertexIDHolder extends DAGIDHolder {
+    TezVertexID getVertexID();
+
+    @Override
+    default TezDAGID getDAGId() {
+        return getVertexID().getDAGId();
+    }
+}

tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java

@@ -2187,12 +2187,12 @@ private class TaskEventDispatcher implements EventHandler<TaskEvent> {
     public void handle(TaskEvent event) {
       DAG dag = context.getCurrentDAG();
       int eventDagIndex =
-          event.getTaskID().getVertexID().getDAGId().getId();
+          event.getDAGId().getId();
       if (dag == null || eventDagIndex != dag.getID().getId()) {
         return; // event not relevant any more
       }
       Task task =
-          dag.getVertex(event.getTaskID().getVertexID()).
+          dag.getVertex(event.getVertexID()).
               getTask(event.getTaskID());
       ((EventHandler<TaskEvent>)task).handle(event);
     }
@@ -2217,13 +2217,13 @@ private class TaskAttemptEventDispatcher
     public void handle(TaskAttemptEvent event) {
       DAG dag = context.getCurrentDAG();
       int eventDagIndex =
-          event.getTaskAttemptID().getTaskID().getVertexID().getDAGId().getId();
+          event.getDAGId().getId();
       if (dag == null || eventDagIndex != dag.getID().getId()) {
         return; // event not relevant any more
       }
       Task task =
-          dag.getVertex(event.getTaskAttemptID().getTaskID().getVertexID()).
-              getTask(event.getTaskAttemptID().getTaskID());
+          dag.getVertex(event.getVertexID()).
+              getTask(event.getTaskID());
       TaskAttempt attempt = task.getAttempt(event.getTaskAttemptID());
       ((EventHandler<TaskAttemptEvent>) attempt).handle(event);
     }
@@ -2236,13 +2236,13 @@ private class VertexEventDispatcher
     public void handle(VertexEvent event) {
       DAG dag = context.getCurrentDAG();
       int eventDagIndex =
-          event.getVertexId().getDAGId().getId();
+          event.getDAGId().getId();
       if (dag == null || eventDagIndex != dag.getID().getId()) {
         return; // event not relevant any more
       }
 
       Vertex vertex =
-          dag.getVertex(event.getVertexId());
+          dag.getVertex(event.getVertexID());
       ((EventHandler<VertexEvent>) vertex).handle(event);
     }
   }

tez-dag/src/main/java/org/apache/tez/dag/app/RecoveryParser.java

@@ -846,19 +846,19 @@ public DAGRecoveryData parseRecoveryData() throws IOException {
           case TASK_STARTED:
           {
             TaskStartedEvent taskStartedEvent = (TaskStartedEvent) event;
-            VertexRecoveryData vertexRecoveryData = recoveredDAGData.vertexRecoveryDataMap.get(taskStartedEvent.getTaskID().getVertexID());
+            VertexRecoveryData vertexRecoveryData = recoveredDAGData.vertexRecoveryDataMap.get(taskStartedEvent.getVertexID());
             Preconditions.checkArgument(vertexRecoveryData != null,
-                "Invalid TaskStartedEvent, its vertex does not exist:" + taskStartedEvent.getTaskID().getVertexID());
+                "Invalid TaskStartedEvent, its vertex does not exist:" + taskStartedEvent.getVertexID());
             TaskRecoveryData taskRecoveryData = vertexRecoveryData.maybeCreateTaskRecoveryData(taskStartedEvent.getTaskID());
             taskRecoveryData.taskStartedEvent = taskStartedEvent;
             break;
           }
           case TASK_FINISHED:
           {
             TaskFinishedEvent taskFinishedEvent = (TaskFinishedEvent) event;
-            VertexRecoveryData vertexRecoveryData = recoveredDAGData.vertexRecoveryDataMap.get(taskFinishedEvent.getTaskID().getVertexID());
+            VertexRecoveryData vertexRecoveryData = recoveredDAGData.vertexRecoveryDataMap.get(taskFinishedEvent.getVertexID());
             Preconditions.checkArgument(vertexRecoveryData != null,
-                "Invalid TaskFinishedEvent, its vertex does not exist:" + taskFinishedEvent.getTaskID().getVertexID());
+                "Invalid TaskFinishedEvent, its vertex does not exist:" + taskFinishedEvent.getVertexID());
             TaskRecoveryData taskRecoveryData = vertexRecoveryData.maybeCreateTaskRecoveryData(taskFinishedEvent.getTaskID());
             taskRecoveryData.taskFinishedEvent = taskFinishedEvent;
             break;
@@ -867,7 +867,7 @@ public DAGRecoveryData parseRecoveryData() throws IOException {
           {
             TaskAttemptStartedEvent taStartedEvent = (TaskAttemptStartedEvent)event;
             VertexRecoveryData vertexRecoveryData = recoveredDAGData.vertexRecoveryDataMap.get(
-                taStartedEvent.getTaskAttemptID().getTaskID().getVertexID());
+                taStartedEvent.getVertexID());
             Preconditions.checkArgument(vertexRecoveryData != null,
                 "Invalid TaskAttemptStartedEvent, its vertexId does not exist, taId=" + taStartedEvent.getTaskAttemptID());
             TaskRecoveryData taskRecoveryData = vertexRecoveryData.taskRecoveryDataMap
@@ -882,7 +882,7 @@ public DAGRecoveryData parseRecoveryData() throws IOException {
           {
             TaskAttemptFinishedEvent taFinishedEvent = (TaskAttemptFinishedEvent)event;
             VertexRecoveryData vertexRecoveryData = recoveredDAGData.vertexRecoveryDataMap.get(
-                taFinishedEvent.getTaskAttemptID().getTaskID().getVertexID());
+                taFinishedEvent.getVertexID());
             Preconditions.checkArgument(vertexRecoveryData != null,
                 "Invalid TaskAttemtFinishedEvent, its vertexId does not exist, taId=" + taFinishedEvent.getTaskAttemptID());
             TaskRecoveryData taskRecoveryData = vertexRecoveryData.taskRecoveryDataMap

tez-dag/src/main/java/org/apache/tez/dag/app/TaskCommunicatorManager.java

@@ -350,14 +350,14 @@ public TaskHeartbeatResponse heartbeat(TaskHeartbeatRequest request)
         }
       }
       if (!eventsForVertex.isEmpty()) {
-        TezVertexID vertexId = taskAttemptID.getTaskID().getVertexID();
+        TezVertexID vertexId = taskAttemptID.getVertexID();
         sendEvent(
             new VertexEventRouteEvent(vertexId, Collections.unmodifiableList(eventsForVertex)));
       }
       taskHeartbeatHandler.pinged(taskAttemptID);
       eventInfo = context
           .getCurrentDAG()
-          .getVertex(taskAttemptID.getTaskID().getVertexID())
+          .getVertex(taskAttemptID.getVertexID())
           .getTaskAttemptTezEvents(taskAttemptID, request.getStartIndex(), request.getPreRoutedStartIndex(),
               request.getMaxEvents());
     }
@@ -442,7 +442,7 @@ public boolean canCommit(TezTaskAttemptID taskAttemptId) throws IOException {
 
     DAG job = context.getCurrentDAG();
     Task task =
-        job.getVertex(taskAttemptId.getTaskID().getVertexID()).
+        job.getVertex(taskAttemptId.getVertexID()).
             getTask(taskAttemptId.getTaskID());
     return task.canCommit(taskAttemptId);
   }