arcly_stream/cluster.rs
1//! Multi-node clustering contracts: origin/edge discovery and stream relay.
2//!
3//! These are **contracts only** — the engine defines the seams an edge tier
4//! plugs into, but ships no concrete discovery or transport (those depend on the
5//! deployment's service mesh / gossip / control plane). An edge node implements
6//! [`ClusterRelay`] to locate a stream's origin, pull it locally, and announce
7//! the streams it serves.
8
9use crate::bus::StreamHandle;
10use crate::{MediaFrame, Result, StreamKey};
11use async_trait::async_trait;
12use tokio_util::sync::CancellationToken;
13
14/// Address of a node in the cluster (opaque to the engine).
15#[derive(Debug, Clone, PartialEq, Eq, Hash)]
16pub struct NodeAddr(pub String);
17
18/// A transport-generic source of frames for the **shared mirror loop**
19/// ([`mirror`]). Both relays mirror the same way — replay the instant-start GOP,
20/// then forward live frames — differing only in *where the frames come from*: an
21/// in-process [`StreamHandle`] subscription, or a TCP `read_record` loop in the
22/// network relay. Implement this for the transport and call [`mirror`] instead of
23/// hand-copying the loop.
24#[async_trait]
25pub trait FrameSource: Send {
26 /// Frames to replay immediately for instant start (the origin's cached GOP),
27 /// in order. Defaults to none.
28 fn replay(&mut self) -> Vec<MediaFrame> {
29 Vec::new()
30 }
31
32 /// The next live frame, or `None` when the source has ended.
33 async fn next(&mut self) -> Option<MediaFrame>;
34}
35
36/// Mirror `src` into `dst`: replay the instant-start buffer, then forward live
37/// frames until the source ends, `cancel` fires, or the local mirror is released.
38/// This is the one loop both [`InProcessRelay`] and the network relay share, so
39/// the federation semantics can't drift between them.
40pub async fn mirror<S: FrameSource>(mut src: S, dst: &StreamHandle, cancel: &CancellationToken) {
41 for frame in src.replay() {
42 if dst.publish_frame(frame).is_err() {
43 return; // local mirror released
44 }
45 }
46 loop {
47 tokio::select! {
48 _ = cancel.cancelled() => return,
49 frame = src.next() => match frame {
50 Some(frame) => {
51 if dst.publish_frame(frame).is_err() {
52 return;
53 }
54 }
55 None => return, // source ended
56 },
57 }
58 }
59}
60
61impl<T: Into<String>> From<T> for NodeAddr {
62 fn from(s: T) -> Self {
63 NodeAddr(s.into())
64 }
65}
66
67/// Origin/edge relay contract. Implement to federate streams across nodes:
68/// an edge that lacks a stream locally [`locate`](Self::locate)s its origin and
69/// [`pull`](Self::pull)s it; an origin [`announce`](Self::announce)s what it has.
70#[async_trait]
71pub trait ClusterRelay: Send + Sync + 'static {
72 /// Find a node currently serving `key`, if any (e.g. via the control plane).
73 async fn locate(&self, key: &StreamKey) -> Result<Option<NodeAddr>>;
74
75 /// Begin relaying `key` from `origin` into the local engine, returning once
76 /// the local mirror is publishing.
77 async fn pull(&self, key: &StreamKey, origin: &NodeAddr) -> Result<()>;
78
79 /// Advertise that this node serves `key` so edges can discover it.
80 async fn announce(&self, key: &StreamKey) -> Result<()>;
81
82 /// Withdraw a previous [`announce`](Self::announce) when the stream ends.
83 async fn withdraw(&self, key: &StreamKey) -> Result<()>;
84
85 /// Demand policy on top of the [`locate`](Self::locate)/[`pull`](Self::pull)
86 /// mechanism: make `key` available locally, returning `true` if a mirror is
87 /// now (or already) present and `false` if no origin serves it.
88 ///
89 /// A host serving a play/WHEP/HLS request calls this so it doesn't have to
90 /// hand-write "locate, then pull if not already local." Implementations
91 /// **should de-duplicate concurrent callers** for the same key (single-flight)
92 /// so a burst of N simultaneous viewer requests for a not-yet-mirrored stream
93 /// triggers **one** pull, not N racing `start_publish` calls.
94 ///
95 /// The default is a plain `locate` + `pull` with **no** coalescing or
96 /// locality check; [`NetworkRelay`](../../arcly_stream_cluster/struct.NetworkRelay.html)
97 /// overrides it with a locality fast-path and single-flight.
98 async fn ensure_mirrored(&self, key: &StreamKey) -> Result<bool> {
99 match self.locate(key).await? {
100 Some(origin) => {
101 self.pull(key, &origin).await?;
102 Ok(true)
103 }
104 None => Ok(false),
105 }
106 }
107}
108
109#[cfg(feature = "cluster")]
110pub use relay::{ClusterDirectory, InProcessRelay};
111
112/// A working, in-process reference [`ClusterRelay`] (feature `cluster`).
113///
114/// It federates streams between nodes that live in the same process — the shape
115/// integration tests and single-box multi-engine setups need — over an in-memory
116/// [`ClusterDirectory`] control plane. A production edge swaps the directory for
117/// its gossip/service-mesh and `pull` for a real transport (RTMP/SRT/QUIC), but
118/// the mirror loop (subscribe origin → republish locally) is identical.
119#[cfg(feature = "cluster")]
120mod relay {
121 use super::{ClusterRelay, FrameSource, NodeAddr};
122 use crate::bus::{PlaybackRegistry, PublishRegistry, StreamHandle, Subscription};
123 use crate::{MediaFrame, Result, StreamError, StreamKey};
124 use async_trait::async_trait;
125 use std::collections::HashMap;
126 use std::time::{Duration, Instant};
127 use tokio_util::sync::CancellationToken;
128
129 /// A [`FrameSource`] over an in-process origin [`StreamHandle`]: snapshot the
130 /// GOP first (for instant start), then forward the live subscription — so the
131 /// in-process relay rides the shared [`mirror`](super::mirror) loop.
132 struct HandleSource {
133 replay: Vec<MediaFrame>,
134 sub: Subscription,
135 }
136
137 impl HandleSource {
138 fn new(handle: StreamHandle) -> Self {
139 // Snapshot the replay buffer *before* subscribing, matching the
140 // original ordering so live frames don't duplicate the GOP tail.
141 let replay = handle
142 .replay_buffer()
143 .into_iter()
144 .map(|f| (*f).clone())
145 .collect();
146 let sub = handle.subscribe_resilient();
147 Self { replay, sub }
148 }
149 }
150
151 #[async_trait]
152 impl FrameSource for HandleSource {
153 fn replay(&mut self) -> Vec<MediaFrame> {
154 std::mem::take(&mut self.replay)
155 }
156 async fn next(&mut self) -> Option<MediaFrame> {
157 self.sub.recv().await.map(|f| (*f).clone())
158 }
159 }
160 use std::sync::{Arc, Mutex};
161
162 /// In-memory cluster directory (control plane): which nodes serve which
163 /// streams. Shared (`Arc`) by every node's relay in the process.
164 ///
165 /// By default entries live until an explicit [`withdraw`](Self::withdraw).
166 /// Opting into a TTL with [`with_ttl`](Self::with_ttl) gives **lease parity**
167 /// with the network [`DirectoryServer`](../../arcly_stream_cluster/struct.DirectoryServer.html):
168 /// an [`announce`](Self::announce) is valid only for the TTL, so a peer that
169 /// stops a stream without `withdraw` (e.g. it crashed) is dropped on its own,
170 /// and `locate` then steers to a live replica instead of a dead node. Re-
171 /// announce within the TTL to keep a lease (the engine's `PublishEnded` event
172 /// makes a tiny auto-withdraw adapter easy, too).
173 #[derive(Debug, Default)]
174 pub struct ClusterDirectory {
175 serving: Mutex<HashMap<StreamKey, HashMap<NodeAddr, Option<Instant>>>>,
176 ttl: Option<Duration>,
177 }
178
179 impl ClusterDirectory {
180 /// A fresh, empty directory whose entries never expire (withdraw-only).
181 pub fn new() -> Self {
182 Self::default()
183 }
184
185 /// A directory whose announces carry a `ttl` lease (network-relay parity);
186 /// an entry not re-announced within `ttl` is reaped on the next access.
187 pub fn with_ttl(ttl: Duration) -> Self {
188 Self {
189 serving: Mutex::new(HashMap::new()),
190 ttl: Some(ttl),
191 }
192 }
193
194 /// Record that `node` serves `key` (refreshing its lease if TTL is set).
195 pub fn announce(&self, node: &NodeAddr, key: &StreamKey) {
196 let expires = self.ttl.map(|t| Instant::now() + t);
197 self.serving
198 .lock()
199 .unwrap()
200 .entry(key.clone())
201 .or_default()
202 .insert(node.clone(), expires);
203 }
204
205 /// Drop `node` from `key`'s server set.
206 pub fn withdraw(&self, node: &NodeAddr, key: &StreamKey) {
207 if let Some(set) = self.serving.lock().unwrap().get_mut(key) {
208 set.remove(node);
209 }
210 }
211
212 /// A live node serving `key` other than `exclude`, if any. Expired leases
213 /// are skipped (and lazily reaped).
214 pub fn locate(&self, key: &StreamKey, exclude: &NodeAddr) -> Option<NodeAddr> {
215 let now = Instant::now();
216 let mut map = self.serving.lock().unwrap();
217 let set = map.get_mut(key)?;
218 set.retain(|_, exp| exp.is_none_or(|e| e > now)); // reap expired
219 set.keys().find(|n| *n != exclude).cloned()
220 }
221 }
222
223 /// An in-process [`ClusterRelay`] for one node: mirrors a peer's stream into
224 /// the local engine.
225 pub struct InProcessRelay {
226 node: NodeAddr,
227 local: Arc<dyn PublishRegistry>,
228 directory: Arc<ClusterDirectory>,
229 peers: HashMap<NodeAddr, Arc<dyn PlaybackRegistry>>,
230 reconnect: bool,
231 /// Active mirrors, so [`withdraw`](ClusterRelay::withdraw) can stop a
232 /// reconnecting supervisor.
233 mirrors: Mutex<HashMap<StreamKey, CancellationToken>>,
234 }
235
236 impl InProcessRelay {
237 /// A relay for `node`, mirroring into `local`, discovering via `directory`.
238 pub fn new(
239 node: impl Into<NodeAddr>,
240 local: Arc<dyn PublishRegistry>,
241 directory: Arc<ClusterDirectory>,
242 ) -> Self {
243 Self {
244 node: node.into(),
245 local,
246 directory,
247 peers: HashMap::new(),
248 reconnect: false,
249 mirrors: Mutex::new(HashMap::new()),
250 }
251 }
252
253 /// Register a peer node's playback registry so [`pull`](ClusterRelay::pull)
254 /// can subscribe to its streams.
255 pub fn with_peer(
256 mut self,
257 addr: impl Into<NodeAddr>,
258 playback: Arc<dyn PlaybackRegistry>,
259 ) -> Self {
260 self.peers.insert(addr.into(), playback);
261 self
262 }
263
264 /// Opt into **reconnect/failover parity** with the network relay: when a
265 /// mirrored origin's stream ends, re-[`locate`](ClusterRelay::locate) the
266 /// stream (excluding the node that just failed) and resume from a live
267 /// replica, keeping the same local publish handle — viewers don't see the
268 /// stream drop. Stops when the stream is gone everywhere or on
269 /// [`withdraw`](ClusterRelay::withdraw). Off by default (one-shot), keeping
270 /// the reference minimal. Pair with [`ClusterDirectory::with_ttl`] so a
271 /// crashed origin's lease lapses and failover actually has somewhere to go.
272 pub fn with_reconnect(mut self, reconnect: bool) -> Self {
273 self.reconnect = reconnect;
274 self
275 }
276 }
277
278 #[async_trait]
279 impl ClusterRelay for InProcessRelay {
280 async fn locate(&self, key: &StreamKey) -> Result<Option<NodeAddr>> {
281 Ok(self.directory.locate(key, &self.node))
282 }
283
284 async fn pull(&self, key: &StreamKey, origin: &NodeAddr) -> Result<()> {
285 let peer = self.peers.get(origin).ok_or_else(|| {
286 StreamError::protocol(format!("cluster: unknown origin node {}", origin.0))
287 })?;
288 // Resolve the origin's live handle and claim the local mirror once.
289 let src = peer.get_stream(key)?;
290 let dst = self.local.start_publish(key).await?;
291
292 let cancel = CancellationToken::new();
293 self.mirrors
294 .lock()
295 .unwrap()
296 .insert(key.clone(), cancel.clone());
297
298 // Mirror via the shared kernel loop (replay GOP → forward live →
299 // release), so in-process and network relays stay in lock-step.
300 let local = Arc::clone(&self.local);
301 let directory = Arc::clone(&self.directory);
302 let peers = self.peers.clone();
303 let reconnect = self.reconnect;
304 let key = key.clone();
305 let mut origin = origin.clone();
306 tokio::spawn(async move {
307 super::mirror(HandleSource::new(src), &dst, &cancel).await;
308 // Reconnect/failover supervisor: re-locate (excluding the failed
309 // origin) and resume on the **same** local handle, until the
310 // stream is gone everywhere or this mirror is withdrawn.
311 while reconnect && !cancel.is_cancelled() {
312 let Some(next) = directory.locate(&key, &origin) else {
313 break; // gone everywhere
314 };
315 origin = next;
316 let Some(peer) = peers.get(&origin) else {
317 break;
318 };
319 let Ok(src) = peer.get_stream(&key) else {
320 break;
321 };
322 super::mirror(HandleSource::new(src), &dst, &cancel).await;
323 }
324 let _ = local.end_publish(&key).await;
325 });
326 Ok(())
327 }
328
329 async fn announce(&self, key: &StreamKey) -> Result<()> {
330 self.directory.announce(&self.node, key);
331 Ok(())
332 }
333
334 async fn withdraw(&self, key: &StreamKey) -> Result<()> {
335 if let Some(cancel) = self.mirrors.lock().unwrap().remove(key) {
336 cancel.cancel(); // stop a reconnecting supervisor
337 }
338 self.directory.withdraw(&self.node, key);
339 Ok(())
340 }
341 }
342
343 #[cfg(test)]
344 mod tests {
345 use super::*;
346 use crate::{AppSpec, CodecId, Engine, MediaFrame};
347
348 fn frame(pts: i64) -> MediaFrame {
349 MediaFrame::new_video(
350 pts,
351 pts,
352 bytes::Bytes::from_static(b"x"),
353 CodecId::H264,
354 true,
355 )
356 }
357
358 #[tokio::test]
359 async fn edge_locates_and_mirrors_origin_stream() {
360 let directory = Arc::new(ClusterDirectory::new());
361 let key = StreamKey::new("live", "cam");
362
363 // Origin node publishes the stream and announces it.
364 let origin = Engine::builder()
365 .application(AppSpec::new("live").gop_cache(8))
366 .build();
367 let origin_relay = InProcessRelay::new("origin", origin.clone(), directory.clone());
368 let src_handle = origin.start_publish(&key).await.unwrap();
369 origin_relay.announce(&key).await.unwrap();
370 src_handle.publish_frame(frame(0)).unwrap();
371
372 // Edge node has the stream locally? No — it locates + pulls the origin.
373 let edge = Engine::builder()
374 .application(AppSpec::new("live").gop_cache(8))
375 .build();
376 let edge_relay = InProcessRelay::new("edge", edge.clone(), directory.clone())
377 .with_peer("origin", origin.clone());
378
379 assert!(edge.get_stream(&key).is_err(), "not local yet");
380 let found = edge_relay.locate(&key).await.unwrap();
381 assert_eq!(found, Some(NodeAddr::from("origin")));
382 edge_relay.pull(&key, &found.unwrap()).await.unwrap();
383
384 // The local mirror is now publishing and receives forwarded frames.
385 let mirror = edge.get_stream(&key).expect("local mirror exists");
386 let mut sub = mirror.subscribe_resilient();
387 src_handle.publish_frame(frame(1)).unwrap();
388 let got = tokio::time::timeout(std::time::Duration::from_secs(5), sub.recv())
389 .await
390 .expect("a frame was mirrored")
391 .expect("frame");
392 assert!(got.is_video());
393 }
394
395 #[tokio::test]
396 async fn locate_excludes_self_and_withdraw_clears() {
397 let directory = Arc::new(ClusterDirectory::new());
398 let key = StreamKey::new("live", "s");
399 let engine = Engine::builder().application(AppSpec::new("live")).build();
400 let relay = InProcessRelay::new("only", engine, directory.clone());
401
402 relay.announce(&key).await.unwrap();
403 // The only server is ourselves → locate returns None.
404 assert_eq!(relay.locate(&key).await.unwrap(), None);
405 relay.withdraw(&key).await.unwrap();
406 assert!(directory.locate(&key, &NodeAddr::from("other")).is_none());
407 }
408
409 #[tokio::test]
410 async fn ttl_directory_reaps_an_unrenewed_lease() {
411 // Real (short) TTL: the directory uses std::Instant, not the mock clock.
412 let dir = ClusterDirectory::with_ttl(Duration::from_millis(80));
413 let key = StreamKey::new("live", "cam");
414 dir.announce(&NodeAddr::from("origin"), &key);
415 // Still leased now.
416 assert_eq!(
417 dir.locate(&key, &NodeAddr::from("edge")),
418 Some(NodeAddr::from("origin"))
419 );
420 // Past the TTL without renewal → reaped, locate finds nothing.
421 tokio::time::sleep(Duration::from_millis(160)).await;
422 assert_eq!(dir.locate(&key, &NodeAddr::from("edge")), None);
423 // Re-announce restores the lease (renewal parity).
424 dir.announce(&NodeAddr::from("origin"), &key);
425 assert!(dir.locate(&key, &NodeAddr::from("edge")).is_some());
426 }
427
428 #[tokio::test]
429 async fn reconnect_fails_over_to_a_second_origin() {
430 let directory = Arc::new(ClusterDirectory::new());
431 let key = StreamKey::new("live", "cam");
432
433 // Two origins (A, B) both publish + announce the same stream.
434 let origin_a = Engine::builder()
435 .application(AppSpec::new("live").gop_cache(8))
436 .build();
437 let origin_b = Engine::builder()
438 .application(AppSpec::new("live").gop_cache(8))
439 .build();
440 let a = origin_a.start_publish(&key).await.unwrap();
441 let b = origin_b.start_publish(&key).await.unwrap();
442 directory.announce(&NodeAddr::from("A"), &key);
443 directory.announce(&NodeAddr::from("B"), &key);
444 a.publish_frame(frame(0)).unwrap();
445 b.publish_frame(frame(0)).unwrap();
446
447 // Edge mirrors with reconnect on, starting from A.
448 let edge = Engine::builder()
449 .application(AppSpec::new("live").gop_cache(8))
450 .build();
451 let relay = InProcessRelay::new("edge", edge.clone(), directory.clone())
452 .with_peer("A", origin_a.clone())
453 .with_peer("B", origin_b.clone())
454 .with_reconnect(true);
455 relay.pull(&key, &NodeAddr::from("A")).await.unwrap();
456
457 let mirror = edge.get_stream(&key).expect("local mirror");
458 let mut sub = mirror.subscribe_resilient();
459
460 // A drops out of the directory and ends its stream → supervisor must
461 // re-locate to B on the *same* local handle.
462 directory.withdraw(&NodeAddr::from("A"), &key);
463 origin_a.end_publish(&key).await.unwrap();
464
465 // Frames from B now reach the unchanged local mirror.
466 for n in 1..20 {
467 b.publish_frame(frame(n)).unwrap();
468 }
469 let got = tokio::time::timeout(std::time::Duration::from_secs(5), sub.recv())
470 .await
471 .expect("a frame arrived after failover")
472 .expect("frame");
473 assert!(got.is_video());
474 assert!(edge.get_stream(&key).is_ok(), "local handle stable");
475 }
476 }
477}