mirror of
https://github.com/TheFunny/TelegramTwitterMediaBot.git
synced 2026-09-25 23:52:04 +00:00
refactor(handlers): split urls.rs into workers, pipeline and startup repair
urls.rs carried five reasons to change: the job channel and its worker pool, the single-flight fetch, URL parsing, the per-URL pipeline, and the startup repair of queued retries. The two with their own lifecycle move out: - url_workers.rs: the bounded channel, its supervised pool and start/stop_url_workers (the pipeline stays in urls.rs, which the workers call). - repair.rs: needs_refetch/apply_refresh/refetch/repair_lost_local_media with their tests, moved whole (the live one keeps its #[ignore]). Also folds the two byte-identical render_fields -> CachedPost mappings (urls.rs and repair.rs) into urls::cached_snapshot, and moves the shared permanent_error test fixture into ctx::test_support.
This commit is contained in:
@@ -65,6 +65,13 @@ pub(crate) mod test_support {
|
||||
RequestError::Api(ApiError::Unknown(message.to_string()))
|
||||
}
|
||||
|
||||
/// The API error a caption edit that changes nothing answers with — what
|
||||
/// the mocks script for a permanent send failure. A `fn` pointer, so it can
|
||||
/// be handed to `MockSender::scripted` as-is.
|
||||
pub(crate) fn permanent_error() -> RequestError {
|
||||
api_error("Bad Request: message is not modified")
|
||||
}
|
||||
|
||||
/// One photo payload item: `media` in the two flags the tests vary (no
|
||||
/// smaller variant, since that is the field most tests leave alone —
|
||||
/// `send`'s own tests build that case directly).
|
||||
|
||||
@@ -9,18 +9,20 @@
|
||||
mod callback;
|
||||
mod commands;
|
||||
mod inline;
|
||||
mod repair;
|
||||
mod statics;
|
||||
mod url_workers;
|
||||
mod urls;
|
||||
|
||||
pub use callback::callback_query_handler;
|
||||
pub use commands::register_commands;
|
||||
pub use inline::inline_query_handler;
|
||||
pub(crate) use inline::prune_idle_states;
|
||||
pub(crate) use repair::repair_lost_local_media;
|
||||
/// The resolved `$DATA_DIR/task_queue.db` path, for the startup config line.
|
||||
pub(crate) use statics::db_path;
|
||||
pub use statics::{CHAT_STORE, CONFIG, LINK_CACHE, TASK_QUEUE};
|
||||
pub(crate) use urls::repair_lost_local_media;
|
||||
pub use urls::{start_url_workers, stop_url_workers};
|
||||
pub use url_workers::{start_url_workers, stop_url_workers};
|
||||
|
||||
use crate::ctx::AppContext;
|
||||
use crate::media_sender::MediaSender;
|
||||
@@ -29,7 +31,8 @@ use teloxide::RequestError;
|
||||
use teloxide::prelude::*;
|
||||
use teloxide::types::{ChatId, Message, MessageId};
|
||||
use teloxide::utils::command::BotCommands;
|
||||
use urls::{URL_JOBS, extract_urls};
|
||||
use url_workers::URL_JOBS;
|
||||
use urls::extract_urls;
|
||||
|
||||
/// Reply to a message by id, keeping the reply decoration even if the
|
||||
/// original was already deleted.
|
||||
|
||||
@@ -0,0 +1,355 @@
|
||||
//! Startup repair, run before any queue worker exists: a queued retry whose
|
||||
//! local media (a ugoira MP4, a bsky remux, a downloaded temp file) did not
|
||||
//! survive the restart can never succeed, because the registry that kept those
|
||||
//! files alive (`send::KEEP_ALIVE`) is in memory. Those rows are re-fetched
|
||||
//! from their post instead of dead-lettering the user's link.
|
||||
|
||||
use super::log_key;
|
||||
use super::urls::{cached_snapshot, media_to_payload};
|
||||
use crate::ctx::AppContext;
|
||||
use crate::link_cache::CachedPost;
|
||||
use crate::send::{self, Delivery, MediaItemPayload, Task};
|
||||
|
||||
// ── Startup repair: queued retries whose local media did not survive ───────
|
||||
|
||||
/// A post's fresh media plus the caption and cache snapshot that go with them:
|
||||
/// what [`refetch`] hands [`apply_refresh`]. Plain data, so the rewrite below
|
||||
/// can be tested without a network fetch (which cannot be faked here:
|
||||
/// [`x_media::site::Fetched`] keeps a private field and is not constructible
|
||||
/// outside its crate).
|
||||
struct Refetched {
|
||||
caption: String,
|
||||
items: Vec<MediaItemPayload>,
|
||||
cache_data: Option<CachedPost>,
|
||||
}
|
||||
|
||||
/// Whether a queued task should have its post re-fetched, because it still
|
||||
/// wants a local file (ugoira MP4, a bsky remux, a downloaded temp file) that is
|
||||
/// gone. Those files live in the system temp dir and the registry that keeps
|
||||
/// them alive for the retry (`send::KEEP_ALIVE`) is in memory, so a restart
|
||||
/// takes all of them — a retry that needs one can only dead-letter.
|
||||
///
|
||||
/// A partially delivered album is left alone: its remaining batches cannot be
|
||||
/// reconciled with a fresh media list without risking a second copy of what the
|
||||
/// user already received.
|
||||
fn needs_refetch(task: &Task) -> bool {
|
||||
if let Task::SendMediaSequence {
|
||||
batch_index,
|
||||
sent_message_ids,
|
||||
..
|
||||
} = task
|
||||
&& (*batch_index > 0 || !sent_message_ids.is_empty())
|
||||
{
|
||||
return false;
|
||||
}
|
||||
task.local_media_paths().iter().any(|path| !path.exists())
|
||||
}
|
||||
|
||||
/// Rebuilds the task from the fresh media, keeping its delivery envelope (chat,
|
||||
/// reply, forward/edit settings, notify targets): the retry that was queued must
|
||||
/// still deliver the same way, whoever asked for it.
|
||||
fn apply_refresh(task: &Task, fresh: &Refetched) -> Option<Task> {
|
||||
let chat_id = task.chat_id()?;
|
||||
let (edit_before_forward, forward_channel_id) = match task {
|
||||
Task::SendMediaSequence {
|
||||
edit_before_forward,
|
||||
forward_channel_id,
|
||||
..
|
||||
}
|
||||
| Task::SendAnimation {
|
||||
edit_before_forward,
|
||||
forward_channel_id,
|
||||
..
|
||||
} => (*edit_before_forward, *forward_channel_id),
|
||||
Task::ForwardMessages { .. } => return None,
|
||||
};
|
||||
let reply_to_message_id = match task {
|
||||
Task::SendMediaSequence {
|
||||
reply_to_message_id,
|
||||
..
|
||||
}
|
||||
| Task::SendAnimation {
|
||||
reply_to_message_id,
|
||||
..
|
||||
} => *reply_to_message_id,
|
||||
Task::ForwardMessages { .. } => return None,
|
||||
};
|
||||
let (notify_chat_id, notify_message_id) = task.notify_target();
|
||||
Some(Task::from_items(
|
||||
Delivery {
|
||||
chat_id,
|
||||
reply_to_message_id,
|
||||
edit_before_forward,
|
||||
forward_channel_id,
|
||||
notify_chat_id,
|
||||
notify_message_id,
|
||||
},
|
||||
task.source_url()?.to_string(),
|
||||
fresh.caption.clone(),
|
||||
fresh.items.clone(),
|
||||
fresh.cache_data.clone(),
|
||||
))
|
||||
}
|
||||
|
||||
/// Fetches the post again and maps it into [`Refetched`]: the same mapping the
|
||||
/// fresh-fetch path uses (per-site caption format from the chat, render fields
|
||||
/// for the link-cache snapshot), so a repaired task looks like a first send.
|
||||
async fn refetch(
|
||||
ctx: &AppContext<'_>,
|
||||
chat_id: i64,
|
||||
url: &str,
|
||||
) -> Result<Option<Refetched>, x_media::site::FetchError> {
|
||||
let Some(fetched) = x_media::site::fetch(url).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
if fetched.media.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
let chat_data = ctx.chat_store.get(chat_id).await;
|
||||
let format = chat_data.format_for(fetched.site_id);
|
||||
let caption = fetched.caption_with(&format);
|
||||
let cache_data = cached_snapshot(&fetched);
|
||||
let items: Vec<MediaItemPayload> = fetched
|
||||
.media
|
||||
.iter()
|
||||
.map(|media| media_to_payload(media, fetched.sensitive))
|
||||
.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.keep_alive() {
|
||||
send::KEEP_ALIVE.lock().push(dir);
|
||||
}
|
||||
Ok(Some(Refetched {
|
||||
caption,
|
||||
items,
|
||||
cache_data,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Re-fetches every queued task whose local media did not survive the restart,
|
||||
/// so the user's link is still delivered instead of dead-lettering on a file
|
||||
/// that cannot come back. Returns how many rows were rewritten.
|
||||
///
|
||||
/// Startup only, before the queue workers start: no worker can lease a row while
|
||||
/// this writes, which is what lets it replace payloads without the lease-token
|
||||
/// guard every worker write-back carries.
|
||||
pub(crate) async fn repair_lost_local_media(ctx: &AppContext<'_>) -> usize {
|
||||
let mut repaired = 0;
|
||||
for (id, payload) in ctx.task_queue.runnable_rows().await {
|
||||
let Ok(task) = serde_json::from_str::<Task>(&payload) else {
|
||||
continue;
|
||||
};
|
||||
if !needs_refetch(&task) {
|
||||
continue;
|
||||
}
|
||||
let (Some(url), Some(chat_id)) = (task.source_url().map(str::to_string), task.chat_id())
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
match refetch(ctx, chat_id, &url).await {
|
||||
Ok(Some(fresh)) => {
|
||||
let Some(updated) = apply_refresh(&task, &fresh) else {
|
||||
continue;
|
||||
};
|
||||
let updated = serde_json::to_value(&updated).expect("task serializes");
|
||||
if ctx.task_queue.replace_payload(&id, &updated).await {
|
||||
repaired += 1;
|
||||
log::info!(
|
||||
"startup repair: re-fetched [key={}] for chat={chat_id} (its local media did not survive the restart)",
|
||||
log_key(&url)
|
||||
);
|
||||
}
|
||||
}
|
||||
// The post is gone or withheld now: the retry could not have
|
||||
// delivered anything either, so say why instead of letting it
|
||||
// dead-letter on a missing file.
|
||||
Ok(None) | Err(_) => {
|
||||
let (notify_chat_id, notify_message_id) = task.notify_target();
|
||||
log::warn!(
|
||||
"startup repair: [key={}] for chat={chat_id} needed a re-fetch and none was possible",
|
||||
log_key(&url)
|
||||
);
|
||||
send::notify_failure(
|
||||
ctx.sender,
|
||||
notify_chat_id,
|
||||
notify_message_id,
|
||||
&format!(
|
||||
"{} — the media held for retry was lost when the bot restarted and the post could not be fetched again. Please send the link again.",
|
||||
log_key(&url)
|
||||
),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
repaired
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::ctx::test_support::{TestStores, permanent_error, photo_item};
|
||||
use crate::media_sender::test_support::MockSender;
|
||||
|
||||
fn queued_task(media: &str, batch_index: usize, sent: Vec<i64>) -> Task {
|
||||
Task::SendMediaSequence {
|
||||
chat_id: 1,
|
||||
reply_to_message_id: 2,
|
||||
caption: "cap".into(),
|
||||
media_batches: vec![vec![photo_item(media, false, false)]],
|
||||
batch_index,
|
||||
sent_message_ids: sent,
|
||||
source_url: "https://x.com/u/status/1".into(),
|
||||
edit_before_forward: true,
|
||||
forward_channel_id: Some(2),
|
||||
notify_chat_id: Some(1),
|
||||
notify_message_id: Some(2),
|
||||
cache_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_tasks_missing_a_local_file_need_a_refetch() {
|
||||
// A URL send needs nothing.
|
||||
assert!(!needs_refetch(&queued_task("https://cdn/1.jpg", 0, vec![])));
|
||||
// A local path that is still there (a survived temp file) needs nothing.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let alive = dir.path().join("ugoira.mp4");
|
||||
std::fs::write(&alive, b"x").unwrap();
|
||||
assert!(!needs_refetch(&queued_task(
|
||||
alive.to_str().unwrap(),
|
||||
0,
|
||||
vec![]
|
||||
)));
|
||||
// A local path the restart took away does.
|
||||
assert!(needs_refetch(&queued_task(
|
||||
"/nonexistent-ugoira.mp4",
|
||||
0,
|
||||
vec![]
|
||||
)));
|
||||
// A partially delivered album is left to its own retry path.
|
||||
assert!(!needs_refetch(&queued_task(
|
||||
"/nonexistent-ugoira.mp4",
|
||||
1,
|
||||
vec![7]
|
||||
)));
|
||||
assert!(!needs_refetch(&queued_task(
|
||||
"/nonexistent-ugoira.mp4",
|
||||
0,
|
||||
vec![7]
|
||||
)));
|
||||
// A channel copy holds no media.
|
||||
assert!(!needs_refetch(&Task::ForwardMessages {
|
||||
from_chat_id: 1,
|
||||
to_chat_id: 2,
|
||||
message_ids: vec![3],
|
||||
notify_chat_id: None,
|
||||
notify_message_id: None,
|
||||
}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apply_refresh_keeps_the_delivery_envelope() {
|
||||
let task = queued_task("/nonexistent-ugoira.mp4", 0, vec![]);
|
||||
let fresh = Refetched {
|
||||
caption: "fresh caption".into(),
|
||||
items: vec![photo_item("https://cdn/fresh.jpg", true, false)],
|
||||
cache_data: None,
|
||||
};
|
||||
match apply_refresh(&task, &fresh).expect("a repairable task") {
|
||||
Task::SendMediaSequence {
|
||||
chat_id,
|
||||
reply_to_message_id,
|
||||
caption,
|
||||
media_batches,
|
||||
batch_index,
|
||||
sent_message_ids,
|
||||
source_url,
|
||||
edit_before_forward,
|
||||
forward_channel_id,
|
||||
notify_chat_id,
|
||||
notify_message_id,
|
||||
..
|
||||
} => {
|
||||
// Same delivery: chat, reply, forward/edit settings, notify.
|
||||
assert_eq!((chat_id, reply_to_message_id), (1, 2));
|
||||
assert!(edit_before_forward);
|
||||
assert_eq!(forward_channel_id, Some(2));
|
||||
assert_eq!((notify_chat_id, notify_message_id), (Some(1), Some(2)));
|
||||
assert_eq!(source_url, "https://x.com/u/status/1");
|
||||
// Fresh media, and nothing of it counted as sent yet.
|
||||
assert_eq!(caption, "fresh caption");
|
||||
assert!(
|
||||
matches!(
|
||||
&media_batches[0][0],
|
||||
MediaItemPayload::Photo { media, .. } if media == "https://cdn/fresh.jpg"
|
||||
),
|
||||
"fresh media must replace the lost local file"
|
||||
);
|
||||
assert!(matches!(
|
||||
media_batches[0][0],
|
||||
MediaItemPayload::Photo {
|
||||
has_spoiler: true,
|
||||
..
|
||||
}
|
||||
));
|
||||
assert_eq!((batch_index, sent_message_ids.len()), (0, 0));
|
||||
}
|
||||
other => panic!("expected a media sequence, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// The whole repair against a real post: a queued row whose media is a local
|
||||
/// file the restart took away is re-fetched from its `source_url` and
|
||||
/// rewritten in place, so the retry can still deliver it.
|
||||
#[tokio::test]
|
||||
#[ignore = "live network: requires outbound HTTPS to public.api.bsky.app"]
|
||||
async fn live_repair_refetches_a_lost_local_media_row() {
|
||||
let stores = TestStores::new();
|
||||
// An empty script: the repair must not need to tell the user anything.
|
||||
let sender = MockSender::scripted(vec![], permanent_error);
|
||||
let ctx = stores.ctx(&sender);
|
||||
let mut task = queued_task("/nonexistent-ugoira.mp4", 0, vec![]);
|
||||
if let Task::SendMediaSequence { source_url, .. } = &mut task {
|
||||
*source_url = "https://bsky.app/profile/fu-futa.bsky.social/post/3laoveufjv224".into();
|
||||
}
|
||||
stores
|
||||
.task_queue()
|
||||
.enqueue(serde_json::to_value(&task).unwrap(), crate::db::now_f64())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(repair_lost_local_media(&ctx).await, 1);
|
||||
|
||||
let updated: Task = serde_json::from_value(stores.queued_payload().await).unwrap();
|
||||
match updated {
|
||||
Task::SendMediaSequence {
|
||||
media_batches,
|
||||
batch_index,
|
||||
sent_message_ids,
|
||||
caption,
|
||||
..
|
||||
} => {
|
||||
let media: Vec<String> = media_batches
|
||||
.iter()
|
||||
.flatten()
|
||||
.map(|item| match item {
|
||||
MediaItemPayload::Photo { media, .. }
|
||||
| MediaItemPayload::Video { media, .. }
|
||||
| MediaItemPayload::Animation { media, .. } => media.clone(),
|
||||
})
|
||||
.collect();
|
||||
assert!(!media.is_empty(), "the fresh fetch yielded no media");
|
||||
assert!(
|
||||
media.iter().all(|m| m.starts_with("http")),
|
||||
"the retry must be uploadable from URLs again: {media:?}"
|
||||
);
|
||||
assert_eq!((batch_index, sent_message_ids.len()), (0, 0));
|
||||
assert!(!caption.is_empty());
|
||||
}
|
||||
other => panic!("expected a repaired media sequence, got {other:?}"),
|
||||
}
|
||||
// The post was re-read, not re-delivered: nothing was sent.
|
||||
assert!(sender.calls().is_empty(), "{:?}", sender.calls());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
//! The URL job channel and its worker pool: a bounded queue (backpressure
|
||||
//! instead of unbounded spawns) drained by [`URL_WORKERS`] supervised workers.
|
||||
//!
|
||||
//! teloxide's per-chat workers are sequential, so a batch forward needs its own
|
||||
//! concurrency: this is where a link handed over by `handlers::mod` actually
|
||||
//! reaches the pipeline.
|
||||
|
||||
use super::urls::{PostSend, url_media};
|
||||
use crate::ctx::CONTEXT;
|
||||
use std::sync::LazyLock;
|
||||
use teloxide::types::Message;
|
||||
|
||||
/// One URL job: the message + the extracted URL (the sender and stores come
|
||||
/// from the shared [`AppContext`], assembled from statics inside the worker).
|
||||
type UrlJob = (Message, String);
|
||||
/// Bounded channel of URL jobs drained by [`start_url_workers`]. The bound
|
||||
/// caps both queued memory and shutdown backlog; a full channel applies
|
||||
/// backpressure to the per-chat handler instead of spawning unbounded tasks.
|
||||
pub(crate) static URL_JOBS: LazyLock<
|
||||
parking_lot::Mutex<Option<tokio::sync::mpsc::Sender<UrlJob>>>,
|
||||
> = LazyLock::new(|| parking_lot::Mutex::new(None));
|
||||
/// Set by main's shutdown sequence; workers stop pulling new jobs.
|
||||
pub(crate) static URL_STOP: std::sync::atomic::AtomicBool =
|
||||
std::sync::atomic::AtomicBool::new(false);
|
||||
|
||||
/// JoinHandles of the URL workers, awaited by [`stop_url_workers`].
|
||||
static URL_WORKER_HANDLES: LazyLock<parking_lot::Mutex<Option<Vec<tokio::task::JoinHandle<()>>>>> =
|
||||
LazyLock::new(|| parking_lot::Mutex::new(None));
|
||||
|
||||
/// Worker count draining URL jobs; keeps the old 8-permit concurrency cap
|
||||
/// while bounding how many jobs can be queued at all.
|
||||
const URL_WORKERS: usize = 8;
|
||||
|
||||
/// 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
|
||||
/// processed one at a time (fetch + send each, roughly a second per
|
||||
/// message); the workers add throughput, and FIFO order preserves per-message
|
||||
/// URL order.
|
||||
pub async fn start_url_workers() {
|
||||
let (tx, rx) = tokio::sync::mpsc::channel::<UrlJob>(256);
|
||||
*URL_JOBS.lock() = Some(tx);
|
||||
let rx = std::sync::Arc::new(tokio::sync::Mutex::new(rx));
|
||||
let mut handles = Vec::with_capacity(URL_WORKERS);
|
||||
for _ in 0..URL_WORKERS {
|
||||
let rx = std::sync::Arc::clone(&rx);
|
||||
handles.push(tokio::spawn(async move {
|
||||
// Supervised like the queue workers: a panic inside a worker
|
||||
// (a handler, a poisoned lock) used to kill it for good and
|
||||
// silently shrink the pool — the remaining workers keep the
|
||||
// channel drained, so nothing else surfaces the loss. The job the
|
||||
// panicking worker held is lost; the panic is not.
|
||||
while !URL_STOP.load(std::sync::atomic::Ordering::Relaxed) {
|
||||
let rx = std::sync::Arc::clone(&rx);
|
||||
if let Err(e) = tokio::spawn(async move {
|
||||
while !URL_STOP.load(std::sync::atomic::Ordering::Relaxed) {
|
||||
let job = rx.lock().await.recv().await;
|
||||
match job {
|
||||
Some((message, url)) => {
|
||||
url_media(
|
||||
&CONTEXT,
|
||||
message.chat.id.0,
|
||||
message.id.0 as i64,
|
||||
&url,
|
||||
PostSend::FromChat,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
{
|
||||
log::error!("url worker panicked, restarting: {e}");
|
||||
}
|
||||
}
|
||||
}));
|
||||
}
|
||||
*URL_WORKER_HANDLES.lock() = Some(handles);
|
||||
}
|
||||
|
||||
/// Stops the URL workers: sets the stop flag, drops the job channel (so
|
||||
/// workers blocked in \`recv()\` wake with \`None\` and exit) and awaits the
|
||||
/// worker tasks. Each worker finishes its in-flight job first; jobs still
|
||||
/// queued in the channel are abandoned (the old implementation neither
|
||||
/// drained them nor woke blocked workers — it only set a flag checked
|
||||
/// between jobs).
|
||||
pub async fn stop_url_workers() {
|
||||
URL_STOP.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
// Dropping the sender makes every worker's recv() return None.
|
||||
*URL_JOBS.lock() = None;
|
||||
// Take the handles first so the lock guard drops before the awaits.
|
||||
let handles = URL_WORKER_HANDLES.lock().take();
|
||||
if let Some(handles) = handles {
|
||||
for handle in handles {
|
||||
if let Err(e) = handle.await {
|
||||
log::error!("url worker panicked at shutdown: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2,7 +2,7 @@
|
||||
//! worker pool, link-cache fast path, fetch, task build and send dispatch.
|
||||
|
||||
use super::{log_key, reply};
|
||||
use crate::ctx::{AppContext, CONTEXT};
|
||||
use crate::ctx::AppContext;
|
||||
use crate::link_cache::{CachedMedia, CachedMediaKind, CachedPost};
|
||||
use crate::media_sender::MediaSender;
|
||||
use crate::send::{self, Delivery, MediaItemPayload, Task};
|
||||
@@ -14,27 +14,6 @@ use teloxide::RequestError;
|
||||
use teloxide::types::{ChatAction, ChatId, Message, MessageEntityKind, MessageId};
|
||||
use x_media::media::Media;
|
||||
|
||||
/// One URL job: the message + the extracted URL (the sender and stores come
|
||||
/// from the shared [`AppContext`], assembled from statics inside the worker).
|
||||
type UrlJob = (Message, String);
|
||||
/// Bounded channel of URL jobs drained by [`start_url_workers`]. The bound
|
||||
/// caps both queued memory and shutdown backlog; a full channel applies
|
||||
/// backpressure to the per-chat handler instead of spawning unbounded tasks.
|
||||
pub(crate) static URL_JOBS: LazyLock<
|
||||
parking_lot::Mutex<Option<tokio::sync::mpsc::Sender<UrlJob>>>,
|
||||
> = LazyLock::new(|| parking_lot::Mutex::new(None));
|
||||
/// Set by main's shutdown sequence; workers stop pulling new jobs.
|
||||
pub(crate) static URL_STOP: std::sync::atomic::AtomicBool =
|
||||
std::sync::atomic::AtomicBool::new(false);
|
||||
|
||||
/// JoinHandles of the URL workers, awaited by [`stop_url_workers`].
|
||||
static URL_WORKER_HANDLES: LazyLock<parking_lot::Mutex<Option<Vec<tokio::task::JoinHandle<()>>>>> =
|
||||
LazyLock::new(|| parking_lot::Mutex::new(None));
|
||||
|
||||
/// Worker count draining URL jobs; keeps the old 8-permit concurrency cap
|
||||
/// 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
|
||||
@@ -120,76 +99,6 @@ impl<T> Drop for InFlightFetch<'_, T> {
|
||||
}
|
||||
}
|
||||
|
||||
/// 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
|
||||
/// processed one at a time (fetch + send each, roughly a second per
|
||||
/// message); the workers add throughput, and FIFO order preserves per-message
|
||||
/// URL order.
|
||||
pub async fn start_url_workers() {
|
||||
let (tx, rx) = tokio::sync::mpsc::channel::<UrlJob>(256);
|
||||
*URL_JOBS.lock() = Some(tx);
|
||||
let rx = std::sync::Arc::new(tokio::sync::Mutex::new(rx));
|
||||
let mut handles = Vec::with_capacity(URL_WORKERS);
|
||||
for _ in 0..URL_WORKERS {
|
||||
let rx = std::sync::Arc::clone(&rx);
|
||||
handles.push(tokio::spawn(async move {
|
||||
// Supervised like the queue workers: a panic inside a worker
|
||||
// (a handler, a poisoned lock) used to kill it for good and
|
||||
// silently shrink the pool — the remaining workers keep the
|
||||
// channel drained, so nothing else surfaces the loss. The job the
|
||||
// panicking worker held is lost; the panic is not.
|
||||
while !URL_STOP.load(std::sync::atomic::Ordering::Relaxed) {
|
||||
let rx = std::sync::Arc::clone(&rx);
|
||||
if let Err(e) = tokio::spawn(async move {
|
||||
while !URL_STOP.load(std::sync::atomic::Ordering::Relaxed) {
|
||||
let job = rx.lock().await.recv().await;
|
||||
match job {
|
||||
Some((message, url)) => {
|
||||
url_media(
|
||||
&CONTEXT,
|
||||
message.chat.id.0,
|
||||
message.id.0 as i64,
|
||||
&url,
|
||||
PostSend::FromChat,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
{
|
||||
log::error!("url worker panicked, restarting: {e}");
|
||||
}
|
||||
}
|
||||
}));
|
||||
}
|
||||
*URL_WORKER_HANDLES.lock() = Some(handles);
|
||||
}
|
||||
|
||||
/// Stops the URL workers: sets the stop flag, drops the job channel (so
|
||||
/// workers blocked in \`recv()\` wake with \`None\` and exit) and awaits the
|
||||
/// worker tasks. Each worker finishes its in-flight job first; jobs still
|
||||
/// queued in the channel are abandoned (the old implementation neither
|
||||
/// drained them nor woke blocked workers — it only set a flag checked
|
||||
/// between jobs).
|
||||
pub async fn stop_url_workers() {
|
||||
URL_STOP.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
// Dropping the sender makes every worker's recv() return None.
|
||||
*URL_JOBS.lock() = None;
|
||||
// Take the handles first so the lock guard drops before the awaits.
|
||||
let handles = URL_WORKER_HANDLES.lock().take();
|
||||
if let Some(handles) = handles {
|
||||
for handle in handles {
|
||||
if let Err(e) = handle.await {
|
||||
log::error!("url worker panicked at shutdown: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Extracts URL and text-link entities (text + caption), deduped in order.
|
||||
///
|
||||
/// The offset work (`parse_entities` turning entities into slices of the
|
||||
@@ -253,7 +162,27 @@ fn thumbnail_for(media: &Media) -> Option<String> {
|
||||
}
|
||||
}
|
||||
|
||||
fn media_to_payload(media: &Media, sensitive: bool) -> MediaItemPayload {
|
||||
/// The link-cache snapshot of a freshly fetched post: its caption and raw
|
||||
/// render fields, with no media yet — the send fills in the Telegram file ids
|
||||
/// and persists the entry. `None` for a post that carries no render data
|
||||
/// (nothing to rebuild a caption format from later).
|
||||
pub(super) fn cached_snapshot(fetched: &x_media::site::Fetched) -> Option<CachedPost> {
|
||||
fetched
|
||||
.render_fields()
|
||||
.map(|(author, author_url, title, content, tags)| CachedPost {
|
||||
url: fetched.source_url.clone(),
|
||||
caption: fetched.caption.clone(),
|
||||
title: title.to_string(),
|
||||
content: content.to_string(),
|
||||
author: author.to_string(),
|
||||
author_url: author_url.to_string(),
|
||||
tags: tags.to_string(),
|
||||
sensitive: fetched.sensitive,
|
||||
media: vec![],
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn media_to_payload(media: &Media, sensitive: bool) -> MediaItemPayload {
|
||||
let fallback_url = media.smaller_url().map(str::to_string);
|
||||
match media {
|
||||
// A gif inside a group becomes a video item; a lone gif takes the
|
||||
@@ -707,20 +636,7 @@ async fn url_media_inner(
|
||||
let caption = fetched.caption_with(&format);
|
||||
// Raw render data for the link cache; the send fills in the
|
||||
// Telegram file ids and persists the entry.
|
||||
let cache_data =
|
||||
fetched
|
||||
.render_fields()
|
||||
.map(|(author, author_url, title, content, tags)| CachedPost {
|
||||
url: fetched.source_url.clone(),
|
||||
caption: fetched.caption.clone(),
|
||||
title: title.to_string(),
|
||||
content: content.to_string(),
|
||||
author: author.to_string(),
|
||||
author_url: author_url.to_string(),
|
||||
tags: tags.to_string(),
|
||||
sensitive: fetched.sensitive,
|
||||
media: vec![],
|
||||
});
|
||||
let cache_data = cached_snapshot(fetched);
|
||||
let items: Vec<MediaItemPayload> = fetched
|
||||
.media
|
||||
.iter()
|
||||
@@ -752,207 +668,12 @@ async fn url_media_inner(
|
||||
}
|
||||
}
|
||||
|
||||
// ── Startup repair: queued retries whose local media did not survive ───────
|
||||
|
||||
/// A post's fresh media plus the caption and cache snapshot that go with them:
|
||||
/// what [`refetch`] hands [`apply_refresh`]. Plain data, so the rewrite below
|
||||
/// can be tested without a network fetch (which cannot be faked here:
|
||||
/// [`x_media::site::Fetched`] keeps a private field and is not constructible
|
||||
/// outside its crate).
|
||||
struct Refetched {
|
||||
caption: String,
|
||||
items: Vec<MediaItemPayload>,
|
||||
cache_data: Option<CachedPost>,
|
||||
}
|
||||
|
||||
/// Whether a queued task should have its post re-fetched, because it still
|
||||
/// wants a local file (ugoira MP4, a bsky remux, a downloaded temp file) that is
|
||||
/// gone. Those files live in the system temp dir and the registry that keeps
|
||||
/// them alive for the retry (`send::KEEP_ALIVE`) is in memory, so a restart
|
||||
/// takes all of them — a retry that needs one can only dead-letter.
|
||||
///
|
||||
/// A partially delivered album is left alone: its remaining batches cannot be
|
||||
/// reconciled with a fresh media list without risking a second copy of what the
|
||||
/// user already received.
|
||||
fn needs_refetch(task: &Task) -> bool {
|
||||
if let Task::SendMediaSequence {
|
||||
batch_index,
|
||||
sent_message_ids,
|
||||
..
|
||||
} = task
|
||||
&& (*batch_index > 0 || !sent_message_ids.is_empty())
|
||||
{
|
||||
return false;
|
||||
}
|
||||
task.local_media_paths().iter().any(|path| !path.exists())
|
||||
}
|
||||
|
||||
/// Rebuilds the task from the fresh media, keeping its delivery envelope (chat,
|
||||
/// reply, forward/edit settings, notify targets): the retry that was queued must
|
||||
/// still deliver the same way, whoever asked for it.
|
||||
fn apply_refresh(task: &Task, fresh: &Refetched) -> Option<Task> {
|
||||
let chat_id = task.chat_id()?;
|
||||
let (edit_before_forward, forward_channel_id) = match task {
|
||||
Task::SendMediaSequence {
|
||||
edit_before_forward,
|
||||
forward_channel_id,
|
||||
..
|
||||
}
|
||||
| Task::SendAnimation {
|
||||
edit_before_forward,
|
||||
forward_channel_id,
|
||||
..
|
||||
} => (*edit_before_forward, *forward_channel_id),
|
||||
Task::ForwardMessages { .. } => return None,
|
||||
};
|
||||
let reply_to_message_id = match task {
|
||||
Task::SendMediaSequence {
|
||||
reply_to_message_id,
|
||||
..
|
||||
}
|
||||
| Task::SendAnimation {
|
||||
reply_to_message_id,
|
||||
..
|
||||
} => *reply_to_message_id,
|
||||
Task::ForwardMessages { .. } => return None,
|
||||
};
|
||||
let (notify_chat_id, notify_message_id) = task.notify_target();
|
||||
Some(Task::from_items(
|
||||
Delivery {
|
||||
chat_id,
|
||||
reply_to_message_id,
|
||||
edit_before_forward,
|
||||
forward_channel_id,
|
||||
notify_chat_id,
|
||||
notify_message_id,
|
||||
},
|
||||
task.source_url()?.to_string(),
|
||||
fresh.caption.clone(),
|
||||
fresh.items.clone(),
|
||||
fresh.cache_data.clone(),
|
||||
))
|
||||
}
|
||||
|
||||
/// Fetches the post again and maps it into [`Refetched`]: the same mapping the
|
||||
/// fresh-fetch path uses (per-site caption format from the chat, render fields
|
||||
/// for the link-cache snapshot), so a repaired task looks like a first send.
|
||||
async fn refetch(
|
||||
ctx: &AppContext<'_>,
|
||||
chat_id: i64,
|
||||
url: &str,
|
||||
) -> Result<Option<Refetched>, x_media::site::FetchError> {
|
||||
let Some(fetched) = x_media::site::fetch(url).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
if fetched.media.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
let chat_data = ctx.chat_store.get(chat_id).await;
|
||||
let format = chat_data.format_for(fetched.site_id);
|
||||
let caption = fetched.caption_with(&format);
|
||||
let cache_data = fetched
|
||||
.render_fields()
|
||||
.map(|(author, author_url, title, content, tags)| CachedPost {
|
||||
url: fetched.source_url.clone(),
|
||||
caption: fetched.caption.clone(),
|
||||
title: title.to_string(),
|
||||
content: content.to_string(),
|
||||
author: author.to_string(),
|
||||
author_url: author_url.to_string(),
|
||||
tags: tags.to_string(),
|
||||
sensitive: fetched.sensitive,
|
||||
media: vec![],
|
||||
});
|
||||
let items: Vec<MediaItemPayload> = fetched
|
||||
.media
|
||||
.iter()
|
||||
.map(|media| media_to_payload(media, fetched.sensitive))
|
||||
.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.keep_alive() {
|
||||
send::KEEP_ALIVE.lock().push(dir);
|
||||
}
|
||||
Ok(Some(Refetched {
|
||||
caption,
|
||||
items,
|
||||
cache_data,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Re-fetches every queued task whose local media did not survive the restart,
|
||||
/// so the user's link is still delivered instead of dead-lettering on a file
|
||||
/// that cannot come back. Returns how many rows were rewritten.
|
||||
///
|
||||
/// Startup only, before the queue workers start: no worker can lease a row while
|
||||
/// this writes, which is what lets it replace payloads without the lease-token
|
||||
/// guard every worker write-back carries.
|
||||
pub(crate) async fn repair_lost_local_media(ctx: &AppContext<'_>) -> usize {
|
||||
let mut repaired = 0;
|
||||
for (id, payload) in ctx.task_queue.runnable_rows().await {
|
||||
let Ok(task) = serde_json::from_str::<Task>(&payload) else {
|
||||
continue;
|
||||
};
|
||||
if !needs_refetch(&task) {
|
||||
continue;
|
||||
}
|
||||
let (Some(url), Some(chat_id)) = (task.source_url().map(str::to_string), task.chat_id())
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
match refetch(ctx, chat_id, &url).await {
|
||||
Ok(Some(fresh)) => {
|
||||
let Some(updated) = apply_refresh(&task, &fresh) else {
|
||||
continue;
|
||||
};
|
||||
let updated = serde_json::to_value(&updated).expect("task serializes");
|
||||
if ctx.task_queue.replace_payload(&id, &updated).await {
|
||||
repaired += 1;
|
||||
log::info!(
|
||||
"startup repair: re-fetched [key={}] for chat={chat_id} (its local media did not survive the restart)",
|
||||
log_key(&url)
|
||||
);
|
||||
}
|
||||
}
|
||||
// The post is gone or withheld now: the retry could not have
|
||||
// delivered anything either, so say why instead of letting it
|
||||
// dead-letter on a missing file.
|
||||
Ok(None) | Err(_) => {
|
||||
let (notify_chat_id, notify_message_id) = task.notify_target();
|
||||
log::warn!(
|
||||
"startup repair: [key={}] for chat={chat_id} needed a re-fetch and none was possible",
|
||||
log_key(&url)
|
||||
);
|
||||
send::notify_failure(
|
||||
ctx.sender,
|
||||
notify_chat_id,
|
||||
notify_message_id,
|
||||
&format!(
|
||||
"{} — the media held for retry was lost when the bot restarted and the post could not be fetched again. Please send the link again.",
|
||||
log_key(&url)
|
||||
),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
repaired
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::ctx::test_support::{TestStores, api_error, cached_photo, photo_item};
|
||||
use crate::ctx::test_support::{TestStores, cached_photo, permanent_error};
|
||||
use crate::media_sender::test_support::{MockSender, Outcome};
|
||||
use std::time::Duration;
|
||||
use teloxide::RequestError;
|
||||
|
||||
/// The API error a caption edit that changes nothing answers with — what
|
||||
/// the mocks script for a permanent send failure. A `fn` pointer, so it
|
||||
/// can be handed to `MockSender::scripted` as-is.
|
||||
fn permanent_error() -> RequestError {
|
||||
api_error("Bad Request: message is not modified")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cache_hit_sends_file_ids_and_degrades_on_permanent_failure() {
|
||||
@@ -1359,168 +1080,6 @@ mod tests {
|
||||
assert!(dedupe_urls(vec![]).is_empty());
|
||||
}
|
||||
|
||||
fn queued_task(media: &str, batch_index: usize, sent: Vec<i64>) -> Task {
|
||||
Task::SendMediaSequence {
|
||||
chat_id: 1,
|
||||
reply_to_message_id: 2,
|
||||
caption: "cap".into(),
|
||||
media_batches: vec![vec![photo_item(media, false, false)]],
|
||||
batch_index,
|
||||
sent_message_ids: sent,
|
||||
source_url: "https://x.com/u/status/1".into(),
|
||||
edit_before_forward: true,
|
||||
forward_channel_id: Some(2),
|
||||
notify_chat_id: Some(1),
|
||||
notify_message_id: Some(2),
|
||||
cache_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_tasks_missing_a_local_file_need_a_refetch() {
|
||||
// A URL send needs nothing.
|
||||
assert!(!needs_refetch(&queued_task("https://cdn/1.jpg", 0, vec![])));
|
||||
// A local path that is still there (a survived temp file) needs nothing.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let alive = dir.path().join("ugoira.mp4");
|
||||
std::fs::write(&alive, b"x").unwrap();
|
||||
assert!(!needs_refetch(&queued_task(
|
||||
alive.to_str().unwrap(),
|
||||
0,
|
||||
vec![]
|
||||
)));
|
||||
// A local path the restart took away does.
|
||||
assert!(needs_refetch(&queued_task(
|
||||
"/nonexistent-ugoira.mp4",
|
||||
0,
|
||||
vec![]
|
||||
)));
|
||||
// A partially delivered album is left to its own retry path.
|
||||
assert!(!needs_refetch(&queued_task(
|
||||
"/nonexistent-ugoira.mp4",
|
||||
1,
|
||||
vec![7]
|
||||
)));
|
||||
assert!(!needs_refetch(&queued_task(
|
||||
"/nonexistent-ugoira.mp4",
|
||||
0,
|
||||
vec![7]
|
||||
)));
|
||||
// A channel copy holds no media.
|
||||
assert!(!needs_refetch(&Task::ForwardMessages {
|
||||
from_chat_id: 1,
|
||||
to_chat_id: 2,
|
||||
message_ids: vec![3],
|
||||
notify_chat_id: None,
|
||||
notify_message_id: None,
|
||||
}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apply_refresh_keeps_the_delivery_envelope() {
|
||||
let task = queued_task("/nonexistent-ugoira.mp4", 0, vec![]);
|
||||
let fresh = Refetched {
|
||||
caption: "fresh caption".into(),
|
||||
items: vec![photo_item("https://cdn/fresh.jpg", true, false)],
|
||||
cache_data: None,
|
||||
};
|
||||
match apply_refresh(&task, &fresh).expect("a repairable task") {
|
||||
Task::SendMediaSequence {
|
||||
chat_id,
|
||||
reply_to_message_id,
|
||||
caption,
|
||||
media_batches,
|
||||
batch_index,
|
||||
sent_message_ids,
|
||||
source_url,
|
||||
edit_before_forward,
|
||||
forward_channel_id,
|
||||
notify_chat_id,
|
||||
notify_message_id,
|
||||
..
|
||||
} => {
|
||||
// Same delivery: chat, reply, forward/edit settings, notify.
|
||||
assert_eq!((chat_id, reply_to_message_id), (1, 2));
|
||||
assert!(edit_before_forward);
|
||||
assert_eq!(forward_channel_id, Some(2));
|
||||
assert_eq!((notify_chat_id, notify_message_id), (Some(1), Some(2)));
|
||||
assert_eq!(source_url, "https://x.com/u/status/1");
|
||||
// Fresh media, and nothing of it counted as sent yet.
|
||||
assert_eq!(caption, "fresh caption");
|
||||
assert!(
|
||||
matches!(
|
||||
&media_batches[0][0],
|
||||
MediaItemPayload::Photo { media, .. } if media == "https://cdn/fresh.jpg"
|
||||
),
|
||||
"fresh media must replace the lost local file"
|
||||
);
|
||||
assert!(matches!(
|
||||
media_batches[0][0],
|
||||
MediaItemPayload::Photo {
|
||||
has_spoiler: true,
|
||||
..
|
||||
}
|
||||
));
|
||||
assert_eq!((batch_index, sent_message_ids.len()), (0, 0));
|
||||
}
|
||||
other => panic!("expected a media sequence, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// The whole repair against a real post: a queued row whose media is a local
|
||||
/// file the restart took away is re-fetched from its `source_url` and
|
||||
/// rewritten in place, so the retry can still deliver it.
|
||||
#[tokio::test]
|
||||
#[ignore = "live network: requires outbound HTTPS to public.api.bsky.app"]
|
||||
async fn live_repair_refetches_a_lost_local_media_row() {
|
||||
let stores = TestStores::new();
|
||||
// An empty script: the repair must not need to tell the user anything.
|
||||
let sender = MockSender::scripted(vec![], permanent_error);
|
||||
let ctx = stores.ctx(&sender);
|
||||
let mut task = queued_task("/nonexistent-ugoira.mp4", 0, vec![]);
|
||||
if let Task::SendMediaSequence { source_url, .. } = &mut task {
|
||||
*source_url = "https://bsky.app/profile/fu-futa.bsky.social/post/3laoveufjv224".into();
|
||||
}
|
||||
stores
|
||||
.task_queue()
|
||||
.enqueue(serde_json::to_value(&task).unwrap(), crate::db::now_f64())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(repair_lost_local_media(&ctx).await, 1);
|
||||
|
||||
let updated: Task = serde_json::from_value(stores.queued_payload().await).unwrap();
|
||||
match updated {
|
||||
Task::SendMediaSequence {
|
||||
media_batches,
|
||||
batch_index,
|
||||
sent_message_ids,
|
||||
caption,
|
||||
..
|
||||
} => {
|
||||
let media: Vec<String> = media_batches
|
||||
.iter()
|
||||
.flatten()
|
||||
.map(|item| match item {
|
||||
MediaItemPayload::Photo { media, .. }
|
||||
| MediaItemPayload::Video { media, .. }
|
||||
| MediaItemPayload::Animation { media, .. } => media.clone(),
|
||||
})
|
||||
.collect();
|
||||
assert!(!media.is_empty(), "the fresh fetch yielded no media");
|
||||
assert!(
|
||||
media.iter().all(|m| m.starts_with("http")),
|
||||
"the retry must be uploadable from URLs again: {media:?}"
|
||||
);
|
||||
assert_eq!((batch_index, sent_message_ids.len()), (0, 0));
|
||||
assert!(!caption.is_empty());
|
||||
}
|
||||
other => panic!("expected a repaired media sequence, got {other:?}"),
|
||||
}
|
||||
// The post was re-read, not re-delivered: nothing was sent.
|
||||
assert!(sender.calls().is_empty(), "{:?}", sender.calls());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fetch_errors_map_to_distinct_user_messages() {
|
||||
use x_media::site::FetchError;
|
||||
|
||||
Reference in New Issue
Block a user