Fix Java sampler live upload closing thread pool

Author: lucko

spark-common/src/main/java/me/lucko/spark/common/monitor/MonitoringExecutor.java

@@ -20,20 +20,14 @@
 
 package me.lucko.spark.common.monitor;
 
+import me.lucko.spark.common.util.SparkScheduledThreadPoolExecutor;
 import me.lucko.spark.common.util.SparkThreadFactory;
 
-import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 
 public enum MonitoringExecutor {
     ;
 
     /** The executor used to monitor & calculate rolling averages. */
-    public static final ScheduledExecutorService INSTANCE = Executors.newSingleThreadScheduledExecutor(r -> {
-        Thread thread = Executors.defaultThreadFactory().newThread(r);
-        thread.setName("spark-monitoring-thread");
-        thread.setUncaughtExceptionHandler(SparkThreadFactory.EXCEPTION_HANDLER);
-        thread.setDaemon(true);
-        return thread;
-    });
+    public static final ScheduledExecutorService INSTANCE = new SparkScheduledThreadPoolExecutor(1, new SparkThreadFactory("spark-monitoring", true));
 }

spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncProfilerJob.java

@@ -25,7 +25,6 @@
 import me.lucko.spark.common.sampler.ThreadDumper;
 import me.lucko.spark.common.sampler.async.jfr.JfrReader;
 import one.profiler.AsyncProfiler;
-import org.checkerframework.checker.lock.qual.GuardedBy;
 
 import java.io.IOException;
 import java.nio.file.Files;

spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java

@@ -20,7 +20,6 @@
 
 package me.lucko.spark.common.sampler.async;
 
-import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import me.lucko.spark.common.SparkPlatform;
 import me.lucko.spark.common.platform.PlatformInfo;
 import me.lucko.spark.common.sampler.AbstractSampler;
@@ -29,12 +28,12 @@
 import me.lucko.spark.common.sampler.SamplerType;
 import me.lucko.spark.common.sampler.window.ProfilingWindowUtils;
 import me.lucko.spark.common.tick.TickHook;
+import me.lucko.spark.common.util.SparkScheduledThreadPoolExecutor;
 import me.lucko.spark.common.util.SparkThreadFactory;
 import me.lucko.spark.common.ws.ViewerSocket;
 import me.lucko.spark.proto.SparkSamplerProtos.SamplerData;
 
 import java.util.Locale;
