TEZ-4466: Measure exact stream read time while fetching

tez-runtime-library/src/main/java/org/apache/tez/http/MeasuredDataInputStream.java

@@ -0,0 +1,79 @@
+/*
+ * 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.http;
+
+import java.io.DataInputStream;
+import java.io.FilterInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.concurrent.TimeUnit;
+
+public class MeasuredDataInputStream extends DataInputStream {
+
+  private final MeasuredInputStream measuredIn;
+
+  private MeasuredDataInputStream(MeasuredInputStream measuredIn) {
+    super(measuredIn);
+    this.measuredIn = measuredIn;
+  }
+
+  public MeasuredDataInputStream(InputStream in) {
+    this(new MeasuredInputStream(in));
+  }
+
+  public long getElapsedTimeMs() {
+    return measuredIn.getElapsedTimeMs();
+  }
+
+  private static class MeasuredInputStream extends FilterInputStream {
+    private long elapsedTimeNanos = 0;
+
+    MeasuredInputStream(InputStream in) {
+      super(in);
+    }
+
+    @Override
+    public int read() throws IOException {
+      long start = System.nanoTime();
+      int ret = super.read();
+      elapsedTimeNanos += (System.nanoTime() - start);
+      return ret;
+    }
+
+    @Override
+    public int read(byte[] b) throws IOException {
+      long start = System.nanoTime();
+      int ret = super.read(b);
+      elapsedTimeNanos += (System.nanoTime() - start);
+      return ret;
+    }
+
+    @Override
+    public int read(byte[] b, int off, int len) throws IOException {
+      long start = System.nanoTime();
+      int ret = super.read(b, off, len);
+      elapsedTimeNanos += (System.nanoTime() - start);
+      return ret;
+    }
+
+    public long getElapsedTimeMs() {
+      return TimeUnit.NANOSECONDS.toMillis(elapsedTimeNanos);
+    }
+  }
+}

tez-runtime-library/src/main/java/org/apache/tez/runtime/library/api/TezRuntimeConfiguration.java

@@ -416,6 +416,13 @@ private TezRuntimeConfiguration() {}
   public static final float TEZ_RUNTIME_SHUFFLE_FETCH_BUFFER_PERCENT_DEFAULT =
       0.90f;
 
+  /**
+   * Enables measuring network IO time in shuffle fetchers.
+   */
+  @ConfigurationProperty(type = "boolean")
+  public static final String TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME = TEZ_RUNTIME_PREFIX + "shuffle.measure.io.time";
+  public static final boolean TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME_DEFAULT = false;
+
   /**
    * Enables fetch failures by a configuration. Should be used for testing only.
    */
@@ -639,6 +646,7 @@ private TezRuntimeConfiguration() {}
     TEZ_RUNTIME_KEYS.add(TEZ_RUNTIME_SHUFFLE_ENABLE_SSL);
     TEZ_RUNTIME_KEYS.add(TEZ_RUNTIME_SHUFFLE_FETCH_VERIFY_DISK_CHECKSUM);
     TEZ_RUNTIME_KEYS.add(TEZ_RUNTIME_SHUFFLE_FETCH_BUFFER_PERCENT);
+    TEZ_RUNTIME_KEYS.add(TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME);
     TEZ_RUNTIME_KEYS.add(TEZ_RUNTIME_SHUFFLE_MEMORY_LIMIT_PERCENT);
     TEZ_RUNTIME_KEYS.add(TEZ_RUNTIME_SHUFFLE_MERGE_PERCENT);
     TEZ_RUNTIME_KEYS.add(TEZ_RUNTIME_SHUFFLE_MEMTOMEM_SEGMENTS);

tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java

@@ -55,6 +55,7 @@
 import org.apache.tez.dag.api.TezUncheckedException;
 import org.apache.tez.http.BaseHttpConnection;
 import org.apache.tez.http.HttpConnectionParams;
+import org.apache.tez.http.MeasuredDataInputStream;
 import org.apache.tez.runtime.api.InputContext;
 import org.apache.tez.runtime.library.api.TezRuntimeConfiguration;
 import org.apache.tez.runtime.library.common.CompositeInputAttemptIdentifier;
@@ -565,6 +566,10 @@ private HostFetchResult setupConnection(Collection<InputAttemptIdentifier> attem
   protected void setupConnectionInternal(String host, Collection<InputAttemptIdentifier> attempts)
       throws IOException, InterruptedException {
     input = httpConnection.getInputStream();
+    if (conf.getBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME,
+        TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME_DEFAULT)) {
+      input = new MeasuredDataInputStream(input);
+    }
     httpConnection.validate();
   }
 

tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java

@@ -44,7 +44,9 @@
 import org.apache.tez.common.security.JobTokenSecretManager;
 import org.apache.tez.http.BaseHttpConnection;
 import org.apache.tez.http.HttpConnectionParams;
+import org.apache.tez.http.MeasuredDataInputStream;
 import org.apache.tez.runtime.api.InputContext;
+import org.apache.tez.runtime.library.api.TezRuntimeConfiguration;
 import org.apache.tez.runtime.library.common.Constants;
 import org.apache.tez.runtime.library.common.InputAttemptIdentifier;
 import org.apache.tez.runtime.library.common.shuffle.InputAttemptFetchFailure;
@@ -392,6 +394,10 @@ boolean setupConnection(MapHost host, Collection<InputAttemptIdentifier> attempt
   protected void setupConnectionInternal(MapHost host, Collection<InputAttemptIdentifier> attempts)
       throws IOException, InterruptedException {
     input = httpConnection.getInputStream();
+    if (conf.getBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME,
+        TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME_DEFAULT)) {
+      input = new MeasuredDataInputStream(input);
+    }
     httpConnection.validate();
   }