/*!
 * \file snapShotMain.cpp
 * \brief This file is used to take snapshot.
 * Copyright (C) 2011-2015 Nippon Telegraph and Telephone Corporation
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 *
 */

#include <jni.h>
#include <jvmti.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>

#include <execinfo.h>
#include <pthread.h>

#include "snapShotMain.hpp"
#include "libmain.hpp"
#include "classContainer.hpp"
#include "snapShotContainer.hpp"
#include "snapShotProcessor.hpp"
#include "gcWatcher.hpp"
#include "timer.hpp"
#include "util.hpp"
#include "oopUtil.hpp"
#include "elapsedTimer.hpp"
#include "jvmInfo.hpp"
#include "bitMapMarker.hpp"

/*!
 * \brief ClassData Container.
 */
TClassContainer *clsContainer;
/*!
 * \brief SnapShot Processor.
 */
TSnapShotProcessor *snapShotProcessor;
/*!
 * \brief GC Watcher.
 */
TGCWatcher *gcWatcher;
/*!
 * \brief Timer Thread.
 */
TTimer *timer;

/*!
 * \brief Pthread mutex for user data dump request.<br>
 *        E.g. continue pushing dump key.<br>
 * <br>
 * This mutex used in below process.<br>
 *   - OnDataDumpRequest @ snapShotMain.cpp<br>
 *     To avoid a lot of data dump request parallel processing.<br>
 */
pthread_mutex_t dumpMutex = PTHREAD_ADAPTIVE_MUTEX_INITIALIZER_NP;

/*!
 * \brief Pthread mutex for change snapshot queue.<br>
 * <br>
 * This mutex used in below process.<br>
 *   - TakeSnapShot @ snapShotMain.cpp<br>
 *     To add snapshot by JVMTI to output wait queue.<br>
 *   - TakeSnapShot @ snapShotMain.cpp<br>
 *     To get snapshot from output wait queue.<br>
 *   - outputSnapShotByGC @ snapShotMain.cpp<br>
 *     To add snapshot by GC to output wait queue.<br>
 */
pthread_mutex_t queueMutex = PTHREAD_ADAPTIVE_MUTEX_INITIALIZER_NP;

/*!
 * \brief Pthread mutex for JVMTI IterateOverHeap calling.<br>
 * <br>
 * This mutex used in below process.<br>
 *   - TakeSnapShot @ snapShotMain.cpp<br>
 *     To avoid taking double snapshot use JVMTI.<br>
 *     Because agent cann't distinguish called from where in hook function.<br>
 */
pthread_mutex_t jvmtiMutex = PTHREAD_ADAPTIVE_MUTEX_INITIALIZER_NP;

/*!
 * \brief Snapshot container instance queue to waiting output.
 */
TSnapShotQueue snapStockQueue;
/*!
 * \brief Container instance stored got snapshot by GC.
 */
TSnapShotContainer *snapshotByGC    = NULL;
/*!
 * \brief Container instance stored got snapshot by CMSGC.
 */
TSnapShotContainer *snapshotByCMS   = NULL;
/*!
 * \brief Container instance stored got snapshot by interval or dump request.
 */
TSnapShotContainer *snapshotByJvmti = NULL;

/*!
 * \brief Count object size in Heap.
 * \param clsTag   [in]     Tag of object's class.
 * \param size     [in]     Object size.
 * \param objTag   [in,out] Object's tag.
 * \param userData [in,out] User's data.
 */
jvmtiIterationControl JNICALL HeapObjectCallBack(jlong clsTag,
                              jlong size, jlong *objTag, void *userData) {
  
  /* This callback is dummy. */
  return JVMTI_ITERATION_ABORT;
}

/*!
 * \brief New class loaded event.
 * \param jvmti  [in] JVMTI environment object.
 * \param env    [in] JNI environment object.
 * \param thread [in] Java thread object.
 * \param klass  [in] Newly loaded class object.
 */
