Skip to content

Commit 4904cc3

Browse files
sypharsyphar
committed
ensure current tracing span is passed into streaming html rewrite
1 parent e247adb commit 4904cc3

File tree

4 files changed

+145
-102
lines changed

4 files changed

+145
-102
lines changed

‎Cargo.lock‎

Lines changed: 29 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ derive_builder = "0.20.2"
6666
# Async
6767
tokio = { version = "1.0", features = ["rt-multi-thread", "signal", "macros"] }
6868
tokio-util = { version = "0.7.15", default-features = false, features = ["io"] }
69+
tracing-futures= { version = "0.2.5", features = ["std-future", "futures-03"] }
6970
futures-util = "0.3.5"
7071
async-stream = "0.3.5"
7172
async-compression = { version = "0.4.25", features = ["tokio", "bzip2", "zstd", "gzip"] }

‎src/utils/html.rs‎

Lines changed: 113 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ use std::sync::Arc;
1717
use tokio::{io::AsyncRead, task::JoinHandle};
1818
use tokio_util::io::ReaderStream;
1919
use tracing::error;
20+
use tracing_futures::Instrument as _;
2021

2122
#[derive(thiserror::Error, Debug)]
2223
pub(crate) enum RustdocRewritingError {
@@ -42,118 +43,127 @@ pub(crate) fn rewrite_rustdoc_html_stream<R>(
4243
where
4344
R: AsyncRead + Unpin + 'static,
4445
{
46+
let stream_span = tracing::info_span!("rewrite_rustdoc_html_stream");
47+
4548
stream!({
4649
let (input_sender, input_receiver) = std::sync::mpsc::channel::<Option<Vec<u8>>>();
4750
let (result_sender, mut result_receiver) = tokio::sync::mpsc::unbounded_channel::<Bytes>();
4851

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

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

63-
// Before: <body> ... rustdoc content ... </body>
64-
// After:
65-
// ```html
66-
// <div id="rustdoc_body_wrapper" class="{{ rustdoc_body_class }}" tabindex="-1">
67-
// ... rustdoc content ...
68-
// </div>
69-
// ```
70-
let body_handler = |rustdoc_body_class: &mut Element| {
71-
// Add the `rustdoc` classes to the html body
72-
let mut tmp;
73-
let klass = if let Some(classes) = rustdoc_body_class.get_attribute("class")
74-
{
75-
tmp = classes;
76-
tmp.push_str(" container-rustdoc");
77-
&tmp
78-
} else {
79-
"container-rustdoc"
80-
};
81-
rustdoc_body_class.set_attribute("class", klass)?;
82-
rustdoc_body_class.set_attribute("id", "rustdoc_body_wrapper")?;
83-
rustdoc_body_class.set_attribute("tabindex", "-1")?;
84-
// Change the `body` to a `div`
85-
rustdoc_body_class.set_tag_name("div")?;
86-
// Prepend the askama content
87-
rustdoc_body_class.prepend(&body_html, ContentType::Html);
88-
// Wrap the transformed body and topbar into a <body> element
89-
rustdoc_body_class
90-
.before(r#"<body class="rustdoc-page">"#, ContentType::Html);
91-
// Insert the topbar outside of the rustdoc div
92-
rustdoc_body_class.before(&topbar_html, ContentType::Html);
93-
// Finalize body with </body>
94-
rustdoc_body_class.after("</body>", ContentType::Html);
65+
let head_html = Head::new(&data).render().unwrap();
66+
let vendored_html = Vendored.render().unwrap();
67+
let body_html = Body.render().unwrap();
68+
let topbar_html = data.render().unwrap();
9569

96-
Ok(())
97-
};
70+
// Before: <body> ... rustdoc content ... </body>
71+
// After:
72+
// ```html
73+
// <div id="rustdoc_body_wrapper" class="{{ rustdoc_body_class }}" tabindex="-1">
74+
// ... rustdoc content ...
75+
// </div>
76+
// ```
77+
let body_handler = |rustdoc_body_class: &mut Element| {
78+
// Add the `rustdoc` classes to the html body
79+
let mut tmp;
80+
let klass =
81+
if let Some(classes) = rustdoc_body_class.get_attribute("class") {
82+
tmp = classes;
83+
tmp.push_str(" container-rustdoc");
84+
&tmp
85+
} else {
86+
"container-rustdoc"
87+
};
88+
rustdoc_body_class.set_attribute("class", klass)?;
89+
rustdoc_body_class.set_attribute("id", "rustdoc_body_wrapper")?;
90+
rustdoc_body_class.set_attribute("tabindex", "-1")?;
91+
// Change the `body` to a `div`
92+
rustdoc_body_class.set_tag_name("div")?;
93+
// Prepend the askama content
94+
rustdoc_body_class.prepend(&body_html, ContentType::Html);
95+
// Wrap the transformed body and topbar into a <body> element
96+
rustdoc_body_class
97+
.before(r#"<body class="rustdoc-page">"#, ContentType::Html);
98+
// Insert the topbar outside of the rustdoc div
99+
rustdoc_body_class.before(&topbar_html, ContentType::Html);
100+
// Finalize body with </body>
101+
rustdoc_body_class.after("</body>", ContentType::Html);
98102

99-
let settings = Settings {
100-
element_content_handlers: vec![
101-
// Append `style.css` stylesheet after all head elements.
102-
element!("head", |head: &mut Element| {
103-
head.append(&head_html, ContentType::Html);
104-
Ok(())
105-
}),
106-
element!("body", body_handler),
107-
// Append `vendored.css` before `rustdoc.css`, so that the duplicate copy of
108-
// `normalize.css` will be overridden by the later version.
109-
//
110-
// Later rustdoc has `#mainThemeStyle` that could be used, but pre-2018 docs
111-
// don't have this:
112-
//
113-
// https://github.com/rust-lang/rust/commit/003b2bc1c65251ec2fc80b78ed91c43fb35402ec
114-
//
115-
// Pre-2018 rustdoc also didn't have the resource suffix, but docs.rs was using a fork
116-
// that had implemented it already then, so we can assume the css files are
117-
// `<some path>/rustdoc-<some suffix>.css` and use the `-` to distinguish from the
118-
// `rustdoc.static` path.
119-
element!(
120-
"link[rel='stylesheet'][href*='rustdoc-']",
121-
move |rustdoc_css: &mut Element| {
122-
rustdoc_css.before(&vendored_html, ContentType::Html);
103+
Ok(())
104+
};
105+
106+
let settings = Settings {
107+
element_content_handlers: vec![
108+
// Append `style.css` stylesheet after all head elements.
109+
element!("head", |head: &mut Element| {
110+
head.append(&head_html, ContentType::Html);
123111
Ok(())
124-
}
125-
),
126-
],
127-
memory_settings: MemorySettings {
128-
max_allowed_memory_usage,
129-
..MemorySettings::default()
130-
},
131-
..Settings::default()
132-
};
112+
}),
113+
element!("body", body_handler),
114+
// Append `vendored.css` before `rustdoc.css`, so that the duplicate copy of
115+
// `normalize.css` will be overridden by the later version.
116+
//
117+
// Later rustdoc has `#mainThemeStyle` that could be used, but pre-2018 docs
118+
// don't have this:
119+
//
120+
// https://github.com/rust-lang/rust/commit/003b2bc1c65251ec2fc80b78ed91c43fb35402ec
121+
//
122+
// Pre-2018 rustdoc also didn't have the resource suffix, but docs.rs was using a fork
123+
// that had implemented it already then, so we can assume the css files are
124+
// `<some path>/rustdoc-<some suffix>.css` and use the `-` to distinguish from the
125+
// `rustdoc.static` path.
126+
element!(
127+
"link[rel='stylesheet'][href*='rustdoc-']",
128+
move |rustdoc_css: &mut Element| {
129+
rustdoc_css.before(&vendored_html, ContentType::Html);
130+
Ok(())
131+
}
132+
),
133+
],
134+
memory_settings: MemorySettings {
135+
max_allowed_memory_usage,
136+
..MemorySettings::default()
137+
},
138+
..Settings::default()
139+
};
133140

134-
let mut rewriter = HtmlRewriter::new(settings, move |chunk: &[u8]| {
135-
// send the result back to the main rewriter when its coming in.
136-
// this can fail only when the receiver is dropped, in which case
137-
// we exit this thread anyways.
138-
let _ = result_sender.send(Bytes::from(chunk.to_vec()));
139-
});
140-
while let Some(chunk) = input_receiver.recv()? {
141-
// receive data from the input receiver.
142-
// `input_receiver` is a non-async one.
143-
// Since we're in a normal background thread, we can use the blocking `.recv`
144-
// here.
145-
// We will get `None` when the reader is done reading,
146-
// so that's our signal to exit this loop and call `rewriter.end()` below.
147-
rewriter.write(&chunk)?;
148-
}
149-
// finalize everything. Will trigger the output sink (and through that,
150-
// sending data to the `result_sender`).
151-
rewriter.end()?;
152-
Ok(())
153-
})
154-
.await?;
155-
Ok(())
156-
});
141+
let mut rewriter = HtmlRewriter::new(settings, move |chunk: &[u8]| {
142+
// send the result back to the main rewriter when its coming in.
143+
// this can fail only when the receiver is dropped, in which case
144+
// we exit this thread anyways.
145+
let _ = result_sender.send(Bytes::from(chunk.to_vec()));
146+
});
147+
while let Some(chunk) = input_receiver.recv()? {
148+
// receive data from the input receiver.
149+
// `input_receiver` is a non-async one.
150+
// Since we're in a normal background thread, we can use the blocking `.recv`
151+
// here.
152+
// We will get `None` when the reader is done reading,
153+
// so that's our signal to exit this loop and call `rewriter.end()` below.
154+
rewriter.write(&chunk)?;
155+
}
156+
// finalize everything. Will trigger the output sink (and through that,
157+
// sending data to the `result_sender`).
158+
rewriter.end()?;
159+
Ok(())
160+
})
161+
.instrument(render_span)
162+
.await?;
163+
Ok(())
164+
}
165+
.instrument(producer_span),
166+
);
157167

158168
let mut reader_stream = ReaderStream::new(&mut reader);
159169
while let Some(chunk) = reader_stream.next().await {
@@ -208,6 +218,7 @@ where
208218
}
209219
})?;
210220
})
221+
.instrument(stream_span)
211222
}
212223

213224
#[cfg(test)]

‎src/web/page/templates.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,9 +69,11 @@ impl TemplateData {
6969
F: FnOnce() -> Result<R> + Send + 'static,
7070
R: Send + 'static,
7171
{
72+
let span = tracing::Span::current();
7273
let (send, recv) = tokio::sync::oneshot::channel();
7374
self.rendering_threadpool.spawn({
7475
move || {
76+
let _guard = span.enter();
7577
// the job may have been queued on the thread-pool for a while,
7678
// if the request was closed in the meantime the receiver should have
7779
// dropped and we don't need to bother rendering the template

0 commit comments

Comments
 (0)