OpenSource/teloxide
OpenSource/teloxide
1
0
Fork 0
mirror of https://github.com/teloxide/teloxide.git synced 2024-10-24 09:57:18 +02:00
Code Issues Projects Releases Packages Wiki Activity

Merge pull request #130 from teloxide/chrottle_retries_and_freeze

Browse source 
`Throttle` retries and freeze
This commit is contained in:
Waffle Maybe 2022-04-13 13:58:59 +04:00 committed by GitHub
commit 4e35f6de31
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
15 changed files with 1032 additions and 837 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
CHANGELOG.md
View file

@ -7,8 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## unreleased
### Fixed
- Fix never ending loop that caused programs that used `Throttling` to never stop, see issue [#535][issue535] ([#130][pr130])
[issue535]: https://github.com/teloxide/teloxide/issues/535
[pr130]: https://github.com/teloxide/teloxide-core/pull/130
### Added
- `errors` module and `errors::AsResponseParameters` trait ([#130][pr130])
- `UserId::{url, is_anonymous, is_channel, is_telegram}` convenience functions ([#197][pr197])
- `User::{tme_url, preferably_tme_url}` convenience functions ([#197][pr197])
- `Me::username` and `Deref<Target = User>` implementation for `Me` ([#197][pr197])
@ -21,6 +29,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Changed
- `user.id` now uses `UserId` type, `ChatId` now represents only _chat id_, not channel username, all `chat_id` function parameters now accept `Recipient` [**BC**]
- Improve `Throttling` adoptor ([#130][pr130])
- Freeze when getting `RetryAfter(_)` error
- Retry requests that previously returned `RetryAfter(_)` error
- `RequestError::RetryAfter` now has a `Duration` field instead of `i32`
## 0.4.5 - 2022-04-03

4
Cargo.toml
View file

@ -22,7 +22,7 @@ exclude = [
[dependencies]
futures = "0.3.5"
tokio = { version = "1.8.0", features = ["fs"] }
tokio = { version = "1.12.0", features = ["fs"] }
tokio-util = "0.6.0"
pin-project = "1.0.3"
bytes = "1.0.0"
@ -66,7 +66,7 @@ native-tls = ["reqwest/native-tls"]
nightly = []
# Throttling bot adaptor
throttle = ["vecrem"]
throttle = ["vecrem", "tokio/macros"]
# Trace bot adaptor
trace_adaptor = []

872
src/adaptors/throttle.rs
View file

@ -1,43 +1,49 @@
/// `ThrottlingRequest` and `ThrottlingSend` structures
mod request;
/// Lock that allows requests to wait until they are allowed to be sent
mod request_lock;
/// `impl Requester for Throttle<_>`
mod requester_impl;
/// `Settings` and `Limits` structures
mod settings;
/// "Worker" that checks the limits
mod worker;
use std::{
collections::{hash_map::Entry, HashMap, VecDeque},
future::Future,
hash::{Hash, Hasher},
pin::Pin,
time::{Duration, Instant},
};
use futures::{
future::ready,
task::{Context, Poll},
FutureExt,
};
use never::Never;
use tokio::sync::{
mpsc,
oneshot::{self, Receiver, Sender},
oneshot::{self},
};
use url::Url;
use vecrem::VecExt;
use crate::{
adaptors::throttle::chan_send::{ChanSend, MpscSend},
requests::{HasPayload, Output, Request, Requester},
types::*,
use crate::{errors::AsResponseParameters, requests::Requester, types::*};
use self::{
request_lock::{channel, RequestLock},
worker::{worker, FreezeUntil, InfoMessage},
};
pub use request::{ThrottlingRequest, ThrottlingSend};
pub use settings::{Limits, Settings};
/// Automatic request limits respecting mechanism.
///
/// Telegram has strict [limits], which, if exceeded will sooner or later cause
/// `RequestError::RetryAfter(_)` errors. These errors can cause users of your
/// bot to never receive responds from the bot or receive them in wrong order.
/// bot to never receive responses from the bot or receive them in a wrong
/// order.
///
/// This bot wrapper automatically checks for limits, suspending requests until
/// they could be sent without exceeding limits (request order in chats is not
/// changed).
///
/// It's recommended to use this wrapper before other wrappers (i.e.:
/// `SomeWrapper<Throttle<Bot>>`) because if done otherwise inner wrappers may
/// cause `Throttle` to miscalculate limits usage.
/// `SomeWrapper<Throttle<Bot>>` not `Throttle<SomeWrapper<Bot>>`) because if
/// done otherwise inner wrappers may cause `Throttle` to miscalculate limits
/// usage.
///
/// [limits]: https://core.telegram.org/bots/faq#my-bot-is-hitting-limits-how-do-i-avoid-this
///
@ -46,7 +52,8 @@ use crate::{
/// ```no_run (throttle fails to spawn task without tokio runtime)
/// use teloxide_core::{adaptors::throttle::Limits, requests::RequesterExt, Bot};
///
/// let bot = Bot::new("TOKEN").throttle(Limits::default());
/// let bot = Bot::new("TOKEN")
/// .throttle(Limits::default());
///
/// /* send many requests here */
/// ```
@ -76,7 +83,11 @@ impl<B> Throttle<B> {
///
/// Note: [`Throttle`] will only send requests if returned worker is
/// polled/spawned/awaited.
pub fn new(bot: B, limits: Limits) -> (Self, impl Future<Output = ()>) {
pub fn new(bot: B, limits: Limits) -> (Self, impl Future<Output = ()>)
where
B: Requester + Clone,
B::Err: AsResponseParameters,
{
let settings = Settings {
limits,
..<_>::default()
@ -88,11 +99,15 @@ impl<B> Throttle<B> {
///
/// Note: [`Throttle`] will only send requests if returned worker is
/// polled/spawned/awaited.
pub fn with_settings(bot: B, settings: Settings) -> (Self, impl Future<Output = ()>) {
pub fn with_settings(bot: B, settings: Settings) -> (Self, impl Future<Output = ()>)
where
B: Requester + Clone,
B::Err: AsResponseParameters,
{
let (tx, rx) = mpsc::channel(settings.limits.messages_per_sec_overall as usize);
let (info_tx, info_rx) = mpsc::channel(2);
let worker = worker(settings, rx, info_rx);
let worker = worker(settings, rx, info_rx, bot.clone());
let this = Self {
bot,
queue: tx,
@ -107,44 +122,27 @@ impl<B> Throttle<B> {
/// Note: it's recommended to use [`RequesterExt::throttle`] instead.
///
/// [`RequesterExt::throttle`]: crate::requests::RequesterExt::throttle
pub fn new_spawn(bot: B, limits: Limits) -> Self {
// new/with_settings copy-pasted here to avoid [rust-lang/#76882]
//
// [rust-lang/#76882]: https://github.com/rust-lang/rust/issues/76882
let (tx, rx) = mpsc::channel(limits.messages_per_sec_overall as usize);
let (info_tx, info_rx) = mpsc::channel(2);
let settings = Settings {
limits,
..<_>::default()
};
let worker = worker(settings, rx, info_rx);
let this = Self {
bot,
queue: tx,
info_tx,
};
pub fn new_spawn(bot: B, limits: Limits) -> Self
where
B: Requester + Clone + Send + Sync + 'static,
B::Err: AsResponseParameters,
B::GetChat: Send,
{
let (this, worker) = Self::new(bot, limits);
tokio::spawn(worker);
this
}
/// Creates new [`Throttle`] spawning the worker with `tokio::spawn`
pub fn spawn_with_settings(bot: B, settings: Settings) -> Self {
// with_settings copy-pasted here to avoid [rust-lang/#76882]
//
// [rust-lang/#76882]: https://github.com/rust-lang/rust/issues/76882
let (tx, rx) = mpsc::channel(settings.limits.messages_per_sec_overall as usize);
let (info_tx, info_rx) = mpsc::channel(2);
let worker = worker(settings, rx, info_rx);
let this = Self {
bot,
queue: tx,
info_tx,
};
pub fn spawn_with_settings(bot: B, settings: Settings) -> Self
where
B: Requester + Clone + Send + Sync + 'static,
B::Err: AsResponseParameters,
B::GetChat: Send,
{
let (this, worker) = Self::with_settings(bot, settings);
tokio::spawn(worker);
this
@ -162,6 +160,8 @@ impl<B> Throttle<B> {
/// Returns currently used [`Limits`].
pub async fn limits(&self) -> Limits {
const WORKER_DIED: &str = "worker died before last `Throttle` instance";
let (tx, rx) = oneshot::channel();
self.info_tx
@ -187,451 +187,6 @@ impl<B> Throttle<B> {
}
}
/// Telegram request limits.
///
/// This struct is used in [`Throttle`].
///
/// Note that you may ask telegram [@BotSupport] to increase limits for your
/// particular bot if it has a lot of users (but they may or may not do that).
///
/// [@BotSupport]: https://t.me/botsupport
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub struct Limits {
/// Allowed messages in one chat per second.
pub messages_per_sec_chat: u32,
/// Allowed messages in one chat per minute.
pub messages_per_min_chat: u32,
/// Allowed messages in one channel per minute.
pub messages_per_min_channel: u32,
/// Allowed messages per second.
pub messages_per_sec_overall: u32,
}
/// Defaults are taken from [telegram documentation][tgdoc] (except for
/// `messages_per_min_channel`).
///
/// [tgdoc]: https://core.telegram.org/bots/faq#my-bot-is-hitting-limits-how-do-i-avoid-this
impl Default for Limits {
fn default() -> Self {
Self {
messages_per_sec_chat: 1,
messages_per_sec_overall: 30,
messages_per_min_chat: 20,
messages_per_min_channel: 10,
}
}
}
/// Settings used by [`Throttle`] adaptor.
///
/// ## Examples
///
/// ```
/// use teloxide_core::adaptors::throttle;
///
/// let settings = throttle::Settings::default()
/// .on_queue_full(|pending| async move { /* do something when internal queue is full */ });
/// // use settings in `Throttle::with_settings` or other constructors
/// # let _ = settings;
/// ```
#[non_exhaustive]
pub struct Settings {
pub limits: Limits,
pub on_queue_full: BoxedFnMut<usize, BoxedFuture>,
}
impl Settings {
pub fn limits(mut self, val: Limits) -> Self {
self.limits = val;
self
}
pub fn on_queue_full<F, Fut>(mut self, mut val: F) -> Self
where
F: FnMut(usize) -> Fut + Send + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
self.on_queue_full = Box::new(move |pending| Box::pin(val(pending)));
self
}
}
impl Default for Settings {
fn default() -> Self {
Self {
limits: <_>::default(),
on_queue_full: Box::new(|pending| {
log::warn!("Throttle queue is full ({} pending requests)", pending);
Box::pin(ready(()))
}),
}
}
}
// Required to not trigger `clippy::type-complexity` lint
type BoxedFnMut<I, O> = Box<dyn FnMut(I) -> O + Send>;
type BoxedFuture = Pin<Box<dyn Future<Output = ()> + Send>>;
const WORKER_DIED: &str = "worker died before last `Throttle` instance";
const MINUTE: Duration = Duration::from_secs(60);
const SECOND: Duration = Duration::from_secs(1);
// Delay between worker iterations.
//
// For now it's `second/4`, but that number is chosen pretty randomly, we may
// want to change this.
const DELAY: Duration = Duration::from_millis(250);
/// Minimal time between calls to queue_full function
const QUEUE_FULL_DELAY: Duration = Duration::from_secs(4);
#[derive(Debug)]
enum InfoMessage {
GetLimits { response: Sender<Limits> },
SetLimits { new: Limits, response: Sender<()> },
}
type RequestsSent = u32;
// I wish there was special data structure for history which removed the
// need in 2 hashmaps
// (waffle)
#[derive(Default)]
struct RequestsSentToChats {
per_min: HashMap<ChatIdHash, RequestsSent>,
per_sec: HashMap<ChatIdHash, RequestsSent>,
}
// Throttling is quite complicated. This comment describes the algorithm of the
// current implementation.
//
// ### Request
//
// When a throttling request is sent, it sends a tuple of `ChatId` and
// `Sender<()>` to the worker. Then the request waits for a notification from
// the worker. When notification is received, it sends the underlying request.
//
// ### Worker
//
// The worker does the most important job -- it ensures that the limits are
// never exceeded.
//
// The worker stores a history of requests sent in the last minute (and to which
// chats they were sent) and a queue of pending updates.
//
// The worker does the following algorithm loop:
//
// 1. If the queue is empty, wait for the first message in incoming channel (and
// add it to the queue).
//
// 2. Read all present messages from an incoming channel and transfer them to
// the queue.
//
// 3. Record the current time.
//
// 4. Clear the history from records whose time < (current time - minute).
//
// 5. Count all requests which were sent last second, `allowed =
// limit.messages_per_sec_overall - count`.
//
// 6. If `allowed == 0` wait a bit and `continue` to the next iteration.
//
// 7. Count how many requests were sent to which chats (i.e.: create
// `Map<ChatId, Count>`). (Note: the same map, but for last minute also exists,
// but it's updated, instead of recreation.)
//
// 8. While `allowed >= 0` search for requests which chat haven't exceed the
// limits (i.e.: map[chat] < limit), if one is found, decrease `allowed`, notify
// the request that it can be now executed, increase counts, add record to the
// history.
async fn worker(
Settings {
mut limits,
mut on_queue_full,
}: Settings,
mut rx: mpsc::Receiver<(ChatIdHash, RequestLock)>,
mut info_rx: mpsc::Receiver<InfoMessage>,
) {
// FIXME(waffle): Make an research about data structures for this queue.
// Currently this is O(n) removing (n = number of elements
// stayed), amortized O(1) push (vec+vecrem).
let mut queue: Vec<(ChatIdHash, RequestLock)> =
Vec::with_capacity(limits.messages_per_sec_overall as usize);
let mut history: VecDeque<(ChatIdHash, Instant)> = VecDeque::new();
let mut requests_sent = RequestsSentToChats::default();
let mut rx_is_closed = false;
let mut last_queue_full = Instant::now()
.checked_sub(QUEUE_FULL_DELAY)
.unwrap_or_else(Instant::now);
while !rx_is_closed || !queue.is_empty() {
// FIXME(waffle):
// 1. If the `queue` is empty, `read_from_rx` call down below will 'block'
// execution until a request is sent. While the execution is 'blocked' no
// `InfoMessage`s could be answered.
//
// 2. If limits are decreased, ideally we want to shrink queue.
//
// *blocked in asynchronous way
answer_info(&mut info_rx, &mut limits);
read_from_rx(&mut rx, &mut queue, &mut rx_is_closed).await;
//debug_assert_eq!(queue.capacity(), limits.messages_per_sec_overall as usize);
if queue.len() == queue.capacity() && last_queue_full.elapsed() > QUEUE_FULL_DELAY {
last_queue_full = Instant::now();
tokio::spawn(on_queue_full(queue.len()));
}
// _Maybe_ we need to use `spawn_blocking` here, because there is
// decent amount of blocking work. However _for now_ I've decided not
// to use it here.
//
// Reasons (not to use `spawn_blocking`):
//
// 1. The work seems not very CPU-bound, it's not heavy computations,
// it's more like light computations.
//
// 2. `spawn_blocking` is not zero-cost — it spawns a new system thread
// + do so other work. This may actually be *worse* then current
// "just do everything in this async fn" approach.
//
// 3. With `rt-threaded` feature, tokio uses [`num_cpus()`] threads
// which should be enough to work fine with one a-bit-blocking task.
// Crucially current behaviour will be problem mostly with
// single-threaded runtimes (and in case you're using one, you
// probably don't want to spawn unnecessary threads anyway).
//
// I think if we'll ever change this behaviour, we need to make it
// _configurable_.
//
// See also [discussion (ru)].
//
// NOTE: If you are reading this because you have any problems because
// of this worker, open an [issue on github]
//
// [`num_cpus()`]: https://vee.gg/JGwq2
// [discussion (ru)]: https://t.me/rust_async/27891
// [issue on github]: https://github.com/teloxide/teloxide/issues/new
//
// (waffle)
let now = Instant::now();
let min_back = now - MINUTE;
let sec_back = now - SECOND;
// make history and requests_sent up-to-date
while let Some((_, time)) = history.front() {
// history is sorted, we found first up-to-date thing
if time >= &min_back {
break;
}
if let Some((chat, _)) = history.pop_front() {
let entry = requests_sent.per_min.entry(chat).and_modify(|count| {
*count -= 1;
});
if let Entry::Occupied(entry) = entry {
if *entry.get() == 0 {
entry.remove_entry();
}
}
}
}
// as truncates which is ok since in case of truncation it would always be >=
// limits.overall_s
let used = history
.iter()
.take_while(|(_, time)| time > &sec_back)
.count() as u32;
let mut allowed = limits.messages_per_sec_overall.saturating_sub(used);
if allowed == 0 {
requests_sent.per_sec.clear();
tokio::time::sleep(DELAY).await;
continue;
}
for (chat, _) in history.iter().take_while(|(_, time)| time > &sec_back) {
*requests_sent.per_sec.entry(*chat).or_insert(0) += 1;
}
let mut queue_removing = queue.removing();
while let Some(entry) = queue_removing.next() {
let chat = &entry.value().0;
let requests_sent_per_sec_count = requests_sent.per_sec.get(chat).copied().unwrap_or(0);
let requests_sent_per_min_count = requests_sent.per_min.get(chat).copied().unwrap_or(0);
let messages_per_min_limit = if chat.is_channel() {
limits.messages_per_min_channel
} else {
limits.messages_per_min_chat
};
let limits_not_exceeded = requests_sent_per_sec_count < limits.messages_per_sec_chat
&& requests_sent_per_min_count < messages_per_min_limit;
if limits_not_exceeded {
*requests_sent.per_sec.entry(*chat).or_insert(0) += 1;
*requests_sent.per_min.entry(*chat).or_insert(0) += 1;
history.push_back((*chat, Instant::now()));
// Close the channel and unlock the associated request.
let (_, lock) = entry.remove();
lock.unlock();
// We have "sent" one request, so now we can send one less.
allowed -= 1;
if allowed == 0 {
break;
}
}
}
// It's easier to just recompute last second stats, instead of keeping
// track of it alongside with minute stats, so we just throw this away.
requests_sent.per_sec.clear();
tokio::time::sleep(DELAY).await;
}
}
fn answer_info(rx: &mut mpsc::Receiver<InfoMessage>, limits: &mut Limits) {
// FIXME(waffle): https://github.com/tokio-rs/tokio/issues/3350
while let Some(Some(req)) = tokio::task::unconstrained(rx.recv()).now_or_never() {
// Errors are ignored with .ok(). Error means that the response channel
// is closed and the response isn't needed.
match req {
InfoMessage::GetLimits { response } => response.send(*limits).ok(),
InfoMessage::SetLimits { new, response } => {
*limits = new;
response.send(()).ok()
}
};
}
}
async fn read_from_rx<T>(rx: &mut mpsc::Receiver<T>, queue: &mut Vec<T>, rx_is_closed: &mut bool) {
if queue.is_empty() {
match rx.recv().await {
Some(req) => queue.push(req),
None => *rx_is_closed = true,
}
}
// Don't grow queue bigger than the capacity to limit DOS possibility
while queue.len() < queue.capacity() {
// FIXME(waffle): https://github.com/tokio-rs/tokio/issues/3350
match tokio::task::unconstrained(rx.recv()).now_or_never() {
Some(Some(req)) => queue.push(req),
Some(None) => *rx_is_closed = true,
// There are no items in queue.
None => break,
}
}
}
macro_rules! f {
($m:ident $this:ident ($($arg:ident : $T:ty),*)) => {
ThrottlingRequest {
request: $this.inner().$m($($arg),*),
chat_id: |p| (&p.payload_ref().chat_id).into(),
worker: $this.queue.clone(),
}
};
}
macro_rules! fty {
($T:ident) => {
ThrottlingRequest<B::$T>
};
}
macro_rules! fid {
($m:ident $this:ident ($($arg:ident : $T:ty),*)) => {
$this.inner().$m($($arg),*)
};
}
macro_rules! ftyid {
($T:ident) => {
B::$T
};
}
impl<B: Requester> Requester for Throttle<B>
where
B::SendMessage: Send,
B::ForwardMessage: Send,
B::CopyMessage: Send,
B::SendPhoto: Send,
B::SendAudio: Send,
B::SendDocument: Send,
B::SendVideo: Send,
B::SendAnimation: Send,
B::SendVoice: Send,
B::SendVideoNote: Send,
B::SendMediaGroup: Send,
B::SendLocation: Send,
B::SendVenue: Send,
B::SendContact: Send,
B::SendPoll: Send,
B::SendDice: Send,
B::SendSticker: Send,
B::SendInvoice: Send,
{
type Err = B::Err;
requester_forward! {
send_message, forward_message, copy_message, send_photo, send_audio,
send_document, send_video, send_animation, send_voice, send_video_note,
send_media_group, send_location, send_venue, send_contact, send_poll,
send_dice, send_sticker, send_invoice => f, fty
}
requester_forward! {
get_me, log_out, close, get_updates, set_webhook, delete_webhook, get_webhook_info,
edit_message_live_location, edit_message_live_location_inline,
stop_message_live_location, stop_message_live_location_inline,
send_chat_action, get_user_profile_photos, get_file, kick_chat_member, ban_chat_member,
unban_chat_member, restrict_chat_member, promote_chat_member,
set_chat_administrator_custom_title,
ban_chat_sender_chat, unban_chat_sender_chat, set_chat_permissions,
export_chat_invite_link, create_chat_invite_link, edit_chat_invite_link,
revoke_chat_invite_link, set_chat_photo, delete_chat_photo, set_chat_title,
set_chat_description, pin_chat_message, unpin_chat_message, unpin_all_chat_messages,
leave_chat, get_chat, get_chat_administrators, get_chat_members_count, get_chat_member_count,
get_chat_member, set_chat_sticker_set, delete_chat_sticker_set,
answer_callback_query, set_my_commands, get_my_commands, delete_my_commands, answer_inline_query,
edit_message_text, edit_message_text_inline, edit_message_caption,
edit_message_caption_inline, edit_message_media, edit_message_media_inline,
edit_message_reply_markup, edit_message_reply_markup_inline, stop_poll,
delete_message, get_sticker_set, upload_sticker_file, create_new_sticker_set,
add_sticker_to_set, set_sticker_position_in_set, delete_sticker_from_set,
set_sticker_set_thumb, answer_shipping_query, answer_pre_checkout_query,
set_passport_data_errors, send_game, set_game_score, set_game_score_inline,
approve_chat_join_request, decline_chat_join_request,
get_game_high_scores => fid, ftyid
}
}
download_forward! {
'w
B
Throttle<B>
{ this => this.inner() }
}
/// An ID used in the worker.
///
/// It is used instead of `ChatId` to make copying cheap even in case of
@ -656,6 +211,8 @@ impl From<&Recipient> for ChatIdHash {
match value {
Recipient::Id(id) => ChatIdHash::Id(*id),
Recipient::ChannelUsername(username) => {
// FIXME: this could probably use a faster hasher, `DefaultHasher` is known to
// be slow (it's not like we _need_ this to be fast, but still)
let mut hasher = std::collections::hash_map::DefaultHasher::new();
username.hash(&mut hasher);
let hash = hasher.finish();
@ -664,318 +221,3 @@ impl From<&Recipient> for ChatIdHash {
}
}
}
#[must_use = "Requests are lazy and do nothing unless sent"]
pub struct ThrottlingRequest<R: HasPayload> {
request: R,
chat_id: fn(&R::Payload) -> ChatIdHash,
worker: mpsc::Sender<(ChatIdHash, RequestLock)>,
}
impl<R: HasPayload> HasPayload for ThrottlingRequest<R> {
type Payload = R::Payload;
fn payload_mut(&mut self) -> &mut Self::Payload {
self.request.payload_mut()
}
fn payload_ref(&self) -> &Self::Payload {
self.request.payload_ref()
}
}
impl<R> Request for ThrottlingRequest<R>
where
R: Request + Send,
{
type Err = R::Err;
type Send = ThrottlingSend<R>;
type SendRef = ThrottlingSendRef<R>;
fn send(self) -> Self::Send {
let (tx, rx) = channel();
let chat_id = (self.chat_id)(self.payload_ref());
let send = self.worker.send1((chat_id, tx));
let inner = ThrottlingSendInner::Registering {
request: self.request,
send,
wait: rx,
};
ThrottlingSend(inner)
}
fn send_ref(&self) -> Self::SendRef {
let (tx, rx) = channel();
let chat_id = (self.chat_id)(self.payload_ref());
let send = self.worker.clone().send1((chat_id, tx));
// As we can't move self.0 (request) out, as we do in `send` we are
// forced to call `send_ref()`. This may have overhead and/or lead to
// wrong results because `R::send_ref` does the send.
//
// However `Request` documentation explicitly notes that `send{,_ref}`
// should **not** do any kind of work, so it's ok.
let request = self.request.send_ref();
let inner = ThrottlingSendRefInner::Registering {
request,
send,
wait: rx,
};
ThrottlingSendRef(inner)
}
}
#[pin_project::pin_project]
pub struct ThrottlingSend<R: Request>(#[pin] ThrottlingSendInner<R>);
#[pin_project::pin_project(project = SendProj, project_replace = SendRepl)]
enum ThrottlingSendInner<R: Request> {
Registering {
request: R,
#[pin]
send: ChanSend<(ChatIdHash, RequestLock)>,
wait: RequestWaiter,
},
Pending {
request: R,
#[pin]
wait: RequestWaiter,
},
Sent {
#[pin]
fut: R::Send,
},
Done,
}
impl<R: Request> Future for ThrottlingSend<R> {
type Output = Result<Output<R>, R::Err>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut this = self.as_mut().project().0;
match this.as_mut().project() {
SendProj::Registering {
request: _,
send,
wait: _,
} => match send.poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(res) => {
if let SendRepl::Registering {
request,
send: _,
wait,
} = this.as_mut().project_replace(ThrottlingSendInner::Done)
{
match res {
Ok(()) => this
.as_mut()
.project_replace(ThrottlingSendInner::Pending { request, wait }),
// The worker is unlikely to drop queue before sending all requests,
// but just in case it has dropped the queue, we want to just send the
// request.
Err(_) => this.as_mut().project_replace(ThrottlingSendInner::Sent {
fut: request.send(),
}),
};
}
self.poll(cx)
}
},
SendProj::Pending { request: _, wait } => match wait.poll(cx) {
Poll::Pending => Poll::Pending,
// Worker pass "message" to unlock us by closing the channel,
// and thus we can safely ignore this result as we know it will
// always be `Err(_)` (because `Ok(Never)` is uninhibited)
// and that's what we want.
Poll::Ready(_) => {
if let SendRepl::Pending { request, wait: _ } =
this.as_mut().project_replace(ThrottlingSendInner::Done)
{
this.as_mut().project_replace(ThrottlingSendInner::Sent {
fut: request.send(),
});
}
self.poll(cx)
}
},
SendProj::Sent { fut } => {
let res = futures::ready!(fut.poll(cx));
this.set(ThrottlingSendInner::Done);
Poll::Ready(res)
}
SendProj::Done => Poll::Pending,
}
}
}
#[pin_project::pin_project]
pub struct ThrottlingSendRef<R: Request>(#[pin] ThrottlingSendRefInner<R>);
#[pin_project::pin_project(project = SendRefProj, project_replace = SendRefRepl)]
enum ThrottlingSendRefInner<R: Request> {
Registering {
request: R::SendRef,
#[pin]
send: ChanSend<(ChatIdHash, RequestLock)>,
wait: RequestWaiter,
},
Pending {
request: R::SendRef,
#[pin]
wait: RequestWaiter,
},
Sent {
#[pin]
fut: R::SendRef,
},
Done,
}
impl<R: Request> Future for ThrottlingSendRef<R> {
type Output = Result<Output<R>, R::Err>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut this = self.as_mut().project().0;
match this.as_mut().project() {
SendRefProj::Registering {
request: _,
send,
wait: _,
} => match send.poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(res) => {
if let SendRefRepl::Registering {
request,
send: _,
wait,
} = this.as_mut().project_replace(ThrottlingSendRefInner::Done)
{
match res {
Ok(()) => this
.as_mut()
.project_replace(ThrottlingSendRefInner::Pending { request, wait }),
// The worker is unlikely to drop queue before sending all requests,
// but just in case it has dropped the queue, we want to just send the
// request.
Err(_) => this
.as_mut()
.project_replace(ThrottlingSendRefInner::Sent { fut: request }),
};
}
self.poll(cx)
}
},
SendRefProj::Pending { request: _, wait } => match wait.poll(cx) {
Poll::Pending => Poll::Pending,
// Worker pass "message" to unlock us by closing the channel,
// and thus we can safely ignore this result as we know it will
// always be `Err(_)` (because `Ok(Never)` is uninhibited)
// and that's what we want.
Poll::Ready(_) => {
if let SendRefRepl::Pending { request, wait: _ } =
this.as_mut().project_replace(ThrottlingSendRefInner::Done)
{
this.as_mut()
.project_replace(ThrottlingSendRefInner::Sent { fut: request });
}
self.poll(cx)
}
},
SendRefProj::Sent { fut } => {
let res = futures::ready!(fut.poll(cx));
this.set(ThrottlingSendRefInner::Done);
Poll::Ready(res)
}
SendRefProj::Done => Poll::Pending,
}
}
}
fn channel() -> (RequestLock, RequestWaiter) {
let (tx, rx) = oneshot::channel();
let tx = RequestLock(tx);
let rx = RequestWaiter(rx);
(tx, rx)
}
#[must_use]
struct RequestLock(Sender<Never>);
impl RequestLock {
fn unlock(self) {
// Unlock request by closing oneshot channel
}
}
#[must_use]
#[pin_project::pin_project]
struct RequestWaiter(#[pin] Receiver<Never>);
impl Future for RequestWaiter {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let this = self.project();
match this.0.poll(cx) {
Poll::Ready(_) => Poll::Ready(()),
Poll::Pending => Poll::Pending,
}
}
}
mod chan_send {
use std::{future::Future, pin::Pin};
use futures::task::{Context, Poll};
use tokio::sync::{mpsc, mpsc::error::SendError};
pub(super) trait MpscSend<T> {
fn send1(self, val: T) -> ChanSend<T>;
}
#[pin_project::pin_project]
pub(super) struct ChanSend<T>(#[pin] Inner<T>);
#[cfg(not(feature = "nightly"))]
type Inner<T> = Pin<Box<dyn Future<Output = Result<(), SendError<T>>> + Send>>;
#[cfg(feature = "nightly")]
type Inner<T> = impl Future<Output = Result<(), SendError<T>>>;
impl<T: Send + 'static> MpscSend<T> for mpsc::Sender<T> {
// `return`s trick IDEA not to show errors
#[allow(clippy::needless_return)]
fn send1(self, val: T) -> ChanSend<T> {
#[cfg(feature = "nightly")]
{
fn def<T>(sender: mpsc::Sender<T>, val: T) -> Inner<T> {
async move { sender.send(val).await }
}
return ChanSend(def(self, val));
}
#[cfg(not(feature = "nightly"))]
{
let this = self;
return ChanSend(Box::pin(async move { this.send(val).await }));
}
}
}
impl<T> Future for ChanSend<T> {
type Output = Result<(), SendError<T>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.project().0.poll(cx)
}
}
}

204
src/adaptors/throttle/request.rs Normal file
View file

@ -0,0 +1,204 @@
use std::{future::Future, pin::Pin, sync::Arc, time::Instant};
use futures::{
future::BoxFuture,
task::{Context, Poll},
};
use tokio::sync::mpsc;
use crate::{
adaptors::throttle::{channel, ChatIdHash, FreezeUntil, RequestLock},
errors::AsResponseParameters,
requests::{HasPayload, Output, Request},
};
/// Request returned by [`Throttling`](crate::adaptors::Throttle) methods.
#[must_use = "Requests are lazy and do nothing unless sent"]
pub struct ThrottlingRequest<R: HasPayload> {
pub(super) request: Arc<R>,
pub(super) chat_id: fn(&R::Payload) -> ChatIdHash,
pub(super) worker: mpsc::Sender<(ChatIdHash, RequestLock)>,
}
/// Future returned by [`ThrottlingRequest`]s.
#[pin_project::pin_project]
pub struct ThrottlingSend<R: Request>(#[pin] BoxFuture<'static, Result<Output<R>, R::Err>>);
enum ShareableRequest<R> {
Shared(Arc<R>),
// Option is used to `take` ownership
Owned(Option<R>),
}
impl<R: HasPayload + Clone> HasPayload for ThrottlingRequest<R> {
type Payload = R::Payload;
/// Note that if this request was already executed via `send_ref` and it
/// didn't yet completed, this method will clone the underlying request.
fn payload_mut(&mut self) -> &mut Self::Payload {
Arc::make_mut(&mut self.request).payload_mut()
}
fn payload_ref(&self) -> &Self::Payload {
self.request.payload_ref()
}
}
impl<R> Request for ThrottlingRequest<R>
where
R: Request + Clone + Send + Sync + 'static, // TODO: rem static
R::Err: AsResponseParameters + Send,
Output<R>: Send,
{
type Err = R::Err;
type Send = ThrottlingSend<R>;
type SendRef = ThrottlingSend<R>;
fn send(self) -> Self::Send {
let chat = (self.chat_id)(self.payload_ref());
let request = match Arc::try_unwrap(self.request) {
Ok(owned) => ShareableRequest::Owned(Some(owned)),
Err(shared) => ShareableRequest::Shared(shared),
};
let fut = send(request, chat, self.worker);
ThrottlingSend(Box::pin(fut))
}
fn send_ref(&self) -> Self::SendRef {
let chat = (self.chat_id)(self.payload_ref());
let request = ShareableRequest::Shared(Arc::clone(&self.request));
let fut = send(request, chat, self.worker.clone());
ThrottlingSend(Box::pin(fut))
}
}
impl<R: Request> Future for ThrottlingSend<R>
where
R::Err: AsResponseParameters,
{
type Output = Result<Output<R>, R::Err>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.as_mut().project().0.poll(cx)
}
}
// This diagram explains how `ThrottlingRequest` works/what `send` does
//
// │
// ThrottlingRequest │ worker()
// │
// ┌───────────────┐ │ ┌────────────────────────┐
// ┌──────────────────►│request is sent│ │ │see worker documentation│
// │ └───────┬───────┘ │ │and comments for more │
// │ │ │ │information on how it │
// │ ▼ │ │actually works │
// │ ┌─────────┐ │ └────────────────────────┘
// │ ┌────────────────┐ │send lock│ │
// │ │has worker died?│◄──┤to worker├─────►:───────────┐
// │ └─┬─────────────┬┘ └─────────┘ │ ▼
// │ │ │ │ ┌──────────────────┐
// │ Y └─N───────┐ │ │ *magic* │
// │ │ │ │ └────────┬─────────┘
// │ ▼ ▼ │ │
// │ ┌───────────┐ ┌────────────────┐ │ ▼
// │ │send inner │ │wait for worker │ │ ┌─────────────────┐
// │ │request │ │to allow sending│◄──:◄─┤ `lock.unlock()` │
// │ └───┬───────┘ │this request │ │ └─────────────────┘
// │ │ └────────┬───────┘ │
// │ │ │ │
// │ ▼ ▼ │
// │ ┌──────┐ ┌────────────────────┐ │
// │ │return│ │send inner request │ │
// │ │result│ │and check its result│ │
// │ └──────┘ └─┬─────────┬────────┘ │
// │ ▲ ▲ │ │ │
// │ │ │ │ Err(RetryAfter(n)) │
// │ │ │ else │ │
// │ │ │ │ ▼ │
// │ │ └─────┘ ┌───────────────┐ │
// │ │ │are retries on?│ │
// │ │ └┬─────────────┬┘ │
// │ │ │ │ │
// │ └────────────N─┘ Y │
// │ │ │ ┌──────────────────┐
// │ ▼ │ │ *magic* │
// │ ┌──────────────────┐ │ └──────────────────┘
// ┌┴────────────┐ │notify worker that│ │ ▲
// │retry request│◄──┤RetryAfter error ├──►:───────────┘
// └─────────────┘ │has happened │ │
// └──────────────────┘ │
// │
/// Actual implementation of the `ThrottlingSend` future
async fn send<R>(
mut request: ShareableRequest<R>,
chat: ChatIdHash,
worker: mpsc::Sender<(ChatIdHash, RequestLock)>,
) -> Result<Output<R>, R::Err>
where
R: Request + Send + Sync + 'static,
R::Err: AsResponseParameters + Send,
Output<R>: Send,
{
// We use option in `ShareableRequest` to `take` when sending by value.
//
// All unwraps down below will succeed because we always return immediately
// after taking.
loop {
let (lock, wait) = channel();
// The worker is unlikely to drop queue before sending all requests,
// but just in case it has dropped the queue, we want to just send the
// request.
if worker.send((chat, lock)).await.is_err() {
log::error!("Worker dropped the queue before sending all requests");
let res = match &mut request {
ShareableRequest::Shared(shared) => shared.send_ref().await,
ShareableRequest::Owned(owned) => owned.take().unwrap().send().await,
};
return res;
};
let (retry, freeze) = wait.await;
let res = match (retry, &mut request) {
// Retries are turned on, use `send_ref` even if we have owned access
(true, request) => {
let request = match request {
ShareableRequest::Shared(shared) => &**shared,
ShareableRequest::Owned(owned) => owned.as_ref().unwrap(),
};
request.send_ref().await
}
(false, ShareableRequest::Shared(shared)) => shared.send_ref().await,
(false, ShareableRequest::Owned(owned)) => owned.take().unwrap().send().await,
};
let retry_after = res.as_ref().err().and_then(<_>::retry_after);
if let Some(retry_after) = retry_after {
let after = retry_after;
let until = Instant::now() + after;
// If we'll retry, we check that worker hasn't died at the start of the loop
// otherwise we don't care if the worker is alive or not
let _ = freeze.send(FreezeUntil { until, after, chat }).await;
if retry {
log::warn!("Freezing, before retrying: {:?}", retry_after);
tokio::time::sleep_until(until.into()).await;
}
}
match res {
Err(_) if retry && retry_after.is_some() => continue,
res => break res,
};
}
}

45
src/adaptors/throttle/request_lock.rs Normal file
View file

@ -0,0 +1,45 @@
use std::pin::Pin;
use futures::{
task::{Context, Poll},
Future,
};
use tokio::sync::{
mpsc,
oneshot::{self, Receiver, Sender},
};
use crate::adaptors::throttle::FreezeUntil;
pub(super) fn channel() -> (RequestLock, RequestWaiter) {
let (tx, rx) = oneshot::channel();
let tx = RequestLock(tx);
let rx = RequestWaiter(rx);
(tx, rx)
}
#[must_use]
pub(super) struct RequestLock(Sender<(bool, mpsc::Sender<FreezeUntil>)>);
#[must_use]
#[pin_project::pin_project]
pub(super) struct RequestWaiter(#[pin] Receiver<(bool, mpsc::Sender<FreezeUntil>)>);
impl RequestLock {
pub(super) fn unlock(self, retry: bool, freeze: mpsc::Sender<FreezeUntil>) -> Result<(), ()> {
self.0.send((retry, freeze)).map_err(drop)
}
}
impl Future for RequestWaiter {
type Output = (bool, mpsc::Sender<FreezeUntil>);
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
match this.0.poll(cx) {
Poll::Ready(Ok(ret)) => Poll::Ready(ret),
Poll::Ready(Err(_)) => panic!("`RequestLock` is dropped by the throttle worker"),
Poll::Pending => Poll::Pending,
}
}
}

103
src/adaptors/throttle/requester_impl.rs Normal file
View file

@ -0,0 +1,103 @@
use std::sync::Arc;
use url::Url;
use crate::{
adaptors::{throttle::ThrottlingRequest, Throttle},
errors::AsResponseParameters,
requests::{HasPayload, Requester},
types::*,
};
macro_rules! f {
($m:ident $this:ident ($($arg:ident : $T:ty),*)) => {
ThrottlingRequest {
request: Arc::new($this.inner().$m($($arg),*)),
chat_id: |p| (&p.payload_ref().chat_id).into(),
worker: $this.queue.clone(),
}
};
}
macro_rules! fty {
($T:ident) => {
ThrottlingRequest<B::$T>
};
}
macro_rules! fid {
($m:ident $this:ident ($($arg:ident : $T:ty),*)) => {
$this.inner().$m($($arg),*)
};
}
macro_rules! ftyid {
($T:ident) => {
B::$T
};
}
impl<B: Requester> Requester for Throttle<B>
where
B::Err: AsResponseParameters,
B::SendMessage: Clone + Send + Sync + 'static,
B::ForwardMessage: Clone + Send + Sync + 'static,
B::CopyMessage: Clone + Send + Sync + 'static,
B::SendPhoto: Clone + Send + Sync + 'static,
B::SendAudio: Clone + Send + Sync + 'static,
B::SendDocument: Clone + Send + Sync + 'static,
B::SendVideo: Clone + Send + Sync + 'static,
B::SendAnimation: Clone + Send + Sync + 'static,
B::SendVoice: Clone + Send + Sync + 'static,
B::SendVideoNote: Clone + Send + Sync + 'static,
B::SendMediaGroup: Clone + Send + Sync + 'static,
B::SendLocation: Clone + Send + Sync + 'static,
B::SendVenue: Clone + Send + Sync + 'static,
B::SendContact: Clone + Send + Sync + 'static,
B::SendPoll: Clone + Send + Sync + 'static,
B::SendDice: Clone + Send + Sync + 'static,
B::SendSticker: Clone + Send + Sync + 'static,
B::SendInvoice: Clone + Send + Sync + 'static,
{
type Err = B::Err;
requester_forward! {
send_message, forward_message, copy_message, send_photo, send_audio,
send_document, send_video, send_animation, send_voice, send_video_note,
send_media_group, send_location, send_venue, send_contact, send_poll,
send_dice, send_sticker, send_invoice => f, fty
}
requester_forward! {
get_me, log_out, close, get_updates, set_webhook, delete_webhook, get_webhook_info,
edit_message_live_location, edit_message_live_location_inline,
stop_message_live_location, stop_message_live_location_inline,
send_chat_action, get_user_profile_photos, get_file, kick_chat_member, ban_chat_member,
unban_chat_member, restrict_chat_member, promote_chat_member,
set_chat_administrator_custom_title,
ban_chat_sender_chat, unban_chat_sender_chat, set_chat_permissions,
export_chat_invite_link, create_chat_invite_link, edit_chat_invite_link,
revoke_chat_invite_link, set_chat_photo, delete_chat_photo, set_chat_title,
set_chat_description, pin_chat_message, unpin_chat_message, unpin_all_chat_messages,
leave_chat, get_chat, get_chat_administrators, get_chat_members_count, get_chat_member_count,
get_chat_member, set_chat_sticker_set, delete_chat_sticker_set,
answer_callback_query, set_my_commands, get_my_commands, delete_my_commands, answer_inline_query,
edit_message_text, edit_message_text_inline, edit_message_caption,
edit_message_caption_inline, edit_message_media, edit_message_media_inline,
edit_message_reply_markup, edit_message_reply_markup_inline, stop_poll,
delete_message, get_sticker_set, upload_sticker_file, create_new_sticker_set,
add_sticker_to_set, set_sticker_position_in_set, delete_sticker_from_set,
set_sticker_set_thumb, answer_shipping_query, answer_pre_checkout_query,
set_passport_data_errors, send_game, set_game_score, set_game_score_inline,
approve_chat_join_request, decline_chat_join_request,
get_game_high_scores => fid, ftyid
}
}
download_forward! {
'w
B
Throttle<B>
{ this => this.inner() }
}

106
src/adaptors/throttle/settings.rs Normal file
View file

@ -0,0 +1,106 @@
use std::pin::Pin;
use futures::{future::ready, Future};
// Required to not trigger `clippy::type-complexity` lint
type BoxedFnMut<I, O> = Box<dyn FnMut(I) -> O + Send>;
type BoxedFuture = Pin<Box<dyn Future<Output = ()> + Send>>;
/// Settings used by [`Throttle`] adaptor.
///
/// ## Examples
///
/// ```
/// use teloxide_core::adaptors::throttle;
///
/// let settings = throttle::Settings::default()
/// .on_queue_full(|pending| async move { /* do something when internal queue is full */ });
///
/// // use settings in `Throttle::with_settings` or other constructors
/// # let _ = settings;
/// ```
#[non_exhaustive]
pub struct Settings {
pub limits: Limits,
pub on_queue_full: BoxedFnMut<usize, BoxedFuture>,
pub retry: bool,
pub check_slow_mode: bool,
}
/// Telegram request limits.
///
/// This struct is used in [`Throttle`].
///
/// Note that you may ask telegram [@BotSupport] to increase limits for your
/// particular bot if it has a lot of users (but they may or may not do that).
///
/// [@BotSupport]: https://t.me/botsupport
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub struct Limits {
/// Allowed messages in one chat per second.
pub messages_per_sec_chat: u32,
/// Allowed messages in one chat per minute.
pub messages_per_min_chat: u32,
/// Allowed messages in one channel per minute.
pub messages_per_min_channel: u32,
/// Allowed messages per second.
pub messages_per_sec_overall: u32,
}
impl Settings {
pub fn limits(mut self, val: Limits) -> Self {
self.limits = val;
self
}
pub fn on_queue_full<F, Fut>(mut self, mut val: F) -> Self
where
F: FnMut(usize) -> Fut + Send + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
self.on_queue_full = Box::new(move |pending| Box::pin(val(pending)));
self
}
pub fn no_retry(mut self) -> Self {
self.retry = false;
self
}
pub fn check_slow_mode(mut self) -> Self {
self.check_slow_mode = true;
self
}
}
impl Default for Settings {
fn default() -> Self {
Self {
limits: <_>::default(),
on_queue_full: Box::new(|pending| {
log::warn!("Throttle queue is full ({} pending requests)", pending);
Box::pin(ready(()))
}),
retry: true,
check_slow_mode: false,
}
}
}
/// Defaults are taken from [telegram documentation][tgdoc] (except for
/// `messages_per_min_channel`).
///
/// [tgdoc]: https://core.telegram.org/bots/faq#my-bot-is-hitting-limits-how-do-i-avoid-this
impl Default for Limits {
fn default() -> Self {
Self {
messages_per_sec_chat: 1,
messages_per_sec_overall: 30,
messages_per_min_chat: 20,
messages_per_min_channel: 10,
}
}
}

395
src/adaptors/throttle/worker.rs Normal file
View file

@ -0,0 +1,395 @@
use std::{
collections::{hash_map::Entry, HashMap, VecDeque},
time::{Duration, Instant},
};
use tokio::sync::{mpsc, mpsc::error::TryRecvError, oneshot::Sender};
use vecrem::VecExt;
use crate::{
adaptors::throttle::{request_lock::RequestLock, ChatIdHash, Limits, Settings},
errors::AsResponseParameters,
requests::{Request, Requester},
};
const MINUTE: Duration = Duration::from_secs(60);
const SECOND: Duration = Duration::from_secs(1);
// Delay between worker iterations.
//
// For now it's `second/4`, but that number is chosen pretty randomly, we may
// want to change this.
const DELAY: Duration = Duration::from_millis(250);
/// Minimal time between calls to queue_full function
const QUEUE_FULL_DELAY: Duration = Duration::from_secs(4);
#[derive(Debug)]
pub(super) enum InfoMessage {
GetLimits { response: Sender<Limits> },
SetLimits { new: Limits, response: Sender<()> },
}
type RequestsSent = u32;
// I wish there was special data structure for history which removed the
// need in 2 hashmaps
// (waffle)
#[derive(Default)]
struct RequestsSentToChats {
per_min: HashMap<ChatIdHash, RequestsSent>,
per_sec: HashMap<ChatIdHash, RequestsSent>,
}
pub(super) struct FreezeUntil {
pub(super) until: Instant,
pub(super) after: Duration,
pub(super) chat: ChatIdHash,
}
// Throttling is quite complicated. This comment describes the algorithm of the
// current implementation.
//
// ### Request
//
// When a throttling request is sent, it sends a tuple of `ChatId` and
// `Sender<()>` to the worker. Then the request waits for a notification from
// the worker. When notification is received, it sends the underlying request.
//
// ### Worker
//
// The worker does the most important job -- it ensures that the limits are
// never exceeded.
//
// The worker stores a history of requests sent in the last minute (and to which
// chats they were sent) and a queue of pending updates.
//
// The worker does the following algorithm loop:
//
// 1. If the queue is empty, wait for the first message in incoming channel (and
// add it to the queue).
//
// 2. Read all present messages from an incoming channel and transfer them to
// the queue.
//
// 3. Record the current time.
//
// 4. Clear the history from records whose time < (current time - minute).
//
// 5. Count all requests which were sent last second, `allowed =
// limit.messages_per_sec_overall - count`.
//
// 6. If `allowed == 0` wait a bit and `continue` to the next iteration.
//
// 7. Count how many requests were sent to which chats (i.e.: create
// `Map<ChatId, Count>`). (Note: the same map, but for last minute also exists,
// but it's updated, instead of recreation.)
//
// 8. While `allowed >= 0` search for requests which chat haven't exceed the
// limits (i.e.: map[chat] < limit), if one is found, decrease `allowed`, notify
// the request that it can be now executed, increase counts, add record to the
// history.
pub(super) async fn worker<B>(
Settings {
mut limits,
mut on_queue_full,
retry,
check_slow_mode,
}: Settings,
mut rx: mpsc::Receiver<(ChatIdHash, RequestLock)>,
mut info_rx: mpsc::Receiver<InfoMessage>,
bot: B,
) where
B: Requester,
B::Err: AsResponseParameters,
{
// FIXME(waffle): Make an research about data structures for this queue.
// Currently this is O(n) removing (n = number of elements
// stayed), amortized O(1) push (vec+vecrem).
let mut queue: Vec<(ChatIdHash, RequestLock)> =
Vec::with_capacity(limits.messages_per_sec_overall as usize);
let mut history: VecDeque<(ChatIdHash, Instant)> = VecDeque::new();
let mut requests_sent = RequestsSentToChats::default();
let mut slow_mode: Option<HashMap<ChatIdHash, (Duration, Instant)>> =
check_slow_mode.then(HashMap::new);
let mut rx_is_closed = false;
let mut last_queue_full = Instant::now()
.checked_sub(QUEUE_FULL_DELAY)
.unwrap_or_else(Instant::now);
let (freeze_tx, mut freeze_rx) = mpsc::channel::<FreezeUntil>(1);
while !rx_is_closed || !queue.is_empty() {
// FIXME(waffle):
// 1. If the `queue` is empty, `read_from_rx` call down below will 'block'
// execution until a request is sent. While the execution is 'blocked' no
// `InfoMessage`s could be answered.
//
// 2. If limits are decreased, ideally we want to shrink queue.
//
// *blocked in asynchronous way
answer_info(&mut info_rx, &mut limits);
loop {
tokio::select! {
freeze_until = freeze_rx.recv() => {
freeze(
&mut freeze_rx,
slow_mode.as_mut(),
&bot,
freeze_until
)
.await;
},
() = read_from_rx(&mut rx, &mut queue, &mut rx_is_closed) => break,
}
}
//debug_assert_eq!(queue.capacity(), limits.messages_per_sec_overall as usize);
if queue.len() == queue.capacity() && last_queue_full.elapsed() > QUEUE_FULL_DELAY {
last_queue_full = Instant::now();
tokio::spawn(on_queue_full(queue.len()));
}
// _Maybe_ we need to use `spawn_blocking` here, because there is
// decent amount of blocking work. However _for now_ I've decided not
// to use it here.
//
// Reasons (not to use `spawn_blocking`):
//
// 1. The work seems not very CPU-bound, it's not heavy computations,
// it's more like light computations.
//
// 2. `spawn_blocking` is not zero-cost — it spawns a new system thread
// + do so other work. This may actually be *worse* then current
// "just do everything in this async fn" approach.
//
// 3. With `rt-threaded` feature, tokio uses [`num_cpus()`] threads
// which should be enough to work fine with one a-bit-blocking task.
// Crucially current behaviour will be problem mostly with
// single-threaded runtimes (and in case you're using one, you
// probably don't want to spawn unnecessary threads anyway).
//
// I think if we'll ever change this behaviour, we need to make it
// _configurable_.
//
// See also [discussion (ru)].
//
// NOTE: If you are reading this because you have any problems because
// of this worker, open an [issue on github]
//
// [`num_cpus()`]: https://vee.gg/JGwq2
// [discussion (ru)]: https://t.me/rust_async/27891
// [issue on github]: https://github.com/teloxide/teloxide/issues/new
//
// (waffle)
let now = Instant::now();
let min_back = now - MINUTE;
let sec_back = now - SECOND;
// make history and requests_sent up-to-date
while let Some((_, time)) = history.front() {
// history is sorted, we found first up-to-date thing
if time >= &min_back {
break;
}
if let Some((chat, _)) = history.pop_front() {
let entry = requests_sent.per_min.entry(chat).and_modify(|count| {
*count -= 1;
});
if let Entry::Occupied(entry) = entry {
if *entry.get() == 0 {
entry.remove_entry();
}
}
}
}
// as truncates which is ok since in case of truncation it would always be >=
// limits.overall_s
let used = history
.iter()
.take_while(|(_, time)| time > &sec_back)
.count() as u32;
let mut allowed = limits.messages_per_sec_overall.saturating_sub(used);
if allowed == 0 {
requests_sent.per_sec.clear();
tokio::time::sleep(DELAY).await;
continue;
}
for (chat, _) in history.iter().take_while(|(_, time)| time > &sec_back) {
*requests_sent.per_sec.entry(*chat).or_insert(0) += 1;
}
let mut queue_removing = queue.removing();
while let Some(entry) = queue_removing.next() {
let chat = &entry.value().0;
let slow_mode = slow_mode.as_mut().and_then(|sm| sm.get_mut(chat));
if let Some(&mut (delay, last)) = slow_mode {
if last + delay > Instant::now() {
continue;
}
}
let requests_sent_per_sec_count = requests_sent.per_sec.get(chat).copied().unwrap_or(0);
let requests_sent_per_min_count = requests_sent.per_min.get(chat).copied().unwrap_or(0);
let messages_per_min_limit = if chat.is_channel() {
limits.messages_per_min_channel
} else {
limits.messages_per_min_chat
};
let limits_not_exceeded = requests_sent_per_sec_count < limits.messages_per_sec_chat
&& requests_sent_per_min_count < messages_per_min_limit;
if limits_not_exceeded {
// Unlock the associated request.
let chat = *chat;
let (_, lock) = entry.remove();
// Only count request as sent if the request wasn't dropped before unlocked
if lock.unlock(retry, freeze_tx.clone()).is_ok() {
*requests_sent.per_sec.entry(chat).or_insert(0) += 1;
*requests_sent.per_min.entry(chat).or_insert(0) += 1;
history.push_back((chat, Instant::now()));
if let Some((_, last)) = slow_mode {
*last = Instant::now();
}
// We have "sent" one request, so now we can send one less.
allowed -= 1;
if allowed == 0 {
break;
}
}
}
}
// It's easier to just recompute last second stats, instead of keeping
// track of it alongside with minute stats, so we just throw this away.
requests_sent.per_sec.clear();
tokio::time::sleep(DELAY).await;
}
}
fn answer_info(rx: &mut mpsc::Receiver<InfoMessage>, limits: &mut Limits) {
while let Ok(req) = rx.try_recv() {
// Errors are ignored with .ok(). Error means that the response channel
// is closed and the response isn't needed.
match req {
InfoMessage::GetLimits { response } => response.send(*limits).ok(),
InfoMessage::SetLimits { new, response } => {
*limits = new;
response.send(()).ok()
}
};
}
}
async fn freeze(
rx: &mut mpsc::Receiver<FreezeUntil>,
mut slow_mode: Option<&mut HashMap<ChatIdHash, (Duration, Instant)>>,
bot: &impl Requester,
mut imm: Option<FreezeUntil>,
) {
while let Some(freeze_until) = imm.take().or_else(|| rx.try_recv().ok()) {
let FreezeUntil { until, after, chat } = freeze_until;
// Clippy thinks that this `.as_deref_mut()` doesn't change the type (&mut
// HashMap -> &mut HashMap), but it's actually a reborrow (the lifetimes
// differ), since we are in a loop, simply using `slow_mode` would produce a
// moved-out error.
#[allow(clippy::needless_option_as_deref)]
if let Some(slow_mode) = slow_mode.as_deref_mut() {
// TODO: do something with channels?...
if let hash @ ChatIdHash::Id(id) = chat {
// TODO: maybe not call `get_chat` every time?
// At this point there isn't much we can do with the error besides ignoring
if let Ok(chat) = bot.get_chat(id).send().await {
match chat.slow_mode_delay() {
Some(delay) => {
let now = Instant::now();
let new_delay = Duration::from_secs(delay.into());
slow_mode.insert(hash, (new_delay, now));
}
None => {
slow_mode.remove(&hash);
}
};
}
}
}
// slow mode is enabled and it is <= to the delay asked by telegram
let slow_mode_enabled_and_likely_the_cause = slow_mode
.as_ref()
.and_then(|m| m.get(&chat).map(|(delay, _)| delay <= &after))
.unwrap_or(false);
// Do not sleep if slow mode is enabled since the freeze is most likely caused
// by the said slow mode and not by the global limits.
if !slow_mode_enabled_and_likely_the_cause {
log::warn!(
"freezing the bot for approximately {:?} due to `RetryAfter` error from telegram",
after
);
tokio::time::sleep_until(until.into()).await;
log::warn!("unfreezing the bot");
}
}
}
async fn read_from_rx<T>(rx: &mut mpsc::Receiver<T>, queue: &mut Vec<T>, rx_is_closed: &mut bool) {
if queue.is_empty() {
log::warn!("A-blocking on queue");
match rx.recv().await {
Some(req) => queue.push(req),
None => *rx_is_closed = true,
}
}
// Don't grow queue bigger than the capacity to limit DOS possibility
while queue.len() < queue.capacity() {
match rx.try_recv() {
Ok(req) => queue.push(req),
Err(TryRecvError::Disconnected) => {
*rx_is_closed = true;
break;
}
// There are no items in queue.
Err(TryRecvError::Empty) => break,
}
}
}
#[cfg(test)]
mod tests {
#[tokio::test]
async fn issue_535() {
let (tx, mut rx) = tokio::sync::mpsc::channel(1);
// Close channel
drop(tx);
// Previously this caused an infinite loop
super::read_from_rx::<()>(&mut rx, &mut Vec::new(), &mut false).await;
}
}

60
src/errors.rs
View file

@ -1,20 +1,9 @@
use std::io;
use std::{io, time::Duration};
use serde::Deserialize;
use thiserror::Error;
/// An error caused by downloading a file.
#[derive(Debug, Error)]
pub enum DownloadError {
/// A network error while downloading a file from Telegram.
#[error("A network error: {0}")]
// NOTE: this variant must not be created by anything except the From impl
Network(#[source] reqwest::Error),
/// An I/O error while writing a file to destination.
#[error("An I/O error: {0}")]
Io(#[from] std::io::Error),
}
use crate::types::ResponseParameters;
/// An error caused by sending a request to Telegram.
#[derive(Debug, Error)]
@ -30,8 +19,8 @@ pub enum RequestError {
/// In case of exceeding flood control, the number of seconds left to wait
/// before the request can be repeated.
#[error("Retry after {0} seconds")]
RetryAfter(i32),
#[error("Retry after {0:?}")]
RetryAfter(Duration),
/// Network error while sending a request to Telegram.
#[error("A network error: {0}")]
@ -57,6 +46,47 @@ pub enum RequestError {
Io(#[source] io::Error),
}
/// An error caused by downloading a file.
#[derive(Debug, Error)]
pub enum DownloadError {
/// A network error while downloading a file from Telegram.
#[error("A network error: {0}")]
// NOTE: this variant must not be created by anything except the From impl
Network(#[source] reqwest::Error),
/// An I/O error while writing a file to destination.
#[error("An I/O error: {0}")]
Io(#[from] std::io::Error),
}
pub trait AsResponseParameters {
fn response_parameters(&self) -> Option<ResponseParameters>;
fn retry_after(&self) -> Option<Duration> {
self.response_parameters().and_then(|rp| match rp {
ResponseParameters::RetryAfter(n) => Some(n),
_ => None,
})
}
fn migrate_to_chat_id(&self) -> Option<i64> {
self.response_parameters().and_then(|rp| match rp {
ResponseParameters::MigrateToChatId(id) => Some(id),
_ => None,
})
}
}
impl AsResponseParameters for crate::RequestError {
fn response_parameters(&self) -> Option<ResponseParameters> {
match *self {
Self::RetryAfter(n) => Some(ResponseParameters::RetryAfter(n)),
Self::MigrateToChatId(id) => Some(ResponseParameters::MigrateToChatId(id)),
_ => None,
}
}
}
/// A kind of an API error.
#[derive(Debug, Error, Deserialize, PartialEq, Hash, Eq, Clone)]
#[serde(field_identifier)]

2
src/lib.rs
View file

@ -101,6 +101,7 @@ pub use self::{
};
pub mod adaptors;
pub mod errors;
pub mod net;
pub mod payloads;
pub mod prelude;
@ -109,7 +110,6 @@ pub mod types;
// reexported
mod bot;
mod errors;
// implementation details
mod serde_multipart;

1
src/requests/json.rs
View file

@ -10,6 +10,7 @@ use crate::{
///
/// [JSON]: https://core.telegram.org/bots/api#making-requests
#[must_use = "Requests are lazy and do nothing unless sent"]
#[derive(Clone)]
pub struct JsonRequest<P> {
bot: Bot,
payload: P,

1
src/requests/multipart.rs
View file

@ -11,6 +11,7 @@ use crate::{
///
/// [multipart/form-data]: https://core.telegram.org/bots/api#making-requests
#[must_use = "Requests are lazy and do nothing unless sent"]
#[derive(Clone)]
pub struct MultipartRequest<P> {
bot: Bot,
payload: P,

8
src/requests/requester_ext.rs
View file

@ -1,4 +1,6 @@
use crate::{adaptors::DefaultParseMode, requests::Requester, types::ParseMode};
use crate::{
adaptors::DefaultParseMode, errors::AsResponseParameters, requests::Requester, types::ParseMode,
};
#[cfg(feature = "cache_me")]
use crate::adaptors::CacheMe;
@ -60,7 +62,9 @@ pub trait RequesterExt: Requester {
#[cfg(feature = "throttle")]
fn throttle(self, limits: Limits) -> Throttle<Self>
where
Self: Sized,
Self: Sized + Clone + Send + Sync + 'static,
Self::Err: AsResponseParameters,
Self::GetChat: Send,
{
Throttle::new_spawn(self, limits)
}

50
src/types.rs
View file

@ -345,3 +345,53 @@ pub(crate) mod option_url_from_string {
}
}
}
pub(crate) mod duration_secs {
use std::time::Duration;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
pub(crate) fn serialize<S>(this: &Duration, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
// match this {
// Some(url) => url.serialize(serializer),
// None => "".serialize(serializer),
// }
this.as_secs().serialize(serializer)
}
pub(crate) fn deserialize<'de, D>(deserializer: D) -> Result<Duration, D::Error>
where
D: Deserializer<'de>,
{
u64::deserialize(deserializer).map(Duration::from_secs)
}
#[test]
fn test() {
#[derive(Serialize, Deserialize)]
struct Struct {
#[serde(with = "crate::types::duration_secs")]
duration: Duration,
}
{
let json = r#"{"duration":0}"#;
let duration: Struct = serde_json::from_str(json).unwrap();
assert_eq!(duration.duration, Duration::from_secs(0));
assert_eq!(serde_json::to_string(&duration).unwrap(), json.to_owned());
let json = r#"{"duration":12}"#;
let duration: Struct = serde_json::from_str(json).unwrap();
assert_eq!(duration.duration, Duration::from_secs(12));
assert_eq!(serde_json::to_string(&duration).unwrap(), json.to_owned());
let json = r#"{"duration":1234}"#;
let duration: Struct = serde_json::from_str(json).unwrap();
assert_eq!(duration.duration, Duration::from_secs(1234));
assert_eq!(serde_json::to_string(&duration).unwrap(), json.to_owned());
}
}
}

6
src/types/response_parameters.rs
View file

@ -1,3 +1,5 @@
use std::time::Duration;
use serde::{Deserialize, Serialize};
/// Contains information about why a request was unsuccessful.
@ -16,7 +18,7 @@ pub enum ResponseParameters {
/// In case of exceeding flood control, the number of seconds left to wait
/// before the request can be repeated.
RetryAfter(i32),
RetryAfter(#[serde(with = "crate::types::duration_secs")] Duration),
}
#[cfg(test)]
@ -34,7 +36,7 @@ mod tests {
#[test]
fn retry_after_deserialization() {
let expected = ResponseParameters::RetryAfter(123_456);
let expected = ResponseParameters::RetryAfter(Duration::from_secs(123_456));
let actual: ResponseParameters = serde_json::from_str(r#"{"retry_after":123456}"#).unwrap();
assert_eq!(expected, actual);