more efficient buffer usage, add tests

Author: wolfchimneyrock

bmq-sdk/src/main/java/com/bloomberg/bmq/impl/infr/io/ByteBufferOutputStream.java

@@ -16,31 +16,67 @@
 package com.bloomberg.bmq.impl.infr.io;
 
 import com.bloomberg.bmq.impl.infr.util.Argument;
-import com.bloomberg.bmq.impl.infr.util.Limits;
 import java.io.DataOutput;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.lang.invoke.MethodHandles;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+/**
+ * An output stream of ByteBuffers
+ *
+ * <p>invariants of this structure are listed below.
+ *
+ * <p>1. a stream that has never been written to will have no buffers yet.
+ *
+ * <p>2. buffers are added on demand.
+ *
+ * <p>3. writeBuffers() or big array write(), the current buffer is sliced and the remainder is
+ * added after as the current buffer.
+ *
+ * <p>4. writeBuffers() appends duplicated buffers wholesale instead of copying.
+ *
+ * <p>5. writeBuffers() ByteBuffers from outside should always be either unflipped or a wrapped
+ * array.
+ *
+ * <p>7. write() byte arrays larger than the buffer size get wrapped, smaller ones get copied in one
+ * piece.
+ *
+ * <p>8. as a result of 7, byte arrays can always be read fully in one read().
+ *
+ * <p>9. totalBytes is kept up to date as fields / buffers are written in.
+ *
+ * <p>10. the current append buffer is always the last buffer in bbArray.
+ */
 public class ByteBufferOutputStream extends OutputStream implements DataOutput {
     static final Logger logger = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
 
     private ArrayList<ByteBuffer> bbArray;
-    private ByteBuffer currentBuffer;
-
-    private int currentBufferIndex = 0;
     private int bufSize;
-    private int prevBuffersNumBytes;
-
+    private int totalBytes;
     private boolean isOpen;
 
     private static final int KB = 1024;
     private static final int DEFAULT_BUF_SIZE = 4 * KB;
+    private static final int BIG_BUF_SIZE = DEFAULT_BUF_SIZE;
+
+    // small buffer size should be adequete for many event types.
+    private static final int SMALL_BUF_SIZE = 512;
+
+    public static ByteBufferOutputStream smallBlocks() {
+        return new ByteBufferOutputStream(SMALL_BUF_SIZE);
+    }
+
+    public static ByteBufferOutputStream bigBlocks() {
+        return new ByteBufferOutputStream(BIG_BUF_SIZE);
+    }
 
     public ByteBufferOutputStream() {
         init(DEFAULT_BUF_SIZE);
@@ -53,20 +89,39 @@ public ByteBufferOutputStream(int bufSize) {
     private void init(int bufSize) {
         bbArray = new ArrayList<>();
         this.bufSize = bufSize;
-        currentBuffer = ByteBuffer.allocate(bufSize);
-        bbArray.add(currentBuffer);
-        currentBufferIndex = 0;
         isOpen = true;
-        prevBuffersNumBytes = 0;
+        totalBytes = 0;
+    }
+
+    private int availableCapacity() {
+        if (bbArray.isEmpty()) {
+            return 0;
+        }
+        return getCurrent().remaining();
+    }
+
+    private void ensureCapacity(int size) {
+        int capacity = availableCapacity();
+        if (size > capacity) {
+            addBuffer(Math.max(bufSize, size));
+        }
+    }
+
+    private void addRemainderOrNew(ByteBuffer remainder) {
+        if (remainder != null) {
+            bbArray.add(remainder);
+        } else {
+            addBuffer();
+        }
     }
 
     @Override
     public void write(int b) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
-        if (currentBuffer.remaining() < Limits.BYTE_SIZE) getNewBuffer();
-
-        currentBuffer.put((byte) b);
+        ensureCapacity(1);
+        getCurrent().put((byte) b);
+        totalBytes += 1;
     }
 
     @Override
@@ -80,16 +135,18 @@ public void write(byte[] ba, int offset, int length) throws IOException {
 
         if (length <= 0 || (length > ba.length - offset)) return;
 
-        int numBytesLeft = length;
-        while (true) {
-            int numToWrite = Math.min(numBytesLeft, currentBuffer.remaining());
-            currentBuffer.put(ba, offset, numToWrite);
-            numBytesLeft -= numToWrite;
-            offset += numToWrite;
-
-            if (numBytesLeft > 0) getNewBuffer();
-            else break;
+        // if its too big for a pre-defined buffer, just wrap it instead
+        if (length > bufSize) {
+            ByteBuffer remainder = maybeSliceCurrent();
+            ByteBuffer buf = ByteBuffer.wrap(ba, offset, length);
+            buf.position(buf.limit());
+            bbArray.add(buf);
+            addRemainderOrNew(remainder);
+        } else {
+            ensureCapacity(length);
+            getCurrent().put(ba, offset, length);
         }
+        totalBytes += length;
     }
 
     /**
@@ -98,9 +155,12 @@ public void write(byte[] ba, int offset, int length) throws IOException {
      * <p>The bbos can continue to be written to.
      */
     public ByteBuffer[] peek() {
-        return bbArray.stream()
-                .map(bb -> (ByteBuffer) (bb.duplicate().flip()))
-                .toArray(ByteBuffer[]::new);
+        ByteBuffer[] duplicates =
+                bbArray.stream().map(ByteBuffer::duplicate).toArray(ByteBuffer[]::new);
+        for (ByteBuffer b : duplicates) {
+            b.flip();
+        }
+        return duplicates;
     }
 
     private ArrayList<ByteBuffer> buffers() {
@@ -116,47 +176,68 @@ public void writeBoolean(boolean v) throws IOException {
     public void writeByte(int v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
-        if (currentBuffer.remaining() < Limits.BYTE_SIZE) getNewBuffer();
-
-        currentBuffer.put((byte) v);
+        ensureCapacity(1);
+        getCurrent().put((byte) v);
+        totalBytes += 1;
     }
 
     @Override
     public void writeBytes(String s) throws IOException {
         throw new UnsupportedOperationException();
     }
 
+    // a buffer that has never been put() to nor flipped
     private boolean bufferIsFresh(ByteBuffer b) {
-        // a buffer that has never been put() to nor flipped
         return b.position() == 0 && b.limit() == b.capacity();
     }
 
-    public void writeBuffer(ByteBuffer b) throws IOException {
-        if (!isOpen) throw new IOException("Stream closed");
-        if (bufferIsFresh(b)) return;
-
-        boolean currentIsFresh = bufferIsFresh(currentBuffer);
-        // remove the currentBuffer if it is fresh
-        // we'll put it back at the end after adding other's buffers
-        // to avoid allocating a new one in that case
-        if (currentIsFresh) {
-            bbArray.remove(currentBufferIndex);
-        }
-        ByteBuffer buf = b.duplicate();
-        if (buf.limit() != buf.capacity()) {
-            // it has already been flipped - unflip it
-            int newPosition = buf.limit();
-            buf.limit(buf.capacity());
-            buf.position(newPosition);
+    private ByteBuffer getCurrent() {
+        if (bbArray.size() == 0) return null;
+        return bbArray.get(bbArray.size() - 1);
+    }
+
+    private ByteBuffer maybeSliceCurrent() {
+        ByteBuffer current = getCurrent();
+        if (current != null) {
+            if (bufferIsFresh(current)) {
+                // current has never been written to,
+                // remove it from bbArray and return it
+                // whole as the remainder
+                bbArray.remove(bbArray.size() - 1);
+                return current;
+            }
+            // a remainder slice should be meaningfully sized - at least as big as the ByteBuffer
+            // overhead
+            if (current.remaining() >= 16) {
+                ByteBuffer remainder = current.slice();
+                return remainder;
+            }
         }
-        prevBuffersNumBytes += buf.position();
-        bbArray.add(buf);
-        if (currentIsFresh) {
-            bbArray.add(currentBuffer);
-        } else {
-            addBuffer();
+        return null;
+    }
+
+    public void writeBuffer(ByteBuffer buffer) throws IOException {
+        writeBuffers(Collections.singletonList(buffer));
+    }
+
+    public void writeBuffers(ByteBuffer... buffers) throws IOException {
+        writeBuffers(Arrays.asList(buffers));
+    }
+
+    public void writeBuffers(Collection<ByteBuffer> buffers) throws IOException {
+        if (!isOpen) throw new IOException("Stream closed");
+        ByteBuffer remainder = maybeSliceCurrent();
+        bbArray.ensureCapacity(bbArray.size() + buffers.size() + 1 /* remainder or new buffer */);
+        for (ByteBuffer b : buffers) {
+            if (bufferIsFresh(b)) continue;
+            ByteBuffer dup = b.duplicate();
+            if (dup.position() == 0) {
+                dup.position(dup.limit());
+            }
+            bbArray.add(dup);
+            totalBytes += dup.position();
         }
-        currentBufferIndex = bbArray.size() - 1;
+        addRemainderOrNew(remainder);
     }
 
     /**
@@ -168,18 +249,16 @@ public void writeBuffer(ByteBuffer b) throws IOException {
      */
     public void writeBuffers(ByteBufferOutputStream other) throws IOException {
         if (!isOpen || !other.isOpen) throw new IOException("Stream closed");
-        for (ByteBuffer b : other.buffers()) {
-            writeBuffer(b);
-        }
+        writeBuffers(other.buffers());
     }
 
     @Override
     public void writeChar(int v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
-        if (currentBuffer.remaining() < Limits.CHAR_SIZE) getNewBuffer();
-
-        currentBuffer.putChar((char) v);
+        ensureCapacity(2);
+        getCurrent().putChar((char) v);
+        totalBytes += 2;
     }
 
     @Override
@@ -191,54 +270,56 @@ public void writeChars(String s) throws IOException {
     public void writeDouble(double v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
-        if (currentBuffer.remaining() < Limits.DOUBLE_SIZE) getNewBuffer();
-
-        currentBuffer.putDouble(v);
+        ensureCapacity(8);
+        getCurrent().putDouble(v);
+        totalBytes += 8;
     }
 
     @Override
     public void writeFloat(float v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
-        if (currentBuffer.remaining() < Limits.FLOAT_SIZE) getNewBuffer();
-
-        currentBuffer.putFloat(v);
+        ensureCapacity(4);
+        getCurrent().putFloat(v);
+        totalBytes += 4;
     }
 
     @Override
     public void writeInt(int v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
-        if (currentBuffer.remaining() < Limits.INT_SIZE) getNewBuffer();
-
-        currentBuffer.putInt(v);
+        ensureCapacity(4);
+        getCurrent().putInt(v);
+        totalBytes += 4;
     }
 
     @Override
     public void writeLong(long v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
-        if (currentBuffer.remaining() < Limits.LONG_SIZE) getNewBuffer();
-
-        currentBuffer.putLong(v);
+        ensureCapacity(8);
+        getCurrent().putLong(v);
+        totalBytes += 8;
     }
 
     @Override
     public void writeShort(int v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
-        if (currentBuffer.remaining() < Limits.SHORT_SIZE) getNewBuffer();
-
-        currentBuffer.putShort((short) v);
+        ensureCapacity(2);
+        getCurrent().putShort((short) v);
+        totalBytes += 2;
     }
 
     @Override
     public void writeUTF(String str) throws IOException {
-        write(str.getBytes(StandardCharsets.UTF_8));
+        byte[] bytes = str.getBytes(StandardCharsets.UTF_8);
+        write(bytes);
     }
 
     public void writeAscii(String str) throws IOException {
-        write(str.getBytes(StandardCharsets.US_ASCII));
+        byte[] bytes = str.getBytes(StandardCharsets.US_ASCII);
+        write(bytes);
     }
 
     /**
@@ -250,10 +331,7 @@ public ByteBuffer[] reset() {
         ByteBuffer[] bbArrayCopy = peek();
 
         bbArray.clear();
-        prevBuffersNumBytes = 0;
-        currentBuffer = null;
-        currentBufferIndex = 0;
-
+        totalBytes = 0;
         isOpen = false;
 
         return bbArrayCopy;
@@ -264,11 +342,7 @@ public int numByteBuffers() {
     }
 
     public int size() {
-        return (isOpen ? (prevBuffersNumBytes + currentBuffer.position()) : (0));
-    }
-
-    private void getNewBuffer() {
-        addBuffer();
+        return (isOpen ? (totalBytes) : (0));
     }
 
     private void addBuffer() {
@@ -278,10 +352,7 @@ private void addBuffer() {
     // allocate a buffer which is large enough to store data
     // of specified size
     private void addBuffer(int size) {
-        prevBuffersNumBytes += currentBuffer.position();
-        int allocationSize = Math.max(size, bufSize);
-        currentBuffer = ByteBuffer.allocate(allocationSize);
-        bbArray.add(currentBuffer);
-        currentBufferIndex = bbArray.size() - 1;
+        ByteBuffer buf = ByteBuffer.allocate(size);
+        bbArray.add(buf);
     }
 }