1use std::{
2 collections::{BTreeMap, VecDeque},
3 fmt::Debug,
4 net::SocketAddr,
5};
6
7use atm0s_sdn_identity::{ConnId, NodeId};
8use atm0s_sdn_router::{RouteRule, ServiceBroadcastLevel};
9use sans_io_runtime::collections::DynamicDeque;
10use serde::{de::DeserializeOwned, Deserialize, Serialize};
11
12use crate::{
13 base::{
14 ConnectionEvent, NetOutgoingMeta, Service, ServiceBuilder, ServiceControlActor, ServiceCtx, ServiceInput, ServiceOutput, ServiceSharedInput, ServiceWorker, ServiceWorkerCtx,
15 ServiceWorkerInput, ServiceWorkerOutput, Ttl,
16 },
17 features::{data, FeaturesControl, FeaturesEvent},
18};
19
20pub const SERVICE_ID: u8 = 1;
21pub const SERVICE_NAME: &str = "visualization";
22
23const NODE_TIMEOUT_MS: u64 = 10000; const NODE_PING_MS: u64 = 5000;
25const NODE_PING_TTL: u8 = 5;
26
27const DATA_PORT: u16 = 0;
28
29fn data_cmd<UserData, SE, TW>(cmd: data::Control) -> ServiceOutput<UserData, FeaturesControl, SE, TW> {
30 ServiceOutput::FeatureControl(FeaturesControl::Data(cmd))
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
34pub struct ConnectionInfo {
35 pub conn: ConnId,
36 pub dest: NodeId,
37 pub local: SocketAddr,
38 pub remote: SocketAddr,
39 pub rtt_ms: u32,
40}
41
42struct NodeInfo<Info> {
43 last_ping_ms: u64,
44 info: Info,
45 conns: Vec<ConnectionInfo>,
46}
47
48#[derive(Debug, Clone)]
49pub enum Control<Info> {
50 Subscribe,
51 GetAll,
52 UpdateInfo(Info),
53}
54
55#[derive(Debug, Clone, PartialEq, Eq)]
56pub enum Event<Info> {
57 GotAll(Vec<(NodeId, Info, Vec<ConnectionInfo>)>),
58 NodeChanged(NodeId, Info, Vec<ConnectionInfo>),
59 NodeRemoved(NodeId),
60}
61
62#[derive(Debug, Serialize, Deserialize)]
63enum Message<Info> {
64 Snapshot(NodeId, Info, Vec<ConnectionInfo>),
65}
66
67pub struct VisualizationService<UserData, SC, SE, TC, TW, Info> {
68 info: Info,
69 last_ping: u64,
70 broadcast_seq: u16,
71 queue: VecDeque<ServiceOutput<UserData, FeaturesControl, SE, TW>>,
72 conns: BTreeMap<ConnId, ConnectionInfo>,
73 network_nodes: BTreeMap<NodeId, NodeInfo<Info>>,
74 subscribers: Vec<ServiceControlActor<UserData>>,
75 shutdown: bool,
76 _tmp: std::marker::PhantomData<(SC, TC)>,
77}
78
79impl<UserData: Copy, SC, SE, TC, TW, Info: Clone> VisualizationService<UserData, SC, SE, TC, TW, Info>
80where
81 SC: From<Control<Info>> + TryInto<Control<Info>>,
82 SE: From<Event<Info>> + TryInto<Event<Info>>,
83{
84 pub fn new(info: Info) -> Self {
85 Self {
86 info,
87 broadcast_seq: 0,
88 last_ping: 0,
89 conns: BTreeMap::new(),
90 network_nodes: BTreeMap::new(),
91 queue: VecDeque::from([ServiceOutput::FeatureControl(FeaturesControl::Data(data::Control::DataListen(DATA_PORT)))]),
92 subscribers: Vec::new(),
93 shutdown: false,
94 _tmp: std::marker::PhantomData,
95 }
96 }
97
98 fn fire_event(&mut self, event: Event<Info>) {
99 for sub in self.subscribers.iter() {
100 self.queue.push_back(ServiceOutput::Event(*sub, event.clone().into()));
101 }
102 }
103}
104
105impl<UserData: Copy + Eq, SC, SE, TC, TW, Info> Service<UserData, FeaturesControl, FeaturesEvent, SC, SE, TC, TW> for VisualizationService<UserData, SC, SE, TC, TW, Info>
106where
107 Info: Debug + Clone + Serialize + DeserializeOwned,
108 SC: From<Control<Info>> + TryInto<Control<Info>>,
109 SE: From<Event<Info>> + TryInto<Event<Info>>,
110{
111 fn is_service_empty(&self) -> bool {
112 self.shutdown && self.queue.is_empty()
113 }
114
115 fn service_id(&self) -> u8 {
116 SERVICE_ID
117 }
118
119 fn service_name(&self) -> &str {
120 SERVICE_NAME
121 }
122
123 fn on_shared_input<'a>(&mut self, ctx: &ServiceCtx, now: u64, input: ServiceSharedInput) {
124 match input {
125 ServiceSharedInput::Tick(_) => {
126 let mut to_remove = Vec::new();
127 for (node, info) in self.network_nodes.iter() {
128 if now >= NODE_TIMEOUT_MS + info.last_ping_ms {
129 log::warn!("[Visualization] Node {} is dead after timeout {NODE_TIMEOUT_MS} ms", node);
130 to_remove.push(*node);
131 }
132 }
133 for node in to_remove {
134 self.fire_event(Event::NodeRemoved(node));
135 self.network_nodes.remove(&node);
136 }
137
138 if now >= self.last_ping + NODE_PING_MS {
139 log::debug!("[Visualization] Sending Snapshot to collector with interval {NODE_PING_MS} ms with {} conns", self.conns.len());
140 self.last_ping = now;
141 let msg = Message::Snapshot(ctx.node_id, self.info.clone(), self.conns.values().cloned().collect::<Vec<_>>());
142 let seq = self.broadcast_seq;
143 self.broadcast_seq = self.broadcast_seq.wrapping_add(1);
144 self.queue.push_back(data_cmd(data::Control::DataSendRule(
145 DATA_PORT,
146 RouteRule::ToServices(SERVICE_ID, ServiceBroadcastLevel::Global, seq),
147 NetOutgoingMeta::new(false, Ttl(NODE_PING_TTL), 0, true),
148 bincode::serialize(&msg).expect("Should to bytes"),
149 )));
150 }
151 }
152 ServiceSharedInput::Connection(ConnectionEvent::Connecting(_ctx)) => {}
153 ServiceSharedInput::Connection(ConnectionEvent::ConnectError(_ctx, _err)) => {}
154 ServiceSharedInput::Connection(ConnectionEvent::Connected(ctx, _)) => {
155 log::info!("[Visualization] New connection from {} to {}, set default rtt_ms to 1000ms", ctx.pair, ctx.node);
156 self.conns.insert(
157 ctx.conn,
158 ConnectionInfo {
159 conn: ctx.conn,
160 dest: ctx.node,
161 local: ctx.pair.local,
162 remote: ctx.pair.remote,
163 rtt_ms: 1000,
164 },
165 );
166 }
167 ServiceSharedInput::Connection(ConnectionEvent::Stats(ctx, stats)) => {
168 log::debug!("[Visualization] Update rtt_ms for connection from {} to {} to {}ms", ctx.pair, ctx.node, stats.rtt_ms);
169 let entry = self.conns.entry(ctx.conn).or_insert(ConnectionInfo {
170 conn: ctx.conn,
171 dest: ctx.node,
172 local: ctx.pair.local,
173 remote: ctx.pair.remote,
174 rtt_ms: 1000,
175 });
176 entry.rtt_ms = stats.rtt_ms;
177 }
178 ServiceSharedInput::Connection(ConnectionEvent::Disconnected(ctx)) => {
179 log::info!("[Visualization] Connection from {} to {} is disconnected", ctx.pair, ctx.node);
180 self.conns.remove(&ctx.conn);
181 }
182 }
183 }
184
185 fn on_input(&mut self, _ctx: &ServiceCtx, now: u64, input: ServiceInput<UserData, FeaturesEvent, SC, TC>) {
186 match input {
187 ServiceInput::FeatureEvent(FeaturesEvent::Data(data::Event::Recv(_port, meta, buf))) => {
188 if !meta.secure {
189 log::warn!("[Visualization] reject unsecure message");
190 return;
191 }
192 if let Ok(msg) = bincode::deserialize::<Message<Info>>(&buf) {
193 match msg {
194 Message::Snapshot(from, info, conns) => {
195 log::debug!("[Visualization] Got snapshot from {} with info {:?} {} connections", from, info, conns.len());
196 self.fire_event(Event::NodeChanged(from, info.clone(), conns.clone()));
197 self.network_nodes.insert(from, NodeInfo { last_ping_ms: now, info, conns });
198 }
199 }
200 }
201 }
202 ServiceInput::Control(actor, control) => {
203 let mut push_all = || {
204 let all = self.network_nodes.iter().map(|(k, v)| (*k, v.info.clone(), v.conns.clone())).collect();
205 self.queue.push_back(ServiceOutput::Event(actor, Event::GotAll(all).into()));
206 };
207 if let Ok(control) = control.try_into() {
208 match control {
209 Control::GetAll => {
210 push_all();
211 }
212 Control::Subscribe => {
213 if !self.subscribers.contains(&actor) {
214 self.subscribers.push(actor);
215 log::info!("[Visualization] New subscriber, sending snapshot with {} nodes", self.network_nodes.len());
216 push_all();
217 }
218 }
219 Control::UpdateInfo(info) => {
220 self.info = info;
221 }
222 }
223 }
224 }
225 _ => {}
226 }
227 }
228
229 fn on_shutdown(&mut self, _ctx: &ServiceCtx, _now: u64) {
230 log::info!("[VisualizationService] Shutdown");
231 self.shutdown = true;
232 }
233
234 fn pop_output2(&mut self, _now: u64) -> Option<ServiceOutput<UserData, FeaturesControl, SE, TW>> {
235 self.queue.pop_front()
236 }
237}
238
239pub struct VisualizationServiceWorker<UserData, SC, SE, TC> {
240 queue: DynamicDeque<ServiceWorkerOutput<UserData, FeaturesControl, FeaturesEvent, SC, SE, TC>, 8>,
241 shutdown: bool,
242}
243
244impl<UserData, SC, SE, TC, TW> ServiceWorker<UserData, FeaturesControl, FeaturesEvent, SC, SE, TC, TW> for VisualizationServiceWorker<UserData, SC, SE, TC> {
245 fn is_service_empty(&self) -> bool {
246 self.shutdown && self.queue.is_empty()
247 }
248
249 fn service_id(&self) -> u8 {
250 SERVICE_ID
251 }
252
253 fn service_name(&self) -> &str {
254 SERVICE_NAME
255 }
256
257 fn on_tick(&mut self, _ctx: &ServiceWorkerCtx, _now: u64, _tick_count: u64) {}
258
259 fn on_input(&mut self, _ctx: &ServiceWorkerCtx, _now: u64, input: ServiceWorkerInput<UserData, FeaturesEvent, SC, TW>) {
260 match input {
261 ServiceWorkerInput::Control(actor, control) => self.queue.push_back(ServiceWorkerOutput::ForwardControlToController(actor, control)),
262 ServiceWorkerInput::FeatureEvent(event) => self.queue.push_back(ServiceWorkerOutput::ForwardFeatureEventToController(event)),
263 ServiceWorkerInput::FromController(_) => {}
264 }
265 }
266
267 fn pop_output2(&mut self, _now: u64) -> Option<ServiceWorkerOutput<UserData, FeaturesControl, FeaturesEvent, SC, SE, TC>> {
268 self.queue.pop_front()
269 }
270
271 fn on_shutdown(&mut self, _ctx: &ServiceWorkerCtx, _now: u64) {
272 self.shutdown = true;
273 }
274}
275
276pub struct VisualizationServiceBuilder<UserData, SC, SE, TC, TW, Info> {
277 info: Info,
278 collector: bool,
279 _tmp: std::marker::PhantomData<(UserData, SC, SE, TC, TW, Info)>,
280}
281
282impl<UserData, SC, SE, TC, TW, Info> VisualizationServiceBuilder<UserData, SC, SE, TC, TW, Info> {
283 pub fn new(info: Info, collector: bool) -> Self {
284 log::info!("[Visualization] started as collector node => will receive metric from all other nodes");
285 Self {
286 info,
287 collector,
288 _tmp: std::marker::PhantomData,
289 }
290 }
291}
292
293impl<UserData, SC, SE, TC, TW, Info> ServiceBuilder<UserData, FeaturesControl, FeaturesEvent, SC, SE, TC, TW> for VisualizationServiceBuilder<UserData, SC, SE, TC, TW, Info>
294where
295 UserData: 'static + Debug + Send + Sync + Copy + Eq,
296 Info: 'static + Debug + Serialize + DeserializeOwned + Clone + Send + Sync,
297 SC: 'static + Debug + Send + Sync + From<Control<Info>> + TryInto<Control<Info>>,
298 SE: 'static + Debug + Send + Sync + From<Event<Info>> + TryInto<Event<Info>>,
299 TC: 'static + Debug + Send + Sync,
300 TW: 'static + Debug + Send + Sync,
301{
302 fn service_id(&self) -> u8 {
303 SERVICE_ID
304 }
305
306 fn service_name(&self) -> &str {
307 SERVICE_NAME
308 }
309
310 fn discoverable(&self) -> bool {
311 self.collector
312 }
313
314 fn create(&self) -> Box<dyn Service<UserData, FeaturesControl, FeaturesEvent, SC, SE, TC, TW>> {
315 Box::new(VisualizationService::new(self.info.clone()))
316 }
317
318 fn create_worker(&self) -> Box<dyn ServiceWorker<UserData, FeaturesControl, FeaturesEvent, SC, SE, TC, TW>> {
319 Box::new(VisualizationServiceWorker {
320 queue: Default::default(),
321 shutdown: false,
322 })
323 }
324}
325
326#[cfg(test)]
327mod test {
328 use atm0s_sdn_identity::{ConnId, NodeId};
329 use atm0s_sdn_router::{RouteRule, ServiceBroadcastLevel};
330 use serde::{Deserialize, Serialize};
331
332 use crate::{
333 base::{ConnectionCtx, ConnectionEvent, MockDecryptor, MockEncryptor, NetIncomingMeta, NetOutgoingMeta, SecureContext, Service, ServiceCtx, ServiceInput, ServiceSharedInput, Ttl},
334 data_plane::NetPair,
335 features::{
336 data::{Control as DataControl, Event as DataEvent},
337 FeaturesEvent,
338 },
339 services::visualization::{data_cmd, Message, DATA_PORT, NODE_PING_MS, NODE_PING_TTL, NODE_TIMEOUT_MS},
340 };
341
342 use super::{Control, Event, VisualizationService, SERVICE_ID};
343
344 #[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone)]
345 struct Info(u8);
346
347 fn data_event(event: DataEvent) -> ServiceInput<(), FeaturesEvent, Control<Info>, ()> {
348 ServiceInput::FeatureEvent(FeaturesEvent::Data(event))
349 }
350
351 fn connected_event(node: NodeId) -> ConnectionEvent {
352 ConnectionEvent::Connected(
353 ConnectionCtx {
354 conn: ConnId::from_in(0, node as u64),
355 node,
356 pair: NetPair::new_str("1.1.1.1:1000", "2.2.2.2:2000").expect("Should parse pair"),
357 },
358 SecureContext {
359 encryptor: Box::new(MockEncryptor::new()),
360 decryptor: Box::new(MockDecryptor::new()),
361 },
362 )
363 }
364
365 fn disconnected_event(node: NodeId) -> ConnectionEvent {
366 ConnectionEvent::Disconnected(ConnectionCtx {
367 conn: ConnId::from_in(0, node as u64),
368 node,
369 pair: NetPair::new_str("1.1.1.1:1000", "2.2.2.2:2000").expect("Should parse pair"),
370 })
371 }
372
373 #[test]
374 fn agent_should_prediotic_sending_snapshot() {
375 let node_info = Info(1);
376 let node_id = 1;
377 let ctx = ServiceCtx { node_id, session: 0 };
378 let mut service = VisualizationService::<(), Control<Info>, Event<Info>, (), (), _>::new(node_info.clone());
379
380 assert_eq!(service.pop_output2(0), Some(data_cmd(DataControl::DataListen(DATA_PORT))));
381 assert_eq!(service.pop_output2(0), None);
382
383 service.on_shared_input(&ctx, NODE_PING_MS, ServiceSharedInput::Tick(0));
384 assert_eq!(
385 service.pop_output2(NODE_PING_MS),
386 Some(data_cmd(DataControl::DataSendRule(
387 DATA_PORT,
388 RouteRule::ToServices(SERVICE_ID, ServiceBroadcastLevel::Global, 0),
389 NetOutgoingMeta::new(false, Ttl(NODE_PING_TTL), 0, true),
390 bincode::serialize(&Message::Snapshot(node_id, node_info.clone(), vec![])).expect("Should to bytes")
391 )))
392 );
393
394 service.on_shared_input(&ctx, NODE_PING_MS * 2, ServiceSharedInput::Tick(0));
395 assert_eq!(
396 service.pop_output2(NODE_PING_MS * 2),
397 Some(data_cmd(DataControl::DataSendRule(
398 DATA_PORT,
399 RouteRule::ToServices(SERVICE_ID, ServiceBroadcastLevel::Global, 1),
400 NetOutgoingMeta::new(false, Ttl(NODE_PING_TTL), 0, true),
401 bincode::serialize(&Message::Snapshot(node_id, node_info.clone(), vec![])).expect("Should to bytes")
402 )))
403 );
404 }
405
406 #[test]
407 fn agent_handle_connection_event() {
408 let node_info = Info(1);
409 let node_id = 1;
410 let mut service = VisualizationService::<(), Control<Info>, Event<Info>, (), (), _>::new(node_info);
411
412 let node2 = 2;
413 let node3 = 3;
414
415 let ctx = ServiceCtx { node_id, session: 0 };
416 service.on_shared_input(&ctx, 110, ServiceSharedInput::Connection(connected_event(node2)));
417 service.on_shared_input(&ctx, 110, ServiceSharedInput::Connection(connected_event(node3)));
418
419 assert_eq!(service.conns.len(), 2);
420
421 service.on_shared_input(&ctx, 210, ServiceSharedInput::Connection(disconnected_event(node3)));
422 assert_eq!(service.conns.len(), 1);
423
424 }
426
427 #[test]
428 fn collector_handle_snapshot_correct() {
429 let node_info = Info(1);
430 let node_id = 1;
431 let ctx = ServiceCtx { node_id, session: 0 };
432 let mut service = VisualizationService::<(), Control<Info>, Event<Info>, (), (), _>::new(node_info.clone());
433
434 let node2_info = Info(2);
435 let node2 = 2;
436
437 let snapshot = Message::Snapshot(node2, node2_info, vec![]);
438 let buf = bincode::serialize(&snapshot).expect("Should to bytes");
439 service.on_input(&ctx, 100, data_event(DataEvent::Recv(DATA_PORT, NetIncomingMeta::new(None, NODE_PING_TTL.into(), 0, true), buf)));
440
441 assert_eq!(service.network_nodes.len(), 1);
442
443 service.on_shared_input(&ctx, 100 + NODE_TIMEOUT_MS, ServiceSharedInput::Tick(0));
445 assert_eq!(service.network_nodes.len(), 0);
446 }
447}