Skip to content

Commit

Permalink
Perf test first refactored jni multiGet function
Browse files Browse the repository at this point in the history 
Bring java/jmh MultiGet tests back to life, and update
Compare old, intermediate and new implementations of simplest multiGet for performance
Do we need linux for it to make any difference ? (io_uring)
  • Loading branch information
@alanpaxton
alanpaxton committed Sep 26, 2023
1 parent baed07f commit 5632af9
Show file tree
Hide file tree
Showing 5 changed files with 125 additions and 14 deletions.
6 changes: 3 additions & 3 deletions java/jmh/pom.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,8 @@
</scm>

<properties>
<project.build.source>1.7</project.build.source>
<project.build.target>1.7</project.build.target>
<project.build.source>17</project.build.source>
<project.build.target>17</project.build.target>
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>

<jmh.version>1.22</jmh.version>
Expand All @@ -50,7 +50,7 @@
<dependency>
<groupId>org.rocksdb</groupId>
<artifactId>rocksdbjni</artifactId>
<version>7.9.0-SNAPSHOT</version>
<version>8.7.0</version>
</dependency>

<dependency>
Expand Down
60 changes: 55 additions & 5 deletions java/jmh/src/main/java/org/rocksdb/jmh/MultiGetBenchmarks.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,12 @@ public void setup() throws IOException, RocksDBException {
cfHandles = cfHandlesList.toArray(new ColumnFamilyHandle[0]);

// store initial data for retrieving via get
for (int j = 0; j < keyCount; j++) {
final byte[] paddedValue = Arrays.copyOf(ba("value" + j), valueSize);
db.put(ba("key" + j), paddedValue);
}

// store initial data for retrieving via get - column families
for (int i = 0; i < cfs; i++) {
for (int j = 0; j < keyCount; j++) {
final byte[] paddedValue = Arrays.copyOf(ba("value" + j), valueSize);
Expand Down Expand Up @@ -168,10 +174,8 @@ public void allocateSliceBuffers() {
valueBuffersList = new ArrayList<>();
keyBuffersList = new ArrayList<>();
for (int i = 0; i < keyCount; i++) {
valueBuffersList.add(valuesBuffer.slice());
valuesBuffer.position(i * valueSize);
keyBuffersList.add(keysBuffer.slice());
keysBuffer.position(i * keySize);
valueBuffersList.add(valuesBuffer.slice(i * valueSize, valueSize));
keyBuffersList.add(keysBuffer.slice(i * keySize, keySize));
}
}

Expand All @@ -181,7 +185,7 @@ public void freeSliceBuffers() {
}

@Benchmark
public List<byte[]> multiGet10() throws RocksDBException {
public List<byte[]> multiGetList10() throws RocksDBException {
final int fromKeyIdx = next(multiGetSize, keyCount);
if (fromKeyIdx >= 0) {
final List<byte[]> keys = keys(fromKeyIdx, fromKeyIdx + multiGetSize);
Expand All @@ -194,6 +198,52 @@ public List<byte[]> multiGet10() throws RocksDBException {
return new ArrayList<>();
}

@Benchmark
public List<byte[]> multiGetList10_intermediate() throws RocksDBException {
final int fromKeyIdx = next(multiGetSize, keyCount);
if (fromKeyIdx >= 0) {
final List<byte[]> keys = keys(fromKeyIdx, fromKeyIdx + multiGetSize);
final List<byte[]> valueResults = db.multiGetAsList_intermediate(keys);
for (final byte[] result : valueResults) {
if (result.length != valueSize)
throw new RuntimeException("Test valueSize assumption wrong");
}
}
return new ArrayList<>();
}

@Benchmark
public List<byte[]> multiGetList10_old() throws RocksDBException {
final int fromKeyIdx = next(multiGetSize, keyCount);
if (fromKeyIdx >= 0) {
final List<byte[]> keys = keys(fromKeyIdx, fromKeyIdx + multiGetSize);
final List<byte[]> valueResults = db.multiGetAsList_old(keys);
for (final byte[] result : valueResults) {
if (result.length != valueSize)
throw new RuntimeException("Test valueSize assumption wrong");
}
}
return new ArrayList<>();
}

@Benchmark
public List<byte[]> multiGet20() throws RocksDBException {
final int fromKeyIdx = next(multiGetSize, keyCount);
if (fromKeyIdx >= 0) {
final List<ByteBuffer> keys = keys(keyBuffersList, fromKeyIdx, fromKeyIdx + multiGetSize);
final List<ByteBuffer> values = valueBuffersList.subList(fromKeyIdx, fromKeyIdx + multiGetSize);
final List<ByteBufferGetStatus> statusResults =
db.multiGetByteBuffers(keys, values);
for (final ByteBufferGetStatus result : statusResults) {
if (result.status.getCode() != Status.Code.Ok)
throw new RuntimeException("Test status not OK: " + result.status);
if (result.value.limit() != valueSize)
throw new RuntimeException("Test valueSize assumption wrong");
}
}
return new ArrayList<>();
}

public static void main(final String[] args) throws RunnerException {
final org.openjdk.jmh.runner.options.Options opt =
new OptionsBuilder()
Expand Down
6 changes: 5 additions & 1 deletion java/rocksjni/jni_get_helpers.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,7 @@ jobjectArray MultiGetValues::byteArrays(
for (std::vector<ROCKSDB_NAMESPACE::Status>::size_type i = 0; i != s.size();
i++) {
if (s[i].ok()) {
std::string* value = &values[i];
TValue* value = &values[i];
const jsize jvalue_len = static_cast<jsize>(value->size());
jbyteArray jentry_value = env->NewByteArray(jvalue_len);
if (jentry_value == nullptr) {
Expand Down Expand Up @@ -207,4 +207,8 @@ template jobjectArray MultiGetValues::byteArrays<std::string>(
JNIEnv* env, std::vector<std::string>& values,
std::vector<ROCKSDB_NAMESPACE::Status>& s);

template jobjectArray MultiGetValues::byteArrays<ROCKSDB_NAMESPACE::PinnableSlice>(
JNIEnv* env, std::vector<ROCKSDB_NAMESPACE::PinnableSlice>& values,
std::vector<ROCKSDB_NAMESPACE::Status>& s);

}; // namespace ROCKSDB_NAMESPACE
29 changes: 26 additions & 3 deletions java/rocksjni/rocksjni.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1986,7 +1986,7 @@ void multi_get_helper_direct(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
* Method: multiGet
* Signature: (J[[B[I[I)[[B
*/
jobjectArray Java_org_rocksdb_RocksDB_multiGet000__J_3_3B_3I_3I(
JNIEXPORT jobjectArray JNICALL Java_org_rocksdb_RocksDB_multiGetOld(
JNIEnv* env, jobject jdb, jlong jdb_handle, jobjectArray jkeys,
jintArray jkey_offs, jintArray jkey_lens) {
return multi_get_helper(
Expand All @@ -1995,13 +1995,13 @@ jobjectArray Java_org_rocksdb_RocksDB_multiGet000__J_3_3B_3I_3I(
}

/*
* TODO (AP) - not finished
* TODO (AP) - intermediate version
*
* Class: org_rocksdb_RocksDB
* Method: multiGet
* Signature: (J[[B[I[I)[[B
*/
jobjectArray Java_org_rocksdb_RocksDB_multiGet__J_3_3B_3I_3I(
jobjectArray Java_org_rocksdb_RocksDB_multiGetIntermediate(
JNIEnv* env, jobject, jlong jdb_handle, jobjectArray jkeys,
jintArray jkey_offs, jintArray jkey_lens) {
ROCKSDB_NAMESPACE::MultiGetKeys keys;
Expand All @@ -2017,6 +2017,29 @@ jobjectArray Java_org_rocksdb_RocksDB_multiGet__J_3_3B_3I_3I(
return ROCKSDB_NAMESPACE::MultiGetValues::byteArrays(env, values, statuses);
}

/*
* Use the efficient/optimized variant of MultiGet()
*
* Class: org_rocksdb_RocksDB
* Method: multiGet
* Signature: (J[[B[I[I)[[B
*/
jobjectArray Java_org_rocksdb_RocksDB_multiGet__J_3_3B_3I_3I(
JNIEnv* env, jobject, jlong jdb_handle, jobjectArray jkeys,
jintArray jkey_offs, jintArray jkey_lens) {
ROCKSDB_NAMESPACE::MultiGetKeys keys;
if (!keys.fromByteArrays(env, jkeys, jkey_offs, jkey_lens)) {
return nullptr;
}
std::vector<ROCKSDB_NAMESPACE::PinnableSlice> values(keys.size());
std::vector<ROCKSDB_NAMESPACE::Status> statuses(keys.size());
auto* db = reinterpret_cast<ROCKSDB_NAMESPACE::DB*>(jdb_handle);
db->MultiGet(ROCKSDB_NAMESPACE::ReadOptions(), db->DefaultColumnFamily(),
keys.size(), keys.array(), values.data(), statuses.data(),
false /* sorted_input*/);
return ROCKSDB_NAMESPACE::MultiGetValues::byteArrays(env, values, statuses);
}

/*
* Class: org_rocksdb_RocksDB
* Method: multiGet
Expand Down
38 changes: 36 additions & 2 deletions java/src/main/java/org/rocksdb/RocksDB.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2147,6 +2147,36 @@ public List<byte[]> multiGetAsList(final List<byte[]> keys)
keyLengths));
}

public List<byte[]> multiGetAsList_intermediate(final List<byte[]> keys)
throws RocksDBException {
assert(keys.size() != 0);

final byte[][] keysArray = keys.toArray(new byte[keys.size()][]);
final int[] keyOffsets = new int[keysArray.length];
final int[] keyLengths = new int[keysArray.length];
for(int i = 0; i < keyLengths.length; i++) {
keyLengths[i] = keysArray[i].length;
}

return Arrays.asList(multiGetIntermediate(nativeHandle_, keysArray, keyOffsets,
keyLengths));
}

public List<byte[]> multiGetAsList_old(final List<byte[]> keys)
throws RocksDBException {
assert(keys.size() != 0);

final byte[][] keysArray = keys.toArray(new byte[keys.size()][]);
final int[] keyOffsets = new int[keysArray.length];
final int[] keyLengths = new int[keysArray.length];
for(int i = 0; i < keyLengths.length; i++) {
keyLengths[i] = keysArray[i].length;
}

return Arrays.asList(multiGetOld(nativeHandle_, keysArray, keyOffsets,
keyLengths));
}

/**
* Returns a list of values for the given list of keys. List will contain
* null for keys which could not be found.
Expand Down Expand Up @@ -2359,7 +2389,7 @@ public List<ByteBufferGetStatus> multiGetByteBuffers(final ReadOptions readOptio
// Check if key size equals cfList size. If not a exception must be
// thrown. If not a Segmentation fault happens.
if (values.size() != keys.size()) {
throw new IllegalArgumentException("For each key there must be a corresponding value.");
throw new IllegalArgumentException("Mismatch " + keys.size() + " keys, but not the same as " + values.size() + " values.");
}

// TODO (AP) support indirect buffers
Expand Down Expand Up @@ -4498,7 +4528,11 @@ private native byte[] get(final long handle,
final long readOptHandle, final byte[] key, final int keyOffset,
final int keyLength, final long cfHandle) throws RocksDBException;
private native byte[][] multiGet(final long dbHandle, final byte[][] keys,
final int[] keyOffsets, final int[] keyLengths);
final int[] keyOffsets, final int[] keyLengths);
private native byte[][] multiGetIntermediate(final long dbHandle, final byte[][] keys,
final int[] keyOffsets, final int[] keyLengths);
private native byte[][] multiGetOld(final long dbHandle, final byte[][] keys,
final int[] keyOffsets, final int[] keyLengths);
private native byte[][] multiGet(final long dbHandle, final byte[][] keys,
final int[] keyOffsets, final int[] keyLengths,
final long[] columnFamilyHandles);
Expand Down

0 comments on commit 5632af9

Please sign in to comment.