diff --git a/crates/x-media/src/site/bsky/interface.rs b/crates/x-media/src/site/bsky/interface.rs index 0bdcacf..00c3a8c 100644 --- a/crates/x-media/src/site/bsky/interface.rs +++ b/crates/x-media/src/site/bsky/interface.rs @@ -81,7 +81,7 @@ pub async fn fetch_from_url(url: &str) -> Result { url: mp4_path.to_string_lossy().into_owned(), thumbnail_url, }); - fetched._keep_alive = Some(keep_alive); + fetched._keep_alive = Some(std::sync::Arc::new(keep_alive)); } // No ffmpeg: a deployment gap, not a bad moment — retrying it // would only waste the fetch budget, so the post degrades (and an diff --git a/crates/x-media/src/site/mod.rs b/crates/x-media/src/site/mod.rs index cb4ca8a..af2b0f1 100644 --- a/crates/x-media/src/site/mod.rs +++ b/crates/x-media/src/site/mod.rs @@ -52,8 +52,11 @@ pub struct Fetched { /// Raw values (pre-escaped) for user-customizable caption formats. pub(crate) render_data: Option, /// Keeps temp files (e.g. an encoded ugoira MP4) alive until the caller - /// finishes uploading; not part of the public contract. - pub(crate) _keep_alive: Option, + /// finishes uploading; not part of the public contract. Shared rather than + /// owned because one fetched post can serve several sends — the bot shares + /// one in-flight fetch between concurrent duplicates of the same link — and + /// the files have to outlive every one of them. + pub(crate) _keep_alive: Option>, } /// Values for the `{url} {author} {author_url} {title} {content} {tags}` @@ -133,12 +136,14 @@ impl Fetched { }) } - /// Hands over the temp dir keeping locally produced media (ugoira MP4, + /// A reference to the temp dir keeping locally produced media (ugoira MP4, /// bsky remux MP4) alive. The bot keeps it while its task may still be /// retried by the queue, which runs after this [`Fetched`] is dropped and - /// its temp files would otherwise be gone. `None` when no such dir exists. - pub fn take_keep_alive(&mut self) -> Option { - self._keep_alive.take() + /// its temp files would otherwise be gone. `None` when no such dir exists; + /// each clone keeps the directory alive for as long as it lives, so two + /// sends of one post can each hold the same files. + pub fn keep_alive(&self) -> Option> { + self._keep_alive.clone() } } @@ -742,6 +747,38 @@ pub async fn download_media_to_file( mod tests { use super::*; + /// Locally produced media (ugoira MP4, bsky remux MP4) lives in a temp dir + /// whose lifetime is refcounted: one fetch result can serve several sends + /// (the bot shares one in-flight fetch between concurrent duplicates), and + /// the files must outlive all of them — but no longer than the last one. + #[test] + fn a_keep_alive_clone_outlives_the_fetched() { + let dir = tempfile::tempdir().unwrap(); + let file = dir.path().join("media.mp4"); + std::fs::write(&file, b"mp4").unwrap(); + let fetched = Fetched { + source_url: "https://x.com/u/status/1".into(), + caption: String::new(), + title: String::new(), + content: String::new(), + media: Vec::new(), + sensitive: false, + site_id: "twitter", + render_data: None, + _keep_alive: Some(std::sync::Arc::new(dir)), + }; + + let shared = fetched.keep_alive().expect("a temp dir to share"); + drop(fetched); + assert!(file.exists(), "the file must survive the fetched post"); + + let second = shared.clone(); + drop(shared); + assert!(file.exists(), "another holder keeps it alive"); + drop(second); + assert!(!file.exists(), "the last holder releases the directory"); + } + #[test] fn cache_key_normalizes_domain_variants() { assert_eq!( diff --git a/crates/x-media/src/site/pixiv/api.rs b/crates/x-media/src/site/pixiv/api.rs index 28f3ae5..588361d 100644 --- a/crates/x-media/src/site/pixiv/api.rs +++ b/crates/x-media/src/site/pixiv/api.rs @@ -153,7 +153,7 @@ impl PixivAPI { url: mp4_path, thumbnail_url: model.image_urls.medium.clone(), }); - illustration._keep_alive = Some(_keep_alive); + illustration._keep_alive = Some(std::sync::Arc::new(_keep_alive)); } Ok(None) => {} Err(e) => { diff --git a/crates/x-media/src/site/pixiv/interface.rs b/crates/x-media/src/site/pixiv/interface.rs index 98aba91..3d46aa8 100644 --- a/crates/x-media/src/site/pixiv/interface.rs +++ b/crates/x-media/src/site/pixiv/interface.rs @@ -181,7 +181,7 @@ pub struct Illustration { pub(crate) media: Vec, nsfw: bool, /// Keeps a temp dir (ugoira MP4) alive until the send completes. - pub(crate) _keep_alive: Option, + pub(crate) _keep_alive: Option>, } impl Illustration { diff --git a/crates/xmedia-bot/src/handlers/urls.rs b/crates/xmedia-bot/src/handlers/urls.rs index 359abd4..c6f9850 100644 --- a/crates/xmedia-bot/src/handlers/urls.rs +++ b/crates/xmedia-bot/src/handlers/urls.rs @@ -7,7 +7,7 @@ use crate::link_cache::{CachedMediaKind, CachedPost}; use crate::media_sender::MediaSender; use crate::send::{self, MediaItemPayload, Task}; use crate::state::ChatData; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::sync::LazyLock; use teloxide::RequestError; @@ -35,6 +35,97 @@ static URL_WORKER_HANDLES: LazyLock>>> = + LazyLock::new(Default::default); + +/// A fetched post (or the error that stopped it), shared as-is: the error side +/// is not `Clone`, so callers read it through the `Arc` — the same shape the +/// send paths use for `&Fetched`. +type FetchOutcome = Result, x_media::site::FetchError>; + +/// The channel a sharer publishes its result on, and waiters subscribe to. +type SharedFetch = tokio::sync::broadcast::Sender>; + +/// Fetches `url`, sharing one in-flight fetch per `key` (its site cache key) +/// with every other caller asking for the same post meanwhile. +async fn fetch_shared(key: &str, url: &str) -> std::sync::Arc { + shared_fetch(&IN_FLIGHT_FETCHES, key, || x_media::site::fetch(url)).await +} + +/// [`fetch_shared`]'s core, over the caller's own map so the sharing rules can +/// be tested without a network fetch. +/// +/// A caller that finds a live entry subscribes to it and waits; the caller that +/// created the entry runs `fetch` and publishes the result. Two things keep +/// that from stranding a request: the entry is removed by a guard (so a +/// cancelled fetch cannot leave waiters subscribed to a channel nothing will +/// ever write to), and a waiter whose sharer vanished fetches for itself. +async fn shared_fetch( + map: &parking_lot::Mutex>>, + key: &str, + fetch: F, +) -> std::sync::Arc +where + T: Send + Sync + 'static, + F: FnOnce() -> Fut, + Fut: Future, +{ + let (sender, leader) = { + let mut map = map.lock(); + match map.get(key) { + Some(sender) => (sender.clone(), false), + None => { + let (sender, _) = tokio::sync::broadcast::channel(1); + map.insert(key.to_string(), sender.clone()); + (sender, true) + } + } + }; + if !leader { + // `Err` means the entry is gone without a value: the sharer was + // cancelled, or it finished just as this caller subscribed (the + // message predates the subscription). Fetch for ourselves instead of + // failing a link that is perfectly fetchable. + let mut receiver = sender.subscribe(); + // The sender clone taken from the map is dropped first: held, it would + // keep the channel open past the sharer's exit (a broadcast channel + // closes when *all* senders are gone), and `recv` would wait forever + // instead of reporting that the sharer vanished. + drop(sender); + match receiver.recv().await { + Ok(shared) => return shared, + Err(_) => return std::sync::Arc::new(fetch().await), + } + } + // Removes the entry on every exit path, cancellation included. + let _guard = InFlightFetch { map, key }; + let outcome = std::sync::Arc::new(fetch().await); + // No receiver is the common case, not an error: a lone caller has nobody + // to publish to. + let _ = sender.send(std::sync::Arc::clone(&outcome)); + outcome +} + +/// Drops the in-flight entry it was created for, however the fetch ends. +struct InFlightFetch<'a, T> { + map: &'a parking_lot::Mutex>>, + key: &'a str, +} + +impl Drop for InFlightFetch<'_, T> { + fn drop(&mut self) { + self.map.lock().remove(self.key); + } +} + /// Starts the URL job workers (called once from main after the queue starts). /// teloxide dispatches updates to a per-chat worker that handles them /// sequentially, so a batch-forward of many messages would otherwise be @@ -585,7 +676,15 @@ async fn url_media_inner( log::debug!("fetching [key={}]", log_key(url)); log::trace!("fetching {url}"); - match x_media::site::fetch(url).await { + // One fetch per post at a time: a concurrent duplicate of this link waits + // for *this* fetch instead of running its own (see [`fetch_shared`]). + let outcome = match x_media::site::cache_key(url) { + Some(key) => fetch_shared(&key, url).await, + // A URL no site claims (reached only through `/test`): nothing to key + // the sharing on, and the dispatcher answers without a request. + None => std::sync::Arc::new(x_media::site::fetch(url).await), + }; + match &*outcome { // Unsupported links are ignored silently (Python parity). Ok(None) => { // The URL itself is user data, so only `trace` names the link; @@ -596,9 +695,9 @@ async fn url_media_inner( // Retries exhausted: notify the user (Rust-only requirement 3). Err(e) => { log::error!("fetch [key={}]: {e}", log_key(url)); - let _ = reply(ctx.sender, chat_id, reply_to, fetch_error_message(&e)).await; + let _ = reply(ctx.sender, chat_id, reply_to, fetch_error_message(e)).await; } - Ok(Some(mut fetched)) => { + Ok(Some(fetched)) => { if fetched.media.is_empty() { let _ = reply( ctx.sender, @@ -654,8 +753,9 @@ async fn url_media_inner( // Hand the keep-alive temp dir (ugoira / bsky remux MP4) to the // retry registry: a queued retry runs after this function returns // and the fetch's own TempDir is dropped, so without this the - // local file would be gone by the time the retry sends it. - if let Some(dir) = fetched.take_keep_alive() { + // local file would be gone by the time the retry sends it. Shared + // (Arc), so a second send of the same post holds its own reference. + if let Some(dir) = fetched.keep_alive() { send::KEEP_ALIVE.lock().push(dir); } dispatch_send(ctx, chat_id, reply_to, &task, url, started).await; @@ -772,7 +872,7 @@ async fn refetch( chat_id: i64, url: &str, ) -> Result, x_media::site::FetchError> { - let Some(mut fetched) = x_media::site::fetch(url).await? else { + let Some(fetched) = x_media::site::fetch(url).await? else { return Ok(None); }; if fetched.media.is_empty() { @@ -805,7 +905,7 @@ async fn refetch( .collect(); // The re-fetch may produce a fresh local file (ugoira / bsky remux): hand it // to the same keep-alive registry the first fetch uses. - if let Some(dir) = fetched.take_keep_alive() { + if let Some(dir) = fetched.keep_alive() { send::KEEP_ALIVE.lock().push(dir); } Ok(Some(Refetched { @@ -964,6 +1064,90 @@ mod tests { assert_eq!(sender.calls(), vec!["send_chat_action"]); } + // ── One fetch per post ─────────────────────────────────────────────── + + /// Two callers asking for the same post while its fetch is in flight run + /// one fetch between them: the duplicate (a second chat, a batch forward + /// and a retry) waits for that result instead of paying for its own. + #[tokio::test] + async fn concurrent_callers_share_one_fetch() { + let map = parking_lot::Mutex::new(HashMap::new()); + let calls = std::sync::atomic::AtomicUsize::new(0); + let fetch = || async { + calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(50)).await; + 7u32 + }; + + let (first, second) = tokio::join!( + shared_fetch(&map, "twitter:1", fetch), + shared_fetch(&map, "twitter:1", fetch) + ); + + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); + assert_eq!(*first, 7); + assert!(std::sync::Arc::ptr_eq(&first, &second), "one shared result"); + assert!( + map.lock().is_empty(), + "the entry must not outlive the fetch" + ); + } + + /// The dedup is *concurrent* only. A caller arriving after the fetch + /// settled fetches again: the source may have changed, and a failure is + /// deliberately not cached (the user is told to try again). + #[tokio::test] + async fn a_later_call_fetches_again() { + let map = parking_lot::Mutex::new(HashMap::new()); + let calls = std::sync::atomic::AtomicUsize::new(0); + let counting = |value: u32| { + let calls = &calls; + async move { + calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + value + } + }; + + let first = shared_fetch(&map, "twitter:1", || counting(1)).await; + let second = shared_fetch(&map, "twitter:1", || counting(2)).await; + + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2); + assert_eq!((*first, *second), (1, 2)); + assert!(!std::sync::Arc::ptr_eq(&first, &second)); + } + + /// A cancelled fetch must not strand the callers that joined it: a live + /// entry whose sharer is gone holds a sender, and the waiters would wait + /// for a value that can never come. They fetch for themselves. + #[tokio::test] + async fn a_cancelled_fetch_does_not_strand_waiters() { + let map = parking_lot::Mutex::new(HashMap::new()); + let calls = std::sync::atomic::AtomicUsize::new(0); + + // A fetch that never finishes, cancelled by the timeout below once it + // has installed its entry. + let slow = shared_fetch(&map, "twitter:1", || async { + calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + std::future::pending::<()>().await; + 0u32 + }); + assert!( + tokio::time::timeout(Duration::from_millis(50), slow) + .await + .is_err(), + "the sharer must still be waiting when it is cancelled" + ); + + // Its entry is gone, and the next caller fetches its own value. + let value = shared_fetch(&map, "twitter:1", || async { + calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + 5u32 + }) + .await; + assert_eq!(*value, 5); + assert!(map.lock().is_empty()); + } + // ── Send modes: the URL flow vs `/test` ───────────────────────────── /// A chat that has both post-send actions configured. diff --git a/crates/xmedia-bot/src/send/mod.rs b/crates/xmedia-bot/src/send/mod.rs index b055485..4a1624e 100644 --- a/crates/xmedia-bot/src/send/mod.rs +++ b/crates/xmedia-bot/src/send/mod.rs @@ -1495,7 +1495,7 @@ mod tests { *notify_message_id = None; } let dir_path = dir.path().to_path_buf(); - KEEP_ALIVE.lock().push(dir); + KEEP_ALIVE.lock().push(std::sync::Arc::new(dir)); // No chat to notify → the notify path sends nothing (its mock would // have no scripted outcome left). diff --git a/crates/xmedia-bot/src/send/post_send.rs b/crates/xmedia-bot/src/send/post_send.rs index 4597ef2..bb910cc 100644 --- a/crates/xmedia-bot/src/send/post_send.rs +++ b/crates/xmedia-bot/src/send/post_send.rs @@ -80,12 +80,14 @@ async fn invalidate_cache(cache: &LinkCache, task: &Task) { /// Locally produced media files (ugoira MP4, bsky remux MP4) whose temp dirs /// must stay alive while their task may be retried by the queue. The fetch -/// pipeline hands ownership here via -/// [`x_media::site::Fetched::take_keep_alive`] before that -/// [`x_media::site::Fetched`] is dropped; a queued retry runs after that drop, -/// so without this the local file would be gone by the time the retry sends -/// it. Entries are removed when the task settles (see [`release_keep_alive`]). -pub(crate) static KEEP_ALIVE: LazyLock>> = +/// pipeline hands a reference here via [`x_media::site::Fetched::keep_alive`] +/// before that [`x_media::site::Fetched`] is dropped; a queued retry runs after +/// that drop, so without this the local file would be gone by the time the +/// retry sends it. `Arc` because one fetch can serve several tasks (a +/// concurrent duplicate of the same link shares it): each holder keeps the +/// directory alive until its own task settles. Entries are removed when the +/// task settles (see [`release_keep_alive`]). +pub(crate) static KEEP_ALIVE: LazyLock>>> = LazyLock::new(|| parking_lot::Mutex::new(Vec::new())); /// Drops the keep-alive temp dirs holding media referenced by `task` (matched