StatsTask Cancellation Fix

LavalinkServer/src/main/java/lavalink/server/io/SocketContext.java

@@ -41,10 +41,7 @@
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.*;
 
 public class SocketContext {
 
@@ -59,6 +56,7 @@ public class SocketContext {
     private final Map<Integer, Core> cores = new HashMap<>();
     private final Map<String, Player> players = new ConcurrentHashMap<>();
     private ScheduledExecutorService statsExecutor;
+    private ScheduledFuture<?> statsFuture;
     public final ScheduledExecutorService playerUpdateService;
     private final ConcurrentHashMap<Integer, IAudioSendFactory> sendFactories = new ConcurrentHashMap<>();
 
@@ -73,7 +71,7 @@ public class SocketContext {
         this.shardCount = shardCount;
 
         statsExecutor = Executors.newSingleThreadScheduledExecutor();
-        statsExecutor.scheduleAtFixedRate(new StatsTask(this, socketServer), 0, 1, TimeUnit.MINUTES);
+        this.statsFuture = statsExecutor.scheduleAtFixedRate(new StatsTask(this, socketServer), 0, 1, TimeUnit.MINUTES);
 
         playerUpdateService = Executors.newScheduledThreadPool(2, r -> {
             Thread thread = new Thread(r);
@@ -100,6 +98,10 @@ Player getPlayer(String guildId) {
         );
     }
 
+    ScheduledFuture<?> getStatsFuture() {
+        return statsFuture;
+    }
+
     int getShardCount() {
         return shardCount;
     }
@@ -133,8 +135,9 @@ void shutdown() {
                 }
             }
         });
-
         players.values().forEach(Player::stop);
+        if (!statsFuture.isCancelled())
+            statsFuture.cancel(true);
     }
 
     private IAudioSendFactory getAudioSendFactory(int shardId) {

LavalinkServer/src/main/java/lavalink/server/io/SocketServer.java

@@ -128,6 +128,8 @@ private void close(WebSocket webSocket, int code, String reason) {
             log.info("Connection closed from {} with protocol {} with reason {} with code {}",
                     webSocket.getRemoteSocketAddress().toString(), webSocket.getDraft(), reason, code);
             context.shutdown();
+            if (!context.getStatsFuture().isCancelled())
+                context.getStatsFuture().cancel(true);
         }
     }
 
@@ -225,6 +227,8 @@ public void onStart() {
     }
 
     public static void sendPlayerUpdate(WebSocket webSocket, Player player) {
+        if (webSocket.isClosed())
+            return;
         JSONObject json = new JSONObject();
         json.put("op", "playerUpdate");
         json.put("guildId", player.getGuildId());

LavalinkServer/src/main/java/lavalink/server/io/StatsTask.java

@@ -116,7 +116,8 @@ public void sendStats() {
             out.put("frameStats", frames);
         }
 
-        context.getSocket().send(out.toString());
+        if (context.getSocket().isOpen())
+            context.getSocket().send(out.toString());
     }
 
     private double uptime = 0;

LavalinkServer/src/main/java/lavalink/server/player/AudioLoader.java

@@ -31,7 +31,6 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 
 public class AudioLoader implements AudioLoadResultHandler {
@@ -43,6 +42,7 @@ public class AudioLoader implements AudioLoadResultHandler {
     private boolean isPlaylist = false;
     private String playlistName = null;
     private Integer selectedTrack = null;
+    private ResultStatus status = ResultStatus.UNKNOWN;
     private boolean used = false;
 
     public AudioLoader(AudioPlayerManager audioPlayerManager) {
@@ -61,13 +61,14 @@ LoadResult loadSync(String identifier) throws InterruptedException {
             this.wait();
         }
 
-        return new LoadResult(loadedItems, isPlaylist, playlistName, selectedTrack);
+        return new LoadResult(loadedItems, isPlaylist, playlistName, status, selectedTrack);
     }
 
     @Override
     public void trackLoaded(AudioTrack audioTrack) {
         loadedItems = new ArrayList<>();
         loadedItems.add(audioTrack);
+        status = ResultStatus.TRACK_LOADED;
         log.info("Loaded track " + audioTrack.getInfo().title);
         synchronized (this) {
             this.notify();
@@ -83,6 +84,7 @@ public void playlistLoaded(AudioPlaylist audioPlaylist) {
         }
 
         log.info("Loaded playlist " + audioPlaylist.getName());
+        status = audioPlaylist.isSearchResult() ? ResultStatus.SEARCH_RESULT : ResultStatus.PLAYLIST_LOADED;
         loadedItems = audioPlaylist.getTracks();
         synchronized (this) {
             this.notify();
@@ -92,6 +94,7 @@ public void playlistLoaded(AudioPlaylist audioPlaylist) {
     @Override
     public void noMatches() {
         log.info("No matches found");
+        status = ResultStatus.NO_MATCHES;
         loadedItems = new ArrayList<>();
         synchronized (this) {
             this.notify();
@@ -101,24 +104,11 @@ public void noMatches() {
     @Override
     public void loadFailed(FriendlyException e) {
         log.error("Load failed", e);
+        status = ResultStatus.LOAD_FAILED;
         loadedItems = new ArrayList<>();
         synchronized (this) {
             this.notify();
         }
     }
 
-}
-
-class LoadResult {
-    public List<AudioTrack> tracks;
-    public boolean isPlaylist;
-    public String playlistName;
-    public Integer selectedTrack;
-
-    public LoadResult(List<AudioTrack> tracks, boolean isPlaylist, String playlistName, Integer selectedTrack) {
-        this.tracks = Collections.unmodifiableList(tracks);
-        this.isPlaylist = isPlaylist;
-        this.playlistName = playlistName;
-        this.selectedTrack = selectedTrack;
-    }
-}
+}

LavalinkServer/src/main/java/lavalink/server/player/AudioLoaderRestHandler.java

@@ -119,6 +119,7 @@ public String getLoadTracks(HttpServletRequest request, HttpServletResponse resp
 
         json.put("isPlaylist", result.isPlaylist);
         json.put("playlistInfo", playlist);
+        json.put("loadType", result.loadResultType);
         json.put("tracks", tracks);
 
         return json.toString();

LavalinkServer/src/main/java/lavalink/server/player/LoadResult.java

@@ -0,0 +1,22 @@
+package lavalink.server.player;
+
+import com.sedmelluq.discord.lavaplayer.track.AudioTrack;
+
+import java.util.Collections;
+import java.util.List;
+
+class LoadResult {
+    public List<AudioTrack> tracks;
+    public boolean isPlaylist;
+    public String playlistName;
+    public String loadResultType;
+    public Integer selectedTrack;
+
+    LoadResult(List<AudioTrack> tracks, boolean isPlaylist, String playlistName, ResultStatus loadResultType, Integer selectedTrack) {
+        this.tracks = Collections.unmodifiableList(tracks);
+        this.isPlaylist = isPlaylist;
+        this.playlistName = playlistName;
+        this.loadResultType = loadResultType.name();
+        this.selectedTrack = selectedTrack;
+    }
+}

LavalinkServer/src/main/java/lavalink/server/player/ResultStatus.java

@@ -0,0 +1,10 @@
+package lavalink.server.player;
+
+public enum ResultStatus {
+    TRACK_LOADED,
+    PLAYLIST_LOADED,
+    SEARCH_RESULT,
+    NO_MATCHES,
+    LOAD_FAILED,
+    UNKNOWN
+}