Java client support for plasma

--------------------------
This commit includes the support of java client for plasma, which is part of the java worker support of Ray.
In addition to some minor changes in build system, it consists of the following modules:
 - java/plasma: java client support for plasma
 - cpp/src/plasma/lib/java: JNI support for plasma client

Author: salah-man

cpp/src/plasma/CMakeLists.txt

@@ -160,6 +160,30 @@ install(
   FILES "${CMAKE_CURRENT_BINARY_DIR}/plasma.pc"
   DESTINATION "${CMAKE_INSTALL_LIBDIR}/pkgconfig/")
 
+# Plasma java client support
+add_compile_options("-I$ENV{JAVA_HOME}/include/")
+if(WIN32)
+  add_compile_options("-I$ENV{JAVA_HOME}/include/win32")
+elseif(APPLE)
+  add_compile_options("-I$ENV{JAVA_HOME}/include/darwin")
+else() # linux
+  add_compile_options("-I$ENV{JAVA_HOME}/include/linux")
+endif()
+
+include_directories("${CMAKE_CURRENT_LIST_DIR}/lib/java")
+
+file(GLOB PLASMA_LIBRARY_EXT_java_SRC
+  lib/java/*.cc lib/*.cc)
+add_library(plasma_java SHARED
+  ${PLASMA_LIBRARY_EXT_java_SRC}
+  )
+
+if(APPLE)
+  target_link_libraries(plasma_java plasma_static ${PLASMA_LINK_LIBS} "-undefined dynamic_lookup" -Wl,-force_load,${FLATBUFFERS_STATIC_LIB} ${FLATBUFFERS_STATIC_LIB} ${PTHREAD_LIBRARY})
+else(APPLE)
+  target_link_libraries(plasma_java plasma_static ${PLASMA_LINK_LIBS} -Wl,--whole-archive ${FLATBUFFERS_STATIC_LIB} -Wl,--no-whole-archive ${FLATBUFFERS_STATIC_LIB} ${PTHREAD_LIBRARY})
+endif(APPLE)
+
 #######################################
 # Unit tests
 #######################################

cpp/src/plasma/lib/java/org_apache_arrow_plasma_PlasmaClientJNI.cc

@@ -0,0 +1,322 @@
+#include <string>
+#include <cstring>
+#include <iostream>
+#include <stdlib.h>
+#include <pthread.h>
+#include <unistd.h>
+#include "plasma/client.h"
+#include "org_apache_arrow_plasma_PlasmaClientJNI.h"
+
+using namespace std;
+using namespace plasma;
+using namespace arrow;
+
+const jsize LEN_OF_OBJECTID = sizeof(ObjectID) / sizeof(jbyte);
+
+inline void jbyteArray_to_object_id(JNIEnv *env, jbyteArray a, ObjectID *oid)
+{
+    env->GetByteArrayRegion(a, 0, LEN_OF_OBJECTID, (jbyte *) oid);
+}
+
+inline void object_id_to_jbyteArray(JNIEnv *env, jbyteArray a, ObjectID *oid)
+{
+    env->SetByteArrayRegion(a, 0, LEN_OF_OBJECTID, (jbyte *) oid);
+}
+
+class JByteArrayGetter {
+private:
+    JNIEnv *_env;
+    jbyteArray _a;
+    jbyte *bp;
+
+public:
+
+    JByteArrayGetter(JNIEnv *env, jbyteArray a, jbyte **out) {
+        _env = env;
+        _a = a;
+
+        bp = _env->GetByteArrayElements(_a, NULL);
+        *out = bp;
+    }
+
+    ~JByteArrayGetter() {
+        _env->ReleaseByteArrayElements(_a, bp, 0);
+    }
+};
+
+JNIEXPORT jlong JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_connect
+  (JNIEnv *env, jclass cls, jstring config, jstring overwrites, jstring store_socket_name, jstring manager_socket_name, jint release_delay)
+{
+    const char *s_name = env->GetStringUTFChars(store_socket_name, NULL);
+    const char *m_name = env->GetStringUTFChars(manager_socket_name, NULL);
+
+    PlasmaClient *client = new PlasmaClient();
+    ARROW_CHECK_OK(
+        client->Connect(s_name, m_name, release_delay));
+
+    //fprintf (stdout, "JNI plasma client init, fd = %d\n", client->get_store_fd());
+    //fflush (stdout);
+
+    env->ReleaseStringUTFChars(store_socket_name, s_name);
+    env->ReleaseStringUTFChars(manager_socket_name, m_name);
+    return (long) client;
+}
+
+JNIEXPORT void JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_disconnect
+  (JNIEnv *env, jclass cls, jlong conn)
+{
+    PlasmaClient *client = (PlasmaClient *) conn;
+
+    //fprintf (stdout, "JNI plasma client disconnect, fd = %d\n", client->get_store_fd());
+    //fflush (stdout);
+
+    ARROW_CHECK_OK(client->Disconnect());
+    delete client;
+    return;
+}
+
+JNIEXPORT jobject JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_create
+  (JNIEnv *env, jclass cls, jlong conn, jbyteArray object_id, jint size, jbyteArray metadata)
+{
+    PlasmaClient *client = (PlasmaClient *) conn;
+    ObjectID oid;
+    jbyteArray_to_object_id(env, object_id, &oid);
+
+    // prepare metadata buffer
+    uint8_t *md = NULL;
+    jsize md_size = 0;
+    std::unique_ptr<JByteArrayGetter> md_getter;
+    if (metadata != NULL) {
+        md_size = env->GetArrayLength(metadata);
+    }
+    if (md_size > 0) {
+        md_getter.reset(new JByteArrayGetter(env, metadata, (jbyte **)&md));
+    }
+
+    std::shared_ptr<Buffer> data;
+    Status s = client->Create(oid, size, md, md_size, &data);
+    if (s.IsPlasmaObjectExists()) {
+        jclass Exception = env->FindClass("org/ray/spi/impl/PlasmaObjectExistsException");
+        env->ThrowNew(Exception, "An object with this ID already exists in the plasma store.");
+        return NULL;
+    }
+    if (s.IsPlasmaStoreFull()) {
+        jclass Exception = env->FindClass("org/ray/spi/impl/PlasmaOutOfMemoryException");
+        env->ThrowNew(Exception, "The plasma store ran out of memory and could not create this object.");
+        return NULL;
+    }
+    ARROW_CHECK(s.ok());
+
+    return env->NewDirectByteBuffer(data->mutable_data(), size);
+}
+
+JNIEXPORT jbyteArray JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_hash
+  (JNIEnv *env, jclass cls, jlong conn, jbyteArray object_id)
+{
+    PlasmaClient *client = (PlasmaClient *) conn;
+    ObjectID oid;
+    jbyteArray_to_object_id(env, object_id, &oid);
+
+    unsigned char digest[kDigestSize];
+    bool success = client->Hash(oid, digest).ok();
+
+    if (success) {
+        jbyteArray ret = env->NewByteArray(kDigestSize);
+        env->SetByteArrayRegion(ret, 0, kDigestSize, (jbyte *) digest);
+        return ret;
+    }
+    else {
+        return NULL;
+    }
+}
+
+JNIEXPORT void JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_seal
+  (JNIEnv *env, jclass cls, jlong conn, jbyteArray object_id)
+{
+    PlasmaClient *client = (PlasmaClient *) conn;
+    ObjectID oid;
+    jbyteArray_to_object_id(env, object_id, &oid);
+
+    ARROW_CHECK_OK(client->Seal(oid));
+}
+
+JNIEXPORT void JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_release
+  (JNIEnv *env, jclass cls, jlong conn, jbyteArray object_id)
+{
+    PlasmaClient *client = (PlasmaClient *) conn;
+    ObjectID oid;
+    jbyteArray_to_object_id(env, object_id, &oid);
+
+    //fprintf (stdout, "JNI plasma client release, fd = %d\n", client->get_store_fd());
+    //fflush (stdout);
+
+    ARROW_CHECK_OK(client->Release(oid));
+}
+
+JNIEXPORT jobjectArray JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_get
+  (JNIEnv *env, jclass cls, jlong conn, jobjectArray object_ids, jint timeout_ms)
+{
+    PlasmaClient *client = (PlasmaClient *) conn;
+    //fprintf (stdout, "JNI plasma client get, fd = %d\n", client->get_store_fd());
+    //fflush (stdout);
+
+    jsize num_oids = env->GetArrayLength(object_ids);
+    std::vector<ObjectID> oids(num_oids);
+    std::vector<ObjectBuffer> obufs(num_oids);
+    for (int i = 0; i < num_oids; ++i) {
+        jbyteArray_to_object_id(env,
+                (jbyteArray) env->GetObjectArrayElement(object_ids, i), &oids[i]);
+    }
+    // TODO: may be blocked. consider to add the thread support
+    ARROW_CHECK_OK(
+            client->Get(oids.data(), num_oids, timeout_ms, obufs.data()));
+
+    jclass clsByteBuffer = env->FindClass("java/nio/ByteBuffer");
+    jclass clsByteBufferArray = env->FindClass("[Ljava/nio/ByteBuffer;");
+
+    jobjectArray ret = env->NewObjectArray(num_oids, clsByteBufferArray, NULL);
+    jobjectArray o = NULL;
+    jobject dataBuf, metadataBuf;
+    for (int i = 0; i < num_oids; ++i) {
+        o = env->NewObjectArray(2, clsByteBuffer, NULL);
+        if (obufs[i].data && obufs[i].data->size() != -1 ) {
+            dataBuf = env->NewDirectByteBuffer(const_cast<uint8_t *>(obufs[i].data->data()), obufs[i].data->size());
+            if (obufs[i].metadata && obufs[i].metadata->size() > 0) {
+                metadataBuf = env->NewDirectByteBuffer(const_cast<uint8_t *>(obufs[i].metadata->data()), obufs[i].metadata->size());
+            }
+            else {
+                metadataBuf = NULL;
+            }
+        }
+        else {
+            dataBuf = NULL;
+            metadataBuf = NULL;
+        }
+        
+        env->SetObjectArrayElement(o, 0, dataBuf);
+        env->SetObjectArrayElement(o, 1, metadataBuf);
+        env->SetObjectArrayElement(ret, i, o);
+    }
+    return ret;
+}
+
+JNIEXPORT jboolean JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_contains
+  (JNIEnv *env, jclass cls, jlong conn, jbyteArray object_id)
+{
+    PlasmaClient *client = (PlasmaClient *) conn;
+    ObjectID oid;
+    jbyteArray_to_object_id(env, object_id, &oid);
+
+    bool has_object;
+    ARROW_CHECK_OK(client->Contains(oid, &has_object));
+
+    if (has_object) {
+        return true;
+    }
+    else {
+        return false;
+    }
+}
+
+JNIEXPORT void JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_fetch
+  (JNIEnv *env, jclass cls, jlong conn, jobjectArray object_ids)
+{
+    PlasmaClient *client = (PlasmaClient *) conn;
+    jsize num_oids = env->GetArrayLength(object_ids);
+
+    /*
+    if (!plasma_manager_is_connected(client)) {
+        jclass Exception = env->FindClass("java/lang/RuntimeException");
+        env->ThrowNew(Exception, "Not connected to the plasma manager.");
+        return;
+    }
+    */
+
+    std::vector<ObjectID> oids(num_oids);
+    for (int i = 0; i < num_oids; ++i) {
+        jbyteArray_to_object_id(env,
+                (jbyteArray) env->GetObjectArrayElement(object_ids, i), &oids[i]);
+    }
+
+    ARROW_CHECK_OK(client->Fetch((int) num_oids, oids.data()));
+
+    return;
+}
+
+JNIEXPORT jobjectArray JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_wait
+  (JNIEnv *env, jclass cls, jlong conn, jobjectArray object_ids, jint timeout_ms, jint num_returns)
+{
+    PlasmaClient *client = (PlasmaClient *) conn;
+    jsize num_oids = env->GetArrayLength(object_ids);
+
+    /*
+    if (!plasma_manager_is_connected(client)) {
+        jclass Exception = env->FindClass("java/lang/RuntimeException");
+        env->ThrowNew(Exception, "Not connected to the plasma manager.");
+        return NULL;
+    }
+    */
+
+    if (num_returns < 0) {
+        jclass Exception = env->FindClass("java/lang/RuntimeException");
+        env->ThrowNew(Exception, "The argument num_returns cannot be less than zero.");
+        return NULL;
+    }
+    if (num_returns > num_oids) {
+        jclass Exception = env->FindClass("java/lang/RuntimeException");
+        env->ThrowNew(Exception,
+                "The argument num_returns cannot be greater than len(object_ids).");
+        return NULL;
+    }
+
+    std::vector<ObjectRequest> oreqs(num_oids);
+
+    for (int i = 0; i < num_oids; ++i) {
+        jbyteArray_to_object_id(env,
+                (jbyteArray) env->GetObjectArrayElement(object_ids, i), &oreqs[i].object_id);
+        oreqs[i].type = PLASMA_QUERY_ANYWHERE;
+    }
+
+    int num_return_objects;
+    // TODO: may be blocked. consider to add the thread support
+    ARROW_CHECK_OK(client->Wait((int) num_oids, oreqs.data(), num_returns,
+                (uint64_t) timeout_ms, &num_return_objects));
+
+    int num_to_return = min(num_return_objects, num_returns);
+    jclass clsByteArray = env->FindClass("[B");
+    jobjectArray ret = env->NewObjectArray(num_to_return, clsByteArray, NULL);
+
+    int num_returned = 0;
+    jbyteArray oid = NULL;
+    for (int i = 0; i < num_oids; ++i) {
+        if (num_returned >= num_to_return) {
+            break;
+        }
+
+        if (oreqs[i].status == ObjectStatusLocal
+                || oreqs[i].status == ObjectStatusRemote) {
+            oid = env->NewByteArray(LEN_OF_OBJECTID);
+            object_id_to_jbyteArray(env, oid, &oreqs[i].object_id);
+            env->SetObjectArrayElement(ret, num_returned, oid);
+            num_returned++;
+        }
+        else {
+            //ARROW_CHECK(oreqs[i].status == ObjectStatus_Nonexistent);
+        }
+    }
+    ARROW_CHECK(num_returned == num_to_return);    
+
+    return ret;
+}
+
+JNIEXPORT jlong JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_evict
+  (JNIEnv *env, jclass cls, jlong conn, jlong num_bytes)
+{
+    PlasmaClient *client = (PlasmaClient *) conn;
+
+    int64_t evicted_bytes;
+    ARROW_CHECK_OK(client->Evict((int64_t) num_bytes, evicted_bytes));
+
+    return (jlong) evicted_bytes;
+}
+