Use WeakReference for command senders (#417)

5 files changed, 31 insertions, 21 deletions

diff --git a/spark-common/src/main/java/me/lucko/spark/common/activitylog/Activity.java b/spark-common/src/main/java/me/lucko/spark/common/activitylog/Activity.java
index 0b69826..c66888a 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/activitylog/Activity.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/activitylog/Activity.java
@@ -38,12 +38,12 @@ public final class Activity {
     private final String dataType;
     private final String dataValue;
 
-    public static Activity urlActivity(CommandSender user, long time, String type, String url) {
-        return new Activity(user.toData(), time, type, DATA_TYPE_URL, url);
+    public static Activity urlActivity(CommandSender.Data user, long time, String type, String url) {
+        return new Activity(user, time, type, DATA_TYPE_URL, url);
     }
 
-    public static Activity fileActivity(CommandSender user, long time, String type, String filePath) {
-        return new Activity(user.toData(), time, type, DATA_TYPE_FILE, filePath);
+    public static Activity fileActivity(CommandSender.Data user, long time, String type, String filePath) {
+        return new Activity(user, time, type, DATA_TYPE_FILE, filePath);
     }
 
     private Activity(CommandSender.Data user, long time, String type, String dataType, String dataValue) {
diff --git a/spark-common/src/main/java/me/lucko/spark/common/command/CommandResponseHandler.java b/spark-common/src/main/java/me/lucko/spark/common/command/CommandResponseHandler.java
index c2e8223..3a894ca 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/command/CommandResponseHandler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/command/CommandResponseHandler.java
@@ -26,6 +26,7 @@ import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.JoinConfiguration;
 import net.kyori.adventure.text.TextComponent;
 
+import java.lang.ref.WeakReference;
 import java.util.Set;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
@@ -48,20 +49,22 @@ public class CommandResponseHandler {
             .build();
 
     private final SparkPlatform platform;
-    private final CommandSender sender;
+    private final CommandSender.Data senderData;
+    private final WeakReference<CommandSender> sender;
     private String commandPrimaryAlias;
 
     public CommandResponseHandler(SparkPlatform platform, CommandSender sender) {
         this.platform = platform;
-        this.sender = sender;
+        this.senderData = sender.toData();
+        this.sender = new WeakReference<>(sender);
     }
 
     public void setCommandPrimaryAlias(String commandPrimaryAlias) {
         this.commandPrimaryAlias = commandPrimaryAlias;
     }
 
-    public CommandSender sender() {
-        return this.sender;
+    public CommandSender.Data senderData() {
+        return this.senderData;
     }
 
     public void allSenders(Consumer<? super CommandSender> action) {
@@ -73,17 +76,24 @@ public class CommandResponseHandler {
                 .filter(s -> s.hasPermission("spark") || s.hasPermission("spark." + this.commandPrimaryAlias))
                 .collect(Collectors.toSet());
 
-        senders.add(this.sender);
+        CommandSender sender = this.sender.get();
+        if (sender != null) {
+            senders.add(sender);
+        }
+
         senders.forEach(action);
     }
 
     public void reply(Component message) {
-        this.sender.sendMessage(message);
+        CommandSender sender = this.sender.get();
+        if (sender != null) {
+            sender.sendMessage(message);
+        }
     }
 
     public void reply(Iterable<Component> message) {
         Component joinedMsg = Component.join(JoinConfiguration.separator(Component.newline()), message);
-        this.sender.sendMessage(joinedMsg);
+        reply(joinedMsg);
     }
 
     public void broadcast(Component message) {
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 503e862..54f7df1 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
@@ -89,7 +89,7 @@ public class HeapAnalysisModule implements CommandModule {
             return;
         }
 
-        SparkHeapProtos.HeapData output = heapDump.toProto(platform, sender);
+        SparkHeapProtos.HeapData output = heapDump.toProto(platform, resp.senderData());
 
         boolean saveToFile = false;
         if (arguments.boolFlag("save-to-file")) {
@@ -107,7 +107,7 @@ public class HeapAnalysisModule implements CommandModule {
                         .build()
                 );
 
-                platform.getActivityLog().addToLog(Activity.urlActivity(sender, System.currentTimeMillis(), "Heap dump summary", url));
+                platform.getActivityLog().addToLog(Activity.urlActivity(resp.senderData(), System.currentTimeMillis(), "Heap dump summary", url));
             } catch (Exception e) {
                 resp.broadcastPrefixed(text("An error occurred whilst uploading the data. Attempting to save to disk instead.", RED));
                 e.printStackTrace();
@@ -128,7 +128,7 @@ public class HeapAnalysisModule implements CommandModule {
                 );
                 resp.broadcastPrefixed(text("You can read the heap dump summary file using the viewer web-app - " + platform.getViewerUrl(), GRAY));
 
-                platform.getActivityLog().addToLog(Activity.fileActivity(sender, System.currentTimeMillis(), "Heap dump summary", file.toString()));
+                platform.getActivityLog().addToLog(Activity.fileActivity(resp.senderData(), System.currentTimeMillis(), "Heap dump summary", file.toString()));
             } catch (IOException e) {
                 resp.broadcastPrefixed(text("An error occurred whilst saving the data.", RED));
                 e.printStackTrace();
@@ -163,7 +163,7 @@ public class HeapAnalysisModule implements CommandModule {
                 .append(text(file.toString(), GRAY))
                 .build()
         );
-        platform.getActivityLog().addToLog(Activity.fileActivity(sender, System.currentTimeMillis(), "Heap dump", file.toString()));
+        platform.getActivityLog().addToLog(Activity.fileActivity(resp.senderData(), System.currentTimeMillis(), "Heap dump", file.toString()));
 
 
         Compression compressionMethod = null;
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 ae8f869..9ce66dc 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
@@ -436,7 +436,7 @@ public class SamplerModule implements CommandModule {
                         .build()
                 );
 
-                platform.getActivityLog().addToLog(Activity.urlActivity(resp.sender(), System.currentTimeMillis(), "Profiler", url));
+                platform.getActivityLog().addToLog(Activity.urlActivity(resp.senderData(), System.currentTimeMillis(), "Profiler", url));
             } catch (Exception e) {
                 resp.broadcastPrefixed(text("An error occurred whilst uploading the results. Attempting to save to disk instead.", RED));
                 e.printStackTrace();
@@ -453,7 +453,7 @@ public class SamplerModule implements CommandModule {
                 resp.broadcastPrefixed(text("Data has been written to: " + file));
                 resp.broadcastPrefixed(text("You can view the profile file using the web app @ " + platform.getViewerUrl(), GRAY));
 
-                platform.getActivityLog().addToLog(Activity.fileActivity(resp.sender(), System.currentTimeMillis(), "Profiler", file.toString()));
+                platform.getActivityLog().addToLog(Activity.fileActivity(resp.senderData(), System.currentTimeMillis(), "Profiler", file.toString()));
             } catch (IOException e) {
                 resp.broadcastPrefixed(text("An error occurred whilst saving the data.", RED));
                 e.printStackTrace();
@@ -495,7 +495,7 @@ public class SamplerModule implements CommandModule {
                     .build()
             );
 
-            platform.getActivityLog().addToLog(Activity.urlActivity(resp.sender(), System.currentTimeMillis(), "Profiler (live)", url));
+            platform.getActivityLog().addToLog(Activity.urlActivity(resp.senderData(), System.currentTimeMillis(), "Profiler (live)", url));
         } catch (Exception e) {
             resp.replyPrefixed(text("An error occurred whilst opening the live profiler.", RED));
             e.printStackTrace();
@@ -504,7 +504,7 @@ public class SamplerModule implements CommandModule {
 
     private Sampler.ExportProps getExportProps(SparkPlatform platform, CommandResponseHandler resp, Arguments arguments) {
         return new Sampler.ExportProps()
-                .creator(resp.sender().toData())
+                .creator(resp.senderData())
                 .comment(Iterables.getFirst(arguments.stringFlag("comment"), null))
                 .mergeMode(() -> {
                     MethodDisambiguator methodDisambiguator = new MethodDisambiguator();
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 c7f4f51..364edd6 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
@@ -123,10 +123,10 @@ public final class HeapDumpSummary {
         this.entries = entries;
     }
 
-    public HeapData toProto(SparkPlatform platform, CommandSender creator) {
+    public HeapData toProto(SparkPlatform platform, CommandSender.Data creator) {
         HeapMetadata.Builder metadata = HeapMetadata.newBuilder()
                 .setPlatformMetadata(platform.getPlugin().getPlatformInfo().toData().toProto())
-                .setCreator(creator.toData().toProto());
+                .setCreator(creator.toProto());
         try {
             metadata.setPlatformStatistics(platform.getStatisticsProvider().getPlatformStatistics(null, true));
         } catch (Exception e) {