/*

 *

 * Redistribution and use in source and binary forms, with or without

 * modification, are permitted provided that the following conditions

 * are met:

 *

 *   - Redistributions of source code must retain the above copyright

 *     notice, this list of conditions and the following disclaimer.

 *

 *   - Redistributions in binary form must reproduce the above copyright

 *     notice, this list of conditions and the following disclaimer in the

 *     documentation and/or other materials provided with the distribution.

 *

 *   - Neither the name of Oracle nor the names of its

 *     contributors may be used to endorse or promote products derived

 *     from this software without specific prior written permission.

 *

 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS

 * IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,

 * THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR

 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR

 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,

 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,

 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR

 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF

 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING

 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS

 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

 */

/*

 * This source code is provided to illustrate the usage of a given feature

 * or technique and has been deliberately simplified. Additional steps

 * required for a production-quality application, such as security checks,

 * input validation and proper error handling, might not be present in

 * this sample code.

 */

#include "stdlib.h"

#include "heapTracker.h"

#include "java_crw_demo.h"

#include "jni.h"

#include "jvmti.h"

#include "agent_util.h"

/* -------------------------------------------------------------------

 * Some constant names that tie to Java class/method names.

 *    We assume the Java class whose static methods we will be calling

 *    looks like:

 *

 * public class HeapTracker {

 *     private static int engaged;

 *     private static native void _newobj(Object thr, Object o);

 *     public static void newobj(Object o)

 *     {

 *              if ( engaged != 0 ) {

 *               _newobj(Thread.currentThread(), o);

 *           }

 *     }

 *     private static native void _newarr(Object thr, Object a);

 *     public static void newarr(Object a)

 *     {

 *            if ( engaged != 0 ) {

 *               _newarr(Thread.currentThread(), a);

 *           }

 *     }

 * }

 *

 *    The engaged field allows us to inject all classes (even system classes)

 *    and delay the actual calls to the native code until the VM has reached

 *    a safe time to call native methods (Past the JVMTI VM_START event).

 *

 */

#define HEAP_TRACKER_class           HeapTracker /* Name of class we are using */

#define HEAP_TRACKER_newobj        newobj   /* Name of java init method */

#define HEAP_TRACKER_newarr        newarr   /* Name of java newarray method */

#define HEAP_TRACKER_native_newobj _newobj  /* Name of java newobj native */

#define HEAP_TRACKER_native_newarr _newarr  /* Name of java newarray native */

#define HEAP_TRACKER_engaged       engaged  /* Name of static field switch */

/* C macros to create strings from tokens */

#define _STRING(s) #s

#define STRING(s) _STRING(s)

/* ------------------------------------------------------------------- */

/* Flavors of traces (to separate out stack traces) */

typedef enum {

    TRACE_FIRST                        = 0,

    TRACE_USER                        = 0,

    TRACE_BEFORE_VM_START        = 1,

    TRACE_BEFORE_VM_INIT        = 2,

    TRACE_VM_OBJECT                = 3,

    TRACE_MYSTERY                = 4,

    TRACE_LAST                        = 4

} TraceFlavor;

static char * flavorDesc[] = {

    "",

    "before VM_START",

    "before VM_INIT",

    "VM_OBJECT",

    "unknown"

};

/* Trace (Stack Trace) */

#define MAX_FRAMES 6

typedef struct Trace {

    /* Number of frames (includes HEAP_TRACKER methods) */

    jint           nframes;

    /* Frames from GetStackTrace() (2 extra for HEAP_TRACKER methods) */

    jvmtiFrameInfo frames[MAX_FRAMES+2];

    /* Used to make some traces unique */

    TraceFlavor    flavor;

} Trace;

/* Trace information (more than one object will have this as a tag) */

typedef struct TraceInfo {

    /* Trace where this object was allocated from */

    Trace             trace;

    /* 64 bit hash code that attempts to identify this specific trace */

    jlong             hashCode;

    /* Total space taken up by objects allocated from this trace */

    jlong             totalSpace;

    /* Total count of objects ever allocated from this trace */

    int               totalCount;

    /* Total live objects that were allocated from this trace */

    int               useCount;

    /* The next TraceInfo in the hash bucket chain */

    struct TraceInfo *next;

} TraceInfo;

/* Global agent data structure */

