Include player/entity/chunk counts in window statistics

13 files changed, 199 insertions, 63 deletions

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 00cd4fa..f576eac 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
@@ -145,7 +145,7 @@ public class SamplerModule implements CommandModule {
             if (previousSampler.isRunningInBackground()) {
                 // there is a background profiler running - stop that first
                 resp.replyPrefixed(text("Stopping the background profiler before starting... please wait"));
-                previousSampler.stop();
+                previousSampler.stop(true);
                 platform.getSamplerContainer().unsetActiveSampler(previousSampler);
             } else {
                 // there is a non-background profiler running - tell the user
@@ -310,7 +310,7 @@ public class SamplerModule implements CommandModule {
         if (sampler == null) {
             resp.replyPrefixed(text("There isn't an active profiler running."));
         } else {
-            platform.getSamplerContainer().stopActiveSampler();
+            platform.getSamplerContainer().stopActiveSampler(true);
             resp.broadcastPrefixed(text("Profiler has been cancelled.", GOLD));
         }
     }
@@ -322,7 +322,7 @@ public class SamplerModule implements CommandModule {
             resp.replyPrefixed(text("There isn't an active profiler running."));
         } else {
             platform.getSamplerContainer().unsetActiveSampler(sampler);
-            sampler.stop();
+            sampler.stop(false);
 
             boolean saveToFile = arguments.boolFlag("save-to-file");
             if (saveToFile) {
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 1eb9753..fc7e78a 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
@@ -31,7 +31,7 @@ import me.lucko.spark.common.monitor.net.NetworkMonitor;
 import me.lucko.spark.common.monitor.os.OperatingSystemInfo;
 import me.lucko.spark.common.monitor.ping.PingStatistics;
 import me.lucko.spark.common.monitor.tick.TickStatistics;
-import me.lucko.spark.common.platform.world.WorldInfoProvider;
+import me.lucko.spark.common.platform.world.AsyncWorldInfoProvider;
 import me.lucko.spark.common.platform.world.WorldStatisticsProvider;
 import me.lucko.spark.proto.SparkProtos.PlatformStatistics;
 import me.lucko.spark.proto.SparkProtos.SystemStatistics;
@@ -188,8 +188,9 @@ public class PlatformStatisticsProvider {
         }
 
         try {
-            WorldInfoProvider worldInfo = this.platform.getPlugin().createWorldInfoProvider();
-            WorldStatisticsProvider worldStatisticsProvider = new WorldStatisticsProvider(this.platform, worldInfo);
+            WorldStatisticsProvider worldStatisticsProvider = new WorldStatisticsProvider(
+                    new AsyncWorldInfoProvider(this.platform, this.platform.getPlugin().createWorldInfoProvider())
+            );
             WorldStatistics worldStatistics = worldStatisticsProvider.getWorldStatistics();
             if (worldStatistics != null) {
                 builder.setWorld(worldStatistics);
diff --git a/spark-common/src/main/java/me/lucko/spark/common/platform/world/AsyncWorldInfoProvider.java b/spark-common/src/main/java/me/lucko/spark/common/platform/world/AsyncWorldInfoProvider.java
new file mode 100644
index 0000000..82cddef
--- /dev/null
+++ b/spark-common/src/main/java/me/lucko/spark/common/platform/world/AsyncWorldInfoProvider.java
@@ -0,0 +1,90 @@
+/*
+ * This file is part of spark.
+ *
+ *  Copyright (c) lucko (Luck) <luck@lucko.me>
+ *  Copyright (c) contributors
+ *
+ *  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 me.lucko.spark.common.platform.world;
+
+import me.lucko.spark.common.SparkPlatform;
+import me.lucko.spark.common.SparkPlugin;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.function.Function;
+import java.util.logging.Level;
+
+/**
+ * Async-friendly wrapper around {@link WorldInfoProvider}.
+ */
+public class AsyncWorldInfoProvider {
+    private static final int TIMEOUT_SECONDS = 5;
+
+    private final SparkPlatform platform;
+    private final WorldInfoProvider provider;
+
+    public AsyncWorldInfoProvider(SparkPlatform platform, WorldInfoProvider provider) {
+        this.platform = platform;
+        this.provider = provider == WorldInfoProvider.NO_OP ? null : provider;
+    }
+
+    private <T> CompletableFuture<T> async(Function<WorldInfoProvider, T> function) {
+        if (this.provider == null) {
+            return null;
+        }
+
+        if (this.provider.mustCallSync()) {
+            SparkPlugin plugin = this.platform.getPlugin();
+            return CompletableFuture.supplyAsync(() -> function.apply(this.provider), plugin::executeSync);
+        } else {
+            return CompletableFuture.completedFuture(function.apply(this.provider));
+        }
+    }
+
+    private <T> T get(CompletableFuture<T> future) {
+        if (future == null) {
+            return null;
+        }
+
+        try {
+            return future.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
+        } catch (InterruptedException | ExecutionException e) {
+            throw new RuntimeException(e);
+        } catch (TimeoutException e) {
+            this.platform.getPlugin().log(Level.WARNING, "Timed out waiting for world statistics");
+            return null;
+        }
+    }
+
+    public CompletableFuture<WorldInfoProvider.CountsResult> pollCounts() {
+        return async(WorldInfoProvider::pollCounts);
+    }
+
+    public CompletableFuture<WorldInfoProvider.ChunksResult<? extends ChunkInfo<?>>> pollChunks() {
+        return async(WorldInfoProvider::pollChunks);
+    }
+
+    public WorldInfoProvider.CountsResult getCounts() {
+        return get(pollCounts());
+    }
+
+    public WorldInfoProvider.ChunksResult<? extends ChunkInfo<?>> getChunks() {
+        return get(pollChunks());
+    }
+}
diff --git a/spark-common/src/main/java/me/lucko/spark/common/platform/world/WorldInfoProvider.java b/spark-common/src/main/java/me/lucko/spark/common/platform/world/WorldInfoProvider.java
index 9494816..7fb581d 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/platform/world/WorldInfoProvider.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/platform/world/WorldInfoProvider.java
@@ -29,20 +29,37 @@ import java.util.Map;
  */
 public interface WorldInfoProvider {
 
-    WorldInfoProvider NO_OP = () -> null;
+    WorldInfoProvider NO_OP = new WorldInfoProvider() {
+        @Override
+        public CountsResult pollCounts() {
+            return null;
+        }
+
+        @Override
+        public ChunksResult<? extends ChunkInfo<?>> pollChunks() {
+            return null;
+        }
+    };
+
+    /**
+     * Polls for counts.
+     *
+     * @return the counts
+     */
+    CountsResult pollCounts();
 
     /**
-     * Polls for information.
+     * Polls for chunk information.
      *
-     * @return the information
+     * @return the chunk information
      */
-    Result<? extends ChunkInfo<?>> poll();
+    ChunksResult<? extends ChunkInfo<?>> pollChunks();
 
     default boolean mustCallSync() {
         return true;
     }
 
-    final class Result<T> {
+    final class ChunksResult<T extends ChunkInfo<?>> {
         private final Map<String, List<T>> worlds = new HashMap<>();
 
         public void put(String worldName, List<T> chunks) {
@@ -54,4 +71,34 @@ public interface WorldInfoProvider {
         }
     }
 
+    final class CountsResult {
+        private final int players;
+        private final int entities;
+        private final int tileEntities;
+        private final int chunks;
+
+        public CountsResult(int players, int entities, int tileEntities, int chunks) {
+            this.players = players;
+            this.entities = entities;
+            this.tileEntities = tileEntities;
+            this.chunks = chunks;
+        }
+
+        public int players() {
+            return this.players;
+        }
+
+        public int entities() {
+            return this.entities;
+        }
+
+        public int tileEntities() {
+            return this.tileEntities;
+        }
+
+        public int chunks() {
+            return this.chunks;
+        }
+    }
+
 }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/platform/world/WorldStatisticsProvider.java b/spark-common/src/main/java/me/lucko/spark/common/platform/world/WorldStatisticsProvider.java
index 80c35a6..7e63222 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/platform/world/WorldStatisticsProvider.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/platform/world/WorldStatisticsProvider.java
@@ -20,8 +20,6 @@
 
 package me.lucko.spark.common.platform.world;
 
-import me.lucko.spark.common.SparkPlatform;
-import me.lucko.spark.common.SparkPlugin;
 import me.lucko.spark.proto.SparkProtos.WorldStatistics;
 
 import java.util.ArrayList;
@@ -30,46 +28,17 @@ import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Set;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.logging.Level;
 
 public class WorldStatisticsProvider {
-    private final SparkPlatform platform;
-    private final WorldInfoProvider provider;
+    private final AsyncWorldInfoProvider provider;
 
-    public WorldStatisticsProvider(SparkPlatform platform, WorldInfoProvider provider) {
-        this.platform = platform;
+    public WorldStatisticsProvider(AsyncWorldInfoProvider provider) {
         this.provider = provider;
     }
 
     public WorldStatistics getWorldStatistics() {
-        if (this.provider == WorldInfoProvider.NO_OP) {
-            return null;
-        }
-
-        CompletableFuture<WorldInfoProvider.Result<? extends ChunkInfo<?>>> future;
-
-        if (this.provider.mustCallSync()) {
-            SparkPlugin plugin = this.platform.getPlugin();
-            future = CompletableFuture.supplyAsync(this.provider::poll, plugin::executeSync);
-        } else {
-            future = CompletableFuture.completedFuture(this.provider.poll());
-        }
-
-        WorldInfoProvider.Result<? extends ChunkInfo<?>> result;
-        try {
-            result = future.get(5, TimeUnit.SECONDS);
-        } catch (InterruptedException | ExecutionException e) {
-            throw new RuntimeException(e);
-        } catch (TimeoutException e) {
-            this.platform.getPlugin().log(Level.WARNING, "Timed out waiting for world statistics");
-            return null;
-        }
-
+        WorldInfoProvider.ChunksResult<? extends ChunkInfo<?>> result = provider.getChunks();
         if (result == null) {
             return null;
         }
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 59e873c..e324fd3 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
@@ -120,7 +120,7 @@ public abstract class AbstractSampler implements Sampler {
     }
 
     @Override
-    public void stop() {
+    public void stop(boolean cancelled) {
         this.windowStatisticsCollector.stop();
     }
 
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 5d2026d..36a63f1 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
@@ -41,7 +41,7 @@ public interface Sampler {
     /**
      * Stops the sampler.
      */
-    void stop();
+    void stop(boolean cancelled);
 
     /**
      * Gets the time when the sampler started (unix timestamp in millis)
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/SamplerContainer.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/SamplerContainer.java
index f56dee5..d55909c 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/SamplerContainer.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/SamplerContainer.java
@@ -66,10 +66,10 @@ public class SamplerContainer implements AutoCloseable {
     /**
      * Stops the active sampler, if there is one.
      */
-    public void stopActiveSampler() {
+    public void stopActiveSampler(boolean cancelled) {
         Sampler sampler = this.activeSampler.getAndSet(null);
         if (sampler != null) {
-            sampler.stop();
+            sampler.stop(cancelled);
         }
     }
 
@@ -79,7 +79,7 @@ public class SamplerContainer implements AutoCloseable {
 
     @Override
     public void close() {
-        stopActiveSampler();
+        stopActiveSampler(true);
     }
 
 }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncProfilerJob.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncProfilerJob.java
index db1808c..d74b75f 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncProfilerJob.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncProfilerJob.java
@@ -224,13 +224,15 @@ public class AsyncProfilerJob {
             }
         }
 
-        // delete the output file after reading
+        deleteOutputFile();
+    }
+
+    public void deleteOutputFile() {
         try {
             Files.deleteIfExists(this.outputFile);
         } catch (IOException e) {
             // ignore
         }
-
     }
 
     private void readSegments(JfrReader reader, Predicate<String> threadFilter, AsyncDataAggregator dataAggregator, int window) throws IOException {
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 f2e7191..178f055 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
@@ -144,7 +144,7 @@ public class AsyncSampler extends AbstractSampler {
         }
 
         this.scheduler.schedule(() -> {
-            stop();
+            stop(false);
             this.future.complete(this);
         }, delay, TimeUnit.MILLISECONDS);
     }
@@ -153,13 +153,17 @@ public class AsyncSampler extends AbstractSampler {
      * Stops the profiler.
      */
     @Override
-    public void stop() {
-        super.stop();
+    public void stop(boolean cancelled) {
+        super.stop(cancelled);
 
         synchronized (this.currentJobMutex) {
             this.currentJob.stop();
-            this.windowStatisticsCollector.measureNow(this.currentJob.getWindow());
-            this.currentJob.aggregate(this.dataAggregator);
+            if (!cancelled) {
+                this.windowStatisticsCollector.measureNow(this.currentJob.getWindow());
+                this.currentJob.aggregate(this.dataAggregator);
+            } else {
+                this.currentJob.deleteOutputFile();
+            }
             this.currentJob = null;
         }
 
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java
index 42a457d..72a37e8 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java
@@ -94,13 +94,15 @@ public class JavaSampler extends AbstractSampler implements Runnable {
     }
 
     @Override
-    public void stop() {
-        super.stop();
+    public void stop(boolean cancelled) {
+        super.stop(cancelled);
 
         this.task.cancel(false);
 
-        // collect statistics for the final window
-        this.windowStatisticsCollector.measureNow(this.lastWindow.get());
+        if (!cancelled) {
+            // collect statistics for the final window
+            this.windowStatisticsCollector.measureNow(this.lastWindow.get());
+        }
     }
 
     @Override
@@ -111,7 +113,7 @@ public class JavaSampler extends AbstractSampler implements Runnable {
             long time = System.currentTimeMillis();
 
             if (this.autoEndTime != -1 && this.autoEndTime <= time) {
-                stop();
+                stop(false);
                 this.future.complete(this);
                 return;
             }
@@ -120,7 +122,7 @@ public class JavaSampler extends AbstractSampler implements Runnable {
             ThreadInfo[] threadDumps = this.threadDumper.dumpThreads(this.threadBean);
             this.workerPool.execute(new InsertDataTask(threadDumps, window));
         } catch (Throwable t) {
-            stop();
+            stop(false);
             this.future.completeExceptionally(t);
         }
     }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/window/WindowStatisticsCollector.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/window/WindowStatisticsCollector.java
index 7da62fa..ce65013 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/window/WindowStatisticsCollector.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/window/WindowStatisticsCollector.java
@@ -23,6 +23,8 @@ package me.lucko.spark.common.sampler.window;
 import me.lucko.spark.common.SparkPlatform;
 import me.lucko.spark.common.monitor.cpu.CpuMonitor;
 import me.lucko.spark.common.monitor.tick.TickStatistics;
+import me.lucko.spark.common.platform.world.AsyncWorldInfoProvider;
+import me.lucko.spark.common.platform.world.WorldInfoProvider;
 import me.lucko.spark.common.tick.TickHook;
 import me.lucko.spark.common.util.RollingAverage;
 import me.lucko.spark.proto.SparkProtos;
@@ -152,6 +154,19 @@ public class WindowStatisticsCollector {
         builder.setCpuProcess(CpuMonitor.processLoad1MinAvg());
         builder.setCpuSystem(CpuMonitor.systemLoad1MinAvg());
 
+        try {
+            AsyncWorldInfoProvider worldInfoProvider = new AsyncWorldInfoProvider(this.platform, this.platform.getPlugin().createWorldInfoProvider());
+            WorldInfoProvider.CountsResult counts = worldInfoProvider.getCounts();
+            if (counts != null) {
+                builder.setPlayers(counts.players());
+                builder.setEntities(counts.entities());
+                builder.setTileEntities(counts.tileEntities());
+                builder.setChunks(counts.chunks());
+            }
+        } catch (Exception e) {
+            e.printStackTrace();
+        }
+
         return builder.build();
     }
 
diff --git a/spark-common/src/main/proto/spark/spark.proto b/spark-common/src/main/proto/spark/spark.proto
index be76bd7..f61e585 100644
--- a/spark-common/src/main/proto/spark/spark.proto
+++ b/spark-common/src/main/proto/spark/spark.proto
@@ -159,6 +159,12 @@ message WindowStatistics {
   double tps = 4;
   double mspt_median = 5;
   double mspt_max = 6;
+
+  // world
+  int32 players = 7;
+  int32 entities = 8;
+  int32 tile_entities = 9;
+  int32 chunks = 10;
 }
 
 message RollingAverageValues {