[GEODE-10646] Reduce up-front buffer allocation in the Lucene file output stream

geode-lucene/src/main/java/org/apache/geode/cache/lucene/internal/filesystem/FileOutputStream.java

@@ -17,41 +17,52 @@
 
 import java.io.IOException;
 import java.io.OutputStream;
-import java.nio.ByteBuffer;
-import java.util.Arrays;
 
+/**
+ * Buffers the current chunk in fixed-size segments that are allocated as they are first needed and
+ * reused for every later chunk of the stream. A chunk is copied into an array of its exact size
+ * when it is written.
+ */
 class FileOutputStream extends OutputStream {
 
+  private static final int SEGMENT_SIZE = 8 * 1024;
+
   private final File file;
-  private ByteBuffer buffer;
+  private final int chunkSize;
+  private final int segmentSize;
+  private byte[][] segments;
+  private int position;
   private boolean open = true;
   private long length;
   private int chunks;
 
   public FileOutputStream(final File file) {
     this.file = file;
-    buffer = ByteBuffer.allocate(file.getChunkSize());
+    chunkSize = file.getChunkSize();
+    segmentSize = Math.min(SEGMENT_SIZE, chunkSize);
+    segments = new byte[(chunkSize + segmentSize - 1) / segmentSize][];
     length = file.length;
     chunks = file.chunks;
-    if (chunks > 0 && file.length % file.getChunkSize() != 0) {
+    if (chunks > 0 && file.length % chunkSize != 0) {
       // If the last chunk was incomplete, we're going to update it
       // rather than add a new chunk. This guarantees that all chunks
       // are full except for the last chunk.
       chunks--;
       byte[] previousChunkData = file.getFileSystem().getChunk(file, chunks);
-      buffer.put(previousChunkData);
+      buffer(previousChunkData, 0, previousChunkData.length);
     }
   }
 
   @Override
   public void write(final int b) throws IOException {
     assertOpen();
 
-    if (buffer.remaining() == 0) {
+    if (position == chunkSize) {
       flushBuffer();
     }
 
-    buffer.put((byte) b);
+    segment(position / segmentSize)[position % segmentSize] = (byte) b;
+    position++;
     length++;
   }
 
@@ -60,15 +71,14 @@ public void write(final byte[] b, int off, int len) throws IOException {
     assertOpen();
 
     while (len > 0) {
-      if (buffer.remaining() == 0) {
+      if (position == chunkSize) {
         flushBuffer();
       }
 
-      final int min = Math.min(buffer.remaining(), len);
-      buffer.put(b, off, min);
-      off += min;
-      len -= min;
-      length += min;
+      final int copied = buffer(b, off, len);
+      off += copied;
+      len -= copied;
+      length += copied;
     }
   }
 
@@ -81,14 +91,46 @@ public void close() throws IOException {
       file.chunks = chunks;
       file.getFileSystem().updateFile(file);
       open = false;
-      buffer = null;
+      segments = null;
+    }
+  }
+
+  /**
+   * Copies bytes into the current chunk, up to the end of the chunk.
+   *
+   * @return the number of bytes copied
+   */
+  private int buffer(final byte[] b, int off, final int len) {
+    final int limit = Math.min(len, chunkSize - position);
+    int copied = 0;
+    while (copied < limit) {
+      final int offsetInSegment = position % segmentSize;
+      final int count = Math.min(limit - copied, segmentSize - offsetInSegment);
+      System.arraycopy(b, off, segment(position / segmentSize), offsetInSegment, count);
+      off += count;
+      copied += count;
+      position += count;
     }
+    return copied;
+  }
+
+  private byte[] segment(final int index) {
+    byte[] segment = segments[index];
+    if (segment == null) {
+      segment = new byte[segmentSize];
+      segments[index] = segment;
+    }
+    return segment;
   }
 
   private void flushBuffer() {
-    byte[] chunk = Arrays.copyOfRange(buffer.array(), buffer.arrayOffset(), buffer.position());
+    final byte[] chunk = new byte[position];
+    for (int copied = 0; copied < position; copied += segmentSize) {
+      System.arraycopy(segments[copied / segmentSize], 0, chunk, copied,
+          Math.min(segmentSize, position - copied));
+    }
     file.getFileSystem().putChunk(file, chunks++, chunk);
-    buffer.rewind();
+    position = 0;
   }
 
   private void assertOpen() throws IOException {

geode-lucene/src/test/java/org/apache/geode/cache/lucene/internal/filesystem/FileOutputStreamJUnitTest.java

@@ -0,0 +1,215 @@
+/*
+ * 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.geode.cache.lucene.internal.filesystem;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assume.assumeTrue;
+import static org.mockito.Mockito.mock;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.lang.management.ManagementFactory;
+import java.util.Random;
+import java.util.concurrent.ConcurrentHashMap;
+
+import com.sun.management.ThreadMXBean;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+
+import org.apache.geode.test.junit.categories.LuceneTest;
+
+@Category({LuceneTest.class})
+public class FileOutputStreamJUnitTest {
+
+  private static final int CHUNK_SIZE = FileSystem.CHUNK_SIZE;
+
+  private FileSystem system;
+
+  @Before
+  public void setUp() {
+    system = new FileSystem(new ConcurrentHashMap<>(), mock(FileSystemStats.class));
+  }
+
+  /**
+   * A test that writing a small file allocates far less than one chunk.
+   */
+  @Test
+  public void testSmallFileAllocatesLessThanChunkSize() throws IOException {
+    ThreadMXBean threadBean = (ThreadMXBean) ManagementFactory.getThreadMXBean();
+    assumeTrue(threadBean.isThreadAllocatedMemorySupported()
+        && threadBean.isThreadAllocatedMemoryEnabled());
+
+    // Load classes and warm up the mocks before measuring
+    writeSmallFile(system.createFile("warmup"));
+
+    File file = system.createFile("small");
+    long threadId = Thread.currentThread().getId();
+    long before = threadBean.getThreadAllocatedBytes(threadId);
+    writeSmallFile(file);
+    long allocated = threadBean.getThreadAllocatedBytes(threadId) - before;
+
+    assertTrue("Allocated " + allocated + " bytes to write a small file",
+        allocated < CHUNK_SIZE / 4);
+  }
+
+  /**
+   * A test that writing a file larger than one chunk allocates no more than its length plus one
+   * chunk of buffer space.
+   */
+  @Test
+  public void testLargeFileAllocatesAtMostOneChunkBeyondItsLength() throws IOException {
+    ThreadMXBean threadBean = (ThreadMXBean) ManagementFactory.getThreadMXBean();
+    assumeTrue(threadBean.isThreadAllocatedMemorySupported()
+        && threadBean.isThreadAllocatedMemoryEnabled());
+    int fileLength = CHUNK_SIZE + CHUNK_SIZE / 2;
+
+    // Load classes and warm up the mocks before measuring
+    writeFile(system.createFile("warmup"), fileLength);
+
+    File file = system.createFile("large");
+    long threadId = Thread.currentThread().getId();
+    long before = threadBean.getThreadAllocatedBytes(threadId);
+    writeFile(file, fileLength);
+    long allocated = threadBean.getThreadAllocatedBytes(threadId) - before;
+
+    assertTrue("Allocated " + allocated + " bytes to write a " + fileLength + " byte file",
+        allocated < fileLength + CHUNK_SIZE + CHUNK_SIZE / 16);
+  }
+
+  /**
+   * A test that files written with a random mix of single bytes, arrays and appends read back
+   * correctly, with every chunk except the last one full.
+   */
+  @Test
+  public void testRandomWritesReadBackWithFullChunks() throws IOException {
+    long seed = System.nanoTime();
+    Random random = new Random(seed);
+
+    for (int iteration = 0; iteration < 50; iteration++) {
+      FileSystem fileSystem =
+          new FileSystem(new ConcurrentHashMap<>(), mock(FileSystemStats.class));
+      File file = fileSystem.createFile("random");
+      ByteArrayOutputStream expected = new ByteArrayOutputStream();
+
+      int sessions = 1 + random.nextInt(4);
+      for (int session = 0; session < sessions; session++) {
+        OutputStream outputStream = file.getOutputStream();
+        int writes = 1 + random.nextInt(12);
+        for (int i = 0; i < writes; i++) {
+          writeRandomly(random, outputStream, expected);
+        }
+        outputStream.close();
+
+        assertFileContents("seed " + seed + ", iteration " + iteration + ", session " + session,
+            fileSystem, file, expected.toByteArray());
+      }
+    }
+  }
+
+  private void writeSmallFile(File file) throws IOException {
+    OutputStream outputStream = file.getOutputStream();
+    outputStream.write(new byte[100]);
+    outputStream.close();
+  }
+
+  private void writeFile(File file, int fileLength) throws IOException {
+    byte[] data = new byte[1000];
+    OutputStream outputStream = file.getOutputStream();
+    for (int written = 0; written < fileLength; written += data.length) {
+      outputStream.write(data, 0, Math.min(data.length, fileLength - written));
+    }
+    outputStream.close();
+  }
+
+  /**
+   * Writes at least one byte to both streams, choosing among single bytes, arrays with offsets,
+   * zero-length writes and writes that end on or next to a chunk boundary.
+   */
+  private void writeRandomly(Random random, OutputStream outputStream,
+      ByteArrayOutputStream expected) throws IOException {
+    int toBoundary = CHUNK_SIZE - expected.size() % CHUNK_SIZE;
+    boolean large = expected.size() < 4 * CHUNK_SIZE;
+    int kind = random.nextInt(8);
+    if (!large && kind >= 4) {
+      kind = random.nextInt(4);
+    }
+
+    int length;
+    switch (kind) {
+      case 0:
+        int b = random.nextInt(256);
+        outputStream.write(b);
+        expected.write(b);
+        return;
+      case 1:
+        outputStream.write(new byte[10], 3, 0);
+        length = 1;
+        break;
+      case 2:
+      case 3:
+        length = 1 + random.nextInt(20_000);
+        break;
+      case 4:
+        length = toBoundary;
+        break;
+      case 5:
+        length = toBoundary > 1 ? toBoundary - 1 : 1;
+        break;
+      case 6:
+        length = toBoundary + 1;
+        break;
+      default:
+        length = 1 + random.nextInt(CHUNK_SIZE + 20_000);
+        break;
+    }
+
+    int offset = random.nextInt(100);
+    byte[] data = new byte[offset + length + random.nextInt(100)];
+    random.nextBytes(data);
+    outputStream.write(data, offset, length);
+    expected.write(data, offset, length);
+  }
+
+  private void assertFileContents(String context, FileSystem fileSystem, File file,
+      byte[] expected) throws IOException {
+    assertEquals(context, expected.length, file.getLength());
+
+    int expectedChunks = Math.max(1, (expected.length + CHUNK_SIZE - 1) / CHUNK_SIZE);
+    assertEquals(context, expectedChunks, file.chunks);
+    for (int i = 0; i < file.chunks; i++) {
+      int expectedLength =
+          i < file.chunks - 1 ? CHUNK_SIZE : expected.length - (file.chunks - 1) * CHUNK_SIZE;
+      assertEquals(context + ", chunk " + i, expectedLength, fileSystem.getChunk(file, i).length);
+    }
+
+    byte[] actual = new byte[expected.length];
+    try (InputStream inputStream = file.getInputStream()) {
+      int read = 0;
+      int count;
+      while (read < actual.length
+          && (count = inputStream.read(actual, read, actual.length - read)) > 0) {
+        read += count;
+      }
+      assertEquals(context, expected.length, read);
+      assertEquals(context, -1, inputStream.read());
+    }
+    assertArrayEquals(context, expected, actual);
+  }
+}