This is an automated email from the ASF dual-hosted git repository.
RexXiong pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new f28246dbd8 [CELEBORN-2430] Expose unexpected errors in client
background threads
f28246dbd8 is described below
commit f28246dbd89074b58221ef4b2e2c542605f325f2
Author: Kalvin2077 <[email protected]>
AuthorDate: Tue Sep 15 13:55:28 2026 +0800
[CELEBORN-2430] Expose unexpected errors in client background threads
### What changes were proposed in this pull request?
Catch and log unexpected `Throwable`s in client background tasks. Partition
readers also propagate background failures through their existing exception
channel. No recovery behavior is added.
### Why are the changes needed?
`ExecutorService.submit` stores task failures in the returned `Future`. If
that `Future` is not inspected, the failures can remain invisible. Catching
only `Exception` also misses `Error`s.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce *any* user-facing change?
- [ ] Yes
### How was this patch tested?
Added unit tests verifying that DFS and local partition readers propagate
background errors, and that change- and release-partition tasks log them. The
focused tests passed, including red/green verification for the manager tasks.
Formatting checks also passed.
Closes #3829 from Kalvin2077/fix/CELEBORN-2430.
Authored-by: Kalvin2077 <[email protected]>
Signed-off-by: Shuang <[email protected]>
---
.../celeborn/client/read/DfsPartitionReader.java | 13 +-
.../celeborn/client/read/LocalPartitionReader.java | 10 +-
.../celeborn/client/ChangePartitionManager.scala | 52 +++--
.../celeborn/client/ReleasePartitionManager.scala | 49 +++--
.../client/ClientThreadErrorLoggingSuiteJ.java | 187 +++++++++++++++++
.../read/PartitionReaderErrorHandlingSuiteJ.java | 224 +++++++++++++++++++++
6 files changed, 487 insertions(+), 48 deletions(-)
diff --git
a/client/src/main/java/org/apache/celeborn/client/read/DfsPartitionReader.java
b/client/src/main/java/org/apache/celeborn/client/read/DfsPartitionReader.java
index 735a532321..d277db9fad 100644
---
a/client/src/main/java/org/apache/celeborn/client/read/DfsPartitionReader.java
+++
b/client/src/main/java/org/apache/celeborn/client/read/DfsPartitionReader.java
@@ -39,6 +39,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.ShuffleClient;
import
org.apache.celeborn.client.read.checkpoint.PartitionReaderCheckpointMetadata;
import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.network.client.TransportClient;
import org.apache.celeborn.common.network.client.TransportClientFactory;
import org.apache.celeborn.common.network.protocol.TransportMessage;
@@ -283,10 +284,16 @@ public class DfsPartitionReader implements
PartitionReader {
results.put(Pair.of(currentChunkIndex,
Unpooled.wrappedBuffer(buffer)));
logger.debug("add index {} to results", currentChunkIndex++);
}
- } catch (Exception e) {
- logger.warn("Fetch thread is cancelled.", e);
+ } catch (InterruptedException e) {
+ logger.warn("Read thread is interrupted.", e);
exception.set(e);
- // cancel a task for speculative, ignore this exception
+ } catch (Throwable t) {
+ logger.error("Read thread encountered error.", t);
+ if (t instanceof Exception) {
+ exception.set((Exception) t);
+ } else {
+ exception.set(new CelebornIOException("Fetch thread
encountered an error", t));
+ }
}
logger.debug("fetch {} is done.",
location.getStorageInfo().getFilePath());
});
diff --git
a/client/src/main/java/org/apache/celeborn/client/read/LocalPartitionReader.java
b/client/src/main/java/org/apache/celeborn/client/read/LocalPartitionReader.java
index 0e79503761..127cbbd144 100644
---
a/client/src/main/java/org/apache/celeborn/client/read/LocalPartitionReader.java
+++
b/client/src/main/java/org/apache/celeborn/client/read/LocalPartitionReader.java
@@ -179,12 +179,12 @@ public class LocalPartitionReader implements
PartitionReader {
} catch (InterruptedException e) {
// cancel a task for speculative, ignore this exception
logger.warn("Read thread is interrupted.", e);
- } catch (Exception ioe) {
- logger.error("Read thread encountered error.", ioe);
- if (ioe instanceof CelebornIOException) {
- exception.set((IOException) ioe);
+ } catch (Throwable t) {
+ logger.error("Read thread encountered error.", t);
+ if (t instanceof CelebornIOException) {
+ exception.set((IOException) t);
} else {
- exception.set(new CelebornIOException("Read thread encountered error",
ioe));
+ exception.set(new CelebornIOException("Read thread encountered error",
t));
}
}
pendingFetchTask.compareAndSet(true, false);
diff --git
a/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
b/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
index c0af4e9113..ec843e8f54 100644
---
a/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
+++
b/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
@@ -88,26 +88,38 @@ class ChangePartitionManager(
batchHandleChangePartitionExecutors.submit {
new Runnable {
override def run(): Unit = {
- val distinctPartitions = {
- val requestSet = inBatchPartitions.get(shuffleId)
- val locksForShuffle = locks.computeIfAbsent(shuffleId,
locksRegisterFunc)
- requests.asScala.map { case (partitionId, request) =>
- locksForShuffle(partitionId %
locksForShuffle.length).synchronized {
- if (!requestSet.contains(partitionId) &&
requests.containsKey(
- partitionId)) {
- requestSet.add(partitionId)
- Some(request.asScala.toArray.maxBy(_.epoch))
- } else {
- None
+ try {
+ val distinctPartitions = {
+ val requestSet = inBatchPartitions.get(shuffleId)
+ val locksForShuffle =
locks.computeIfAbsent(shuffleId, locksRegisterFunc)
+ requests.asScala.map { case (partitionId, request) =>
+ locksForShuffle(partitionId %
locksForShuffle.length).synchronized {
+ if (!requestSet.contains(partitionId) &&
requests.containsKey(
+ partitionId)) {
+ requestSet.add(partitionId)
+ Some(request.asScala.toArray.maxBy(_.epoch))
+ } else {
+ None
+ }
}
- }
- }.filter(_.isDefined).map(_.get).toArray
- }
- if (distinctPartitions.nonEmpty) {
- handleRequestPartitions(
- shuffleId,
- distinctPartitions,
-
lifecycleManager.commitManager.isSegmentGranularityVisible(shuffleId))
+ }.filter(_.isDefined).map(_.get).toArray
+ }
+ if (distinctPartitions.nonEmpty) {
+ handleRequestPartitions(
+ shuffleId,
+ distinctPartitions,
+
lifecycleManager.commitManager.isSegmentGranularityVisible(shuffleId))
+ }
+ } catch {
+ case e: InterruptedException =>
+ logError(
+ s"Batch handle change partition for shuffle
$shuffleId interrupted.",
+ e)
+ throw e
+ case t: Throwable =>
+ logError(
+ s"Batch handle change partition for shuffle
$shuffleId failed.",
+ t)
}
}
}
@@ -117,6 +129,8 @@ class ChangePartitionManager(
case e: InterruptedException =>
logError("Partition split scheduler thread is shutting down,
detail: ", e)
throw e
+ case t: Throwable =>
+ logError("Batch handle change partition scheduler failed.", t)
}
}
},
diff --git
a/client/src/main/scala/org/apache/celeborn/client/ReleasePartitionManager.scala
b/client/src/main/scala/org/apache/celeborn/client/ReleasePartitionManager.scala
index a800c8c9ce..be50c71c1a 100644
---
a/client/src/main/scala/org/apache/celeborn/client/ReleasePartitionManager.scala
+++
b/client/src/main/scala/org/apache/celeborn/client/ReleasePartitionManager.scala
@@ -62,29 +62,36 @@ class ReleasePartitionManager(
batchHandleReleasePartitionExecutors.submit {
new Runnable {
override def run(): Unit = {
- val unReleasePartitionIds = new util.HashSet[Int]
- unReleasedPartitionIdRequestSet.synchronized {
-
unReleasePartitionIds.addAll(unReleasedPartitionIdRequestSet)
- unReleasedPartitionIdRequestSet.clear()
- }
+ try {
+ val unReleasePartitionIds = new util.HashSet[Int]
+ unReleasedPartitionIdRequestSet.synchronized {
+
unReleasePartitionIds.addAll(unReleasedPartitionIdRequestSet)
+ unReleasedPartitionIdRequestSet.clear()
+ }
-
lifecycleManager.workerSnapshots(shuffleId).asScala.foreach {
- case (_, partitionLocationInfo) =>
- val destroyResource = new WorkerResource
- unReleasePartitionIds.asScala.foreach {
- partitionId =>
- addDestroyResource(
- destroyResource,
- partitionLocationInfo,
- partitionId)
- }
+
lifecycleManager.workerSnapshots(shuffleId).asScala.foreach {
+ case (_, partitionLocationInfo) =>
+ val destroyResource = new WorkerResource
+ unReleasePartitionIds.asScala.foreach {
+ partitionId =>
+ addDestroyResource(
+ destroyResource,
+ partitionLocationInfo,
+ partitionId)
+ }
- if (!destroyResource.isEmpty) {
- lifecycleManager.destroySlotsWithRetry(
- shuffleId,
- destroyResource)
- logTrace(s"Destroyed partition resource for
shuffle $shuffleId $destroyResource")
- }
+ if (!destroyResource.isEmpty) {
+ lifecycleManager.destroySlotsWithRetry(
+ shuffleId,
+ destroyResource)
+ logTrace(s"Destroyed partition resource for
shuffle $shuffleId $destroyResource")
+ }
+ }
+ } catch {
+ case t: Throwable =>
+ logError(
+ s"Error releasing partition resource for shuffle
$shuffleId",
+ t)
}
}
}
diff --git
a/client/src/test/java/org/apache/celeborn/client/ClientThreadErrorLoggingSuiteJ.java
b/client/src/test/java/org/apache/celeborn/client/ClientThreadErrorLoggingSuiteJ.java
new file mode 100644
index 0000000000..55d2b22e93
--- /dev/null
+++
b/client/src/test/java/org/apache/celeborn/client/ClientThreadErrorLoggingSuiteJ.java
@@ -0,0 +1,187 @@
+/*
+ * 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.celeborn.client;
+
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.lang.reflect.Field;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import scala.Option;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.core.LogEvent;
+import org.apache.logging.log4j.core.Logger;
+import org.apache.logging.log4j.core.appender.AbstractAppender;
+import org.junit.Test;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.protocol.message.StatusCode;
+
+public class ClientThreadErrorLoggingSuiteJ {
+
+ @Test
+ public void testChangePartitionBackgroundErrorIsLogged() throws Exception {
+ int shuffleId = 1;
+ int partitionId = 2;
+ CelebornConf conf = new CelebornConf();
+
conf.set(CelebornConf.CLIENT_BATCH_HANDLE_CHANGE_PARTITION_INTERVAL().key(),
"10ms");
+ LifecycleManager lifecycleManager = mock(LifecycleManager.class);
+ CommitManager commitManager = mock(CommitManager.class);
+ when(lifecycleManager.commitManager()).thenReturn(commitManager);
+ when(lifecycleManager.latestPartitionLocation()).thenReturn(new
ConcurrentHashMap<>());
+
+ AssertionError expected = new AssertionError("change partition failed");
+ FailingChangePartitionManager manager =
+ new FailingChangePartitionManager(conf, lifecycleManager, expected);
+ manager.handleRequestPartitionLocation(
+ new NoOpRequestLocationCallContext(),
+ shuffleId,
+ partitionId,
+ 0,
+ newPartition(partitionId),
+ Option.empty(),
+ false);
+ MessageAppender appender =
+ new MessageAppender("Batch handle change partition for shuffle 1
failed.");
+
+ Logger rootLogger = (Logger) LogManager.getRootLogger();
+ rootLogger.addAppender(appender);
+ try {
+ manager.start();
+ assertTrue("The change-partition error was not logged",
appender.await());
+ assertSame(expected, appender.loggingEvent.getThrown());
+ } finally {
+ manager.stop();
+ shutdownExecutor(manager, "batchHandleChangePartitionExecutors");
+ rootLogger.removeAppender(appender);
+ appender.stop();
+ }
+ }
+
+ @Test
+ public void testReleasePartitionBackgroundErrorIsLogged() throws Exception {
+ int shuffleId = 1;
+ CelebornConf conf = new CelebornConf();
+
conf.set(CelebornConf.CLIENT_BATCH_HANDLED_RELEASE_PARTITION_INTERVAL().key(),
"10ms");
+ LifecycleManager lifecycleManager = mock(LifecycleManager.class);
+ AssertionError expected = new AssertionError("release partition failed");
+ when(lifecycleManager.workerSnapshots(shuffleId)).thenThrow(expected);
+
+ ReleasePartitionManager manager = new ReleasePartitionManager(conf,
lifecycleManager);
+ manager.releasePartition(shuffleId, 2);
+ MessageAppender appender =
+ new MessageAppender("Error releasing partition resource for shuffle
1");
+
+ Logger rootLogger = (Logger) LogManager.getRootLogger();
+ rootLogger.addAppender(appender);
+ try {
+ manager.start();
+ assertTrue("The release-partition error was not logged",
appender.await());
+ assertSame(expected, appender.loggingEvent.getThrown());
+ } finally {
+ manager.stop();
+ shutdownExecutor(manager, "batchHandleReleasePartitionExecutors");
+ rootLogger.removeAppender(appender);
+ appender.stop();
+ }
+ }
+
+ private static PartitionLocation newPartition(int partitionId) {
+ return new PartitionLocation(
+ partitionId, 0, "localhost", 1, 2, 3, 4,
PartitionLocation.Mode.PRIMARY);
+ }
+
+ private static void shutdownExecutor(Object target, String fieldSuffix)
throws Exception {
+ Class<?> currentClass = target.getClass();
+ while (currentClass != null) {
+ for (Field field : currentClass.getDeclaredFields()) {
+ if (field.getName().endsWith(fieldSuffix)) {
+ field.setAccessible(true);
+ ((ExecutorService) field.get(target)).shutdownNow();
+ return;
+ }
+ }
+ currentClass = currentClass.getSuperclass();
+ }
+ throw new NoSuchFieldException(fieldSuffix);
+ }
+
+ private static class FailingChangePartitionManager extends
ChangePartitionManager {
+ private final Throwable failure;
+
+ FailingChangePartitionManager(
+ CelebornConf conf, LifecycleManager lifecycleManager, Throwable
failure) {
+ super(conf, lifecycleManager);
+ this.failure = failure;
+ }
+
+ @Override
+ public void handleRequestPartitions(
+ int shuffleId,
+ ChangePartitionRequest[] changePartitions,
+ boolean isSegmentGranularityVisible) {
+ ClientThreadErrorLoggingSuiteJ.<RuntimeException>throwUnchecked(failure);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private static <T extends Throwable> void throwUnchecked(Throwable failure)
throws T {
+ throw (T) failure;
+ }
+
+ private static class NoOpRequestLocationCallContext implements
RequestLocationCallContext {
+ @Override
+ public void reply(
+ int partitionId,
+ StatusCode status,
+ Option<PartitionLocation> partitionLocation,
+ boolean available) {}
+ }
+
+ private static class MessageAppender extends AbstractAppender {
+ private final String expectedMessage;
+ private final CountDownLatch logged = new CountDownLatch(1);
+ private volatile LogEvent loggingEvent;
+
+ MessageAppender(String expectedMessage) {
+ super("MessageAppender", null, null, false);
+ this.expectedMessage = expectedMessage;
+ start();
+ }
+
+ @Override
+ public void append(LogEvent event) {
+ if (expectedMessage.equals(event.getMessage().getFormattedMessage())) {
+ loggingEvent = event.toImmutable();
+ logged.countDown();
+ }
+ }
+
+ boolean await() throws InterruptedException {
+ return logged.await(5, TimeUnit.SECONDS);
+ }
+ }
+}
diff --git
a/client/src/test/java/org/apache/celeborn/client/read/PartitionReaderErrorHandlingSuiteJ.java
b/client/src/test/java/org/apache/celeborn/client/read/PartitionReaderErrorHandlingSuiteJ.java
new file mode 100644
index 0000000000..09119f4603
--- /dev/null
+++
b/client/src/test/java/org/apache/celeborn/client/read/PartitionReaderErrorHandlingSuiteJ.java
@@ -0,0 +1,224 @@
+/*
+ * 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.celeborn.client.read;
+
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.io.DataOutputStream;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.lang.reflect.Field;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import io.netty.buffer.ByteBuf;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import org.apache.celeborn.client.ShuffleClient;
+import
org.apache.celeborn.client.read.checkpoint.PartitionReaderCheckpointMetadata;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.exception.CelebornIOException;
+import org.apache.celeborn.common.network.client.TransportClient;
+import org.apache.celeborn.common.network.client.TransportClientFactory;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.protocol.PbStreamHandler;
+import org.apache.celeborn.common.protocol.StorageInfo;
+import org.apache.celeborn.common.util.Utils;
+
+public class PartitionReaderErrorHandlingSuiteJ {
+
+ private static final int FETCH_PORT = 10001;
+
+ @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+ @Before
+ public void setUp() {
+ ShuffleClient.reset();
+ }
+
+ @After
+ public void tearDown() {
+ ShuffleClient.reset();
+ }
+
+ @Test
+ public void testDfsBackgroundErrorFailsNext() throws Exception {
+ PartitionReaderCheckpointMetadata checkpointMetadata =
+ mock(PartitionReaderCheckpointMetadata.class);
+
when(checkpointMetadata.getReturnedChunks()).thenReturn(Collections.emptySet());
+ AssertionError expected = new AssertionError("dfs fetch failed");
+ when(checkpointMetadata.isCheckpointed(anyInt())).thenThrow(expected);
+
+ DfsPartitionReader reader = newDfsReader(checkpointMetadata);
+ try {
+ Throwable failure = awaitNextFailure(reader);
+ assertTrue(failure instanceof CelebornIOException);
+ assertSame(expected, failure.getCause());
+ } finally {
+ reader.close();
+ }
+ }
+
+ @Test
+ public void testLocalBackgroundErrorFailsNext() throws Exception {
+ File dataFile = temporaryFolder.newFile("local-partition.data");
+ LocalPartitionReader reader = newLocalReader(dataFile);
+ FileChannel channel = mock(FileChannel.class);
+ AssertionError expected = new AssertionError("local fetch failed");
+ when(channel.read(any(ByteBuffer.class))).thenThrow(expected);
+ setField(reader, "shuffleChannel", channel);
+
+ try {
+ Throwable failure = awaitNextFailure(reader);
+ assertTrue(failure instanceof CelebornIOException);
+ assertSame(expected, failure.getCause());
+ verify(channel).read(any(ByteBuffer.class));
+ } finally {
+ setField(reader, "shuffleChannel", null);
+ reader.close();
+ }
+ }
+
+ private DfsPartitionReader newDfsReader(PartitionReaderCheckpointMetadata
checkpointMetadata)
+ throws Exception {
+ File dataFile = temporaryFolder.newFile("partition.data");
+ try (FileOutputStream output = new FileOutputStream(dataFile)) {
+ output.write(1);
+ }
+ try (DataOutputStream output =
+ new DataOutputStream(new
FileOutputStream(Utils.getIndexFilePath(dataFile.getPath())))) {
+ output.writeInt(2);
+ output.writeLong(0L);
+ output.writeLong(1L);
+ }
+
+ Map<StorageInfo.Type, FileSystem> hadoopFs =
+ Collections.singletonMap(StorageInfo.Type.HDFS,
FileSystem.getLocal(new Configuration()));
+ setStaticField(ShuffleClient.class, "hadoopFs", hadoopFs);
+
+ TransportClientFactory clientFactory = mock(TransportClientFactory.class);
+ when(clientFactory.createClient(anyString(),
anyInt())).thenReturn(mock(TransportClient.class));
+ return new DfsPartitionReader(
+ new CelebornConf(),
+ "app-1",
+ newLocation(StorageInfo.Type.HDFS, dataFile.getAbsolutePath()),
+ PbStreamHandler.newBuilder().setStreamId(1L).build(),
+ clientFactory,
+ 0,
+ Integer.MAX_VALUE,
+ new NoOpMetricsCallback(),
+ -1,
+ -1,
+ Optional.of(checkpointMetadata));
+ }
+
+ private LocalPartitionReader newLocalReader(File dataFile) throws Exception {
+ TransportClientFactory clientFactory = mock(TransportClientFactory.class);
+ when(clientFactory.createClient(anyString(), anyInt(), anyInt()))
+ .thenReturn(mock(TransportClient.class));
+ PbStreamHandler streamHandler =
+ PbStreamHandler.newBuilder()
+ .setStreamId(1L)
+ .setNumChunks(1)
+ .addChunkOffsets(0L)
+ .addChunkOffsets(1L)
+ .setFullPath(dataFile.getAbsolutePath())
+ .build();
+ return new LocalPartitionReader(
+ new CelebornConf(),
+ "app-1",
+ newLocation(StorageInfo.Type.HDD, dataFile.getAbsolutePath()),
+ streamHandler,
+ clientFactory,
+ 0,
+ Integer.MAX_VALUE,
+ new NoOpMetricsCallback(),
+ -1,
+ -1);
+ }
+
+ private static Throwable awaitNextFailure(PartitionReader reader) throws
Exception {
+ ExecutorService caller = Executors.newSingleThreadExecutor();
+ Future<ByteBuf> result = caller.submit(reader::next);
+ try {
+ result.get(5, TimeUnit.SECONDS);
+ fail("PartitionReader.next() should fail when its background task
fails");
+ } catch (ExecutionException e) {
+ return e.getCause();
+ } catch (TimeoutException e) {
+ throw new AssertionError(
+ "PartitionReader.next() did not terminate after its background task
failed", e);
+ } finally {
+ result.cancel(true);
+ caller.shutdownNow();
+ }
+ throw new AssertionError("unreachable");
+ }
+
+ private static PartitionLocation newLocation(StorageInfo.Type storageType,
String path) {
+ PartitionLocation location =
+ new PartitionLocation(
+ 0, 0, "localhost", 10000, 10002, FETCH_PORT, 10003,
PartitionLocation.Mode.PRIMARY);
+ location.setStorageInfo(new StorageInfo(storageType, true, path));
+ return location;
+ }
+
+ private static void setField(Object target, String name, Object value)
throws Exception {
+ Field field = target.getClass().getDeclaredField(name);
+ field.setAccessible(true);
+ field.set(target, value);
+ }
+
+ private static void setStaticField(Class<?> targetClass, String name, Object
value)
+ throws Exception {
+ Field field = targetClass.getDeclaredField(name);
+ field.setAccessible(true);
+ field.set(null, value);
+ }
+
+ private static class NoOpMetricsCallback implements MetricsCallback {
+ @Override
+ public void incBytesRead(long bytesRead) {}
+
+ @Override
+ public void incReadTime(long time) {}
+ }
+}
