Skip to content

Commit de318c5

Browse files
sypharsyphar
committed
ensure current tracing span is passed into streaming html rewrite
1 parent 5ab3742 commit de318c5

File tree

4 files changed

+150
-106
lines changed

4 files changed

+150
-106
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: 118 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,8 @@ use lol_html::{element, errors::RewritingError};
1818
use std::sync::Arc;
1919
use tokio::{io::AsyncRead, task::JoinHandle};
2020
use tokio_util::io::ReaderStream;
21-
use tracing::error;
21+
use tracing::{Span, error, instrument};
22+
use tracing_futures::Instrument as _;
2223

2324
const CHANNEL_SIZE: usize = 64;
2425

@@ -36,6 +37,7 @@ pub(crate) enum RustdocRewritingError {
3637
/// render the `rustdoc/` templates with the `html`.
3738
/// The output is an HTML page which has not yet been UTF-8 validated.
3839
/// In practice, the output should always be valid UTF-8.
40+
#[instrument(skip_all, fields(memory_limit = max_allowed_memory_usage))]
3941
pub(crate) fn rewrite_rustdoc_html_stream<R>(
4042
template_data: Arc<TemplateData>,
4143
mut reader: R,
@@ -46,123 +48,132 @@ pub(crate) fn rewrite_rustdoc_html_stream<R>(
4648
where
4749
R: AsyncRead + Unpin + Send + 'static,
4850
{
51+
let stream_span = Span::current();
52+
4953
stream!({
5054
let (input_sender, mut input_receiver) =
5155
tokio::sync::mpsc::channel::<Option<Bytes>>(CHANNEL_SIZE);
5256
let (result_sender, mut result_receiver) =
5357
tokio::sync::mpsc::channel::<Bytes>(CHANNEL_SIZE);
5458

55-
let join_handle: JoinHandle<anyhow::Result<_>> = tokio::spawn(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-
template_data
60-
.render_in_threadpool(move || {
61-
use lol_html::html_content::{ContentType, Element};
62-
use lol_html::{HtmlRewriter, MemorySettings, Settings};
59+
let producer_span = tracing::info_span!("producer_task");
6360

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

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

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

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

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

167178
let mut reader_stream = ReaderStream::new(&mut reader);
168179
while let Some(chunk) = reader_stream.next().await {
@@ -221,6 +232,7 @@ where
221232
}
222233
})?;
223234
})
235+
.instrument(stream_span)
224236
}
225237

226238
#[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)