risotto_lib/state_store/
memory.rs

1use serde::{Deserialize, Serialize};
2use std::collections::{HashMap, HashSet};
3use std::net::IpAddr;
4
5use crate::state::{RouterPeerUpdate, TimedPrefix};
6use crate::state_store::store::StateStore;
7use crate::update::Update;
8
9#[derive(Serialize, Deserialize)]
10pub struct MemoryStore {
11    routers: HashMap<IpAddr, Router>,
12}
13
14impl MemoryStore {
15    pub fn new() -> MemoryStore {
16        MemoryStore {
17            routers: HashMap::new(),
18        }
19    }
20
21    fn _get_router(&mut self, router_addr: &IpAddr) -> &mut Router {
22        let router = self.routers.entry(*router_addr).or_insert(Router::new());
23        router
24    }
25}
26
27impl StateStore for MemoryStore {
28    fn get_all(&self) -> Vec<RouterPeerUpdate> {
29        let mut res = Vec::new();
30        for (router_addr, router) in &self.routers {
31            for (peer_addr, peer) in &router.peers {
32                for update in &peer.updates {
33                    res.push((*router_addr, *peer_addr, update.clone()));
34                }
35            }
36        }
37        res
38    }
39
40    fn get_updates_by_peer(&self, router_addr: &IpAddr, peer_addr: &IpAddr) -> Vec<TimedPrefix> {
41        let router_binding = Router::new();
42        let router = self.routers.get(router_addr).unwrap_or(&router_binding);
43        let peer = match router.peers.get(&*peer_addr) {
44            Some(peer) => peer,
45            None => return Vec::new(),
46        };
47
48        peer.updates.iter().cloned().collect()
49    }
50
51    fn remove_peer(&mut self, router_addr: &IpAddr, peer_addr: &IpAddr) {
52        let router = self._get_router(router_addr);
53        router.remove_peer(peer_addr);
54    }
55
56    fn update(&mut self, router_addr: &IpAddr, peer_addr: &IpAddr, update: &Update) -> bool {
57        let router = self._get_router(router_addr);
58        router.update(peer_addr, update)
59    }
60}
61
62#[derive(Serialize, Deserialize, Clone)]
63pub struct Peer {
64    pub peer_addr: IpAddr,
65    pub updates: HashSet<TimedPrefix>,
66}
67
68#[derive(Serialize, Deserialize)]
69pub struct Router {
70    peers: HashMap<IpAddr, Peer>,
71}
72
73impl Router {
74    fn new() -> Router {
75        Router {
76            peers: HashMap::new(),
77        }
78    }
79
80    fn add_peer(&mut self, peer_addr: &IpAddr) {
81        self.peers.entry(*peer_addr).or_insert_with(|| Peer {
82            peer_addr: peer_addr.clone(),
83            updates: HashSet::new(),
84        });
85    }
86
87    fn remove_peer(&mut self, peer_addr: &IpAddr) {
88        self.peers.remove(&peer_addr);
89    }
90
91    fn update(&mut self, peer_addr: &IpAddr, update: &Update) -> bool {
92        self.add_peer(peer_addr);
93        let peer = self.peers.get_mut(&peer_addr).unwrap();
94
95        let now: i64 = chrono::Utc::now().timestamp_millis();
96        let timed_prefix = TimedPrefix {
97            prefix_addr: update.prefix_addr,
98            prefix_len: update.prefix_len,
99            is_post_policy: update.is_post_policy,
100            is_adj_rib_out: update.is_adj_rib_out,
101            timestamp: now,
102        };
103
104        // Will emit the update only if (1) announced + not present or (2) withdrawn + present
105        // Which is a XOR operation
106        let emit = update.announced ^ peer.updates.contains(&timed_prefix);
107
108        if update.announced {
109            // Announced prefix: add the update or overwrite it if present
110            peer.updates.replace(timed_prefix);
111        } else {
112            // Withdrawn prefix: remove the update if present
113            peer.updates.remove(&timed_prefix);
114        }
115        emit
116    }
117}
risotto_lib/state_store/memory.rs

risotto_lib/state_store/
memory.rs
