distributed_topic_tracker/
dht.rs

1//! Mainline BitTorrent DHT client for mutable record operations.
2//!
3//! Provides async interface for DHT get/put operations with automatic
4//! retry logic and connection management.
5
6use std::time::Duration;
7
8use actor_helper::{Action, Actor, Handle, Receiver, act};
9use anyhow::{Context, Result, bail};
10use ed25519_dalek::VerifyingKey;
11use futures_lite::StreamExt;
12use mainline::{MutableItem, SigningKey};
13
14const RETRY_DEFAULT: usize = 3;
15
16/// DHT client wrapper with actor-based concurrency.
17///
18/// Manages connections to the mainline DHT and handles
19/// mutable record get/put operations with automatic retries.
20#[derive(Debug, Clone)]
21pub struct Dht {
22    api: Handle<DhtActor, anyhow::Error>,
23}
24
25#[derive(Debug)]
26struct DhtActor {
27    rx: Receiver<Action<Self>>,
28    dht: Option<mainline::async_dht::AsyncDht>,
29}
30
31impl Dht {
32    /// Create a new DHT client.
33    ///
34    /// Spawns a background actor for handling DHT operations.
35    pub fn new() -> Self {
36        let (api, rx) = Handle::channel();
37
38        tokio::spawn(async move {
39            let mut actor = DhtActor { rx, dht: None };
40            let _ = actor.run().await;
41        });
42
43        Self { api }
44    }
45
46    /// Retrieve mutable records from the DHT.
47    ///
48    /// # Arguments
49    ///
50    /// * `pub_key` - Ed25519 public key for the record
51    /// * `salt` - Optional salt for record lookup
52    /// * `more_recent_than` - Sequence number filter (get records newer than this)
53    /// * `timeout` - Maximum time to wait for results
54    pub async fn get(
55        &self,
56        pub_key: VerifyingKey,
57        salt: Option<Vec<u8>>,
58        more_recent_than: Option<i64>,
59        timeout: Duration,
60    ) -> Result<Vec<MutableItem>> {
61        self.api
62            .call(act!(actor => actor.get(pub_key, salt, more_recent_than, timeout)))
63            .await
64    }
65
66    /// Publish a mutable record to the DHT.
67    ///
68    /// # Arguments
69    ///
70    /// * `signing_key` - Ed25519 secret key for signing
71    /// * `pub_key` - Ed25519 public key (used for routing)
72    /// * `salt` - Optional salt for record slot
73    /// * `data` - Record value to publish
74    /// * `retry_count` - Number of retry attempts (default: 3)
75    /// * `timeout` - Per-request timeout
76    pub async fn put_mutable(
77        &self,
78        signing_key: SigningKey,
79        pub_key: VerifyingKey,
80        salt: Option<Vec<u8>>,
81        data: Vec<u8>,
82        retry_count: Option<usize>,
83        timeout: Duration,
84    ) -> Result<()> {
85        self.api.call(act!(actor => actor.put_mutable(signing_key, pub_key, salt, data, retry_count, timeout))).await
86    }
87}
88
89impl Actor<anyhow::Error> for DhtActor {
90    async fn run(&mut self) -> Result<()> {
91        loop {
92            tokio::select! {
93                Ok(action) = self.rx.recv_async() => {
94                    action(self).await;
95                }
96                _ = tokio::signal::ctrl_c() => {
97                    break;
98                }
99            }
100        }
101        Err(anyhow::anyhow!("actor stopped"))
102    }
103}
104
105impl DhtActor {
106    pub async fn get(
107        &mut self,
108        pub_key: VerifyingKey,
109        salt: Option<Vec<u8>>,
110        more_recent_than: Option<i64>,
111        timeout: Duration,
112    ) -> Result<Vec<MutableItem>> {
113        if self.dht.is_none() {
114            self.reset().await?;
115        }
116
117        let dht = self.dht.as_mut().context("DHT not initialized")?;
118        Ok(tokio::time::timeout(
119            timeout,
120            dht.get_mutable(pub_key.as_bytes(), salt.as_deref(), more_recent_than)
121                .collect::<Vec<_>>(),
122        )
123        .await?)
124    }
125
126    pub async fn put_mutable(
127        &mut self,
128        signing_key: SigningKey,
129        pub_key: VerifyingKey,
130        salt: Option<Vec<u8>>,
131        data: Vec<u8>,
132        retry_count: Option<usize>,
133        timeout: Duration,
134    ) -> Result<()> {
135        if self.dht.is_none() {
136            self.reset().await?;
137        }
138
139        for i in 0..retry_count.unwrap_or(RETRY_DEFAULT) {
140            let dht = self.dht.as_mut().context("DHT not initialized")?;
141
142            let most_recent_result = tokio::time::timeout(
143                timeout,
144                dht.get_mutable_most_recent(pub_key.as_bytes(), salt.as_deref()),
145            )
146            .await?;
147
148            let item = if let Some(mut_item) = most_recent_result {
149                MutableItem::new(
150                    signing_key.clone(),
151                    &data,
152                    mut_item.seq() + 1,
153                    salt.as_deref(),
154                )
155            } else {
156                MutableItem::new(signing_key.clone(), &data, 0, salt.as_deref())
157            };
158
159            let put_result = match tokio::time::timeout(
160                Duration::from_secs(10),
161                dht.put_mutable(item.clone(), Some(item.seq())),
162            )
163            .await
164            {
165                Ok(result) => result.ok(),
166                Err(_) => None,
167            };
168
169            if put_result.is_some() {
170                break;
171            } else if i == retry_count.unwrap_or(RETRY_DEFAULT) - 1 {
172                bail!("failed to publish record")
173            }
174
175            self.reset().await?;
176
177            tokio::time::sleep(Duration::from_millis(rand::random::<u64>() % 2000)).await;
178        }
179        Ok(())
180    }
181
182    async fn reset(&mut self) -> Result<()> {
183        self.dht = Some(mainline::Dht::builder().build()?.as_async());
184        Ok(())
185    }
186}