switch to critical mem access for checksums (#878)

Author: DmitriyMusatkin

src/native/checksums.c

@@ -9,19 +9,29 @@
 #include "crt.h"
 #include "java_class_ids.h"
 
+/**
+ * Note: we use critical mem access in below functions to speed up checksums.
+ * This approach is the same as what OpenJDK uses for their checksum implementation.
+ * Think twice before using similar approach elsewhere as it can lead to stalling.
+ */
+
 jint crc32_common(
     JNIEnv *env,
     jbyteArray input,
     jint previous,
     const size_t start,
     size_t length,
     uint32_t (*checksum_fn)(const uint8_t *, size_t, uint32_t)) {
-    struct aws_byte_cursor c_byte_array = aws_jni_byte_cursor_from_jbyteArray_acquire(env, input);
+    struct aws_byte_cursor c_byte_array = aws_jni_byte_cursor_from_jbyteArray_critical_acquire(env, input);
+    if (AWS_UNLIKELY(c_byte_array.ptr == NULL)) {
+        return previous;
+    }
+
     struct aws_byte_cursor cursor = c_byte_array;
     aws_byte_cursor_advance(&cursor, start);
     cursor.len = aws_min_size(length, cursor.len);
     jint res_signed = (jint)checksum_fn(cursor.ptr, cursor.len, previous);
-    aws_jni_byte_cursor_from_jbyteArray_release(env, input, c_byte_array);
+    aws_jni_byte_cursor_from_jbyteArray_critical_release(env, input, c_byte_array);
     return res_signed;
 }
 
@@ -61,11 +71,15 @@ JNIEXPORT jlong JNICALL Java_software_amazon_awssdk_crt_checksums_CRC64NVME_crc6
     (void)jni_class;
     aws_cache_jni_ids(env);
 
-    struct aws_byte_cursor c_byte_array = aws_jni_byte_cursor_from_jbyteArray_acquire(env, input);
+    struct aws_byte_cursor c_byte_array = aws_jni_byte_cursor_from_jbyteArray_critical_acquire(env, input);
+    if (AWS_UNLIKELY(c_byte_array.ptr == NULL)) {
+        return previous;
+    }
+
     struct aws_byte_cursor cursor = c_byte_array;
     aws_byte_cursor_advance(&cursor, offset);
     cursor.len = aws_min_size(length, cursor.len);
     jlong res_signed = (jlong)aws_checksums_crc64nvme_ex(cursor.ptr, cursor.len, previous);
-    aws_jni_byte_cursor_from_jbyteArray_release(env, input, c_byte_array);
+    aws_jni_byte_cursor_from_jbyteArray_critical_release(env, input, c_byte_array);
     return res_signed;
 }

src/native/crt.c

@@ -273,6 +273,15 @@ void aws_jni_throw_null_pointer_exception(JNIEnv *env, const char *msg, ...) {
     (*env)->ThrowNew(env, (*env)->FindClass(env, "java/lang/NullPointerException"), buf);
 }
 
