Fix race condition in AudioLoader and adjust AudioLoaderRestHandler's usage of it

Author: schnapster

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

@@ -31,70 +31,61 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 public class AudioLoader implements AudioLoadResultHandler {
 
     private static final Logger log = LoggerFactory.getLogger(AudioLoader.class);
     private final AudioPlayerManager audioPlayerManager;
 
-    private List<AudioTrack> loadedItems;
-    private boolean used = false;
+    private final CompletableFuture<List<AudioTrack>> loadedItems;
+    private final AtomicBoolean used = new AtomicBoolean(false);
 
     public AudioLoader(AudioPlayerManager audioPlayerManager) {
         this.audioPlayerManager = audioPlayerManager;
+        this.loadedItems = new CompletableFuture<>();
     }
 
-    List<AudioTrack> loadSync(String identifier) throws InterruptedException {
-        if(used)
+    public CompletionStage<List<AudioTrack>> load(String identifier) {
+        boolean isUsed = this.used.getAndSet(true);
+        if (isUsed) {
             throw new IllegalStateException("This loader can only be used once per instance");
-
-        used = true;
-
-        audioPlayerManager.loadItem(identifier, this);
-
-        synchronized (this) {
-            this.wait();
         }
 
+        log.trace("Loading item with identifier {}", identifier);
+        this.audioPlayerManager.loadItem(identifier, this);
+
         return loadedItems;
     }
 
     @Override
     public void trackLoaded(AudioTrack audioTrack) {
-        loadedItems = new ArrayList<>();
-        loadedItems.add(audioTrack);
         log.info("Loaded track " + audioTrack.getInfo().title);
-        synchronized (this) {
-            this.notify();
-        }
+        ArrayList<AudioTrack> result = new ArrayList<>();
+        result.add(audioTrack);
+        this.loadedItems.complete(result);
     }
 
     @Override
     public void playlistLoaded(AudioPlaylist audioPlaylist) {
         log.info("Loaded playlist " + audioPlaylist.getName());
-        loadedItems = audioPlaylist.getTracks();
-        synchronized (this) {
-            this.notify();
-        }
+        this.loadedItems.complete(audioPlaylist.getTracks());
     }
 
     @Override
     public void noMatches() {
         log.info("No matches found");
-        loadedItems = new ArrayList<>();
-        synchronized (this) {
-            this.notify();
-        }
+        this.loadedItems.complete(Collections.emptyList());
     }
 
     @Override
     public void loadFailed(FriendlyException e) {
         log.error("Load failed", e);
-        loadedItems = new ArrayList<>();
-        synchronized (this) {
-            this.notify();
-        }
+        this.loadedItems.complete(Collections.emptyList());
     }
 
 }

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

@@ -45,6 +45,8 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
 
 @RestController
 public class AudioLoaderRestHandler {
@@ -91,21 +93,9 @@ private JSONObject trackToJSON(AudioTrack audioTrack) {
                 .put("position", audioTrack.getPosition());
     }
 
-    @GetMapping(value = "/loadtracks", produces = "application/json")
-    @ResponseBody
-    public ResponseEntity<String> getLoadTracks(HttpServletRequest request, HttpServletResponse response, @RequestParam String identifier)
-            throws IOException, InterruptedException {
-        log(request);
-
-        Optional<ResponseEntity<String>> notAuthed = checkAuthorization(request);
-        if (notAuthed.isPresent()) {
-            return notAuthed.get();
-        }
-
+    private JSONArray encodeTrackList(List<AudioTrack> trackList) {
         JSONArray tracks = new JSONArray();
-        List<AudioTrack> list = new AudioLoader(audioPlayerManager).loadSync(identifier);
-
-        list.forEach(track -> {
+        trackList.forEach(track -> {
             JSONObject object = new JSONObject();
             object.put("info", trackToJSON(track));
 
@@ -114,11 +104,25 @@ public ResponseEntity<String> getLoadTracks(HttpServletRequest request, HttpServ
                 object.put("track", encoded);
                 tracks.put(object);
             } catch (IOException e) {
-                throw new RuntimeException();
+                log.warn("Failed to encode a track {}, skipping", track.getIdentifier(), e);
             }
         });
+        return tracks;
+    }
+
+    @GetMapping(value = "/loadtracks", produces = "application/json")
+    @ResponseBody
+    public CompletionStage<ResponseEntity<String>> getLoadTracks(HttpServletRequest request, @RequestParam String identifier) {
+        log(request);
+
+        Optional<ResponseEntity<String>> notAuthed = checkAuthorization(request);
+        if (notAuthed.isPresent()) {
+            return CompletableFuture.completedFuture(notAuthed.get());
+        }
 
-        return new ResponseEntity<>(tracks.toString(), HttpStatus.OK);
+        return new AudioLoader(audioPlayerManager).load(identifier)
+                .thenApply(this::encodeTrackList)
+                .thenApply(tracksArray -> new ResponseEntity<>(tracksArray.toString(), HttpStatus.OK));
     }
 
     @GetMapping(value = "/decodetrack", produces = "application/json")