void JNICALL OnClassPrepare(jvmtiEnv *jvmti, JNIEnv *env,
                                     jthread thread, jclass klass) {
  
  /* Get klassOop. */
  void *mirror   = *(void**)klass;
  void *klassOop = asKlassOop(mirror);
  
  if (likely(klassOop != NULL)) {
    /* Push new loaded class. */
    clsContainer->pushNewClass(klassOop);
  }
}

/*!
 * \brief Class unload event.
 * \param jvmti  [in] JVMTI environment object.
 * \param env    [in] JNI environment object.
 * \param thread [in] Java thread object.
 * \param klass  [in] Unload class object.
 * \sa    from: hotspot/src/share/vm/prims/jvmti.xml
 */
void JNICALL OnClassUnload(jvmtiEnv *jvmti, JNIEnv *env,
                                      jthread thread, jclass klass) {
  
  /* Get klassOop. */
  void *mirror   = *(void**)klass;
  void *klassOop = asKlassOop(mirror);
  
  if (likely(klassOop != NULL)) {
    /* Search class. */
    TObjectData *counter = clsContainer->findClass(klassOop);
      
    if (likely(counter != NULL)) {
      /* Push to remove pending queue. */
      clsContainer->pushToPendingQueue(counter);
    }

  }
}

/*!
 * \brief Setting JVM information to snapshot.
 * \param snapshot [in] Snapshot instance.
 * \param cause    [in] Cause of taking a snapshot.<br>
 *                      e.g. GC, DumpRequest or Interval.
 */
inline void setSnapShotInfo(TInvokeCause cause, TSnapShotContainer *snapshot) {
  /* Get now date and time. */
  struct timeval tv;
  gettimeofday(&tv, NULL);
  
  /* Set snapshot information. */
  snapshot->setSnapShotTime((jlong)tv.tv_sec * 1000 + (jlong)tv.tv_usec / 1000);
  snapshot->setSnapShotCause(cause);
  snapshot->setJvmInfo(jvmInfo);
}

/*!
 * \brief Set information and push waiting queue.
 * \param snapshot [in] Snapshot instance
 */
inline void outputSnapShotByGC(TSnapShotContainer *snapshot) {
  
  setSnapShotInfo(GC, snapshot);
  
  /* Standby for next GC. */
  jvmInfo->resumeGCinfo();
  
  /* Push output waiting queue. */
  ENTER_PTHREAD_SECTION(&queueMutex) {
    snapStockQueue.push(snapshot);
  } EXIT_PTHREAD_SECTION(&queueMutex)
  
  /* Send notification. */
  gcWatcher->notify();
}

/*!
 * \brief Interrupt inner garbage collection event.
 *        Call this function when JVM invoke many times of GC process
 *        during single JVMTI event of GC start and GC finish.
 */
void onInnerGarbageCollectionInterrupt(void) {

  /* Standby for next GC. */
  jvmInfo->resumeGCinfo();

  /* Clear unfinished snapshot data. */
  snapshotByGC->clear(false);
}


/*!
 * \brief Before garbage collection event.
 * \param jvmti [in] JVMTI environment object.
 */
void JNICALL OnGarbageCollectionStart(jvmtiEnv *jvmti) {
  snapshotByGC = TSnapShotContainer::getInstance();

  /* Enable inner GC event. */
  setupHookForInnerGCEvent(true, &onInnerGarbageCollectionInterrupt);
}

/*!
 * \brief After garbage collection event.
 * \param jvmti [in] JVMTI environment object.
 */
void JNICALL OnGarbageCollectionFinish(jvmtiEnv *jvmti) {
  /* Disable inner GC event. */
  setupHookForInnerGCEvent(false, NULL);

  /* If need getting snapshot. */
  if(gcWatcher->needToStartGCTrigger()){
    /* Set information and push waiting queue. */
    outputSnapShotByGC(snapshotByGC);
  }
  else{
    TSnapShotContainer::releaseInstance(snapshotByGC);
  }

  snapshotByGC = NULL;
}

/*!
 * \brief After G1 garbage collection event.
 */
