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
74
75
76
77
78
|
use std::collections::HashMap;
use futures::{Stream, StreamExt};
use tokio::sync::watch::Receiver;
use tokio_stream::wrappers::WatchStream;
use crate::{
service::{IntoService, ServiceConfig, ServiceKind},
Status,
};
#[derive(Clone)]
pub struct AppState {
rx_map: HashMap<String, Receiver<Status>>,
indexes: Vec<String>,
}
impl AppState {
pub fn new(configs: Vec<ServiceConfig>) -> AppState {
let mut indexes = Vec::new();
let rx_map = configs
.into_iter()
.map(|ServiceConfig { name, kind }| {
indexes.push(name.clone());
let (tx, rx) = tokio::sync::watch::channel(Status::default());
tokio::spawn(Self::spawn_service(kind, tx.clone()));
(name, rx)
})
.collect();
AppState { rx_map, indexes }
}
#[tracing::instrument(skip(tx))]
async fn spawn_service(kind: ServiceKind, tx: tokio::sync::watch::Sender<Status>) {
let mut stream = kind.into_service();
while let Some(res) = stream.next().await {
let status = res.into();
tx.send_if_modified(|s| {
if *s != status {
tracing::debug!(?status, "Updated service status");
*s = status;
true
} else {
false
}
});
}
}
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.indexes
.iter()
.cloned()
.map(|s| {
let status = self.status(&s).expect("Service was unexpectedly removed");
(s, status)
})
.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.indexes.iter().cloned().map(|s| {
let stream = self.stream(&s).expect("Service was unexpectedly removed");
stream.map(move |status| (s.to_owned(), status))
});
futures::stream::select_all(iter)
}
}
|
