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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
use std::collections::HashMap;
use std::fmt::Debug;
use std::net::IpAddr;
use std::net::Ipv4Addr;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use serde::de::DeserializeOwned;
use serde::Deserialize;
use serde::Serialize;
use serde_big_array::BigArray;
use tokio::net::UdpSocket;
use tokio::sync::broadcast;

mod interval;
use interval::Interval;
use tracing::trace;

mod notify;
pub use notify::Notify;

use crate::Id;
mod builder;
use builder::Port;

pub use builder::ChartBuilder;

pub mod to_vec;

use self::interval::Until;

#[derive(Serialize, Deserialize, Clone, Copy, Debug)]
pub struct DiscoveryMsg<const N: usize, T>
where
    T: Serialize + DeserializeOwned,
{
    header: u64,
    id: Id,
    #[serde(with = "BigArray")]
    msg: [T; N],
}

/// A chart entry representing a discovered node. The msg is an array of
/// ports or a custom struct if you used [`custom_msg`](ChartBuilder::custom_msg()).
///
/// You probably do not want to use one of the [iterator methods](iter) instead
#[derive(Debug, Clone)]
pub struct Entry<Msg: Debug + Clone> {
    pub ip: IpAddr,
    pub msg: Msg,
}

/// The chart keeping track of the discoverd nodes. That a node appears in the
/// chart is no guarentee that it is reachable at this moment.
#[derive(Debug, Clone)]
pub struct Chart<const N: usize, T: Debug + Clone + Serialize> {
    header: u64,
    service_id: Id,
    msg: [T; N],
    sock: Arc<UdpSocket>,
    interval: Interval,
    map: Arc<std::sync::Mutex<HashMap<Id, Entry<[T; N]>>>>,
    broadcast: broadcast::Sender<(Id, Entry<[T; N]>)>,
}

impl<const N: usize, T: Serialize + Debug + Clone> Chart<N, T> {
    fn insert(&self, id: Id, entry: Entry<[T; N]>) -> bool {
        let old_key = {
            let mut map = self.map.lock().unwrap();
            map.insert(id, entry.clone())
        };
        if old_key.is_none() {
            // errors if there are no active recievers which is
            // the default and not a problem
            let _ig_err = self.broadcast.send((id, entry));
            true
        } else {
            false
        }
    }

    #[tracing::instrument(skip(self, buf))]
    fn process_buf<'de>(&self, buf: &'de [u8], addr: SocketAddr) -> bool
    where
        T: Serialize + DeserializeOwned + Debug,
    {
        let DiscoveryMsg::<N, T> { header, id, msg } = bincode::deserialize(buf).unwrap();
        if header != self.header {
            return false;
        }
        if id == self.service_id {
            return false;
        }
        self.insert(id, Entry { ip: addr.ip(), msg })
    }
}

/// The array of ports set for this chart instance, set in `ChartBuilder::with_service_ports`.
impl<const N: usize> Chart<N, Port> {
    #[must_use]
    pub fn our_service_ports(&self) -> &[u16] {
        &self.msg
    }
}

/// The port set for this chart instance, set in `ChartBuilder::with_service_port`.
impl Chart<1, Port> {
    #[must_use]
    pub fn our_service_port(&self) -> u16 {
        self.msg[0]
    }
}

/// The msg struct for this chart instance, set in `ChartBuilder::custom_msg`.
impl<T: Debug + Clone + Serialize> Chart<1, T> {
    #[must_use]
    pub fn our_msg(&self) -> &T {
        &self.msg[0]
    }
}

