From 860dd8e7a34e4ae46fbeed86175f02296f79416a Mon Sep 17 00:00:00 2001
From: Jonathan Eagles <jeagles@apache.org>
Date: Tue, 7 May 2024 16:48:35 -0500
Subject: [PATCH] TEZ-4562. Fix Tez Job Analyzer after TEZ_DAG_EXTRA_INFO

---
 .../tez/history/parser/ATSFileParser.java     | 21 +++++++++++++++++++
 .../history/parser/datamodel/BaseInfo.java    |  9 ++++++--
 .../history/parser/datamodel/Constants.java   |  1 +
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/tez-plugins/tez-history-parser/src/main/java/org/apache/tez/history/parser/ATSFileParser.java b/tez-plugins/tez-history-parser/src/main/java/org/apache/tez/history/parser/ATSFileParser.java
index caeb406f21..124ca864f6 100644
--- a/tez-plugins/tez-history-parser/src/main/java/org/apache/tez/history/parser/ATSFileParser.java
+++ b/tez-plugins/tez-history-parser/src/main/java/org/apache/tez/history/parser/ATSFileParser.java
@@ -190,7 +190,10 @@ private void parseATSZipFile(File atsFile)
       throws IOException, JSONException, TezException, InterruptedException {
     final ZipFile atsZipFile = new ZipFile(atsFile);
     try {
+      // Need to keep around the dag to handle merging dag extra info
+      JSONObject dagJsonSave = null;
       Enumeration<? extends ZipEntry> zipEntries = atsZipFile.entries();
+
       while (zipEntries.hasMoreElements()) {
         ZipEntry zipEntry = zipEntries.nextElement();
         LOG.debug("Processing " + zipEntry.getName());
@@ -202,6 +205,24 @@ private void parseATSZipFile(File atsFile)
         if (dagJson != null) {
           //TODO: support for multiple dags per ATS file later.
           dagInfo = DagInfo.create(dagJson);
+          dagJsonSave = dagJson;
+        }
+        JSONObject dagExtraInfoJson = jsonObject.optJSONObject(Constants.DAG_EXTRA_INFO);
+        // We need to merge DAG and DAG_EXTRA_INFO together before processing
+        if (dagJsonSave != null) {
+          if (dagExtraInfoJson != null) {
+            JSONObject dagOtherInfo = dagJsonSave.getJSONObject(Constants.OTHER_INFO);
+            JSONObject extraOtherInfo = dagExtraInfoJson.getJSONObject(Constants.OTHER_INFO);
+            @SuppressWarnings("unchecked")
+            Iterator<String> iter = extraOtherInfo.keys();
+            while (iter.hasNext()) {
+              String key = iter.next();
+              dagOtherInfo.put(key, extraOtherInfo.get(key));
+            }
+            // Recreate the dagInfo with the merged DAG and DAG_EXTRA_INFO
+            // Needs created before vertex processing can begin
+            dagInfo = DagInfo.create(dagJsonSave);
+          }
         }
 
         //Process vertex
diff --git a/tez-plugins/tez-history-parser/src/main/java/org/apache/tez/history/parser/datamodel/BaseInfo.java b/tez-plugins/tez-history-parser/src/main/java/org/apache/tez/history/parser/datamodel/BaseInfo.java
index 783f486a15..35c437912a 100644
--- a/tez-plugins/tez-history-parser/src/main/java/org/apache/tez/history/parser/datamodel/BaseInfo.java
+++ b/tez-plugins/tez-history-parser/src/main/java/org/apache/tez/history/parser/datamodel/BaseInfo.java
@@ -55,9 +55,14 @@ public abstract class BaseInfo {
        * already generated events should be parsed correctly, hence this workaround.
        * Will be investigated in the scope of TEZ-4324.
        */
-      countersObj = new JSONObject(otherInfoNode.optString(Constants.COUNTERS));
+      String countersStr = otherInfoNode.optString(Constants.COUNTERS);
+      if (!countersStr.isEmpty()) {
+        countersObj = new JSONObject(otherInfoNode.optString(Constants.COUNTERS));
+      }
+    }
+    if (countersObj != null) {
+      tezCounters = Utils.parseTezCountersFromJSON(countersObj);
     }
-    tezCounters = Utils.parseTezCountersFromJSON(countersObj);
 
     //parse events
     eventList = Lists.newArrayList();
diff --git a/tez-plugins/tez-history-parser/src/main/java/org/apache/tez/history/parser/datamodel/Constants.java b/tez-plugins/tez-history-parser/src/main/java/org/apache/tez/history/parser/datamodel/Constants.java
index dce79e2f78..e6cbd67194 100644
--- a/tez-plugins/tez-history-parser/src/main/java/org/apache/tez/history/parser/datamodel/Constants.java
+++ b/tez-plugins/tez-history-parser/src/main/java/org/apache/tez/history/parser/datamodel/Constants.java
@@ -54,6 +54,7 @@ public class Constants extends ATSConstants {
 
   //constants for ATS data export
   public static final String DAG = "dag";
+  public static final String DAG_EXTRA_INFO = "dag-extra-info";
   public static final String VERTICES = "vertices";
   public static final String TASKS = "tasks";
   public static final String TASK_ATTEMPTS = "task_attempts";