#![doc = include_str!("../Documentation.md")]
#![forbid(unsafe_code)]
use std::{
collections::{HashMap, VecDeque},
future::Future,
panic::AssertUnwindSafe,
sync::Arc,
};
use futures::{FutureExt, future::BoxFuture};
use teloxide::types::{Message, Update, UpdateKind};
use tokio::sync::{Mutex, Semaphore};
const PER_PRINCIPAL_QUEUE_CAPACITY: usize = 32;
const MAX_ACTIVE_PRINCIPAL_WORKERS: usize = 256;
const GLOBAL_PROCESSING_CONCURRENCY: usize = 16;
type UpdateProcessor = Arc<dyn Fn(Update) -> BoxFuture<'static, anyhow::Result<()>> + Send + Sync>;
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
enum UpdateQueueKey {
Private(i64),
GroupUser(i64, i64),
GroupControl(i64),
Other,
}
struct DispatcherInner {
processor: UpdateProcessor,
workers: Mutex<HashMap<UpdateQueueKey, VecDeque<Update>>>,
concurrency: Arc<Semaphore>,
}
#[derive(Clone)]
pub struct UpdateDispatcher {
inner: Arc<DispatcherInner>,
}
impl UpdateDispatcher {
pub fn new<F, Fut>(processor: F) -> Self
where
F: Fn(Update) -> Fut + Send + Sync + 'static,
Fut: Future<Output = anyhow::Result<()>> + Send + 'static,
{
let processor: UpdateProcessor = Arc::new(move |update| Box::pin(processor(update)));
Self {
inner: Arc::new(DispatcherInner {
processor,
workers: Mutex::new(HashMap::new()),
concurrency: Arc::new(Semaphore::new(GLOBAL_PROCESSING_CONCURRENCY)),
}),
}
}
pub async fn enqueue(&self, update: Update) -> anyhow::Result<()> {
let update_id = i64::from(update.id.0);
let key = update_queue_key(&update);
let spawn_worker = {
let mut workers = self.inner.workers.lock().await;
if let Some(queue) = workers.get_mut(&key) {
if queue.len() >= PER_PRINCIPAL_QUEUE_CAPACITY {
tracing::warn!(
update_id,
"dropping Telegram update because its principal queue is full"
);
return Ok(());
}
queue.push_back(update);
false
} else {
if workers.len() >= MAX_ACTIVE_PRINCIPAL_WORKERS {
tracing::warn!(
update_id,
active_principals = workers.len(),
"dropping Telegram update because the principal-worker limit is full"
);
return Ok(());
}
workers.insert(key.clone(), VecDeque::from([update]));
true
}
};
if spawn_worker {
tokio::spawn(run_worker(self.inner.clone(), key));
}
Ok(())
}
}
async fn remove_worker(inner: &DispatcherInner, key: &UpdateQueueKey) {
inner.workers.lock().await.remove(key);
}
async fn run_worker(inner: Arc<DispatcherInner>, key: UpdateQueueKey) {
loop {
let update = {
let mut workers = inner.workers.lock().await;
let Some(queue) = workers.get_mut(&key) else {
return;
};
if let Some(update) = queue.pop_front() {
update
} else {
workers.remove(&key);
return;
}
};
let update_id = i64::from(update.id.0);
let permit = match inner.concurrency.clone().acquire_owned().await {
Ok(permit) => permit,
Err(_) => {
remove_worker(&inner, &key).await;
return;
}
};
match AssertUnwindSafe((inner.processor)(update))
.catch_unwind()
.await
{
Ok(Ok(())) => {}
Ok(Err(_)) => {
tracing::warn!(
update_id,
"Telegram update failed locally; later work will continue"
);
}
Err(_) => {
tracing::error!(
update_id,
"Telegram update processor panicked; later work will continue"
);
}
}
drop(permit);
}
}
fn update_queue_key(update: &Update) -> UpdateQueueKey {
match &update.kind {
UpdateKind::Message(message) | UpdateKind::EditedMessage(message) => {
if message.chat.is_private() {
message
.from
.as_ref()
.and_then(|user| i64::try_from(user.id.0).ok())
.map(UpdateQueueKey::Private)
.unwrap_or(UpdateQueueKey::Other)
} else if message.chat.is_group() || message.chat.is_supergroup() {
if message.migrate_to_chat_id().is_some()
|| message.migrate_from_chat_id().is_some()
|| is_group_authored_message(message)
{
UpdateQueueKey::GroupControl(message.chat.id.0)
} else {
message
.from
.as_ref()
.and_then(|user| i64::try_from(user.id.0).ok())
.map(|user_id| UpdateQueueKey::GroupUser(message.chat.id.0, user_id))
.unwrap_or(UpdateQueueKey::GroupControl(message.chat.id.0))
}
} else {
UpdateQueueKey::Other
}
}
UpdateKind::ChatMember(change) | UpdateKind::MyChatMember(change) => {
UpdateQueueKey::GroupControl(change.chat.id.0)
}
_ => UpdateQueueKey::Other,
}
}
fn is_group_authored_message(message: &Message) -> bool {
message
.sender_chat
.as_ref()
.is_some_and(|sender| sender.id == message.chat.id)
|| message
.from
.as_ref()
.is_some_and(|user| user.is_anonymous())
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use std::sync::Mutex as StdMutex;
use teloxide::types::UpdateId;
use tokio::sync::{Semaphore, mpsc};
fn private_message(message_id: i32, user_id: i64, text: &str) -> Message {
serde_json::from_value(json!({
"message_id":message_id,
"date":1629404938,
"from":{"id":user_id,"is_bot":false,"first_name":"User"},
"chat":{"id":user_id,"first_name":"User","type":"private"},
"text":text
}))
.unwrap()
}
fn private_update(update_id: u32, user_id: i64) -> Update {
Update {
id: UpdateId(update_id),
kind: UpdateKind::Message(private_message(
i32::try_from(update_id).unwrap(),
user_id,
"test",
)),
}
}
fn edited_private_update(update_id: u32, user_id: i64) -> Update {
Update {
id: UpdateId(update_id),
kind: UpdateKind::EditedMessage(private_message(1, user_id, "edited")),
}
}
fn group_message(message_id: i32, chat_id: i64, user_id: i64, group_authored: bool) -> Message {
let mut value = json!({
"message_id":message_id,
"date":1629404938,
"chat":{"id":chat_id,"title":"Friends","type":"supergroup"},
"text":"test"
});
if group_authored {
value["sender_chat"] = json!({"id":chat_id,"title":"Friends","type":"supergroup"});
} else {
value["from"] = json!({"id":user_id,"is_bot":false,"first_name":"User"});
}
serde_json::from_value(value).unwrap()
}
fn group_update(update_id: u32, chat_id: i64, user_id: i64, group_authored: bool) -> Update {
Update {
id: UpdateId(update_id),
kind: UpdateKind::Message(group_message(
i32::try_from(update_id).unwrap(),
chat_id,
user_id,
group_authored,
)),
}
}
#[tokio::test]
async fn principals_overlap_while_each_principal_remains_ordered() {
let slow_gate = Arc::new(Semaphore::new(0));
let completed = Arc::new(StdMutex::new(Vec::new()));
let (started_sender, mut started_receiver) = mpsc::unbounded_channel();
let (completion_sender, mut completion_receiver) = mpsc::unbounded_channel();
let processor = {
let slow_gate = slow_gate.clone();
let completed = completed.clone();
move |update: Update| {
let slow_gate = slow_gate.clone();
let completed = completed.clone();
let started_sender = started_sender.clone();
let completion_sender = completion_sender.clone();
async move {
let update_id = i64::from(update.id.0);
if update_id == 1 {
let _ = started_sender.send(());
let permit = slow_gate.acquire().await?;
permit.forget();
}
completed.lock().unwrap().push(update_id);
let _ = completion_sender.send(update_id);
Ok(())
}
}
};
let dispatcher = UpdateDispatcher::new(processor);
dispatcher.enqueue(private_update(1, 42)).await.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(1), started_receiver.recv())
.await
.unwrap()
.unwrap();
dispatcher.enqueue(private_update(3, 42)).await.unwrap();
dispatcher.enqueue(private_update(2, 77)).await.unwrap();
assert_eq!(
tokio::time::timeout(
std::time::Duration::from_secs(1),
completion_receiver.recv(),
)
.await
.unwrap(),
Some(2)
);
slow_gate.add_permits(1);
assert_eq!(
tokio::time::timeout(
std::time::Duration::from_secs(1),
completion_receiver.recv(),
)
.await
.unwrap(),
Some(1)
);
assert_eq!(
tokio::time::timeout(
std::time::Duration::from_secs(1),
completion_receiver.recv(),
)
.await
.unwrap(),
Some(3)
);
assert_eq!(*completed.lock().unwrap(), vec![2, 1, 3]);
}
#[tokio::test]
async fn processor_errors_are_local_to_one_update() {
let (completion_sender, mut completion_receiver) = mpsc::unbounded_channel();
let dispatcher = UpdateDispatcher::new(move |update: Update| {
let completion_sender = completion_sender.clone();
async move {
let update_id = i64::from(update.id.0);
if update_id == 1 {
anyhow::bail!("test processor error");
}
let _ = completion_sender.send(update_id);
Ok(())
}
});
dispatcher.enqueue(private_update(1, 42)).await.unwrap();
dispatcher.enqueue(private_update(2, 42)).await.unwrap();
assert_eq!(
tokio::time::timeout(
std::time::Duration::from_secs(1),
completion_receiver.recv(),
)
.await
.unwrap(),
Some(2)
);
}
#[tokio::test]
async fn processor_panics_are_local_to_one_update() {
let (completion_sender, mut completion_receiver) = mpsc::unbounded_channel();
let dispatcher = UpdateDispatcher::new(move |update: Update| {
let completion_sender = completion_sender.clone();
async move {
let update_id = i64::from(update.id.0);
if update_id == 1 {
panic!("test processor panic");
}
let _ = completion_sender.send(update_id);
Ok(())
}
});
dispatcher.enqueue(private_update(1, 42)).await.unwrap();
dispatcher.enqueue(private_update(2, 42)).await.unwrap();
assert_eq!(
tokio::time::timeout(
std::time::Duration::from_secs(1),
completion_receiver.recv(),
)
.await
.unwrap(),
Some(2)
);
}
#[tokio::test]
async fn a_drained_principal_is_removed_before_a_successor_worker_starts() {
let completed = Arc::new(StdMutex::new(Vec::new()));
let (completion_sender, mut completion_receiver) = mpsc::unbounded_channel();
let processor = {
let completed = completed.clone();
move |update: Update| {
let completed = completed.clone();
let completion_sender = completion_sender.clone();
async move {
let update_id = i64::from(update.id.0);
completed.lock().unwrap().push(update_id);
let _ = completion_sender.send(update_id);
Ok(())
}
}
};
let dispatcher = UpdateDispatcher::new(processor);
dispatcher.enqueue(private_update(1, 42)).await.unwrap();
assert_eq!(
tokio::time::timeout(
std::time::Duration::from_secs(1),
completion_receiver.recv(),
)
.await
.unwrap(),
Some(1)
);
loop {
if dispatcher.inner.workers.lock().await.is_empty() {
break;
}
tokio::task::yield_now().await;
}
dispatcher.enqueue(private_update(2, 42)).await.unwrap();
assert_eq!(
tokio::time::timeout(
std::time::Duration::from_secs(1),
completion_receiver.recv(),
)
.await
.unwrap(),
Some(2)
);
assert_eq!(*completed.lock().unwrap(), vec![1, 2]);
}
#[test]
fn edits_share_their_source_principals_queue() {
let original = private_update(10, 42);
let edited = edited_private_update(11, 42);
assert_eq!(update_queue_key(&original), update_queue_key(&edited));
assert_ne!(
update_queue_key(&original),
update_queue_key(&private_update(12, 77))
);
}
#[test]
fn group_user_and_control_updates_use_separate_keys() {
let ordinary = group_update(1, -100, 42, false);
let control = group_update(2, -100, 42, true);
assert_eq!(
update_queue_key(&ordinary),
UpdateQueueKey::GroupUser(-100, 42)
);
assert_eq!(
update_queue_key(&control),
UpdateQueueKey::GroupControl(-100)
);
}
#[test]
fn unkeyed_updates_share_one_bounded_worker_class() {
let first = Update {
id: UpdateId(1),
kind: UpdateKind::Error("unknown".into()),
};
let second = Update {
id: UpdateId(2),
kind: UpdateKind::Error("unknown".into()),
};
assert_eq!(update_queue_key(&first), UpdateQueueKey::Other);
assert_eq!(update_queue_key(&first), update_queue_key(&second));
}
}