JVMTI Agent for Bytecode Instrumentation of Hot Methods

agent.c

/*
 * agent.c
 *
 *  Created on: Dec 9, 2016
 *      Author: Saqib Ahmed
 */
#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <string.h>
#include "class.h"
#include "jvmti.h"
#include "jni.h"

/* Global agent data structure */

typedef struct {
	/* JVMTI Environment */
	jvmtiEnv *jvmti;
	JNIEnv * jni;
	jboolean vm_is_started;
	jboolean vmDead;

	/* Data access Lock */
	jrawMonitorID lock;
	JavaVM* jvm;
} GlobalAgentData;

static jrawMonitorID lock;

static GlobalAgentData *gdata;

void fatal_error(const char * format, ...) {
	va_list ap;

	va_start(ap, format);
	(void) vfprintf(stderr, format, ap);
	(void) fflush(stderr);
	va_end(ap);
	exit(3);
}

void check_jvmti_error(jvmtiEnv *jvmti, jvmtiError errnum, const char *str) {
	if (errnum != JVMTI_ERROR_NONE) {
		char *errnum_str;

		errnum_str = NULL;
		(void) (*jvmti)->GetErrorName(jvmti, errnum, &errnum_str);

		fatal_error("ERROR: JVMTI: %d(%s): %s\n", errnum,
				(errnum_str == NULL ? "Unknown" : errnum_str),
				(str == NULL ? "" : str));
	}
}

static int x = 1;
void JNICALL
compiled_method_load(jvmtiEnv *jvmti, jmethodID method, jint code_size,
        const void* code_addr, jint map_length, const jvmtiAddrLocationMap* map,
        const void* compile_info) {
    jvmtiError err;
    jclass klass;

    char* name = NULL;
    char* signature = NULL;
    char* generic_ptr = NULL;

    err = (*jvmti)->RawMonitorEnter(jvmti, lock);
    check_jvmti_error(jvmti, err, "raw monitor enter");

    err = (*jvmti)->GetMethodName(jvmti, method, &name, &signature,
            &generic_ptr);
    check_jvmti_error(jvmti, err, "Get Method Name");

    printf("\nCompiled method load event\n");
    printf("Method name %s %s %s\n\n", name, signature,
            generic_ptr == NULL ? "" : generic_ptr);

    if (strstr(name, "main") != NULL && x == 1) {
        x++;
        err = (*jvmti)->GetMethodDeclaringClass(jvmti, method, &klass);
        check_jvmti_error(jvmti, err, "Get Declaring Class");

        err = (*jvmti)->RetransformClasses(jvmti, 1, &klass);
        check_jvmti_error(jvmti, err, "Retransform class");

    }

    if (name != NULL) {
        err = (*jvmti)->Deallocate(jvmti, (unsigned char*) name);
        check_jvmti_error(jvmti, err, "deallocate name");
    }
    if (signature != NULL) {
        err = (*jvmti)->Deallocate(jvmti, (unsigned char*) signature);
        check_jvmti_error(jvmti, err, "deallocate signature");
    }
    if (generic_ptr != NULL) {
        err = (*jvmti)->Deallocate(jvmti, (unsigned char*) generic_ptr);
        check_jvmti_error(jvmti, err, "deallocate generic_ptr");
    }

    err = (*jvmti)->RawMonitorExit(jvmti, lock);
    check_jvmti_error(jvmti, err, "raw monitor exit");
}