void OnG1GarbageCollectionFinish(void) {
  //jvmInfo->loadGCCause();
  jvmInfo->SetUnknownGCCause();

  /* Set information and push waiting queue. */
  outputSnapShotByGC(snapshotByGC);
  snapshotByGC = TSnapShotContainer::getInstance();
}

 /*!
 * \brief Count object size in Heap.
 * \param snapshot [in] Snapshot instance.
 * \param klassOop [in] Pointer of java class object(KlassOopDesc).
 * \param size     [in] Size of object instance.
 */
void calculateObjectUsage(TSnapShotContainer *snapshot,
                                       void *klassOop, jlong size){
  
  /* Sanity check. */
  if (unlikely(snapshot == NULL)) {
    return;
  }

  snapshot->setIsCleared(false);
  
  TObjectCounter *counter = NULL;
  TObjectData *objData = NULL;

  TClassContainer *my_classContainer = clsContainer->getTLSContainer();
  
  /* Search class. */
  objData = my_classContainer->findClass(klassOop);

  if(objData == NULL){
    /* Search class from rootset. */
    objData = clsContainer->findClass(klassOop);
    
    if (unlikely(objData == NULL)) {
      /* Push new loaded class. */
      objData = clsContainer->pushNewClass(klassOop);
    }

  }

  TSnapShotContainer *my_container = NULL;

  if (likely(objData != NULL)) {

    my_container = snapshot->getTLSContainer();

    /* Search class. */
    counter = my_container->findClass(objData);
      
    if (unlikely(counter == NULL)) {
      /* Push new loaded class. */
      counter = my_container->pushNewClass(objData);
    }

    /* Count class size and instance count. */
    if(likely(counter != NULL)){
      my_container->FastInc(counter, size);
    }

  }

}

/*!
 * \brief Count object size in Heap.
 * \param klassOop [in] Pointer of java class object(KlassOopDesc).
 * \param size     [in] Size of object instance.
 */
void HeapObjectCallbackOnGC(void *klassOop, jlong size) {
  calculateObjectUsage(snapshotByGC, klassOop, size);
}

/*!
 * \brief Count object size in Heap by CMSGC.
 * \param klassOop [in] Pointer of java class object(KlassOopDesc).
 * \param size     [in] Size of object instance.
 */
void HeapObjectCallbackOnCMS(void *klassOop, jlong size) {
  calculateObjectUsage(snapshotByCMS, klassOop, size);
}

/*!
 * \brief Count object size in Heap by JVMTI iterateOverHeap.
 * \param klassOop [in] Pointer of java class object(KlassOopDesc).
 * \param size     [in] Size of object instance.
 */
void HeapObjectCallbackOnJvmti(void *klassOop, jlong size) {
  calculateObjectUsage(snapshotByJvmti, klassOop, size);
}

/*!
 * \brief This function is for class oop adjust callback by GC.
 * \param oldOop [in] Old pointer of java class object(KlassOopDesc).
 * \param newOop [in] New pointer of java class object(KlassOopDesc).
 */
void HeapKlassAdjustCallback(void *oldOop, void *newOop) {
  /* Class information update. */
  clsContainer->updateClass(oldOop, newOop);
}

/*!
 * \brief Event of before garbage collection by CMS collector.
 * \param jvmti [in] JVMTI environment object.
 */
void JNICALL OnCMSGCStart(jvmtiEnv *jvmti) {
  
  /* Get CMS state. */
  bool needShapShot = false;
  int cmsState = checkCMSState(gcStart, &needShapShot);
  
  /* If occurred snapshot target GC. */
  if(needShapShot){

    /* If need getting snapshot. */
    if(gcWatcher->needToStartGCTrigger()){
      /* Set information and push waiting queue. */
      outputSnapShotByGC(snapshotByCMS);
      snapshotByCMS = NULL;
    }
  }
  
  if(likely(snapshotByGC == NULL)){
    snapshotByGC = TSnapShotContainer::getInstance();
  }
  else{
    snapshotByGC->clear(false);
  }

  if(likely(snapshotByCMS == NULL)){
    snapshotByCMS = TSnapShotContainer::getInstance();
  }
  else if(cmsState == CMS_FINALMARKING){
    snapshotByCMS->clear(false);
  }

  /* Enable inner GC event. */
  setupHookForInnerGCEvent(true, &onInnerGarbageCollectionInterrupt);
}

