Skip to content

Commit

Permalink
more efficient buffer usage, add tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@wolfchimneyrock
wolfchimneyrock committed Sep 12, 2023
1 parent ae2cd67 commit f20e57d
Show file tree
Hide file tree
Showing 7 changed files with 377 additions and 103 deletions.
251 changes: 161 additions & 90 deletions bmq-sdk/src/main/java/com/bloomberg/bmq/impl/infr/io/ByteBufferOutputStream.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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
Expand All @@ -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;
}

/**
Expand All @@ -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() {
Expand All @@ -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);
}

/**
Expand All @@ -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
Expand All @@ -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);
}

/**
Expand All @@ -250,10 +331,7 @@ public ByteBuffer[] reset() {
ByteBuffer[] bbArrayCopy = peek();

bbArray.clear();
prevBuffersNumBytes = 0;
currentBuffer = null;
currentBufferIndex = 0;

totalBytes = 0;
isOpen = false;

return bbArrayCopy;
Expand All @@ -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() {
Expand All @@ -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);
}
}
Loading

0 comments on commit f20e57d

Please sign in to comment.