TEZ-4466: Measure exact stream read time while fetching

tez-api/src/main/java/org/apache/tez/common/counters/TaskCounter.java

@@ -189,6 +189,21 @@ public enum TaskCounter {
    */
   SHUFFLE_BYTES_DISK_DIRECT,
 
+  /**
+   * Time spent waiting on network I/O during shuffle. Represented in milliseconds.
+   * Only populated if "tez.runtime.shuffle.measure.io.time" is enabled.
+   * Warning: enabling this counter might have a slight overhead.
+   */
+  SHUFFLE_IO_STREAM_TIME_MILLISECONDS,
+
+  /**
+   * Actual bytes read from the network I/O streams during shuffle.
+   * Should be equal to SHUFFLE_BYTES if no issues occur.
+   * Only populated if "tez.runtime.shuffle.measure.io.time" is enabled.
+   * Warning: enabling this counter might have a slight overhead.
+   */
+  SHUFFLE_IO_STREAM_BYTES,
+
   /**
    * Number of Memory to Disk merges performed during sort-merge.
    * Used by ShuffledMergedInput

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

@@ -0,0 +1,97 @@
+/*
+ * 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();
+  }
+
+  public long getBytesRead() {
+    return measuredIn.getBytesRead();
+  }
+
+  private static class MeasuredInputStream extends FilterInputStream {
+    private long elapsedTimeNanos = 0;
+    private long bytesRead = 0;
+
+    MeasuredInputStream(InputStream in) {
+      super(in);
+    }
+
+    @Override
+    public int read() throws IOException {
+      long start = System.nanoTime();
+      int val = super.read();
+      elapsedTimeNanos += (System.nanoTime() - start);
+      if (val != -1) {
+        bytesRead += 1;
+      }
+      return val;
+    }
+
+    @Override
+    public int read(byte[] b) throws IOException {
+      long start = System.nanoTime();
+      int bytes = super.read(b);
+      elapsedTimeNanos += (System.nanoTime() - start);
+      if (bytes > 0) {
+        bytesRead += bytes;
+      }
+      return bytes;
+    }
+
+    @Override
+    public int read(byte[] b, int off, int len) throws IOException {
+      long start = System.nanoTime();
+      int bytes = super.read(b, off, len);
+      elapsedTimeNanos += (System.nanoTime() - start);
+      if (bytes > 0) {
+        bytesRead += bytes;
+      }
+      return bytes;
+    }
+
+    public long getElapsedTimeMs() {
+      return TimeUnit.NANOSECONDS.toMillis(elapsedTimeNanos);
+    }
+
+    public long getBytesRead() {
+      return bytesRead;
+    }
+  }
+}

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

@@ -416,6 +416,15 @@ private TezRuntimeConfiguration() {}
   public static final float TEZ_RUNTIME_SHUFFLE_FETCH_BUFFER_PERCENT_DEFAULT =
       0.90f;
 
+  /**
+   * Enables measuring network IO time in shuffle fetchers.
+   * Warning: enabling this counter might have a slight overhead.
+   * See {@link org.apache.tez.http.MeasuredDataInputStream} for more details.
+   */
+  @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 +648,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

@@ -51,10 +51,13 @@
 import org.apache.tez.common.CallableWithNdc;
 import org.apache.tez.common.Preconditions;
 import org.apache.tez.common.TezUtilsInternal;
+import org.apache.tez.common.counters.TaskCounter;
+import org.apache.tez.common.counters.TezCounter;
 import org.apache.tez.common.security.JobTokenSecretManager;
 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;
@@ -177,6 +180,8 @@ public String getHost() {
 
   BaseHttpConnection httpConnection;
   private HttpConnectionParams httpConnectionParams;
+  private final TezCounter ioTimeCounter;
+  private final TezCounter ioBytesCounter;
 
   private final boolean localDiskFetchEnabled;
   private final boolean sharedFetchEnabled;
@@ -219,6 +224,13 @@ protected Fetcher(FetcherCallback fetcherCallback, HttpConnectionParams params,
     this.localDiskFetchEnabled = localDiskFetchEnabled;
     this.sharedFetchEnabled = sharedFetchEnabled;
 
+    boolean measureTime = conf.getBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME,
+        TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME_DEFAULT);
+    this.ioTimeCounter =
+        measureTime ? inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_STREAM_TIME_MILLISECONDS) : null;
+    this.ioBytesCounter =
+        measureTime ? inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_STREAM_BYTES) : null;
+
     this.fetcherIdentifier = fetcherIdGen.getAndIncrement();
 
     String sourceDestNameTrimmed = TezUtilsInternal.cleanVertexName(inputContext.getSourceVertexName()) + " -> "