/*!
 * \brief Event of after garbage collection by CMS collector.
 * \param jvmti [in] JVMTI environment object.
 */
void JNICALL OnCMSGCFinish(jvmtiEnv *jvmti) {
  /* Disable inner GC event. */
  setupHookForInnerGCEvent(false, NULL);
  
  /* Get CMS state. */
  bool needShapShot = false;
  checkCMSState(gcFinish, &needShapShot);
  
  /* If occurred snapshot target GC. */
  if(needShapShot){
    
    /* If need getting snapshot. */
    if(gcWatcher->needToStartGCTrigger()){
      /* Set information and push waiting queue. */
      outputSnapShotByGC(snapshotByGC);
      snapshotByGC = NULL;
      snapshotByCMS->clear(false);
    }
  }

}

/*!
 * \brief Data dump request event.
 * \param jvmti [in] JVMTI environment object.
 */
void JNICALL OnDataDumpRequest(jvmtiEnv *jvmti) {
  
  /* Avoid the plural simultaneous take snapshot by dump-request.        */
  /* E.g. keeping pushed dump key.                                       */
  /* Because classContainer register a redundancy class in TakeSnapShot. */
  ENTER_PTHREAD_SECTION(&dumpMutex) {
    
    /* Make snapshot. */
    TakeSnapShot(jvmti, NULL, DataDumpRequest);
    
  } EXIT_PTHREAD_SECTION(&dumpMutex)
}

/*!
 * \brief Take a heap information snapshot.
 * \param jvmti [in] JVMTI environment object.
 * \param env   [in] JNI environment object.
 * \param cause [in] Cause of taking a snapshot.<br>
 *                   e.g. GC, DumpRequest or Interval.
 */
void TakeSnapShot(jvmtiEnv *jvmti, JNIEnv *env, TInvokeCause cause) {

  if(useCMS && (*CMS_collectorState > CMS_IDLING)){
    PRINT_WARN_MSG("CMS GC is working. Skip to take a SnapShot.");
    return;
  }

  /* Count working time. */
  static const char *label = "Take SnapShot";
  TElapsedTimer elapsedTime(label);
  
  /* Phase1: Heap Walking. */
  jvmtiError error = JVMTI_ERROR_INTERNAL;
  if (cause != GC) {
    TSnapShotContainer *snapshot = TSnapShotContainer::getInstance();
    
    if (likely(snapshot != NULL)) {
      
      /* Lock to avoid doubling call JVMTI. */
      ENTER_PTHREAD_SECTION(&jvmtiMutex) {
        snapshotByJvmti = snapshot;
        
        /* Enable JVMTI hooking. */
        if (likely(setJvmtiHookState(true))) {
          
          /* Count object size on heap. */
          error = jvmti->IterateOverHeap(JVMTI_HEAP_OBJECT_EITHER,
            &HeapObjectCallBack, NULL);
          
          /* Disable JVMTI hooking. */
          setJvmtiHookState(false);
        }
        
        snapshotByJvmti = NULL;
      } EXIT_PTHREAD_SECTION(&jvmtiMutex)
    }
    
    if (likely(error == JVMTI_ERROR_NONE)) {
      
      setSnapShotInfo(cause, snapshot);
      
      /* Push waiting queue. */
      ENTER_PTHREAD_SECTION(&queueMutex) {
        snapStockQueue.push(snapshot);
      } EXIT_PTHREAD_SECTION(&queueMutex)

    } else {
      /* Process failure. So data is junk. */
      TSnapShotContainer::releaseInstance(snapshot);
    }
  } else {
    /* Failure hook JVMTI function. */
    error = JVMTI_ERROR_NONE;
  }
  
  /* If failure count object size. */
  if (unlikely(isError(jvmti, error))) {
    PRINT_WARN_MSG("Heap SnapShot failed!");
  } else {
    TSnapShotContainer *snapshot = NULL;
    
    /* Get snapshot data form waiting queue. */
    ENTER_PTHREAD_SECTION(&queueMutex) {
      snapshot = snapStockQueue.front();
      snapStockQueue.pop();
    } EXIT_PTHREAD_SECTION(&queueMutex)

    /* Set total memory */
    snapshot->setTotalSize(jvmInfo->getTotalMemory());
    
    try {
      /* Sending notification means able to output file. */
      snapShotProcessor->notify(snapshot);
    } catch(...) {
      PRINT_WARN_MSG("Snapshot processeor notify failed!.");
    }
  }
  
  /* Phase2: Reset Timer. */
  if (arg.TimerInterval > 0) {
    /* Sending notification means reset timer. */
    timer->notify();
  }
}