void JNICALL
Class_File_Load_Hook(jvmtiEnv *jvmti_env, JNIEnv* jni_env,
        jclass class_being_redefined, jobject loader, const char* name,
        jobject protection_domain, jint class_data_len,
        const unsigned char* class_data, jint* new_class_data_len,
        unsigned char** new_class_data) {
    jvmtiError err;
    unsigned char* jvmti_space = NULL;

    if (strstr(name, "Test") != NULL && x == 2) {
        char* args = "op";

        javab_main(2, args, class_data, class_data_len);

        err = (*jvmti_env)->Allocate(jvmti_env, (jlong)global_pos, &jvmti_space);
        check_jvmti_error(jvmti_env, err, "Allocate new class Buffer.");

        (void)memcpy((void*)jvmti_space, (void*)new_class_ptr, (int)global_pos);

        *new_class_data_len = (jint)global_pos;
        *new_class_data = jvmti_space;

        if ( new_class_ptr != NULL ) {
            (void)free((void*)new_class_ptr);
        }

#if DEBUG
        printf("Size of the class is: %d\n", class_data_len);
        for (int i = 0; i < class_data_len; i += 4) {
            if (i % 16 == 0)
                printf("\n");

            printf("%02x%02x  %02x%02x  ", new_class_data[i],
                    new_class_data[i + 1], new_class_data[i + 2],
                    new_class_data[i + 3]);
        }
        printf("\n");
        system("javap -c -v Test_debug");
#endif
        x++;
    }
}

#if DEBUG_THREADS
static void JNICALL callbackThreadStart(jvmtiEnv *jvmti, JNIEnv* jni_env,
		jthread thread) {

	jvmtiError err1, err2;
	jvmtiThreadInfo info1;

	/* Make sure the stack variables are garbage free */
	(void) memset(&info1, 0, sizeof(info1));
	err1 = (*jvmti)->GetThreadInfo(jvmti, thread, &info1);
//	check_jvmti_error(jvmti, err1, "Get Thread Information");

	if (err1 == JVMTI_ERROR_NONE)
		printf("Running Thread: %s, Priority: %d, context class loader:%s\n",
				info1.name, info1.priority,
				(info1.context_class_loader == NULL ? ": NULL" : "Not Null"));

	/* Every string allocated by JVMTI needs to be freed */
	err2 = (*jvmti)->Deallocate(jvmti, (unsigned char*) info1.name);
	check_jvmti_error(jvmti, err2, "Dealocate Thread Name");
}
#endif

JNIEXPORT jint JNICALL Agent_OnLoad(JavaVM *jvm, char *options, void *reserved) {
	static jvmtiEnv *jvmti = NULL;
	static jvmtiCapabilities capa;
	jvmtiError error;
	jvmtiEventCallbacks callbacks;
	static GlobalAgentData data;
	jint res;

	memset((void*) &data, 0, sizeof(data));
	gdata = &data;

	gdata->jvm = jvm;
	res = (*jvm)->GetEnv(jvm, (void **) &jvmti, JVMTI_VERSION_1_0);

	if (res != JNI_OK || jvmti == NULL) {
		/* This means that the VM was unable to obtain this version of the
		 *   JVMTI interface, this is a fatal error.
		 */
		printf("ERROR: Unable to access JVMTI Version 1 (0x%x),"
				" is your J2SE a 1.5 or newer version?"
				" JNIEnv's GetEnv() returned %d\n", JVMTI_VERSION_1, res);

	}

	memset(&capa, 0, sizeof(jvmtiCapabilities));

	capa.can_generate_compiled_method_load_events = 1;
#if DEBUG_THREADS
	capa.can_signal_thread = 1;
#endif
	capa.can_redefine_classes = 1;
	capa.can_redefine_any_class = 1;
	capa.can_retransform_classes = 1;
	capa.can_retransform_any_class = 1;
	capa.can_get_bytecodes = 1;
	capa.can_get_constant_pool = 1;
	capa.can_access_local_variables = 1;
	capa.can_generate_all_class_hook_events = 1;

	error = (*jvmti)->AddCapabilities(jvmti, &capa);
	check_jvmti_error(jvmti, error, "Add Capabilities");

#if DEBUG_THREADS
	error = (*jvmti)->SetEventNotificationMode (jvmti, JVMTI_ENABLE,
			JVMTI_EVENT_THREAD_START, (jthread)NULL);
	check_jvmti_error(jvmti, error, "Set Event for Thread Start");
#endif

	/* enable JVMTI events */
	error = (*jvmti)->SetEventNotificationMode(jvmti, JVMTI_ENABLE,
			JVMTI_EVENT_COMPILED_METHOD_LOAD, NULL);
	check_jvmti_error(jvmti, error, "Set Event for Compiled Method Load");

	error = (*jvmti)->SetEventNotificationMode(jvmti, JVMTI_ENABLE,
			JVMTI_EVENT_CLASS_FILE_LOAD_HOOK, NULL);
	check_jvmti_error(jvmti, error, "Set Event for Compiled Method Load");

	memset(&callbacks, 0, sizeof(callbacks));

#if DEBUG_THREADS
	callbacks.ThreadStart = &callbackThreadStart;/* JVMTI_EVENT_THREAD_START */
#endif

	callbacks.CompiledMethodLoad = &compiled_method_load; /* JVMTI_COMPILED_METHOD_LOAD */

	callbacks.ClassFileLoadHook = &Class_File_Load_Hook;

	error = (*jvmti)->SetEventCallbacks(jvmti, &callbacks,
			(jint) sizeof(callbacks));
	check_jvmti_error(jvmti, error, "Set Event for CallBacks");

	/* create coordination monitor */
	error = (*jvmti)->CreateRawMonitor(jvmti, "agent lock", &lock);
	check_jvmti_error(jvmti, error, "Create raw Monitor");

	return JNI_OK;
}

