rsky-relay: publisher state/mio

Author: utkarshgupta137

rsky-relay/Cargo.toml

@@ -6,7 +6,6 @@ edition = "2024"
 [dependencies]
 # external
 bstr = "1"
-bus = "2"
 chrono = { version = "0.4", features = ["serde"] }
 ciborium = "0.2"
 cid = { version = "0.11", features = ["serde-codec"] }
@@ -19,8 +18,8 @@ httparse = "1.10"
 k256 = "0.13"
 magnetic = "2"
 mimalloc = "0.1"
+mio = { version = "1", features = ["os-ext", "os-poll"] }
 multibase = "0.9"
-nix = { version = "0.29", features = ["event", "socket"] }
 p256 = "0.13"
 reqwest = { version = "0.12", default-features = false, features = ["gzip", "hickory-dns", "http2", "json", "rustls-tls-webpki-roots-no-provider"] }
 rs-car-sync = "0.4"
@@ -38,12 +37,11 @@ thiserror = "2"
 tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] }
 tracing = { version = "0.1", features = ["release_max_level_debug"] }
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
-tungstenite = { version = "0.26", features = ["rustls-tls-webpki-roots"] }
+tungstenite = { version = "0.26", features = ["rustls-tls-webpki-roots", "url"] }
 url = "2"
+urlencoding = "2"
 zerocopy = { version = "0.8", features = ["derive"] }
 
 # internal
 rsky-common = { workspace = true }
-rsky-crypto = { workspace = true }
 rsky-identity = { workspace = true }
-rsky-lexicon = { workspace = true }

rsky-relay/bootstrap.sh

@@ -35,3 +35,5 @@ curl -X POST -H 'Content-Type: application/json' -d '{"hostname":"verpa.us-west.
 curl -X POST -H 'Content-Type: application/json' -d '{"hostname":"witchesbutter.us-west.host.bsky.network"}' 'http://localhost:9000/xrpc/com.atproto.sync.requestCrawl'
 curl -X POST -H 'Content-Type: application/json' -d '{"hostname":"woodear.us-west.host.bsky.network"}' 'http://localhost:9000/xrpc/com.atproto.sync.requestCrawl'
 curl -X POST -H 'Content-Type: application/json' -d '{"hostname":"atproto.brid.gy"}' 'http://localhost:9000/xrpc/com.atproto.sync.requestCrawl'
+curl -X POST -H 'Content-Type: application/json' -d '{"hostname":"esnoticia.online"}' 'http://localhost:9000/xrpc/com.atproto.sync.requestCrawl'
+curl -X POST -H 'Content-Type: application/json' -d '{"hostname":"at.app.wafrn.net"}' 'http://localhost:9000/xrpc/com.atproto.sync.requestCrawl'

rsky-relay/src/crawler/connection.rs

@@ -1,13 +1,14 @@
 use std::io;
-use std::os::fd::{AsFd, BorrowedFd};
+use std::net::TcpStream;
+use std::os::fd::{AsRawFd, RawFd};
 
 use thingbuf::mpsc;
 use thiserror::Error;
-use tungstenite::Message;
 use tungstenite::stream::MaybeTlsStream;
+use tungstenite::{Message, WebSocket};
+use url::Url;
 
