Skip to content

Commit 1b9711b

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

File tree

4 files changed

+148
-108
lines changed

4 files changed

+148
-108
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: 116 additions & 108 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,8 @@ use lol_html::{element, errors::RewritingError};
1616
use std::sync::Arc;
1717
use tokio::{io::AsyncRead, task::JoinHandle};
1818
use tokio_util::io::ReaderStream;
19-
use tracing::error;
19+
use tracing::{Span, error, instrument};
20+
use tracing_futures::Instrument as _;
2021

2122
#[derive(thiserror::Error, Debug)]
2223
pub(crate) enum RustdocRewritingError {
@@ -32,6 +33,7 @@ pub(crate) enum RustdocRewritingError {
3233
/// render the `rustdoc/` templates with the `html`.
3334
/// The output is an HTML page which has not yet been UTF-8 validated.
3435
/// In practice, the output should always be valid UTF-8.
36+
#[instrument(skip_all, fields(memory_limit = max_allowed_memory_usage))]
3537
pub(crate) fn rewrite_rustdoc_html_stream<R>(
3638
template_data: Arc<TemplateData>,
3739
mut reader: R,
@@ -42,118 +44,127 @@ pub(crate) fn rewrite_rustdoc_html_stream<R>(
4244
where
4345
R: AsyncRead + Unpin + 'static,
4446
{
47+
let stream_span = Span::current();
48+
4549
stream!({
4650
let (input_sender, input_receiver) = std::sync::mpsc::channel::<Option<Vec<u8>>>();
4751
let (result_sender, mut result_receiver) = tokio::sync::mpsc::unbounded_channel::<Bytes>();
4852

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};
53+
let producer_span = tracing::info_span!("producer_task");
5754

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();
55+
let join_handle: JoinHandle<anyhow::Result<_>> = tokio::spawn(
56+
async move {
57+
// we're using the rendering threadpool to limit CPU usage on the server, and to
58+
// offload potentially CPU intensive stuff from the tokio runtime.
59+
// Also this lets us limit the threadpool size and through that the CPU usage.
60+
let render_span = tracing::info_span!("render_task");
61+
template_data
62+
.render_in_threadpool(move || {
63+
use lol_html::html_content::{ContentType, Element};
64+
use lol_html::{HtmlRewriter, MemorySettings, Settings};
6265

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);
66+
let head_html = Head::new(&data).render().unwrap();
67+
let vendored_html = Vendored.render().unwrap();
68+
let body_html = Body.render().unwrap();
69+
let topbar_html = data.render().unwrap();
9570

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

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

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

158169
let mut reader_stream = ReaderStream::new(&mut reader);
159170
while let Some(chunk) = reader_stream.next().await {
@@ -188,11 +199,7 @@ where
188199
}
189200

190201
join_handle.await.expect("Task panicked").map_err(|e| {
191-
error!(
192-
?e,
193-
memory_limit = max_allowed_memory_usage,
194-
"error while rewriting rustdoc HTML"
195-
);
202+
error!(?e, "error while rewriting rustdoc HTML");
196203
// our `render_in_threadpool` and so the async tokio task return an `anyhow::Result`.
197204
// In most cases this will be an error from the `HtmlRewriter`, which we'll get as a
198205
// `RewritingError` which we extract here again. The other cases remain an
@@ -208,6 +215,7 @@ where
208215
}
209216
})?;
210217
})
218+
.instrument(stream_span)
211219
}
212220

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