1use std::collections::HashSet;
9use std::future::Future;
10use std::pin::Pin;
11use std::sync::Arc;
12
13use spvirit_codec::spvd_decode::{DecodedValue, StructureDesc};
14use spvirit_types::NtPayload;
15use tokio::sync::{RwLock, mpsc};
16use tracing::debug;
17
18#[derive(Debug, Clone)]
24pub struct PvInfo {
25 pub descriptor: StructureDesc,
27 pub writable: bool,
29}
30
31pub trait Source: Send + Sync {
56 fn claim(&self, name: &str) -> Pin<Box<dyn Future<Output = Option<PvInfo>> + Send + '_>>;
60
61 fn get(&self, name: &str) -> Pin<Box<dyn Future<Output = Option<NtPayload>> + Send + '_>>;
65
66 fn put(
71 &self,
72 name: &str,
73 value: &DecodedValue,
74 ) -> Pin<Box<dyn Future<Output = Result<Vec<(String, NtPayload)>, String>> + Send + '_>>;
75
76 fn subscribe(
80 &self,
81 name: &str,
82 ) -> Pin<Box<dyn Future<Output = Option<mpsc::Receiver<NtPayload>>> + Send + '_>>;
83
84 fn rpc(
92 &self,
93 _name: &str,
94 _args: &DecodedValue,
95 ) -> Pin<Box<dyn Future<Output = Result<NtPayload, String>> + Send + '_>> {
96 Box::pin(async { Err("RPC not supported".to_string()) })
97 }
98
99 fn names(&self) -> Pin<Box<dyn Future<Output = Vec<String>> + Send + '_>>;
101}
102
103pub trait StoreSource: Source {
115 fn record_names(&self) -> Vec<String>;
116}
117
118struct SourceEntry {
123 label: String,
125 order: i32,
127 source: Arc<dyn Source>,
129 is_store: bool,
134}
135
136pub struct SourceRegistry {
145 sources: RwLock<Vec<SourceEntry>>,
146 shadow_checked: RwLock<HashSet<String>>,
150}
151
152impl SourceRegistry {
153 pub fn new() -> Self {
155 Self {
156 sources: RwLock::new(Vec::new()),
157 shadow_checked: RwLock::new(HashSet::new()),
158 }
159 }
160
161 pub async fn add(&self, label: impl Into<String>, order: i32, source: Arc<dyn Source>) {
165 self.insert(label.into(), order, source, false).await;
166 }
167
168 pub async fn add_store(&self, label: impl Into<String>, order: i32, source: Arc<dyn Source>) {
171 self.insert(label.into(), order, source, true).await;
172 }
173
174 async fn insert(&self, label: String, order: i32, source: Arc<dyn Source>, is_store: bool) {
175 debug!(
176 "SourceRegistry: adding source '{}' at order {} (store: {})",
177 label, order, is_store
178 );
179 let mut sources = self.sources.write().await;
180 sources.push(SourceEntry {
181 label,
182 order,
183 source,
184 is_store,
185 });
186 sources.sort_by_key(|e| e.order);
187 }
188
189 pub async fn remove(&self, label: &str) {
191 debug!("SourceRegistry: removing source '{}'", label);
192 let mut sources = self.sources.write().await;
193 sources.retain(|e| e.label != label);
194 }
195
196 pub async fn claim(&self, name: &str) -> Option<PvInfo> {
200 let sources = self.sources.read().await;
201 for entry in sources.iter() {
202 if let Some(info) = entry.source.claim(name).await {
203 if !entry.is_store {
204 self.warn_if_shadowing_a_store(&sources, &entry.label, name)
205 .await;
206 }
207 return Some(info);
208 }
209 }
210 None
211 }
212
213 async fn warn_if_shadowing_a_store(&self, sources: &[SourceEntry], winner: &str, name: &str) {
223 if self.shadow_checked.read().await.contains(name) {
224 return;
225 }
226 if !self.shadow_checked.write().await.insert(name.to_string()) {
227 return;
229 }
230 for entry in sources.iter().filter(|e| e.is_store) {
231 if entry.source.claim(name).await.is_some() {
232 tracing::warn!(
233 "source '{winner}' shadows store '{}' for PV '{name}': the store's \
234 value will never be served",
235 entry.label
236 );
237 return;
238 }
239 }
240 }
241
242 pub async fn has_pv(&self, name: &str) -> bool {
244 self.claim(name).await.is_some()
245 }
246
247 pub async fn get(&self, name: &str) -> Option<NtPayload> {
249 let sources = self.sources.read().await;
250 for entry in sources.iter() {
251 if entry.source.claim(name).await.is_some() {
252 return entry.source.get(name).await;
253 }
254 }
255 None
256 }
257
258 pub async fn get_descriptor(&self, name: &str) -> Option<StructureDesc> {
260 self.claim(name).await.map(|info| info.descriptor)
261 }
262
263 pub async fn is_writable(&self, name: &str) -> bool {
265 self.claim(name).await.is_some_and(|info| info.writable)
266 }
267
268 pub async fn put(
270 &self,
271 name: &str,
272 value: &DecodedValue,
273 ) -> Result<Vec<(String, NtPayload)>, String> {
274 let sources = self.sources.read().await;
275 for entry in sources.iter() {
276 if entry.source.claim(name).await.is_some() {
277 return entry.source.put(name, value).await;
278 }
279 }
280 Err(format!("PV '{}' not found", name))
281 }
282
283 pub async fn subscribe(&self, name: &str) -> Option<mpsc::Receiver<NtPayload>> {
285 let sources = self.sources.read().await;
286 for entry in sources.iter() {
287 if entry.source.claim(name).await.is_some() {
288 return entry.source.subscribe(name).await;
289 }
290 }
291 None
292 }
293
294 pub async fn rpc(&self, name: &str, args: &DecodedValue) -> Result<NtPayload, String> {
296 let sources = self.sources.read().await;
297 for entry in sources.iter() {
298 if entry.source.claim(name).await.is_some() {
299 return entry.source.rpc(name, args).await;
300 }
301 }
302 Err(format!("RPC channel '{}' not found", name))
303 }
304
305 pub async fn names(&self) -> Vec<String> {
307 let sources = self.sources.read().await;
308 let mut seen = HashSet::new();
309 let mut all = Vec::new();
310 for entry in sources.iter() {
311 for name in entry.source.names().await {
312 if seen.insert(name.clone()) {
313 all.push(name);
314 }
315 }
316 }
317 all.sort();
318 all
319 }
320}
321
322impl Default for SourceRegistry {
323 fn default() -> Self {
324 Self::new()
325 }
326}
327
328#[cfg(test)]
329mod tests {
330 use super::*;
331
332 struct StubSource {
334 names: Vec<String>,
335 claims: std::sync::atomic::AtomicUsize,
336 }
337
338 impl StubSource {
339 fn new(names: &[&str]) -> Self {
340 Self {
341 names: names.iter().map(|s| s.to_string()).collect(),
342 claims: std::sync::atomic::AtomicUsize::new(0),
343 }
344 }
345
346 fn claim_count(&self) -> usize {
347 self.claims.load(std::sync::atomic::Ordering::SeqCst)
348 }
349 }
350
351 impl Source for StubSource {
352 fn claim(&self, name: &str) -> Pin<Box<dyn Future<Output = Option<PvInfo>> + Send + '_>> {
353 self.claims.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
354 let claimed = self.names.iter().any(|n| n == name);
355 Box::pin(async move {
356 claimed.then(|| PvInfo {
357 descriptor: StructureDesc::default(),
358 writable: true,
359 })
360 })
361 }
362
363 fn get(&self, _name: &str) -> Pin<Box<dyn Future<Output = Option<NtPayload>> + Send + '_>> {
364 Box::pin(async { None })
365 }
366
367 fn put(
368 &self,
369 _name: &str,
370 _value: &DecodedValue,
371 ) -> Pin<Box<dyn Future<Output = Result<Vec<(String, NtPayload)>, String>> + Send + '_>>
372 {
373 Box::pin(async { Ok(vec![]) })
374 }
375
376 fn subscribe(
377 &self,
378 _name: &str,
379 ) -> Pin<Box<dyn Future<Output = Option<mpsc::Receiver<NtPayload>>> + Send + '_>> {
380 Box::pin(async { None })
381 }
382
383 fn names(&self) -> Pin<Box<dyn Future<Output = Vec<String>> + Send + '_>> {
384 let names = self.names.clone();
385 Box::pin(async move { names })
386 }
387 }
388
389 #[tokio::test]
390 async fn stores_are_recorded_as_stores_and_sources_are_not() {
391 let reg = SourceRegistry::new();
392 reg.add_store("builtin", 0, Arc::new(StubSource::new(&["A"]))).await;
393 reg.add("custom", 10, Arc::new(StubSource::new(&["B"]))).await;
394 let flags: Vec<(String, bool)> = reg
395 .sources
396 .read()
397 .await
398 .iter()
399 .map(|e| (e.label.clone(), e.is_store))
400 .collect();
401 assert_eq!(
402 flags,
403 vec![("builtin".to_string(), true), ("custom".to_string(), false)]
404 );
405 }
406
407 #[tokio::test]
408 async fn a_store_added_late_still_sorts_by_order() {
409 let reg = SourceRegistry::new();
410 reg.add("custom", 10, Arc::new(StubSource::new(&["B"]))).await;
411 reg.add_store("builtin", 0, Arc::new(StubSource::new(&["A"]))).await;
412 let labels: Vec<String> = reg
413 .sources
414 .read()
415 .await
416 .iter()
417 .map(|e| e.label.clone())
418 .collect();
419 assert_eq!(labels, vec!["builtin".to_string(), "custom".to_string()]);
420 }
421
422 #[tokio::test]
425 async fn a_source_shadowing_a_store_still_wins_the_claim() {
426 let reg = SourceRegistry::new();
427 reg.add("override", -1, Arc::new(StubSource::new(&["PV:X"]))).await;
428 reg.add_store("builtin", 0, Arc::new(StubSource::new(&["PV:X"]))).await;
429 assert!(reg.claim("PV:X").await.is_some());
430 }
431
432 #[tokio::test]
435 async fn the_shadow_check_runs_once_per_pv() {
436 let reg = SourceRegistry::new();
437 let store = Arc::new(StubSource::new(&["PV:X"]));
438 reg.add("override", -1, Arc::new(StubSource::new(&["PV:X"]))).await;
439 reg.add_store("builtin", 0, store.clone()).await;
440 let before = store.claim_count();
441 for _ in 0..5 {
442 reg.claim("PV:X").await;
443 }
444 assert_eq!(
445 store.claim_count() - before,
446 1,
447 "the shadowed store must be consulted exactly once"
448 );
449 }
450
451 #[tokio::test]
454 async fn an_unshadowed_source_claim_is_also_checked_only_once() {
455 let reg = SourceRegistry::new();
456 let store = Arc::new(StubSource::new(&["PV:OTHER"]));
457 reg.add("plain", -1, Arc::new(StubSource::new(&["PV:X"]))).await;
458 reg.add_store("builtin", 0, store.clone()).await;
459 let before = store.claim_count();
460 for _ in 0..5 {
461 reg.claim("PV:X").await;
462 }
463 assert_eq!(store.claim_count() - before, 1);
464 }
465
466 #[tokio::test]
468 async fn a_store_winning_its_own_claim_consults_nothing_else() {
469 let reg = SourceRegistry::new();
470 let other = Arc::new(StubSource::new(&["PV:X"]));
471 reg.add_store("builtin", 0, Arc::new(StubSource::new(&["PV:X"]))).await;
472 reg.add_store("second", 5, other.clone()).await;
473 let before = other.claim_count();
474 reg.claim("PV:X").await;
475 assert_eq!(other.claim_count() - before, 0);
476 }
477
478 #[derive(Clone, Default)]
482 struct CaptureWriter(Arc<std::sync::Mutex<Vec<u8>>>);
483
484 impl std::io::Write for CaptureWriter {
485 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
486 self.0.lock().unwrap().extend_from_slice(buf);
487 Ok(buf.len())
488 }
489
490 fn flush(&mut self) -> std::io::Result<()> {
491 Ok(())
492 }
493 }
494
495 #[tokio::test]
500 async fn the_shadow_warning_is_emitted_once_not_just_counted() {
501 let buffer = CaptureWriter::default();
502 let writer = buffer.clone();
503 let subscriber = tracing_subscriber::fmt()
504 .with_max_level(tracing::Level::WARN)
505 .with_ansi(false)
506 .without_time()
507 .with_writer(move || writer.clone())
508 .finish();
509
510 let _subscriber_guard = tracing::subscriber::set_default(subscriber);
511
512 let reg = SourceRegistry::new();
513 reg.add("override", -1, Arc::new(StubSource::new(&["PV:X"]))).await;
514 reg.add_store("builtin", 0, Arc::new(StubSource::new(&["PV:X"]))).await;
515
516 reg.claim("PV:X").await;
517 reg.claim("PV:X").await;
518
519 let captured = String::from_utf8(buffer.0.lock().unwrap().clone()).unwrap();
520 let warnings: Vec<&str> = captured.lines().filter(|l| !l.is_empty()).collect();
521 assert_eq!(
522 warnings.len(),
523 1,
524 "expected exactly one warning event, got: {captured:?}"
525 );
526 assert!(warnings[0].contains("PV:X"), "missing PV name: {captured:?}");
527 assert!(warnings[0].contains("override"), "missing source label: {captured:?}");
528 assert!(warnings[0].contains("builtin"), "missing store label: {captured:?}");
529 }
530}