1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
|
use std::collections::HashMap;
use futures::{Stream, StreamExt};
use tokio::sync::watch::Receiver;
use tokio_stream::wrappers::WatchStream;
use crate::{
service::{IntoService, ServiceConfig},
Status,
};
#[derive(Clone)]
pub struct AppState {
rx_map: HashMap<String, Receiver<Status>>,
}
impl AppState {
pub fn spawn_services(configs: HashMap<String, ServiceConfig>) -> AppState {
let mut rx_map = HashMap::new();
let mut tx_map = HashMap::new();
for name in configs.keys() {
let (tx, rx) = tokio::sync::watch::channel(Default::default());
rx_map.insert(name.clone(), rx);
tx_map.insert(name.clone(), tx);
}
for (name, config) in configs {
let (tx, rx) = tokio::sync::watch::channel(Status::default());
rx_map.insert(name.clone(), rx);
tokio::spawn(async move {
let mut stream = config.into_service();
while let Some(res) = stream.next().await {
let status = res.into();
tx.send_if_modified(|s| {
if *s != status {
tracing::debug!(name, ?status, "Updated service status");
*s = status;
true
} else {
false
}
});
}
});
}
AppState { rx_map }
}
pub fn status(&self, k: &str) -> Option<Status> {
self.rx_map.get(k).map(|rx| rx.borrow().clone())
}
pub fn statuses(&self) -> Vec<(String, Status)> {
self.rx_map
.iter()
.map(|(k, v)| (k.clone(), v.borrow().clone()))
.collect()
}
pub fn stream(&self, k: &str) -> Option<impl Stream<Item = Status>> {
self.rx_map.get(k).cloned().map(WatchStream::new)
}
pub fn streams(&self) -> impl Stream<Item = (String, Status)> {
let iter =
self.rx_map.clone().into_iter().map(|(name, rx)| {
WatchStream::new(rx).map(move |status| (name.to_owned(), status))
});
futures::stream::select_all(iter)
}
}
|
