move responsewriter to its own module

lovasoa · lovasoa · commit 60ada622c154 · 2024-10-25T23:53:36.000+03:00
diff --git a/src/app_config.rs b/src/app_config.rs
@@ -230,6 +230,7 @@ pub struct AppConfig {
     pub site_prefix: String,
 
     /// Maximum number of messages that can be stored in memory before sending them to the client.
+    /// This prevents a single request from using up all available memory.
     #[serde(default = "default_max_pending_rows")]
     pub max_pending_rows: usize,
 
@@ -470,6 +471,8 @@ fn default_https_acme_directory_url() -> String {
     "https://acme-v02.api.letsencrypt.org/directory".to_string()
 }
 
+/// If the sending queue exceeds this number of outgoing messages, an error will be thrown
+/// This prevents a single request from using up all available memory
 fn default_max_pending_rows() -> usize {
     256
 }
diff --git a/src/render.rs b/src/render.rs
@@ -1,5 +1,6 @@
 use crate::templates::SplitTemplate;
-use crate::webserver::http::{AsyncResponseWriter, RequestContext, ResponseWriter};
+use crate::webserver::http::RequestContext;
+use crate::webserver::response_writer::{AsyncResponseWriter, ResponseWriter};
 use crate::webserver::ErrorWithStatus;
 use crate::AppState;
 use actix_web::cookie::time::format_description::well_known::Rfc3339;
diff --git a/src/webserver/http.rs b/src/webserver/http.rs
@@ -11,167 +11,32 @@ use actix_web::http::header::{ContentType, Header, HttpDate, IfModifiedSince, La
 use actix_web::http::{header, StatusCode, Uri};
 use actix_web::web::PayloadConfig;
 use actix_web::{
-    dev::ServiceResponse, middleware, middleware::Logger, web, web::Bytes, App, HttpResponse,
-    HttpServer,
+    dev::ServiceResponse, middleware, middleware::Logger, web, App, HttpResponse, HttpServer,
 };
 use actix_web::{HttpResponseBuilder, ResponseError};
 
 use super::https::make_auto_rustls_config;
+use super::response_writer::ResponseWriter;
 use super::static_content;
 use actix_web::body::MessageBody;
 use anyhow::{bail, Context};
 use chrono::{DateTime, Utc};
 use futures_util::stream::Stream;
 use futures_util::StreamExt;
 use std::borrow::Cow;
-use std::io::Write;
 use std::mem;
 use std::path::PathBuf;
 use std::pin::Pin;
 use std::sync::Arc;
-use std::task::Poll;
 use std::time::SystemTime;
 use tokio::sync::mpsc;
 
-/// If the sending queue exceeds this number of outgoing messages, an error will be thrown
-/// This prevents a single request from using up all available memory
-
-#[derive(Clone)]
-pub struct ResponseWriter {
-    buffer: Vec<u8>,
-    response_bytes: mpsc::Sender<Bytes>,
-}
-
 #[derive(Clone)]
 pub struct RequestContext {
     pub is_embedded: bool,
     pub content_security_policy: ContentSecurityPolicy,
 }
 
-impl ResponseWriter {
-    fn new(response_bytes: mpsc::Sender<Bytes>) -> Self {
-        Self {
-            response_bytes,
-            buffer: Vec::new(),
-        }
-    }
-    async fn close_with_error(&mut self, mut msg: String) {
-        if !self.response_bytes.is_closed() {
-            if let Err(e) = self.async_flush().await {
-                use std::fmt::Write;
-                write!(&mut msg, "Unable to flush data: {e}").unwrap();
-            }
-            if let Err(e) = self.response_bytes.send(msg.into()).await {
-                log::error!("Unable to send error back to client: {e}");
-            }
-        }
-    }
-
-    pub async fn async_flush(&mut self) -> std::io::Result<()> {
-        if self.buffer.is_empty() {
-            return Ok(());
-        }
-        log::trace!(
-            "Flushing data to client: {}",
-            String::from_utf8_lossy(&self.buffer)
-        );
-        let sender = self
-            .response_bytes
-            .reserve()
-            .await
-            .map_err(|_| std::io::ErrorKind::WouldBlock)?;
-        sender.send(std::mem::take(&mut self.buffer).into());
-        Ok(())
-    }
-}
-
-impl Write for ResponseWriter {
-    #[inline]
-    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
-        self.buffer.extend_from_slice(buf);
-        Ok(buf.len())
-    }
-    fn flush(&mut self) -> std::io::Result<()> {
-        if self.buffer.is_empty() {
-            return Ok(());
-        }
-        log::trace!(
-            "Flushing data to client: {}",
-            String::from_utf8_lossy(&self.buffer)
-        );
-        self.response_bytes
-            .try_send(mem::take(&mut self.buffer).into())
-            .map_err(|e|
-                std::io::Error::new(
-                    std::io::ErrorKind::WouldBlock,
-                    format!("{e}: Row limit exceeded. The server cannot store more than {} pending messages in memory. Try again later or increase max_pending_rows in the configuration.", self.response_bytes.max_capacity())
-                )
-            )
-    }
-}
-
-pub struct AsyncResponseWriter {
-    poll_sender: tokio_util::sync::PollSender<Bytes>,
-    writer: ResponseWriter,
-}
-
-impl AsyncResponseWriter {
-    #[must_use]
-    pub fn new(writer: ResponseWriter) -> Self {
-        let sender = writer.response_bytes.clone();
-        Self {
-            poll_sender: tokio_util::sync::PollSender::new(sender),
-            writer,
-        }
-    }
-    #[must_use]
-    pub fn into_inner(self) -> ResponseWriter {
-        self.writer
-    }
-}
-impl tokio::io::AsyncWrite for AsyncResponseWriter {
-    fn poll_write(
-        mut self: Pin<&mut Self>,
-        _cx: &mut std::task::Context<'_>,
-        buf: &[u8],
-    ) -> std::task::Poll<std::io::Result<usize>> {
-        Poll::Ready(self.as_mut().writer.write(buf))
-    }
-
-    fn poll_flush(
-        self: Pin<&mut Self>,
-        cx: &mut std::task::Context<'_>,
-    ) -> std::task::Poll<std::io::Result<()>> {
-        let Self {
-            poll_sender,
-            writer,
-        } = self.get_mut();
-        match poll_sender.poll_reserve(cx) {
-            Poll::Ready(Ok(())) => {
-                let res = poll_sender.send_item(std::mem::take(&mut writer.buffer).into());
-                Poll::Ready(res.map_err(|_| std::io::ErrorKind::BrokenPipe.into()))
-            }
-            Poll::Pending => Poll::Pending,
-            Poll::Ready(Err(_e)) => Poll::Ready(Err(std::io::ErrorKind::BrokenPipe.into())),
-        }
-    }
-
-    fn poll_shutdown(
-        self: Pin<&mut Self>,
-        cx: &mut std::task::Context<'_>,
-    ) -> std::task::Poll<Result<(), std::io::Error>> {
-        self.poll_flush(cx)
-    }
-}
-
-impl Drop for ResponseWriter {
-    fn drop(&mut self) {
-        if let Err(e) = std::io::Write::flush(self) {
-            log::error!("Could not flush data to client: {e}");
-        }
-    }
-}
-
 async fn stream_response(stream: impl Stream<Item = DbItem>, mut renderer: AnyRenderBodyContext) {
     let mut stream = Box::pin(stream);
 
diff --git a/src/webserver/mod.rs b/src/webserver/mod.rs
@@ -11,4 +11,5 @@ pub use error_with_status::ErrorWithStatus;
 
 pub use database::make_placeholder;
 pub use database::migrations::apply;
+pub mod response_writer;
 mod static_content;
diff --git a/src/webserver/response_writer.rs b/src/webserver/response_writer.rs
@@ -0,0 +1,149 @@
+use actix_web::web::Bytes;
+use std::io::Write;
+use std::mem;
+use std::pin::Pin;
+use tokio::sync::mpsc;
+
+/// The response writer is a buffered async writer that sends data to the client.
+/// Writing to it just appends to an in-memory buffer, which is flushed to the client asynchronously
+/// when `async_flush()` is called.
+/// This allows streaming data to the client without blocking, and has built-in back-pressure:
+/// if the client cannot keep up with the data, `async_flush()` will fill the sending queue,
+/// then block until the client has consumed some data.
+#[derive(Clone)]
+pub struct ResponseWriter {
+    buffer: Vec<u8>,
+    response_bytes: mpsc::Sender<Bytes>,
+}
+
+impl ResponseWriter {
+    #[must_use]
+    pub fn new(response_bytes: mpsc::Sender<Bytes>) -> Self {
+        Self {
+            response_bytes,
+            buffer: Vec::new(),
+        }
+    }
+
+    pub async fn close_with_error(&mut self, mut msg: String) {
+        if !self.response_bytes.is_closed() {
+            if let Err(e) = self.async_flush().await {
+                use std::fmt::Write;
+                write!(&mut msg, "Unable to flush data: {e}").unwrap();
+            }
+            if let Err(e) = self.response_bytes.send(msg.into()).await {
+                log::error!("Unable to send error back to client: {e}");
+            }
+        }
+    }
+
+    pub async fn async_flush(&mut self) -> std::io::Result<()> {
+        if self.buffer.is_empty() {
+            return Ok(());
+        }
+        log::trace!(
+            "Flushing data to client: {}",
+            String::from_utf8_lossy(&self.buffer)
+        );
+        let sender = self
+            .response_bytes
+            .reserve()
+            .await
+            .map_err(|_| std::io::ErrorKind::WouldBlock)?;
+        sender.send(std::mem::take(&mut self.buffer).into());
+        Ok(())
+    }
+}
+
+impl Write for ResponseWriter {
+    #[inline]
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.buffer.extend_from_slice(buf);
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        if self.buffer.is_empty() {
+            return Ok(());
+        }
+        log::trace!(
+            "Flushing data to client: {}",
+            String::from_utf8_lossy(&self.buffer)
+        );
+        self.response_bytes
+            .try_send(mem::take(&mut self.buffer).into())
+            .map_err(|e|
+                std::io::Error::new(
+                    std::io::ErrorKind::WouldBlock,
+                    format!("{e}: Row limit exceeded. The server cannot store more than {} pending messages in memory. Try again later or increase max_pending_rows in the configuration.", self.response_bytes.max_capacity())
+                )
+            )
+    }
+}
+
+#[allow(clippy::module_name_repetitions)]
+pub struct AsyncResponseWriter {
+    poll_sender: tokio_util::sync::PollSender<Bytes>,
+    writer: ResponseWriter,
+}
+
+impl AsyncResponseWriter {
+    #[must_use]
+    pub fn new(writer: ResponseWriter) -> Self {
+        let sender = writer.response_bytes.clone();
+        Self {
+            poll_sender: tokio_util::sync::PollSender::new(sender),
+            writer,
+        }
+    }
+
+    #[must_use]
+    pub fn into_inner(self) -> ResponseWriter {
+        self.writer
+    }
+}
+
+impl tokio::io::AsyncWrite for AsyncResponseWriter {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        _cx: &mut std::task::Context<'_>,
+        buf: &[u8],
+    ) -> std::task::Poll<std::io::Result<usize>> {
+        std::task::Poll::Ready(self.as_mut().writer.write(buf))
+    }
+
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<std::io::Result<()>> {
+        let Self {
+            poll_sender,
+            writer,
+        } = self.get_mut();
+        match poll_sender.poll_reserve(cx) {
+            std::task::Poll::Ready(Ok(())) => {
+                let res = poll_sender.send_item(std::mem::take(&mut writer.buffer).into());
+                std::task::Poll::Ready(res.map_err(|_| std::io::ErrorKind::BrokenPipe.into()))
+            }
+            std::task::Poll::Pending => std::task::Poll::Pending,
+            std::task::Poll::Ready(Err(_e)) => {
+                std::task::Poll::Ready(Err(std::io::ErrorKind::BrokenPipe.into()))
+            }
+        }
+    }
+
+    fn poll_shutdown(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), std::io::Error>> {
+        self.poll_flush(cx)
+    }
+}
+
+impl Drop for ResponseWriter {
+    fn drop(&mut self) {
+        if let Err(e) = std::io::Write::flush(self) {
+            log::error!("Could not flush data to client: {e}");
+        }
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -230,6 +230,7 @@ pub struct AppConfig {`
`230`	`230`	`pub site_prefix: String,`
`231`	`231`
`232`	`232`	`/// Maximum number of messages that can be stored in memory before sending them to the client.`
	`233`	`+ /// This prevents a single request from using up all available memory.`
`233`	`234`	`#[serde(default = "default_max_pending_rows")]`
`234`	`235`	`pub max_pending_rows: usize,`
`235`	`236`
`@@ -470,6 +471,8 @@ fn default_https_acme_directory_url() -> String {`
`470`	`471`	`"https://acme-v02.api.letsencrypt.org/directory".to_string()`
`471`	`472`	`}`
`472`	`473`
	`474`	`+/// If the sending queue exceeds this number of outgoing messages, an error will be thrown`
	`475`	`+/// This prevents a single request from using up all available memory`
`473`	`476`	`fn default_max_pending_rows() -> usize {`
`474`	`477`	`256`
`475`	`478`	`}`