This is an automated email from the ASF dual-hosted git repository.
fapifta pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/ozone.git
The following commit(s) were added to refs/heads/master by this push:
new 67fd51904cf HDDS-15170. Add mock-based unit tests for DataStream write
path (#10230)
67fd51904cf is described below
commit 67fd51904cf4eb0704e3e63385c5c9c133818dfa
Author: Istvan Fajth <[email protected]>
AuthorDate: Tue Jul 7 09:37:30 2026 +0200
HDDS-15170. Add mock-based unit tests for DataStream write path (#10230)
---
.../hdds/scm/storage/MockDatanodePipeline.java | 288 +++++++++++++++++++
.../scm/storage/TestBlockDataStreamOutput.java | 266 +++++++++++++++++
.../ozone/client/io/TestKeyDataStreamOutput.java | 313 +++++++++++++++++++++
3 files changed, 867 insertions(+)
diff --git
a/hadoop-hdds/client/src/test/java/org/apache/hadoop/hdds/scm/storage/MockDatanodePipeline.java
b/hadoop-hdds/client/src/test/java/org/apache/hadoop/hdds/scm/storage/MockDatanodePipeline.java
new file mode 100644
index 00000000000..85341fd2721
--- /dev/null
+++
b/hadoop-hdds/client/src/test/java/org/apache/hadoop/hdds/scm/storage/MockDatanodePipeline.java
@@ -0,0 +1,288 @@
+/*
+ * 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.hadoop.hdds.scm.storage;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Supplier;
+import org.apache.hadoop.hdds.client.BlockID;
+import
org.apache.hadoop.hdds.protocol.datanode.proto.ContainerProtos.ContainerCommandRequestProto;
+import
org.apache.hadoop.hdds.protocol.datanode.proto.ContainerProtos.ContainerCommandResponseProto;
+import
org.apache.hadoop.hdds.protocol.datanode.proto.ContainerProtos.GetCommittedBlockLengthResponseProto;
+import
org.apache.hadoop.hdds.protocol.datanode.proto.ContainerProtos.PutBlockResponseProto;
+import org.apache.hadoop.hdds.protocol.datanode.proto.ContainerProtos.Result;
+import org.apache.hadoop.hdds.protocol.datanode.proto.ContainerProtos.Type;
+import org.apache.hadoop.hdds.scm.XceiverClientFactory;
+import org.apache.hadoop.hdds.scm.XceiverClientManager;
+import org.apache.hadoop.hdds.scm.XceiverClientRatis;
+import org.apache.hadoop.hdds.scm.XceiverClientReply;
+import org.apache.hadoop.hdds.scm.pipeline.MockPipeline;
+import org.apache.hadoop.hdds.scm.pipeline.Pipeline;
+import org.apache.ratis.client.api.DataStreamApi;
+import org.apache.ratis.client.api.DataStreamOutput;
+import org.apache.ratis.io.FilePositionCount;
+import org.apache.ratis.io.StandardWriteOption;
+import org.apache.ratis.io.WriteOption;
+import org.apache.ratis.protocol.DataStreamReply;
+import org.apache.ratis.protocol.RoutingTable;
+
+/**
+ * A stateful test harness that simulates a datanode pipeline for {@link
BlockDataStreamOutput} unit tests.
+ * Replaces a real Ratis pipeline with mocked {@link XceiverClientRatis} and a
concrete {@link DataStreamOutput}
+ * implementation.
+ *
+ * <p>Tracks all chunks written, putBlock calls, and watchForCommit calls.
+ * Configurable failure injection for each operation.
+ */
+public class MockDatanodePipeline {
+
+ private final Pipeline pipeline;
+ private final XceiverClientRatis xceiverClient;
+ private final XceiverClientFactory clientFactory;
+ private final BlockID blockID;
+
+ // Recorded state
+ private final List<byte[]> receivedChunks = Collections.synchronizedList(new
ArrayList<>());
+ private final List<ContainerCommandRequestProto> receivedPutBlocks =
Collections.synchronizedList(new ArrayList<>());
+ private final AtomicInteger watchForCommitCount = new AtomicInteger(0);
+
+ // Commit tracking
+ private final AtomicLong nextLogIndex = new AtomicLong(1);
+
+ // Failure injection
+ private volatile Supplier<Throwable> chunkFailure = null;
+ private volatile int chunkFailAfter = Integer.MAX_VALUE;
+ private final AtomicInteger chunkCount = new AtomicInteger(0);
+
+ private volatile Supplier<Throwable> putBlockFailure = null;
+ private volatile int putBlockFailAfter = Integer.MAX_VALUE;
+ private final AtomicInteger putBlockCount = new AtomicInteger(0);
+
+ private volatile Supplier<Throwable> watchFailure = null;
+ private volatile int watchFailAfter = Integer.MAX_VALUE;
+
+ public MockDatanodePipeline() throws IOException {
+ this(new BlockID(1, 1));
+ }
+
+ public MockDatanodePipeline(BlockID blockID) throws IOException {
+ this.blockID = blockID;
+ this.pipeline = MockPipeline.createRatisPipeline();
+
+ // Ensure metrics are initialized
+ XceiverClientManager.getXceiverClientMetrics();
+
+ // Create concrete DataStreamOutput
+ DataStreamOutput mockDataStreamOutput = spy(DataStreamOutput.class);
+ doThrow(new UnsupportedOperationException()).
+ when(mockDataStreamOutput).writeAsync(any(FilePositionCount.class),
any(WriteOption[].class));
+ doThrow(new
UnsupportedOperationException()).when(mockDataStreamOutput).getRaftClientReplyFuture();
+ doThrow(new
UnsupportedOperationException()).when(mockDataStreamOutput).getWritableByteChannel();
+
doReturn(CompletableFuture.completedFuture(dataStreamReply(0))).when(mockDataStreamOutput).closeAsync();
+
+ // Mock XceiverClientRatis
+ this.xceiverClient = mock(XceiverClientRatis.class);
+ when(xceiverClient.getPipeline()).thenReturn(pipeline);
+ doReturn(0L).when(xceiverClient).getReplicatedMinCommitIndex();
+
+ // Mock DataStreamApi to return our concrete DataStreamOutput
+ // Both overloads must be stubbed: stream(ByteBuffer) and
stream(ByteBuffer, RoutingTable) — the pipeline-mode
+ // default is true, so the 2-arg overload is what
BlockDataStreamOutput.setupStream calls.
+ DataStreamApi dataStreamApi = mock(DataStreamApi.class);
+
doReturn(mockDataStreamOutput).when(dataStreamApi).stream(any(ByteBuffer.class));
+
doReturn(mockDataStreamOutput).when(dataStreamApi).stream(any(ByteBuffer.class),
any(RoutingTable.class));
+ doReturn(dataStreamApi).when(xceiverClient).getDataStreamApi();
+
+ // Setup sendCommandAsync (putBlock) behavior
+ doAnswer(invocation -> {
+ ContainerCommandRequestProto request = invocation.getArgument(0);
+ if (request.getCmdType() == Type.PutBlock) {
+ receivedPutBlocks.add(request);
+ int count = putBlockCount.incrementAndGet();
+ CompletableFuture<ContainerCommandResponseProto> f = new
CompletableFuture<>();
+ if (count > putBlockFailAfter && putBlockFailure != null) {
+ f.completeExceptionally(putBlockFailure.get());
+ } else {
+ ContainerCommandResponseProto response =
buildPutBlockResponse(blockID);
+ f.complete(response);
+ }
+ XceiverClientReply reply = new XceiverClientReply(f);
+ reply.setLogIndex(nextLogIndex.getAndIncrement());
+ return reply;
+ }
+ // Default: return success
+ ContainerCommandResponseProto response =
+ ContainerCommandResponseProto.newBuilder()
+ .setCmdType(request.getCmdType())
+ .setResult(Result.SUCCESS)
+ .build();
+ XceiverClientReply reply = new
XceiverClientReply(CompletableFuture.completedFuture(response));
+ reply.setLogIndex(0);
+ return reply;
+ }).when(xceiverClient).sendCommandAsync(any());
+
+ // Setup watchForCommit behavior
+ doAnswer(invocation -> {
+ long index = invocation.getArgument(0);
+ int count = watchForCommitCount.incrementAndGet();
+ CompletableFuture<XceiverClientReply> f = new CompletableFuture<>();
+ if (count > watchFailAfter && watchFailure != null) {
+ f.completeExceptionally(watchFailure.get());
+ } else {
+ XceiverClientReply watchReply = new XceiverClientReply(null);
+ watchReply.setLogIndex(index);
+ f.complete(watchReply);
+ }
+ return f;
+ }).when(xceiverClient).watchForCommit(anyLong());
+
+ // Setup updateCommitInfosMap — no-op
+ doAnswer(invocation -> {
+ ByteBuffer src = invocation.getArgument(0);
+ Iterable<WriteOption> options = invocation.getArgument(1);
+ int size = src.remaining();
+ for (WriteOption option : options) {
+ if (option == StandardWriteOption.CLOSE) {
+ if (!receivedChunks.isEmpty()) {
+ receivedChunks.remove(receivedChunks.size() - 1);
+ }
+ src.position(src.limit());
+ return CompletableFuture.completedFuture(dataStreamReply(size));
+ }
+ }
+ int count = chunkCount.incrementAndGet();
+ if (count > chunkFailAfter && chunkFailure != null) {
+ CompletableFuture<DataStreamReply> failed = new CompletableFuture<>();
+ failed.completeExceptionally(chunkFailure.get());
+ return failed;
+ }
+ byte[] data = new byte[size];
+ src.get(data);
+ receivedChunks.add(data);
+ return CompletableFuture.completedFuture(dataStreamReply(data.length));
+ }).when(mockDataStreamOutput).writeAsync(any(ByteBuffer.class),
any(Iterable.class));
+
+ // Mock XceiverClientFactory
+ this.clientFactory = mock(XceiverClientFactory.class);
+
doReturn(xceiverClient).when(clientFactory).acquireClient(any(Pipeline.class),
anyBoolean());
+
doReturn(xceiverClient).when(clientFactory).acquireClient(any(Pipeline.class));
+ }
+
+ // --- Accessors ---
+
+ public Pipeline getPipeline() {
+ return pipeline;
+ }
+
+ public XceiverClientRatis getXceiverClient() {
+ return xceiverClient;
+ }
+
+ public XceiverClientFactory getClientFactory() {
+ return clientFactory;
+ }
+
+ public BlockID getBlockID() {
+ return blockID;
+ }
+
+ public List<byte[]> getReceivedChunks() {
+ return receivedChunks;
+ }
+
+ public List<ContainerCommandRequestProto> getReceivedPutBlocks() {
+ return receivedPutBlocks;
+ }
+
+ public int getWatchForCommitCount() {
+ return watchForCommitCount.get();
+ }
+
+ /** Concatenate all received chunks into a single byte array. */
+ public byte[] getAllReceivedData() {
+ int total = receivedChunks.stream().mapToInt(c -> c.length).sum();
+ byte[] result = new byte[total];
+ int pos = 0;
+ for (byte[] chunk : receivedChunks) {
+ System.arraycopy(chunk, 0, result, pos, chunk.length);
+ pos += chunk.length;
+ }
+ return result;
+ }
+
+ // --- Failure injection ---
+
+ public MockDatanodePipeline failChunkAfter(int n, Supplier<Throwable> err) {
+ this.chunkFailAfter = n;
+ this.chunkFailure = err;
+ return this;
+ }
+
+ public MockDatanodePipeline failPutBlockAfter(int n, Supplier<Throwable>
err) {
+ this.putBlockFailAfter = n;
+ this.putBlockFailure = err;
+ return this;
+ }
+
+ public MockDatanodePipeline failWatchAfter(int n, Supplier<Throwable> err) {
+ this.watchFailAfter = n;
+ this.watchFailure = err;
+ return this;
+ }
+
+ // --- Helpers ---
+
+ private static ContainerCommandResponseProto buildPutBlockResponse(BlockID
blockID) {
+ return ContainerCommandResponseProto.newBuilder()
+ .setCmdType(Type.PutBlock)
+ .setResult(Result.SUCCESS)
+ .setPutBlock(PutBlockResponseProto.newBuilder()
+ .setCommittedBlockLength(
+ GetCommittedBlockLengthResponseProto.newBuilder()
+ .setBlockID(blockID.getDatanodeBlockIDProtobuf())
+ .setBlockLength(0)
+ .build())
+ .build())
+ .build();
+ }
+
+ private static DataStreamReply dataStreamReply(long bytesWritten) {
+ DataStreamReply reply = mock(DataStreamReply.class);
+ when(reply.isSuccess()).thenReturn(true);
+ when(reply.getBytesWritten()).thenReturn(bytesWritten);
+ when(reply.getDataLength()).thenReturn(bytesWritten);
+ when(reply.getCommitInfos()).thenReturn(Collections.emptyList());
+ return reply;
+ }
+}
diff --git
a/hadoop-hdds/client/src/test/java/org/apache/hadoop/hdds/scm/storage/TestBlockDataStreamOutput.java
b/hadoop-hdds/client/src/test/java/org/apache/hadoop/hdds/scm/storage/TestBlockDataStreamOutput.java
new file mode 100644
index 00000000000..7634328c56f
--- /dev/null
+++
b/hadoop-hdds/client/src/test/java/org/apache/hadoop/hdds/scm/storage/TestBlockDataStreamOutput.java
@@ -0,0 +1,266 @@
+/*
+ * 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.hadoop.hdds.scm.storage;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletionException;
+import org.apache.commons.lang3.RandomUtils;
+import org.apache.hadoop.hdds.protocol.datanode.proto.ContainerProtos;
+import org.apache.hadoop.hdds.scm.OzoneClientConfig;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Unit tests for {@link BlockDataStreamOutput} exercised through the {@link
ByteBufferStreamOutput} interface with a
+ * mocked datanode pipeline.
+ */
+class TestBlockDataStreamOutput {
+
+ // Config: CHUNK=100, flush boundary=4 chunks (400B), window=5 chunks (500B)
+ private static final int CHUNK_SIZE = 100;
+ private static final long DS_FLUSH_SIZE = 400;
+ private static final long STREAM_WINDOW = 500;
+
+ private static OzoneClientConfig createConfig() {
+ OzoneClientConfig config = new OzoneClientConfig();
+ config.setDataStreamMinPacketSize(CHUNK_SIZE);
+ config.setDataStreamBufferFlushSize(DS_FLUSH_SIZE);
+ config.setStreamWindowSize(STREAM_WINDOW);
+ config.setStreamBufferSize(CHUNK_SIZE);
+ config.setStreamBufferFlushSize(DS_FLUSH_SIZE);
+ config.setStreamBufferMaxSize(2 * DS_FLUSH_SIZE);
+ config.setStreamBufferFlushDelay(false);
+ config.setChecksumType(ContainerProtos.ChecksumType.NONE);
+ config.setBytesPerChecksum(CHUNK_SIZE);
+ return config;
+ }
+
+ private BlockDataStreamOutput createStream(MockDatanodePipeline pipeline)
throws IOException {
+ List<StreamBuffer> bufferList = new ArrayList<>();
+ return new BlockDataStreamOutput(
+ pipeline.getBlockID(),
+ pipeline.getClientFactory(),
+ pipeline.getPipeline(),
+ createConfig(),
+ null, // no token
+ bufferList);
+ }
+
+ @Test
+ void writeSubChunkThenClose() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ try (BlockDataStreamOutput stream = createStream(pipeline)) {
+ byte[] data = randomBytes(50);
+ stream.write(ByteBuffer.wrap(data), 0, data.length);
+ // No chunk shipped yet — data is in currentBuffer
+ assertEquals(0, pipeline.getReceivedChunks().size(),
+ "No chunk should be shipped before close for sub-chunk write");
+ }
+ // After close: 1 chunk flushed + 1 putBlock
+ assertEquals(1, pipeline.getReceivedChunks().size());
+ assertEquals(1, pipeline.getReceivedPutBlocks().size());
+ }
+
+ @Test
+ void writeExactChunkThenClose() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ byte[] data = randomBytes(CHUNK_SIZE);
+ try (BlockDataStreamOutput stream = createStream(pipeline)) {
+ stream.write(ByteBuffer.wrap(data), 0, data.length);
+ // Exact chunk → shipped immediately (currentBuffer full)
+ assertEquals(1, pipeline.getReceivedChunks().size());
+ }
+ assertEquals(1, pipeline.getReceivedPutBlocks().size());
+ assertArrayEquals(data, pipeline.getAllReceivedData());
+ }
+
+ @Test
+ void writeFlushBoundaryTriggersPutBlock() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ byte[] data = randomBytes(400);
+ try (BlockDataStreamOutput stream = createStream(pipeline)) {
+ stream.write(ByteBuffer.wrap(data), 0, data.length);
+ // 4 chunks of 100B each → hits flush boundary → 1 putBlock
+ assertEquals(4, pipeline.getReceivedChunks().size());
+ assertEquals(1, pipeline.getReceivedPutBlocks().size(), "PutBlock should
trigger at flush boundary (400B)");
+ }
+ // Close adds another putBlock
+ assertEquals(2, pipeline.getReceivedPutBlocks().size());
+ }
+
+ @Test
+ void writeAcrossStreamWindowTriggersBackPressure() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ // Window = 500B = 5 chunks. Writing 500B should trigger back-pressure.
+ byte[] data = randomBytes(500);
+ try (BlockDataStreamOutput stream = createStream(pipeline)) {
+ stream.write(ByteBuffer.wrap(data), 0, data.length);
+ // 5 chunks written, putBlock at 400B boundary, back-pressure at 500B
+ // should have triggered watchForCommit
+ assertEquals(1, pipeline.getWatchForCommitCount(), "watchForCommit
should be called for back-pressure");
+ }
+ assertArrayEquals(data, pipeline.getAllReceivedData());
+ }
+
+ @Test
+ void hsyncFlushesAndWaitsForCommit() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ try (BlockDataStreamOutput stream = createStream(pipeline)) {
+ byte[] data = randomBytes(200);
+ stream.write(ByteBuffer.wrap(data), 0, data.length);
+ stream.hsync();
+ // 2 chunks, 1 putBlock (from hsync), watch called
+ assertEquals(2, pipeline.getReceivedChunks().size());
+ assertEquals(1, pipeline.getReceivedPutBlocks().size());
+ assertEquals(1, pipeline.getWatchForCommitCount());
+ }
+ }
+
+ //@Test - skipped as it fails now.
+ void hsyncPropagatesIOException() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ // Fail the first putBlock
+ pipeline.failPutBlockAfter(0, () -> new IOException("simulated putBlock
fail"));
+
+ BlockDataStreamOutput stream = createStream(pipeline);
+ byte[] data = randomBytes(200);
+ stream.write(ByteBuffer.wrap(data), 0, data.length);
+
+ // hsync should propagate the IOException from the failed putBlock
+ assertThrows(IOException.class, stream::hsync, "hsync() must propagate
IOException from failed putBlock");
+ stream.close();
+ }
+
+ //@Test - skipped as it fails now
+ void hsyncPropagatesWatchFailure() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ // Fail the first watchForCommit
+ pipeline.failWatchAfter(0,
+ () -> new IOException("simulated watch timeout"));
+
+ BlockDataStreamOutput stream = createStream(pipeline);
+ byte[] data = randomBytes(200);
+ stream.write(ByteBuffer.wrap(data), 0, data.length);
+
+ // hsync should propagate the watch failure
+ assertThrows(IOException.class, stream::hsync, "hsync() must propagate
IOException from failed watchForCommit");
+ stream.close();
+ }
+
+ @Test
+ void closeAfterWriteFailureThrows() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ // Fail on the 2nd chunk write
+ pipeline.failChunkAfter(1, () -> new IOException("chunk write failed due
to injected failure"));
+
+ BlockDataStreamOutput stream = createStream(pipeline);
+ byte[] data = randomBytes(50);
+ stream.write(ByteBuffer.wrap(data), 0, data.length); // ok, stays in buffer
+
+ byte[] data2 = randomBytes(CHUNK_SIZE + 50);
+ // This write will fill the buffer and trigger a chunk write that may fail.
+ // Close will surface the exception, which will be caused by the injected
exception,
+ // however based on thread scheduling, the actual exception we get back
from the stream
+ // is either a CompletionException caused by the injected exception or the
+ // injected exception itself so check for both.
+ stream.write(ByteBuffer.wrap(data2), 0, data2.length);
+ Throwable e = assertThrows(IOException.class, () -> stream.close());
+ if (e instanceof CompletionException) {
+ assertThat(e.getMessage()).contains("Failed to write chunk ");
+ boolean foundExpectedCause = false;
+ while (e.getCause() != null) {
+ e = e.getCause();
+ if (e instanceof IOException && e.getMessage().contains("chunk write
failed due to injected failure")) {
+ foundExpectedCause = true;
+ }
+ }
+ assertTrue(foundExpectedCause);
+ } else {
+ assertThat(e.getMessage().contains("chunk write failed due to injected
failure"));
+ }
+ }
+
+ @Test
+ void writeAfterCloseThrows() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ BlockDataStreamOutput stream = createStream(pipeline);
+ byte[] data = randomBytes(CHUNK_SIZE);
+ stream.write(ByteBuffer.wrap(data), 0, data.length);
+ stream.close();
+
+ IOException e = assertThrows(IOException.class,
+ () -> stream.write(ByteBuffer.wrap(data), 0, data.length),
+ "write() after close() should throw IOException");
+ assertThat(e.getMessage()).contains("has been closed");
+ }
+
+ @Test
+ void ackDataLengthTracksCommittedData() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ BlockDataStreamOutput stream = createStream(pipeline);
+ byte[] data = randomBytes(400);
+ stream.write(ByteBuffer.wrap(data), 0, data.length);
+ stream.close();
+
+ assertEquals(400, stream.getTotalAckDataLength(), "After close, all
written data should be acknowledged");
+ }
+
+ @Test
+ void chunkDataIntegrity() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ byte[] data = randomBytes(350);
+ try (BlockDataStreamOutput stream = createStream(pipeline)) {
+ stream.write(ByteBuffer.wrap(data), 0, data.length);
+ }
+ // 3 full chunks of 100B + 1 partial chunk of 50B
+ assertEquals(4, pipeline.getReceivedChunks().size());
+ assertArrayEquals(data, pipeline.getAllReceivedData(), "Concatenated chunk
data must match original input");
+ }
+
+ @Test
+ void putBlockContainsAllChunkMetadata() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ byte[] data = randomBytes(300);
+ try (BlockDataStreamOutput stream = createStream(pipeline)) {
+ stream.write(ByteBuffer.wrap(data), 0, data.length);
+ }
+
+ // One putBlock should have been sent
+ assertEquals(1, pipeline.getReceivedPutBlocks().size());
+
+ // 3 chunks with 300 bytes were sent and all chunk file name is correct.
+ List<ContainerProtos.ChunkInfo> chunksList =
+
pipeline.getReceivedPutBlocks().get(0).getPutBlock().getBlockData().getChunksList();
+
+ assertEquals(3, chunksList.size());
+ assertEquals(300,
chunksList.stream().mapToLong(ContainerProtos.ChunkInfo::getLen).sum());
+ assertTrue(chunksList.stream().allMatch(c ->
c.getChunkName().contains("_chunk_")));
+ }
+
+ private static byte[] randomBytes(int length) {
+ return RandomUtils.secure().randomBytes(length);
+ }
+}
diff --git
a/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/io/TestKeyDataStreamOutput.java
b/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/io/TestKeyDataStreamOutput.java
new file mode 100644
index 00000000000..4295c853c96
--- /dev/null
+++
b/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/io/TestKeyDataStreamOutput.java
@@ -0,0 +1,313 @@
+/*
+ * 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.hadoop.ozone.client.io;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.apache.commons.lang3.RandomUtils;
+import org.apache.hadoop.hdds.client.BlockID;
+import org.apache.hadoop.hdds.client.RatisReplicationConfig;
+import org.apache.hadoop.hdds.client.ReplicationConfig;
+import org.apache.hadoop.hdds.protocol.datanode.proto.ContainerProtos;
+import org.apache.hadoop.hdds.protocol.proto.HddsProtos.ReplicationFactor;
+import org.apache.hadoop.hdds.scm.OzoneClientConfig;
+import org.apache.hadoop.hdds.scm.XceiverClientFactory;
+import org.apache.hadoop.hdds.scm.container.common.helpers.ExcludeList;
+import
org.apache.hadoop.hdds.scm.container.common.helpers.StorageContainerException;
+import org.apache.hadoop.hdds.scm.pipeline.Pipeline;
+import org.apache.hadoop.hdds.scm.storage.MockDatanodePipeline;
+import org.apache.hadoop.ozone.om.helpers.OmKeyArgs;
+import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
+import org.apache.hadoop.ozone.om.helpers.OmKeyLocationInfo;
+import org.apache.hadoop.ozone.om.helpers.OmKeyLocationInfoGroup;
+import org.apache.hadoop.ozone.om.helpers.OpenKeySession;
+import org.apache.hadoop.ozone.om.protocol.OzoneManagerProtocol;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Unit tests for {@link KeyDataStreamOutput} exercised through the
+ * {@link org.apache.hadoop.hdds.scm.storage.ByteBufferStreamOutput} interface
with mocked datanode pipeline and OM
+ * client.
+ *
+ * <p>These tests verify the key-level stream behavior: block allocation,
hsync→OM integration, retry on container
+ * close, and atomic key commit.
+ *
+ */
+class TestKeyDataStreamOutput {
+
+ private static final int CHUNK_SIZE = 100;
+ private static final long DS_FLUSH_SIZE = 400;
+ private static final long STREAM_WINDOW = 500;
+ private static final long BLOCK_SIZE = 800;
+
+ private static OzoneClientConfig createConfig() {
+ OzoneClientConfig config = new OzoneClientConfig();
+ config.setDataStreamMinPacketSize(CHUNK_SIZE);
+ config.setDataStreamBufferFlushSize(DS_FLUSH_SIZE);
+ config.setStreamWindowSize(STREAM_WINDOW);
+ config.setStreamBufferSize(CHUNK_SIZE);
+ config.setStreamBufferFlushSize(DS_FLUSH_SIZE);
+ config.setStreamBufferMaxSize(2 * DS_FLUSH_SIZE);
+ config.setStreamBufferFlushDelay(false);
+ config.setChecksumType(ContainerProtos.ChecksumType.NONE);
+ config.setBytesPerChecksum(CHUNK_SIZE);
+ return config;
+ }
+
+ /**
+ * Creates a shared XceiverClientFactory that routes acquireClient calls
+ * to the correct MockDatanodePipeline based on pipeline ID.
+ */
+ private XceiverClientFactory
createSharedClientFactory(MockDatanodePipeline... pipelines) throws IOException
{
+ XceiverClientFactory factory = mock(XceiverClientFactory.class);
+ doAnswer(invocation -> {
+ Pipeline p = invocation.getArgument(0);
+ for (MockDatanodePipeline pipeline : pipelines) {
+ if (pipeline.getPipeline().getId().equals(p.getId())) {
+ return pipeline.getXceiverClient();
+ }
+ }
+ throw new IOException("Unknown pipeline: " + p.getId());
+ }).when(factory).acquireClient(any(Pipeline.class), anyBoolean());
+
+ doAnswer(invocation -> {
+ Pipeline p = invocation.getArgument(0);
+ for (MockDatanodePipeline pipeline : pipelines) {
+ if (pipeline.getPipeline().getId().equals(p.getId())) {
+ return pipeline.getXceiverClient();
+ }
+ }
+ throw new IOException("Unknown pipeline: " + p.getId());
+ }).when(factory).acquireClient(any(Pipeline.class));
+
+ return factory;
+ }
+
+ /**
+ * Creates a KeyDataStreamOutput with a mocked OM client that allocates
blocks from the given mocked pipelines.
+ * Each call to allocateBlock returns a block on the next pipeline in the
list.
+ */
+ private KeyDataStreamOutput createKeyStream(OzoneManagerProtocol omClient,
MockDatanodePipeline... pipelines)
+ throws Exception {
+
+ OzoneClientConfig config = createConfig();
+ ReplicationConfig replicationConfig =
RatisReplicationConfig.getInstance(ReplicationFactor.THREE);
+
+ OmKeyInfo keyInfo = new OmKeyInfo.Builder()
+ .setVolumeName("vol")
+ .setBucketName("bucket")
+ .setKeyName("testkey")
+ .setDataSize(BLOCK_SIZE)
+ .setReplicationConfig(replicationConfig)
+ .build();
+
+ OpenKeySession session = new OpenKeySession(1L, keyInfo, 0L);
+
+ XceiverClientFactory sharedFactory = createSharedClientFactory(pipelines);
+
+ KeyDataStreamOutput keyStream = new KeyDataStreamOutput(
+ config,
+ session,
+ sharedFactory,
+ omClient,
+ CHUNK_SIZE,
+ "test-request-id",
+ replicationConfig,
+ null, // uploadID
+ 0, // partNumber
+ false, // isMultipart
+ false // unsafeByteBufferConversion
+ );
+
+ // Pre-allocate the first block on mocked pipelines[0]
+ OmKeyLocationInfo firstBlock = new OmKeyLocationInfo.Builder()
+ .setBlockID(pipelines[0].getBlockID())
+ .setPipeline(pipelines[0].getPipeline())
+ .setLength(BLOCK_SIZE)
+ .build();
+ OmKeyLocationInfoGroup version = new OmKeyLocationInfoGroup(0,
Collections.singletonList(firstBlock));
+ keyStream.addPreallocateBlocks(version, 0);
+
+ return keyStream;
+ }
+
+ /**
+ * Creates a mock OM client that allocates blocks from mocked pipelines,
starting from the given index.
+ */
+ private OzoneManagerProtocol createOmClient(MockDatanodePipeline...
pipelines) throws IOException {
+ OzoneManagerProtocol omClient = mock(OzoneManagerProtocol.class);
+ AtomicInteger allocIndex = new AtomicInteger(0);
+ doAnswer(invocation -> {
+ int idx = allocIndex.getAndIncrement();
+ if (idx >= pipelines.length) {
+ throw new IOException("No more blocks to allocate");
+ }
+ MockDatanodePipeline pipeline = pipelines[idx];
+ return new OmKeyLocationInfo.Builder()
+ .setBlockID(pipeline.getBlockID())
+ .setPipeline(pipeline.getPipeline())
+ .setLength(BLOCK_SIZE)
+ .build();
+ }).when(omClient).allocateBlock(any(OmKeyArgs.class), anyLong(),
any(ExcludeList.class));
+ return omClient;
+ }
+
+ @Test
+ void writeAndCloseCommitsKey() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ OzoneManagerProtocol omClient = createOmClient(pipeline);
+
+ try (KeyDataStreamOutput stream = createKeyStream(omClient, pipeline)) {
+ writeRandom(stream, 300);
+ }
+
+ verify(omClient, times(1)).commitKey(any(OmKeyArgs.class), anyLong());
+ }
+
+ @Test
+ void writeCrossBlockBoundary() throws Exception {
+ MockDatanodePipeline pipeline1 = new MockDatanodePipeline(new BlockID(1,
1));
+ MockDatanodePipeline pipeline2 = new MockDatanodePipeline(new BlockID(2,
2));
+
+ // OM returns pipeline2 when allocateBlock is called
+ OzoneManagerProtocol omClient = createOmClient(pipeline2);
+
+ // The first block (pipeline1) has BLOCK_SIZE=800 capacity. Both mocks
must be known to the shared client factory.
+ try (KeyDataStreamOutput stream = createKeyStream(omClient, pipeline1,
pipeline2)) {
+ writeRandom(stream, 850);
+ }
+
+ // allocateBlock should have been called for the second block
+ verify(omClient, times(1)).allocateBlock(any(OmKeyArgs.class), anyLong(),
any(ExcludeList.class));
+ verify(omClient, times(1)).commitKey(any(OmKeyArgs.class), anyLong());
+
+ // pipeline1 should have received 800 bytes, pipeline2 should have
received 50
+ assertEquals(800, totalReceived(pipeline1));
+ assertEquals(50, totalReceived(pipeline2));
+ }
+
+ @Test
+ void hsyncCallsOmHsyncKey() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ OzoneManagerProtocol omClient = createOmClient(pipeline);
+
+ try (KeyDataStreamOutput stream = createKeyStream(omClient, pipeline)) {
+ writeRandom(stream, 200);
+ stream.hsync();
+
+ verify(omClient, times(1)).hsyncKey(any(OmKeyArgs.class), anyLong());
+ }
+ }
+
+// @Test - skipped as it fails now
+ void hsyncWithBlockErrorDoesNotCallOmHsync() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ // First putBlock will fail
+ pipeline.failPutBlockAfter(0, () -> new IOException("putBlock failed"));
+
+ OzoneManagerProtocol omClient = createOmClient(pipeline);
+
+ KeyDataStreamOutput stream = createKeyStream(omClient, pipeline);
+ writeRandom(stream, 200);
+
+ // hsync should throw because the block-level flush failed
+ assertThrows(IOException.class, stream::hsync, "hsync() must throw when
block-level flush fails");
+
+ // OM hsyncKey must NOT have been called — data was not committed
+ verify(omClient, never()).hsyncKey(any(OmKeyArgs.class), anyLong());
+
+ stream.close();
+ }
+
+ @Test
+ void containerCloseTriggersRetryOnNewBlock() throws Exception {
+ MockDatanodePipeline pipeline1 = new MockDatanodePipeline(new BlockID(1,
1));
+ MockDatanodePipeline pipeline2 = new MockDatanodePipeline(new BlockID(2,
2));
+
+ // First pipeline: putBlock fails with ContainerNotOpen
+ pipeline1.failPutBlockAfter(0,
+ () -> new StorageContainerException("Container closed",
ContainerProtos.Result.CLOSED_CONTAINER_IO));
+
+ OzoneManagerProtocol omClient = createOmClient(pipeline2);
+
+ try (KeyDataStreamOutput stream = createKeyStream(omClient, pipeline1,
pipeline2)) {
+ writeRandom(stream, 200);
+ // The flush on close will hit the container closed error, trigger
exception handling, allocate a new block on
+ // pipeline2, and retry to write there.
+ }
+
+ // allocateBlock should have been called (for the retry block)
+ verify(omClient).allocateBlock(any(OmKeyArgs.class), anyLong(),
any(ExcludeList.class));
+ verify(omClient).commitKey(any(OmKeyArgs.class), anyLong());
+ }
+
+ @Test
+ void multipleHsyncsCallOmAtLeastOnce() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ OzoneManagerProtocol omClient = createOmClient(pipeline);
+
+ try (KeyDataStreamOutput stream = createKeyStream(omClient, pipeline)) {
+ writeRandom(stream, 200);
+ stream.hsync();
+
+ writeRandom(stream, 200);
+ stream.hsync();
+
+ // hsyncKey is called at least once; the second call is skipped because
the block ID hasn't changed
+ // (OM optimization at BlockDataStreamOutputEntryPool.hsyncKey line 172).
+ verify(omClient, times(1)).hsyncKey(any(OmKeyArgs.class), anyLong());
+
+ // But both hsyncs should have flushed data to the datanode
+ assertEquals(400, totalReceived(pipeline));
+ }
+ }
+
+ @Test
+ void writeAfterCloseThrows() throws Exception {
+ MockDatanodePipeline pipeline = new MockDatanodePipeline();
+ KeyDataStreamOutput stream = createKeyStream(createOmClient(pipeline),
pipeline);
+
+ writeRandom(stream, 100);
+ stream.close();
+
+ assertThrows(IOException.class, () -> writeRandom(stream, 100), "write()
after close() should throw");
+ }
+
+ // --- Helpers ---
+
+ private static int totalReceived(MockDatanodePipeline pipeline) {
+ return pipeline.getReceivedChunks().stream().mapToInt(c -> c.length).sum();
+ }
+
+ private static void writeRandom(KeyDataStreamOutput stream, int length)
throws IOException {
+ stream.write(ByteBuffer.wrap(RandomUtils.secure().randomBytes(length)), 0,
length);
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
