refactor(iroh-dns-server): Move db ops into an actor and implement wr…

…ite batching (#2995)

## Description

Batch to make write go brrr....

- move all db interactions into an actor
- inside the actor, have two nested loops, inner loop keeps write txn
open
- txn is closed when a number of msgs have been written or some time has
elapsed, whatever comes first

## Breaking Changes

None

## Benches

This branch:

```
Benchmarking dns_server_writes/1000: Collecting 10 samples in estimated 9.4698 s
dns_server_writes/1000  time:   [171.28 ms 174.05 ms 176.58 ms]
                        thrpt:  [5.6631 Kelem/s 5.7455 Kelem/s 5.8384 Kelem/s]
                 change:
                        time:   [-6.8425% -3.3112% -0.0384%] (p = 0.10 > 0.05)
                        thrpt:  [+0.0384% +3.4246% +7.3451%]
                        No change in performance detected.
```

arqu/dns-bench:

```
Benchmarking dns_server_writes/1000: Warming up for 3.0000 s
dns_server_writes/1000  time:   [6.2530 s 6.3920 s 6.5351 s]
                        thrpt:  [153.02  elem/s 156.45  elem/s 159.92  elem/s]
                 change:
                        time:   [+3881.0% +3974.8% +4080.8%] (p = 0.00 < 0.05)
                        thrpt:  [-97.608% -97.546% -97.488%]
                        Performance has regressed.
```

## Downsides

The downside of write batching is that in case of a crash or hard
program termination we lose the last MAX_BATCH_TIME seconds and
MAX_BATCH_SIZE writes. I think that this is acceptable since discovery
information is republished.

---------

Co-authored-by: dignifiedquire <[email protected]>
Co-authored-by: Asmir Avdicevic <[email protected]>

iroh-dns-server/Cargo.toml

@@ -41,7 +41,7 @@ rustls-pemfile = { version = "2.1" }
 serde = { version = "1", features = ["derive"] }
 struct_iterable = "0.1.1"
 strum = { version = "0.26", features = ["derive"] }
-tokio = { version = "1.36.0", features = ["full"] }
+tokio = { version = "1", features = ["full"] }
 tokio-rustls = { version = "0.26", default-features = false, features = [
     "logging",
     "ring",
@@ -59,10 +59,15 @@ url = "2.5"
 z32 = "1.1.1"
 
 [dev-dependencies]
+criterion = "0.5.1"
 hickory-resolver = "=0.25.0-alpha.2"
 iroh = { version = "0.29.0", path = "../iroh" }
 iroh-test = { version = "0.29.0", path = "../iroh-test" }
 pkarr = { version = "2.2.0", features = ["rand"] }
 
+[[bench]]
+name = "write"
+harness = false
+
 [package.metadata.docs.rs]
 all-features = true

iroh-dns-server/benches/write.rs

@@ -0,0 +1,51 @@
+use anyhow::Result;
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
+use iroh::{discovery::pkarr::PkarrRelayClient, dns::node_info::NodeInfo, key::SecretKey};
+use iroh_dns_server::{config::Config, server::Server, ZoneStore};
+use tokio::runtime::Runtime;
+
+const LOCALHOST_PKARR: &str = "http://localhost:8080/pkarr";
+
+async fn start_dns_server(config: Config) -> Result<Server> {
+    let store = ZoneStore::persistent(Config::signed_packet_store_path()?)?;
+    Server::spawn(config, store).await
+}
+
+fn benchmark_dns_server(c: &mut Criterion) {
+    let mut group = c.benchmark_group("dns_server_writes");
+    group.sample_size(10);
+    for iters in [10_u64, 100_u64, 250_u64, 1000_u64].iter() {
+        group.throughput(Throughput::Elements(*iters));
+        group.bench_with_input(BenchmarkId::from_parameter(iters), iters, |b, &iters| {
+            b.iter(|| {
+                let rt = Runtime::new().unwrap();
+                rt.block_on(async move {
+                    let config = Config::load("./config.dev.toml").await.unwrap();
+                    let server = start_dns_server(config).await.unwrap();
+
+                    let secret_key = SecretKey::generate();
+                    let node_id = secret_key.public();
+
+                    let pkarr_relay = LOCALHOST_PKARR.parse().expect("valid url");
+                    let relay_url = Some("http://localhost:8080".parse().unwrap());
+                    let pkarr = PkarrRelayClient::new(pkarr_relay);
+                    let node_info = NodeInfo::new(node_id, relay_url, Default::default());
+                    let signed_packet = node_info.to_pkarr_signed_packet(&secret_key, 30).unwrap();
+
+                    let start = std::time::Instant::now();
+                    for _ in 0..iters {
+                        pkarr.publish(&signed_packet).await.unwrap();
+                    }
+                    let duration = start.elapsed();
+
+                    server.shutdown().await.unwrap();
+
+                    duration
+                })
+            });
+        });
+    }
+}
+
+criterion_group!(benches, benchmark_dns_server);
+criterion_main!(benches);

iroh-dns-server/src/lib.rs

@@ -11,6 +11,9 @@ pub mod state;
 mod store;
 mod util;
 
+// Re-export to be able to construct your own dns-server
+pub use store::ZoneStore;
+
 #[cfg(test)]
 mod tests {
     use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};

iroh-dns-server/src/store/signed_packets.rs

@@ -1,20 +1,145 @@
-use std::{path::Path, sync::Arc};
+use std::{path::Path, result, time::Duration};
 
 use anyhow::{Context, Result};
 use iroh_metrics::inc;
 use pkarr::SignedPacket;
 use redb::{backends::InMemoryBackend, Database, ReadableTable, TableDefinition};
+use tokio::sync::{mpsc, oneshot};
+use tokio_util::sync::CancellationToken;
 use tracing::info;
 
 use crate::{metrics::Metrics, util::PublicKeyBytes};
 
 pub type SignedPacketsKey = [u8; 32];
 const SIGNED_PACKETS_TABLE: TableDefinition<&SignedPacketsKey, &[u8]> =
     TableDefinition::new("signed-packets-1");
+const MAX_BATCH_SIZE: usize = 1024 * 64;
+const MAX_BATCH_TIME: Duration = Duration::from_secs(1);
 
 #[derive(Debug)]
 pub struct SignedPacketStore {
-    db: Arc<Database>,
+    send: mpsc::Sender<Message>,
+    cancel: CancellationToken,
+    thread: Option<std::thread::JoinHandle<()>>,
+}
+
+impl Drop for SignedPacketStore {
+    fn drop(&mut self) {
+        // cancel the actor
+        self.cancel.cancel();
+        // join the thread. This is important so that Drop implementations that
+        // are called from the actor thread can complete before we return.
+        if let Some(thread) = self.thread.take() {
+            let _ = thread.join();
+        }
+    }
+}
+
+enum Message {
+    Upsert {
+        packet: SignedPacket,
+        res: oneshot::Sender<bool>,
+    },
+    Get {
+        key: PublicKeyBytes,
+        res: oneshot::Sender<Option<SignedPacket>>,
+    },
+    Remove {
+        key: PublicKeyBytes,
+        res: oneshot::Sender<bool>,
+    },
+}
+
+struct Actor {
+    db: Database,
+    recv: mpsc::Receiver<Message>,
+    cancel: CancellationToken,
+    max_batch_size: usize,
+    max_batch_time: Duration,
+}
+
+impl Actor {
+    async fn run(mut self) {
+        match self.run0().await {
+            Ok(()) => {}
+            Err(e) => {
+                self.cancel.cancel();
+                tracing::error!("packet store actor failed: {:?}", e);
+            }
+        }
+    }
+
+    async fn run0(&mut self) -> anyhow::Result<()> {
+        loop {
+            let transaction = self.db.begin_write()?;
+            let mut tables = Tables::new(&transaction)?;
+            let timeout = tokio::time::sleep(self.max_batch_time);
+            tokio::pin!(timeout);
+            for _ in 0..self.max_batch_size {
+                tokio::select! {
+                    _ = self.cancel.cancelled() => {
+                        drop(tables);
+                        transaction.commit()?;
+                        return Ok(());
+                    }
+                    _ = &mut timeout => break,
+                    Some(msg) = self.recv.recv() => {
+                        match msg {
+                            Message::Get { key, res } => {
+                                let packet = get_packet(&tables.signed_packets, &key)?;
+                                res.send(packet).ok();
+                            }
+                            Message::Upsert { packet, res } => {
+                                let key = PublicKeyBytes::from_signed_packet(&packet);
+                                let mut replaced = false;
+                                if let Some(existing) = get_packet(&tables.signed_packets, &key)? {
+                                    if existing.more_recent_than(&packet) {
+                                        res.send(false).ok();
+                                        continue;
+                                    } else {
+                                        replaced = true;
+                                    }
+                                }
+                                let value = packet.as_bytes();
+                                tables.signed_packets.insert(key.as_bytes(), &value[..])?;
+                                if replaced {
+                                    inc!(Metrics, store_packets_updated);
+                                } else {
+                                    inc!(Metrics, store_packets_inserted);
+                                }
+                                res.send(true).ok();
+                            }
+                            Message::Remove { key, res } => {
+                                let updated =
+                                    tables.signed_packets.remove(key.as_bytes())?.is_some()
+                                ;
+                                if updated {
+                                    inc!(Metrics, store_packets_removed);
+                                }
+                                res.send(updated).ok();
+                            }
+                        }
+                    }
+                }
+            }
+            drop(tables);
+            transaction.commit()?;
+        }
+    }
+}
+
+/// A struct similar to [`redb::Table`] but for all tables that make up the
+/// signed packet store.
+pub(super) struct Tables<'a> {
+    pub signed_packets: redb::Table<'a, &'static SignedPacketsKey, &'static [u8]>,
+}
+
+impl<'txn> Tables<'txn> {
+    pub fn new(tx: &'txn redb::WriteTransaction) -> result::Result<Self, redb::TableError> {
+        Ok(Self {
+            signed_packets: tx.open_table(SIGNED_PACKETS_TABLE)?,
+        })
+    }
 }
 
 impl SignedPacketStore {
@@ -42,73 +167,53 @@ impl SignedPacketStore {
     }
 
     pub fn open(db: Database) -> Result<Self> {
+        // create tables
         let write_tx = db.begin_write()?;
-        {
-            let _table = write_tx.open_table(SIGNED_PACKETS_TABLE)?;
-        }
+        let _ = Tables::new(&write_tx)?;
         write_tx.commit()?;
-        Ok(Self { db: Arc::new(db) })
+        let (send, recv) = mpsc::channel(1024);
+        let cancel = CancellationToken::new();
+        let cancel2 = cancel.clone();
+        let actor = Actor {
+            db,
+            recv,
+            cancel: cancel2,
+            max_batch_size: MAX_BATCH_SIZE,
+            max_batch_time: MAX_BATCH_TIME,
+        };
+        // start an io thread and donate it to the tokio runtime so we can do blocking IO
+        // inside the thread despite being in a tokio runtime
+        let handle = tokio::runtime::Handle::try_current()?;
+        let thread = std::thread::Builder::new()
+            .name("packet-store-actor".into())
+            .spawn(move || {
+                handle.block_on(actor.run());
+            })?;
+        Ok(Self {
+            send,
+            cancel,
+            thread: Some(thread),
+        })
     }
 
     pub async fn upsert(&self, packet: SignedPacket) -> Result<bool> {
-        let key = PublicKeyBytes::from_signed_packet(&packet);
-        let db = self.db.clone();
-        tokio::task::spawn_blocking(move || {
-            let tx = db.begin_write()?;
-            let mut replaced = false;
-            {
-                let mut table = tx.open_table(SIGNED_PACKETS_TABLE)?;
-                if let Some(existing) = get_packet(&table, &key)? {
-                    if existing.more_recent_than(&packet) {
-                        return Ok(false);
-                    } else {
-                        replaced = true;
-                    }
-                }
-                let value = packet.as_bytes();
-                table.insert(key.as_bytes(), &value[..])?;
-            }
-            tx.commit()?;
-            if replaced {
-                inc!(Metrics, store_packets_updated);
-            } else {
-                inc!(Metrics, store_packets_inserted);
-            }
-            Ok(true)
-        })
-        .await?
+        let (tx, rx) = oneshot::channel();
+        self.send.send(Message::Upsert { packet, res: tx }).await?;
+        Ok(rx.await?)
     }
 
     pub async fn get(&self, key: &PublicKeyBytes) -> Result<Option<SignedPacket>> {
-        let db = self.db.clone();
-        let key = *key;
-        let res = tokio::task::spawn_blocking(move || {
-            let tx = db.begin_read()?;
-            let table = tx.open_table(SIGNED_PACKETS_TABLE)?;
-            get_packet(&table, &key)
-        })
-        .await??;
-        Ok(res)
+        let (tx, rx) = oneshot::channel();
+        self.send.send(Message::Get { key: *key, res: tx }).await?;
+        Ok(rx.await?)
     }
 
     pub async fn remove(&self, key: &PublicKeyBytes) -> Result<bool> {
-        let db = self.db.clone();
-        let key = *key;
-        tokio::task::spawn_blocking(move || {
-            let tx = db.begin_write()?;
-            let updated = {
-                let mut table = tx.open_table(SIGNED_PACKETS_TABLE)?;
-                let did_remove = table.remove(key.as_bytes())?.is_some();
-                #[allow(clippy::let_and_return)]
-                did_remove
-            };
-            tx.commit()?;
-            if updated {
-                inc!(Metrics, store_packets_removed)
-            }
-            Ok(updated)
-        })
-        .await?
+        let (tx, rx) = oneshot::channel();
+        self.send
+            .send(Message::Remove { key: *key, res: tx })
+            .await?;
+        Ok(rx.await?)
     }
 }