1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
use std::{
    future::{Future, IntoFuture},
    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"]
#[derive(Clone)]
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> IntoFuture for ThrottlingRequest<R>
where
    R: Request + Clone + Send + Sync + 'static,
    R::Err: AsResponseParameters + Send,
    Output<R>: Send,
{
    type Output = Result<Output<Self>, <Self as Request>::Err>;
    type IntoFuture = <Self as Request>::Send;

    fn into_future(self) -> Self::IntoFuture {
        self.send()
    }
}

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().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().await,
        };

        let retry_after = res.as_ref().err().and_then(<_>::retry_after);
        if let Some(retry_after) = retry_after {
            let after = retry_after.duration();
            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,
        };
    }
}