1use std::collections::VecDeque;
4use std::path::{Path, PathBuf};
5use std::sync::{Arc, Mutex};
6use std::time::{SystemTime, UNIX_EPOCH};
7
8use serde::Serialize;
9
10use crate::{Pane, Result, Rmux, RmuxError};
11
12const DEFAULT_MAX_TRACE_EVENTS: usize = 100_000;
13
14#[derive(Debug)]
16pub struct RmuxTraceBuilder<'a> {
17 rmux: &'a Rmux,
18 max_events: usize,
19}
20
21impl<'a> RmuxTraceBuilder<'a> {
22 pub(crate) const fn new(rmux: &'a Rmux) -> Self {
23 Self {
24 rmux,
25 max_events: DEFAULT_MAX_TRACE_EVENTS,
26 }
27 }
28
29 pub const fn max_events(mut self, max_events: usize) -> Self {
35 self.max_events = max_events;
36 self
37 }
38
39 pub async fn start(self) -> Result<TraceSession> {
45 let session = TraceSession {
46 endpoint: format!("{:?}", self.rmux.resolved_endpoint()?),
47 max_events: self.max_events,
48 events: Arc::new(Mutex::new(VecDeque::new())),
49 };
50 session.record("trace.start", None::<TracePayload>)?;
51 Ok(session)
52 }
53}
54
55#[derive(Debug, Clone)]
57pub struct TraceSession {
58 endpoint: String,
59 max_events: usize,
60 events: Arc<Mutex<VecDeque<TraceEvent>>>,
61}
62
63impl TraceSession {
64 pub fn record_action(&self, action: impl Into<String>) -> Result<()> {
66 self.record(
67 "action",
68 Some(TracePayload {
69 action: Some(action.into()),
70 ..TracePayload::default()
71 }),
72 )
73 }
74
75 pub fn record_input(&self, pane: &Pane, input: impl Into<String>) -> Result<()> {
77 self.record(
78 "input",
79 Some(TracePayload {
80 pane: Some(format!("{}", pane.target().to_proto())),
81 input: Some(input.into()),
82 ..TracePayload::default()
83 }),
84 )
85 }
86
87 pub async fn record_snapshot(&self, pane: &Pane) -> Result<()> {
89 let snapshot = pane.snapshot().await?;
90 self.record(
91 "snapshot",
92 Some(TracePayload {
93 pane: Some(format!("{}", pane.target().to_proto())),
94 revision: Some(snapshot.revision),
95 snapshot: Some(snapshot.visible_text()),
96 ..TracePayload::default()
97 }),
98 )
99 }
100
101 pub async fn stop(self, directory: impl AsRef<Path>) -> Result<PathBuf> {
103 self.record("trace.stop", None::<TracePayload>)?;
104 let directory = directory.as_ref();
105 tokio::fs::create_dir_all(directory)
106 .await
107 .map_err(trace_io_error)?;
108 let path = directory.join("trace.jsonl");
109 let events = {
110 let events = self.events.lock().map_err(lock_error)?;
111 events.iter().cloned().collect::<Vec<_>>()
112 };
113 let lines = events
114 .iter()
115 .map(serde_json::to_string)
116 .collect::<core::result::Result<Vec<_>, _>>()
117 .map_err(|error| {
118 RmuxError::protocol(rmux_proto::RmuxError::Server(format!(
119 "failed to encode rmux trace event: {error}"
120 )))
121 })?
122 .join("\n");
123 tokio::fs::write(&path, format!("{lines}\n"))
124 .await
125 .map_err(trace_io_error)?;
126 Ok(path)
127 }
128
129 fn record(&self, kind: &'static str, payload: Option<TracePayload>) -> Result<()> {
130 let mut events = self.events.lock().map_err(lock_error)?;
131 if self.max_events == 0 {
132 return Ok(());
133 }
134 if events.len() == self.max_events {
135 events.pop_front();
136 }
137 events.push_back(TraceEvent {
138 timestamp_ms: timestamp_ms(),
139 endpoint: self.endpoint.clone(),
140 kind,
141 payload,
142 });
143 Ok(())
144 }
145}
146
147#[derive(Debug, Clone, Serialize)]
148struct TraceEvent {
149 timestamp_ms: u128,
150 endpoint: String,
151 kind: &'static str,
152 #[serde(skip_serializing_if = "Option::is_none")]
153 payload: Option<TracePayload>,
154}
155
156#[derive(Debug, Clone, Default, Serialize)]
157struct TracePayload {
158 #[serde(skip_serializing_if = "Option::is_none")]
159 action: Option<String>,
160 #[serde(skip_serializing_if = "Option::is_none")]
161 pane: Option<String>,
162 #[serde(skip_serializing_if = "Option::is_none")]
163 input: Option<String>,
164 #[serde(skip_serializing_if = "Option::is_none")]
165 revision: Option<u64>,
166 #[serde(skip_serializing_if = "Option::is_none")]
167 snapshot: Option<String>,
168}
169
170fn timestamp_ms() -> u128 {
171 SystemTime::now()
172 .duration_since(UNIX_EPOCH)
173 .map_or(0, |duration| duration.as_millis())
174}
175
176fn trace_io_error(error: std::io::Error) -> RmuxError {
177 RmuxError::transport("write rmux trace", error)
178}
179
180fn lock_error<T>(error: std::sync::PoisonError<T>) -> RmuxError {
181 RmuxError::protocol(rmux_proto::RmuxError::Server(format!(
182 "rmux trace lock poisoned: {error}"
183 )))
184}