Line data Source code
1 : use http::header::{HeaderMap, HeaderValue};
2 : use std::collections::HashMap;
3 : use std::fmt;
4 : use std::fmt::Display;
5 : use std::sync::Arc;
6 : use std::sync::RwLock;
7 : use std::sync::atomic::AtomicUsize;
8 : use std::time::Duration;
9 : use std::time::Instant;
10 : use tokio::sync::Mutex;
11 : use tokio::time::timeout;
12 :
13 : use crate::status::Status;
14 :
15 : #[derive(Debug)]
16 : pub enum HttpError {
17 : ErrorStatus(reqwest::StatusCode),
18 : ConnectError(String),
19 : OtherError(String),
20 : }
21 :
22 : impl Display for HttpError {
23 1 : fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
24 1 : match &self {
25 0 : HttpError::ErrorStatus(status) => {
26 0 : write!(f, "{}", &status)
27 : }
28 1 : HttpError::ConnectError(e) => {
29 1 : write!(f, "{}", e)
30 : }
31 0 : HttpError::OtherError(e) => {
32 0 : write!(f, "{}", e)
33 : },
34 : }
35 1 : }
36 : }
37 :
38 : impl From<&reqwest::Error> for HttpError {
39 1 : fn from(e: &reqwest::Error) -> HttpError {
40 : // Is a status?
41 1 : if let Some(s) = e.status() {
42 0 : return HttpError::ErrorStatus(s);
43 1 : }
44 :
45 : // Inpack inner error
46 1 : let mut inner: &dyn std::error::Error = &e;
47 4 : while let Some(i) = inner.source() {
48 3 : inner = i;
49 3 : }
50 :
51 1 : if e.is_connect() {
52 1 : return HttpError::ConnectError(inner.to_string());
53 : } else {
54 0 : return HttpError::OtherError(e.to_string());
55 : }
56 1 : }
57 : }
58 :
59 : #[derive(Debug)]
60 : pub enum HttpServerStreamEvent {
61 : RequestFailed(HttpError),
62 : Header(reqwest::StatusCode, HeaderMap<HeaderValue>),
63 : BodyChunk(Vec<u8>),
64 : BodyEnd,
65 : BodyReadFailed(String),
66 : }
67 :
68 : #[derive(Debug, PartialEq)]
69 : pub enum HttpServerStreamBatchEvent {
70 : StreamFailed,
71 : StreamFinished,
72 : FlushAfterClose,
73 : FlushAfterTimeout,
74 : FlushAfterBytes,
75 : }
76 :
77 : impl HttpServerStreamBatchEvent {
78 10 : pub fn to_str(&self) -> &'static str {
79 10 : match self {
80 0 : HttpServerStreamBatchEvent::StreamFailed => "StreamFailed",
81 5 : HttpServerStreamBatchEvent::StreamFinished => "StreamFinished",
82 0 : HttpServerStreamBatchEvent::FlushAfterClose => "FlushAfterClose",
83 3 : HttpServerStreamBatchEvent::FlushAfterTimeout => "FlushAfterTimeout",
84 2 : HttpServerStreamBatchEvent::FlushAfterBytes => "FlushAfterBytes",
85 : }
86 10 : }
87 : }
88 :
89 : /// A batch of response data read from an HTTP server stream in one read call
90 : #[derive(Debug)]
91 : pub struct HttpServerStreamBatch {
92 : /// The event that emitted the batch
93 : pub event: HttpServerStreamBatchEvent,
94 : /// The status
95 : pub status: Option<reqwest::StatusCode>,
96 : /// The headers
97 : pub headers: HeaderMap<HeaderValue>,
98 : /// The body chunks
99 : pub body_chunks: Vec<Vec<u8>>,
100 : /// The total body bytes in the batch
101 : pub total_body_bytes: usize,
102 : }
103 :
104 : impl Default for HttpServerStreamBatch {
105 13 : fn default() -> Self {
106 13 : Self {
107 13 : event: HttpServerStreamBatchEvent::StreamFailed,
108 13 : status: None,
109 13 : headers: HeaderMap::new(),
110 13 : body_chunks: Vec::default(),
111 13 : total_body_bytes: 0,
112 13 : }
113 13 : }
114 : }
115 :
116 : /// An HTTP response stream whose events are buffered by a background reader task
117 : pub struct HttpServerStream {
118 : /// The response queue reading. Dropping it stops the reader task, which cancels the request.
119 : response_receiver: Mutex<tokio::sync::mpsc::Receiver<HttpServerStreamEvent>>,
120 : }
121 :
122 : /// The registry of open HTTP server streams, keyed by stream id
123 : #[derive(Default)]
124 : pub struct HttpStreamManager {
125 : /// The next message id
126 : pub next_stream_id: AtomicUsize,
127 : /// The server streams
128 : pub server_streams: RwLock<HashMap<usize, Arc<HttpServerStream>>>,
129 : }
130 :
131 : impl HttpStreamManager {
132 : /// Start a server stream
133 : #[allow(dead_code)]
134 16 : pub fn start_server_stream(
135 16 : self: &Arc<Self>,
136 16 : client: reqwest::Client,
137 16 : request: reqwest::Request,
138 16 : ) -> Result<usize, Status> {
139 :
140 : // Allocate an identifier for the stream
141 16 : let reg = self.clone();
142 16 : let stream_id = reg
143 16 : .next_stream_id
144 16 : .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
145 :
146 : // Register the receiver
147 : const STREAM_CAPACITY: usize = 10;
148 16 : let (sender, receiver) = tokio::sync::mpsc::channel(STREAM_CAPACITY);
149 16 : let stream_entry = Arc::new(HttpServerStream {
150 16 : response_receiver: Mutex::new(receiver),
151 16 : });
152 16 : if let Ok(mut streams) = self.server_streams.write() {
153 16 : streams.insert(stream_id, stream_entry);
154 16 : }
155 :
156 : // Spawn the async reader
157 16 : tokio::spawn(async move {
158 16 : let reader = async {
159 : // Execute a request
160 16 : let mut response = match client.execute(request).await {
161 10 : Ok(response) => response,
162 1 : Err(e) => {
163 1 : let mapped = HttpError::from(&e);
164 1 : let queue_error = sender
165 1 : .send(HttpServerStreamEvent::RequestFailed(mapped))
166 1 : .await;
167 1 : if let Err(e) = queue_error {
168 0 : log::warn!("Failed forwarding request error: {}", e);
169 0 : return;
170 1 : }
171 1 : return;
172 : }
173 : };
174 :
175 : // Send the header
176 10 : if let Err(e) = sender
177 10 : .send(HttpServerStreamEvent::Header(response.status(), std::mem::take(response.headers_mut())))
178 10 : .await
179 : {
180 0 : log::warn!("Failed writing header response to stream: {}", e);
181 0 : return;
182 10 : }
183 :
184 : // Poll the result stream
185 10 : let mut stream_alive = true;
186 35 : while stream_alive {
187 : // Read the next message from the body stream
188 29 : let response = match response.chunk().await {
189 : // Received bytes, forward the message as-is to receivers
190 20 : Ok(Some(m)) => {
191 20 : let buffer: Vec<u8> = m.into();
192 20 : log::debug!(
193 : "Receiving chunk from server stream {}, bytes={}",
194 : stream_id,
195 0 : buffer.len()
196 : );
197 20 : HttpServerStreamEvent::BodyChunk(buffer)
198 : }
199 : // Stream was closed, delete stream and return trailers
200 : Ok(None) => {
201 5 : stream_alive = false;
202 5 : log::debug!("Server stream {} ended", stream_id);
203 5 : HttpServerStreamEvent::BodyEnd
204 : }
205 : // Stream failed, send error to receivers and close the stream
206 1 : Err(e) => {
207 1 : stream_alive = false;
208 1 : log::warn!("Failed reading from server stream: {}", e);
209 1 : HttpServerStreamEvent::BodyReadFailed(e.to_string())
210 : }
211 : };
212 :
213 : // Push the response into the queue
214 26 : if let Err(e) = sender
215 26 : .send(response)
216 26 : .await
217 : {
218 0 : stream_alive = false;
219 0 : log::warn!("Failed writing response to stream: {}", e);
220 25 : }
221 : }
222 7 : };
223 : // A deleted stream drops the receiver, dropping the reader then cancels the request
224 16 : tokio::select! {
225 16 : _ = sender.closed() => log::debug!("Server stream {} was deleted", stream_id),
226 16 : _ = reader => {}
227 : }
228 10 : });
229 :
230 : // Return initial metadata
231 16 : Ok(stream_id)
232 16 : }
233 :
234 : /// Read from a server stream
235 : #[allow(dead_code)]
236 17 : pub async fn read_server_stream(
237 17 : self: &Arc<Self>,
238 17 : stream_id: usize,
239 17 : timeout_read_after: Duration,
240 17 : flush_batch_after: Duration,
241 17 : flush_batch_bytes: usize,
242 17 : ) -> Result<HttpServerStreamBatch, Status> {
243 17 : let started_at = Instant::now();
244 :
245 : // Try to find the stream.
246 : // If there is none, the request might have finished already.
247 17 : let reg = self.clone();
248 17 : let stream = if let Some(streams) = reg.server_streams.read().unwrap().get(&stream_id) {
249 13 : streams.clone()
250 : } else {
251 4 : log::debug!("Unknown stream id: {}", stream_id);
252 4 : return Err(Status::HttpStreamIsUnknown { stream_id });
253 : };
254 :
255 : // Acquire the receiver lock.
256 : // There can only be one active reader at a time.
257 13 : let mut receiver = stream.response_receiver.lock().await;
258 :
259 : // Receive messages and return them in batches
260 13 : let mut batch = HttpServerStreamBatch::default();
261 : loop {
262 : // Receive the next message from the reader
263 30 : let elapsed_since_start = started_at.elapsed();
264 30 : let response = if batch.body_chunks.len() == 0 && batch.status.is_none() {
265 : // Initially, we check the `timeout_read_after` timeout.
266 13 : let receive_timeout = timeout_read_after.checked_sub(elapsed_since_start).unwrap_or_default();
267 13 : match timeout(receive_timeout, receiver.recv()).await {
268 12 : Ok(Some(response)) => {
269 12 : log::debug!("Http stream {} receiving event", stream_id);
270 12 : response
271 : },
272 : Ok(None) => {
273 0 : log::debug!("Http stream {} closed", stream_id);
274 0 : return Err(Status::HttpStreamClosed { stream_id });
275 : },
276 : Err(_) => {
277 1 : log::debug!("Http stream {} read timed out", stream_id);
278 1 : return Err(Status::HttpStreamReadTimedOut { stream_id });
279 : }
280 : }
281 : } else {
282 : // After receiving one batch or the status, we check the `flush_batch_after` timeout.
283 : // A status must not get lost when the body takes longer than the read timeout.
284 17 : let receive_timeout = match flush_batch_after.checked_sub(elapsed_since_start) {
285 17 : Some(timeout) => {
286 17 : timeout
287 : },
288 : None => {
289 0 : batch.event = HttpServerStreamBatchEvent::FlushAfterTimeout;
290 0 : return Ok(batch);
291 : }
292 : };
293 17 : match timeout(receive_timeout, receiver.recv()).await {
294 14 : Ok(Some(response)) => {
295 14 : log::debug!("Http stream {} returning event", stream_id);
296 14 : response
297 : },
298 : Ok(None) => {
299 0 : log::debug!("Http stream {} closed", stream_id);
300 0 : batch.event = HttpServerStreamBatchEvent::FlushAfterClose;
301 0 : return Ok(batch);
302 : },
303 : Err(_) => {
304 3 : log::debug!("Http stream {} read timed out", stream_id);
305 3 : batch.event = HttpServerStreamBatchEvent::FlushAfterTimeout;
306 3 : return Ok(batch);
307 : }
308 : }
309 : };
310 :
311 26 : match response {
312 : // Request failed?
313 1 : HttpServerStreamEvent::RequestFailed(e) => {
314 1 : log::debug!("Http stream {} request failed: {}", stream_id, e);
315 1 : if let Ok(mut streams) = self.server_streams.write() {
316 1 : streams.remove(&stream_id);
317 1 : }
318 1 : return Err(Status::HttpRequestFailed { stream_id, error: e.to_string() });
319 : }
320 : // Received headers
321 9 : HttpServerStreamEvent::Header(status, headers) => {
322 9 : log::debug!("Http stream {} receiving header", stream_id);
323 9 : batch.status = Some(status);
324 9 : batch.headers = headers;
325 : },
326 : // Return the message
327 10 : HttpServerStreamEvent::BodyChunk(m) => {
328 10 : log::debug!("Http stream {} receiving body chunk", stream_id);
329 : // Add the message to the current batch and check if we should flush
330 10 : batch.total_body_bytes += m.len();
331 10 : batch.body_chunks.push(m);
332 :
333 : // Flush if we hit the message size
334 10 : if batch.total_body_bytes > flush_batch_bytes {
335 2 : batch.event = HttpServerStreamBatchEvent::FlushAfterBytes;
336 2 : return Ok(batch);
337 8 : }
338 : },
339 : // An error occurred,.
340 : // Throw away any intermediate messages that we held back.
341 1 : HttpServerStreamEvent::BodyReadFailed(e) => {
342 1 : log::debug!("Http stream {} failed reading body", stream_id);
343 1 : if let Ok(mut streams) = self.server_streams.write() {
344 1 : streams.remove(&stream_id);
345 1 : }
346 1 : return Err(Status::HttpStreamReadFailed{
347 1 : stream_id,
348 1 : error: e
349 1 : });
350 : },
351 : // We reached the end of the stream.
352 : // Flush the current batch with trailers and an EOS marker.
353 : HttpServerStreamEvent::BodyEnd => {
354 5 : log::debug!("Http stream {} ended", stream_id);
355 5 : if let Ok(mut streams) = self.server_streams.write() {
356 5 : streams.remove(&stream_id);
357 5 : }
358 5 : batch.event = HttpServerStreamBatchEvent::StreamFinished;
359 5 : return Ok(batch);
360 : }
361 : }
362 : }
363 17 : }
364 :
365 : /// Cancellation is done by just dropping everything, see `start_server_stream`
366 : #[allow(dead_code)]
367 5 : pub async fn destroy_server_stream(self: &Arc<Self>, stream_id: usize) {
368 5 : if let Ok(mut streams) = self.server_streams.write() {
369 5 : streams.remove(&stream_id);
370 5 : }
371 5 : }
372 : }
|
