Skip to main content

nxd_plugin_protocol/
progress.rs

1//! Product-neutral progress mechanics shared by linked and external providers.
2//!
3//! These helpers own no lifecycle, vendor, retry, cleanup, or transport policy.
4
5use crate::{SensitiveTextRedactor, truncate_progress_message, v1};
6use std::sync::{
7	Arc, Mutex,
8	atomic::{AtomicU32, Ordering},
9};
10use std::time::{Duration, Instant};
11
12/// Emit progress chunks often enough for long Nix builds without flooding the
13/// gRPC stream (and without retaining every line as a separate event).
14pub const DEFAULT_PROGRESS_CHUNK_BYTES: usize = 32 * 1024;
15const TIMED_FLUSH_INTERVAL: Duration = Duration::from_secs(2);
16/// Unlimited timed flushes so multi-minute builds keep streaming.
17
18#[derive(Clone)]
19pub struct ProgressReporter {
20	milestone: Arc<dyn Fn(String) + Send + Sync>,
21	detail: Arc<dyn Fn(String) + Send + Sync>,
22}
23
24impl ProgressReporter {
25	pub fn new(
26		milestone: impl Fn(String) + Send + Sync + 'static,
27		detail: impl Fn(String) + Send + Sync + 'static,
28	) -> Self {
29		Self { milestone: Arc::new(milestone), detail: Arc::new(detail) }
30	}
31
32	pub fn report(&self, message: impl Into<String>) {
33		(self.milestone)(message.into());
34	}
35
36	pub fn detail(&self, message: impl Into<String>) {
37		(self.detail)(message.into());
38	}
39}
40
41#[derive(Clone)]
42pub struct EventSequencer {
43	action_id: Arc<str>,
44	next: Arc<AtomicU32>,
45	redactor: Arc<Mutex<SensitiveTextRedactor>>,
46}
47
48impl EventSequencer {
49	pub fn new(action_id: impl Into<Arc<str>>) -> Self {
50		Self::with_sensitive_values(action_id, std::iter::empty())
51	}
52
53	pub fn with_sensitive_values<'a>(
54		action_id: impl Into<Arc<str>>,
55		values: impl IntoIterator<Item = &'a [u8]>,
56	) -> Self {
57		Self {
58			action_id: action_id.into(),
59			next: Arc::new(AtomicU32::new(1)),
60			redactor: Arc::new(Mutex::new(SensitiveTextRedactor::with_sensitive_values(values))),
61		}
62	}
63
64	pub fn progress(&self, message: &str) -> v1::ApplyResponse {
65		self.event(v1::EventPhase::Progress, message, false, None)
66	}
67
68	pub fn detail(&self, message: &str) -> v1::ApplyResponse {
69		self.event(v1::EventPhase::Progress, message, true, None)
70	}
71
72	pub fn completed(&self, message: &str) -> v1::ApplyResponse {
73		self.event(v1::EventPhase::Completed, message, false, None)
74	}
75
76	pub fn failed(&self, message: &str, error: v1::ProviderError) -> v1::ApplyResponse {
77		self.event(v1::EventPhase::Failed, message, false, Some(error))
78	}
79
80	fn event(
81		&self,
82		phase: v1::EventPhase,
83		message: &str,
84		detail: bool,
85		error: Option<v1::ProviderError>,
86	) -> v1::ApplyResponse {
87		let redacted = self.redactor.lock().expect("provider event redactor lock").redact(message);
88		let safe_message = truncate_progress_message(&redacted);
89		v1::ApplyResponse {
90			sequence: self.next.fetch_add(1, Ordering::Relaxed),
91			phase: phase.into(),
92			action_id: self.action_id.to_string(),
93			safe_message,
94			error,
95			detail,
96			confidential_outputs: Default::default(),
97			public_outputs: Default::default(),
98		}
99	}
100}
101
102#[derive(Clone)]
103pub struct ProgressBatcher {
104	label: Arc<str>,
105	reporter: ProgressReporter,
106	state: Arc<Mutex<BatchState>>,
107	chunk_bytes: usize,
108}
109
110struct BatchState {
111	text: String,
112	last_emit: Instant,
113	label_emitted: bool,
114	redactor: SensitiveTextRedactor,
115}
116
117impl Default for BatchState {
118	fn default() -> Self {
119		Self {
120			text: String::new(),
121			last_emit: Instant::now(),
122			label_emitted: false,
123			redactor: SensitiveTextRedactor::default(),
124		}
125	}
126}
127
128impl ProgressBatcher {
129	pub fn new(label: impl Into<Arc<str>>, reporter: ProgressReporter) -> Self {
130		Self {
131			label: label.into(),
132			reporter,
133			state: Arc::new(Mutex::new(BatchState::default())),
134			chunk_bytes: DEFAULT_PROGRESS_CHUNK_BYTES,
135		}
136	}
137
138	pub fn push_line(&self, line: &str) {
139		let mut state = self.state.lock().expect("provider progress batch lock");
140		let safe = state.redactor.redact(line.trim_end_matches(['\r', '\n']));
141		state.text.push_str(&safe);
142		state.text.push('\n');
143		let timed = state.last_emit.elapsed() >= TIMED_FLUSH_INTERVAL;
144		if state.text.len() >= self.chunk_bytes || timed {
145			state.last_emit = Instant::now();
146			self.emit(&mut state);
147		}
148	}
149
150	pub fn finish(&self) {
151		let mut state = self.state.lock().expect("provider progress batch lock");
152		self.emit(&mut state);
153	}
154
155	fn emit(&self, state: &mut BatchState) {
156		let chunk = std::mem::take(&mut state.text);
157		let detail = chunk.trim_end();
158		if !detail.is_empty() {
159			let message = if state.label_emitted {
160				detail.to_string()
161			} else {
162				state.label_emitted = true;
163				format!("{}:\n{detail}", self.label)
164			};
165			// Clamp before gRPC so DATA frames stay bounded.
166			self.reporter.detail(truncate_progress_message(&message));
167		}
168	}
169}
170
171#[cfg(test)]
172mod tests {
173	use super::*;
174
175	#[test]
176	fn sequencer_orders_and_sanitizes_events() {
177		let events =
178			EventSequencer::with_sensitive_values("action/test", [b"opaque-provider-value".as_slice()]);
179		let progress = events.progress("token=not-safe opaque-provider-value");
180		let detail = events.detail("ordinary command detail");
181		let completed = events.completed("ordinary completion");
182		assert_eq!(progress.sequence, 1);
183		assert_eq!(progress.safe_message, "token=[REDACTED] [REDACTED]");
184		assert!(detail.detail);
185		assert!(!progress.detail);
186		assert_eq!(completed.sequence, 3);
187		assert_eq!(completed.action_id, "action/test");
188	}
189
190	#[test]
191	fn batcher_preserves_full_normal_output_and_redacts_values() {
192		let messages = Arc::new(Mutex::new(Vec::new()));
193		let captured = messages.clone();
194		let batcher = ProgressBatcher::new(
195			"Nix copy",
196			ProgressReporter::new(
197				|_| {},
198				move |message| {
199					captured.lock().unwrap().push(message);
200				},
201			),
202		);
203		for index in 0..100 {
204			batcher.push_line(&format!("copying /nix/store/path-{index}"));
205		}
206		batcher.push_line("copying /nix/store/final from https://user:[email protected]");
207		batcher.push_line("token=not-safe");
208		batcher.finish();
209		let joined = messages.lock().unwrap().join("\n");
210		assert!(joined.contains("copying /nix/store/path-0"));
211		assert!(joined.contains("copying /nix/store/path-99"));
212		assert!(joined.contains("copying /nix/store/final"));
213		assert!(joined.contains("https://[REDACTED]@example.test"));
214		assert!(joined.contains("token=[REDACTED]"));
215		assert!(!joined.contains("user:pass"));
216		assert!(!joined.contains("not-safe"));
217	}
218
219	#[test]
220	fn batcher_labels_only_the_first_bounded_chunk() {
221		let messages = Arc::new(Mutex::new(Vec::new()));
222		let captured = messages.clone();
223		let mut batcher = ProgressBatcher::new(
224			"Nix activation",
225			ProgressReporter::new(|_| {}, move |message| captured.lock().unwrap().push(message)),
226		);
227		batcher.chunk_bytes = 8;
228		for line in ["first", "second", "third"] {
229			batcher.push_line(line);
230		}
231		batcher.finish();
232		let joined = messages.lock().unwrap().join("\n");
233		assert_eq!(joined.matches("Nix activation:").count(), 1);
234		for line in ["first", "second", "third"] {
235			assert!(joined.contains(line));
236		}
237	}
238}