typedef struct {

    /* JVMTI Environment */

    jvmtiEnv      *jvmti;

    /* State of the VM flags */

    jboolean       vmStarted;

    jboolean       vmInitialized;

    jboolean       vmDead;

    /* Options */

    int            maxDump;

    /* Data access Lock */

    jrawMonitorID  lock;

    /* Counter on classes where BCI has been applied */

    jint           ccount;

    /* Hash table to lookup TraceInfo's via Trace's */

    #define HASH_INDEX_BIT_WIDTH 12 /* 4096 */

    #define HASH_BUCKET_COUNT (1<<HASH_INDEX_BIT_WIDTH)

    #define HASH_INDEX_MASK (HASH_BUCKET_COUNT-1)

    TraceInfo     *hashBuckets[HASH_BUCKET_COUNT];

    /* Count of TraceInfo's allocated */

    int            traceInfoCount;

    /* Pre-defined traces for the system and mystery situations */

    TraceInfo     *emptyTrace[TRACE_LAST+1];

} GlobalAgentData;

static GlobalAgentData *gdata;

/* Enter a critical section by doing a JVMTI Raw Monitor Enter */

static void

enterCriticalSection(jvmtiEnv *jvmti)

{

    jvmtiError error;

    error = (*jvmti)->RawMonitorEnter(jvmti, gdata->lock);

    check_jvmti_error(jvmti, error, "Cannot enter with raw monitor");

}

/* Exit a critical section by doing a JVMTI Raw Monitor Exit */

static void

exitCriticalSection(jvmtiEnv *jvmti)

{

    jvmtiError error;

    error = (*jvmti)->RawMonitorExit(jvmti, gdata->lock);

    check_jvmti_error(jvmti, error, "Cannot exit with raw monitor");

}

/* Update stats on a TraceInfo */

static TraceInfo *

updateStats(TraceInfo *tinfo)

{

    tinfo->totalCount++;

    tinfo->useCount++;

    return tinfo;

}

/* Get TraceInfo for empty stack */

static TraceInfo *

emptyTrace(TraceFlavor flavor)

{

    return updateStats(gdata->emptyTrace[flavor]);

}

/* Allocate new TraceInfo */

static TraceInfo *

newTraceInfo(Trace *trace, jlong hashCode, TraceFlavor flavor)

{

    TraceInfo *tinfo;

    tinfo = (TraceInfo*)calloc(1, sizeof(TraceInfo));

    if ( tinfo == NULL ) {

        fatal_error("ERROR: Ran out of malloc() space\n");

    } else {

        int hashIndex;

        tinfo->trace = *trace;

        tinfo->trace.flavor = flavor;

        tinfo->hashCode = hashCode;

        gdata->traceInfoCount++;

        hashIndex = (int)(hashCode & HASH_INDEX_MASK);

        tinfo->next = gdata->hashBuckets[hashIndex];

        gdata->hashBuckets[hashIndex] = tinfo;

    }

    return tinfo;

}

/* Create hash code for a Trace */

static jlong

hashTrace(Trace *trace)

{

    jlong hashCode;

    int   i;

    hashCode = 0;

    for ( i = 0 ; i < trace->nframes ; i++ ) {

        hashCode = (hashCode << 3) +

                (jlong)(ptrdiff_t)(void*)(trace->frames[i].method);

        hashCode = (hashCode << 2) +

                (jlong)(trace->frames[i].location);

    }

    hashCode = (hashCode << 3) + trace->nframes;

    hashCode += trace->flavor;

    return hashCode;

}

/* Lookup or create a new TraceInfo */

static TraceInfo *

lookupOrEnter(jvmtiEnv *jvmti, Trace *trace, TraceFlavor flavor)

{

    TraceInfo *tinfo;

    jlong      hashCode;

    /* Calculate hash code (outside critical section to lessen contention) */

    hashCode = hashTrace(trace);

    /* Do a lookup in the hash table */

    enterCriticalSection(jvmti); {

        TraceInfo *prev;

        int        hashIndex;

        /* Start with first item in hash buck chain */

        prev = NULL;

        hashIndex = (int)(hashCode & HASH_INDEX_MASK);

        tinfo = gdata->hashBuckets[hashIndex];

        while ( tinfo != NULL ) {

            if ( tinfo->hashCode == hashCode &&

                 memcmp(trace, &(tinfo->trace), sizeof(Trace))==0 ) {

                 /* We found one that matches, move to head of bucket chain */

                 if ( prev != NULL ) {

                     /* Remove from list and add to head of list */

                     prev->next = tinfo->next;

                     tinfo->next = gdata->hashBuckets[hashIndex];

                     gdata->hashBuckets[hashIndex] = tinfo;

                 }

                 /* Break out */

                 break;

            }

            prev = tinfo;

            tinfo = tinfo->next;

        }

        /* If we didn't find anything we need to enter a new entry */

        if ( tinfo == NULL ) {

            /* Create new hash table element */

            tinfo = newTraceInfo(trace, hashCode, flavor);

        }

        /* Update stats */

        (void)updateStats(tinfo);

    } exitCriticalSection(jvmti);

    return tinfo;

}