+void aws_jni_throw_out_of_memory_exception(JNIEnv *env, const char *msg, ...) {
+    va_list args;
+    va_start(args, msg);
+    char buf[1024];
+    vsnprintf(buf, sizeof(buf), msg, args);
+    va_end(args);
+    (*env)->ThrowNew(env, (*env)->FindClass(env, "java/lang/OutOfMemoryError"), buf);
+}
+
 void aws_jni_throw_illegal_argument_exception(JNIEnv *env, const char *msg, ...) {
     va_list args;
     va_start(args, msg);
@@ -441,12 +450,34 @@ struct aws_byte_cursor aws_jni_byte_cursor_from_jbyteArray_acquire(JNIEnv *env,
     return aws_byte_cursor_from_array(bytes, len);
 }
 
+struct aws_byte_cursor aws_jni_byte_cursor_from_jbyteArray_critical_acquire(JNIEnv *env, jbyteArray array) {
+    if (array == NULL) {
+        aws_jni_throw_null_pointer_exception(env, "byte[] is null");
+        return aws_byte_cursor_from_array(NULL, 0);
+    }
+
+    jbyte *bytes = (*env)->GetPrimitiveArrayCritical(env, array, NULL);
+    if (bytes == NULL) {
+        aws_jni_throw_out_of_memory_exception(env, "failed to acquire array memory");
+        return aws_byte_cursor_from_array(NULL, 0);
+    }
+
+    size_t len = (*env)->GetArrayLength(env, array);
+    return aws_byte_cursor_from_array(bytes, len);
+}
+
 void aws_jni_byte_cursor_from_jbyteArray_release(JNIEnv *env, jbyteArray array, struct aws_byte_cursor cur) {
     if (cur.ptr != NULL) {
         (*env)->ReleaseByteArrayElements(env, array, (jbyte *)cur.ptr, JNI_ABORT);
     }
 }
 
+void aws_jni_byte_cursor_from_jbyteArray_critical_release(JNIEnv *env, jbyteArray array, struct aws_byte_cursor cur) {
+    if (cur.ptr != NULL) {
+        (*env)->ReleasePrimitiveArrayCritical(env, array, (jbyte *)cur.ptr, JNI_ABORT);
+    }
+}
+
 struct aws_byte_cursor aws_jni_byte_cursor_from_direct_byte_buffer(JNIEnv *env, jobject byte_buffer) {
     jlong payload_size = (*env)->GetDirectBufferCapacity(env, byte_buffer);
     if (payload_size == -1) {
@@ -619,7 +650,7 @@ void JNICALL Java_software_amazon_awssdk_crt_CRT_awsCrtInit(
  * AWS-LC always defines FIPS_mode() that you can call and check what the library was built with. It does not define
  * a public OPENSSL_FIPS/AWSLC_FIPS macro that we can (or need to) check here
  *
- * Safeguard with macro's, for example because Libressl dosn't define
+ * Safeguard with macro's, for example because Libressl doesn't define
  * FIPS_mode() by default.
  * */
 #if defined(OPENSSL_FIPS) || defined(OPENSSL_IS_AWSLC)

src/native/crt.h

@@ -46,6 +46,11 @@ void aws_jni_throw_runtime_exception(JNIEnv *env, const char *msg, ...);
  ******************************************************************************/
 void aws_jni_throw_null_pointer_exception(JNIEnv *env, const char *msg, ...);
 
+/*******************************************************************************
+ * Throws java OutOfMemoryError
+ ******************************************************************************/
+void aws_jni_throw_out_of_memory_exception(JNIEnv *env, const char *msg, ...);
+
 /*******************************************************************************
  * Throws java IllegalArgumentException
  ******************************************************************************/
@@ -172,19 +177,40 @@ void aws_jni_byte_cursor_from_jstring_release(JNIEnv *env, jstring str, struct a
 /*******************************************************************************
  * aws_jni_byte_cursor_from_jbyteArray_acquire - Creates an aws_byte_cursor from the
  * bytes extracted from the supplied jbyteArray.
- * The aws_byte_cursor MUST be given to aws_jni_byte_cursor_from jstring_release() when
+ * The aws_byte_cursor MUST be given to aws_jni_byte_cursor_from_jbyteArray_release() when
  * it's no longer needed, or it will leak.
  *
  * If there is an error, the returned aws_byte_cursor.ptr will be NULL and
  * and a java exception is being thrown.
  ******************************************************************************/
 struct aws_byte_cursor aws_jni_byte_cursor_from_jbyteArray_acquire(JNIEnv *env, jbyteArray str);
 
+/*******************************************************************************
+ * aws_jni_byte_cursor_from_jbyteArray_critical_acquire - Creates an aws_byte_cursor from the
+ * bytes extracted from the supplied jbyteArray.
+ * The aws_byte_cursor MUST be given to aws_jni_byte_cursor_from_jbyteArray_critical_release() when
+ * it's no longer needed, or it will leak.
+ *
+ * If there is an error, the returned aws_byte_cursor.ptr will be NULL and
+ * and a java exception is being thrown.
+ * **** WARNING: Here be dragons ****
+ * This function uses critical jvm functions to access memory, which is faster, but comes
+ * with a litany of limitations, so use at your own risk.
+ * Spending too much time in a critical section will stall gc and mem functionality, so under
+ * no conditions use jni functions while in critical section on current thread (ok, on other threads).
+ ******************************************************************************/
+struct aws_byte_cursor aws_jni_byte_cursor_from_jbyteArray_critical_acquire(JNIEnv *env, jbyteArray str);
+
 /********************************************************************************
  * aws_jni_byte_cursor_from_jbyteArray_release - Releases the array back to the JVM
  ********************************************************************************/
 void aws_jni_byte_cursor_from_jbyteArray_release(JNIEnv *env, jbyteArray str, struct aws_byte_cursor cur);
 
+/********************************************************************************
+ * aws_jni_byte_cursor_from_jbyteArray_critical_release - Releases critical array access back to JVM
+ ********************************************************************************/
+void aws_jni_byte_cursor_from_jbyteArray_critical_release(JNIEnv *env, jbyteArray str, struct aws_byte_cursor cur);
+
 /*******************************************************************************
  * aws_jni_byte_cursor_from_direct_byte_buffer - Creates an aws_byte_cursor from the
  * direct byte buffer. Note that the buffer is not reference pinned, so the cursor