From 65ef7ac1db62e91af706a179a91ab20ef74b0b3c Mon Sep 17 00:00:00 2001 From: Raghav Aggarwal Date: Wed, 5 Aug 2026 23:26:09 +0530 Subject: [PATCH 1/4] TEZ-4466: Measure exact stream read time while fetching --- .../tez/common/counters/TaskCounter.java | 5 ++ .../tez/http/MeasuredDataInputStream.java | 79 +++++++++++++++++++ .../library/api/TezRuntimeConfiguration.java | 8 ++ .../library/common/shuffle/Fetcher.java | 13 +++ .../orderedgrouped/FetcherOrderedGrouped.java | 12 +++ .../library/common/shuffle/TestFetcher.java | 5 +- .../shuffle/orderedgrouped/TestFetcher.java | 59 ++++++++++++++ 7 files changed, 177 insertions(+), 4 deletions(-) create mode 100644 tez-runtime-library/src/main/java/org/apache/tez/http/MeasuredDataInputStream.java diff --git a/tez-api/src/main/java/org/apache/tez/common/counters/TaskCounter.java b/tez-api/src/main/java/org/apache/tez/common/counters/TaskCounter.java index 56cafbc162..cabccd2ea1 100644 --- a/tez-api/src/main/java/org/apache/tez/common/counters/TaskCounter.java +++ b/tez-api/src/main/java/org/apache/tez/common/counters/TaskCounter.java @@ -189,6 +189,11 @@ public enum TaskCounter { */ SHUFFLE_BYTES_DISK_DIRECT, + /** + * Time spent waiting on network I/O during shuffle. Represented in milliseconds. + */ + SHUFFLE_IO_TIME_MILLISECONDS, + /** * Number of Memory to Disk merges performed during sort-merge. * Used by ShuffledMergedInput diff --git a/tez-runtime-library/src/main/java/org/apache/tez/http/MeasuredDataInputStream.java b/tez-runtime-library/src/main/java/org/apache/tez/http/MeasuredDataInputStream.java new file mode 100644 index 0000000000..c365accf44 --- /dev/null +++ b/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); + } + } +} diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/api/TezRuntimeConfiguration.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/api/TezRuntimeConfiguration.java index 569cde6367..df94b8d17d 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/api/TezRuntimeConfiguration.java +++ b/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); diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java index f31140a316..d4a9b317f2 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java +++ b/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,7 @@ public String getHost() { BaseHttpConnection httpConnection; private HttpConnectionParams httpConnectionParams; + private final TezCounter ioTimeCounter; private final boolean localDiskFetchEnabled; private final boolean sharedFetchEnabled; @@ -219,6 +223,8 @@ protected Fetcher(FetcherCallback fetcherCallback, HttpConnectionParams params, this.localDiskFetchEnabled = localDiskFetchEnabled; this.sharedFetchEnabled = sharedFetchEnabled; + this.ioTimeCounter = inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_TIME_MILLISECONDS); + this.fetcherIdentifier = fetcherIdGen.getAndIncrement(); String sourceDestNameTrimmed = TezUtilsInternal.cleanVertexName(inputContext.getSourceVertexName()) + " -> " @@ -565,6 +571,10 @@ private HostFetchResult setupConnection(Collection attem protected void setupConnectionInternal(String host, Collection 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(); } @@ -813,6 +823,9 @@ private void shutdownInternal(boolean disconnect) { synchronized (isShutDown) { try { if (httpConnection != null) { + if (input instanceof MeasuredDataInputStream && ioTimeCounter != null) { + ioTimeCounter.increment(((MeasuredDataInputStream) input).getElapsedTimeMs()); + } httpConnection.cleanup(disconnect); } } catch (IOException e) { diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java index 7f8f98f3ea..5d7cc758f8 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java +++ b/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,7 @@ class FetcherOrderedGrouped extends CallableWithNdc { private final TezCounter wrongLengthErrs; private final TezCounter badIdErrs; private final TezCounter wrongReduceErrs; + private final TezCounter ioTimeCounter; private final FetchedInputAllocatorOrderedGrouped allocator; private final ShuffleScheduler scheduler; private final ExceptionReporter exceptionReporter; @@ -151,6 +155,7 @@ public FetcherOrderedGrouped(HttpConnectionParams httpConnectionParams, this.badIdErrs = badIdErrsCounter; this.connectionErrs = connectionErrsCounter; this.wrongReduceErrs = wrongReduceErrsCounter; + this.ioTimeCounter = inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_TIME_MILLISECONDS); this.applicationId = inputContext.getApplicationId().toString(); this.dagId = inputContext.getDagIdentifier(); @@ -227,6 +232,9 @@ private void cleanupCurrentConnection(boolean disconnect) { synchronized (cleanupLock) { try { if (httpConnection != null) { + if (input instanceof MeasuredDataInputStream && ioTimeCounter != null) { + ioTimeCounter.increment(((MeasuredDataInputStream) input).getElapsedTimeMs()); + } httpConnection.cleanup(disconnect); httpConnection = null; } @@ -392,6 +400,10 @@ boolean setupConnection(MapHost host, Collection attempt protected void setupConnectionInternal(MapHost host, Collection 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(); } diff --git a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/TestFetcher.java b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/TestFetcher.java index fb0e34f50c..b53f1ce446 100644 --- a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/TestFetcher.java +++ b/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; } diff --git a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/TestFetcher.java b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/TestFetcher.java index c9b2473742..5569aa7a70 100644 --- a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/TestFetcher.java +++ b/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,62 @@ 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 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_TIME_MILLISECONDS); + long ioTime = ioTimeCounter.getValue(); + assertTrue(ioTime >= 10, "IO Time should be at least 10ms, but was " + ioTime); + } } From b1dd37b417b36eb575085fb422430084fea76cf4 Mon Sep 17 00:00:00 2001 From: Raghav Aggarwal Date: Thu, 6 Aug 2026 23:50:04 +0530 Subject: [PATCH 2/4] Fix SHUFFLE_IO_TIME off counter stat behaviour --- .../apache/tez/runtime/library/common/shuffle/Fetcher.java | 7 ++++--- .../shuffle/orderedgrouped/FetcherOrderedGrouped.java | 7 ++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java index d4a9b317f2..bf51163c2f 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java @@ -223,7 +223,9 @@ protected Fetcher(FetcherCallback fetcherCallback, HttpConnectionParams params, this.localDiskFetchEnabled = localDiskFetchEnabled; this.sharedFetchEnabled = sharedFetchEnabled; - this.ioTimeCounter = inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_TIME_MILLISECONDS); + this.ioTimeCounter = conf.getBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME, + TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME_DEFAULT) ? + inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_TIME_MILLISECONDS) : null; this.fetcherIdentifier = fetcherIdGen.getAndIncrement(); @@ -571,8 +573,7 @@ private HostFetchResult setupConnection(Collection attem protected void setupConnectionInternal(String host, Collection 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)) { + if (ioTimeCounter != null) { input = new MeasuredDataInputStream(input); } httpConnection.validate(); diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java index 5d7cc758f8..a679273e68 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java @@ -155,7 +155,9 @@ public FetcherOrderedGrouped(HttpConnectionParams httpConnectionParams, this.badIdErrs = badIdErrsCounter; this.connectionErrs = connectionErrsCounter; this.wrongReduceErrs = wrongReduceErrsCounter; - this.ioTimeCounter = inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_TIME_MILLISECONDS); + this.ioTimeCounter = conf.getBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME, + TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME_DEFAULT) ? + inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_TIME_MILLISECONDS) : null; this.applicationId = inputContext.getApplicationId().toString(); this.dagId = inputContext.getDagIdentifier(); @@ -400,8 +402,7 @@ boolean setupConnection(MapHost host, Collection attempt protected void setupConnectionInternal(MapHost host, Collection 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)) { + if (ioTimeCounter != null) { input = new MeasuredDataInputStream(input); } httpConnection.validate(); From 6b05939bbcc67b5971ccc273fbbf1c639c2594c6 Mon Sep 17 00:00:00 2001 From: Raghav Aggarwal Date: Mon, 17 Aug 2026 16:39:55 +0530 Subject: [PATCH 3/4] Address review comments --- .../tez/common/counters/TaskCounter.java | 12 +++++++- .../tez/http/MeasuredDataInputStream.java | 30 +++++++++++++++---- .../library/api/TezRuntimeConfiguration.java | 2 ++ .../library/common/shuffle/Fetcher.java | 11 +++++-- .../orderedgrouped/FetcherOrderedGrouped.java | 11 +++++-- .../shuffle/orderedgrouped/TestFetcher.java | 9 ++++-- 6 files changed, 60 insertions(+), 15 deletions(-) diff --git a/tez-api/src/main/java/org/apache/tez/common/counters/TaskCounter.java b/tez-api/src/main/java/org/apache/tez/common/counters/TaskCounter.java index cabccd2ea1..0c4f15674a 100644 --- a/tez-api/src/main/java/org/apache/tez/common/counters/TaskCounter.java +++ b/tez-api/src/main/java/org/apache/tez/common/counters/TaskCounter.java @@ -191,8 +191,18 @@ public enum TaskCounter { /** * 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_TIME_MILLISECONDS, + 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. diff --git a/tez-runtime-library/src/main/java/org/apache/tez/http/MeasuredDataInputStream.java b/tez-runtime-library/src/main/java/org/apache/tez/http/MeasuredDataInputStream.java index c365accf44..74f2dda2fa 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/http/MeasuredDataInputStream.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/http/MeasuredDataInputStream.java @@ -41,8 +41,13 @@ 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); @@ -51,29 +56,42 @@ private static class MeasuredInputStream extends FilterInputStream { @Override public int read() throws IOException { long start = System.nanoTime(); - int ret = super.read(); + int val = super.read(); elapsedTimeNanos += (System.nanoTime() - start); - return ret; + if (val != -1) { + bytesRead += 1; + } + return val; } @Override public int read(byte[] b) throws IOException { long start = System.nanoTime(); - int ret = super.read(b); + int bytes = super.read(b); elapsedTimeNanos += (System.nanoTime() - start); - return ret; + if (bytes > 0) { + bytesRead += bytes; + } + return bytes; } @Override public int read(byte[] b, int off, int len) throws IOException { long start = System.nanoTime(); - int ret = super.read(b, off, len); + int bytes = super.read(b, off, len); elapsedTimeNanos += (System.nanoTime() - start); - return ret; + if (bytes > 0) { + bytesRead += bytes; + } + return bytes; } public long getElapsedTimeMs() { return TimeUnit.NANOSECONDS.toMillis(elapsedTimeNanos); } + + public long getBytesRead() { + return bytesRead; + } } } diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/api/TezRuntimeConfiguration.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/api/TezRuntimeConfiguration.java index df94b8d17d..115753a112 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/api/TezRuntimeConfiguration.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/api/TezRuntimeConfiguration.java @@ -418,6 +418,8 @@ private TezRuntimeConfiguration() {} /** * 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"; diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java index bf51163c2f..fd7b6fbd38 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java @@ -181,6 +181,7 @@ public String getHost() { BaseHttpConnection httpConnection; private HttpConnectionParams httpConnectionParams; private final TezCounter ioTimeCounter; + private final TezCounter ioBytesCounter; private final boolean localDiskFetchEnabled; private final boolean sharedFetchEnabled; @@ -223,9 +224,12 @@ protected Fetcher(FetcherCallback fetcherCallback, HttpConnectionParams params, this.localDiskFetchEnabled = localDiskFetchEnabled; this.sharedFetchEnabled = sharedFetchEnabled; - this.ioTimeCounter = conf.getBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME, - TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME_DEFAULT) ? - inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_TIME_MILLISECONDS) : null; + 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(); @@ -826,6 +830,7 @@ private void shutdownInternal(boolean disconnect) { if (httpConnection != null) { if (input instanceof MeasuredDataInputStream && ioTimeCounter != null) { ioTimeCounter.increment(((MeasuredDataInputStream) input).getElapsedTimeMs()); + ioBytesCounter.increment(((MeasuredDataInputStream) input).getBytesRead()); } httpConnection.cleanup(disconnect); } diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java index a679273e68..ffccf93520 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java @@ -79,6 +79,7 @@ class FetcherOrderedGrouped extends CallableWithNdc { 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; @@ -155,9 +156,12 @@ public FetcherOrderedGrouped(HttpConnectionParams httpConnectionParams, this.badIdErrs = badIdErrsCounter; this.connectionErrs = connectionErrsCounter; this.wrongReduceErrs = wrongReduceErrsCounter; - this.ioTimeCounter = conf.getBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME, - TezRuntimeConfiguration.TEZ_RUNTIME_SHUFFLE_MEASURE_IO_TIME_DEFAULT) ? - inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_TIME_MILLISECONDS) : null; + 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(); @@ -236,6 +240,7 @@ private void cleanupCurrentConnection(boolean disconnect) { if (httpConnection != null) { if (input instanceof MeasuredDataInputStream && ioTimeCounter != null) { ioTimeCounter.increment(((MeasuredDataInputStream) input).getElapsedTimeMs()); + ioBytesCounter.increment(((MeasuredDataInputStream) input).getBytesRead()); } httpConnection.cleanup(disconnect); httpConnection = null; diff --git a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/TestFetcher.java b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/TestFetcher.java index 5569aa7a70..f9c818c4e2 100644 --- a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/TestFetcher.java +++ b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/TestFetcher.java @@ -853,8 +853,13 @@ public int read(byte[] b, int off, int len) { fetcher.shutDown(); // Check if io time counter is updated - TezCounter ioTimeCounter = inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_TIME_MILLISECONDS); + TezCounter ioTimeCounter = inputContext.getCounters().findCounter(TaskCounter.SHUFFLE_IO_STREAM_TIME_MILLISECONDS); long ioTime = ioTimeCounter.getValue(); - assertTrue(ioTime >= 10, "IO Time should be at least 10ms, but was " + ioTime); + 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); } } From 1caccbc5922d8f70b70e26f8f60b42e141bea0f9 Mon Sep 17 00:00:00 2001 From: Raghav Aggarwal Date: Tue, 18 Aug 2026 00:46:08 +0530 Subject: [PATCH 4/4] review comments --- .../org/apache/tez/runtime/library/common/shuffle/Fetcher.java | 2 +- .../common/shuffle/orderedgrouped/FetcherOrderedGrouped.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java index fd7b6fbd38..6739a36038 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/Fetcher.java @@ -828,7 +828,7 @@ private void shutdownInternal(boolean disconnect) { synchronized (isShutDown) { try { if (httpConnection != null) { - if (input instanceof MeasuredDataInputStream && ioTimeCounter != null) { + if (ioTimeCounter != null && input != null) { ioTimeCounter.increment(((MeasuredDataInputStream) input).getElapsedTimeMs()); ioBytesCounter.increment(((MeasuredDataInputStream) input).getBytesRead()); } diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java index ffccf93520..7f367d5ad2 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/FetcherOrderedGrouped.java @@ -238,7 +238,7 @@ private void cleanupCurrentConnection(boolean disconnect) { synchronized (cleanupLock) { try { if (httpConnection != null) { - if (input instanceof MeasuredDataInputStream && ioTimeCounter != null) { + if (ioTimeCounter != null && input != null) { ioTimeCounter.increment(((MeasuredDataInputStream) input).getElapsedTimeMs()); ioBytesCounter.increment(((MeasuredDataInputStream) input).getBytesRead()); }