1extern crate proc_macro;
2
3use quote::quote;
4use syn::{DeriveInput, parse_macro_input};
5
6#[proc_macro_derive(AutoNet)]
7pub fn auto_net(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
8 let input = parse_macro_input!(input as DeriveInput);
9 let name = input.ident;
10
11 let expanded = quote! {
12 #[async_trait::async_trait]
13 impl Network for #name {
14 async fn heartbeat(&self) {
15 let clients = self.clients.clone();
16 let peer_addresses = self.peer_addresses.clone();
17 let signer = self.signer.clone();
18 let id = self.id.to_string();
19 let interval = self.heartbeat_interval;
20
21 tokio::spawn(async move {
22 loop {
23 for (peer_id, client) in &clients {
24 let msg = Message::ping(&id, peer_id, 0);
25 let result = {
26 let client = client.lock().await;
27 client.send(msg).await
28 };
29
30 if let Err(e) = result {
31 debug!("Heartbeat failed to {peer_id}: {e}");
32
33 if let Some(addr) = peer_addresses.get(peer_id) {
34 debug!("Attempting to reconnect to {peer_id} at {addr}...");
35
36 match Client::connect(addr, signer.clone()).await {
37 Ok(new_client) => {
38 debug!("Reconnected to {peer_id}");
39 let mut locked = client.lock().await;
40 *locked = new_client;
41 }
42 Err(err) => {
43 debug!("Failed to reconnect to {peer_id}: {err}");
44 }
45 }
46 } else {
47 debug!("No known address for {peer_id}, cannot reconnect.");
48 }
49 }
50 }
51
52 tokio::time::sleep(interval).await;
53 }
54 });
55 }
56
57 async fn broadcast(&self, payload: &str) -> anyhow::Result<()> {
58 let tasks = self.clients.iter().map(|(peer_id, client)| {
59 let mut msg = Message::broadcast(&self.id, payload, 0);
60 msg.to = peer_id.clone();
61 let client = client.clone();
62 async move {
63 let send_result = {
64 let client_guard = client.lock().await;
65 let client_ref = client_guard.clone();
66 client_ref
67 }.send(msg).await;
68
69 if let Err(e) = send_result {
70 debug!("Broadcast to {peer_id} failed: {e}");
71 } else {
72 debug!("Broadcast to {peer_id} succeeded");
73 }
74 }
75 });
76
77 futures::future::join_all(tasks).await;
78 Ok(())
79 }
80 }
81 };
82
83 proc_macro::TokenStream::from(expanded)
84}