From 4ad0480e9843b653397acaf08b36109e3decffa1 Mon Sep 17 00:00:00 2001 From: Raghav Aggarwal Date: Wed, 5 Aug 2026 23:26:09 +0530 Subject: [PATCH] TEZ-4466: Measure exact stream read time while fetching --- .../tez/http/MeasuredDataInputStream.java | 79 +++++++++++++++++++ .../library/api/TezRuntimeConfiguration.java | 8 ++ .../library/common/shuffle/Fetcher.java | 5 ++ .../orderedgrouped/FetcherOrderedGrouped.java | 6 ++ 4 files changed, 98 insertions(+) create mode 100644 tez-runtime-library/src/main/java/org/apache/tez/http/MeasuredDataInputStream.java 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..dd7b7efd94 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 @@ -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 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(); } 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..2ce59928f5 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 @@ -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 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(); }