TEZ-4733: Fix flaky TestHistoryParser.testParserWithSuccessfulJob

tez-plugins/tez-history-parser/src/main/java/org/apache/tez/history/parser/ATSFileParser.java

@@ -24,6 +24,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
 import java.util.Enumeration;
 import java.util.Iterator;
 import java.util.List;
@@ -174,11 +175,30 @@ private void processApplication(JSONObject tezApplicationJson) throws JSONExcept
     }
   }
 
-  private JSONObject readJson(InputStream in) throws IOException, JSONException {
-    //Read entire content to memory
-    final NonSyncByteArrayOutputStream bout = new NonSyncByteArrayOutputStream();
-    IOUtils.copy(in, bout);
-    return new JSONObject(new String(bout.toByteArray(), "UTF-8"));
+  /**
+   * Read a zip entry's payload and parse it as JSON.
+   * Returns null if the payload contains only whitespace (including a zero-length payload) —
+   * callers should skip such entries.
+   *
+   * @throws JSONException if the payload is non-blank but not valid JSON
+   */
+  private JSONObject readJson(InputStream inputStream, String entryName)
+      throws IOException, JSONException {
+    NonSyncByteArrayOutputStream bout = new NonSyncByteArrayOutputStream();
+    IOUtils.copy(inputStream, bout);
+    String text = new String(bout.toByteArray(), StandardCharsets.UTF_8);
+    if (text.trim().isEmpty()) {
+      LOG.warn("Skipping zip entry '{}' - payload is whitespace only (length={})",
+          entryName, text.length());
+      return null;
+    }
+    try {
+      return new JSONObject(text);
+    } catch (JSONException e) {
+      String snippet = text.length() > 200 ? text.substring(0, 200) + "..." : text;
+      throw new JSONException("Failed to parse JSON from zip entry '" + entryName
+          + "' (length=" + text.length() + ", snippet=" + snippet + "): " + e.getMessage());
+    }
   }
 
   /**
@@ -190,14 +210,18 @@ private JSONObject readJson(InputStream in) throws IOException, JSONException {
    */
   private void parseATSZipFile(File atsFile)
       throws IOException, JSONException, TezException, InterruptedException {
-    final ZipFile atsZipFile = new ZipFile(atsFile);
-    try {
+    try (ZipFile atsZipFile = new ZipFile(atsFile)) {
       Enumeration<? extends ZipEntry> zipEntries = atsZipFile.entries();
       while (zipEntries.hasMoreElements()) {
         ZipEntry zipEntry = zipEntries.nextElement();
         LOG.debug("Processing " + zipEntry.getName());
-        InputStream inputStream = atsZipFile.getInputStream(zipEntry);
-        JSONObject jsonObject = readJson(inputStream);
+        JSONObject jsonObject;
+        try (InputStream inputStream = atsZipFile.getInputStream(zipEntry)) {
+          jsonObject = readJson(inputStream, zipEntry.getName());
+        }
+        if (jsonObject == null) {
+          continue;
+        }
 
         //This json can contain dag, vertices, tasks, task_attempts
         JSONObject dagJson = jsonObject.optJSONObject(Constants.DAG);
@@ -230,8 +254,6 @@ private void parseATSZipFile(File atsFile)
           processApplication(tezAppJson);
         }
       }
-    } finally {
-      atsZipFile.close();
     }
   }
 }

tez-plugins/tez-history-parser/src/test/java/org/apache/tez/history/TestHistoryParser.java

@@ -33,6 +33,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 import org.apache.commons.cli.ParseException;
 import org.apache.commons.io.FileUtils;
