Skip to content

Small speedups in RecyclerBytesStreamOutput #140257

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
DaveCTurner wants to merge 1 commit into from
Closed

Small speedups in RecyclerBytesStreamOutput #140257

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 42 additions & 14 deletions ...s/src/main/java/org/elasticsearch/benchmark/bytes/RecyclerBytesStreamOutputBenchmark.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import org.openjdk.jmh.annotations.Warmup;

import java.io.IOException;
import java.util.Random;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
Expand All @@ -48,6 +49,7 @@ public class RecyclerBytesStreamOutputBenchmark {
private byte[] bytes3;
private byte[] multiPageBytes;
private int[] vints;
private long[] vlongs;

@Setup
public void initResults() throws IOException {
Expand All @@ -70,25 +72,43 @@ public void initResults() throws IOException {
longString = UTF8StringBytesBenchmark.generateAsciiString(100);
nonAsciiString = UTF8StringBytesBenchmark.generateUTF8String(200);
veryLongString = UTF8StringBytesBenchmark.generateAsciiString(800);
// vint values for benchmarking

vints = new int[1000];
for (int i = 0; i < vints.length; i++) {
vints[i] = randomVInt(random);
}
vlongs = new long[1000];
for (int i = 0; i < vlongs.length; i++) {
vlongs[i] = randomVLong(random);
}
}

private static int randomVInt(Random random) {
if (random.nextBoolean()) {
// 1-byte 50% of the time
return random.nextInt(1 << 7);
}
for (int maxBits = 14; maxBits < 32; maxBits += 7) {
// 2-byte 25% of the time, 3-byte 12.5% of the time etc.
if (random.nextBoolean()) {
// 1-byte 50% of the time
vints[i] = random.nextInt(128);
} else if (random.nextBoolean()) {
// 2-byte 25% of the time
vints[i] = random.nextInt(128, 16384);
} else {
if (random.nextBoolean()) {
// 3-byte vints
vints[i] = random.nextInt(16384, 2097152);
} else {
// All vint variants
vints[i] = random.nextInt();
}
return random.nextInt(1 << (maxBits - 7), 1 << maxBits);
}
}
return random.nextInt();
}

private static long randomVLong(Random random) {
if (random.nextBoolean()) {
// 1-byte 50% of the time
return random.nextLong(1L << 7);
}
for (int maxBits = 14; maxBits < 63; maxBits += 7) {
// 2-byte 25% of the time, 3-byte 12.5% of the time etc.
if (random.nextBoolean()) {
return random.nextLong(1L << (maxBits - 7), 1L << maxBits);
}
}
return random.nextLong() & 0x7FFF_FFFF_FFFF_FFFFL;
}

@Benchmark
Expand Down Expand Up @@ -144,6 +164,14 @@ public void writeVInt() throws IOException {
}
}

@Benchmark
public void writeVLong() throws IOException {
streamOutput.seek(1);
for (long vlong : vlongs) {
streamOutput.writeVLong(vlong);
}
}

private record BenchmarkRecycler(AtomicReference<BytesRef> bytesRef) implements Recycler<BytesRef> {

@Override
Expand Down
23 changes: 18 additions & 5 deletions server/src/main/java/org/elasticsearch/common/io/stream/RecyclerBytesStreamOutput.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,16 +132,15 @@ public void writeVInt(int i) throws IOException {
return;
}

int bytesNeeded = vIntLength(i);
if (bytesNeeded > remainingBytesInPage) {
if (MAX_VINT_BYTES > remainingBytesInPage && vIntLength(i) > remainingBytesInPage) {
super.writeVInt(i);
} else {
BytesRef currentPage = currentBytesRef;
putVInt(i, bytesNeeded, currentPage.bytes, currentPage.offset + currentPageOffset);
this.currentPageOffset = currentPageOffset + bytesNeeded;
this.currentPageOffset = StreamOutputHelper.putMultiByteVInt(currentBytesRef.bytes, i, currentPageOffset);
}
}

private static final int MAX_VINT_BYTES = 5;

public static int vIntLength(int value) {
int leadingZeros = Integer.numberOfLeadingZeros(value);
if (leadingZeros >= 25) {
Expand All @@ -164,6 +163,20 @@ private void putVInt(int i, int bytesNeeded, byte[] page, int offset) {
}
}

private static final int MAX_VLONG_BYTES = 9;

@Override
void writeVLongNoCheck(long i) throws IOException {
final int currentPageOffset = this.currentPageOffset;
final int remainingBytesInPage = pageSize - currentPageOffset;

if (MAX_VLONG_BYTES < remainingBytesInPage) {
this.currentPageOffset = StreamOutputHelper.putVLong(currentBytesRef.bytes, i, currentPageOffset);
} else {
super.writeVLongNoCheck(i);
}
}

@Override
public void writeInt(int i) throws IOException {
final int currentPageOffset = this.currentPageOffset;
Expand Down
45 changes: 37 additions & 8 deletions server/src/main/java/org/elasticsearch/common/io/stream/StreamOutputHelper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ public static void writeString(String str, OutputStream outputStream) throws IOE
*/
public static void writeString(String str, byte[] buffer, int prefixLength, OutputStream outputStream) throws IOException {
final int charCount = str.length();
int offset = prefixLength + putVInt(buffer, charCount, prefixLength);
int offset = putVInt(buffer, charCount, prefixLength);
for (int i = 0; i < charCount; i++) {
final int c = str.charAt(i);
if (c <= 0x007F) {
Expand Down Expand Up @@ -130,12 +130,12 @@ public static void writeGenericString(String value, OutputStream outputStream) t
* Put the integer {@code i} into the given {@code buffer} starting at the given {@code offset}, formatted as per
* {@link StreamOutput#writeVInt}. Performs no bounds checks: callers must verify that there is enough space in {@code buffer} first.
*
* @return number of bytes written.
* @return updated offset.
*/
public static int putVInt(byte[] buffer, int i, int offset) {
if (Integer.numberOfLeadingZeros(i) >= 25) {
buffer[offset] = (byte) i;
return 1;
return offset + 1;
}
return putMultiByteVInt(buffer, i, offset);
}
Expand All @@ -144,17 +144,46 @@ public static int putVInt(byte[] buffer, int i, int offset) {
* Put the integer {@code i} into the given {@code buffer} starting at the given {@code offset}, formatted as per
* {@link StreamOutput#writeVInt}. Performs no bounds checks: callers must verify that there is enough space in {@code buffer} first.
*
* @return number of bytes written.
* @return updated offset.
*/
// extracted from putVInt() to allow the hot single-byte path to be inlined
public static int putMultiByteVInt(byte[] buffer, int i, int offset) {
int index = offset;
do {
buffer[index++] = ((byte) ((i & 0x7f) | 0x80));
buffer[offset++] = ((byte) ((i & 0x7f) | 0x80));
i >>>= 7;
} while ((i & ~0x7F) != 0);
buffer[index++] = (byte) i;
return index - offset;
buffer[offset++] = (byte) i;
return offset;
}

/**
* Put the long {@code l} into the given {@code buffer} starting at the given {@code offset}, formatted as per
* {@link StreamOutput#writeVLong}. Performs no bounds checks: callers must verify that there is enough space in {@code buffer} first
* and that {@code l} is non-negative.
*
* @return updated offset.
*/
public static int putVLong(byte[] buffer, long l, int offset) {
if (Long.numberOfLeadingZeros(l) >= 57) {
buffer[offset] = (byte) l;
return offset + 1;
}
return putMultiByteVLong(buffer, l, offset);
}

/**
* Put the long {@code l} into the given {@code buffer} starting at the given {@code offset}, formatted as per
* {@link StreamOutput#writeVLong}. Performs no bounds checks: callers must verify that there is enough space in {@code buffer} first
* and that {@code l} is non-negative.
*
* @return updated offset.
*/
public static int putMultiByteVLong(byte[] buffer, long l, int offset) {
do {
buffer[offset++] = ((byte) ((l & 0x7f) | 0x80));
l >>>= 7;
} while ((l & 0xFFFF_FFFF_FFFF_FF80L) != 0);
buffer[offset++] = (byte) l;
return offset;
}
}