impl<const N: usize, T: Debug + Clone + Serialize + DeserializeOwned> Chart<N, T> {
    /// Wait for new discoveries. Use one of the methods on the [notify object](notify::Notify)
    /// to _await_ a new discovery and get the data.
    /// # Examples
    /// ```rust
    /// # use std::error::Error;
    /// # use instance_chart::{discovery, ChartBuilder};
    /// #
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn Error>> {
    /// # let full_size = 4u16;
    /// # let handles: Vec<_> = (1..=full_size)
    /// #     .into_iter()
    /// #     .map(|id|
    /// #         ChartBuilder::new()
    /// #             .with_id(id.into())
    /// #             .with_service_port(8042+id)
    /// #             .with_discovery_port(8080)
    /// #             .local_discovery(true)
    /// #             .finish()
    /// #             .unwrap()
    /// #     )
    /// #     .map(discovery::maintain)
    /// #     .map(tokio::spawn)
    /// #     .collect();
    /// #
    /// let chart = ChartBuilder::new()
    ///     .with_id(1)
    ///     .with_service_port(8042)
    /// #   .with_discovery_port(8080)
    ///     .local_discovery(true)
    ///     .finish()?;
    /// let mut node_discoverd = chart.notify();
    /// let maintain = discovery::maintain(chart.clone());
    /// let _ = tokio::spawn(maintain); // maintain task will run forever
    ///
    /// while chart.size() < full_size as usize {
    ///     let new = node_discoverd.recv().await.unwrap();
    ///     println!("discoverd new node: {:?}", new);
    /// }
    ///
    /// #   Ok(())
    /// # }
    /// ```
    #[must_use]
    pub fn notify(&self) -> Notify<N, T> {
        Notify(self.broadcast.subscribe())
    }

    /// number of instances discoverd including self
    
    // lock poisoning happens only on crash in another thread, in which 
    // case panicing here is expected
    #[allow(clippy::missing_panics_doc)] 
    #[must_use]
    pub fn size(&self) -> usize {
        self.map.lock().unwrap().len() + 1
    }

    /// The id set for this chart instance
    #[must_use]
    pub fn our_id(&self) -> u64 {
        self.service_id
    }

    /// The port this instance is using for discovery
    #[allow(clippy::missing_panics_doc)] // socket is set during building
    #[must_use]
    pub fn discovery_port(&self) -> u16 {
        self.sock.local_addr().unwrap().port()
    }

    #[must_use]
    fn discovery_msg(&self) -> DiscoveryMsg<N, T> {
        DiscoveryMsg {
            header: self.header,
            id: self.service_id,
            msg: self.msg.clone(),
        }
    }

    #[must_use]
    fn discovery_buf(&self) -> Vec<u8> {
        let msg = self.discovery_msg();
        bincode::serialize(&msg).unwrap()
    }

    #[must_use]
    fn broadcast_soon(&mut self) -> bool {
        let next = self.interval.next();
        next.until() < Duration::from_millis(100)
    }
}

#[tracing::instrument]
pub(crate) async fn handle_incoming<const N: usize, T>(mut chart: Chart<N, T>)
where
    T: Debug + Clone + Serialize + DeserializeOwned,
{
    loop {
        let mut buf = [0; 1024];
        let (_len, addr) = chart.sock.recv_from(&mut buf).await.unwrap();
        trace!("got msg from: {addr:?}");
        let was_uncharted = chart.process_buf(&buf, addr);
        if was_uncharted && !chart.broadcast_soon() {
            chart
                .sock
                .send_to(&chart.discovery_buf(), addr)
                .await
                .unwrap();
        }
    }
}

#[tracing::instrument]
pub(crate) async fn broadcast_periodically<const N: usize, T>(
    mut chart: Chart<N, T>,
    period: Duration,
) where
    T: Debug + Serialize + DeserializeOwned + Clone,
{
    loop {
        chart.interval.sleep_till_next().await;
        trace!("sending discovery msg");
        broadcast(&chart.sock, chart.discovery_port(), &chart.discovery_buf()).await;
    }
}

#[tracing::instrument]
async fn broadcast(sock: &Arc<UdpSocket>, port: u16, msg: &[u8]) {
    let multiaddr = Ipv4Addr::from([224, 0, 0, 251]);
    let _len = sock
        .send_to(msg, (multiaddr, port))
        .await
        .unwrap_or_else(|e| panic!("broadcast failed with port: {port}, error: {e:?}"));
}