Log HTML rewriting errors

Author: Kixiron

Cargo.lock

@@ -1,5 +1,6 @@
 use crate::web::page::TemplateData;
 use lol_html::errors::RewritingError;
+use std::path::Path;
 use tera::Context;
 
 /// Rewrite a rustdoc page to have the docs.rs header
@@ -13,6 +14,7 @@ pub(crate) fn rewrite_lol(
     max_allowed_memory_usage: usize,
     ctx: Context,
     templates: &TemplateData,
+    file_path: &Path,
 ) -> Result<Vec<u8>, RewritingError> {
     use lol_html::html_content::{ContentType, Element};
     use lol_html::{ElementContentHandlers, HtmlRewriter, MemorySettings, Settings};
@@ -23,6 +25,7 @@ pub(crate) fn rewrite_lol(
 
     let head_handler = |head: &mut Element| {
         head.append(&tera_head, ContentType::Html);
+
         Ok(())
     };
 
@@ -77,11 +80,33 @@ pub(crate) fn rewrite_lol(
 
     // The input and output are always strings, we just use `&[u8]` so we only have to validate once.
     let mut buffer = Vec::new();
+    // TODO: Make the rewriter persistent?
     let mut writer = HtmlRewriter::try_new(settings, |bytes: &[u8]| {
         buffer.extend_from_slice(bytes);
     })
     .expect("utf8 is a valid encoding");
-    writer.write(html)?;
-    writer.end()?;
+
+    log_memory_errors(writer.write(html), file_path)?;
+    log_memory_errors(writer.end(), file_path)?;
+
     Ok(buffer)
 }
+
+/// Logs OOM errors encountered while rewriting a html file
+fn log_memory_errors<T>(
+    result: Result<T, RewritingError>,
+    file_path: &Path,
+) -> Result<T, RewritingError> {
+    match result {
+        err @ Err(RewritingError::MemoryLimitExceeded(..)) => {
+            crate::web::metrics::HTML_REWRITE_OOMS.inc();
+            log::error!(
+                "Failed to serve the rustdoc file '{}' because rewriting it hit a memory limit",
+                file_path.display(),
+            );
+
+            err
+        }
+        result => result,
+    }
+}

src/utils/html.rs

@@ -147,6 +147,14 @@ pub static CURRENTLY_RUNNING_THREADS: Lazy<IntGauge> = Lazy::new(|| {
     .unwrap()
 });
 
+pub static HTML_REWRITE_OOMS: Lazy<IntGauge> = Lazy::new(|| {
+    register_int_gauge!(
+        "docsrs_html_rewrite_ooms",
+        "The number of attempted files that failed due to a memory limit"
+    )
+    .unwrap()
+});
+
 pub fn metrics_handler(req: &mut Request) -> IronResult<Response> {
     let pool = extension!(req, Pool);
     let queue = extension!(req, BuildQueue);