rsky-relay: crawler

Author: utkarshgupta137

rsky-relay/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "rsky-relay"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+# external
+bus = "2"
+color-eyre = "0.6"
+derive_more = { version = "2", features = ["full"] }
+hashbrown = "0.15"
+http = "1"
+httparse = "1.10"
+magnetic = "2"
+nix = { version = "0.29", features = ["event", "socket"] }
+rtrb = "0.3"
+rustls = "0.23"
+serde = { version = "1", features = ["derive"] }
+serde_json = "1"
+thingbuf = "0.1"
+thiserror = "2"
+tungstenite = { version = "0.26", features = ["rustls-tls-native-roots"] }
+url = "2"

rsky-relay/src/client/connection.rs

@@ -0,0 +1,83 @@
+use std::io;
+use std::os::fd::{AsFd, BorrowedFd};
+
+use thingbuf::mpsc;
+use thiserror::Error;
+use tungstenite::Message;
+use tungstenite::stream::MaybeTlsStream;
+
+use crate::client::types::{Client, Config, StatusSender};
+use crate::types::MessageSender;
+
+#[derive(Debug, Error)]
+pub enum ConnectionError {
+    #[error("io error: {0}")]
+    Io(#[from] io::Error),
+    #[error("tungstenite error: {0}")]
+    Tungstenite(#[from] tungstenite::Error),
+    #[error("thingbuf error: {0}")]
+    Thingbuf(#[from] mpsc::errors::TrySendError),
+}
+
+pub struct Connection {
+    client: Client,
+    config: Config,
+    message_tx: MessageSender,
+    status_tx: StatusSender,
+}
+
+impl AsFd for Connection {
+    fn as_fd(&self) -> BorrowedFd<'_> {
+        match self.client.get_ref() {
+            MaybeTlsStream::Plain(stream) => stream.as_fd(),
+            MaybeTlsStream::Rustls(stream) => stream.get_ref().as_fd(),
+            _ => todo!(),
+        }
+    }
+}
+
+impl Connection {
+    pub fn connect(
+        config: Config, message_tx: MessageSender, status_tx: StatusSender,
+    ) -> Result<Self, ConnectionError> {
+        let (client, _) = tungstenite::connect(&config.uri)?;
+        match client.get_ref() {
+            MaybeTlsStream::Rustls(stream) => {
+                stream.get_ref().set_nonblocking(true)?;
+            }
+            MaybeTlsStream::Plain(stream) => {
+                stream.set_nonblocking(true)?;
+            }
+            _ => {}
+        }
+        Ok(Self { client, config, message_tx, status_tx })
+    }
+
+    pub fn close(&mut self) -> Result<(), ConnectionError> {
+        self.client.close(None)?;
+        self.client.flush()?;
+        Ok(())
+    }
+
+    pub fn poll(&mut self) -> Result<(), ConnectionError> {
+        loop {
+            let msg = match self.client.read() {
+                Ok(msg) => msg,
+                Err(tungstenite::Error::Io(e)) if e.kind() == io::ErrorKind::WouldBlock => {
+                    return Ok(());
+                }
+                Err(err) => Err(err)?,
+            };
+
+            let bytes = match msg {
+                Message::Binary(bytes) => bytes,
+                Message::Close(_) => todo!(),
+                _ => unreachable!("{msg}"),
+            };
+
+            let mut slot = self.message_tx.try_send_ref()?;
+            slot.data = bytes.into();
+            slot.uri = self.config.uri.clone();
+        }
+    }
+}

rsky-relay/src/client/manager.rs

@@ -0,0 +1,111 @@
+use std::thread;
+
+use hashbrown::HashMap;
+use http::Uri;
+use magnetic::Consumer;
+use magnetic::buffer::dynamic::DynamicBufferP2;
+use thiserror::Error;
+
+use crate::client::types::{
+    Command, CommandSender, Config, LocalId, Status, StatusReceiver, WorkerId,
+};
+use crate::client::worker::{Worker, WorkerError};
+use crate::types::{CrawlRequestReceiver, Cursor, MessageSender};
+
+const CAPACITY: usize = 1024;
+
+#[derive(Debug, Error)]
+pub enum ManagerError {
+    #[error("worker error: {0}")]
+    WorkerError(#[from] WorkerError),
+    #[error("rtrb error: {0}")]
+    PushError(#[from] rtrb::PushError<Command>),
+    #[error("join error")]
+    JoinError,
+}
+
+#[derive(Debug)]
+struct WorkerHandle {
+    pub configs: Vec<Config>,
+    pub command_tx: CommandSender,
+    pub thread_handle: thread::JoinHandle<()>,
+}
+
+pub struct Manager {
+    workers: Box<[WorkerHandle]>,
+    next_id: WorkerId,
+    configs: HashMap<Uri, Config>,
+    status_rx: StatusReceiver,
+    crawl_request_rx: CrawlRequestReceiver,
+}
+
+impl Manager {
+    pub fn new(
+        n_workers: usize, message_tx: MessageSender, crawl_request_rx: CrawlRequestReceiver,
+    ) -> Result<Self, ManagerError> {
+        let (status_tx, status_rx) =
+            magnetic::mpsc::mpsc_queue(DynamicBufferP2::new(CAPACITY).unwrap());
+        let workers = (0..n_workers)
+            .map(|worker_id| {
+                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::spawn(move || {
+                    Worker::new(WorkerId(worker_id), message_tx, status_tx, command_rx).run();
+                });
+                WorkerHandle { configs: Vec::new(), command_tx, thread_handle }
+            })
+            .collect::<Vec<_>>();
+        Ok(Self {
+            workers: workers.into_boxed_slice(),
+            next_id: WorkerId(0),
+            configs: HashMap::new(),
+            status_rx,
+            crawl_request_rx,
+        })
+    }
+
+    pub fn run(mut self) -> Result<(), ManagerError> {
+        while self.update()? {}
+        self.shutdown()
+    }
+
+    pub fn shutdown(mut self) -> Result<(), ManagerError> {
+        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)?;
+        }
+        Ok(())
+    }
+
+    fn handle_status(&mut self, status: Status) -> Result<bool, ManagerError> {
+        match status {}
+        Ok(true)
+    }
+
+    fn update(&mut self) -> Result<bool, ManagerError> {
+        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) {
+                let config = Config {
+                    uri: request.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.workers[config.worker_id.0].command_tx.push(Command::Connect(config)).unwrap();
+            }
+        }
+
+        Ok(true)
+    }
+}

rsky-relay/src/client/mod.rs

@@ -0,0 +1,6 @@
+mod connection;
+mod manager;
+mod types;
+mod worker;
+
+pub use manager::{Manager, ManagerError};

rsky-relay/src/client/types.rs

@@ -0,0 +1,39 @@
+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 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>>;
+
+#[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,
+    pub cursor: Cursor,
+    pub worker_id: WorkerId,
+    pub local_id: LocalId,
+}
+
+#[derive(Debug, Clone)]
+pub enum Command {
+    Connect(Config),
+    Reconnect(LocalId),
+    Shutdown,
+}
+
+#[derive(Debug, Clone)]
+pub enum Status {}