Add pagination to spotify and thread count lock for ytdl

Add pagination to the spotify wrapper to retreive all
songs in a playlist.
Add a thread count lock to the youtube-dl process to avoid
starting too many at once.

Signed-off-by: trivernis <trivernis@protonmail.com>

src/commands/music/mod.rs

@@ -157,7 +157,8 @@ async fn play_next_in_queue(
 /// Returns the list of songs for a given url
 async fn get_songs_for_query(ctx: &Context, msg: &Message, query: &str) -> BotResult<Vec<Song>> {
     lazy_static::lazy_static! {
-        static ref YOUTUBE_URL_REGEX: Regex = Regex::new(r"^(https?(://))?(www\.)?(youtube\.com/watch\?.*v=.*)|(/youtu.be/.*)$").unwrap();
+        // expressions to determine the type of url
+        static ref YOUTUBE_URL_REGEX: Regex = Regex::new(r"^(https?(://))?(www\.)?(youtube\.com/watch\?.*v=.*)|(/youtu.be/.*)|(youtube\.com/playlist\?.*list=.*)$").unwrap();
         static ref SPOTIFY_PLAYLIST_REGEX: Regex = Regex::new(r"^(https?(://))?(www\.|open\.)?spotify\.com/playlist/.*").unwrap();
         static ref SPOTIFY_ALBUM_REGEX: Regex = Regex::new(r"^(https?(://))?(www\.|open\.)?spotify\.com/album/.*").unwrap();
         static ref SPOTIFY_SONG_REGEX: Regex = Regex::new(r"^(https?(://))?(www\.|open\.)?spotify\.com/track/.*").unwrap();
@@ -167,12 +168,14 @@ async fn get_songs_for_query(ctx: &Context, msg: &Message, query: &str) -> BotRe
     let store = data.get::<Store>().unwrap();
 
     if YOUTUBE_URL_REGEX.is_match(query) {
+        // try fetching the url as a playlist
         songs = get_videos_for_playlist(query)
             .await?
             .into_iter()
             .map(Song::from)
             .collect();
 
+        // if no songs were found fetch the song as a video
         if songs.len() == 0 {
             let song: Song = get_video_information(query).await?.into();
             added_one_msg(&ctx, msg, &song).await?;
@@ -181,14 +184,17 @@ async fn get_songs_for_query(ctx: &Context, msg: &Message, query: &str) -> BotRe
             added_multiple_msg(&ctx, msg, &mut songs).await?;
         }
     } else if SPOTIFY_PLAYLIST_REGEX.is_match(query) {
+        // search for all songs in the playlist and search for them on youtube
         let song_names = store.spotify_api.get_songs_in_playlist(query).await?;
         songs = parallel_search_youtube(song_names).await;
         added_multiple_msg(&ctx, msg, &mut songs).await?;
     } else if SPOTIFY_ALBUM_REGEX.is_match(query) {
+        // fetch all songs in the album and search for them on youtube
         let song_names = store.spotify_api.get_songs_in_album(query).await?;
         songs = parallel_search_youtube(song_names).await;
         added_multiple_msg(&ctx, msg, &mut songs).await?;
     } else if SPOTIFY_SONG_REGEX.is_match(query) {
+        // fetch the song name and search it on youtube
         let name = store.spotify_api.get_song_name(query).await?;
         let song: Song = search_video_information(name.clone())
             .await?

src/providers/music/mod.rs

@@ -1,6 +1,9 @@
 use crate::providers::music::responses::{PlaylistEntry, VideoInformation};
 use crate::utils::error::BotResult;
 use std::process::Stdio;
+use std::sync::atomic::{AtomicU8, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
 use tokio::io::AsyncReadExt;
 use tokio::process::Command;
 
@@ -8,6 +11,8 @@ pub(crate) mod queue;
 pub(crate) mod responses;
 pub(crate) mod spotify;
 
+static THREAD_LIMIT: u8 = 64;
+
 /// Returns a list of youtube videos for a given url
 pub(crate) async fn get_videos_for_playlist(url: &str) -> BotResult<Vec<PlaylistEntry>> {
     let output =
@@ -45,13 +50,34 @@ pub(crate) async fn search_video_information(query: String) -> BotResult<Option<
 }
 
 /// Executes youtube-dl asynchronously
+/// An atomic U8 is used to control the number of parallel processes
+/// to avoid using too much memory
 async fn youtube_dl(args: &[&str]) -> BotResult<String> {
+    lazy_static::lazy_static! { static ref THREAD_LOCK: Arc<AtomicU8> = Arc::new(AtomicU8::new(0)); }
+
+    while THREAD_LOCK.load(Ordering::SeqCst) >= THREAD_LIMIT {
+        tokio::time::sleep(Duration::from_millis(100)).await;
+    }
+    THREAD_LOCK.fetch_add(1, Ordering::Relaxed);
+
     let ytdl = Command::new("youtube-dl")
         .args(args)
         .stdout(Stdio::piped())
-        .spawn()?;
+        .spawn()
+        .map_err(|e| {
+            THREAD_LOCK.fetch_sub(1, Ordering::Relaxed);
+            e
+        })?;
     let mut output = String::new();
-    ytdl.stdout.unwrap().read_to_string(&mut output).await?;
+    ytdl.stdout
+        .unwrap()
+        .read_to_string(&mut output)
+        .await
+        .map_err(|e| {
+            THREAD_LOCK.fetch_sub(1, Ordering::Relaxed);
+            e
+        })?;
+    THREAD_LOCK.fetch_sub(1, Ordering::Relaxed);
 
     Ok(output)
 }

src/providers/music/spotify.rs

@@ -1,5 +1,5 @@
 use crate::utils::error::{BotError, BotResult};
-use aspotify::{Client, ClientCredentials, PlaylistItemType};
+use aspotify::{Client, ClientCredentials, PlaylistItem, PlaylistItemType};
 
 pub struct SpotifyApi {
     client: Client,
@@ -21,10 +21,19 @@ impl SpotifyApi {
     /// Returns the song names for a playlist
     pub async fn get_songs_in_playlist(&self, url: &str) -> BotResult<Vec<String>> {
         let id = self.get_id_for_url(url)?;
-        let playlist = self.client.playlists().get_playlist(&*id, None).await?.data;
+        let mut playlist_tracks = Vec::new();
-        let song_names = playlist
+        let mut offset = 0;
-            .tracks
+
-            .items
+        loop {
+            let mut tracks = self.get_tracks_in_playlist(&*id, 100, offset).await?;
+            if tracks.len() == 0 {
+                break;
+            }
+            playlist_tracks.append(&mut tracks);
+            offset += 100;
+        }
+
+        let song_names = playlist_tracks
             .into_iter()
             .filter_map(|item| item.item)
             .map(|t| match t {
@@ -34,7 +43,7 @@ impl SpotifyApi {
                         .into_iter()
                         .map(|a| a.name)
                         .collect::<Vec<String>>()
-                        .join(","),
+                        .join(" & "),
                     t.name
                 ),
                 PlaylistItemType::Episode(e) => e.name,
@@ -44,6 +53,23 @@ impl SpotifyApi {
         Ok(song_names)
     }
 
+    /// Returns the tracks of a playlist with pagination
+    async fn get_tracks_in_playlist(
+        &self,
+        id: &str,
+        limit: usize,
+        offset: usize,
+    ) -> BotResult<Vec<PlaylistItem>> {
+        let tracks = self
+            .client
+            .playlists()
+            .get_playlists_items(id, limit, offset, None)
+            .await?
+            .data;
+
+        Ok(tracks.items)
+    }
+
     /// Returns all song names for a given album
     pub async fn get_songs_in_album(&self, url: &str) -> BotResult<Vec<String>> {
         let id = self.get_id_for_url(url)?;
@@ -59,7 +85,7 @@ impl SpotifyApi {
                         .into_iter()
                         .map(|a| a.name)
                         .collect::<Vec<String>>()
-                        .join(","),
+                        .join(" & "),
                     item.name
                 )
             })
@@ -80,7 +106,7 @@ impl SpotifyApi {
                 .into_iter()
                 .map(|a| a.name)
                 .collect::<Vec<String>>()
-                .join(","),
+                .join(" & "),
             track.name
         ))
     }