code: /**
   * Copyright 2014 Alexey Ragozin
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.gridkit.jvmtool.stacktrace;
  
  import java.io.IOException;
  import java.io.Serializable;
  import java.lang.management.ThreadInfo;
  import java.lang.management.ThreadMXBean;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.List;
  import java.util.concurrent.TimeUnit;
  import java.util.regex.Matcher;
  import java.util.regex.Pattern;
  
  import javax.management.ObjectName;
  
  /**
   * Thread stack sampler.
   *
   * @author Alexey Ragozin (alexey.ragozin@gmail.com)
   */
  public class ThreadDumpSampler {
  
      @SuppressWarnings("restriction")
      public static ThreadMXBean adaptThreadMXBean(ThreadMXBean bean) {
          try {
              if (bean instanceof com.sun.management.ThreadMXBean) {
                  final com.sun.management.ThreadMXBean beanX = (com.sun.management.ThreadMXBean) bean;
                  bean = new ThreadMXBeanEx() {
  
                      @SuppressWarnings("unused")
                      /* Method added in Java 7, required for compilation */
                      public ObjectName getObjectName() {
                          return ThreadMXBeanEx.THREADING_MBEAN;
                      }
  
                      public int getPeakThreadCount() {
                          return beanX.getPeakThreadCount();
                      }
  
                      public int getDaemonThreadCount() {
                          return beanX.getDaemonThreadCount();
                      }
  
                      public long[] getAllThreadIds() {
                          return beanX.getAllThreadIds();
                      }
  
                      public ThreadInfo[] getThreadInfo(long[] ids) {
                          return beanX.getThreadInfo(ids);
                      }
  
                      public long getCurrentThreadCpuTime() {
                          return beanX.getCurrentThreadCpuTime();
                      }
  
                      public long getCurrentThreadUserTime() {
                          return beanX.getCurrentThreadUserTime();
                      }
  
                      public long getThreadUserTime(long id) {
                          return beanX.getThreadUserTime(id);
                      }
  
                      public long[] findMonitorDeadlockedThreads() {
                          return beanX.findMonitorDeadlockedThreads();
                      }
  
                      public long[] findDeadlockedThreads() {
                          return beanX.findDeadlockedThreads();
                      }
  
                      public ThreadInfo[] dumpAllThreads(boolean lockedMonitors, boolean lockedSynchronizers) {
                          return beanX.dumpAllThreads(lockedMonitors, lockedSynchronizers);
                      }
  
                      public long[] getThreadAllocatedBytes(long[] arg0) {
                          return beanX.getThreadAllocatedBytes(arg0);
                      }
  
                      public int getThreadCount() {
                          return beanX.getThreadCount();
                      }
  
                      public long getThreadCpuTime(long id) {
                          return beanX.getThreadCpuTime(id);
                      }
  
                      public long[] getThreadCpuTime(long[] arg0) {
                          return beanX.getThreadCpuTime(arg0);
                      }
  
                      public ThreadInfo getThreadInfo(long id) {
                          return beanX.getThreadInfo(id);
                      }
  
                      public ThreadInfo getThreadInfo(long id, int maxDepth) {
                          return beanX.getThreadInfo(id, maxDepth);
                      }
  
                      public ThreadInfo[] getThreadInfo(long[] ids, int maxDepth) {
                          return beanX.getThreadInfo(ids, maxDepth);
                      }
  
                      public ThreadInfo[] getThreadInfo(long[] ids, boolean lockedMonitors, boolean lockedSynchronizers) {
                          return beanX.getThreadInfo(ids, lockedMonitors, lockedSynchronizers);
                      }
  
                      public long[] getThreadUserTime(long[] arg0) {
                          return beanX.getThreadUserTime(arg0);
                      }
  
                      public long getTotalStartedThreadCount() {
                          return beanX.getTotalStartedThreadCount();
                      }
  
                      public boolean isCurrentThreadCpuTimeSupported() {
                          return beanX.isCurrentThreadCpuTimeSupported();
                      }
  
                      public boolean isObjectMonitorUsageSupported() {
                          return beanX.isObjectMonitorUsageSupported();
                      }
  
                      public boolean isSynchronizerUsageSupported() {
                          return beanX.isSynchronizerUsageSupported();
                      }
  
                      public boolean isThreadContentionMonitoringSupported() {
                          return beanX.isThreadContentionMonitoringSupported();
                      }
  
                      public boolean isThreadContentionMonitoringEnabled() {
                          return beanX.isThreadContentionMonitoringEnabled();
                      }
  
                      public boolean isThreadCpuTimeSupported() {
                          return beanX.isThreadCpuTimeSupported();
                      }
  
                      public boolean isThreadCpuTimeEnabled() {
                          return beanX.isThreadCpuTimeEnabled();
                      }
  
                      public void resetPeakThreadCount() {
                          beanX.resetPeakThreadCount();
                      }
  
                      public void setThreadContentionMonitoringEnabled(boolean enable) {
                          beanX.setThreadContentionMonitoringEnabled(enable);
                      }
  
                      public void setThreadCpuTimeEnabled(boolean enable) {
                          beanX.setThreadCpuTimeEnabled(enable);
                      }
                  };
              }
          }
          catch(NoClassDefFoundError e) {
              // ignore
          }
          catch(Exception e) {
              // ignore
          }
          return bean;
      }
  
      boolean collectTraces = true;
      boolean collectCpu = true;
      boolean collectUserCpu = true;
      boolean collectAllocation = true;
      private ThreadMXBean threading;
  
      private ThreadNameFilter threadFilter;
      private long[] threadSet;
      private List<CounterCollector> collectors = new ArrayList<CounterCollector>();
  
      public ThreadDumpSampler() {
      }
  
      public void setThreadFilter(final String pattern) {
          this.threadFilter = new ThreadNameFilter() {
  
              final Matcher matcher = Pattern.compile(pattern).matcher("");
  
              @Override
              public boolean accept(String threadName) {
                  if (threadName != null) {
                      matcher.reset(threadName);
                      return matcher.matches();
                  }
                  else {
                      return false;
                  }
              }
          };
      }
  
      public void setThreadFilter(final ThreadNameFilter filter) {
          this.threadFilter = filter;
      }
  
      public void enableThreadStackTrace(boolean enable) {
          collectTraces = enable;
      }
  
      public void enableThreadCpu(boolean enable) {
          collectCpu = true;
      }
  
      public void enableThreadUserCpu(boolean enable) {
          collectUserCpu = true;
      }
  
      public void enableThreadAllocation(boolean enable) {
          collectAllocation = true;
      }
  
      public void connect(ThreadMXBean threadingMBean) {
          this.threading = adaptThreadMXBean(threadingMBean);
          collectors.clear();
          if (collectCpu) {
              collectors.add(new GenericMBeanThreadCounter(threading, CounterType.CPU_TIME));
          }
          if (collectUserCpu) {
              collectors.add(new GenericMBeanThreadCounter(threading, CounterType.USER_TIME));
          }
          if (collectAllocation) {
              collectors.add(new GenericMBeanThreadCounter(threading, CounterType.ALLOCATED_BYTES));
          }
      }
  
      /**
       * Find threads according to thread name filters and memorise their IDs.
       * <br/>
       * Optional method to avoid dump all threads at every collection.
       */
      public void prime() {
          ThreadInfo[] ti = threading.dumpAllThreads(false, false);
          long[] tids = new long[ti.length];
          int n = 0;
          for(ThreadInfo t:ti) {
              long tid = t.getThreadId();
              String name = t.getThreadName();
              if (threadFilter == null || threadFilter.accept(name)) {
                  tids[n++] = tid;
              }
          }
          tids = Arrays.copyOf(tids, n);
          threadSet = tids;
      }
  
      public void collect(StackTraceWriter writer) throws IOException {
          long timestamp = System.currentTimeMillis();
          ThreadInfo[] dump;
          if (threadSet != null) {
              if (collectTraces) {
                  dump = compactThreads(threading.getThreadInfo(threadSet, Integer.MAX_VALUE));
              }
              else {
                  dump = compactThreads(threading.getThreadInfo(threadSet));
              }
          }
          else {
              if (collectTraces) {
                  dump = filterThreads(threading.dumpAllThreads(false, false));
              }
              else {
                  dump = filterThreads(threading.getThreadInfo(threading.getAllThreadIds()));
              }
  
          }
          long[] ids = new long[dump.length];
          for(int i = 0; i != dump.length; ++i) {
              ids[i] = dump[i].getThreadId();
          }
          for(CounterCollector cc: collectors) {
              try {
                  cc.collect(ids);
              }
              catch(Exception e) {
                  // ignore
              }
          }
          ThreadCapture ts = new ThreadCapture();
          for(ThreadInfo ti: dump) {
              ts.reset();
              ts.timestamp = timestamp;
              ts.copyFrom(ti);
              for (CounterCollector cc: collectors) {
                  try {
                      cc.fillIntoSnapshot(ts);
                  }
                  catch(Exception e) {
                      // ignore
                  }
              }
              writer.write(ts);
          }
      }
  
      private ThreadInfo[] compactThreads(ThreadInfo[] dumpAllThreads) {
          int n = 0;
          for(int i = 0; i != dumpAllThreads.length; ++i) {
              if (dumpAllThreads[i] != null) {
                  ++n;
              }
          }
          if (n == dumpAllThreads.length) {
              return dumpAllThreads;
          }
          else {
              ThreadInfo[] result = new ThreadInfo[n];
              n = 0;
              for(ThreadInfo ti: dumpAllThreads) {
                  if (ti != null) {
                      result[n++] = ti;
                  }
              }
              return result;
          }
      }
  
      private ThreadInfo[] filterThreads(ThreadInfo[] dumpAllThreads) {
          if (threadFilter == null) {
              return compactThreads(dumpAllThreads);
          }
          else {
              int n = 0;
              for(int i = 0; i != dumpAllThreads.length; ++i) {
                  if (dumpAllThreads[i] != null) {
                      if (threadFilter.accept(dumpAllThreads[i].getThreadName())) {
                          ++n;
                      }
                      else {
                          dumpAllThreads[i] = null;
                      }
                  }
              }
              if (n == dumpAllThreads.length) {
                  return dumpAllThreads;
              }
              else {
                  ThreadInfo[] result = new ThreadInfo[n];
                  n = 0;
                  for(ThreadInfo ti: dumpAllThreads) {
                      if (ti != null) {
                          result[n++] = ti;
                      }
                  }
                  return result;
              }
          }
      }
  
      @SuppressWarnings("serial")
      public static class Trace implements Serializable {
  
          private long threadId;
          private long timestamp;
          private int[] trace;
  
          private StackTraceElement[] traceDictionary;
  
          public Trace(long threadId, long timestamp, int[] trace) {
              this.threadId = threadId;
              this.timestamp = timestamp;
              this.trace = trace;
          }
  
          public long getTimestamp() {
              return timestamp;
          }
  
          public long getThreadId() {
              return threadId;
          }
  
          public StackTraceElement[] getTrace() {
              StackTraceElement[] strace = new StackTraceElement[trace.length];
              for(int i = 0; i != strace.length; ++i) {
                  strace[i] = traceDictionary[trace[i]];
              }
              return strace;
          }
  
          public void copyToSnapshot(ThreadCapture snap) {
              snap.reset();
              snap.threadId = threadId;
              snap.timestamp = timestamp;
              snap.elements = getTrace();
          }
      }
  
      static interface CounterCollector {
  
          public void collect(long[] threadID);
  
          public void fillIntoSnapshot(ThreadCapture snap);
  
      }
  
      private static enum CounterType {
          CPU_TIME,
          USER_TIME,
          ALLOCATED_BYTES
      }
  
      private static class GenericMBeanThreadCounter implements CounterCollector {
  
          private ThreadMXBean slowBean;
          private ThreadMXBeanEx fastBean;
          private CounterType counter;
  
          private long[] threads;
          private long[] counters;
          private int n = 0;
  
          public GenericMBeanThreadCounter(ThreadMXBean bean, CounterType counter) {
              this.slowBean = bean;
              try {
                  if (bean instanceof ThreadMXBeanEx) {
                      fastBean = (ThreadMXBeanEx) bean;
                  }
              }
              catch(NoClassDefFoundError e) {
                  // ignore
              }
              catch(Exception e) {
                  // ignore
              }
              this.counter = counter;
          }
  
          @Override
          public void collect(long[] threadID) {
              n = 0;
              threads = threadID;
              if (fastBean != null) {
                  try {
                      counters = callFast(threads);
                  }
                  catch(Exception e) {
                      // fallback to slow mode
                      fastBean = null;
                      collect(threadID);
                  }
              }
              else {
                  counters = new long[threads.length];
                  for(int i = 0; i != threads.length; ++i) {
                      counters[i] = callSlow(threads[i]);
                  }
              }
          }
  
          private long[] callFast(long[] threads) {
              switch(counter) {
                  case CPU_TIME: return fastBean.getThreadCpuTime(threads);
                  case USER_TIME: return fastBean.getThreadUserTime(threads);
                  case ALLOCATED_BYTES: return fastBean.getThreadAllocatedBytes(threads);
                  default: throw new RuntimeException("Unknown counter: " + counter);
              }
          }
  
          private long callSlow(long threadId) {
              switch (counter) {
                  case CPU_TIME: return slowBean.getThreadCpuTime(threadId);
                  case USER_TIME: return slowBean.getThreadUserTime(threadId);
                  case ALLOCATED_BYTES: return -1;
                  default: throw new RuntimeException("Unknown counter: " + counter);
              }
          }
  
          @Override
          public void fillIntoSnapshot(ThreadCapture snap) {
              int n = indexOf(snap.threadId);
              String counterKey;
              switch(counter) {
                  case CPU_TIME: counterKey = ThreadCounters.CPU_TIME_MS; break;
                  case USER_TIME: counterKey = ThreadCounters.USER_TIME_MS; break;
                  case ALLOCATED_BYTES: counterKey = ThreadCounters.ALLOCATED_BYTES; break;
                  default: throw new RuntimeException("Unknown counter: " + counter);
              }
              long v = counters[n];
              if (v >= 0 && counter != CounterType.ALLOCATED_BYTES) {
                  v = TimeUnit.NANOSECONDS.toMillis(v);
              }
              snap.counters.set(counterKey, v);
          }
  
          private int indexOf(long threadId) {
              if (threads[n] == threadId) {
                  return n;
              }
              n++;
              if (threads[n] == threadId) {
                  return n;
              }
              for(int i = 0; i != threads.length; ++i) {
                  if (threads[i] == threadId) {
                      return n = i;
                  }
              }
              return -1;
          }
      }
  }
  