Stream live data to the viewer using WebSockets (#294)

26 files changed, 1508 insertions, 53 deletions

diff --git a/build.gradle b/build.gradle
index f1a0a8e..8984ca1 100644
--- a/build.gradle
+++ b/build.gradle
@@ -29,6 +29,12 @@ subprojects {
         options.release = 8
     }
 
+    jar {
+        manifest {
+            attributes('Multi-Release': 'true')
+        }
+    }
+
     processResources {
         duplicatesStrategy = DuplicatesStrategy.INCLUDE
     }
diff --git a/spark-common/build.gradle b/spark-common/build.gradle
index c3d960d..742d943 100644
--- a/spark-common/build.gradle
+++ b/spark-common/build.gradle
@@ -4,6 +4,24 @@ plugins {
     id 'com.google.protobuf' version '0.9.1'
 }
 
+sourceSets {
+    java11 {
+        java {
+            srcDirs = ['src/main/java11']
+        }
+    }
+}
+
+compileJava11Java {
+    options.release = 11
+}
+
+jar {
+    into('META-INF/versions/11') {
+        from sourceSets.java11.output
+    }
+}
+
 license {
     exclude '**/sampler/async/jfr/**'
 }
@@ -18,8 +36,11 @@ dependencies {
     api project(':spark-api')
     implementation 'com.github.jvm-profiling-tools:async-profiler:v2.8.3'
     implementation 'org.ow2.asm:asm:9.1'
-    implementation 'com.google.protobuf:protobuf-javalite:3.21.11'
     implementation 'net.bytebuddy:byte-buddy-agent:1.11.0'
+
+    implementation 'com.google.protobuf:protobuf-javalite:3.21.11'
+    java11Implementation 'com.google.protobuf:protobuf-javalite:3.21.11'
+
     api('net.kyori:adventure-api:4.12.0') {
         exclude(module: 'adventure-bom')
         exclude(module: 'checker-qual')
@@ -40,6 +61,8 @@ dependencies {
     compileOnly 'com.google.code.gson:gson:2.7'
     compileOnly 'com.google.guava:guava:19.0'
     compileOnly 'org.checkerframework:checker-qual:3.8.0'
+
+    java11Implementation files(sourceSets.main.output.classesDirs) { builtBy compileJava }
 }
 
 protobuf {
diff --git a/spark-common/src/main/java/me/lucko/spark/common/SparkPlatform.java b/spark-common/src/main/java/me/lucko/spark/common/SparkPlatform.java
index dae04ff..61c6062 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/SparkPlatform.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/SparkPlatform.java
@@ -53,6 +53,8 @@ import me.lucko.spark.common.tick.TickReporter;
 import me.lucko.spark.common.util.BytebinClient;
 import me.lucko.spark.common.util.Configuration;
 import me.lucko.spark.common.util.TemporaryFiles;
+import me.lucko.spark.common.util.ws.BytesocksClient;
+import me.lucko.spark.common.ws.TrustedKeyStore;
 
 import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.event.ClickEvent;
@@ -95,6 +97,8 @@ public class SparkPlatform {
     private final Configuration configuration;
     private final String viewerUrl;
     private final BytebinClient bytebinClient;
+    private final BytesocksClient bytesocksClient;
+    private final TrustedKeyStore trustedKeyStore;
     private final boolean disableResponseBroadcast;
     private final List<CommandModule> commandModules;
     private final List<Command> commands;
@@ -118,8 +122,12 @@ public class SparkPlatform {
         this.configuration = new Configuration(this.plugin.getPluginDirectory().resolve("config.json"));
 
         this.viewerUrl = this.configuration.getString("viewerUrl", "https://spark.lucko.me/");
-        String bytebinUrl = this.configuration.getString("bytebinUrl", "https://bytebin.lucko.me/");
+        String bytebinUrl = this.configuration.getString("bytebinUrl", "https://spark-usercontent.lucko.me/");
+        String bytesocksHost = this.configuration.getString("bytesocksHost", "spark-usersockets.lucko.me");
+
         this.bytebinClient = new BytebinClient(bytebinUrl, "spark-plugin");
+        this.bytesocksClient = BytesocksClient.create(bytesocksHost, "spark-plugin");
+        this.trustedKeyStore = new TrustedKeyStore(this.configuration);
 
         this.disableResponseBroadcast = this.configuration.getBoolean("disableResponseBroadcast", false);
 
@@ -228,6 +236,14 @@ public class SparkPlatform {
         return this.bytebinClient;
     }
 
+    public BytesocksClient getBytesocksClient() {
+        return this.bytesocksClient;
+    }
+
+    public TrustedKeyStore getTrustedKeyStore() {
+        return this.trustedKeyStore;
+    }
+
     public boolean shouldBroadcastResponse() {
         return !this.disableResponseBroadcast;
     }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/command/modules/HeapAnalysisModule.java b/spark-common/src/main/java/me/lucko/spark/common/command/modules/HeapAnalysisModule.java
index 5bd62a8..6ac3b2f 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/command/modules/HeapAnalysisModule.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/command/modules/HeapAnalysisModule.java
@@ -32,6 +32,7 @@ import me.lucko.spark.common.heapdump.HeapDump;
 import me.lucko.spark.common.heapdump.HeapDumpSummary;
 import me.lucko.spark.common.util.Compression;
 import me.lucko.spark.common.util.FormatUtil;
+import me.lucko.spark.common.util.MediaTypes;
 import me.lucko.spark.proto.SparkHeapProtos;
 
 import net.kyori.adventure.text.event.ClickEvent;
@@ -52,7 +53,6 @@ import static net.kyori.adventure.text.format.NamedTextColor.GREEN;
 import static net.kyori.adventure.text.format.NamedTextColor.RED;
 
 public class HeapAnalysisModule implements CommandModule {
-    private static final String SPARK_HEAP_MEDIA_TYPE = "application/x-spark-heap";
 
     @Override
     public void registerCommands(Consumer<Command> consumer) {
@@ -97,7 +97,7 @@ public class HeapAnalysisModule implements CommandModule {
             saveToFile = true;
         } else {
             try {
-                String key = platform.getBytebinClient().postContent(output, SPARK_HEAP_MEDIA_TYPE).key();
+                String key = platform.getBytebinClient().postContent(output, MediaTypes.SPARK_HEAP_MEDIA_TYPE).key();
                 String url = platform.getViewerUrl() + key;
 
                 resp.broadcastPrefixed(text("Heap dump summmary output:", GOLD));
diff --git a/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java b/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java
index 041cacf..049c817 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java
@@ -41,7 +41,10 @@ import me.lucko.spark.common.sampler.node.MergeMode;
 import me.lucko.spark.common.sampler.source.ClassSourceLookup;
 import me.lucko.spark.common.tick.TickHook;
 import me.lucko.spark.common.util.FormatUtil;
+import me.lucko.spark.common.util.MediaTypes;
 import me.lucko.spark.common.util.MethodDisambiguator;
+import me.lucko.spark.common.util.ws.BytesocksClient;
+import me.lucko.spark.common.ws.ViewerSocket;
 import me.lucko.spark.proto.SparkSamplerProtos;
 
 import net.kyori.adventure.text.Component;
@@ -68,7 +71,6 @@ import static net.kyori.adventure.text.format.NamedTextColor.RED;
 import static net.kyori.adventure.text.format.NamedTextColor.WHITE;
 
 public class SamplerModule implements CommandModule {
-    private static final String SPARK_SAMPLER_MEDIA_TYPE = "application/x-spark-sampler";
 
     @Override
     public void registerCommands(Consumer<Command> consumer) {
@@ -76,6 +78,7 @@ public class SamplerModule implements CommandModule {
                 .aliases("profiler", "sampler")
                 .allowSubCommand(true)
                 .argumentUsage("info", "", null)
+                .argumentUsage("open", "", null)
                 .argumentUsage("start", "timeout", "timeout seconds")
                 .argumentUsage("start", "thread *", null)
                 .argumentUsage("start", "thread", "thread name")
@@ -103,7 +106,7 @@ public class SamplerModule implements CommandModule {
                     }
 
                     return TabCompleter.create()
-                            .at(0, CompletionSupplier.startsWith(Arrays.asList("info", "start", "stop", "cancel")))
+                            .at(0, CompletionSupplier.startsWith(Arrays.asList("info", "start", "open", "stop", "cancel")))
                             .from(1, CompletionSupplier.startsWith(opts))
                             .complete(arguments);
                 })
@@ -119,6 +122,16 @@ public class SamplerModule implements CommandModule {
             return;
         }
 
+        if (subCommand.equals("open") || arguments.boolFlag("open")) {
+            profilerOpen(platform, sender, resp, arguments);
+            return;
+        }
+
+        if (subCommand.equals("trust-viewer") || arguments.boolFlag("trust-viewer")) {
+            profilerTrustViewer(platform, sender, resp, arguments);
+            return;
+        }
+
         if (subCommand.equals("cancel") || arguments.boolFlag("cancel")) {
             profilerCancel(platform, resp);
             return;
@@ -254,6 +267,8 @@ public class SamplerModule implements CommandModule {
             resp.broadcastPrefixed(text("It will run in the background until it is stopped by an admin."));
             resp.broadcastPrefixed(text("To stop the profiler and upload the results, run:"));
             resp.broadcastPrefixed(cmdPrompt("/" + platform.getPlugin().getCommandName() + " profiler stop"));
+            resp.broadcastPrefixed(text("To view the profiler while it's running, run:"));
+            resp.broadcastPrefixed(cmdPrompt("/" + platform.getPlugin().getCommandName() + " profiler open"));
         } else {
             resp.broadcastPrefixed(text("The results will be automatically returned after the profiler has been running for " + FormatUtil.formatSeconds(timeoutSeconds) + "."));
         }
@@ -273,13 +288,11 @@ public class SamplerModule implements CommandModule {
 
         // await the result
         if (timeoutSeconds != -1) {
-            String comment = Iterables.getFirst(arguments.stringFlag("comment"), null);
-            MethodDisambiguator methodDisambiguator = new MethodDisambiguator();
-            MergeMode mergeMode = arguments.boolFlag("separate-parent-calls") ? MergeMode.separateParentCalls(methodDisambiguator) : MergeMode.sameMethod(methodDisambiguator);
+            Sampler.ExportProps exportProps = getExportProps(platform, resp, arguments);
             boolean saveToFile = arguments.boolFlag("save-to-file");
             future.thenAcceptAsync(s -> {
                 resp.broadcastPrefixed(text("The active profiler has completed! Uploading results..."));
-                handleUpload(platform, resp, s, comment, mergeMode, saveToFile);
+                handleUpload(platform, resp, s, exportProps, saveToFile);
             });
         }
     }
@@ -306,6 +319,9 @@ public class SamplerModule implements CommandModule {
                 resp.replyPrefixed(text("So far, it has profiled for " + FormatUtil.formatSeconds(runningTime) + "."));
             }
 
+            resp.replyPrefixed(text("To view the profiler while it's running, run:"));
+            resp.replyPrefixed(cmdPrompt("/" + platform.getPlugin().getCommandName() + " profiler open"));
+
             long timeout = sampler.getAutoEndTime();
             if (timeout == -1) {
                 resp.replyPrefixed(text("To stop the profiler and upload the results, run:"));
@@ -320,6 +336,48 @@ public class SamplerModule implements CommandModule {
         }
     }
 
+    private void profilerOpen(SparkPlatform platform, CommandSender sender, CommandResponseHandler resp, Arguments arguments) {
+        BytesocksClient bytesocksClient = platform.getBytesocksClient();
+        if (bytesocksClient == null) {
+            resp.replyPrefixed(text("The live viewer is only supported on Java 11 or newer.", RED));
+            return;
+        }
+
+        Sampler sampler = platform.getSamplerContainer().getActiveSampler();
+        if (sampler == null) {
+            resp.replyPrefixed(text("The profiler isn't running!"));
+            resp.replyPrefixed(text("To start a new one, run:"));
+            resp.replyPrefixed(cmdPrompt("/" + platform.getPlugin().getCommandName() + " profiler start"));
+            return;
+        }
+
+        Sampler.ExportProps exportProps = getExportProps(platform, resp, arguments);
+        handleOpen(platform, bytesocksClient, resp, sampler, exportProps);
+    }
+
+    private void profilerTrustViewer(SparkPlatform platform, CommandSender sender, CommandResponseHandler resp, Arguments arguments) {
+        Set<String> ids = arguments.stringFlag("id");
+        if (ids.isEmpty()) {
+            resp.replyPrefixed(text("Please provide a client id with '--id <client id>'."));
+            return;
+        }
+
+        for (String id : ids) {
+            boolean success = platform.getTrustedKeyStore().trustPendingKey(id);
+            if (success) {
+                Sampler sampler = platform.getSamplerContainer().getActiveSampler();
+                if (sampler != null) {
+                    for (ViewerSocket socket : sampler.getAttachedSockets()) {
+                        socket.sendClientTrustedMessage(id);
+                    }
+                }
+                resp.replyPrefixed(text("Client connected to the viewer using id '" + id + "' is now trusted."));
+            } else {
+                resp.replyPrefixed(text("Unable to find pending client with id '" + id + "'."));
+            }
+        }
+    }
+
     private void profilerCancel(SparkPlatform platform, CommandResponseHandler resp) {
         Sampler sampler = platform.getSamplerContainer().getActiveSampler();
         if (sampler == null) {
@@ -346,10 +404,8 @@ public class SamplerModule implements CommandModule {
                 resp.broadcastPrefixed(text("Stopping the profiler & uploading results, please wait..."));
             }
 
-            String comment = Iterables.getFirst(arguments.stringFlag("comment"), null);
-            MethodDisambiguator methodDisambiguator = new MethodDisambiguator();
-            MergeMode mergeMode = arguments.boolFlag("separate-parent-calls") ? MergeMode.separateParentCalls(methodDisambiguator) : MergeMode.sameMethod(methodDisambiguator);
-            handleUpload(platform, resp, sampler, comment, mergeMode, saveToFile);
+            Sampler.ExportProps exportProps = getExportProps(platform, resp, arguments);
+            handleUpload(platform, resp, sampler, exportProps, saveToFile);
 
             // if the previous sampler was running in the background, create a new one
             if (platform.getBackgroundSamplerManager().restartBackgroundSampler()) {
@@ -362,15 +418,15 @@ public class SamplerModule implements CommandModule {
         }
     }
 
-    private void handleUpload(SparkPlatform platform, CommandResponseHandler resp, Sampler sampler, String comment, MergeMode mergeMode, boolean saveToFileFlag) {
-        SparkSamplerProtos.SamplerData output = sampler.toProto(platform, resp.sender(), comment, mergeMode, ClassSourceLookup.create(platform));
+    private void handleUpload(SparkPlatform platform, CommandResponseHandler resp, Sampler sampler, Sampler.ExportProps exportProps, boolean saveToFileFlag) {
+        SparkSamplerProtos.SamplerData output = sampler.toProto(platform, exportProps);
 
         boolean saveToFile = false;
         if (saveToFileFlag) {
             saveToFile = true;
         } else {
             try {
-                String key = platform.getBytebinClient().postContent(output, SPARK_SAMPLER_MEDIA_TYPE).key();
+                String key = platform.getBytebinClient().postContent(output, MediaTypes.SPARK_SAMPLER_MEDIA_TYPE).key();
                 String url = platform.getViewerUrl() + key;
 
                 resp.broadcastPrefixed(text("Profiler stopped & upload complete!", GOLD));
@@ -406,6 +462,45 @@ public class SamplerModule implements CommandModule {
         }
     }
 
+    private void handleOpen(SparkPlatform platform, BytesocksClient bytesocksClient, CommandResponseHandler resp, Sampler sampler, Sampler.ExportProps exportProps) {
+        try {
+            ViewerSocket socket = new ViewerSocket(platform, bytesocksClient, exportProps);
+            sampler.attachSocket(socket);
+            exportProps.channelInfo(socket.getPayload());
+
+            SparkSamplerProtos.SamplerData data = sampler.toProto(platform, exportProps);
+
+            String key = platform.getBytebinClient().postContent(data, MediaTypes.SPARK_SAMPLER_MEDIA_TYPE, "live").key();
+            String url = platform.getViewerUrl() + key;
+
+            resp.broadcastPrefixed(text("Profiler live viewer:", GOLD));
+            resp.broadcast(text()
+                    .content(url)
+                    .color(GRAY)
+                    .clickEvent(ClickEvent.openUrl(url))
+                    .build()
+            );
+
+            platform.getActivityLog().addToLog(Activity.urlActivity(resp.sender(), System.currentTimeMillis(), "Profiler (live)", url));
+        } catch (Exception e) {
+            resp.replyPrefixed(text("An error occurred whilst opening the live profiler.", RED));
+            e.printStackTrace();
+        }
+    }
+
+    private Sampler.ExportProps getExportProps(SparkPlatform platform, CommandResponseHandler resp, Arguments arguments) {
+        return new Sampler.ExportProps()
+                .creator(resp.sender().toData())
+                .comment(Iterables.getFirst(arguments.stringFlag("comment"), null))
+                .mergeMode(() -> {
+                    MethodDisambiguator methodDisambiguator = new MethodDisambiguator();
+                    return arguments.boolFlag("separate-parent-calls")
+                            ? MergeMode.separateParentCalls(methodDisambiguator)
+                            : MergeMode.sameMethod(methodDisambiguator);
+                })
+                .classSourceLookup(() -> ClassSourceLookup.create(platform));
+    }
+
     private static Component cmdPrompt(String cmd) {
         return text()
                 .append(text("  "))
diff --git a/spark-common/src/main/java/me/lucko/spark/common/heapdump/HeapDumpSummary.java b/spark-common/src/main/java/me/lucko/spark/common/heapdump/HeapDumpSummary.java
index c0980e7..eaedd31 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/heapdump/HeapDumpSummary.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/heapdump/HeapDumpSummary.java
@@ -130,7 +130,7 @@ public final class HeapDumpSummary {
                 .setPlatformMetadata(platform.getPlugin().getPlatformInfo().toData().toProto())
                 .setCreator(creator.toData().toProto());
         try {
-            metadata.setPlatformStatistics(platform.getStatisticsProvider().getPlatformStatistics(null));
+            metadata.setPlatformStatistics(platform.getStatisticsProvider().getPlatformStatistics(null, true));
         } catch (Exception e) {
             e.printStackTrace();
         }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/platform/PlatformStatisticsProvider.java b/spark-common/src/main/java/me/lucko/spark/common/platform/PlatformStatisticsProvider.java
index fc7e78a..059590d 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/platform/PlatformStatisticsProvider.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/platform/PlatformStatisticsProvider.java
@@ -128,7 +128,7 @@ public class PlatformStatisticsProvider {
         return builder.build();
     }
 
-    public PlatformStatistics getPlatformStatistics(Map<String, GarbageCollectorStatistics> startingGcStatistics) {
+    public PlatformStatistics getPlatformStatistics(Map<String, GarbageCollectorStatistics> startingGcStatistics, boolean includeWorld) {
         PlatformStatistics.Builder builder = PlatformStatistics.newBuilder();
 
         MemoryUsage memoryUsage = ManagementFactory.getMemoryMXBean().getHeapMemoryUsage();
@@ -187,19 +187,20 @@ public class PlatformStatisticsProvider {
             builder.setPlayerCount(playerCount);
         }
 
-        try {
-            WorldStatisticsProvider worldStatisticsProvider = new WorldStatisticsProvider(
-                    new AsyncWorldInfoProvider(this.platform, this.platform.getPlugin().createWorldInfoProvider())
-            );
-            WorldStatistics worldStatistics = worldStatisticsProvider.getWorldStatistics();
-            if (worldStatistics != null) {
-                builder.setWorld(worldStatistics);
+        if (includeWorld) {
+            try {
+                WorldStatisticsProvider worldStatisticsProvider = new WorldStatisticsProvider(
+                        new AsyncWorldInfoProvider(this.platform, this.platform.getPlugin().createWorldInfoProvider())
+                );
+                WorldStatistics worldStatistics = worldStatisticsProvider.getWorldStatistics();
+                if (worldStatistics != null) {
+                    builder.setWorld(worldStatistics);
+                }
+            } catch (Exception e) {
+                e.printStackTrace();
             }
-        } catch (Exception e) {
-            e.printStackTrace();
         }
 
-
         return builder.build();
     }
 
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java
index 5abe71f..d814002 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java
@@ -32,9 +32,12 @@ import me.lucko.spark.common.sampler.source.ClassSourceLookup;
 import me.lucko.spark.common.sampler.source.SourceMetadata;
 import me.lucko.spark.common.sampler.window.ProtoTimeEncoder;
 import me.lucko.spark.common.sampler.window.WindowStatisticsCollector;
+import me.lucko.spark.common.ws.ViewerSocket;
+import me.lucko.spark.proto.SparkProtos;
 import me.lucko.spark.proto.SparkSamplerProtos.SamplerData;
 import me.lucko.spark.proto.SparkSamplerProtos.SamplerMetadata;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.List;
@@ -74,6 +77,9 @@ public abstract class AbstractSampler implements Sampler {
     /** The garbage collector statistics when profiling started */
     protected Map<String, GarbageCollectorStatistics> initialGcStats;
 
+    /** A set of viewer sockets linked to the sampler */
+    protected List<ViewerSocket> viewerSockets = new ArrayList<>();
+
     protected AbstractSampler(SparkPlatform platform, SamplerSettings settings) {
         this.platform = platform;
         this.interval = settings.interval();
@@ -122,13 +128,54 @@ public abstract class AbstractSampler implements Sampler {
     @Override
     public void stop(boolean cancelled) {
         this.windowStatisticsCollector.stop();
+        for (ViewerSocket viewerSocket : this.viewerSockets) {
+            viewerSocket.processSamplerStopped(this);
+        }
+    }
+
+    @Override
+    public void attachSocket(ViewerSocket socket) {
+        this.viewerSockets.add(socket);
+    }
+
+    @Override
+    public Collection<ViewerSocket> getAttachedSockets() {
+        return this.viewerSockets;
+    }
+
+    protected void processWindowRotate() {
+        this.viewerSockets.removeIf(socket -> {
+            if (!socket.isOpen()) {
+                return true;
+            }
+
+            socket.processWindowRotate(this);
+            return false;
+        });
+    }
+
+    protected void sendStatisticsToSocket() {
+        try {
+            if (this.viewerSockets.isEmpty()) {
+                return;
+            }
+
+            SparkProtos.PlatformStatistics platform = this.platform.getStatisticsProvider().getPlatformStatistics(getInitialGcStats(), false);
+            SparkProtos.SystemStatistics system = this.platform.getStatisticsProvider().getSystemStatistics();
+
+            for (ViewerSocket viewerSocket : this.viewerSockets) {
+                viewerSocket.sendUpdatedStatistics(platform, system);
+            }
+        } catch (Exception e) {
+            e.printStackTrace();
+        }
     }
 
-    protected void writeMetadataToProto(SamplerData.Builder proto, SparkPlatform platform, CommandSender creator, String comment, DataAggregator dataAggregator) {
+    protected void writeMetadataToProto(SamplerData.Builder proto, SparkPlatform platform, CommandSender.Data creator, String comment, DataAggregator dataAggregator) {
         SamplerMetadata.Builder metadata = SamplerMetadata.newBuilder()
                 .setSamplerMode(getMode().asProto())
                 .setPlatformMetadata(platform.getPlugin().getPlatformInfo().toData().toProto())
-                .setCreator(creator.toData().toProto())
+                .setCreator(creator.toProto())
                 .setStartTime(this.startTime)
                 .setEndTime(System.currentTimeMillis())
                 .setInterval(this.interval)
@@ -145,7 +192,7 @@ public abstract class AbstractSampler implements Sampler {
         }
 
         try {
-            metadata.setPlatformStatistics(platform.getStatisticsProvider().getPlatformStatistics(getInitialGcStats()));
+            metadata.setPlatformStatistics(platform.getStatisticsProvider().getPlatformStatistics(getInitialGcStats(), true));
         } catch (Exception e) {
             e.printStackTrace();
         }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/Sampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/Sampler.java
index aaf4f38..844ab0b 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/Sampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/Sampler.java
@@ -24,9 +24,13 @@ import me.lucko.spark.common.SparkPlatform;
 import me.lucko.spark.common.command.sender.CommandSender;
 import me.lucko.spark.common.sampler.node.MergeMode;
 import me.lucko.spark.common.sampler.source.ClassSourceLookup;
+import me.lucko.spark.common.ws.ViewerSocket;
 import me.lucko.spark.proto.SparkSamplerProtos.SamplerData;
+import me.lucko.spark.proto.SparkSamplerProtos.SocketChannelInfo;
 
+import java.util.Collection;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Supplier;
 
 /**
  * Abstract superinterface for all sampler implementations.
@@ -44,6 +48,20 @@ public interface Sampler {
     void stop(boolean cancelled);
 
     /**
+     * Attaches a viewer socket to this sampler.
+     *
+     * @param socket the socket
+     */
+    void attachSocket(ViewerSocket socket);
+
+    /**
+     * Gets the sockets attached to this sampler.
+     *
+     * @return the attached sockets
+     */
+    Collection<ViewerSocket> getAttachedSockets();
+
+    /**
      * Gets the time when the sampler started (unix timestamp in millis)
      *
      * @return the start time
@@ -79,6 +97,62 @@ public interface Sampler {
     CompletableFuture<Sampler> getFuture();
 
     // Methods used to export the sampler data to the web viewer.
-    SamplerData toProto(SparkPlatform platform, CommandSender creator, String comment, MergeMode mergeMode, ClassSourceLookup classSourceLookup);
+    SamplerData toProto(SparkPlatform platform, ExportProps exportProps);
+
+    final class ExportProps {
+        private CommandSender.Data creator;
+        private String comment;
+        private Supplier<MergeMode> mergeMode;
+        private Supplier<ClassSourceLookup> classSourceLookup;
+        private SocketChannelInfo channelInfo;
+
+        public ExportProps() {
+        }
+
+        public CommandSender.Data creator() {
+            return this.creator;
+        }
+
+        public String comment() {
+            return this.comment;
+        }
+
+        public Supplier<MergeMode> mergeMode() {
+            return this.mergeMode;
+        }
+
+        public Supplier<ClassSourceLookup> classSourceLookup() {
+            return this.classSourceLookup;
+        }
+
+        public SocketChannelInfo channelInfo() {
+            return this.channelInfo;
+        }
+
+        public ExportProps creator(CommandSender.Data creator) {
+            this.creator = creator;
+            return this;
+        }
+
+        public ExportProps comment(String comment) {
+            this.comment = comment;
+            return this;
+        }
+
+        public ExportProps mergeMode(Supplier<MergeMode> mergeMode) {
+            this.mergeMode = mergeMode;
+            return this;
+        }
+
+        public ExportProps classSourceLookup(Supplier<ClassSourceLookup> classSourceLookup) {
+            this.classSourceLookup = classSourceLookup;
+            return this;
+        }
+
+        public ExportProps channelInfo(SocketChannelInfo channelInfo) {
+            this.channelInfo = channelInfo;
+            return this;
+        }
+    }
 
 }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java
index 2328582..ec88677 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java
@@ -23,18 +23,17 @@ package me.lucko.spark.common.sampler.async;