diff --git a/tez-common/src/main/java/org/apache/tez/common/io/NonSyncBufferedOutputStream.java b/tez-common/src/main/java/org/apache/tez/common/io/NonSyncBufferedOutputStream.java new file mode 100644 index 0000000000..b7f8254b6e --- /dev/null +++ b/tez-common/src/main/java/org/apache/tez/common/io/NonSyncBufferedOutputStream.java @@ -0,0 +1,89 @@ +/* + * 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.common.io; + +import java.io.FilterOutputStream; +import java.io.IOException; +import java.io.OutputStream; + +/** + * A thread-not-safe version of BufferedOutputStream, which removes all + * synchronized modifiers. Use it where a stream is confined to one thread and + * the writes are small and frequent: the lock is a measurable share of the cost + * when the work per call is tiny. + */ +public class NonSyncBufferedOutputStream extends FilterOutputStream { + private final byte[] buf; + private int count; + + public NonSyncBufferedOutputStream(OutputStream stream, int size) { + super(stream); + if (size <= 0) { + throw new IllegalArgumentException("Buffer size <= 0"); + } + this.buf = new byte[size]; + } + + /** + * {@inheritDoc} + */ + @Override + public void write(int b) throws IOException { + if (count == buf.length) { + drain(); + } + buf[count++] = (byte) b; + } + + /** + * {@inheritDoc} + * + *

A write at least as large as the buffer goes straight to the underlying + * stream, so a large record is not copied twice. + */ + @Override + public void write(byte[] b, int off, int len) throws IOException { + if (len >= buf.length) { + drain(); + out.write(b, off, len); + return; + } + if (len > buf.length - count) { + drain(); + } + System.arraycopy(b, off, buf, count, len); + count += len; + } + + /** + * {@inheritDoc} + */ + @Override + public void flush() throws IOException { + drain(); + out.flush(); + } + + private void drain() throws IOException { + if (count > 0) { + out.write(buf, 0, count); + count = 0; + } + } +} diff --git a/tez-common/src/test/java/org/apache/tez/common/io/TestNonSyncBufferedOutputStream.java b/tez-common/src/test/java/org/apache/tez/common/io/TestNonSyncBufferedOutputStream.java new file mode 100644 index 0000000000..a1b21ae031 --- /dev/null +++ b/tez-common/src/test/java/org/apache/tez/common/io/TestNonSyncBufferedOutputStream.java @@ -0,0 +1,136 @@ +/* + * 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.common.io; + +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 java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +import org.junit.jupiter.api.Test; + +public class TestNonSyncBufferedOutputStream { + + private static final int SIZE = 64; + + /** Records the size of every write that reaches the underlying stream. */ + private static final class RecordingStream extends OutputStream { + private final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + private final List writes = new ArrayList<>(); + + @Override + public void write(int b) { + bytes.write(b); + writes.add(1); + } + + @Override + public void write(byte[] b, int off, int len) { + bytes.write(b, off, len); + writes.add(len); + } + } + + private static byte[] pattern(int len, int seed) { + byte[] b = new byte[len]; + for (int i = 0; i < len; i++) { + b[i] = (byte) (seed + i); + } + return b; + } + + @Test + public void testWriteOfExactlyBufferSizeBypassesTheBuffer() throws IOException { + RecordingStream sink = new RecordingStream(); + byte[] data = pattern(SIZE, 1); + try (NonSyncBufferedOutputStream out = new NonSyncBufferedOutputStream(sink, SIZE)) { + out.write(data, 0, data.length); + // len >= buffer: handed to the underlying stream as it is, before any flush + assertEquals(List.of(SIZE), sink.writes); + } + assertArrayEquals(data, sink.bytes.toByteArray()); + } + + @Test + public void testWriteLargerThanTheBufferBypassesTheBuffer() throws IOException { + RecordingStream sink = new RecordingStream(); + byte[] data = pattern(SIZE * 3 + 7, 2); + try (NonSyncBufferedOutputStream out = new NonSyncBufferedOutputStream(sink, SIZE)) { + out.write(data, 0, data.length); + assertEquals(List.of(data.length), sink.writes); + } + assertArrayEquals(data, sink.bytes.toByteArray()); + } + + @Test + public void testWritesThatCrossTheBufferAreDrainedOnce() throws IOException { + RecordingStream sink = new RecordingStream(); + byte[] first = pattern(SIZE - 1, 3); + byte[] second = pattern(2, 4); + try (NonSyncBufferedOutputStream out = new NonSyncBufferedOutputStream(sink, SIZE)) { + out.write(first, 0, first.length); + assertEquals(List.of(), sink.writes, "the first write fits, so nothing reaches the stream"); + out.write(second, 0, second.length); + // does not fit in what is left: the buffer is drained, then the new bytes are kept + assertEquals(List.of(SIZE - 1), sink.writes); + out.flush(); + assertEquals(List.of(SIZE - 1, 2), sink.writes); + } + byte[] expected = Arrays.copyOf(first, first.length + second.length); + System.arraycopy(second, 0, expected, first.length, second.length); + assertArrayEquals(expected, sink.bytes.toByteArray()); + } + + @Test + public void testSingleByteWritesFillAndDrainTheBuffer() throws IOException { + RecordingStream sink = new RecordingStream(); + byte[] data = pattern(SIZE + 1, 5); + try (NonSyncBufferedOutputStream out = new NonSyncBufferedOutputStream(sink, SIZE)) { + for (byte b : data) { + out.write(b); + } + // the byte after the buffer is full drains the full buffer first + assertEquals(List.of(SIZE), sink.writes); + } + assertArrayEquals(data, sink.bytes.toByteArray()); + } + + @Test + public void testCloseDrainsTheBuffer() throws IOException { + RecordingStream sink = new RecordingStream(); + byte[] data = pattern(3, 6); + NonSyncBufferedOutputStream out = new NonSyncBufferedOutputStream(sink, SIZE); + out.write(data, 0, data.length); + assertEquals(List.of(), sink.writes); + out.close(); + assertArrayEquals(data, sink.bytes.toByteArray()); + } + + @Test + public void testNonPositiveBufferSizeIsRejected() { + assertThrows(IllegalArgumentException.class, + () -> new NonSyncBufferedOutputStream(new ByteArrayOutputStream(), 0)); + } +} diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/sort/impl/IFile.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/sort/impl/IFile.java index 484fca6ea7..90d3cd3c67 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/sort/impl/IFile.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/sort/impl/IFile.java @@ -46,6 +46,8 @@ import org.apache.hadoop.io.serializer.Serialization; import org.apache.hadoop.io.serializer.Serializer; import org.apache.tez.common.counters.TezCounter; +import org.apache.tez.common.io.NonSyncBufferedOutputStream; +import org.apache.tez.common.io.NonSyncDataOutputStream; import org.apache.tez.runtime.library.common.task.local.output.TezTaskOutput; import org.apache.tez.runtime.library.utils.BufferUtils; import org.apache.tez.runtime.library.utils.CodecUtils; @@ -261,6 +263,8 @@ public Path getOutputPath() { @InterfaceStability.Unstable @SuppressWarnings({"unchecked", "rawtypes"}) public static class Writer { + private static final int WRITE_BUFFER_SIZE = 64 * 1024; + protected DataOutputStream out; boolean ownOutputStream = false; long start = 0; @@ -356,17 +360,27 @@ void setupOutputStream(CompressionCodec codec) throws IOException { if (this.compressor != null) { this.compressor.reset(); this.compressedOut = CodecUtils.createOutputStream(codec, checksumOut, compressor); - this.out = new FSDataOutputStream(this.compressedOut, null); + this.out = bufferedOutput(this.compressedOut); this.compressOutput = true; } else { LOG.warn("Could not obtain compressor from CodecPool"); - this.out = new FSDataOutputStream(checksumOut,null); + this.out = bufferedOutput(checksumOut); } } else { - this.out = new FSDataOutputStream(checksumOut,null); + this.out = bufferedOutput(checksumOut); } } + /** + * The writer emits several small writes per record (vint lengths, key, value). Unbuffered, + * each one is a full call into the codec stream (synchronized Compressor.setInput and + * friends) or the checksum stream. The buffer is unsynchronized: a Writer is confined to + * one thread. + */ + private static DataOutputStream bufferedOutput(OutputStream stream) { + return new NonSyncDataOutputStream(new NonSyncBufferedOutputStream(stream, WRITE_BUFFER_SIZE)); + } + public Writer(Serialization keySerialization, Serialization valSerialization, FileSystem fs, Path file) throws IOException { this(keySerialization, valSerialization, fs, file, null, null, null, null, null); } @@ -406,6 +420,8 @@ public void close() throws IOException { if (ownOutputStream) { out.close(); } else { + // Drain the write buffer before finishing the streams below it. + out.flush(); if (compressOutput) { // Flush compressedOut.finish(); @@ -634,6 +650,7 @@ public long getCompressedLength() { } } + /** * IFile.Reader to read intermediate map-outputs. */