/*!
 * \brief Setting JVMTI functional capabilities.
 * \param capabilities [out] Event capabilities structure.
 */
void setCapabilitiesForSnapShot(jvmtiCapabilities *capabilities) {
  /* Enable for take snapshop. */
  capabilities->can_generate_garbage_collection_events = 1;
  capabilities->can_tag_objects = 1;
}

/*!
 * \brief Setting enable of JVMTI and extension events for snapshot function.
 * \param jvmti  [in] JVMTI environment object.
 * \param enable [in] Event notification is enable.
 * \return Setting process result.
 */
jint setEventEnableForSnapShot(jvmtiEnv *jvmti, bool enable) {
  jvmtiEventMode mode = JVMTI_ENABLE;
  if (unlikely(!enable)) {
    mode = JVMTI_DISABLE;
  }
  
  /* Switch date dump event. */
  if (arg.triggerOnDump) {
    SWITCH_EVENT_NOTIFICATION(jvmti, JVMTI_EVENT_DATA_DUMP_REQUEST, mode)
  }
  
  /* Switch gc events. */
  if (arg.triggerOnFullGC) {
    
    /* Switch JVMTI GC events. */
    SWITCH_EVENT_NOTIFICATION(jvmti, JVMTI_EVENT_GARBAGE_COLLECTION_START,
      mode)
    SWITCH_EVENT_NOTIFICATION(jvmti, JVMTI_EVENT_GARBAGE_COLLECTION_FINISH,
      mode)
  }
  
  return SUCCESS;
}

/*!
 * \brief Setting enable of agent each threads for snapshot function.
 * \param jvmti  [in] JVMTI environment object.
 * \param env    [in] JNI environment object.
 * \param enable [in] Event notification is enable.
 */
void setThreadEnableForSnapShot(jvmtiEnv *jvmti, JNIEnv *env, bool enable) {
  
  /* Start or suspend HeapStats agent threads. */
  try{
    
    /* Switch GC watcher state. */
    if (arg.triggerOnFullGC) {
      if (enable) {
        gcWatcher->start(jvmti, env);
      } else {
        gcWatcher->stop();
      }
      
      /* Switch GC hooking state. */
      setGCHookState(enable);
    }
    
    /* Switch interval snapshot timer state. */
    if (arg.TimerInterval > 0) {
      if (enable) {
        timer->start(jvmti, env);
      } else {
        timer->stop();
      }
    }
    
    /* Switch snapshot processor state. */
    if (enable) {
      snapShotProcessor->start(jvmti, env);
    } else {
      snapShotProcessor->stop();
    }
  } catch (const char *errMsg) {
    PRINT_WARN_MSG(errMsg);
  }

}

/*!
 * \brief JVM initialization event for snapshot function.
 * \param jvmti  [in] JVMTI environment object.
 * \param env    [in] JNI environment object.
 */
