Avoid long TTSP pauses caused by Unsafe.allocateMemory/freeMemory

unmalloc.c

// Replaces implementation of Unsafe.allocateMemory/freeMemory
// to avoid long time-to-safepoint pauses.
//
// Compile: gcc -O3 -fno-omit-frame-pointer -fPIC -shared -olibunmalloc.so unmalloc.c
//
// Usage:   java -agentpath:/path/to/libunmalloc.so ...

#include <jvmti.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

#define AGENT_NAME "unmalloc"


jlong JNICALL unmalloc_AllocateMemory(JNIEnv* env, jobject unsafe, jlong size) {
    return (intptr_t) malloc((size_t)size);
}

jlong JNICALL unmalloc_ReallocateMemory(JNIEnv* env, jobject unsafe, jlong addr, jlong size) {
    return (intptr_t) realloc((void*)(intptr_t)addr, (size_t)size);
}

void JNICALL unmalloc_FreeMemory(JNIEnv* env, jobject unsafe, jlong addr) {
    free((void*)(intptr_t)addr);
}


static const JNINativeMethod new_unsafe_methods[] = {
    {(char*)"allocateMemory0",   (char*)"(J)J",  (void*)unmalloc_AllocateMemory},
    {(char*)"reallocateMemory0", (char*)"(JJ)J", (void*)unmalloc_ReallocateMemory},
    {(char*)"freeMemory0",       (char*)"(J)V",  (void*)unmalloc_FreeMemory}
};

static const JNINativeMethod old_unsafe_methods[] = {
    {(char*)"allocateMemory",   (char*)"(J)J",  (void*)unmalloc_AllocateMemory},
    {(char*)"reallocateMemory", (char*)"(JJ)J", (void*)unmalloc_ReallocateMemory},
    {(char*)"freeMemory",       (char*)"(J)V",  (void*)unmalloc_FreeMemory}
};

int redefine_unsafe_methods(JNIEnv* env) {
    jclass unsafe;
    jint result;

    if ((unsafe = (*env)->FindClass(env, "jdk/internal/misc/Unsafe")) != NULL) {
        result = (*env)->RegisterNatives(env, unsafe, new_unsafe_methods, 3);
    } else if ((unsafe = (*env)->FindClass(env, "sun/misc/Unsafe")) != NULL) {
        result = (*env)->RegisterNatives(env, unsafe, old_unsafe_methods, 3);
    } else {
        printf("[" AGENT_NAME "] could not find Unsafe class\n");
        (*env)->ExceptionClear(env);
        return 1;
    }

    if (result) {
        printf("[" AGENT_NAME "] failed to redefine Unsafe memory methods\n");
    } else {
        printf("[" AGENT_NAME "] Unsafe memory methods redefined\n");
    }

    (*env)->ExceptionClear(env);
    return result;
}

void JNICALL VMInit(jvmtiEnv* jvmti, JNIEnv* jni, jthread thread) {
    redefine_unsafe_methods(jni);
}


// Entry point for -agentpath
JNIEXPORT jint JNICALL Agent_OnLoad(JavaVM* vm, char* options, void* reserved) {
    jvmtiEnv* jvmti;
    (*vm)->GetEnv(vm, (void**)&jvmti, JVMTI_VERSION_1_0);

    jvmtiEventCallbacks callbacks = {0};
    callbacks.VMInit = VMInit;
    (*jvmti)->SetEventCallbacks(jvmti, &callbacks, sizeof(callbacks));
    (*jvmti)->SetEventNotificationMode(jvmti, JVMTI_ENABLE, JVMTI_EVENT_VM_INIT, NULL);

    return 0;
}

// Entry point for attaching agent in runtime
JNIEXPORT jint JNICALL Agent_OnAttach(JavaVM* vm, char* options, void* reserved) {
    JNIEnv* jni;
    (*vm)->GetEnv(vm, (void**)&jni, JNI_VERSION_1_6);

    return redefine_unsafe_methods(jni);
}