TEZ-4760: IFile.Writer: buffer per-record writes before the compression and checksum streams

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}
+   *
+   * <p>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;
+    }
+  }
+}

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<Integer> 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));
+  }
+}

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() {
     }
   }
 
+
   /**
    * <code>IFile.Reader</code> to read intermediate map-outputs.
    */