@@ -565,6 +577,9 @@ private HostFetchResult setupConnection(Collection<InputAttemptIdentifier> attem
   protected void setupConnectionInternal(String host, Collection<InputAttemptIdentifier> attempts)
       throws IOException, InterruptedException {
     input = httpConnection.getInputStream();
+    if (ioTimeCounter != null) {
+      input = new MeasuredDataInputStream(input);
+    }
     httpConnection.validate();
   }
 
@@ -813,6 +828,10 @@ private void shutdownInternal(boolean disconnect) {
     synchronized (isShutDown) {
       try {
         if (httpConnection != null) {
+          if (input instanceof MeasuredDataInputStream && ioTimeCounter != null) {
+            ioTimeCounter.increment(((MeasuredDataInputStream) input).getElapsedTimeMs());
+            ioBytesCounter.increment(((MeasuredDataInputStream) input).getBytesRead());
+          }
           httpConnection.cleanup(disconnect);
         }
       } catch (IOException e) {

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

@@ -40,11 +40,14 @@
 import org.apache.tez.common.CallableWithNdc;
 import org.apache.tez.common.TezRuntimeFrameworkConfigs;
 import org.apache.tez.common.TezUtilsInternal;
+import org.apache.tez.common.counters.TaskCounter;
 import org.apache.tez.common.counters.TezCounter;
 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;
@@ -75,6 +78,8 @@ class FetcherOrderedGrouped extends CallableWithNdc<Void> {
   private final TezCounter wrongLengthErrs;
   private final TezCounter badIdErrs;
   private final TezCounter wrongReduceErrs;
+  private final TezCounter ioTimeCounter;
+  private final TezCounter ioBytesCounter;
   private final FetchedInputAllocatorOrderedGrouped allocator;
   private final ShuffleScheduler scheduler;
   private final ExceptionReporter exceptionReporter;
@@ -151,6 +156,12 @@ public FetcherOrderedGrouped(HttpConnectionParams httpConnectionParams,
     this.badIdErrs = badIdErrsCounter;
     this.connectionErrs = connectionErrsCounter;
     this.wrongReduceErrs = wrongReduceErrsCounter;
+    boolean measureTime = conf.getBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME,
+        TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME_DEFAULT);
+    this.ioTimeCounter =
+        measureTime ? inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_STREAM_TIME_MILLISECONDS) : null;
+    this.ioBytesCounter =
+        measureTime ? inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_STREAM_BYTES) : null;
     this.applicationId = inputContext.getApplicationId().toString();
     this.dagId = inputContext.getDagIdentifier();
 
@@ -227,6 +238,10 @@ private void cleanupCurrentConnection(boolean disconnect) {
     synchronized (cleanupLock) {
       try {
         if (httpConnection != null) {
+          if (input instanceof MeasuredDataInputStream && ioTimeCounter != null) {
+            ioTimeCounter.increment(((MeasuredDataInputStream) input).getElapsedTimeMs());
+            ioBytesCounter.increment(((MeasuredDataInputStream) input).getBytesRead());
+          }
           httpConnection.cleanup(disconnect);
           httpConnection = null;
         }
@@ -392,6 +407,9 @@ boolean setupConnection(MapHost host, Collection<InputAttemptIdentifier> attempt
   protected void setupConnectionInternal(MapHost host, Collection<InputAttemptIdentifier> attempts)
       throws IOException, InterruptedException {
     input = httpConnection.getInputStream();
+    if (ioTimeCounter != null) {
+      input = new MeasuredDataInputStream(input);
+    }
     httpConnection.validate();
   }
 

tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/TestFetcher.java

@@ -333,10 +333,6 @@ public void testShuffleHandlerDiskErrorUnordered()
       throws Exception {
     Configuration conf = new Configuration();
 
-    InputContext inputContext = mock(InputContext.class);
-    doReturn(new TezCounters()).when(inputContext).getCounters();
-    doReturn("vertex").when(inputContext).getSourceVertexName();
-
     Fetcher.FetcherBuilder builder = new Fetcher.FetcherBuilder(mock(ShuffleManager.class), null,
         null, createMockInputContext(), null, conf, true, HOST, PORT,
         false, true, false);
@@ -361,6 +357,7 @@ private InputContext createMockInputContext() {
     doReturn(1).when(inputContext).getDagIdentifier();
     doReturn("sourceVertex").when(inputContext).getSourceVertexName();
     doReturn("taskVertex").when(inputContext).getTaskVertexName();
+    doReturn(new TezCounters()).when(inputContext).getCounters();
 
     return inputContext;
   }

tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/TestFetcher.java

@@ -60,6 +60,7 @@
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.fs.RawLocalFileSystem;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
+import org.apache.tez.common.counters.TaskCounter;
 import org.apache.tez.common.counters.TezCounter;
 import org.apache.tez.common.counters.TezCounters;
 import org.apache.tez.common.security.JobTokenSecretManager;
@@ -798,4 +799,67 @@ private InputContext createMockInputContext() {
 
     return inputContext;
   }
+
+  @Test
+  public void testShuffleMeasureIOTime() throws Exception {
+    Configuration conf = new TezConfiguration();
+    conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME, true);
+
+    ShuffleScheduler scheduler = mock(ShuffleScheduler.class);
+    MergeManager merger = mock(MergeManager.class);
+    Shuffle shuffle = mock(Shuffle.class);
+
+    final MapHost host = new MapHost(HOST, PORT, 1, 1);
+    InputContext inputContext = createMockInputContext();
+    FetcherOrderedGrouped mockFetcher =
+        new FetcherOrderedGrouped(null, scheduler, merger, shuffle, null, false, 0, null, conf, getRawFs(conf), false,
+            HOST, PORT, host, ioErrsCounter, wrongLengthErrsCounter, badIdErrsCounter, wrongMapErrsCounter,
+            connectionErrsCounter, wrongReduceErrsCounter, false, false, true, false, inputContext);
+    final FetcherOrderedGrouped fetcher = spy(mockFetcher);
+
+    final List<InputAttemptIdentifier> srcAttempts =
+        List.of(new InputAttemptIdentifier(0, 1, InputAttemptIdentifier.PATH_PREFIX + "pathComponent_0"));
+    doReturn(srcAttempts).when(scheduler).getMapsForHost(host);
+
+    URL url =
+        ShuffleUtils.constructInputURL("http" + "://" + HOST + ":" + PORT + "/mapOutput?job=job_123&&reduce=1&map=",
+            srcAttempts, false);
+    fetcher.httpConnection = new FakeHttpConnection(url, null, "", null) {
+      @Override
+      public DataInputStream getInputStream() {
+        ByteArrayInputStream bin = new ByteArrayInputStream(new byte[1024]) {
+          @Override
+          public int read(byte[] b, int off, int len) {
+            try {
+              Thread.sleep(10);
+            } catch (InterruptedException e) {
+              Thread.currentThread().interrupt();
+            }
+            return super.read(b, off, len);
+          }
+        };
+        return new DataInputStream(bin);
+      }
+    };
+
+    fetcher.setupConnectionInternal(host, srcAttempts);
+
+    // Read some bytes to trigger the elapsed time measurement
+    byte[] buffer = new byte[10];
+    int bytesRead = fetcher.input.read(buffer, 0, buffer.length);
+    assertEquals(10, bytesRead);
+
+    // shutDown will update the counter
+    fetcher.shutDown();
+
+    // Check if io time counter is updated
+    TezCounter ioTimeCounter = inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_STREAM_TIME_MILLISECONDS);
+    long ioTime = ioTimeCounter.getValue();
+    assertTrue(ioTime >= 0, "IO Time should be measured and >= 0, but was " + ioTime);
+
+    // Check if io bytes counter is updated
+    TezCounter ioBytesCounter = inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_STREAM_BYTES);
+    long ioBytes = ioBytesCounter.getValue();
+    assertEquals(10, ioBytes, "IO Bytes should be exactly 10, but was " + ioBytes);
+  }
 }