Test.java

class Test {
	static final int SIZE = 700;
	static int dum = 0;
	static int x = 0;
	static final int start = 400;
	public static void main(String arg[]) {

		long t1, t2;
		int[][] a, b, c, c2, c3;

		a = new int[SIZE][SIZE];
		b = new int[SIZE][SIZE];
		c = new int[SIZE][SIZE];
		c2 = new int[SIZE][SIZE];
		c3 = new int[SIZE][SIZE];
		// INIT

		for (int i = 0; i < SIZE; i++)
			for (int j = 0; j < SIZE; j++) {
				a[i][j] = 1;
				b[i][j] = 2;
				dum++; // prevents parallelization
			}

		for (int s = start; s < SIZE; s += 100) {
			
			t1 = System.currentTimeMillis();
			for (int i = 0; i < s; i++) {
				for (int j = 0; j < s; j++) {
					c[i][j] = 0;
					for (int k = 0; k < s; k++) {
						c[i][j] += a[i][k] * b[k][j];
						dum++; // prevents parallelization
					}
				}
			}			
			t2 = System.currentTimeMillis();
			
			double d = (t2 - t1) / 1000.0d;

			System.out.println("Serial: "+ d);
			
			t1 = System.currentTimeMillis();
			for (int i = 0; i < s; i++)
				for (int j = 0; j < s; j++) {
					c2[i][j] = 0;
					for (int k = 0; k < s; k++)
						c2[i][j] += a[i][k] * b[k][j];
				}
			
			t2 = System.currentTimeMillis();
			
			d = (t2 - t1) / 1000.0d;

			System.out.println("Parallel: " + d);				

			
			t1 = System.currentTimeMillis();
			mult(a, b, c3, s);
			t2 = System.currentTimeMillis();
					
			
			for (int i = 0; i < s; i++)
			for (int j = 0; j < s; j++)
				if (c2[i][j] != c[i][j] || c2[i][j] != c3[i][j]) {
					System.out.println("ERROR!");
					System.exit(1);
				}
			
			d = (t2 - t1) / 1000.0d;

			System.out.println("Parallel Function call: " + d);

		}
	}
	
	static void mult(int[][] a, int[][] b, int[][] c, int seed)
	{
		
		for (int i = 0; i < seed; i++)
			for (int j = 0; j < seed; j++) {
				c[i][j] = 0;
				for (int k = 0; k < seed; k++)
				{
					c[i][j] += a[i][k] * b[k][j];
				}
			}
	}
}