-import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
@@ -84,12 +83,7 @@ private AsyncSampler(SparkPlatform platform, SamplerSettings settings, SampleCol
         this.dataAggregator = dataAggregator;
         this.forceNanoTime = forceNanoTime;
         this.profilerAccess = AsyncProfilerAccess.getInstance(platform);
-        this.scheduler = Executors.newSingleThreadScheduledExecutor(
-                new ThreadFactoryBuilder()
-                        .setNameFormat("spark-async-sampler-worker-thread")
-                        .setUncaughtExceptionHandler(SparkThreadFactory.EXCEPTION_HANDLER)
-                        .build()
-        );
+        this.scheduler = new SparkScheduledThreadPoolExecutor(1, new SparkThreadFactory("spark-async-sampler-worker", false));
     }
 
     /**

spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaDataAggregator.java

@@ -27,9 +27,7 @@
 import me.lucko.spark.common.sampler.node.ThreadNode;
 
 import java.lang.management.ThreadInfo;
-import java.util.List;
 import java.util.concurrent.ExecutorService;
-import java.util.concurrent.TimeUnit;
 
 /**
  * Abstract {@link DataAggregator} for the {@link JavaSampler}.
@@ -75,19 +73,6 @@ protected void writeData(ThreadInfo threadInfo, int window) {
         }
     }
 
-    @Override
-    public List<ThreadNode> exportData() {
-        // wait for all pending data to be inserted
-        this.workerPool.shutdown();
-        try {
-            this.workerPool.awaitTermination(15, TimeUnit.SECONDS);
-        } catch (InterruptedException e) {
-            e.printStackTrace();
-        }
-
-        return super.exportData();
-    }
-
     static boolean isSleeping(ThreadInfo thread) {
         if (thread.getThreadState() == Thread.State.WAITING || thread.getThreadState() == Thread.State.TIMED_WAITING) {
             return true;

spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java

@@ -20,7 +20,6 @@
 
 package me.lucko.spark.common.sampler.java;
 
-import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import me.lucko.spark.common.SparkPlatform;
 import me.lucko.spark.common.sampler.AbstractSampler;
 import me.lucko.spark.common.sampler.SamplerMode;
@@ -30,14 +29,14 @@
 import me.lucko.spark.common.sampler.window.WindowStatisticsCollector;
 import me.lucko.spark.common.tick.TickHook;
 import me.lucko.spark.common.util.MethodDisambiguator;
+import me.lucko.spark.common.util.SparkScheduledThreadPoolExecutor;
 import me.lucko.spark.common.util.SparkThreadFactory;
 import me.lucko.spark.common.ws.ViewerSocket;
 import me.lucko.spark.proto.SparkSamplerProtos.SamplerData;
 
 import java.lang.management.ManagementFactory;
 import java.lang.management.ThreadInfo;
 import java.lang.management.ThreadMXBean;
-import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
@@ -48,15 +47,9 @@
  * A sampler implementation using Java (WarmRoast).
  */
 public class JavaSampler extends AbstractSampler implements Runnable {
-    private static final AtomicInteger THREAD_ID = new AtomicInteger(0);
 
     /** The worker pool for inserting stack nodes */
-    private final ScheduledExecutorService workerPool = Executors.newScheduledThreadPool(
-            6, new ThreadFactoryBuilder()
-                    .setNameFormat("spark-java-sampler-" + THREAD_ID.getAndIncrement() + "-%d")
-                    .setUncaughtExceptionHandler(SparkThreadFactory.EXCEPTION_HANDLER)
-                    .build()
-    );
+    private final ScheduledExecutorService workerPool = new SparkScheduledThreadPoolExecutor(6, new SparkThreadFactory("spark-java-sampler", false));
 
     /** The main sampling task */
     private ScheduledFuture<?> task;

spark-common/src/main/java/me/lucko/spark/common/util/SparkScheduledThreadPoolExecutor.java

@@ -0,0 +1,62 @@
+/*
+ * 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.util;
+
+import org.jspecify.annotations.NonNull;
+
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.ThreadFactory;
+
+/**
+ * A {@link ScheduledThreadPoolExecutor} which logs uncaught exceptions.
+ */
+public class SparkScheduledThreadPoolExecutor extends ScheduledThreadPoolExecutor {
+    public SparkScheduledThreadPoolExecutor(int corePoolSize) {
+        super(corePoolSize, new SparkThreadFactory());
+    }
+
+    public SparkScheduledThreadPoolExecutor(int corePoolSize, @NonNull ThreadFactory threadFactory) {
+        super(corePoolSize, threadFactory);
+    }
+
+    @Override
+    protected void afterExecute(Runnable r, Throwable t) {
+        super.afterExecute(r, t);
+        if (t == null && r instanceof Future<?> && ((Future<?>) r).isDone()) {
+            try {
+                ((Future<?>) r).get();
+            } catch (CancellationException ce) {
+                t = ce;
+            } catch (ExecutionException ee) {
+                t = ee.getCause();
+            } catch (InterruptedException ie) {
+                // ignore/reset
+                Thread.currentThread().interrupt();
+            }
+        }
+        if (t != null) {
+            SparkThreadFactory.EXCEPTION_HANDLER.uncaughtException(Thread.currentThread(), t);
+        }
+    }
+}

spark-common/src/main/java/me/lucko/spark/common/util/SparkThreadFactory.java

@@ -28,24 +28,30 @@
 
 public class SparkThreadFactory implements ThreadFactory {
 
-    public static final Thread.UncaughtExceptionHandler EXCEPTION_HANDLER = (t, e) -> {
-        SparkStaticLogger.log(Level.SEVERE, "Uncaught exception thrown by thread " + t.getName(), e);
+    static final Thread.UncaughtExceptionHandler EXCEPTION_HANDLER = (t, e) -> {
+        SparkStaticLogger.log(Level.SEVERE, "Uncaught exception thrown in thread " + t.getName(), e);
     };
 
     private static final AtomicInteger poolNumber = new AtomicInteger(1);
     private final AtomicInteger threadNumber = new AtomicInteger(1);
     private final String namePrefix;
+    private final boolean daemon;
 
     public SparkThreadFactory() {
-        this.namePrefix = "spark-worker-pool-" +
+        this("spark-worker-pool", true);
+    }
+
+    public SparkThreadFactory(String prefix, boolean daemon) {
+        this.namePrefix = prefix + "-" +
                 poolNumber.getAndIncrement() +
                 "-thread-";
+        this.daemon = daemon;
     }
 
     public Thread newThread(Runnable r) {
         Thread t = new Thread(r, this.namePrefix + this.threadNumber.getAndIncrement());
         t.setUncaughtExceptionHandler(EXCEPTION_HANDLER);
-        t.setDaemon(true);
+        t.setDaemon(this.daemon);
         return t;
     }
 

spark-common/src/test/java/me/lucko/spark/test/plugin/TestSparkPlugin.java

@@ -24,14 +24,14 @@
 import me.lucko.spark.common.SparkPlugin;
 import me.lucko.spark.common.command.sender.CommandSender;
 import me.lucko.spark.common.platform.PlatformInfo;
+import me.lucko.spark.common.util.SparkScheduledThreadPoolExecutor;
 import me.lucko.spark.common.util.classfinder.ClassFinder;
 import me.lucko.spark.common.util.classfinder.FallbackClassFinder;
 
 import java.nio.file.Path;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -40,7 +40,7 @@
 public class TestSparkPlugin implements SparkPlugin, AutoCloseable {
 
     private static final Logger LOGGER = Logger.getLogger("spark-test");
-    private static final ScheduledExecutorService EXECUTOR_SERVICE = Executors.newScheduledThreadPool(16);
+    private static final ScheduledExecutorService EXECUTOR_SERVICE = new SparkScheduledThreadPoolExecutor(16);
 
     private final Path directory;
     private final Map<String, String> props;

spark-fabric/src/main/java/me/lucko/spark/fabric/placeholder/SparkFabricPlaceholderApi.java

@@ -23,7 +23,6 @@
 import com.google.gson.JsonParseException;
 import com.mojang.serialization.JsonOps;
 import eu.pb4.placeholders.api.PlaceholderContext;
-import eu.pb4.placeholders.api.PlaceholderHandler;
 import eu.pb4.placeholders.api.PlaceholderResult;
 import eu.pb4.placeholders.api.Placeholders;
 import me.lucko.spark.common.SparkPlatform;

spark-minecraft/src/main/java/me/lucko/spark/minecraft/plugin/MinecraftSparkPlugin.java

@@ -33,15 +33,14 @@
 import me.lucko.spark.common.SparkPlatform;
 import me.lucko.spark.common.SparkPlugin;
 import me.lucko.spark.common.command.sender.CommandSender;
-import me.lucko.spark.common.util.SparkThreadFactory;
+import me.lucko.spark.common.util.SparkScheduledThreadPoolExecutor;
 import me.lucko.spark.minecraft.SparkMinecraftMod;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
 import java.nio.file.Path;
 import java.util.Arrays;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.logging.Level;
 
@@ -56,7 +55,7 @@ public abstract class MinecraftSparkPlugin<M extends SparkMinecraftMod> implemen
     protected MinecraftSparkPlugin(M mod) {
         this.mod = mod;
         this.logger = LogManager.getLogger("spark");
-        this.scheduler = Executors.newScheduledThreadPool(4, new SparkThreadFactory());
+        this.scheduler = new SparkScheduledThreadPoolExecutor(4);
     }
 
     public void enable() {