rsky-relay: fixes/perf

Author: utkarshgupta137

rsky-relay/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2024"
 [dependencies]
 # external
 bstr = "1"
+bytes = "1"
 chrono = { version = "0.4", features = ["serde"] }
 ciborium = "0.2"
 cid = { version = "0.11", features = ["serde-codec"] }
@@ -16,6 +17,7 @@ hashbrown = "0.15"
 http = "1"
 httparse = "1.10"
 k256 = "0.13"
+lru = "0.14"
 magnetic = "2"
 mimalloc = "0.1"
 mio = { version = "1", features = ["os-ext", "os-poll"] }

rsky-relay/src/crawler/connection.rs

@@ -17,7 +17,7 @@ pub enum ConnectionError {
     #[error("tungstenite error: {0}")]
     Tungstenite(#[from] tungstenite::Error),
     #[error("thingbuf error: {0}")]
-    Thingbuf(#[from] mpsc::errors::TrySendError),
+    Thingbuf(#[from] mpsc::errors::Closed),
 }
 
 pub struct Connection {
@@ -66,8 +66,14 @@ impl Connection {
         Ok(())
     }
 
-    pub fn poll(&mut self) -> Result<(), ConnectionError> {
+    // false: not polled
+    // true: polled
+    pub fn poll(&mut self) -> Result<bool, ConnectionError> {
         for _ in 0..1024 {
+            if self.message_tx.remaining() < 16 {
+                return Ok(false);
+            }
+
             let msg = match self.client.read() {
                 Ok(msg) => msg,
                 Err(tungstenite::Error::Io(e)) if e.kind() == io::ErrorKind::WouldBlock => {
@@ -91,10 +97,10 @@ impl Connection {
                 }
             };
 
-            let mut slot = self.message_tx.try_send_ref()?;
+            let mut slot = self.message_tx.send_ref()?;
             slot.data = bytes;
             slot.hostname.clone_from(&self.hostname);
         }
-        Ok(())
+        Ok(true)
     }
 }

rsky-relay/src/crawler/worker.rs

@@ -23,6 +23,7 @@ pub enum WorkerError {
 pub struct Worker {
     id: usize,
     connections: Vec<Option<Connection>>,
+    next_idx: usize,
     message_tx: MessageSender,
     command_rx: CommandReceiver,
     poll: Poll,
@@ -35,7 +36,7 @@ impl Worker {
     ) -> Result<Self, WorkerError> {
         let poll = Poll::new()?;
         let events = Events::with_capacity(1024);
-        Ok(Self { id, connections: Vec::new(), message_tx, command_rx, poll, events })
+        Ok(Self { id, connections: Vec::new(), next_idx: 0, message_tx, command_rx, poll, events })
     }
 
     #[expect(clippy::unnecessary_wraps)]
@@ -101,24 +102,29 @@ impl Worker {
                 }
             }
 
+            if self.message_tx.remaining() < 16 {
+                break;
+            }
+
             let mut events = std::mem::replace(&mut self.events, Events::with_capacity(0));
-            for _ in 0..32 {
+            'outer: for _ in 0..32 {
                 #[expect(clippy::expect_used)]
                 self.poll
                     .poll(&mut events, Some(Duration::from_millis(1)))
                     .expect("failed to poll");
                 for ev in &events {
                     if !self.poll(ev.token().0) {
-                        return false;
+                        break 'outer;
                     }
                 }
             }
             self.events = events;
         }
 
-        for idx in 0..self.connections.len() {
-            if !self.poll(idx) {
-                return false;
+        for _ in 0..self.connections.len() {
+            self.next_idx = (self.next_idx + 1) % self.connections.len();
+            if !self.poll(self.next_idx) {
+                break;
             }
         }
 
@@ -127,16 +133,21 @@ impl Worker {
 
     fn poll(&mut self, idx: usize) -> bool {
         if let Some(conn) = &mut self.connections[idx] {
-            if let Err(err) = conn.poll() {
-                tracing::info!("[{}] disconnected: {err}", conn.hostname);
-                #[expect(clippy::expect_used)]
-                self.poll
-                    .registry()
-                    .deregister(&mut SourceFd(&conn.as_raw_fd()))
-                    .expect("failed to deregister");
-                self.connections[idx] = None;
+            match conn.poll() {
+                Ok(true) => {}
+                Ok(false) => return false,
+                Err(err) => {
+                    tracing::info!("[{}] disconnected: {err}", conn.hostname);
+                    #[expect(clippy::expect_used)]
+                    self.poll
+                        .registry()
+                        .deregister(&mut SourceFd(&conn.as_raw_fd()))
+                        .expect("failed to deregister");
+                    self.connections[idx] = None;
+                }
             }
         }
+
         true
     }
 }

rsky-relay/src/main.rs

@@ -39,11 +39,11 @@ pub async fn main() -> Result<()> {
     let (request_crawl_tx, request_crawl_rx) = rtrb::RingBuffer::new(CAPACITY2);
     let (subscribe_repos_tx, subscribe_repos_rx) = rtrb::RingBuffer::new(CAPACITY2);
     let validator = ValidatorManager::new(message_rx)?;
+    tokio::spawn(validator.run());
     let crawler = CrawlerManager::new(WORKERS, &message_tx, request_crawl_rx)?;
     let publisher = PublisherManager::new(WORKERS, subscribe_repos_rx)?;
     let server = Server::new(request_crawl_tx, subscribe_repos_tx)?;
     thread::scope(move |s| {
-        tokio::spawn(validator.run());
         thread::Builder::new().name("rsky-crawl".into()).spawn_scoped(s, move || crawler.run())?;
         thread::Builder::new().name("rsky-pub".into()).spawn_scoped(s, move || publisher.run())?;
         thread::Builder::new().name("rsky-server".into()).spawn_scoped(s, move || server.run())?;

rsky-relay/src/publisher/connection.rs

@@ -5,13 +5,17 @@ use std::os::fd::{AsRawFd, RawFd};
 use sled::Tree;
 use thiserror::Error;
 use tungstenite::handshake::server::NoCallback;
+use tungstenite::protocol::CloseFrame;
+use tungstenite::protocol::frame::coding::CloseCode;
 use tungstenite::stream::MaybeTlsStream;
-use tungstenite::{Bytes, HandshakeError, Message, ServerHandshake, WebSocket};
+use tungstenite::{Bytes, HandshakeError, Message, ServerHandshake, Utf8Bytes, WebSocket};
 
 use crate::types::Cursor;
 
-const OUTDATED: &[u8] = b"\xa2ate#infobop\x01\xa2dnamenOutdatedCursorgmessagex8Requested cursor exceeded limit. Possibly missing events.";
-const FUTURE: &[u8] = b"\xa1bop \xa2eerrorlFutureCursorgmessageuCursor in the future.";
+const OUTDATED_MSG: &[u8] = b"\xa2ate#infobop\x01\xa2dnamenOutdatedCursorgmessagex8Requested cursor exceeded limit. Possibly missing events.";
+const FUTURE_MSG: &[u8] = b"\xa1bop \xa2eerrorlFutureCursorgmessageuCursor in the future.";
+const FUTURE_CLOSE: CloseFrame =
+    CloseFrame { code: CloseCode::Policy, reason: Utf8Bytes::from_static("FutureCursor") };
 
 #[derive(Debug, Error)]
 pub enum ConnectionError {
@@ -59,12 +63,14 @@ impl Connection {
         Ok(Self { addr, client, cursor })
     }
 
-    pub fn close(&mut self) -> Result<(), ConnectionError> {
-        self.client.close(None)?;
+    pub fn close(&mut self, code: CloseFrame) -> Result<(), ConnectionError> {
+        self.client.close(Some(code))?;
         self.client.flush()?;
         Ok(())
     }
 
+    /// false: not sent
+    /// true: sent
     pub fn send(&mut self, mut seq: Cursor, data: Bytes) -> Result<bool, ConnectionError> {
         if self.cursor != seq {
             return Ok(false);
@@ -83,22 +89,27 @@ impl Connection {
         }
     }
 
-    pub fn poll(&mut self, mut seq: Cursor, firehose: &Tree) -> Result<(), ConnectionError> {
-        if self.cursor.get() > seq.get() {
-            self.send(self.cursor, Bytes::from_static(FUTURE))?;
-            return self.close();
+    /// false: closed
+    /// true: not closed
+    pub fn poll(&mut self, mut seq: Cursor, firehose: &Tree) -> Result<bool, ConnectionError> {
+        if self.cursor.get() != 0 && self.cursor.get() > seq.get() + 1 {
+            self.send(self.cursor, Bytes::from_static(FUTURE_MSG))?;
+            self.close(FUTURE_CLOSE)?;
+            return Ok(false);
         }
         for msg in firehose.range(self.cursor..=seq) {
             let (k, v) = msg?;
             seq = k.into();
             if self.cursor != seq {
-                self.send(self.cursor, Bytes::from_static(OUTDATED))?;
+                if self.cursor.get() != 0 {
+                    self.send(self.cursor, Bytes::from_static(OUTDATED_MSG))?;
+                }
                 self.cursor = seq;
             }
             if !self.send(seq, Bytes::from_owner(v).slice(8..))? {
                 break;
             }
         }
-        Ok(())
+        Ok(true)
     }
 }

rsky-relay/src/publisher/worker.rs

@@ -2,17 +2,22 @@ use std::os::fd::AsRawFd;
 use std::time::Duration;
 use std::{io, thread};
 
+use bytes::Bytes;
 use mio::unix::SourceFd;
 use mio::{Events, Interest, Poll, Token};
 use sled::Tree;
 use thiserror::Error;
-use tungstenite::Bytes;
+use tungstenite::Utf8Bytes;
+use tungstenite::protocol::CloseFrame;
+use tungstenite::protocol::frame::coding::CloseCode;
 
 use crate::publisher::connection::{Connection, ConnectionError};
 use crate::publisher::types::{Command, CommandReceiver};
 use crate::types::{Cursor, DB};
 
 const INTEREST: Interest = Interest::WRITABLE;
+const SHUTDOWN_FRAME: CloseFrame =
+    CloseFrame { code: CloseCode::Restart, reason: Utf8Bytes::from_static("RelayRestart") };
 
 #[derive(Debug, Error)]
 pub enum WorkerError {
@@ -27,6 +32,7 @@ pub enum WorkerError {
 pub struct Worker {
     id: usize,
     connections: Vec<Option<Connection>>,
+    next_idx: usize,
     command_rx: CommandReceiver,
     firehose: Tree,
     poll: Poll,
@@ -38,7 +44,7 @@ impl Worker {
         let firehose = DB.open_tree("firehose")?;
         let poll = Poll::new()?;
         let events = Events::with_capacity(1024);
-        Ok(Self { id, connections: Vec::new(), command_rx, firehose, poll, events })
+        Ok(Self { id, connections: Vec::new(), next_idx: 0, command_rx, firehose, poll, events })
     }
 
     pub fn run(mut self) -> Result<(), WorkerError> {
@@ -53,7 +59,7 @@ impl Worker {
 
     pub fn shutdown(mut self) {
         for conn in self.connections.iter_mut().filter_map(|x| x.as_mut()) {
-            if let Err(err) = conn.close() {
+            if let Err(err) = conn.close(SHUTDOWN_FRAME) {
                 tracing::warn!("publisher conn close error: {err}");
             }
         }
@@ -112,23 +118,24 @@ impl Worker {
             }
 
             let mut events = std::mem::replace(&mut self.events, Events::with_capacity(0));
-            for _ in 0..32 {
+            'outer: for _ in 0..32 {
                 #[expect(clippy::expect_used)]
                 self.poll
                     .poll(&mut events, Some(Duration::from_millis(1)))
                     .expect("failed to poll");
                 for ev in &events {
                     if !self.poll(*seq, ev.token().0) {
-                        return Ok(false);
+                        break 'outer;
                     }
                 }
             }
             self.events = events;
         }
 
-        for idx in 0..self.connections.len() {
-            if !self.poll(*seq, idx) {
-                return Ok(false);
+        for _ in 0..self.connections.len() {
+            self.next_idx = (self.next_idx + 1) % self.connections.len();
+            if !self.poll(*seq, self.next_idx) {
+                break;
             }
         }
 
@@ -154,16 +161,23 @@ impl Worker {
 
     fn poll(&mut self, seq: Cursor, idx: usize) -> bool {
         if let Some(conn) = &mut self.connections[idx] {
-            if let Err(err) = conn.poll(seq, &self.firehose) {
-                tracing::info!("[{}] disconnected: {err}", conn.addr);
-                #[expect(clippy::expect_used)]
-                self.poll
-                    .registry()
-                    .deregister(&mut SourceFd(&conn.as_raw_fd()))
-                    .expect("failed to deregister");
-                self.connections[idx] = None;
+            match conn.poll(seq, &self.firehose) {
+                Ok(true) => return true,
+                Ok(false) => {
+                    tracing::info!("[{}] closed due to invalid cursor", conn.addr);
+                }
+                Err(err) => {
+                    tracing::info!("[{}] disconnected: {err}", conn.addr);
+                }
             }
+            #[expect(clippy::expect_used)]
+            self.poll
+                .registry()
+                .deregister(&mut SourceFd(&conn.as_raw_fd()))
+                .expect("failed to deregister");
+            self.connections[idx] = None;
         }
+
         true
     }
 }

rsky-relay/src/types.rs

@@ -2,9 +2,9 @@ use std::fmt;
 use std::ops::{Add, Sub};
 use std::sync::LazyLock;
 
+use bytes::Bytes;
 use sled::{Db, IVec, Mode};
 use thingbuf::{Recycle, mpsc};
-use tungstenite::Bytes;
 use zerocopy::big_endian::U64;
 use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Unaligned};
 
@@ -32,6 +32,7 @@ impl Recycle<Message> for MessageRecycle {
 
     fn recycle(&self, element: &mut Message) {
         element.data.clear();
+        element.hostname.clear();
     }
 }
 

rsky-relay/src/validator/manager.rs

@@ -124,7 +124,7 @@ impl Manager {
                             }
                             continue;
                         } else if prev + 1 != curr {
-                            tracing::debug!(
+                            tracing::trace!(
                                 "[{}] seq gap: {prev} -> {curr} ({})",
                                 msg.hostname,
                                 curr - prev - 1
@@ -181,7 +181,7 @@ impl Manager {
             let Some((did, key)) = self.resolver.poll().await? else {
                 continue;
             };
-            let Some(msgs) = self.queue.remove(&did)? else {
+            let Some(msgs) = self.queue.remove(did)? else {
                 tracing::debug!("missing queue for did: {did}");
                 continue;
             };
@@ -195,7 +195,7 @@ impl Manager {
                 };
                 #[expect(clippy::unwrap_used)]
                 let commit = event.commit().unwrap().unwrap(); // Already tried parsing
-                match utils::verify_commit_sig(&commit, key) {
+                match utils::verify_commit_sig(&commit, *key) {
                     Ok(res) => {
                         if res {
                             let data = event.serialize(data.len(), seq.next())?;
@@ -230,5 +230,6 @@ impl Drop for Manager {
         if let Err(err) = DB.flush() {
             tracing::warn!("unable to persist cursors: {err}\n{:#?}", self.cursors);
         }
+        tracing::info!("persisted {} cursors", self.cursors.len());
     }
 }