mirror of
https://github.com/teloxide/teloxide.git
synced 2024-10-24 09:57:18 +02:00
Merge pull request #130 from teloxide/chrottle_retries_and_freeze
`Throttle` retries and freeze
This commit is contained in:
commit
4e35f6de31
15 changed files with 1032 additions and 837 deletions
12
CHANGELOG.md
12
CHANGELOG.md
|
@ -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
|
||||
|
||||
|
|
|
@ -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 = []
|
||||
|
|
|
@ -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
204
src/adaptors/throttle/request.rs
Normal 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
45
src/adaptors/throttle/request_lock.rs
Normal 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
103
src/adaptors/throttle/requester_impl.rs
Normal 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
106
src/adaptors/throttle/settings.rs
Normal 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
395
src/adaptors/throttle/worker.rs
Normal 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;
|
||||
}
|
||||
}
|
|
@ -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)]
|
||||
|
|
|
@ -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;
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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
50
src/types.rs
|
@ -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());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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);
|
||||
|
|
Loading…
Reference in a new issue