offload CPU intensive rendering into rayon threadpool

for now only:
* source pages
* rustdoc pages

Author: syphar

Cargo.lock

@@ -16,7 +16,7 @@ exclude = [
 ]
 
 [features]
-consistency_check = ["crates-index", "rayon", "itertools"]
+consistency_check = ["crates-index", "itertools"]
 
 [dependencies]
 sentry = "0.30.0"
@@ -31,7 +31,8 @@ tracing-log = "0.1.3"
 regex = "1"
 clap = { version = "4.0.22", features = [ "derive" ] }
 crates-index = { version = "0.18.5", optional = true }
-rayon = { version = "1", optional = true }
+rayon = "1.6.1"
+num_cpus = "1.15.0"
 crates-index-diff = { version = "16.0.1", features = [ "max-performance" ]}
 reqwest = { version = "0.11", features = ["blocking", "json"] } # TODO: Remove blocking when async is ready
 semver = { version = "1.0.4", features = ["serde"] }

Cargo.toml

@@ -49,6 +49,9 @@ pub struct Config {
     // Time between 'git gc --auto' calls in seconds
     pub(crate) registry_gc_interval: u64,
 
+    /// amout of threads for CPU intensive rendering
+    pub(crate) render_threads: usize,
+
     // random crate search generates a number of random IDs to
     // efficiently find a random crate with > 100 GH stars.
     // The amount depends on the ratio of crates with >100 stars
@@ -153,6 +156,7 @@ impl Config {
             // https://github.com/rust-lang/docs.rs/pull/930#issuecomment-667729380
             max_parse_memory: env("DOCSRS_MAX_PARSE_MEMORY", 5 * 1024 * 1024)?,
             registry_gc_interval: env("DOCSRS_REGISTRY_GC_INTERVAL", 60 * 60)?,
+            render_threads: env("DOCSRS_RENDER_THREADS", num_cpus::get())?,
 
             random_crate_search_view_size: env("DOCSRS_RANDOM_CRATE_SEARCH_VIEW_SIZE", 500)?,
 

src/config.rs

@@ -548,8 +548,13 @@ impl TestFrontend {
         }
 
         debug!("loading template data");
-        let template_data =
-            Arc::new(TemplateData::new(&mut context.pool().unwrap().get().unwrap()).unwrap());
+        let template_data = Arc::new(
+            TemplateData::new(
+                &mut context.pool().unwrap().get().unwrap(),
+                context.config().unwrap().render_threads,
+            )
+            .unwrap(),
+        );
 
         debug!("binding local TCP port for axum");
         let axum_listener =

src/test/mod.rs

@@ -272,7 +272,10 @@ pub(crate) fn build_axum_app(
 
 #[instrument(skip_all)]
 pub fn start_web_server(addr: Option<&str>, context: &dyn Context) -> Result<(), Error> {
-    let template_data = Arc::new(TemplateData::new(&mut *context.pool()?.get()?)?);
+    let template_data = Arc::new(TemplateData::new(
+        &mut *context.pool()?.get()?,
+        context.config()?.render_threads,
+    )?);
 
     let axum_addr: SocketAddr = addr.unwrap_or(DEFAULT_BIND).parse()?;
 

src/web/mod.rs

@@ -7,7 +7,7 @@ use chrono::{DateTime, Utc};
 use path_slash::PathExt;
 use postgres::Client;
 use serde_json::Value;
-use std::{collections::HashMap, fmt, path::PathBuf};
+use std::{collections::HashMap, fmt, path::PathBuf, sync::Arc};
 use tera::{Result as TeraResult, Tera};
 use tracing::{error, trace};
 use walkdir::WalkDir;
@@ -18,20 +18,49 @@ const TEMPLATES_DIRECTORY: &str = "templates";
 #[derive(Debug)]
 pub(crate) struct TemplateData {
     pub templates: Tera,
+    rendering_threadpool: rayon::ThreadPool,
 }
 
 impl TemplateData {
-    pub(crate) fn new(conn: &mut Client) -> Result<Self> {
+    pub(crate) fn new(conn: &mut Client, num_threads: usize) -> Result<Self> {
         trace!("Loading templates");
 
         let data = Self {
             templates: load_templates(conn)?,
+            rendering_threadpool: rayon::ThreadPoolBuilder::new()
+                .num_threads(num_threads)
+                .thread_name(move |idx| format!("docsrs-render {idx}"))
+                .build()?,
         };
 
         trace!("Finished loading templates");
 
         Ok(data)
     }
+
+    /// offload CPU intensive rendering into a rayon threadpool.
+    ///
+    /// This is a thin wrapper around `rayon::spawn` which waits
+    /// sync task to finish.
+    ///
+    /// Use this instead of `spawn_blocking` so we don't block tokio.
+    pub(crate) async fn render_in_threadpool<F, R>(self: &Arc<Self>, render_fn: F) -> Result<R>
+    where
+        F: FnOnce(&TemplateData) -> Result<R> + Send + 'static,
+        R: Send + 'static,
+    {
+        let (send, recv) = tokio::sync::oneshot::channel();
+        self.rendering_threadpool.spawn({
+            let templates = self.clone();
+            move || {
+                // `.send` only fails when the receiver is dropped,
+                // at which point we don't need the result any more.
+                let _ = send.send(render_fn(&templates));
+            }
+        });
+
+        recv.await.context("sender was dropped")?
+    }
 }
 
 fn load_rustc_resource_suffix(conn: &mut Client) -> Result<String> {

src/web/page/templates.rs

@@ -1,8 +1,5 @@
 use super::TemplateData;
-use crate::{
-    utils::spawn_blocking,
-    web::{csp::Csp, error::AxumNope},
-};
+use crate::web::{csp::Csp, error::AxumNope};
 use anyhow::Error;
 use axum::{
     body::{boxed, Body},
@@ -141,11 +138,14 @@ fn render_response(
             context.insert("csp_nonce", &csp_nonce);
 
             let rendered = if cpu_intensive_rendering {
-                spawn_blocking({
-                    let templates = templates.clone();
-                    move || Ok(templates.templates.render(&template, &context)?)
-                })
-                .await
+                templates
+                    .render_in_threadpool(move |templates| {
+                        templates
+                            .templates
+                            .render(&template, &context)
+                            .map_err(Into::into)
+                    })
+                    .await
             } else {
                 templates
                     .templates

src/web/page/web_page.rs

@@ -649,33 +649,33 @@ pub(crate) async fn rustdoc_html_server_handler(
     rendering_time.step("rewrite html");
 
     // Build the page of documentation,
-    // inside `spawn_blocking` since it's CPU intensive.
-    spawn_blocking({
-        let metrics = metrics.clone();
-        move || {
-            Ok(RustdocPage {
-                latest_path,
-                permalink_path,
-                latest_version,
-                target,
-                inner_path,
-                is_latest_version,
-                is_latest_url,
-                is_prerelease,
-                metadata: krate.metadata.clone(),
-                krate: krate.clone(),
+    templates
+        .render_in_threadpool({
+            let metrics = metrics.clone();
+            move |templates| {
+                Ok(RustdocPage {
+                    latest_path,
+                    permalink_path,
+                    latest_version,
+                    target,
+                    inner_path,
+                    is_latest_version,
+                    is_latest_url,
+                    is_prerelease,
+                    metadata: krate.metadata.clone(),
+                    krate: krate.clone(),
+                }
+                .into_response(
+                    &blob.content,
+                    config.max_parse_memory,
+                    templates,
+                    &metrics,
+                    &config,
+                    &storage_path,
+                ))
             }
-            .into_response(
-                &blob.content,
-                config.max_parse_memory,
-                &templates,
-                &metrics,
-                &config,
-                &storage_path,
-            ))
-        }
-    })
-    .await?
+        })
+        .await?
 }
 
 /// Checks whether the given path exists.