risotto_lib/state_store/
memory.rs

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

risotto_lib/state_store/
memory.rs
