perf: batch channel fetching and prefetch earlier

src/youtube.rs

@@ -9,7 +9,7 @@ use std::{
 
 use serde::{Deserialize, Serialize};
 
-pub const PAGE_SIZE: usize = 20;
+pub const PAGE_SIZE: usize = 60;
 
 pub struct Page {
     pub videos: Vec<Video>,
@@ -222,10 +222,23 @@ fn page_end(start: usize) -> io::Result<usize> {
         .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "Invalid playlist page offset"))
 }
 
-// Own the extractor so replacing a request or quitting also cancels its process.
+// Own every extractor so replacing a request or quitting cancels the whole page.
 pub struct Job {
-    child: Child,
+    children: Vec<Extraction>,
+}
+
+struct Extraction {
+    child: Option<Child>,
     output: Receiver<io::Result<(String, String)>>,
+    status: Option<std::process::ExitStatus>,
+    completed: Option<String>,
+}
+
+fn channel_batches(channels: &[String]) -> impl Iterator<Item = &[String]> {
+    let count = channels.len().min(4);
+    (0..count).map(move |index| {
+        &channels[index * channels.len() / count..(index + 1) * channels.len() / count]
+    })
 }
 
 impl Job {
@@ -252,22 +265,26 @@ impl Job {
             ));
         }
         let end = page_end(start)?;
+        let start = start.to_string();
+        let end = end.to_string();
         // ponytail: offsets traverse prior continuations; use direct tokens/per-channel exhaustion
         // tracking if deep archives become slow.
-        Self::spawn(
-            extractor()
+        Self::spawn_all(channel_batches(channels).map(|batch| {
+            let mut command = extractor();
+            command
                 .args([
                     "--flat-playlist",
                     "--playlist-start",
-                    &start.to_string(),
+                    &start,
                     "--playlist-end",
-                    &end.to_string(),
+                    &end,
                     "--extractor-args",
                     "youtubetab:approximate_date",
                     "--dump-single-json",
                 ])
-                .args(channels.iter().map(|channel| format!("{channel}/videos"))),
-        )
+                .args(batch.iter().map(|channel| format!("{channel}/videos")));
+            command
+        }))
     }
 
     pub fn channel(id: &str) -> io::Result<Self> {
@@ -295,6 +312,64 @@ impl Job {
         ]))
     }
 
+    fn spawn(command: &mut Command) -> io::Result<Self> {
+        Ok(Self {
+            children: vec![Extraction::spawn(command)?],
+        })
+    }
+
+    fn spawn_all(commands: impl IntoIterator<Item = Command>) -> io::Result<Self> {
+        // Collect drops already-started children if a later spawn fails.
+        let children = commands
+            .into_iter()
+            .map(|mut command| Extraction::spawn(&mut command))
+            .collect::<io::Result<Vec<_>>>()?;
+        Ok(Self { children })
+    }
+
+    pub fn poll(&mut self) -> Option<Result<String, String>> {
+        if self.children.is_empty() {
+            return None;
+        }
+        let mut pending = false;
+        for child in &mut self.children {
+            if child.completed.is_some() {
+                continue;
+            }
+            match child.poll() {
+                None => pending = true,
+                Some(Ok(output)) => child.completed = Some(output),
+                Some(Err(error)) => {
+                    self.children.clear();
+                    return Some(Err(error));
+                }
+            }
+        }
+        if pending {
+            return None;
+        }
+        let mut children = self.children.drain(..);
+        // Keep the first buffer; only subsequent batches need to be appended.
+        let mut output = children.next()?.completed.take().unwrap();
+        output.reserve(
+            children
+                .as_slice()
+                .iter()
+                .map(|child| child.completed.as_ref().unwrap().len() + 1)
+                .sum(),
+        );
+        for mut child in children {
+            let next = child.completed.take().unwrap();
+            if !output.is_empty() && !next.is_empty() {
+                output.push('\n');
+            }
+            output.push_str(&next);
+        }
+        Some(Ok(output))
+    }
+}
+
+impl Extraction {
     fn spawn(command: &mut Command) -> io::Result<Self> {
         let mut child = command
             .stdin(Stdio::null())
@@ -321,16 +396,27 @@ impl Job {
                 .and_then(|text| stderr_result.map(|_| (text, errors)));
             let _ = sender.send(result);
         });
-        Ok(Self { child, output })
+        Ok(Self {
+            child: Some(child),
+            output,
+            status: None,
+            completed: None,
+        })
     }
 
