Initial commit for spark

1 files changed, 156 insertions, 0 deletions

diff --git a/common/src/main/java/com/sk89q/warmroast/Sampler.java b/common/src/main/java/com/sk89q/warmroast/Sampler.java
new file mode 100644
index 0000000..6c4f60c
--- /dev/null
+++ b/common/src/main/java/com/sk89q/warmroast/Sampler.java
@@ -0,0 +1,156 @@
+/*
+ * WarmRoast
+ * Copyright (C) 2013 Albert Pham <http://www.sk89q.com>
+ *
+ * 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 3 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, see <http://www.gnu.org/licenses/>.
+*/
+
+package com.sk89q.warmroast;
+
+import com.google.gson.JsonArray;
+import com.google.gson.JsonObject;
+
+import java.lang.management.ManagementFactory;
+import java.lang.management.ThreadInfo;
+import java.lang.management.ThreadMXBean;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Timer;
+import java.util.TimerTask;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * Main sampler class.
+ */
+public class Sampler extends TimerTask {
+
+    /**
+     * The thread management interface for the current JVM
+     */
+    private final ThreadMXBean threadBean = ManagementFactory.getThreadMXBean();
+
+    /**
+     * A map of root stack nodes for each thread with sampling data
+     */
+    private final Map<String, StackNode> threadData = new HashMap<>();
+
+    /**
+     * A future to encapsulation the completion of this sampler instance
+     */
+    private final CompletableFuture<Sampler> future = new CompletableFuture<>();
+
+    /** The interval to wait between sampling, in milliseconds */
+    private final int interval;
+    /** The instance used to generate thread information for use in sampling */
+    private final ThreadDumper threadDumper;
+    /** The time when sampling first began */
+    private long startTime = -1;
+    /** The unix timestamp (in millis) when this sampler should automatically complete.*/
+    private final long endTime; // -1 for nothing
+    
+    public Sampler(int interval, ThreadDumper threadDumper, long endTime) {
+        this.interval = interval;
+        this.threadDumper = threadDumper;
+        this.endTime = endTime;
+    }
+
+    /**
+     * Starts the sampler.
+     *
+     * @param timer the timer to schedule the sampling on
+     */
+    public synchronized void start(Timer timer) {
+        timer.scheduleAtFixedRate(this, 0, this.interval);
+        this.startTime = System.currentTimeMillis();
+    }
+
+    /**
+     * Gets the sampling data recorded by this instance.
+     *
+     * @return the data
+     */
+    public Map<String, StackNode> getData() {
+        return this.threadData;
+    }
+
+    public long getStartTime() {
+        if (this.startTime == -1) {
+            throw new IllegalStateException("Not yet started");
+        }
+        return this.startTime;
+    }
+
+    public long getEndTime() {
+        return this.endTime;
+    }
+
+    public CompletableFuture<Sampler> getFuture() {
+        return this.future;
+    }
+
+    private StackNode getRootNode(String threadName) {
+        return this.threadData.computeIfAbsent(threadName, StackNode::new);
+    }
+
+    @Override
+    public synchronized void run() {
+        try {
+            if (this.endTime != -1 && this.endTime <= System.currentTimeMillis()) {
+                this.future.complete(this);
+                cancel();
+                return;
+            }
+
+            ThreadInfo[] threadDumps = this.threadDumper.dumpThreads(this.threadBean);
+            for (ThreadInfo threadInfo : threadDumps) {
+                String threadName = threadInfo.getThreadName();
+                StackTraceElement[] stack = threadInfo.getStackTrace();
+
+                if (threadName == null || stack == null) {
+                    continue;
+                }
+
+                StackNode node = getRootNode(threadName);
+                node.log(stack, this.interval);
+            }
+        } catch (Throwable t) {
+            this.future.completeExceptionally(t);
+            cancel();
+        }
+    }
+
+    public JsonObject formOutput() {
+        JsonObject out = new JsonObject();
+
+        JsonArray threads = new JsonArray();
+
+        List<Map.Entry<String, StackNode>> data = new ArrayList<>(getData().entrySet());
+        data.sort(Map.Entry.comparingByKey());
+
+        for (Map.Entry<String, StackNode> entry : data) {
+            JsonObject o = new JsonObject();
+            o.addProperty("threadName", entry.getKey());
+            o.addProperty("totalTime", entry.getValue().getTotalTime());
+            o.add("rootNode", entry.getValue().serialize());
+
+            threads.add(o);
+        }
+        out.add("threads", threads);
+
+        return out;
+    }
+
+}