Skip to content

Commit

Permalink
Replace JNI method rocksdb_get_helper_direct
Browse files Browse the repository at this point in the history
  • Loading branch information
@alanpaxton
alanpaxton committed Oct 9, 2023
1 parent 618793a commit f20018f
Show file tree
Hide file tree
Showing 4 changed files with 131 additions and 102 deletions.
74 changes: 69 additions & 5 deletions java/rocksjni/jni_get_helpers.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,28 @@ bool GetJNIKey::fromByteArray(JNIEnv* env, jbyteArray jkey, jint jkey_off,
return true;
}

bool GetJNIKey::fromByteBuffer(JNIEnv* env, jobject jkey, jint jkey_off,
jint jkey_len) {
char* key = reinterpret_cast<char*>(env->GetDirectBufferAddress(jkey));
if (key == nullptr) {
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(
env,
"Invalid key argument (argument is not a valid direct ByteBuffer)");
return false;
}
if (env->GetDirectBufferCapacity(jkey) < (jkey_off + jkey_len)) {
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(
env,
"Invalid key argument. Capacity is less than requested region (offset "
"+ length).");
return false;
}

slice_ = Slice(key + jkey_off, jkey_len);

return true;
}

jint GetJNIValue::fillValue(JNIEnv* env, ROCKSDB_NAMESPACE::Status& s,
ROCKSDB_NAMESPACE::PinnableSlice& pinnable_value,
jbyteArray jval, jint jval_off, jint jval_len) {
Expand Down Expand Up @@ -58,17 +80,59 @@ jint GetJNIValue::fillValue(JNIEnv* env, ROCKSDB_NAMESPACE::Status& s,
return pinnable_value_len;
}

jbyteArray GetJNIValue::byteArray(
jint GetJNIValue::fillByteBuffer(
JNIEnv* env, ROCKSDB_NAMESPACE::Status& s,
ROCKSDB_NAMESPACE::PinnableSlice& pinnable_value) {
ROCKSDB_NAMESPACE::PinnableSlice& pinnable_value, jobject jval,
jint jval_off, jint jval_len) {
if (s.IsNotFound()) {
return kNotFound;
} else if (!s.ok()) {
// Here since we are throwing a Java exception from c++ side.
// As a result, c++ does not know calling this function will in fact
// throwing an exception. As a result, the execution flow will
// not stop here, and codes after this throw will still be
// executed.
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(env, s);

// Return a dummy const value to avoid compilation error, although
// java side might not have a chance to get the return value :)
return kStatusError;
}

char* value = reinterpret_cast<char*>(env->GetDirectBufferAddress(jval));
if (value == nullptr) {
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(
env,
"Invalid value argument (argument is not a valid direct ByteBuffer)");
return kArgumentError;
}

if (env->GetDirectBufferCapacity(jval) < (jval_off + jval_len)) {
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(
env,
"Invalid value argument. Capacity is less than requested region "
"(offset + length).");
return kArgumentError;
}

const jint pinnable_value_len = static_cast<jint>(pinnable_value.size());
const jint length = std::min(jval_len, pinnable_value_len);

memcpy(value + jval_off, pinnable_value.data(), length);
pinnable_value.Reset();

return pinnable_value_len;
}

jbyteArray GetJNIValue::byteArray(JNIEnv* env, ROCKSDB_NAMESPACE::Status& s,
ROCKSDB_NAMESPACE::PinnableSlice& value) {
if (s.IsNotFound()) {
return nullptr;
}

if (s.ok()) {
jbyteArray jret_value =
ROCKSDB_NAMESPACE::JniUtil::copyBytes(env, pinnable_value);
pinnable_value.Reset();
jbyteArray jret_value = ROCKSDB_NAMESPACE::JniUtil::copyBytes(env, value);
value.Reset();
if (jret_value == nullptr) {
// exception occurred
return nullptr;
Expand Down
24 changes: 23 additions & 1 deletion java/rocksjni/jni_get_helpers.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ class GetJNIKey {
bool fromByteArray(JNIEnv* env, jbyteArray jkey, jint jkey_off,
jint jkey_len);

bool fromByteBuffer(JNIEnv* env, jobject jkey, jint jkey_off, jint jkey_len);

inline ROCKSDB_NAMESPACE::Slice slice() { return slice_; }
};

Expand Down Expand Up @@ -66,7 +68,7 @@ class MultiGetJNIKeys {
bool fromByteArrays(JNIEnv* env, jobjectArray jkeys);

/**
* @brief Construct helper multiget keys object from array of java keys
* @brief Construct helper multiget keys object from array of java ByteBuffers
*
* @param env JNI environment
* @param jkeys Array of `java.nio.ByteBuffer`, each of which contains a key
Expand Down Expand Up @@ -115,6 +117,7 @@ class GetJNIValue {
public:
static const int kNotFound = -1;
static const int kStatusError = -2;
static const int kArgumentError = -3;

/**
* @brief allocate and fill a byte array from the value in a pinnable slice
Expand Down Expand Up @@ -144,6 +147,25 @@ class GetJNIValue {
static jint fillValue(JNIEnv* env, ROCKSDB_NAMESPACE::Status& s,
ROCKSDB_NAMESPACE::PinnableSlice& value,
jbyteArray jval, jint jval_off, jint jval_len);

/**
* @brief fill an existing direct ByteBuffer from the value in a pinnable
* slice
*
* If the supplied status is faulty, raise an exception instead
*
* @param env JNI environment in which to raise any exception
* @param s status to check before copying the result
* @param value pinnable slice containing a value
* @param jval ByteBuffer target for value
* @param jval_off offset in the array at which to place the value
* @param jval_len length of byte array into which to copy
* @return jint length copied, or a -ve status code
*/

static jint fillByteBuffer(JNIEnv* env, ROCKSDB_NAMESPACE::Status& s,
ROCKSDB_NAMESPACE::PinnableSlice& value,
jobject jval, jint jval_off, jint jval_len);
};

/**
Expand Down
113 changes: 20 additions & 93 deletions java/rocksjni/rocksjni.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1011,93 +1011,6 @@ void Java_org_rocksdb_RocksDB_deleteRange__J_3BII_3BII(
jend_key, jend_key_off, jend_key_len);
}

jint rocksdb_get_helper_direct(
JNIEnv* env, ROCKSDB_NAMESPACE::DB* db,
const ROCKSDB_NAMESPACE::ReadOptions& read_options,
ROCKSDB_NAMESPACE::ColumnFamilyHandle* column_family_handle, jobject jkey,
jint jkey_off, jint jkey_len, jobject jval, jint jval_off, jint jval_len,
bool* has_exception) {
static const int kNotFound = -1;
static const int kStatusError = -2;
static const int kArgumentError = -3;

char* key = reinterpret_cast<char*>(env->GetDirectBufferAddress(jkey));
if (key == nullptr) {
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(
env,
"Invalid key argument (argument is not a valid direct ByteBuffer)");
*has_exception = true;
return kArgumentError;
}
if (env->GetDirectBufferCapacity(jkey) < (jkey_off + jkey_len)) {
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(
env,
"Invalid key argument. Capacity is less than requested region (offset "
"+ length).");
*has_exception = true;
return kArgumentError;
}

char* value = reinterpret_cast<char*>(env->GetDirectBufferAddress(jval));
if (value == nullptr) {
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(
env,
"Invalid value argument (argument is not a valid direct ByteBuffer)");
*has_exception = true;
return kArgumentError;
}

if (env->GetDirectBufferCapacity(jval) < (jval_off + jval_len)) {
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(
env,
"Invalid value argument. Capacity is less than requested region "
"(offset + length).");
*has_exception = true;
return kArgumentError;
}

key += jkey_off;
value += jval_off;

ROCKSDB_NAMESPACE::Slice key_slice(key, jkey_len);

ROCKSDB_NAMESPACE::PinnableSlice pinnable_value;
ROCKSDB_NAMESPACE::Status s;
if (column_family_handle != nullptr) {
s = db->Get(read_options, column_family_handle, key_slice, &pinnable_value);
} else {
// backwards compatibility
s = db->Get(read_options, db->DefaultColumnFamily(), key_slice,
&pinnable_value);
}

if (s.IsNotFound()) {
*has_exception = false;
return kNotFound;
} else if (!s.ok()) {
*has_exception = true;
// Here since we are throwing a Java exception from c++ side.
// As a result, c++ does not know calling this function will in fact
// throwing an exception. As a result, the execution flow will
// not stop here, and codes after this throw will still be
// executed.
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(env, s);

// Return a dummy const value to avoid compilation error, although
// java side might not have a chance to get the return value :)
return kStatusError;
}

const jint pinnable_value_len = static_cast<jint>(pinnable_value.size());
const jint length = std::min(jval_len, pinnable_value_len);

memcpy(value, pinnable_value.data(), length);
pinnable_value.Reset();

*has_exception = false;
return pinnable_value_len;
}

/*
* Class: org_rocksdb_RocksDB
* Method: deleteRange
Expand Down Expand Up @@ -1177,16 +1090,30 @@ jint Java_org_rocksdb_RocksDB_getDirect(JNIEnv* env, jobject /*jdb*/,
jint jkey_len, jobject jval,
jint jval_off, jint jval_len,
jlong jcf_handle) {
auto* db_handle = reinterpret_cast<ROCKSDB_NAMESPACE::DB*>(jdb_handle);
auto* db = reinterpret_cast<ROCKSDB_NAMESPACE::DB*>(jdb_handle);
auto* ro_opt =
reinterpret_cast<ROCKSDB_NAMESPACE::ReadOptions*>(jropt_handle);
auto* cf_handle =
reinterpret_cast<ROCKSDB_NAMESPACE::ColumnFamilyHandle*>(jcf_handle);
bool has_exception = false;
return rocksdb_get_helper_direct(
env, db_handle,
ro_opt == nullptr ? ROCKSDB_NAMESPACE::ReadOptions() : *ro_opt, cf_handle,
jkey, jkey_off, jkey_len, jval, jval_off, jval_len, &has_exception);

ROCKSDB_NAMESPACE::GetJNIKey key;
if (!key.fromByteBuffer(env, jkey, jkey_off, jkey_len)) {
return ROCKSDB_NAMESPACE::GetJNIValue::kStatusError;
}

ROCKSDB_NAMESPACE::PinnableSlice value;
ROCKSDB_NAMESPACE::Status s;
if (cf_handle != nullptr) {
s = db->Get(ro_opt == nullptr ? ROCKSDB_NAMESPACE::ReadOptions() : *ro_opt,
cf_handle, key.slice(), &value);
} else {
// backwards compatibility
s = db->Get(ro_opt == nullptr ? ROCKSDB_NAMESPACE::ReadOptions() : *ro_opt,
db->DefaultColumnFamily(), key.slice(), &value);
}

return ROCKSDB_NAMESPACE::GetJNIValue::fillByteBuffer(env, s, value, jval,
jval_off, jval_len);
}

//////////////////////////////////////////////////////////////////////////////
Expand Down
22 changes: 19 additions & 3 deletions java/src/test/java/org/rocksdb/RocksDBTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -221,15 +221,31 @@ public void put() throws RocksDBException {

key.position(4);

final ByteBuffer result2 = ByteBuffer.allocateDirect(12);
result2.put("abcdefghijkl".getBytes());
result2.flip().position(3);
assertThat(db.get(optr, key, result2)).isEqualTo(4);
assertThat(result2.position()).isEqualTo(3);
assertThat(result2.limit()).isEqualTo(7);
assertThat(key.position()).isEqualTo(8);
assertThat(key.limit()).isEqualTo(8);

final byte[] tmp2 = new byte[12];
result2.position(0).limit(12);
result2.get(tmp2);
assertThat(tmp2).isEqualTo("abcval3hijkl".getBytes());

key.position(4);

result.clear().position(9);
assertThat(db.get(optr, key, result)).isEqualTo(4);
assertThat(result.position()).isEqualTo(9);
assertThat(result.limit()).isEqualTo(12);
assertThat(key.position()).isEqualTo(8);
assertThat(key.limit()).isEqualTo(8);
final byte[] tmp2 = new byte[3];
result.get(tmp2);
assertThat(tmp2).isEqualTo("val".getBytes());
final byte[] tmp3 = new byte[3];
result.get(tmp3);
assertThat(tmp3).isEqualTo("val".getBytes());

// put
final Segment key3 = sliceSegment("key3");
Expand Down

0 comments on commit f20018f

Please sign in to comment.