Skip to main content

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 Default for Dht {
90    fn default() -> Self {
91        Self::new()
92    }
93}
94
95impl Actor<anyhow::Error> for DhtActor {
96    async fn run(&mut self) -> Result<()> {
97        loop {
98            tokio::select! {
99                Ok(action) = self.rx.recv_async() => {
100                    action(self).await;
101                }
102                else => break Ok(()),
103            }
104        }
105    }
106}
107
108impl DhtActor {
109    pub async fn get(
110        &mut self,
111        pub_key: VerifyingKey,
112        salt: Option<Vec<u8>>,
113        more_recent_than: Option<i64>,
114        timeout: Duration,
115    ) -> Result<Vec<MutableItem>> {
116        if self.dht.is_none() {
117            self.reset().await?;
118        }
119
120        let dht = self.dht.as_mut().context("DHT not initialized")?;
121        Ok(tokio::time::timeout(
122            timeout,
123            dht.get_mutable(pub_key.as_bytes(), salt.as_deref(), more_recent_than)
124                .collect::<Vec<_>>(),
125        )
126        .await?)
127    }
128
129    pub async fn put_mutable(
130        &mut self,
131        signing_key: SigningKey,
132        pub_key: VerifyingKey,
133        salt: Option<Vec<u8>>,
134        data: Vec<u8>,
135        retry_count: Option<usize>,
136        timeout: Duration,
137    ) -> Result<()> {
138        if self.dht.is_none() {
139            self.reset().await?;
140        }
141
142        for i in 0..retry_count.unwrap_or(RETRY_DEFAULT) {
143            let dht = self.dht.as_mut().context("DHT not initialized")?;
144
145            let most_recent_result = tokio::time::timeout(
146                timeout,
147                dht.get_mutable_most_recent(pub_key.as_bytes(), salt.as_deref()),
148            )
149            .await?;
150
151            let item = if let Some(mut_item) = most_recent_result {
152                MutableItem::new(
153                    signing_key.clone(),
154                    &data,
155                    mut_item.seq() + 1,
156                    salt.as_deref(),
157                )
158            } else {
159                MutableItem::new(signing_key.clone(), &data, 0, salt.as_deref())
160            };
161
162            let put_result = match tokio::time::timeout(
163                Duration::from_secs(10),
164                dht.put_mutable(item.clone(), Some(item.seq())),
165            )
166            .await
167            {
168                Ok(result) => result.ok(),
169                Err(_) => None,
170            };
171
172            if put_result.is_some() {
173                break;
174            } else if i == retry_count.unwrap_or(RETRY_DEFAULT) - 1 {
175                bail!("failed to publish record")
176            }
177
178            self.reset().await?;
179
180            tokio::time::sleep(Duration::from_millis(rand::random::<u64>() % 2000)).await;
181        }
182        Ok(())
183    }
184
185    async fn reset(&mut self) -> Result<()> {
186        self.dht = Some(mainline::Dht::builder().build()?.as_async());
187        Ok(())
188    }
189}