@@ -118,6 +119,7 @@ public class TestHistoryParser {
   private final static String SUMMATION = "Summation";
   private final static String SIMPLE_HISTORY_DIR = "/tmp/simplehistory/";
   private final static String HISTORY_TXT = "history.txt";
+  private static final long ATS_RETRY_DELAY_MS = 5_000L;
 
   private static Configuration conf = new Configuration();
   private static FileSystem fs;
@@ -209,22 +211,21 @@ public void testParserWithSuccessfulJob() throws Exception {
     String dagId = runWordCount(WordCount.TokenProcessor.class.getName(),
         WordCount.SumProcessor.class.getName(), "WordCount", true);
 
-    //Export the data from ATS
-    String[] args = { "--dagId=" + dagId, "--downloadDir=" + DOWNLOAD_DIR, "--yarnTimelineAddress=" + yarnTimelineAddress };
+    //Retry the ATS export+parse pipeline until the resulting DagInfo actually contains
+    //the expected DAG (two vertices, non-empty vertices/tasks). Under load the AM's async
+    //flush and the timeline server's write path can race the export, leaving empty/partial
+    //entities in the zip.
+    DagInfo dagInfoFromATS = fetchDagInfoFromAtsWithRetry(dagId, 6, 2);
 
-    int result = ATSImportTool.process(args);
-    assertEquals(0, result);
-
-    //Parse ATS data and verify results
-    DagInfo dagInfoFromATS = getDagInfo(dagId);
     verifyDagInfo(dagInfoFromATS, true);
     verifyJobSpecificInfo(dagInfoFromATS);
     checkConfig(dagInfoFromATS);
 
     //Now run with SimpleHistoryLogging
     dagId = runWordCount(WordCount.TokenProcessor.class.getName(),
         WordCount.SumProcessor.class.getName(), "WordCount", false);
-    Thread.sleep(10000); //For all flushes to happen and to avoid half-cooked download.
+
+    waitForHistoryFileReady(dagId, 60_000L);
 
     DagInfo shDagInfo = getDagInfoFromSimpleHistory(dagId);
     verifyDagInfo(shDagInfo, false);
@@ -234,6 +235,83 @@ public void testParserWithSuccessfulJob() throws Exception {
     isDAGEqual(dagInfoFromATS, shDagInfo);
   }
 
+  /**
+   * The ATS write path is async (AM event queue → timeline client → timeline server). Even
+   * after the DAG client reports the job complete, timeline entities may still be in transit,
+   * so an export triggered right after completion can capture a partial snapshot (empty zip
+   * entries, or a DagInfo with fewer vertices / empty task lists). Retry the export+parse
+   * pipeline until {@link #isDagInfoComplete} confirms the snapshot is populated.
+   */
+  private DagInfo fetchDagInfoFromAtsWithRetry(String dagId, int maxAttempts,
+      int expectedNumOfVertices) throws Exception {
+    String[] args = { "--dagId=" + dagId, "--downloadDir=" + DOWNLOAD_DIR,
+        "--yarnTimelineAddress=" + yarnTimelineAddress };
+    Exception lastError = null;
+    DagInfo lastPartial = null;
+    for (int attempt = 0; attempt < maxAttempts; attempt++) {
+      try {
+        // Fresh download every attempt — ATSImportTool overwrites the zip.
+        int result = ATSImportTool.process(args);
+        assertEquals(0, result);
+        DagInfo info = getDagInfo(dagId);
+        if (isDagInfoComplete(info, expectedNumOfVertices)) {
+          return info;
+        }
+        lastPartial = info;
+      } catch (Exception e) {
+        lastError = e;
+      }
+      if (attempt < maxAttempts - 1) {
+        Thread.sleep(ATS_RETRY_DELAY_MS);
+      }
+    }
+    fail("Could not fetch a complete DagInfo for " + dagId + " after " + maxAttempts
+        + " attempts (lastError=" + lastError + ", lastPartial="
+        + (lastPartial == null ? "null"
+            : "vertices=" + lastPartial.getVertices().size()
+            + ", tasks=" + (lastPartial.getVertices().isEmpty() ? 0
+                : lastPartial.getVertices().iterator().next().getTasks().size()))
+        + ")");
+    return null;
+  }
+
+  /**
+   * ATS parsing can succeed on a partially-written zip: the reader returns a DagInfo with
+   * fewer vertices than the DAG actually has, or vertices whose task/attempt lists are still
+   * empty. That's the race we're guarding against — a "successful" parse is not proof the
+   * export was complete. We know the expected vertex count from the DAG under test, so require
+   * it explicitly and require every vertex to have at least one task with at least one attempt.
+   */
+  private static boolean isDagInfoComplete(DagInfo info, int expectedNumOfVertices) {
+    return info != null
+        && info.getVertices().size() >= expectedNumOfVertices
+        && info.getVertices().stream().allMatch(v ->
+            !v.getTasks().isEmpty()
+                && v.getTasks().stream().allMatch(t -> !t.getTaskAttempts().isEmpty()));
+  }
+
+  private void waitForHistoryFileReady(String dagId, long timeoutMs) throws Exception {
+    TezDAGID tezDAGID = TezDAGID.fromString(dagId);
+    ApplicationAttemptId applicationAttemptId = ApplicationAttemptId.newInstance(tezDAGID.getApplicationId(), 1);
+    Path historyPath = new Path(conf.get("fs.defaultFS")
+        + SIMPLE_HISTORY_DIR + HISTORY_TXT + "." + applicationAttemptId);
+    FileSystem hfs = historyPath.getFileSystem(conf);
+    long deadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(timeoutMs);
+    long lastLen = -1L;
+    while (System.nanoTime() < deadlineNanos) {
+      if (hfs.exists(historyPath)) {
+        long len = hfs.getFileStatus(historyPath).getLen();
+        if (len > 0 && len == lastLen) {
+          return;
+        }
+        lastLen = len;
+      }
+      Thread.sleep(500);
+    }
+    fail("Timed out waiting for SimpleHistory file " + historyPath
+        + " to be ready within " + timeoutMs + "ms (lastLen=" + lastLen + ")");
+  }
+
   private DagInfo getDagInfoFromSimpleHistory(String dagId) throws TezException, IOException {
     TezDAGID tezDAGID = TezDAGID.fromString(dagId);
     ApplicationAttemptId applicationAttemptId = ApplicationAttemptId.newInstance(tezDAGID

tez-plugins/tez-history-parser/src/test/java/org/apache/tez/history/parser/TestATSFileParser.java

@@ -0,0 +1,109 @@
+/*
+ * 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.history.parser;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
+import java.util.Collections;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipOutputStream;
+
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.history.parser.datamodel.DagInfo;
+
+import org.codehaus.jettison.json.JSONArray;
+import org.codehaus.jettison.json.JSONException;
+import org.codehaus.jettison.json.JSONObject;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+public class TestATSFileParser {
+
+  private static final String DAG_ID = "dag_1234567890_0001_1";
+
+  private static JSONObject minimalDagJson() throws JSONException {
+    JSONObject dag = new JSONObject();
+    dag.put("entityId", DAG_ID);
+    dag.put("entityType", "TEZ_DAG_ID");
+    JSONObject otherInfo = new JSONObject();
+    otherInfo.put("startTime", 1L);
+    otherInfo.put("endTime", 2L);
+    otherInfo.put("status", "SUCCEEDED");
+    otherInfo.put("counters", new JSONObject().put("counterGroups", new JSONArray()));
+    dag.put("otherInfo", otherInfo);
+    return dag;
+  }
+
+  private static File writeZip(Path dir, String name, ZipContent... entries) throws Exception {
+    File zip = dir.resolve(name).toFile();
+    try (FileOutputStream fos = new FileOutputStream(zip);
+         ZipOutputStream zos = new ZipOutputStream(fos)) {
+      for (ZipContent entry : entries) {
+        zos.putNextEntry(new ZipEntry(entry.name()));
+        zos.write(entry.payload().getBytes(StandardCharsets.UTF_8));
+        zos.closeEntry();
+      }
+    }
+    return zip;
+  }
+
+  @Test
+  public void parserSkipsEmptyZipEntryAndParsesRemaining(@TempDir Path tmp) throws Exception {
+    JSONObject dagRoot = new JSONObject().put("dag", minimalDagJson());
+
+    File zip = writeZip(tmp, "empty-then-good.zip",
+        new ZipContent("empty-part.json", ""),
+        new ZipContent("whitespace-part.json", "   \n\t  "),
+        new ZipContent(DAG_ID, dagRoot.toString()));
+
+    ATSFileParser parser = new ATSFileParser(Collections.singletonList(zip));
+    DagInfo info = parser.getDAGData(DAG_ID);
+
+    assertNotNull(info, "Parser should return DagInfo even when some entries are empty");
+    assertEquals(DAG_ID, info.getDagId());
+    assertEquals("SUCCEEDED", info.getStatus());
+  }
+
+  @Test
+  public void parserReportsOffendingEntryOnMalformedJson(@TempDir Path tmp) throws Exception {
+    // Simulates the timeline server returning an HTML error page instead of JSON.
+    File zip = writeZip(tmp, "malformed.zip",
+        new ZipContent(DAG_ID, "<html><body>internal error</body></html>"));
+
+    ATSFileParser parser = new ATSFileParser(Collections.singletonList(zip));
+    TezException thrown = assertThrows(TezException.class, () -> parser.getDAGData(DAG_ID));
+
+    Throwable cause = thrown.getCause();
+    assertNotNull(cause);
+    String msg = cause.getMessage();
+    assertNotNull(msg);
+    assertTrue(msg.contains(DAG_ID), "Error should name the offending zip entry, got: " + msg);
+    assertTrue(msg.contains("<html>"), "Error should include a snippet of the offending payload, got: " + msg);
+  }
+
+  private record ZipContent(String name, String payload) {
+  }
+}