tawasuyu/brahman

1//! Flow channels: data plane sobre Unix socket por edge enriquecido.
2//!
3//! Cuando un splitter detecta el TypeRef de un edge, además de replicar a
4//! los consumers internos del pipeline, se levanta un FlowChannel que
5//! expone los bytes a subscribers externos (otros módulos del fractal).
6//!
7//! ## Diseño
8//!
9//! - `tokio::sync::broadcast::channel` para fan-out lock-less entre el
10//!   splitter (sender) y los N subscribers conectados.
11//! - `UnixListener` accept-loop: por cada cliente nuevo, spawn una task
12//!   que drena el receiver y escribe al socket.
13//! - Subscribers lentos pueden perder mensajes (broadcast::Receiver::Lagged)
14//!   — se loguea warn y se sigue. Esto es deliberado para no bloquear el
15//!   splitter en consumers lentos.
16//!
17//! ## Lifetime
18//!
19//! `FlowChannel` se construye con `new(path)`. Cuando se drop:
20//! - El `accept_task` se cancela (vía drop del `tokio::task::JoinHandle`
21//!   que tenemos abort-on-drop).
22//! - El socket file se borra del FS (`Drop` impl).
23//!
24//! Sender clones son baratos; los subscribers conectados se enteran del
25//! cierre cuando todos los senders se dropean.
26
27use std::collections::VecDeque;
28use std::path::{Path, PathBuf};
29use std::sync::{Arc, Mutex};
30use tokio::io::AsyncWriteExt;
31use tokio::net::UnixListener;
32use tokio::sync::broadcast;
33use tokio::task::AbortHandle;
34use tracing::{debug, warn};
35
36/// Capacidad del broadcast channel. Si un subscriber está más de N chunks
37/// atrasado, queda `Lagged` y empieza a perder mensajes.
38const BROADCAST_CAP: usize = 64;
39
40/// Chunks default del replay buffer. Cuando un cliente nuevo se conecta,
41/// recibe hasta estos N chunks antes de iniciar el broadcast live.
42/// Override via `FlowChannel::with_replay_cap`.
43pub const DEFAULT_REPLAY_CHUNKS: usize = 32;
44
45pub struct FlowChannel {
46    sender: broadcast::Sender<Arc<Vec<u8>>>,
47    replay: Arc<Mutex<VecDeque<Arc<Vec<u8>>>>>,
48    replay_caps: ReplayCaps,
49    socket_path: PathBuf,
50    _accept_handle: AbortOnDrop,
51}
52
53#[derive(Debug, Clone, Copy)]
54pub struct ReplayCaps {
55    /// Máximo de chunks retenidos.
56    pub chunks: usize,
57    /// Máximo de bytes (sumando len de chunks). `0` = sin tope.
58    pub bytes: usize,
59}
60
61impl ReplayCaps {
62    pub fn chunks_only(chunks: usize) -> Self {
63        Self { chunks: chunks.max(1), bytes: 0 }
64    }
65    pub fn new(chunks: usize, bytes: usize) -> Self {
66        Self { chunks: chunks.max(1), bytes }
67    }
68}
69
70#[derive(Clone)]
71pub struct FlowSender {
72    sender: broadcast::Sender<Arc<Vec<u8>>>,
73    replay: Arc<Mutex<VecDeque<Arc<Vec<u8>>>>>,
74    replay_caps: ReplayCaps,
75}
76
77impl FlowSender {
78    /// Pushea al broadcast y al replay buffer. Si no hay subscribers,
79    /// el broadcast::send retorna Err pero igual guardamos en replay
80    /// (subscribers tarde verán los chunks pasados).
81    pub fn send(&self, data: Arc<Vec<u8>>) {
82        let incoming = data.len();
83        let caps = self.replay_caps;
84        if let Ok(mut g) = self.replay.lock() {
85            evict_for_incoming(&mut g, caps, incoming);
86            g.push_back(data.clone());
87        }
88        let _ = self.sender.send(data);
89    }
90}
91
92/// Evict los chunks más viejos para hacer espacio a un chunk entrante de
93/// `incoming` bytes — el buffer post-push queda dentro de los caps.
94fn evict_for_incoming(buf: &mut VecDeque<Arc<Vec<u8>>>, caps: ReplayCaps, incoming: usize) {
95    // 1) chunks: dejar lugar para 1 más.
96    while buf.len() + 1 > caps.chunks {
97        if buf.pop_front().is_none() {
98            break;
99        }
100    }
101    // 2) bytes (si está activado).
102    if caps.bytes > 0 {
103        let mut current: usize = buf.iter().map(|a| a.len()).sum();
104        while current + incoming > caps.bytes {
105            match buf.pop_front() {
106                Some(c) => current = current.saturating_sub(c.len()),
107                None => break,
108            }
109        }
110    }
111}
112
113impl std::fmt::Debug for FlowChannel {
114    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
115        f.debug_struct("FlowChannel")
116            .field("socket_path", &self.socket_path)
117            .field("subscribers", &self.sender.receiver_count())
118            .finish()
119    }
120}
121
122impl FlowChannel {
123    /// Crea un FlowChannel atado al path `socket_path`. Si el path ya
124    /// existe, lo borra antes de bind (asume restart limpio).
125    pub fn new(socket_path: PathBuf) -> std::io::Result<Self> {
126        Self::with_replay_caps(socket_path, ReplayCaps::chunks_only(DEFAULT_REPLAY_CHUNKS))
127    }
128
129    pub fn with_replay_cap(socket_path: PathBuf, chunks: usize) -> std::io::Result<Self> {
130        Self::with_replay_caps(socket_path, ReplayCaps::chunks_only(chunks))
131    }
132
133    pub fn with_replay_caps(socket_path: PathBuf, caps: ReplayCaps) -> std::io::Result<Self> {
134        if socket_path.exists() {
135            let _ = std::fs::remove_file(&socket_path);
136        }
137        if let Some(parent) = socket_path.parent() {
138            let _ = std::fs::create_dir_all(parent);
139        }
140        let listener = UnixListener::bind(&socket_path)?;
141        let (tx, _rx_unused) = broadcast::channel::<Arc<Vec<u8>>>(BROADCAST_CAP);
142        let replay: Arc<Mutex<VecDeque<Arc<Vec<u8>>>>> =
143            Arc::new(Mutex::new(VecDeque::with_capacity(caps.chunks)));
144        let tx_for_accept = tx.clone();
145        let replay_for_accept = replay.clone();
146        let path_for_log = socket_path.clone();
147
148        let join = tokio::spawn(async move {
149            debug!(path = %path_for_log.display(), "flow channel listening");
150            loop {
151                let (mut stream, _addr) = match listener.accept().await {
152                    Ok(p) => p,
153                    Err(e) => {
154                        warn!(?e, "flow channel accept failed");
155                        return;
156                    }
157                };
158                // Snapshot del replay buffer Y subscribe al broadcast.
159                // El orden es crítico: subscribe ANTES de drenar el replay
160                // para no perder chunks que llegan justo en el medio.
161                let mut rx = tx_for_accept.subscribe();
162                let snapshot: Vec<Arc<Vec<u8>>> = {
163                    let g = replay_for_accept.lock().expect("replay lock");
164                    g.iter().cloned().collect()
165                };
166                tokio::spawn(async move {
167                    // Fase 1: drenar replay snapshot al subscriber.
168                    for chunk in &snapshot {
169                        if let Err(e) = stream.write_all(chunk).await {
170                            debug!(?e, "flow subscriber dropped during replay");
171                            return;
172                        }
173                    }
174                    // Fase 2: live broadcast.
175                    loop {
176                        match rx.recv().await {
177                            Ok(chunk) => {
178                                if let Err(e) = stream.write_all(&chunk).await {
179                                    debug!(?e, "flow subscriber dropped");
180                                    return;
181                                }
182                            }
183                            Err(broadcast::error::RecvError::Closed) => return,
184                            Err(broadcast::error::RecvError::Lagged(n)) => {
185                                warn!(skipped = n, "flow subscriber lagged");
186                            }
187                        }
188                    }
189                });
190            }
191        });
192
193        Ok(Self {
194            sender: tx,
195            replay,
196            replay_caps: caps,
197            socket_path,
198            _accept_handle: AbortOnDrop(join.abort_handle()),
199        })
200    }
201
202    /// Push un chunk al channel. Si no hay subscribers, drop silencioso.
203    /// Siempre se guarda en el replay buffer (con cap rotation por chunks
204    /// y opcionalmente por bytes).
205    pub fn send(&self, data: Vec<u8>) {
206        let incoming = data.len();
207        let arc = Arc::new(data);
208        let caps = self.replay_caps;
209        if let Ok(mut g) = self.replay.lock() {
210            evict_for_incoming(&mut g, caps, incoming);
211            g.push_back(arc.clone());
212        }
213        let _ = self.sender.send(arc);
214    }
215
216    pub fn socket_path(&self) -> &Path {
217        &self.socket_path
218    }
219
220    /// Handle clone-able para que tasks externas (splitter) pushen al
221    /// channel sin tener ownership del FlowChannel. Cada push se guarda
222    /// también en el replay buffer.
223    pub fn sender_handle(&self) -> FlowSender {
224        FlowSender {
225            sender: self.sender.clone(),
226            replay: self.replay.clone(),
227            replay_caps: self.replay_caps,
228        }
229    }
230
231    pub fn subscriber_count(&self) -> usize {
232        self.sender.receiver_count()
233    }
234}
235
236impl Drop for FlowChannel {
237    fn drop(&mut self) {
238        // El AbortOnDrop cancela el accept loop; sólo nos queda limpiar el
239        // socket file.
240        let _ = std::fs::remove_file(&self.socket_path);
241    }
242}
243
244struct AbortOnDrop(AbortHandle);
245impl Drop for AbortOnDrop {
246    fn drop(&mut self) {
247        self.0.abort();
248    }
249}
250
251/// Path canónico para un flow channel: `$XDG_RUNTIME_DIR/shipote-flow-<id>.sock`.
252pub fn default_flow_socket_path(id: &str) -> PathBuf {
253    let base = std::env::var("XDG_RUNTIME_DIR").unwrap_or_else(|_| {
254        let uid = nix::unistd::getuid().as_raw();
255        let p = format!("/run/user/{uid}");
256        if std::path::Path::new(&p).exists() {
257            p
258        } else {
259            "/tmp".into()
260        }
261    });
262    PathBuf::from(base).join(format!("shipote-flow-{id}.sock"))
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268    use tokio::io::AsyncReadExt;
269    use tokio::net::UnixStream;
270
271    #[tokio::test]
272    async fn channel_delivers_to_subscriber() {
273        let tmp = tempfile::tempdir().unwrap();
274        let path = tmp.path().join("flow.sock");
275        let ch = FlowChannel::new(path.clone()).unwrap();
276
277        // Subscriber se conecta.
278        let path_clone = path.clone();
279        let task = tokio::spawn(async move {
280            let mut stream = UnixStream::connect(&path_clone).await.unwrap();
281            let mut buf = vec![0u8; 64];
282            let n = stream.read(&mut buf).await.unwrap();
283            buf.truncate(n);
284            buf
285        });
286
287        // Damos tiempo al accept.
288        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
289        // Hasta que haya 1 receiver_count, el send no llega.
290        for _ in 0..50 {
291            if ch.subscriber_count() >= 1 {
292                break;
293            }
294            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
295        }
296        ch.send(b"hello-flow".to_vec());
297
298        let received = tokio::time::timeout(std::time::Duration::from_secs(2), task)
299            .await
300            .expect("timeout")
301            .unwrap();
302        assert_eq!(received, b"hello-flow");
303    }
304
305    #[tokio::test]
306    async fn replay_buffer_serves_late_subscriber() {
307        let tmp = tempfile::tempdir().unwrap();
308        let path = tmp.path().join("flow.sock");
309        let ch = FlowChannel::new(path.clone()).unwrap();
310
311        // Pushes ANTES de cualquier subscriber: van solo al replay.
312        ch.send(b"chunk-1".to_vec());
313        ch.send(b"chunk-2".to_vec());
314        ch.send(b"chunk-3".to_vec());
315
316        // Subscriber LATE — debe recibir los 3 chunks del replay.
317        let path_clone = path.clone();
318        let task = tokio::spawn(async move {
319            let mut stream = UnixStream::connect(&path_clone).await.unwrap();
320            let mut buf = vec![0u8; 256];
321            // Leemos hasta recibir los 3 chunks (21 bytes esperados).
322            let mut total = Vec::new();
323            for _ in 0..20 {
324                let n = stream.read(&mut buf).await.unwrap();
325                if n == 0 {
326                    break;
327                }
328                total.extend_from_slice(&buf[..n]);
329                if total.len() >= 21 {
330                    break;
331                }
332            }
333            total
334        });
335
336        let received = tokio::time::timeout(std::time::Duration::from_secs(2), task)
337            .await
338            .expect("timeout")
339            .unwrap();
340        let s = String::from_utf8_lossy(&received);
341        assert!(s.contains("chunk-1"), "got: {s:?}");
342        assert!(s.contains("chunk-2"), "got: {s:?}");
343        assert!(s.contains("chunk-3"), "got: {s:?}");
344    }
345
346    #[tokio::test]
347    async fn replay_evicts_by_bytes_cap() {
348        let tmp = tempfile::tempdir().unwrap();
349        let path = tmp.path().join("flow.sock");
350        // chunks=100 (no limita), bytes=20: deberíamos retener sólo los
351        // últimos chunks cuyos bytes sumen ≤ 20.
352        let ch = FlowChannel::with_replay_caps(path.clone(), ReplayCaps::new(100, 20)).unwrap();
353        ch.send(b"AAAAAAAA".to_vec()); // 8 bytes
354        ch.send(b"BBBBBBBB".to_vec()); // 8 → total 16
355        ch.send(b"CCCCCCCC".to_vec()); // 8 → total 24 > 20, evict A → 16
356        ch.send(b"DDDDDDDD".to_vec()); // 8 → total 24 > 20, evict B → 16
357
358        let path_clone = path.clone();
359        let task = tokio::spawn(async move {
360            let mut stream = UnixStream::connect(&path_clone).await.unwrap();
361            let mut buf = vec![0u8; 64];
362            let mut total = Vec::new();
363            for _ in 0..20 {
364                let n = stream.read(&mut buf).await.unwrap();
365                if n == 0 {
366                    break;
367                }
368                total.extend_from_slice(&buf[..n]);
369                if total.len() >= 16 {
370                    break;
371                }
372            }
373            total
374        });
375        let got = tokio::time::timeout(std::time::Duration::from_secs(2), task)
376            .await
377            .expect("timeout")
378            .unwrap();
379        let s = String::from_utf8_lossy(&got);
380        // Sólo C y D (los más viejos A y B fueron evicted).
381        assert!(!s.contains("AAAA"), "should have evicted A: {s:?}");
382        assert!(!s.contains("BBBB"), "should have evicted B: {s:?}");
383        assert!(s.contains("CCCC"), "should keep C: {s:?}");
384        assert!(s.contains("DDDD"), "should keep D: {s:?}");
385    }
386
387    #[tokio::test]
388    async fn drop_removes_socket() {
389        let tmp = tempfile::tempdir().unwrap();
390        let path = tmp.path().join("flow.sock");
391        {
392            let _ch = FlowChannel::new(path.clone()).unwrap();
393            assert!(path.exists());
394        }
395        // Después del drop, el socket file no debe quedar.
396        // Damos un pelín de tiempo al runtime para que el drop corra
397        // mientras estamos en task.
398        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
399        assert!(!path.exists());
400    }
401}