/* Get TraceInfo for this allocation */

static TraceInfo *

findTraceInfo(jvmtiEnv *jvmti, jthread thread, TraceFlavor flavor)

{

    TraceInfo *tinfo;

    jvmtiError error;

    tinfo = NULL;

    if ( thread != NULL ) {

        static Trace  empty;

        Trace         trace;

        /* Before VM_INIT thread could be NULL, watch out */

        trace = empty;

        error = (*jvmti)->GetStackTrace(jvmti, thread, 0, MAX_FRAMES+2,

                            trace.frames, &(trace.nframes));

        /* If we get a PHASE error, the VM isn't ready, or it died */

        if ( error == JVMTI_ERROR_WRONG_PHASE ) {

            /* It is assumed this is before VM_INIT */

            if ( flavor == TRACE_USER ) {

                tinfo = emptyTrace(TRACE_BEFORE_VM_INIT);

            } else {

                tinfo = emptyTrace(flavor);

            }

        } else {

            check_jvmti_error(jvmti, error, "Cannot get stack trace");

            /* Lookup this entry */

            tinfo = lookupOrEnter(jvmti, &trace, flavor);

        }

    } else {

        /* If thread==NULL, it's assumed this is before VM_START */

        if ( flavor == TRACE_USER ) {

            tinfo = emptyTrace(TRACE_BEFORE_VM_START);

        } else {

            tinfo = emptyTrace(flavor);

        }

    }

    return tinfo;

}

/* Tag an object with a TraceInfo pointer. */

static void

tagObjectWithTraceInfo(jvmtiEnv *jvmti, jobject object, TraceInfo *tinfo)

{

    jvmtiError error;

    jlong      tag;

    /* Tag this object with this TraceInfo pointer */

    tag = (jlong)(ptrdiff_t)(void*)tinfo;

    error = (*jvmti)->SetTag(jvmti, object, tag);

    check_jvmti_error(jvmti, error, "Cannot tag object");

}

/* Java Native Method for Object.<init> */

static void JNICALL

HEAP_TRACKER_native_newobj(JNIEnv *env, jclass klass, jthread thread, jobject o)

{

    TraceInfo *tinfo;

    if ( gdata->vmDead ) {

        return;

    }

    tinfo = findTraceInfo(gdata->jvmti, thread, TRACE_USER);

    tagObjectWithTraceInfo(gdata->jvmti, o, tinfo);

}

/* Java Native Method for newarray */

static void JNICALL

HEAP_TRACKER_native_newarr(JNIEnv *env, jclass klass, jthread thread, jobject a)

{

    TraceInfo *tinfo;

    if ( gdata->vmDead ) {

        return;

    }

    tinfo = findTraceInfo(gdata->jvmti, thread, TRACE_USER);

    tagObjectWithTraceInfo(gdata->jvmti, a, tinfo);

}

/* Callback for JVMTI_EVENT_VM_START */

static void JNICALL

cbVMStart(jvmtiEnv *jvmti, JNIEnv *env)

{

    enterCriticalSection(jvmti); {

        jclass klass;

        jfieldID field;

        jint rc;

        /* Java Native Methods for class */

        static JNINativeMethod registry[2] = {

            {STRING(HEAP_TRACKER_native_newobj), "(Ljava/lang/Object;Ljava/lang/Object;)V",

                (void*)&HEAP_TRACKER_native_newobj},

            {STRING(HEAP_TRACKER_native_newarr), "(Ljava/lang/Object;Ljava/lang/Object;)V",

                (void*)&HEAP_TRACKER_native_newarr}

        };

        /* Register Natives for class whose methods we use */

        klass = (*env)->FindClass(env, STRING(HEAP_TRACKER_class));

        if ( klass == NULL ) {

            fatal_error("ERROR: JNI: Cannot find %s with FindClass\n",

                        STRING(HEAP_TRACKER_class));

        }

        rc = (*env)->RegisterNatives(env, klass, registry, 2);

        if ( rc != 0 ) {

            fatal_error("ERROR: JNI: Cannot register natives for class %s\n",

                        STRING(HEAP_TRACKER_class));

        }

        /* Engage calls. */

        field = (*env)->GetStaticFieldID(env, klass, STRING(HEAP_TRACKER_engaged), "I");

        if ( field == NULL ) {

            fatal_error("ERROR: JNI: Cannot get field from %s\n",

                        STRING(HEAP_TRACKER_class));

        }

        (*env)->SetStaticIntField(env, klass, field, 1);

        /* Indicate VM has started */

        gdata->vmStarted = JNI_TRUE;

    } exitCriticalSection(jvmti);

}

