Handle get bytesRead from different thread

Change-Id: I76c6ff84904211e3fae4dcd11772fb7fa5ec503c

Author: jerryshao

core/src/main/scala/org/apache/spark/deploy/SparkHadoopUtil.scala

@@ -21,6 +21,7 @@ import java.io.IOException
 import java.security.PrivilegedExceptionAction
 import java.text.DateFormat
 import java.util.{Arrays, Comparator, Date, Locale}
+import java.util.concurrent.ConcurrentHashMap
 
 import scala.collection.JavaConverters._
 import scala.util.control.NonFatal
@@ -143,14 +144,18 @@ class SparkHadoopUtil extends Logging {
    * Returns a function that can be called to find Hadoop FileSystem bytes read. If
    * getFSBytesReadOnThreadCallback is called from thread r at time t, the returned callback will
    * return the bytes read on r since t.
-   *
-   * @return None if the required method can't be found.
    */
   private[spark] def getFSBytesReadOnThreadCallback(): () => Long = {
-    val threadStats = FileSystem.getAllStatistics.asScala.map(_.getThreadStatistics)
-    val f = () => threadStats.map(_.getBytesRead).sum
-    val baselineBytesRead = f()
-    () => f() - baselineBytesRead
+    val f = () => FileSystem.getAllStatistics.asScala.map(_.getThreadStatistics.getBytesRead).sum
+    val baseline = (Thread.currentThread().getId, f())
+    val bytesReadMap = new ConcurrentHashMap[Long, Long]()
+
+    () => {
+      bytesReadMap.put(Thread.currentThread().getId, f())
+      bytesReadMap.asScala.map { case (k, v) =>
+        v - (if (k == baseline._1) baseline._2 else 0)
+      }.sum
+    }
   }
 
   /**

core/src/main/scala/org/apache/spark/rdd/HadoopRDD.scala

@@ -251,7 +251,13 @@ class HadoopRDD[K, V](
             null
         }
       // Register an on-task-completion callback to close the input stream.
-      context.addTaskCompletionListener{ context => closeIfNeeded() }
+      context.addTaskCompletionListener { context =>
+        // Update the bytes read before closing is to make sure lingering bytesRead statistics in
+        // this thread get correctly added.
+        updateBytesRead()
+        closeIfNeeded()
+      }
+
       private val key: K = if (reader == null) null.asInstanceOf[K] else reader.createKey()
       private val value: V = if (reader == null) null.asInstanceOf[V] else reader.createValue()
 

core/src/main/scala/org/apache/spark/rdd/NewHadoopRDD.scala

@@ -191,7 +191,13 @@ class NewHadoopRDD[K, V](
         }
 
       // Register an on-task-completion callback to close the input stream.
-      context.addTaskCompletionListener(context => close())
+      context.addTaskCompletionListener { context =>
+        // Update the bytesRead before closing is to make sure lingering bytesRead statistics in
+        // this thread get correctly added.
+        updateBytesRead()
+        close()
+      }
+
       private var havePair = false
       private var recordsSinceMetricsUpdate = 0
 

core/src/test/scala/org/apache/spark/metrics/InputOutputMetricsSuite.scala

@@ -319,6 +319,45 @@ class InputOutputMetricsSuite extends SparkFunSuite with SharedSparkContext
     }
     assert(bytesRead >= tmpFile.length())
   }
+
+  test("input metrics with old Hadoop API in different thread") {
+    val bytesRead = runAndReturnBytesRead {
+      sc.textFile(tmpFilePath, 4).mapPartitions { iter =>
+        val buf = new ArrayBuffer[String]()
+        val thread = new Thread() {
+          override def run(): Unit = {
+            iter.flatMap(_.split(" ")).foreach(buf.append(_))
+          }
+        }
+        thread.start()
+        thread.join()
+
+        buf.iterator
+      }.count()
+    }
+    assert(bytesRead != 0)
+    assert(bytesRead >= tmpFile.length())
+  }
+
+  test("input metrics with new Hadoop API in different thread") {
+    val bytesRead = runAndReturnBytesRead {
+      sc.newAPIHadoopFile(tmpFilePath, classOf[NewTextInputFormat], classOf[LongWritable],
+        classOf[Text]).mapPartitions { iter =>
+        val buf = new ArrayBuffer[String]()
+        val thread = new Thread() {
+          override def run(): Unit = {
+            iter.map(_._2.toString).flatMap(_.split(" ")).foreach(buf.append(_))
+          }
+        }
+        thread.start()
+        thread.join()
+
+        buf.iterator
+      }.count()
+    }
+    assert(bytesRead != 0)
+    assert(bytesRead >= tmpFile.length())
+  }
 }
 
 /**