1use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
4use std::sync::{mpsc, Arc, Mutex};
5
6use serde_json::Value;
7
8use crate::client::Client;
9use crate::error::Error;
10use crate::evidence::Lease as EvidenceProviderLease;
11use crate::framing::encode_frame;
12use crate::limits::Limits;
13use crate::marker::encode_marker;
14use crate::messages::{RevisionCommit, SnapshotMessage};
15use crate::tree::Snapshot;
16use crate::validate::validate_snapshot;
17
18struct Publication {
19 revision: i64,
20 tree_frame: Option<Vec<u8>>,
21 commit_frame: Vec<u8>,
22}
23
24pub struct PublicationQueue {
30 sender: Option<mpsc::SyncSender<Publication>>,
31 token: String,
32 session_id: String,
33 limits: Limits,
34 subscribe: String,
35 marker_enabled: bool,
36 revision: i64,
37 evidence_lease: Option<EvidenceProviderLease>,
38 failed: Arc<AtomicBool>,
39 dropped: AtomicU64,
40 fatal: Arc<Mutex<Option<(String, String)>>>,
41 worker: Option<std::thread::JoinHandle<()>>,
42 #[cfg(test)]
43 done: Mutex<mpsc::Receiver<()>>,
44}
45
46impl PublicationQueue {
47 pub fn new(client: Client, capacity: usize) -> Result<Self, Error> {
52 Self::new_inner(client, capacity, None)
53 }
54
55 fn new_inner(
56 mut client: Client,
57 capacity: usize,
58 #[cfg(test)] gate: Option<Arc<TestWorkerGate>>,
59 #[cfg(not(test))] _gate: Option<()>,
60 ) -> Result<Self, Error> {
61 if capacity == 0 || !client.connected() {
62 return Err(Error::PublicationWorkerFailed);
63 }
64 let (token, session_id, limits, subscribe, marker_enabled, revision) = client
65 .publication_config()
66 .ok_or(Error::PublicationWorkerFailed)?;
67 let evidence_lease = client.take_evidence_lease();
68 let (sender, receiver) = mpsc::sync_channel::<Publication>(capacity);
69 #[cfg(test)]
70 let (done_sender, done_receiver) = mpsc::sync_channel(1);
71 let failed = Arc::new(AtomicBool::new(false));
72 let worker_failed = failed.clone();
73 let fatal: Arc<Mutex<Option<(String, String)>>> = Arc::new(Mutex::new(None));
74 let worker_fatal = fatal.clone();
75 let worker = std::thread::Builder::new()
76 .name("termwright-semantic-publication".into())
77 .spawn(move || {
78 while let Ok(publication) = receiver.recv() {
79 #[cfg(test)]
80 if let Some(gate) = gate.as_ref() {
81 gate.enter();
82 }
83 let result = (|| {
84 if let Some(frame) = publication.tree_frame.as_ref() {
85 client.write_frame(frame)?;
86 }
87 client.write_frame(&publication.commit_frame)?;
88 client.accept_queued_publication(
89 publication.revision,
90 publication.tree_frame.is_some(),
91 );
92 Ok::<(), Error>(())
93 })();
94 if result.is_err() {
95 worker_failed.store(true, Ordering::Release);
96 client.close();
97 #[cfg(test)]
98 let _ = done_sender.send(());
99 return;
100 }
101 }
102 let fatal = worker_fatal.lock().ok().and_then(|mut value| value.take());
103 if let Some((code, message)) = fatal {
104 let _ = client.fail(&code, message);
105 } else {
106 client.close();
107 }
108 #[cfg(test)]
109 let _ = done_sender.send(());
110 })
111 .map_err(Error::Io)?;
112 Ok(Self {
113 sender: Some(sender),
114 token,
115 session_id,
116 limits,
117 subscribe,
118 marker_enabled,
119 revision,
120 evidence_lease,
121 failed,
122 dropped: AtomicU64::new(0),
123 fatal,
124 worker: Some(worker),
125 #[cfg(test)]
126 done: Mutex::new(done_receiver),
127 })
128 }
129
130 pub fn limits(&self) -> &Limits {
132 &self.limits
133 }
134
135 pub fn publish(&mut self, snapshot: &mut Snapshot) -> Result<Option<String>, Error> {
137 if self.failed.load(Ordering::Acquire) {
138 return Err(Error::PublicationWorkerFailed);
139 }
140 let revision = self.revision + 1;
141 snapshot.v = 2;
142 snapshot.session_id = self.session_id.clone();
143 snapshot.revision = revision;
144 if let Some(lease) = self.evidence_lease.as_ref() {
145 snapshot.provider_evidence =
146 lease.collect(&self.session_id, revision, snapshot.columns, snapshot.rows);
147 }
148 let body = serde_json::to_string(&snapshot).map_err(|_| {
149 Error::Protocol(crate::error::Violation::new(
150 "frame-malformed",
151 "snapshot is not JSON-serialisable",
152 ))
153 })?;
154 let parsed: Value = serde_json::from_str(&body).expect("just serialised");
155 validate_snapshot(&parsed, &self.limits)?;
156 let marker = if self.marker_enabled {
157 Some(encode_marker(&self.token, &self.session_id, revision)?)
158 } else {
159 None
160 };
161 let tree_frame = if self.subscribe != "revisions" {
162 Some(encode_frame(
163 &SnapshotMessage::new(snapshot),
164 self.limits.max_frame_bytes,
165 )?)
166 } else {
167 None
168 };
169 let commit_frame =
170 encode_frame(&RevisionCommit::new(revision), self.limits.max_frame_bytes)?;
171 let publication = Publication {
172 revision,
173 tree_frame,
174 commit_frame,
175 };
176 let Some(sender) = self.sender.as_ref() else {
177 self.failed.store(true, Ordering::Release);
178 return Err(Error::PublicationWorkerFailed);
179 };
180 match sender.try_send(publication) {
181 Ok(()) => {
182 self.revision = revision;
183 Ok(marker)
184 }
185 Err(mpsc::TrySendError::Full(_)) => {
186 self.dropped.fetch_add(1, Ordering::Relaxed);
187 Err(Error::PublicationQueueFull)
188 }
189 Err(mpsc::TrySendError::Disconnected(_)) => {
190 self.failed.store(true, Ordering::Release);
191 Err(Error::PublicationWorkerFailed)
192 }
193 }
194 }
195
196 pub fn dropped(&self) -> u64 {
198 self.dropped.load(Ordering::Relaxed)
199 }
200
201 pub fn failed(&self) -> bool {
203 self.failed.load(Ordering::Acquire)
204 }
205
206 pub fn fail(&mut self, code: impl Into<String>, message: impl Into<String>) {
209 self.failed.store(true, Ordering::Release);
210 if let Ok(mut fatal) = self.fatal.lock() {
211 *fatal = Some((code.into(), message.into()));
212 }
213 self.sender.take();
216 }
217
218 pub fn shutdown(mut self) -> bool {
223 self.sender.take();
224 if let Some(worker) = self.worker.take() {
225 let _ = worker.join();
226 }
227 !self.failed.load(Ordering::Acquire)
228 }
229}
230
231#[cfg(test)]
232struct TestWorkerGate {
233 entered: mpsc::SyncSender<()>,
234 release: Mutex<mpsc::Receiver<()>>,
235}
236
237#[cfg(test)]
238impl TestWorkerGate {
239 fn enter(&self) {
240 let _ = self.entered.send(());
241 let _ = self.release.lock().expect("gate lock").recv();
242 }
243}
244
245#[cfg(all(test, unix))]
246mod tests {
247 use super::*;
248 use std::io::Read;
249 use std::os::unix::net::UnixStream;
250
251 use crate::FrameDecoder;
252 use crate::{Node, Role};
253
254 fn snapshot() -> Snapshot {
255 let mut snapshot = Snapshot::new(80, 24);
256 snapshot.push(Node::new("root", Role::Application, "fixture"));
257 snapshot
258 }
259
260 fn gated_queue(
261 capacity: usize,
262 ) -> (
263 PublicationQueue,
264 UnixStream,
265 mpsc::Receiver<()>,
266 mpsc::SyncSender<()>,
267 ) {
268 let (client_stream, server_stream) = UnixStream::pair().expect("socket pair");
269 let client = Client::test_connected(client_stream);
270 let (entered_tx, entered_rx) = mpsc::sync_channel(1);
271 let (release_tx, release_rx) = mpsc::sync_channel(1);
272 let gate = Arc::new(TestWorkerGate {
273 entered: entered_tx,
274 release: Mutex::new(release_rx),
275 });
276 let queue = PublicationQueue::new_inner(client, capacity, Some(gate)).expect("queue");
277 (queue, server_stream, entered_rx, release_tx)
278 }
279
280 #[test]
281 fn full_queue_drops_without_revision_gap_or_marker() {
282 let (mut queue, mut server, entered, release) = gated_queue(1);
283 assert!(queue.publish(&mut snapshot()).unwrap().is_some());
284 entered.recv().expect("worker entered first job");
285 assert!(queue.publish(&mut snapshot()).unwrap().is_some());
286 assert!(matches!(
287 queue.publish(&mut snapshot()),
288 Err(Error::PublicationQueueFull)
289 ));
290 assert_eq!(queue.revision, 2);
291 assert_eq!(queue.dropped(), 1);
292 release.send(()).unwrap();
293 entered.recv().expect("worker entered second job");
294 release.send(()).unwrap();
295 queue.fail("test-complete", "close");
296 queue.done.lock().unwrap().recv().expect("worker stopped");
297 let mut bytes = Vec::new();
298 let _ = server.read_to_end(&mut bytes);
299 let mut decoder = FrameDecoder::new(
300 crate::DEFAULT_LIMITS.max_frame_bytes,
301 crate::DEFAULT_LIMITS.max_depth,
302 );
303 let frames = decoder.push(&bytes).expect("ordered frames decode");
304 let names: Vec<_> = frames
305 .iter()
306 .take(4)
307 .map(|frame| frame.value["type"].as_str())
308 .collect();
309 assert_eq!(
310 names,
311 [
312 Some("snapshot"),
313 Some("revision-commit"),
314 Some("snapshot"),
315 Some("revision-commit")
316 ]
317 );
318 }
319
320 #[test]
321 fn worker_failure_permanently_refuses_later_markers() {
322 let (mut queue, server, entered, release) = gated_queue(1);
323 assert!(queue.publish(&mut snapshot()).unwrap().is_some());
324 entered.recv().expect("worker entered");
325 drop(server);
326 release.send(()).unwrap();
327 queue.done.lock().unwrap().recv().expect("worker stopped");
330 assert!(matches!(
331 queue.publish(&mut snapshot()),
332 Err(Error::PublicationWorkerFailed)
333 ));
334 }
335
336 #[test]
337 fn shutdown_drains_an_admitted_one_frame_process_before_returning() {
338 let (client_stream, mut server) = UnixStream::pair().expect("socket pair");
339 let client = Client::test_connected(client_stream);
340 let mut queue = PublicationQueue::new(client, 1).expect("queue");
341 assert!(queue.publish(&mut snapshot()).unwrap().is_some());
342
343 assert!(queue.shutdown(), "clean drain reported worker failure");
344
345 let mut bytes = Vec::new();
346 server
347 .read_to_end(&mut bytes)
348 .expect("worker closed socket");
349 let mut decoder = FrameDecoder::new(
350 crate::DEFAULT_LIMITS.max_frame_bytes,
351 crate::DEFAULT_LIMITS.max_depth,
352 );
353 let frames = decoder.push(&bytes).expect("ordered frames decode");
354 assert_eq!(frames.len(), 2);
355 assert_eq!(frames[0].value["type"], "snapshot");
356 assert_eq!(frames[1].value["type"], "revision-commit");
357 }
358}