Skip to content

ensure current tracing span is passed into streaming html rewrite #2974

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Nov 4, 2025
+150 −106
Merged

ensure current tracing span is passed into streaming html rewrite #2974

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ derive_builder = "0.20.2"
# Async
tokio = { version = "1.0", features = ["rt-multi-thread", "signal", "macros"] }
tokio-util = { version = "0.7.15", default-features = false, features = ["io"] }
tracing-futures= { version = "0.2.5", features = ["std-future", "futures-03"] }
futures-util = "0.3.5"
async-stream = "0.3.5"
async-compression = { version = "0.4.25", features = ["tokio", "bzip2", "zstd", "gzip"] }
Expand Down
224 changes: 118 additions & 106 deletions src/utils/html.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,8 @@ use lol_html::{element, errors::RewritingError};
use std::sync::Arc;
use tokio::{io::AsyncRead, task::JoinHandle};
use tokio_util::io::ReaderStream;
use tracing::error;
use tracing::{Span, error, instrument};
use tracing_futures::Instrument as _;

const CHANNEL_SIZE: usize = 64;

Expand All @@ -36,6 +37,7 @@ pub(crate) enum RustdocRewritingError {
/// render the `rustdoc/` templates with the `html`.
/// The output is an HTML page which has not yet been UTF-8 validated.
/// In practice, the output should always be valid UTF-8.
#[instrument(skip_all, fields(memory_limit = max_allowed_memory_usage))]
pub(crate) fn rewrite_rustdoc_html_stream<R>(
template_data: Arc<TemplateData>,
mut reader: R,
Expand All @@ -46,123 +48,132 @@ pub(crate) fn rewrite_rustdoc_html_stream<R>(
where
R: AsyncRead + Unpin + Send + 'static,
{
let stream_span = Span::current();

stream!({
let (input_sender, mut input_receiver) =
tokio::sync::mpsc::channel::<Option<Bytes>>(CHANNEL_SIZE);
let (result_sender, mut result_receiver) =
tokio::sync::mpsc::channel::<Bytes>(CHANNEL_SIZE);

let join_handle: JoinHandle<anyhow::Result<_>> = tokio::spawn(async move {
// we're using the rendering threadpool to limit CPU usage on the server, and to
// offload potentially CPU intensive stuff from the tokio runtime.
// Also this lets us limit the threadpool size and through that the CPU usage.
template_data
.render_in_threadpool(move || {
use lol_html::html_content::{ContentType, Element};
use lol_html::{HtmlRewriter, MemorySettings, Settings};
let producer_span = tracing::info_span!("producer_task");

let head_html = Head::new(&data).render().unwrap();
let vendored_html = Vendored.render().unwrap();
let body_html = Body.render().unwrap();
let topbar_html = data.render().unwrap();
let join_handle: JoinHandle<anyhow::Result<_>> = tokio::spawn(
async move {
// we're using the rendering threadpool to limit CPU usage on the server, and to
// offload potentially CPU intensive stuff from the tokio runtime.
// Also this lets us limit the threadpool size and through that the CPU usage.
let render_span = tracing::info_span!("render_task");
template_data
.render_in_threadpool(move || {
use lol_html::html_content::{ContentType, Element};
use lol_html::{HtmlRewriter, MemorySettings, Settings};

// Before: <body> ... rustdoc content ... </body>
// After:
// ```html
// <div id="rustdoc_body_wrapper" class="{{ rustdoc_body_class }}" tabindex="-1">
// ... rustdoc content ...
// </div>
// ```
let body_handler = |rustdoc_body_class: &mut Element| {
// Add the `rustdoc` classes to the html body
let mut tmp;
let klass = if let Some(classes) = rustdoc_body_class.get_attribute("class")
{
tmp = classes;
tmp.push_str(" container-rustdoc");
&tmp
} else {
"container-rustdoc"
};
rustdoc_body_class.set_attribute("class", klass)?;
rustdoc_body_class.set_attribute("id", "rustdoc_body_wrapper")?;
rustdoc_body_class.set_attribute("tabindex", "-1")?;
// Change the `body` to a `div`
rustdoc_body_class.set_tag_name("div")?;
// Prepend the askama content
rustdoc_body_class.prepend(&body_html, ContentType::Html);
// Wrap the transformed body and topbar into a <body> element
rustdoc_body_class
.before(r#"<body class="rustdoc-page">"#, ContentType::Html);
// Insert the topbar outside of the rustdoc div
rustdoc_body_class.before(&topbar_html, ContentType::Html);
// Finalize body with </body>
rustdoc_body_class.after("</body>", ContentType::Html);
let head_html = Head::new(&data).render().unwrap();
let vendored_html = Vendored.render().unwrap();
let body_html = Body.render().unwrap();
let topbar_html = data.render().unwrap();

Ok(())
};
// Before: <body> ... rustdoc content ... </body>
// After:
// ```html
// <div id="rustdoc_body_wrapper" class="{{ rustdoc_body_class }}" tabindex="-1">
// ... rustdoc content ...
// </div>
// ```
let body_handler = |rustdoc_body_class: &mut Element| {
// Add the `rustdoc` classes to the html body
let mut tmp;
let klass =
if let Some(classes) = rustdoc_body_class.get_attribute("class") {
tmp = classes;
tmp.push_str(" container-rustdoc");
&tmp
} else {
"container-rustdoc"
};
rustdoc_body_class.set_attribute("class", klass)?;
rustdoc_body_class.set_attribute("id", "rustdoc_body_wrapper")?;
rustdoc_body_class.set_attribute("tabindex", "-1")?;
// Change the `body` to a `div`
rustdoc_body_class.set_tag_name("div")?;
// Prepend the askama content
rustdoc_body_class.prepend(&body_html, ContentType::Html);
// Wrap the transformed body and topbar into a <body> element
rustdoc_body_class
.before(r#"<body class="rustdoc-page">"#, ContentType::Html);
// Insert the topbar outside of the rustdoc div
rustdoc_body_class.before(&topbar_html, ContentType::Html);
// Finalize body with </body>
rustdoc_body_class.after("</body>", ContentType::Html);

let settings = Settings {
element_content_handlers: vec![
// Append `style.css` stylesheet after all head elements.
element!("head", |head: &mut Element| {
head.append(&head_html, ContentType::Html);
Ok(())
}),
element!("body", body_handler),
// Append `vendored.css` before `rustdoc.css`, so that the duplicate copy of
// `normalize.css` will be overridden by the later version.
//
// Later rustdoc has `#mainThemeStyle` that could be used, but pre-2018 docs
// don't have this:
//
// https://github.com/rust-lang/rust/commit/003b2bc1c65251ec2fc80b78ed91c43fb35402ec
//
// Pre-2018 rustdoc also didn't have the resource suffix, but docs.rs was using a fork
// that had implemented it already then, so we can assume the css files are
// `<some path>/rustdoc-<some suffix>.css` and use the `-` to distinguish from the
// `rustdoc.static` path.
element!(
"link[rel='stylesheet'][href*='rustdoc-']",
move |rustdoc_css: &mut Element| {
rustdoc_css.before(&vendored_html, ContentType::Html);
Ok(())
};

let settings = Settings {
element_content_handlers: vec![
// Append `style.css` stylesheet after all head elements.
element!("head", |head: &mut Element| {
head.append(&head_html, ContentType::Html);
Ok(())
}
),
],
memory_settings: MemorySettings {
max_allowed_memory_usage,
..MemorySettings::default()
},
..Settings::default()
};
}),
element!("body", body_handler),
// Append `vendored.css` before `rustdoc.css`, so that the duplicate copy of
// `normalize.css` will be overridden by the later version.
//
// Later rustdoc has `#mainThemeStyle` that could be used, but pre-2018 docs
// don't have this:
//
// https://github.com/rust-lang/rust/commit/003b2bc1c65251ec2fc80b78ed91c43fb35402ec
//
// Pre-2018 rustdoc also didn't have the resource suffix, but docs.rs was using a fork
// that had implemented it already then, so we can assume the css files are
// `<some path>/rustdoc-<some suffix>.css` and use the `-` to distinguish from the
// `rustdoc.static` path.
element!(
"link[rel='stylesheet'][href*='rustdoc-']",
move |rustdoc_css: &mut Element| {
rustdoc_css.before(&vendored_html, ContentType::Html);
Ok(())
}
),
],
memory_settings: MemorySettings {
max_allowed_memory_usage,
..MemorySettings::default()
},
..Settings::default()
};

let mut rewriter = HtmlRewriter::new(settings, move |chunk: &[u8]| {
// send the result back to the main rewriter when its coming in.
// this can fail only when the receiver is dropped, in which case
// we exit this thread anyways.
let _ = result_sender.blocking_send(Bytes::copy_from_slice(chunk));
});
while let Some(chunk) = input_receiver
.blocking_recv()
.ok_or_else(|| anyhow!("couldn't receive from input_receiver"))?
{
// receive data from the input receiver.
// `input_receiver` is a non-async one.
// Since we're in a normal background thread, we can use the blocking `.recv`
// here.
// We will get `None` when the reader is done reading,
// so that's our signal to exit this loop and call `rewriter.end()` below.
rewriter.write(&chunk)?;
}
// finalize everything. Will trigger the output sink (and through that,
// sending data to the `result_sender`).
rewriter.end()?;
Ok(())
})
.await?;
Ok(())
});
let mut rewriter = HtmlRewriter::new(settings, move |chunk: &[u8]| {
// send the result back to the main rewriter when its coming in.
// this can fail only when the receiver is dropped, in which case
// we exit this thread anyways.
let _ = result_sender.blocking_send(Bytes::copy_from_slice(chunk));
});
while let Some(chunk) = input_receiver
.blocking_recv()
.ok_or_else(|| anyhow!("couldn't receive from input_receiver"))?
{
// receive data from the input receiver.
// `input_receiver` is a non-async one.
// Since we're in a normal background thread, we can use the blocking `.recv`
// here.
// We will get `None` when the reader is done reading,
// so that's our signal to exit this loop and call `rewriter.end()` below.
rewriter.write(&chunk)?;
}
// finalize everything. Will trigger the output sink (and through that,
// sending data to the `result_sender`).
rewriter.end()?;
Ok(())
})
.instrument(render_span)
.await?;
Ok(())
}
.instrument(producer_span),
);

let mut reader_stream = ReaderStream::new(&mut reader);
while let Some(chunk) = reader_stream.next().await {
Expand Down Expand Up @@ -221,6 +232,7 @@ where
}
})?;
})
.instrument(stream_span)
}

#[cfg(test)]
Expand Down
2 changes: 2 additions & 0 deletions src/web/page/templates.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,9 +69,11 @@ impl TemplateData {
F: FnOnce() -> Result<R> + Send + 'static,
R: Send + 'static,
{
let span = tracing::Span::current();
let (send, recv) = tokio::sync::oneshot::channel();
self.rendering_threadpool.spawn({
move || {
let _guard = span.enter();
// the job may have been queued on the thread-pool for a while,
// if the request was closed in the meantime the receiver should have
// dropped and we don't need to bother rendering the template
Expand Down
Loading