playwright_cdp/
tracing.rs1use crate::cdp::session::CdpSession;
14use crate::cdp::CdpEvent;
15use crate::error::{Error, Result};
16use crate::options::TracingStartOptions;
17use base64::Engine;
18use parking_lot::Mutex;
19use serde_json::{json, Value};
20use std::sync::Arc;
21use tokio::sync::broadcast;
22
23const SCREENSHOT_CATEGORY: &str = "disabled-by-default-devtools.screenshot";
24
25#[derive(Clone)]
30pub struct Tracing {
31 inner: Arc<TracingInner>,
32}
33
34struct TracingInner {
35 session: CdpSession,
37 path: Mutex<Option<String>>,
39 started: Mutex<bool>,
41}
42
43impl Tracing {
44 pub(crate) fn new(session: CdpSession) -> Self {
45 Self {
46 inner: Arc::new(TracingInner {
47 session,
48 path: Mutex::new(None),
49 started: Mutex::new(false),
50 }),
51 }
52 }
53
54 pub async fn start(&self, options: Option<TracingStartOptions>) -> Result<()> {
60 let opts = options.unwrap_or_default();
61
62 *self.inner.path.lock() = opts.path.clone();
64
65 let mut categories: Vec<String> = opts.categories.unwrap_or_default();
68 if opts.screenshots.unwrap_or(false) && !categories.iter().any(|c| c == SCREENSHOT_CATEGORY) {
69 categories.push(SCREENSHOT_CATEGORY.to_string());
70 }
71
72 let mut params = json!({
73 "transferMode": "ReturnAsStream",
74 "streamFormat": "json",
75 "streamCompression": "none",
76 });
77 if !categories.is_empty() {
78 params["categories"] = json!(categories.join(","));
79 }
80
81 self.inner
82 .session
83 .send("Tracing.start", params)
84 .await
85 .map(|_: Value| ())?;
86 *self.inner.started.lock() = true;
87 Ok(())
88 }
89
90 pub async fn stop(&self) -> Result<Vec<u8>> {
96 {
97 let mut started = self.inner.started.lock();
98 if !*started {
99 return Err(Error::InvalidArgument(
100 "Tracing::stop called without a preceding Tracing::start".into(),
101 ));
102 }
103 *started = false;
104 }
105
106 let mut rx = self.inner.session.subscribe();
109
110 self.inner
111 .session
112 .send("Tracing.end", json!({}))
113 .await
114 .map(|_: Value| ())?;
115
116 let stream = await_tracing_complete(&mut rx).await?;
118
119 let stream_handle = stream
120 .get("stream")
121 .and_then(|v| v.as_str())
122 .ok_or_else(|| {
123 Error::ProtocolError(
124 "Tracing.tracingComplete did not include a stream handle".into(),
125 )
126 })?
127 .to_string();
128
129 let bytes = read_stream(&self.inner.session, &stream_handle).await?;
131
132 if let Some(path) = self.inner.path.lock().clone() {
134 if !path.is_empty() {
135 tokio::fs::write(&path, &bytes).await?;
136 }
137 }
138
139 Ok(bytes)
140 }
141}
142
143async fn await_tracing_complete(rx: &mut broadcast::Receiver<CdpEvent>) -> Result<Value> {
146 let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(60);
147 loop {
148 let recv = tokio::time::timeout_at(deadline, rx.recv()).await;
149 match recv {
150 Ok(Ok(ev)) if ev.method == "Tracing.tracingComplete" => return Ok(ev.params),
151 Ok(Ok(_)) => continue,
152 Ok(Err(broadcast::error::RecvError::Closed)) => {
153 return Err(Error::ChannelClosed);
154 }
155 Ok(Err(broadcast::error::RecvError::Lagged(_))) => continue,
156 Err(_) => {
157 return Err(Error::Timeout(
158 "timed out waiting for Tracing.tracingComplete".into(),
159 ));
160 }
161 }
162 }
163}
164
165async fn read_stream(session: &CdpSession, handle: &str) -> Result<Vec<u8>> {
168 let engine = base64::engine::general_purpose::STANDARD;
169 let mut out = Vec::new();
170 loop {
171 let resp = session
172 .send(
173 "IO.read",
174 json!({ "handle": handle }),
175 )
176 .await?;
177
178 let eof = resp.get("eof").and_then(|v| v.as_bool()).unwrap_or(false);
179 if let Some(data) = resp.get("base64Encoded").and_then(|v| v.as_bool()).unwrap_or(false)
180 .then(|| resp.get("data"))
181 .flatten()
182 .and_then(|v| v.as_str())
183 {
184 let decoded = engine
185 .decode(data)
186 .map_err(|e| Error::ProtocolError(format!("base64 decode failed: {e}")))?;
187 out.extend_from_slice(&decoded);
188 } else if let Some(data) = resp.get("data").and_then(|v| v.as_str()) {
189 out.extend_from_slice(data.as_bytes());
191 }
192
193 if eof {
194 break;
195 }
196 }
197 Ok(out)
198}