-use crate::crawler::types::{Client, Config, StatusSender};
-use crate::types::MessageSender;
+use crate::types::{Cursor, MessageSender};
 
 #[derive(Debug, Error)]
 pub enum ConnectionError {
@@ -20,28 +21,33 @@ pub enum ConnectionError {
 }
 
 pub struct Connection {
-    client: Client,
-    config: Config,
+    pub(crate) hostname: String,
+    client: WebSocket<MaybeTlsStream<TcpStream>>,
     message_tx: MessageSender,
-    _status_tx: StatusSender,
 }
 
-impl AsFd for Connection {
+impl AsRawFd for Connection {
     #[inline]
-    fn as_fd(&self) -> BorrowedFd<'_> {
+    fn as_raw_fd(&self) -> RawFd {
         match self.client.get_ref() {
-            MaybeTlsStream::Plain(stream) => stream.as_fd(),
-            MaybeTlsStream::Rustls(stream) => stream.get_ref().as_fd(),
+            MaybeTlsStream::Plain(stream) => stream.as_raw_fd(),
+            MaybeTlsStream::Rustls(stream) => stream.get_ref().as_raw_fd(),
             _ => todo!(),
         }
     }
 }
 
 impl Connection {
     pub fn connect(
-        config: Config, message_tx: MessageSender, status_tx: StatusSender,
+        hostname: String, cursor: Option<Cursor>, message_tx: MessageSender,
     ) -> Result<Self, ConnectionError> {
-        let (client, _) = tungstenite::connect(&config.uri)?;
+        #[expect(clippy::unwrap_used)]
+        let mut url =
+            Url::parse(&format!("wss://{hostname}/xrpc/com.atproto.sync.subscribeRepos")).unwrap();
+        if let Some(cursor) = cursor {
+            url.query_pairs_mut().append_pair("cursor", &cursor.to_string());
+        }
+        let (client, _) = tungstenite::connect(url)?;
         match client.get_ref() {
             MaybeTlsStream::Rustls(stream) => {
                 stream.get_ref().set_nonblocking(true)?;
@@ -51,7 +57,7 @@ impl Connection {
             }
             _ => {}
         }
-        Ok(Self { client, config, message_tx, _status_tx: status_tx })
+        Ok(Self { hostname, client, message_tx })
     }
 
     pub fn close(&mut self) -> Result<(), ConnectionError> {
@@ -72,16 +78,22 @@ impl Connection {
 
             let bytes = match msg {
                 Message::Binary(bytes) => bytes,
-                Message::Close(_) => todo!(),
+                Message::Ping(_) | Message::Pong(_) => {
+                    continue;
+                }
+                Message::Close(close) => {
+                    tracing::debug!("[{}] received close: {close:?}", self.hostname);
+                    continue;
+                }
                 _ => {
-                    tracing::debug!("[{}] unknown message: {msg}", self.config.hostname);
+                    tracing::debug!("[{}] unknown message: {msg:?}", self.hostname);
                     continue;
                 }
             };
 
             let mut slot = self.message_tx.try_send_ref()?;
-            slot.data = bytes.into();
-            slot.hostname = self.config.hostname.clone();
+            slot.data = bytes;
+            slot.hostname.clone_from(&self.hostname);
         }
         Ok(())
     }

rsky-relay/src/crawler/manager.rs

@@ -2,75 +2,61 @@ use std::sync::atomic::Ordering;
 use std::time::Duration;
 use std::{io, thread};
 
-use hashbrown::HashMap;
-use http::Uri;
-use magnetic::Consumer;
-use magnetic::buffer::dynamic::DynamicBufferP2;
 use thiserror::Error;
 
 use crate::SHUTDOWN;
-use crate::crawler::types::{
-    Command, CommandSender, Config, LocalId, Status, StatusReceiver, WorkerId,
-};
+use crate::crawler::types::{Command, CommandSender, RequestCrawlReceiver};
 use crate::crawler::worker::{Worker, WorkerError};
-use crate::types::{MessageSender, RequestCrawlReceiver};
+use crate::types::MessageSender;
 
 const CAPACITY: usize = 1024;
 const SLEEP: Duration = Duration::from_millis(10);
 
 #[derive(Debug, Error)]
 pub enum ManagerError {
     #[error("spawn error: {0}")]
-    SpawnError(#[from] io::Error),
+    Spawn(#[from] io::Error),
     #[error("worker error: {0}")]
-    WorkerError(#[from] WorkerError),
+    Worker(#[from] WorkerError),
     #[error("rtrb error: {0}")]
-    PushError(#[from] rtrb::PushError<Command>),
+    Push(#[from] Box<rtrb::PushError<Command>>),
     #[error("join error")]
-    JoinError,
+    Join,
+}
+
+impl From<rtrb::PushError<Command>> for ManagerError {
+    fn from(value: rtrb::PushError<Command>) -> Self {
+        Box::new(value).into()
+    }
 }
 
 #[derive(Debug)]
 struct WorkerHandle {
-    pub configs: Vec<Config>,
     pub command_tx: CommandSender,
     pub thread_handle: thread::JoinHandle<Result<(), WorkerError>>,
 }
 
 pub struct Manager {
     workers: Box<[WorkerHandle]>,
-    next_id: WorkerId,
-    configs: HashMap<Uri, Config>,
-    status_rx: StatusReceiver,
+    next_id: usize,
     request_crawl_rx: RequestCrawlReceiver,
 }
 
 impl Manager {
     pub fn new(
-        n_workers: usize, message_tx: MessageSender, request_crawl_rx: RequestCrawlReceiver,
+        n_workers: usize, message_tx: &MessageSender, request_crawl_rx: RequestCrawlReceiver,
     ) -> Result<Self, ManagerError> {
-        let (status_tx, status_rx) =
-            magnetic::mpsc::mpsc_queue(DynamicBufferP2::new(CAPACITY).unwrap());
         let workers = (0..n_workers)
             .map(|worker_id| -> Result<_, ManagerError> {
                 let message_tx = message_tx.clone();
-                let status_tx = status_tx.clone();
                 let (command_tx, command_rx) = rtrb::RingBuffer::new(CAPACITY);
                 let thread_handle = thread::Builder::new()
                     .name(format!("rsky-crawl-{worker_id}"))
-                    .spawn(move || {
-                        Worker::new(WorkerId(worker_id), message_tx, status_tx, command_rx).run()
-                    })?;
-                Ok(WorkerHandle { configs: Vec::new(), command_tx, thread_handle })
+                    .spawn(move || Worker::new(worker_id, message_tx, command_rx)?.run())?;
+                Ok(WorkerHandle { command_tx, thread_handle })
             })
             .collect::<Result<Vec<_>, _>>()?;
-        Ok(Self {
-            workers: workers.into_boxed_slice(),
-            next_id: WorkerId(0),
-            configs: HashMap::new(),
-            status_rx,
-            request_crawl_rx,
-        })
+        Ok(Self { workers: workers.into_boxed_slice(), next_id: 0, request_crawl_rx })
     }
 
     pub fn run(mut self) -> Result<(), ManagerError> {
@@ -87,40 +73,21 @@ impl Manager {
             worker.command_tx.push(Command::Shutdown)?;
         }
         for (id, worker) in self.workers.into_iter().enumerate() {
-            if let Err(err) = worker.thread_handle.join().map_err(|_| ManagerError::JoinError)? {
+            if let Err(err) = worker.thread_handle.join().map_err(|_| ManagerError::Join)? {
                 tracing::warn!("crawler worker {id} error: {err}");
             }
         }
         Ok(())
     }
 
-    fn handle_status(&mut self, _status: Status) -> Result<bool, ManagerError> {
-        Ok(true)
-    }
-
     fn update(&mut self) -> Result<bool, ManagerError> {
         if SHUTDOWN.load(Ordering::Relaxed) {
             return Ok(false);
         }
 
-        if let Ok(status) = self.status_rx.try_pop() {
-            if !self.handle_status(status)? {
-                return Ok(false);
-            }
-        }
-
         if let Ok(request_crawl) = self.request_crawl_rx.pop() {
-            if !self.configs.contains_key(&request_crawl.uri) {
-                let config = Config {
-                    uri: request_crawl.uri.clone(),
-                    hostname: request_crawl.hostname.clone(),
-                    worker_id: self.next_id,
-                    local_id: LocalId(self.workers[self.next_id.0].configs.len()),
-                };
-                self.next_id = WorkerId((self.next_id.0 + 1) % self.workers.len());
-                self.configs.insert(request_crawl.uri, config.clone());
-                self.workers[config.worker_id.0].command_tx.push(Command::Connect(config)).unwrap();
-            }
+            self.workers[self.next_id].command_tx.push(Command::Connect(request_crawl))?;
+            self.next_id = (self.next_id + 1) % self.workers.len();
         }
 
         Ok(true)

rsky-relay/src/crawler/mod.rs

@@ -4,3 +4,4 @@ mod types;
 mod worker;
 
 pub use manager::Manager;
+pub use types::{RequestCrawl, RequestCrawlSender};

rsky-relay/src/crawler/types.rs

@@ -1,36 +1,22 @@
-use std::net::TcpStream;
-
-use http::Uri;
-use magnetic::buffer::dynamic::DynamicBufferP2;
-use magnetic::mpsc::{MPSCConsumer, MPSCProducer};
 use rtrb::{Consumer, Producer};
-use tungstenite::WebSocket;
-use tungstenite::stream::MaybeTlsStream;
+use serde::Deserialize;
+
+use crate::types::Cursor;
 
-pub type Client = WebSocket<MaybeTlsStream<TcpStream>>;
 pub type CommandSender = Producer<Command>;
 pub type CommandReceiver = Consumer<Command>;
-pub type StatusSender = MPSCProducer<Status, DynamicBufferP2<Status>>;
-pub type StatusReceiver = MPSCConsumer<Status, DynamicBufferP2<Status>>;
+pub type RequestCrawlSender = Producer<RequestCrawl>;
+pub type RequestCrawlReceiver = Consumer<RequestCrawl>;
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
-pub struct WorkerId(pub usize);
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
-pub struct LocalId(pub usize);
-
-#[derive(Debug, Clone)]
-pub struct Config {
-    pub uri: Uri,
+#[derive(Debug, Deserialize)]
+pub struct RequestCrawl {
     pub hostname: String,
-    pub worker_id: WorkerId,
-    pub local_id: LocalId,
+    #[serde(skip)]
+    pub cursor: Option<Cursor>,
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub enum Command {
-    Connect(Config),
+    Connect(RequestCrawl),
     Shutdown,
 }
-
-#[derive(Debug, Clone)]
-pub enum Status {}