void onVMInitForSnapShot(jvmtiEnv *jvmti, JNIEnv *env) {
  size_t maxMemSize = jvmInfo->getMaxMemory();
  /* Setup for hooking. */
  setupHook(&HeapObjectCallbackOnGC, &HeapObjectCallbackOnCMS,
                  &HeapObjectCallbackOnJvmti, &HeapKlassAdjustCallback,
                                       &OnG1GarbageCollectionFinish, maxMemSize);
  
  /* JVMTI Extension Event Setup. */
  int eventIdx = GetClassUnloadingExtEventIndex(jvmti);
  
  /* JVMTI extension event is influenced by JVM's implementation.       */
  /* Extension event may no exist , if JVM was a little updated.        */
  /* This is JVMTI's Specifications.                                    */
  /* So result is disregard, even if failure processed extension event. */
  if (eventIdx < 0) {
    
    /* Miss get extension event list. */
    PRINT_WARN_MSG("Couldn't get ClassUnload event.");
    
  } else {
    
    /* Set onClassUnload event. */
    if (isError(jvmti, jvmti->SetExtensionEventCallback(eventIdx,
      (jvmtiExtensionEvent)&OnClassUnload))) {
      
      /* Failure setting extension event. */
      PRINT_WARN_MSG("Couldn't register ClassUnload event.");
    }
  }

}

/*!
 * \brief JVM finalization event for snapshot function.
 * \param jvmti [in] JVMTI environment object.
 * \param env   [in] JNI environment object.
 */
void onVMDeathForSnapShot(jvmtiEnv *jvmti, JNIEnv *env) {
  
  if(useCMS){
    /* Get CMS state. */
    bool needShapShot = false;
    checkCMSState(gcLast, &needShapShot);
    
    /* If occurred snapshot target GC. */
    if (needShapShot && snapshotByCMS != NULL) {
      
      /* Output snapshot. */
      outputSnapShotByGC(snapshotByCMS);
      snapshotByCMS = NULL;
    }
  }
  
  /* Terminate each thread. */
  gcWatcher->terminate();
  
  timer->terminate();
  
  snapShotProcessor->terminate();
}

/*!
 * \brief Agent initialization for snapshot function.
 * \param jvmti [in] JVMTI environment object.
 * \return Initialize process result.
 */
jint onAgentInitForSnapShot(jvmtiEnv **jvmti) {
  
  /* Initialize oop util. */
  if (unlikely(!oopUtilInitialize(*jvmti))) {
    PRINT_CRIT_MSG("Please check installation and version of java and debuginfo packages.");
    return GET_LOW_LEVEL_INFO_FAILED;
  }
  
  /* Initialize snapshot containers. */
  if (unlikely(!TSnapShotContainer::globalInitialize())) {
    PRINT_CRIT_MSG("TSnapshotContainer initialize failed!");
    return CLASSCONTAINER_INITIALIZE_FAILED;
  }
  
  /* Initialize TClassContainer. */
  try {
    clsContainer = new TClassContainer();
    
  } catch(...) {
    PRINT_CRIT_MSG("TClassContainer initialize failed!");
    return CLASSCONTAINER_INITIALIZE_FAILED;
  }
  
  /* Create thread instances that controlled snapshot trigger. */
  try {
    gcWatcher = new TGCWatcher(&TakeSnapShot, jvmInfo);
    
    snapShotProcessor = new TSnapShotProcessor(clsContainer, jvmInfo);
    
    timer = new TTimer(&TakeSnapShot, arg.TimerInterval,
                                              "HeapStats Snapshot Timer");
    
  } catch(const char *errMsg) {
    PRINT_CRIT_MSG(errMsg);
    return AGENT_THREAD_INITIALIZE_FAILED;
    
  } catch(...) {
    PRINT_CRIT_MSG("AgentThread initialize failed!");
    return AGENT_THREAD_INITIALIZE_FAILED;
  }
  
  return SUCCESS;
}

/*!
 * \brief Agent finalization for snapshot function.
 */
void onAgentFinalForSnapShot(void) {
  /* Destroy object that is for snapshot. */
  delete clsContainer;
  clsContainer = NULL;
  
  delete snapShotProcessor;
  snapShotProcessor = NULL;
  
  /* Finalize and deallocate old snapshot containers. */
  TSnapShotContainer::globalFinalize();
  
  /* Destroy object that is each snapshot trigger. */
  delete gcWatcher;
  gcWatcher = NULL;
  
  delete timer;
  timer = NULL;
  
  /* Finalize oop util. */
  oopUtilFinalize();
}