/* Iterate Through Heap callback */

static jint JNICALL

cbObjectTagger(jlong class_tag, jlong size, jlong* tag_ptr, jint length,

               void *user_data)

{

    TraceInfo *tinfo;

    tinfo = emptyTrace(TRACE_BEFORE_VM_INIT);

    *tag_ptr = (jlong)(ptrdiff_t)(void*)tinfo;

    return JVMTI_VISIT_OBJECTS;

}

/* Callback for JVMTI_EVENT_VM_INIT */

static void JNICALL

cbVMInit(jvmtiEnv *jvmti, JNIEnv *env, jthread thread)

{

    jvmtiHeapCallbacks heapCallbacks;

    jvmtiError         error;

    /* Iterate through heap, find all untagged objects allocated before this */

    (void)memset(&heapCallbacks, 0, sizeof(heapCallbacks));

    heapCallbacks.heap_iteration_callback = &cbObjectTagger;

    error = (*jvmti)->IterateThroughHeap(jvmti, JVMTI_HEAP_FILTER_TAGGED,

                                         NULL, &heapCallbacks, NULL);

    check_jvmti_error(jvmti, error, "Cannot iterate through heap");

    enterCriticalSection(jvmti); {

        /* Indicate VM is initialized */

        gdata->vmInitialized = JNI_TRUE;

    } exitCriticalSection(jvmti);

}

/* Iterate Through Heap callback */

static jint JNICALL

cbObjectSpaceCounter(jlong class_tag, jlong size, jlong* tag_ptr, jint length,

                     void *user_data)

{

    TraceInfo *tinfo;

    tinfo = (TraceInfo*)(ptrdiff_t)(*tag_ptr);

    if ( tinfo == NULL ) {

        tinfo = emptyTrace(TRACE_MYSTERY);

        *tag_ptr = (jlong)(ptrdiff_t)(void*)tinfo;

    }

    tinfo->totalSpace += size;

    return JVMTI_VISIT_OBJECTS;

}

/* Qsort compare function */

static int

compareInfo(const void *p1, const void *p2)

{

    TraceInfo *tinfo1, *tinfo2;

    tinfo1 = *((TraceInfo**)p1);

    tinfo2 = *((TraceInfo**)p2);

    return (int)(tinfo2->totalSpace - tinfo1->totalSpace);

}

/* Frame to text */

static void

frameToString(jvmtiEnv *jvmti, char *buf, int buflen, jvmtiFrameInfo *finfo)

{

    jvmtiError           error;

    jclass               klass;

    char                *signature;

    char                *methodname;

    char                *methodsig;

    jboolean             isNative;

    char                *filename;

    int                  lineCount;

    jvmtiLineNumberEntry*lineTable;

    int                  lineNumber;

    /* Initialize defaults */

    buf[0]     = 0;

    klass      = NULL;

    signature  = NULL;

    methodname = NULL;

    methodsig  = NULL;

    isNative   = JNI_FALSE;

    filename   = NULL;

    lineCount  = 0;

    lineTable  = NULL;

    lineNumber = 0;

    /* Get jclass object for the jmethodID */

    error = (*jvmti)->GetMethodDeclaringClass(jvmti, finfo->method, &klass);

    check_jvmti_error(jvmti, error, "Cannot get method's class");

    /* Get the class signature */

    error = (*jvmti)->GetClassSignature(jvmti, klass, &signature, NULL);

    check_jvmti_error(jvmti, error, "Cannot get class signature");

    /* Skip all this if it's our own Tracker method */

    if ( strcmp(signature, "L" STRING(HEAP_TRACKER_class) ";" ) == 0 ) {

        deallocate(jvmti, signature);

        return;

    }

    /* Get the name and signature for the method */

    error = (*jvmti)->GetMethodName(jvmti, finfo->method,

                &methodname, &methodsig, NULL);

    check_jvmti_error(jvmti, error, "Cannot method name");

    /* Check to see if it's a native method, which means no lineNumber */