mirror of
https://github.com/TheFunny/TelegramTwitterMediaBot.git
synced 2026-09-23 23:32:05 +00:00
perf: share one fetch between concurrent duplicates of a link
`link_cache` is only written *after* a send succeeds, so two chats posting the same link at the same moment each ran a full fetch: two sets of source requests, and for an ugoira or a bsky video two ffmpeg encodes of the same post — minutes of CPU for the second one. The same applies to a batch forward racing a queued retry. Within one message `dedupe_urls` already handled the duplicates; across calls nothing did. `fetch_shared` keys an in-flight fetch by the post's cache key: the first caller runs it, the rest subscribe and take its result. The entry is removed by a guard when the fetch settles (cancellation included), so this dedupes what is *concurrent* and never answers from an old result — a repeat later fetches again, and a failure is deliberately not cached: the user is told to try again, and a cached failure would answer that retry from a stale state. Two hazards that shape the code, each with a test: a broadcast channel lives while *any* sender does, so the waiter drops its own clone of the sender before waiting — otherwise a cancelled sharer would leave it waiting forever — and a waiter whose sharer vanished fetches for itself instead of failing a link that is perfectly fetchable. One fetched post can now serve several sends, which the temp files behind `Fetched::keep_alive` had to support: the field is `Arc<TempDir>` and the accessor hands out references (the bot's `KEEP_ALIVE` registry holds the same), so the ugoira/bsky MP4 stays on disk until the *last* task settles rather than the first. `take_keep_alive` is gone — a `take` could only ever serve one of the senders. Verified: 116 bot tests (three new sharing tests, including the cancelled sharer) and 91 x-media tests (a new one pinning the keep-alive refcount) pass, plus a live check that two concurrent `fetch_shared` calls for one tweet return the very same `Arc` after one fetch's worth of wall time. `cargo fmt --check`, `cargo clippy --workspace --all-targets --locked -- -D warnings` and `cargo test --workspace --locked` clean.
This commit is contained in:
@@ -81,7 +81,7 @@ pub async fn fetch_from_url(url: &str) -> Result<Fetched, FetchError> {
|
||||
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
|
||||
|
||||
@@ -52,8 +52,11 @@ pub struct Fetched {
|
||||
/// Raw values (pre-escaped) for user-customizable caption formats.
|
||||
pub(crate) render_data: Option<RenderData>,
|
||||
/// 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<tempfile::TempDir>,
|
||||
/// 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<std::sync::Arc<tempfile::TempDir>>,
|
||||
}
|
||||
|
||||
/// 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<tempfile::TempDir> {
|
||||
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<std::sync::Arc<tempfile::TempDir>> {
|
||||
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!(
|
||||
|
||||
@@ -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) => {
|
||||
|
||||
@@ -181,7 +181,7 @@ pub struct Illustration {
|
||||
pub(crate) media: Vec<Media>,
|
||||
nsfw: bool,
|
||||
/// Keeps a temp dir (ugoira MP4) alive until the send completes.
|
||||
pub(crate) _keep_alive: Option<tempfile::TempDir>,
|
||||
pub(crate) _keep_alive: Option<std::sync::Arc<tempfile::TempDir>>,
|
||||
}
|
||||
|
||||
impl Illustration {
|
||||
|
||||
@@ -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<parking_lot::Mutex<Option<Vec<tokio::task::J
|
||||
/// while bounding how many jobs can be queued at all.
|
||||
const URL_WORKERS: usize = 8;
|
||||
|
||||
/// One fetch per post at a time, keyed by the normalized cache key. Two chats
|
||||
/// posting the same link at the same moment (or a batch forward and a queued
|
||||
/// retry) used to run two full fetches: two sets of source requests, and for an
|
||||
/// ugoira or a bsky video two ffmpeg encodes of the same post. The first caller
|
||||
/// runs it and the rest wait for its result. The entry is dropped the moment
|
||||
/// the fetch settles, so this dedupes what is *concurrent* and never answers
|
||||
/// from an old result: a repeat later fetches again, and a failure is not
|
||||
/// cached (the user may well retry it).
|
||||
static IN_FLIGHT_FETCHES: LazyLock<parking_lot::Mutex<HashMap<String, SharedFetch<FetchOutcome>>>> =
|
||||
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<Option<x_media::site::Fetched>, x_media::site::FetchError>;
|
||||
|
||||
/// The channel a sharer publishes its result on, and waiters subscribe to.
|
||||
type SharedFetch<T> = tokio::sync::broadcast::Sender<std::sync::Arc<T>>;
|
||||
|
||||
/// 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<FetchOutcome> {
|
||||
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<T, F, Fut>(
|
||||
map: &parking_lot::Mutex<HashMap<String, SharedFetch<T>>>,
|
||||
key: &str,
|
||||
fetch: F,
|
||||
) -> std::sync::Arc<T>
|
||||
where
|
||||
T: Send + Sync + 'static,
|
||||
F: FnOnce() -> Fut,
|
||||
Fut: Future<Output = T>,
|
||||
{
|
||||
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<HashMap<String, SharedFetch<T>>>,
|
||||
key: &'a str,
|
||||
}
|
||||
|
||||
impl<T> 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<Option<Refetched>, 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.
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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<parking_lot::Mutex<Vec<tempfile::TempDir>>> =
|
||||
/// 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<parking_lot::Mutex<Vec<std::sync::Arc<tempfile::TempDir>>>> =
|
||||
LazyLock::new(|| parking_lot::Mutex::new(Vec::new()));
|
||||
|
||||
/// Drops the keep-alive temp dirs holding media referenced by `task` (matched
|
||||
|
||||
Reference in New Issue
Block a user