-    pub fn poll(&mut self) -> Option<Result<String, String>> {
-        let status = match self.child.try_wait() {
-            Ok(Some(status)) => status,
-            Ok(None) => return None,
-            Err(error) => return Some(Err(format!("Cannot wait for yt-dlp: {error}"))),
-        };
-        match self.output.try_recv() {
+    fn poll(&mut self) -> Option<Result<String, String>> {
+        if self.status.is_none() {
+            self.status = match self.child.as_mut()?.try_wait() {
+                Ok(status) => status,
+                Err(error) => return Some(Err(format!("Cannot wait for yt-dlp: {error}"))),
+            };
+        }
+        let status = self.status?;
+        if !status.success() {
+            // Do not wait for inherited pipes held open by a failed child's descendants.
+            self.cancel();
+        }
+        let result = match self.output.try_recv() {
             Ok(Ok((output, _))) if status.success() => Some(Ok(output)),
             Ok(Ok((_, errors))) => Some(Err(if errors.trim().is_empty() {
                 format!("yt-dlp exited with {status}")
@@ -342,13 +428,29 @@ impl Job {
             Err(TryRecvError::Disconnected) => {
                 Some(Err("yt-dlp output reader disconnected".into()))
             }
+        };
+        if result.is_some() {
+            self.cancel();
+        }
+        result
+    }
+
+    fn cancel(&mut self) {
+        if let Some(mut child) = self.child.take() {
+            // The leader may have exited while descendants still own its output pipes.
+            // SAFETY: every extraction owns an isolated process group.
+            #[cfg(unix)]
+            unsafe {
+                libc::kill(-(child.id() as i32), libc::SIGKILL);
+            }
+            stop(&mut child);
         }
     }
 }
 
-impl Drop for Job {
+impl Drop for Extraction {
     fn drop(&mut self) {
-        stop(&mut self.child);
+        self.cancel();
     }
 }
 
@@ -554,6 +656,243 @@ impl Drop for PlayerLaunch {
 mod tests {
     use super::*;
 
+    #[test]
+    fn channel_batches_are_bounded_balanced_and_preserve_every_requested_url() {
+        for count in 0..=17 {
+            let channels: Vec<_> = (0..count).map(|index| format!("channel-{index}")).collect();
+            let batches: Vec<_> = channel_batches(&channels).collect();
+            assert_eq!(batches.len(), count.min(4));
+            assert!(batches.iter().all(|batch| !batch.is_empty()));
+            assert_eq!(batches.concat(), channels);
+            if !batches.is_empty() {
+                let smallest = batches.iter().map(|batch| batch.len()).min().unwrap();
+                let largest = batches.iter().map(|batch| batch.len()).max().unwrap();
+                assert!(largest - smallest <= 1);
+            }
+        }
+    }
+
+    #[cfg(unix)]
+    mod subprocesses {
+        use super::*;
+
+        fn controlled(script: &str) -> Command {
+            let mut command = isolated_command(OsStr::new("/bin/sh"));
+            command.args(["-c", &format!("kill -STOP $$; {script}")]);
+            command
+        }
+
+        fn wait_stopped(extraction: &Extraction) {
+            wait_stopped_pid(extraction.child.as_ref().unwrap().id());
+        }
+
+        fn wait_stopped_pid(pid: u32) {
+            let pid = pid as i32;
+            let mut status = 0;
+            loop {
+                // SAFETY: wait only for this test's owned child, without reaping its exit.
+                let result = unsafe { libc::waitpid(pid, &mut status, libc::WUNTRACED) };
+                if result == -1 && io::Error::last_os_error().kind() == io::ErrorKind::Interrupted {
+                    continue;
+                }
+                assert_eq!(result, pid);
+                assert!(libc::WIFSTOPPED(status));
+                break;
+            }
+        }
+
+        fn release(extraction: &mut Extraction) {
+            let child = extraction.child.as_mut().unwrap();
+            // SAFETY: this child is stopped at the deterministic test barrier.
+            assert_eq!(unsafe { libc::kill(child.id() as i32, libc::SIGCONT) }, 0);
+            child.wait().unwrap();
+            // Synchronize pipe draining too, then leave the real result for Job::poll.
+            // Blocking is confined to tests; production polling never waits for output.
+            let result = extraction.output.recv().unwrap();
+            let (sender, receiver) = mpsc::channel();
+            sender.send(result).unwrap();
+            extraction.output = receiver;
+        }
+
+        fn take_output(extraction: &mut Extraction) -> Receiver<io::Result<(String, String)>> {
+            let (_, receiver) = mpsc::channel();
+            std::mem::replace(&mut extraction.output, receiver)
+        }
+
+        fn assert_reaped(pid: u32) {
+            let mut status = 0;
+            // SAFETY: query only this test's child; ECHILD proves cancellation reaped it.
+            assert_eq!(
+                unsafe { libc::waitpid(pid as i32, &mut status, libc::WNOHANG) },
+                -1
+            );
+            assert_eq!(
+                io::Error::last_os_error().raw_os_error(),
+                Some(libc::ECHILD)
+            );
+        }
+
+        #[test]
+        fn a_later_spawn_failure_cancels_and_reaps_already_started_batches() {
+            static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
+            let directory = std::env::temp_dir().join(format!(
+                "tubular-extractor-spawn-{}-{}",
+                std::process::id(),
+                NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
+            ));
+            std::fs::create_dir(&directory).unwrap();
+            let ready = directory.join("ready");
+            let mut first = isolated_command(OsStr::new("/bin/sh"));
+            first
+                .args(["-c", "printf '%s' $$ > \"$1\"; kill -STOP $$", "test"])
+                .arg(&ready);
+            let mut started_pid = None;
+            let commands = std::iter::once(first).chain(std::iter::once_with(|| {
+                // The file only communicates identity; SIGSTOP is the completion barrier.
+                let pid = loop {
+                    if let Ok(text) = std::fs::read_to_string(&ready)
+                        && let Ok(pid) = text.parse::<u32>()
+                    {
+                        break pid;
+                    }
+                    thread::yield_now();
+                };
+                wait_stopped_pid(pid);
+                started_pid = Some(pid);
+                isolated_command(directory.join("missing-extractor").as_os_str())
+            }));
+            let error = Job::spawn_all(commands).err().unwrap();
+            assert_eq!(error.kind(), io::ErrorKind::NotFound);
+            assert_reaped(started_pid.unwrap());
+            std::fs::remove_dir_all(directory).unwrap();
+        }
+
+        #[test]
+        fn aggregate_waits_for_all_batches_and_merges_in_request_not_completion_order() {
+            let mut job = Job::spawn_all([
+                controlled("printf '%s' '{\"batch\":1}'"),
+                controlled("printf '%s' '{\"batch\":2}'"),
+                controlled("printf '%s' '{\"batch\":3}'"),
+            ])
+            .unwrap();
+            for child in &job.children {
+                wait_stopped(child);
+            }
+            assert!(job.poll().is_none());
+            release(&mut job.children[2]);
+            assert!(
+                job.poll().is_none(),
+                "the last batch cannot complete a page"
+            );
+            release(&mut job.children[1]);
+            assert!(job.poll().is_none(), "the first batch is still outstanding");
+            release(&mut job.children[0]);
+            let output = job.poll().unwrap().unwrap();
+            let batches: Vec<_> = serde_json::Deserializer::from_str(&output)
+                .into_iter::<serde_json::Value>()
+                .map(Result::unwrap)
+                .collect();
+            assert_eq!(
+                batches,
+                [
+                    serde_json::json!({"batch":1}),
+                    serde_json::json!({"batch":2}),
+                    serde_json::json!({"batch":3}),
+                ]
+            );
+            assert!(job.poll().is_none());
+        }
+
+        #[test]
+        fn a_failed_batch_discards_completed_output_and_cancels_pending_siblings() {
+            let mut job = Job::spawn_all([
+                controlled("printf '%s' '{\"complete\":true}'"),
+                controlled(
+                    "printf '%s' '{\"partial\":true}'; printf '%s' 'batch failed' >&2; exit 7",
+                ),
+                controlled("/bin/sh -c 'kill -STOP $$' & wait"),
+            ])
+            .unwrap();
+            for child in &job.children {
+                wait_stopped(child);
+            }
+            let pending_pid = job.children[2].child.as_ref().unwrap().id();
+            let pending_output = take_output(&mut job.children[2]);
+            release(&mut job.children[0]);
+            assert!(job.poll().is_none());
+            release(&mut job.children[1]);
+            assert_eq!(job.poll().unwrap().unwrap_err(), "batch failed");
+            assert!(pending_output.recv().unwrap().is_ok());
+            assert_reaped(pending_pid);
+            assert!(job.poll().is_none());
+        }
+
+        #[test]
+        fn unreadable_batch_output_fails_the_whole_page_and_cancels_siblings() {
+            let mut job = Job::spawn_all([
+                controlled("printf '\\377'"),
+                controlled("printf '%s' 'must not finish'"),
+            ])
+            .unwrap();
+            for child in &job.children {
+                wait_stopped(child);
+            }
+            let pending_pid = job.children[1].child.as_ref().unwrap().id();
+            let pending_output = take_output(&mut job.children[1]);
+            release(&mut job.children[0]);
+            assert!(
+                job.poll()
+                    .unwrap()
+                    .unwrap_err()
+                    .starts_with("Cannot read yt-dlp output:")
+            );
+            assert!(pending_output.recv().unwrap().is_ok());
+            assert_reaped(pending_pid);
+        }
+
+        #[test]
+        fn dropping_a_page_cancels_and_reaps_every_batch() {
+            let mut job = Job::spawn_all([
+                controlled("/bin/sh -c 'kill -STOP $$' & wait"),
+                controlled("printf '%s' 'must not finish'"),
+            ])
+            .unwrap();
+            for child in &job.children {
+                wait_stopped(child);
+            }
+            let pids: Vec<_> = job
+                .children
+                .iter()
+                .map(|child| child.child.as_ref().unwrap().id())
+                .collect();
+            let outputs: Vec<_> = job.children.iter_mut().map(take_output).collect();
+            drop(job);
+            for output in outputs {
+                assert!(output.recv().unwrap().is_ok());
+            }
+            for pid in pids {
+                assert_reaped(pid);
+            }
+        }
+
+        #[test]
+        fn dropping_after_the_leader_exits_still_cancels_inherited_pipe_owners() {
+            let mut job =
+                Job::spawn_all([controlled("/bin/sh -c 'kill -STOP $$' & exit 0")]).unwrap();
+            wait_stopped(&job.children[0]);
+            let child = job.children[0].child.as_mut().unwrap();
+            let pid = child.id();
+            // SAFETY: resume this test's stopped leader, which exits without its descendant.
+            assert_eq!(unsafe { libc::kill(pid as i32, libc::SIGCONT) }, 0);
+            assert!(child.wait().unwrap().success());
+            let output = take_output(&mut job.children[0]);
+            drop(job);
+            // EOF proves the descendant that inherited stdout/stderr was also canceled.
+            assert!(output.recv().unwrap().is_ok());
+            assert_reaped(pid);
+        }
+    }
+
     #[test]
     fn player_launch_rejects_unsafe_video_ids_before_spawning() {
         for id in ["", "--help", "jNQXAC9IVRw&list=other", "日本語"] {