rivet_envoy_client/
http.rs

1use std::collections::HashMap;
2
3use rivet_envoy_protocol as protocol;
4use tokio::sync::{mpsc, watch};
5
6pub const HTTP_BODY_STREAM_CHANNEL_CAPACITY: usize = 16;
7pub const HTTP_BODY_MAX_CHUNK_SIZE: usize = 64 * 1024;
8
9#[derive(Clone, Debug)]
10pub struct HttpRequestBodyError {
11	pub reason: protocol::HttpStreamAbortReason,
12}
13
14impl std::fmt::Display for HttpRequestBodyError {
15	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
16		match &self.reason.detail {
17			Some(detail) => write!(f, "{:?}: {detail}", self.reason.kind),
18			None => write!(f, "{:?}", self.reason.kind),
19		}
20	}
21}
22
23impl std::error::Error for HttpRequestBodyError {}
24
25#[derive(Debug)]
26pub struct HttpRequestBodyStream {
27	rx: mpsc::Receiver<Vec<u8>>,
28	abort_rx: watch::Receiver<Option<HttpRequestBodyError>>,
29}
30
31impl HttpRequestBodyStream {
32	pub fn new(
33		rx: mpsc::Receiver<Vec<u8>>,
34		abort_rx: watch::Receiver<Option<HttpRequestBodyError>>,
35	) -> Self {
36		Self { rx, abort_rx }
37	}
38
39	pub async fn recv(&mut self) -> Result<Option<Vec<u8>>, HttpRequestBodyError> {
40		loop {
41			if let Some(error) = self.abort_rx.borrow().clone() {
42				return Err(error);
43			}
44
45			tokio::select! {
46				biased;
47				changed = self.abort_rx.changed() => {
48					if changed.is_ok() {
49						continue;
50					}
51					return Ok(self.rx.recv().await);
52				}
53				chunk = self.rx.recv() => return Ok(chunk),
54			}
55		}
56	}
57}
58
59/// HTTP request/response types used by the envoy client.
60pub struct HttpRequest {
61	pub method: String,
62	pub path: String,
63	pub headers: HashMap<String, String>,
64	pub body: Option<Vec<u8>>,
65	/// If the request is streamed, body chunks arrive on this channel.
66	pub body_stream: Option<HttpRequestBodyStream>,
67}
68
69pub struct HttpResponse {
70	pub status: u16,
71	pub headers: HashMap<String, String>,
72	pub body: Option<Vec<u8>>,
73	/// If set, the response is streamed. The envoy client reads chunks and sends
74	/// `ToRivetResponseChunk` for each one.
75	pub body_stream: Option<HttpResponseBodyStream>,
76}
77
78/// A chunk in a streaming HTTP response.
79pub enum ResponseChunk {
80	Data { data: Vec<u8>, finish: bool },
81	Error(String),
82}
83
84pub struct HttpResponseBodyStream {
85	rx: mpsc::Receiver<ResponseChunk>,
86	on_drop: Option<Box<dyn FnOnce() + Send>>,
87}
88
89impl HttpResponseBodyStream {
90	pub fn set_on_drop(&mut self, on_drop: impl FnOnce() + Send + 'static) {
91		self.on_drop = Some(Box::new(on_drop));
92	}
93
94	pub async fn recv(&mut self) -> Option<ResponseChunk> {
95		self.rx.recv().await
96	}
97}
98
99impl From<mpsc::Receiver<ResponseChunk>> for HttpResponseBodyStream {
100	fn from(rx: mpsc::Receiver<ResponseChunk>) -> Self {
101		Self { rx, on_drop: None }
102	}
103}
104
105impl Drop for HttpResponseBodyStream {
106	fn drop(&mut self) {
107		if let Some(on_drop) = self.on_drop.take() {
108			on_drop();
109		}
110	}
111}
rivet_envoy_client/http.rs

rivet_envoy_client/
http.rs
