Eliminate array copies

bmq-sdk/src/main/java/com/bloomberg/bmq/impl/BrokerSession.java

@@ -59,7 +59,9 @@
 import java.lang.invoke.MethodHandles;
 import java.nio.ByteBuffer;
 import java.time.Duration;
+import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -1219,10 +1221,12 @@ public QueueHandle lookupQueue(QueueId queueId) {
         return queueStateManager.findByQueueId(queueId);
     }
 
-    public void post(QueueHandle queueHandle, PutMessageImpl... msgs) throws BMQException {
+    public void post(QueueHandle queueHandle, Collection<PutMessageImpl> msgs) throws BMQException {
         Argument.expectNonNull(queueHandle, "queueHandle");
         Argument.expectNonNull(msgs, "msgs");
-        Argument.expectPositive(msgs.length, "message array length");
+        if (msgs.isEmpty()) {
+            return;
+        }
 
         // Queue state guard
         QueueState state = queueHandle.getState();
@@ -1249,6 +1253,14 @@ public void post(QueueHandle queueHandle, PutMessageImpl... msgs) throws BMQExce
         }
     }
 
+    public void post(QueueHandle queueHandle, PutMessageImpl msg) throws BMQException {
+        post(queueHandle, Collections.singletonList(msg));
+    }
+
+    public void post(QueueHandle queueHandle, PutMessageImpl... msgs) throws BMQException {
+        post(queueHandle, Arrays.asList(msgs));
+    }
+
     public GenericResult confirm(QueueHandle queueHandle, PushMessageImpl... messages) {
         Argument.expectNonNull(queueHandle, "queueHandle");
         Argument.expectNonNull(messages, "messages");

bmq-sdk/src/main/java/com/bloomberg/bmq/impl/PutPoster.java

@@ -33,7 +33,9 @@
 import java.lang.invoke.MethodHandles;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
@@ -68,9 +70,9 @@ void setMaxEventSize(int val) {
         maxEventSize = Argument.expectNotGreater(val, EventHeader.MAX_SIZE_SOFT, "max event size");
     }
 
-    public void pack(PutMessageImpl... msgs) {
+    public void pack(Collection<PutMessageImpl> msgs) {
         Argument.expectNonNull(msgs, "msgs");
-        Argument.expectPositive(msgs.length, "message array length");
+        Argument.expectPositive(msgs.size(), "message array length");
         for (PutMessageImpl m : msgs) {
             Argument.expectNonNull(m, "put message");
 
@@ -87,11 +89,19 @@ public void flush() {
         }
     }
 
-    public void post(PutMessageImpl... msgs) {
+    public void post(Collection<PutMessageImpl> msgs) {
         pack(msgs);
         flush();
     }
 
+    public void post(PutMessageImpl msg) {
+        post(Collections.singletonList(msg));
+    }
+
+    public void post(PutMessageImpl... msgs) {
+        post(Arrays.asList(msgs));
+    }
+
     private void sendEvent() {
         PutEventBuilder putBuilder = new PutEventBuilder();
         putBuilder.setMaxEventSize(maxEventSize);

bmq-sdk/src/main/java/com/bloomberg/bmq/impl/QueueImpl.java

@@ -37,6 +37,8 @@
 import java.lang.invoke.MethodHandles;
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.Nonnull;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -48,6 +50,7 @@ public class QueueImpl implements QueueHandle {
     static final Logger logger = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
 
     static final int INVALID_QUEUE_ID = -1;
+    private static final int INITIAL_PUTMESSAGES_SIZE = 100;
 
     // Immutable fields
     private final BrokerSession brokerSession;
@@ -61,7 +64,8 @@ public class QueueImpl implements QueueHandle {
     // Fields exposed to user thread
     private final QueueHandleParameters parameters; // mutable object and final field
     private volatile QueueState state;
-    private final ArrayList<PutMessageImpl> putMessages = new ArrayList<>();
+    private final AtomicReference<Collection<PutMessageImpl>> putMessages =
+            new AtomicReference<>(new ArrayList<>(INITIAL_PUTMESSAGES_SIZE));
     private volatile boolean isSuspended = false;
     // Whether the queue is suspended.
     // While suspended, a queue receives no
@@ -261,20 +265,11 @@ public BmqFuture<CloseQueueCode> closeAsync(Duration timeout) {
     }
 
     public void pack(PutMessageImpl message) throws BMQException {
-        synchronized (lock) {
-            putMessages.add(message);
-        }
+        putMessages.get().add(message);
     }
 
-    public PutMessageImpl[] flush() throws BMQException {
-        PutMessageImpl[] msgs;
-        synchronized (lock) {
-            msgs = new PutMessageImpl[putMessages.size()];
-            msgs = putMessages.toArray(msgs);
-            putMessages.clear();
-        }
-        brokerSession.post(this, msgs);
-        return msgs;
+    public void flush() throws BMQException {
+        brokerSession.post(this, putMessages.getAndSet(new ArrayList<>(INITIAL_PUTMESSAGES_SIZE)));
     }
 
     @Override

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

@@ -92,10 +92,27 @@ public void write(byte[] ba, int offset, int length) throws IOException {
         }
     }
 
+    /**
+     * Make a readable view of the underlying data without copying it.
+     *
+     * <p>The bbos can continue to be written to.
+     */
+    public ByteBuffer[] peek() {
+        return bbArray.stream()
+                .map(bb -> (ByteBuffer) (bb.duplicate().flip()))
+                .toArray(ByteBuffer[]::new);
+    }
+
+    private ArrayList<ByteBuffer> buffers() {
+        return bbArray;
+    }
+
+    @Override
     public void writeBoolean(boolean v) throws IOException {
         throw new UnsupportedOperationException();
     }
 
+    @Override
     public void writeByte(int v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
@@ -104,26 +121,42 @@ public void writeByte(int v) throws IOException {
         currentBuffer.put((byte) v);
     }
 
+    @Override
     public void writeBytes(String s) throws IOException {
         throw new UnsupportedOperationException();
     }
 
-    public void writeBytes(ByteBuffer b) throws IOException {
+    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");
-        b.rewind();
-        while (b.hasRemaining()) {
-            if (currentBuffer.hasRemaining()) {
-                if (b.remaining() > currentBuffer.remaining()) {
-                    // Read one byte
-                    currentBuffer.put(b.get());
-                } else {
-                    // Read the whole buffer
-                    currentBuffer.put(b);
-                }
-                continue;
-            }
-            getNewBuffer();
+        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);
+        }
+        prevBuffersNumBytes += buf.position();
+        bbArray.add(buf);
+        if (currentIsFresh) {
+            bbArray.add(currentBuffer);
+        } else {
+            addBuffer();
         }
+        currentBufferIndex = bbArray.size() - 1;
     }
 
     /**
@@ -133,17 +166,14 @@ public void writeBytes(ByteBuffer b) throws IOException {
      * @param bbos stream to write bytes to
      * @throws IOException if the stream is not open
      */
-    public void writeBytes(ByteBufferOutputStream bbos) throws IOException {
-
-        if (!isOpen || !bbos.isOpen) throw new IOException("Stream closed");
-
-        for (int i = 0; i < bbos.bbArray.size(); i++) {
-            ByteBuffer data = bbos.bbArray.get(i);
-            data.flip();
-            writeBytes(data);
+    public void writeBuffers(ByteBufferOutputStream other) throws IOException {
+        if (!isOpen || !other.isOpen) throw new IOException("Stream closed");
+        for (ByteBuffer b : other.buffers()) {
+            writeBuffer(b);
         }
     }
 
+    @Override
     public void writeChar(int v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
@@ -152,10 +182,12 @@ public void writeChar(int v) throws IOException {
         currentBuffer.putChar((char) v);
     }
 
+    @Override
     public void writeChars(String s) throws IOException {
         throw new UnsupportedOperationException();
     }
 
+    @Override
     public void writeDouble(double v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
@@ -164,6 +196,7 @@ public void writeDouble(double v) throws IOException {
         currentBuffer.putDouble(v);
     }
 
+    @Override
     public void writeFloat(float v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
@@ -172,6 +205,7 @@ public void writeFloat(float v) throws IOException {
         currentBuffer.putFloat(v);
     }
 
+    @Override
     public void writeInt(int v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
@@ -180,6 +214,7 @@ public void writeInt(int v) throws IOException {
         currentBuffer.putInt(v);
     }
 
+    @Override
     public void writeLong(long v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
@@ -188,6 +223,7 @@ public void writeLong(long v) throws IOException {
         currentBuffer.putLong(v);
     }
 
+    @Override
     public void writeShort(int v) throws IOException {
         if (!isOpen) throw new IOException("Stream closed");
 
@@ -196,6 +232,7 @@ public void writeShort(int v) throws IOException {
         currentBuffer.putShort((short) v);
     }
 
+    @Override
     public void writeUTF(String str) throws IOException {
         write(str.getBytes(StandardCharsets.UTF_8));
     }
@@ -210,12 +247,7 @@ public void writeAscii(String str) throws IOException {
      * @return ByteBuffer[] previously allocated buffers flipped to the read mode.
      */
     public ByteBuffer[] reset() {
-        ByteBuffer[] bbArrayCopy = new ByteBuffer[currentBufferIndex + 1];
-
-        for (int i = 0; i <= currentBufferIndex; ++i) {
-            bbArrayCopy[i] = bbArray.get(i);
-            bbArrayCopy[i].flip();
-        }
+        ByteBuffer[] bbArrayCopy = peek();
 
         bbArray.clear();
         prevBuffersNumBytes = 0;
@@ -227,22 +259,6 @@ public ByteBuffer[] reset() {
         return bbArrayCopy;
     }
 
-    /**
-     * Clears the Stream for reuse. Clears the previously stored data but leaves the stream in its
-     * previous state. If the stream is still open, then it can be reused to insert new data. Reuses
-     * previously allocated Buffers if possible.
-     */
-    public void clear() {
-        if (!isOpen) return;
-
-        for (int i = 0; i <= currentBufferIndex; i++) {
-            bbArray.get(i).clear();
-        }
-        prevBuffersNumBytes = 0;
-        currentBuffer = bbArray.get(0);
-        currentBufferIndex = 0;
-    }
-
     public int numByteBuffers() {
         return bbArray.size();
     }
@@ -252,13 +268,7 @@ public int size() {
     }
 
     private void getNewBuffer() {
-        // Try to reuse a previously allocated buffer before allocating new.
-        if (currentBufferIndex < bbArray.size() - 1) {
-            prevBuffersNumBytes += currentBuffer.position();
-            currentBuffer = bbArray.get(++currentBufferIndex);
-        } else {
-            addBuffer();
-        }
+        addBuffer();
     }
 
     private void addBuffer() {
@@ -272,6 +282,6 @@ private void addBuffer(int size) {
         int allocationSize = Math.max(size, bufSize);
         currentBuffer = ByteBuffer.allocate(allocationSize);
         bbArray.add(currentBuffer);
-        currentBufferIndex++;
+        currentBufferIndex = bbArray.size() - 1;
     }
 }

bmq-sdk/src/main/java/com/bloomberg/bmq/impl/infr/net/NettyTcpConnection.java

@@ -745,19 +745,13 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) {
 
             ByteBuf byteBuf = (ByteBuf) msg;
             try {
-                // Make a copy of bytes present in 'byteBuf'.  This is needed
-                // so that the lifetime of 'readBuffer' (and ByteBuffer[]
-                // obtained from it via 'readBuffer.reset()') can be
-                // decoupled with the associated 'byteBuf'.  As per netty docs,
-                // a ByteBuf object is refcounted, and user must invoke
-                // 'ByteBuf.release' to decrement the counter when done.  Since
-                // there is no place in the logic where we can confidentally
-                // invoke 'ByteBuf.release' (BlazingMQ user may keep a PUSH message
-                // for ever), we make a copy of bytes present in 'ByteBuf' and
-                // invoke 'ByteBuf.release' right away.  This copying is
-                // unfortunate, and we will investigate ways to avoid this
-                // copy later.
-                readBuffer.writeBytes(byteBuf.nioBuffer());
+                // This still makes a copy of the underlying bytes, but let netty do it
+                // instead of doing it ourselves.
+                //
+                // TODO - can this copy be eliminated?
+                ByteBuffer nioBuf = ByteBuffer.allocateDirect(byteBuf.readableBytes());
+                byteBuf.readBytes(nioBuf);
+                readBuffer.writeBuffer(nioBuf);
             } catch (IOException e) {
                 logger.error("Failed to write data: ", e);
             }

bmq-sdk/src/main/java/com/bloomberg/bmq/impl/infr/proto/AckMessageImpl.java

@@ -154,9 +154,7 @@ public void streamIn(ByteBufferInputStream bbis) throws IOException {
 
     public void streamOut(ByteBufferOutputStream bbos) throws IOException {
         bbos.writeInt(statusAndCorrelationId);
-        for (int i = 0; i < MessageGUID.SIZE_BINARY; i++) {
-            bbos.writeByte(messageGUID[i]);
-        }
+        bbos.write(messageGUID);
         bbos.writeInt(queueId);
     }