rsky-relay: publisher

Author: utkarshgupta137

rsky-relay/Cargo.toml

@@ -5,6 +5,7 @@ edition = "2024"
 
 [dependencies]
 # external
+bstr = "1"
 bus = "2"
 color-eyre = "0.6"
 derive_more = { version = "2", features = ["full"] }
@@ -17,7 +18,10 @@ rtrb = "0.3"
 rustls = "0.23"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
+signal-hook = { version = "0.3", features = ["extended-siginfo"] }
 thingbuf = "0.1"
 thiserror = "2"
+tracing = { version = "0.1", features = ["release_max_level_debug"] }
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 tungstenite = { version = "0.26", features = ["rustls-tls-native-roots"] }
 url = "2"

rsky-relay/src/client/connection.rs renamed to rsky-relay/src/crawler/connection.rs

@@ -6,7 +6,7 @@ use thiserror::Error;
 use tungstenite::Message;
 use tungstenite::stream::MaybeTlsStream;
 
-use crate::client::types::{Client, Config, StatusSender};
+use crate::crawler::types::{Client, Config, StatusSender};
 use crate::types::MessageSender;
 
 #[derive(Debug, Error)]
@@ -27,6 +27,7 @@ pub struct Connection {
 }
 
 impl AsFd for Connection {
+    #[inline]
     fn as_fd(&self) -> BorrowedFd<'_> {
         match self.client.get_ref() {
             MaybeTlsStream::Plain(stream) => stream.as_fd(),
@@ -60,11 +61,11 @@ impl Connection {
     }
 
     pub fn poll(&mut self) -> Result<(), ConnectionError> {
-        loop {
+        for _ in 0..1024 {
             let msg = match self.client.read() {
                 Ok(msg) => msg,
                 Err(tungstenite::Error::Io(e)) if e.kind() == io::ErrorKind::WouldBlock => {
-                    return Ok(());
+                    break;
                 }
                 Err(err) => Err(err)?,
             };
@@ -79,5 +80,6 @@ impl Connection {
             slot.data = bytes.into();
             slot.uri = self.config.uri.clone();
         }
+        Ok(())
     }
 }

rsky-relay/src/client/manager.rs renamed to rsky-relay/src/crawler/manager.rs

@@ -1,3 +1,4 @@
+use std::sync::atomic::Ordering;
 use std::thread;
 
 use hashbrown::HashMap;
@@ -6,11 +7,12 @@ use magnetic::Consumer;
 use magnetic::buffer::dynamic::DynamicBufferP2;
 use thiserror::Error;
 
-use crate::client::types::{
+use crate::SHUTDOWN;
+use crate::crawler::types::{
     Command, CommandSender, Config, LocalId, Status, StatusReceiver, WorkerId,
 };
-use crate::client::worker::{Worker, WorkerError};
-use crate::types::{CrawlRequestReceiver, Cursor, MessageSender};
+use crate::crawler::worker::{Worker, WorkerError};
+use crate::types::{Cursor, MessageSender, RequestCrawlReceiver};
 
 const CAPACITY: usize = 1024;
 
@@ -28,20 +30,20 @@ pub enum ManagerError {
 struct WorkerHandle {
     pub configs: Vec<Config>,
     pub command_tx: CommandSender,
-    pub thread_handle: thread::JoinHandle<()>,
+    pub thread_handle: thread::JoinHandle<Result<(), WorkerError>>,
 }
 
 pub struct Manager {
     workers: Box<[WorkerHandle]>,
     next_id: WorkerId,
     configs: HashMap<Uri, Config>,
     status_rx: StatusReceiver,
-    crawl_request_rx: CrawlRequestReceiver,
+    request_crawl_rx: RequestCrawlReceiver,
 }
 
 impl Manager {
     pub fn new(
-        n_workers: usize, message_tx: MessageSender, crawl_request_rx: CrawlRequestReceiver,
+        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());
@@ -51,7 +53,7 @@ impl Manager {
                 let status_tx = status_tx.clone();
                 let (command_tx, command_rx) = rtrb::RingBuffer::new(CAPACITY);
                 let thread_handle = thread::spawn(move || {
-                    Worker::new(WorkerId(worker_id), message_tx, status_tx, command_rx).run();
+                    Worker::new(WorkerId(worker_id), message_tx, status_tx, command_rx).run()
                 });
                 WorkerHandle { configs: Vec::new(), command_tx, thread_handle }
             })
@@ -61,7 +63,7 @@ impl Manager {
             next_id: WorkerId(0),
             configs: HashMap::new(),
             status_rx,
-            crawl_request_rx,
+            request_crawl_rx,
         })
     }
 
@@ -74,8 +76,10 @@ impl Manager {
         for worker in &mut self.workers {
             worker.command_tx.push(Command::Shutdown)?;
         }
-        for worker in self.workers {
-            worker.thread_handle.join().map_err(|_| ManagerError::JoinError)?;
+        for (id, worker) in self.workers.into_iter().enumerate() {
+            if let Err(err) = worker.thread_handle.join().map_err(|_| ManagerError::JoinError)? {
+                tracing::warn!("crawler worker {id} error: {err}");
+            }
         }
         Ok(())
     }
@@ -86,22 +90,26 @@ impl Manager {
     }
 
     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) = self.crawl_request_rx.pop() {
-            if !self.configs.contains_key(&request.uri) {
+        if let Ok(request_crawl) = self.request_crawl_rx.pop() {
+            if !self.configs.contains_key(&request_crawl.uri) {
                 let config = Config {
-                    uri: request.uri.clone(),
+                    uri: request_crawl.uri.clone(),
                     cursor: Cursor(0),
                     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.uri.clone(), config.clone());
+                self.configs.insert(request_crawl.uri, config.clone());
                 self.workers[config.worker_id.0].command_tx.push(Command::Connect(config)).unwrap();
             }
         }

rsky-relay/src/client/mod.rs renamed to rsky-relay/src/crawler/mod.rs

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

rsky-relay/src/client/types.rs renamed to rsky-relay/src/crawler/types.rs

@@ -31,7 +31,6 @@ pub struct Config {
 #[derive(Debug, Clone)]
 pub enum Command {
     Connect(Config),
-    Reconnect(LocalId),
     Shutdown,
 }
 

rsky-relay/src/client/worker.rs renamed to rsky-relay/src/crawler/worker.rs

@@ -1,10 +1,14 @@
+use std::thread;
+
+#[cfg(target_os = "linux")]
 use nix::sys::epoll::{Epoll, EpollCreateFlags, EpollEvent, EpollFlags};
 use thiserror::Error;
 
-use crate::client::connection::{Connection, ConnectionError};
-use crate::client::types::{Command, CommandReceiver, Config, LocalId, StatusSender, WorkerId};
+use crate::crawler::connection::{Connection, ConnectionError};
+use crate::crawler::types::{Command, CommandReceiver, Config, LocalId, StatusSender, WorkerId};
 use crate::types::MessageSender;
 
+#[cfg(target_os = "linux")]
 const EPOLL_FLAGS: EpollFlags = EpollFlags::EPOLLIN;
 
 #[derive(Debug, Error)]
@@ -20,7 +24,9 @@ pub struct Worker {
     message_tx: MessageSender,
     status_tx: StatusSender,
     command_rx: CommandReceiver,
+    #[cfg(target_os = "linux")]
     epoll: Epoll,
+    #[cfg(target_os = "linux")]
     events: Vec<EpollEvent>,
 }
 
@@ -36,14 +42,28 @@ impl Worker {
             message_tx,
             status_tx,
             command_rx,
+            #[cfg(target_os = "linux")]
             #[expect(clippy::expect_used)]
             epoll: Epoll::new(EpollCreateFlags::empty()).expect("failed to create epoll"),
+            #[cfg(target_os = "linux")]
             events: vec![EpollEvent::empty(); 1024],
         }
     }
 
-    pub fn run(mut self) {
-        while self.update() {}
+    pub fn run(mut self) -> Result<(), WorkerError> {
+        while self.update() {
+            thread::yield_now();
+        }
+        self.shutdown()
+    }
+
+    pub fn shutdown(mut self) -> Result<(), WorkerError> {
+        for conn in self.connections.iter_mut().filter_map(|x| x.as_mut()) {
+            if let Err(err) = conn.close() {
+                tracing::warn!("crawler conn close error: {err}");
+            }
+        }
+        Ok(())
     }
 
     fn handle_command(&mut self, command: Command) -> bool {
@@ -55,44 +75,30 @@ impl Worker {
                     self.status_tx.clone(),
                 ) {
                     Ok(conn) => {
+                        #[cfg(target_os = "linux")]
                         #[expect(clippy::expect_used)]
                         self.epoll
                             .add(&conn, EpollEvent::new(EPOLL_FLAGS, config.local_id.0 as _))
                             .expect("failed to add connection");
                         self.connections.push(Some(conn));
                         self.configs.push(config);
                     }
-                    Err(_) => todo!(),
-                }
-            }
-            Command::Reconnect(local_id) => {
-                if let Some(conn) = &mut self.connections[local_id.0] {
-                    match conn.close() {
-                        Ok(()) => {
-                            self.connections[local_id.0] = None;
-                        }
-                        Err(_) => todo!(),
-                    }
-                }
-                match Connection::connect(
-                    self.configs[local_id.0].clone(),
-                    self.message_tx.clone(),
-                    self.status_tx.clone(),
-                ) {
-                    Ok(conn) => {
-                        self.connections[local_id.0] = Some(conn);
+                    Err(err) => {
+                        tracing::warn!("[{}] unable to requestCrawl: {err}", config.uri);
                     }
-                    Err(_) => todo!(),
                 }
             }
-            Command::Shutdown => return false,
+            Command::Shutdown => {
+                tracing::debug!("shutting down crawler: {}", self.worker_id.0);
+                return false;
+            }
         }
         true
     }
 
+    #[cfg(target_os = "linux")]
     fn update(&mut self) -> bool {
         for _ in 0..32 {
-            // returns false only upon global shutdown, true otherwise
             if let Ok(command) = self.command_rx.pop() {
                 if !self.handle_command(command) {
                     return false;
@@ -128,9 +134,29 @@ impl Worker {
         true
     }
 
+    #[cfg(not(target_os = "linux"))]
+    fn update(&mut self) -> bool {
+        if let Ok(command) = self.command_rx.pop() {
+            if !self.handle_command(command) {
+                return false;
+            }
+        }
+
+        for _ in 0..32 {
+            for local_id in 0..self.connections.len() {
+                if !self.poll(LocalId(local_id)) {
+                    return false;
+                }
+            }
+        }
+
+        true
+    }
+
     fn poll(&mut self, local_id: LocalId) -> bool {
         if let Some(conn) = &mut self.connections[local_id.0] {
             if let Err(_) = conn.poll() {
+                #[cfg(target_os = "linux")]
                 #[expect(clippy::expect_used)]
                 self.epoll.delete(conn).expect("failed to delete connection");
                 self.connections[local_id.0] = None;

rsky-relay/src/lib.rs

@@ -1,4 +1,13 @@
-pub mod client;
+mod crawler;
+mod publisher;
+mod server;
 mod types;
 
-pub use types::{CrawlRequest, MessageRecycle};
+use std::sync::atomic::AtomicBool;
+
+pub static SHUTDOWN: AtomicBool = AtomicBool::new(false);
+
+pub use crawler::Manager as CrawlerManager;
+pub use publisher::Manager as PublisherManager;
+pub use server::Server;
+pub use types::{MessageRecycle, RequestCrawl};