From ba3b68b1fa81d2b860359df06bdf146ed9ad8b3e Mon Sep 17 00:00:00 2001
From: Anthony Hsu <ahsu@linkedin.com>
Date: Tue, 21 May 2019 13:56:12 -0700
Subject: [PATCH] Updated getMaxMemoryBytesUsedForTaskType so that it will
 check if other tasks have metrics rather than stopping on the first task
 without metrics

---
 .../drelephant/tony/util/TonyUtils.java       |  2 +-
 .../drelephant/tony/util/TonyUtilsTest.java   | 33 +++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)
 create mode 100644 test/com/linkedin/drelephant/tony/util/TonyUtilsTest.java

diff --git a/app/com/linkedin/drelephant/tony/util/TonyUtils.java b/app/com/linkedin/drelephant/tony/util/TonyUtils.java
index b7b652a25..55b708d66 100644
--- a/app/com/linkedin/drelephant/tony/util/TonyUtils.java
+++ b/app/com/linkedin/drelephant/tony/util/TonyUtils.java
@@ -35,7 +35,7 @@ public static double getMaxMemoryBytesUsedForTaskType(Map<String, Map<Integer, T
     for (TonyTaskData taskData : taskMap.get(taskType).values()) {
       List<Metric> metrics = taskData.getMetrics();
       if (metrics == null) {
-        break;
+        continue;
       }
       for (Metric metric : metrics) {
         if (metric.getName().equals(Constants.MAX_MEMORY_BYTES)) {
diff --git a/test/com/linkedin/drelephant/tony/util/TonyUtilsTest.java b/test/com/linkedin/drelephant/tony/util/TonyUtilsTest.java
new file mode 100644
index 000000000..116cea7cb
--- /dev/null
+++ b/test/com/linkedin/drelephant/tony/util/TonyUtilsTest.java
@@ -0,0 +1,33 @@
+package com.linkedin.drelephant.tony.util;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.linkedin.drelephant.tony.data.TonyTaskData;
+import com.linkedin.tony.Constants;
+import com.linkedin.tony.events.Metric;
+import java.util.Map;
+import java.util.TreeMap;
+import org.junit.Assert;
+import org.junit.Test;
+
+
+public class TonyUtilsTest {
+  /**
+   * Worker 0 is missing metrics, but worker 1 has metrics; we should use worker 1's
+   * max memory metrics.
+   */
+  @Test
+  public void testGetMaxMemorySomeTasksMissingMetrics() {
+    Map<Integer, TonyTaskData> taskDataMap = new TreeMap<>();
+    TonyTaskData worker0Data = new TonyTaskData("worker", 0);
+    TonyTaskData worker1Data = new TonyTaskData("worker", 1);
+    double worker1MaxMemoryBytes = 123d;
+    worker1Data.setMetrics(ImmutableList.of(new Metric(Constants.MAX_MEMORY_BYTES, worker1MaxMemoryBytes)));
+
+    taskDataMap.put(0, worker0Data);
+    taskDataMap.put(1, worker1Data);
+
+    Assert.assertEquals(worker1MaxMemoryBytes,
+        TonyUtils.getMaxMemoryBytesUsedForTaskType(ImmutableMap.of("worker", taskDataMap), "worker"), 0);
+  }
+}