mirror of
https://github.com/TheFunny/TelegramTwitterMediaBot.git
synced 2026-09-23 23:32:05 +00:00
561 lines
21 KiB
Rust
561 lines
21 KiB
Rust
//! URL extraction and the per-URL media pipeline: bounded job channel +
|
|
//! worker pool, link-cache fast path, fetch, task build and send dispatch.
|
|
|
|
use super::{CHAT_STORE, CONFIG, LINK_CACHE, TASK_QUEUE, log_key, reply};
|
|
use crate::config::Config;
|
|
use crate::db::now_f64;
|
|
use crate::link_cache::{CachedMediaKind, CachedPost, LinkCache};
|
|
use crate::media_sender::MediaSender;
|
|
use crate::queue::PersistentTaskQueue;
|
|
use crate::send::{self, MediaItemPayload, Task};
|
|
use crate::state::{ChatData, ChatStore};
|
|
use std::collections::HashSet;
|
|
use std::sync::LazyLock;
|
|
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;
|
|
|
|
/// Dependencies of the per-URL pipeline, injected so tests can substitute a
|
|
/// mock sender and tempdir-backed stores.
|
|
pub(crate) struct AppContext<'a> {
|
|
pub sender: &'a dyn MediaSender,
|
|
pub chat_store: &'a ChatStore,
|
|
pub task_queue: &'a PersistentTaskQueue,
|
|
pub link_cache: &'a LinkCache,
|
|
pub config: &'a Config,
|
|
}
|
|
|
|
/// Assembles the production context from the process-wide statics.
|
|
fn app_context() -> AppContext<'static> {
|
|
AppContext {
|
|
sender: &*crate::send::BOT,
|
|
chat_store: &CHAT_STORE,
|
|
task_queue: &TASK_QUEUE,
|
|
link_cache: &LINK_CACHE,
|
|
config: &CONFIG,
|
|
}
|
|
}
|
|
|
|
/// 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 {
|
|
let ctx = app_context();
|
|
while !URL_STOP.load(std::sync::atomic::Ordering::Relaxed) {
|
|
let job = rx.lock().await.recv().await;
|
|
match job {
|
|
Some((message, url)) => {
|
|
url_media(&ctx, message.chat.id.0, message.id.0 as i64, &url).await
|
|
}
|
|
None => break,
|
|
}
|
|
}
|
|
}));
|
|
}
|
|
*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 {
|
|
let _ = handle.await;
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Extracts URL and text-link entities (text + caption), deduped in order.
|
|
pub fn extract_urls(message: &Message) -> Vec<String> {
|
|
let mut urls = Vec::new();
|
|
for entity in message.parse_entities().into_iter().flatten() {
|
|
match entity.kind() {
|
|
MessageEntityKind::Url => urls.push(entity.text().to_string()),
|
|
MessageEntityKind::TextLink { url } => urls.push(url.to_string()),
|
|
_ => {}
|
|
}
|
|
}
|
|
for entity in message.parse_caption_entities().into_iter().flatten() {
|
|
match entity.kind() {
|
|
MessageEntityKind::Url => urls.push(entity.text().to_string()),
|
|
MessageEntityKind::TextLink { url } => urls.push(url.to_string()),
|
|
_ => {}
|
|
}
|
|
}
|
|
let mut seen = HashSet::new();
|
|
// Dedup by the normalized post id so variant URLs of the same post
|
|
// (/status/1 vs /status/1/photo/1) are sent once; unsupported URLs fall
|
|
// back to exact-string dedup.
|
|
urls.retain(|url| seen.insert(x_media::site::cache_key(url).unwrap_or_else(|| url.clone())));
|
|
urls
|
|
}
|
|
|
|
/// For locally produced media (encoded ugoira MP4) the thumbnail URL is a
|
|
/// hotlink-protected remote URL Telegram may not fetch; let Telegram generate
|
|
/// its own thumbnail instead.
|
|
fn thumbnail_for(media: &Media) -> Option<String> {
|
|
let url = media.url();
|
|
if url.starts_with("http://") || url.starts_with("https://") {
|
|
media.thumbnail_url().map(str::to_string)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
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
|
|
// animation path (see url_media).
|
|
Media::Illustration { .. } => MediaItemPayload::Photo {
|
|
media: media.url().to_string(),
|
|
has_spoiler: sensitive,
|
|
fallback_url,
|
|
file_id: false,
|
|
},
|
|
Media::Video { .. } => MediaItemPayload::Video {
|
|
media: media.url().to_string(),
|
|
has_spoiler: sensitive,
|
|
thumbnail: thumbnail_for(media),
|
|
fallback_url,
|
|
file_id: false,
|
|
},
|
|
Media::Animated { .. } => MediaItemPayload::Video {
|
|
media: media.url().to_string(),
|
|
has_spoiler: sensitive,
|
|
thumbnail: thumbnail_for(media),
|
|
fallback_url,
|
|
file_id: false,
|
|
},
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn enqueue_retry(queue: &PersistentTaskQueue, task: Task, delay_seconds: f64) {
|
|
let payload = serde_json::to_value(task).expect("task serializes");
|
|
let run_after = now_f64() + delay_seconds;
|
|
if let Err(e) = queue.enqueue(payload, run_after).await {
|
|
log::error!("failed to enqueue retry: {e}");
|
|
}
|
|
}
|
|
|
|
/// Sends a task and handles the outcome: post-send actions on success, retry
|
|
/// enqueue on retryable failure, reply + link-cache invalidation on
|
|
/// permanent failure (a stale cached file id must not repeat forever).
|
|
async fn dispatch_send(
|
|
ctx: &AppContext<'_>,
|
|
chat_id: i64,
|
|
reply_to: MessageId,
|
|
task: &Task,
|
|
url: &str,
|
|
) {
|
|
let result = match task {
|
|
Task::SendAnimation { .. } => send::send_animation(ctx.sender, task).await,
|
|
Task::SendMediaSequence { .. } => send::send_media_sequence(ctx.sender, task).await,
|
|
Task::ForwardMessages { .. } => unreachable!(),
|
|
};
|
|
match result {
|
|
Ok(message_ids) => {
|
|
log::info!(
|
|
"sent {} message(s) for [key={}]",
|
|
message_ids.len(),
|
|
log_key(url)
|
|
);
|
|
send::post_send_actions(ctx.sender, task, message_ids).await;
|
|
// The task settled: drop any keep-alive temp media.
|
|
send::release_keep_alive(task);
|
|
}
|
|
Err(send::SendError::Retryable {
|
|
delay_seconds,
|
|
task,
|
|
}) => {
|
|
log::info!(
|
|
"send for [key={}] failed, queued for retry in {delay_seconds:.1}s",
|
|
log_key(url)
|
|
);
|
|
enqueue_retry(ctx.task_queue, task, delay_seconds).await;
|
|
let _ = reply(
|
|
ctx.sender,
|
|
chat_id,
|
|
reply_to,
|
|
"Send failed. Task queued for retry.",
|
|
)
|
|
.await;
|
|
}
|
|
Err(send::SendError::Permanent {
|
|
message: err_message,
|
|
task,
|
|
}) => {
|
|
send::invalidate_cache_with(ctx.link_cache, &task).await;
|
|
send::release_keep_alive(&task);
|
|
log::error!("send for {url} failed permanently: {err_message}");
|
|
let _ = reply(
|
|
ctx.sender,
|
|
chat_id,
|
|
reply_to,
|
|
format!("Send failed: {err_message}"),
|
|
)
|
|
.await;
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Builds the send task from ready-made items, sharing the payload shape
|
|
/// between the fresh-fetch and link-cache paths.
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn build_send_task(
|
|
chat_data: &ChatData,
|
|
chat_id: i64,
|
|
reply_to_message_id: i64,
|
|
source_url: String,
|
|
caption: String,
|
|
items: Vec<MediaItemPayload>,
|
|
cache_data: Option<CachedPost>,
|
|
) -> Task {
|
|
if items.len() == 1 && matches!(items[0], MediaItemPayload::Animation { .. }) {
|
|
Task::SendAnimation {
|
|
chat_id,
|
|
reply_to_message_id,
|
|
caption,
|
|
animation: items.into_iter().next().unwrap(),
|
|
source_url,
|
|
edit_before_forward: chat_data.edit_before_forward,
|
|
forward_channel_id: chat_data.forward_channel_id,
|
|
notify_chat_id: Some(chat_id),
|
|
notify_message_id: Some(reply_to_message_id),
|
|
cache_data,
|
|
}
|
|
} else {
|
|
Task::SendMediaSequence {
|
|
chat_id,
|
|
reply_to_message_id,
|
|
caption,
|
|
// Photos first so a mixed photo+video group starts with a photo
|
|
// (Telegram's sendMediaGroup rule); order within each kind is kept.
|
|
media_batches: send::chunk_media_items(send::photos_first(items)),
|
|
batch_index: 0,
|
|
sent_message_ids: vec![],
|
|
source_url,
|
|
edit_before_forward: chat_data.edit_before_forward,
|
|
forward_channel_id: chat_data.forward_channel_id,
|
|
notify_chat_id: Some(chat_id),
|
|
notify_message_id: Some(reply_to_message_id),
|
|
cache_data,
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn url_media(ctx: &AppContext<'_>, chat_id: i64, reply_to_message_id: i64, url: &str) {
|
|
let reply_to = MessageId(reply_to_message_id as i32);
|
|
if let Err(e) = ctx
|
|
.sender
|
|
.send_chat_action(ChatId(chat_id), ChatAction::Typing)
|
|
.await
|
|
{
|
|
log::error!("send_chat_action failed: {e}");
|
|
}
|
|
|
|
// Link cache: a post sent before is re-sent from Telegram file ids —
|
|
// no source-site request, no download, no upload. Keyed by the
|
|
// normalized post id so x.com / fxtwitter / /photo/N variants collide.
|
|
if let Some(key) = x_media::site::cache_key(url)
|
|
&& let Some(cached) = ctx.link_cache.get(&key, ctx.config.link_cache_ttl).await
|
|
{
|
|
log::debug!("link cache hit for {key}");
|
|
let chat_data = ctx.chat_store.get(chat_id).await;
|
|
// Cache keys are prefixed with the site id ("twitter:…"), matching
|
|
// the value a fresh fetch would read from Fetched::site_id.
|
|
let site = x_media::site::site_id_from_key(&key);
|
|
let format = chat_data
|
|
.message_format
|
|
.get(site)
|
|
.cloned()
|
|
.unwrap_or_default();
|
|
let caption = if format.is_empty() {
|
|
x_media::site::truncate_caption(&cached.caption)
|
|
} else {
|
|
x_media::site::caption_from_fields(
|
|
&format,
|
|
"",
|
|
&cached.url,
|
|
&cached.author,
|
|
&cached.author_url,
|
|
&cached.title,
|
|
&cached.tags,
|
|
)
|
|
};
|
|
let items: Vec<MediaItemPayload> = cached
|
|
.media
|
|
.iter()
|
|
.map(|m| match m.kind {
|
|
CachedMediaKind::Photo => MediaItemPayload::Photo {
|
|
media: m.file_id.clone(),
|
|
has_spoiler: cached.sensitive,
|
|
fallback_url: None,
|
|
file_id: true,
|
|
},
|
|
CachedMediaKind::Video => MediaItemPayload::Video {
|
|
media: m.file_id.clone(),
|
|
has_spoiler: cached.sensitive,
|
|
thumbnail: None,
|
|
fallback_url: None,
|
|
file_id: true,
|
|
},
|
|
CachedMediaKind::Animation => MediaItemPayload::Animation {
|
|
media: m.file_id.clone(),
|
|
has_spoiler: cached.sensitive,
|
|
file_id: true,
|
|
},
|
|
})
|
|
.collect();
|
|
let task = build_send_task(
|
|
&chat_data,
|
|
chat_id,
|
|
reply_to_message_id,
|
|
cached.url.clone(),
|
|
caption,
|
|
items,
|
|
Some(cached),
|
|
);
|
|
dispatch_send(ctx, chat_id, reply_to, &task, url).await;
|
|
return;
|
|
}
|
|
|
|
log::debug!("fetching {url} [key={}]", log_key(url));
|
|
match x_media::site::fetch(url).await {
|
|
// Unsupported links are ignored silently (Python parity).
|
|
Ok(None) => {
|
|
log::debug!("no site pattern matches {url}; ignoring");
|
|
}
|
|
// Retries exhausted: notify the user (Rust-only requirement 3).
|
|
Err(e) => {
|
|
log::error!("fetch {url}: {e}");
|
|
let _ = reply(
|
|
ctx.sender,
|
|
chat_id,
|
|
reply_to,
|
|
"Failed to fetch media from this link.",
|
|
)
|
|
.await;
|
|
}
|
|
Ok(Some(mut fetched)) => {
|
|
if fetched.media.is_empty() {
|
|
let _ = reply(
|
|
ctx.sender,
|
|
chat_id,
|
|
reply_to,
|
|
"No media found or media type is not supported.",
|
|
)
|
|
.await;
|
|
return;
|
|
}
|
|
let chat_data = ctx.chat_store.get(chat_id).await;
|
|
// Per-site caption format override (empty -> built-in caption).
|
|
let format = chat_data
|
|
.message_format
|
|
.get(fetched.site_name())
|
|
.cloned()
|
|
.unwrap_or_default();
|
|
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, tags)| CachedPost {
|
|
url: fetched.source_url.clone(),
|
|
caption: fetched.caption.clone(),
|
|
title: title.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();
|
|
let task = build_send_task(
|
|
&chat_data,
|
|
chat_id,
|
|
reply_to_message_id,
|
|
fetched.source_url.clone(),
|
|
caption,
|
|
items,
|
|
cache_data,
|
|
);
|
|
// 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() {
|
|
send::KEEP_ALIVE.lock().push(dir);
|
|
}
|
|
dispatch_send(ctx, chat_id, reply_to, &task, url).await;
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::db;
|
|
use crate::link_cache::CachedMedia;
|
|
use crate::media_sender::test_support::{MockSender, Outcome};
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
use teloxide::{ApiError, RequestError};
|
|
|
|
fn permanent_error() -> RequestError {
|
|
RequestError::Api(ApiError::Unknown(
|
|
"Bad Request: message is not modified".into(),
|
|
))
|
|
}
|
|
|
|
fn cached_photo_entry() -> CachedPost {
|
|
CachedPost {
|
|
url: "https://x.com/u/status/1".into(),
|
|
caption: "cap".into(),
|
|
title: "t".into(),
|
|
author: "a".into(),
|
|
author_url: "au".into(),
|
|
tags: "".into(),
|
|
sensitive: false,
|
|
media: vec![CachedMedia {
|
|
kind: CachedMediaKind::Photo,
|
|
file_id: "file-1".into(),
|
|
}],
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn cache_hit_sends_file_ids_and_invalidates_on_permanent_failure() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let pool = db::open_store(dir.path().join("t.db").to_str().unwrap()).unwrap();
|
|
let chat_store = ChatStore::new(Arc::clone(&pool));
|
|
let task_queue = PersistentTaskQueue::new(Arc::clone(&pool));
|
|
let link_cache = LinkCache::new(Arc::clone(&pool));
|
|
let config = Config::load();
|
|
let sender = MockSender::scripted(
|
|
vec![Outcome::GroupErr, Outcome::MessageErr],
|
|
permanent_error,
|
|
);
|
|
let ctx = AppContext {
|
|
sender: &sender,
|
|
chat_store: &chat_store,
|
|
task_queue: &task_queue,
|
|
link_cache: &link_cache,
|
|
config: &config,
|
|
};
|
|
link_cache.put("twitter:1", &cached_photo_entry()).await;
|
|
|
|
url_media(&ctx, 1, 2, "https://x.com/u/status/1").await;
|
|
|
|
// The cached file id went out as a group send; the permanent failure
|
|
// then triggered the fire-and-forget reply (its mock error is fine).
|
|
assert_eq!(
|
|
sender.calls(),
|
|
vec!["send_chat_action", "send_media_group", "send_message"]
|
|
);
|
|
// The stale cache entry was invalidated so the next request re-fetches.
|
|
assert!(
|
|
link_cache
|
|
.get("twitter:1", Duration::from_secs(3600))
|
|
.await
|
|
.is_none()
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn cache_hit_success_keeps_the_cache_entry() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let pool = db::open_store(dir.path().join("t.db").to_str().unwrap()).unwrap();
|
|
let chat_store = ChatStore::new(Arc::clone(&pool));
|
|
let task_queue = PersistentTaskQueue::new(Arc::clone(&pool));
|
|
let link_cache = LinkCache::new(Arc::clone(&pool));
|
|
let config = Config::load();
|
|
let sender = MockSender::scripted(vec![Outcome::GroupOk], permanent_error);
|
|
let ctx = AppContext {
|
|
sender: &sender,
|
|
chat_store: &chat_store,
|
|
task_queue: &task_queue,
|
|
link_cache: &link_cache,
|
|
config: &config,
|
|
};
|
|
link_cache.put("twitter:1", &cached_photo_entry()).await;
|
|
|
|
url_media(&ctx, 1, 2, "https://x.com/u/status/1").await;
|
|
|
|
assert_eq!(sender.calls(), vec!["send_chat_action", "send_media_group"]);
|
|
// Success must not evict the entry.
|
|
assert!(
|
|
link_cache
|
|
.get("twitter:1", Duration::from_secs(3600))
|
|
.await
|
|
.is_some()
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn unsupported_url_is_ignored_silently() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let pool = db::open_store(dir.path().join("t.db").to_str().unwrap()).unwrap();
|
|
let chat_store = ChatStore::new(Arc::clone(&pool));
|
|
let task_queue = PersistentTaskQueue::new(Arc::clone(&pool));
|
|
let link_cache = LinkCache::new(Arc::clone(&pool));
|
|
let config = Config::load();
|
|
let sender = MockSender::scripted(vec![], permanent_error);
|
|
let ctx = AppContext {
|
|
sender: &sender,
|
|
chat_store: &chat_store,
|
|
task_queue: &task_queue,
|
|
link_cache: &link_cache,
|
|
config: &config,
|
|
};
|
|
|
|
// No cache key → the fetch dispatcher returns Ok(None) without any
|
|
// network; nothing is sent or replied.
|
|
url_media(&ctx, 1, 2, "https://example.com/not-a-post").await;
|
|
assert_eq!(sender.calls(), vec!["send_chat_action"]);
|
|
}
|
|
}
|