backend redesign

relay/src/relay.rs

@@ -0,0 +1,323 @@
+use crate::config::Config;
+use crate::proto;
+use anyhow::{Context, Result};
+use prost::Message;
+use tracing::{debug, error, info, warn};
+
+const PROTOCOL_VERSION: u8 = 0x01;
+const MSG_TYPE_SUBMIT_REQUEST: u8 = 0x10;
+const MSG_TYPE_SUBMIT_RESPONSE: u8 = 0x11;
+const MSG_TYPE_DATA_REQUEST: u8 = 0x01;
+const MSG_TYPE_HISTORY_READY: u8 = 0x12;
+
+pub struct Relay {
+    config: Config,
+    context: zmq::Context,
+}
+
+impl Relay {
+    pub fn new(config: Config) -> Result<Self> {
+        let context = zmq::Context::new();
+
+        Ok(Self {
+            config,
+            context,
+        })
+    }
+
+    pub async fn run(self) -> Result<()> {
+        info!("Initializing Stateless ZMQ Relay");
+
+        // Bind sockets
+        let client_request_socket = self.create_client_request_socket()?;
+        let market_data_frontend = self.create_market_data_frontend()?;
+        let market_data_backend = self.create_market_data_backend()?;
+        let ingestor_work_socket = self.create_ingestor_work_socket()?;
+
+        info!("All sockets initialized successfully - relay is STATELESS");
+        info!("No pending requests tracked - all async via pub/sub");
+
+        // Run main loop
+        tokio::task::spawn_blocking(move || {
+            Self::proxy_loop(
+                client_request_socket,
+                market_data_frontend,
+                market_data_backend,
+                ingestor_work_socket,
+            )
+        })
+        .await?
+    }
+
+    fn create_client_request_socket(&self) -> Result<zmq::Socket> {
+        let socket = self.context.socket(zmq::ROUTER)?;
+        socket.set_sndhwm(self.config.high_water_mark)?;
+        socket.set_rcvhwm(self.config.high_water_mark)?;
+        socket.set_linger(1000)?;
+
+        let endpoint = format!("{}:{}", self.config.bind_address, self.config.client_request_port);
+        socket.bind(&endpoint)?;
+        info!("Client request socket (ROUTER) bound to {}", endpoint);
+        info!("  → Accepts SubmitHistoricalRequest, returns SubmitResponse immediately");
+
+        Ok(socket)
+    }
+
+    fn create_market_data_frontend(&self) -> Result<zmq::Socket> {
+        let socket = self.context.socket(zmq::XPUB)?;
+        socket.set_sndhwm(self.config.high_water_mark)?;
+        socket.set_xpub_verbose(true)?;
+
+        let endpoint = format!("{}:{}", self.config.bind_address, self.config.market_data_pub_port);
+        socket.bind(&endpoint)?;
+        info!("Market data frontend (XPUB) bound to {}", endpoint);
+        info!("  → Clients subscribe here for HistoryReadyNotification and market data");
+
+        Ok(socket)
+    }
+
+    fn create_market_data_backend(&self) -> Result<zmq::Socket> {
+        let socket = self.context.socket(zmq::XSUB)?;
+        socket.set_rcvhwm(self.config.high_water_mark)?;
+
+        socket.connect(&self.config.flink_market_data_endpoint)?;
+        info!("Market data backend (XSUB) connected to {}", self.config.flink_market_data_endpoint);
+        info!("  → Receives HistoryReadyNotification and market data from Flink");
+
+        Ok(socket)
+    }
+
+    fn create_ingestor_work_socket(&self) -> Result<zmq::Socket> {
+        let socket = self.context.socket(zmq::PUB)?;
+        socket.set_sndhwm(self.config.high_water_mark)?;
+        socket.set_linger(1000)?;
+
+        let endpoint = format!("{}:{}", self.config.bind_address, self.config.ingestor_work_port);
+        socket.bind(&endpoint)?;
+        info!("Ingestor work queue (PUB) bound to {}", endpoint);
+        info!("  → Publishes DataRequest with exchange prefix");
+
+        Ok(socket)
+    }
+
+    fn proxy_loop(
+        client_request_socket: zmq::Socket,
+        market_data_frontend: zmq::Socket,
+        market_data_backend: zmq::Socket,
+        ingestor_work_socket: zmq::Socket,
+    ) -> Result<()> {
+        let mut items = [
+            client_request_socket.as_poll_item(zmq::POLLIN),
+            market_data_frontend.as_poll_item(zmq::POLLIN),
+            market_data_backend.as_poll_item(zmq::POLLIN),
+        ];
+
+        info!("Entering stateless proxy loop");
+
+        loop {
+            // Poll with 100ms timeout
+            zmq::poll(&mut items, 100)
+                .context("Failed to poll sockets")?;
+
+            // Handle client request submissions
+            if items[0].is_readable() {
+                if let Err(e) = Self::handle_client_submission(
+                    &client_request_socket,
+                    &ingestor_work_socket,
+                ) {
+                    error!("Error handling client submission: {}", e);
+                }
+            }
+
+            // Handle market data subscriptions from clients (XPUB → XSUB)
+            if items[1].is_readable() {
+                if let Err(e) = Self::proxy_subscription(&market_data_frontend, &market_data_backend) {
+                    error!("Error proxying subscription: {}", e);
+                }
+            }
+
+            // Handle market data from Flink (XSUB → XPUB)
+            // This includes HistoryReadyNotification and regular market data
+            if items[2].is_readable() {
+                if let Err(e) = Self::proxy_market_data(&market_data_backend, &market_data_frontend) {
+                    error!("Error proxying market data: {}", e);
+                }
+            }
+        }
+    }
+
+    fn handle_client_submission(
+        client_socket: &zmq::Socket,
+        ingestor_socket: &zmq::Socket,
+    ) -> Result<()> {
+        // Receive from client: [identity][empty][version][message]
+        let identity = client_socket.recv_bytes(0)?;
+        let _empty = client_socket.recv_bytes(0)?;
+        let version_frame = client_socket.recv_bytes(0)?;
+        let message_frame = client_socket.recv_bytes(0)?;
+
+        if version_frame.len() != 1 || version_frame[0] != PROTOCOL_VERSION {
+            warn!("Invalid protocol version from client");
+            return Ok(());
+        }
+
+        if message_frame.is_empty() {
+            warn!("Empty message frame from client");
+            return Ok(());
+        }
+
+        let msg_type = message_frame[0];
+        let payload = &message_frame[1..];
+
+        debug!("Received client submission: type=0x{:02x}, payload_len={}", msg_type, payload.len());
+
+        match msg_type {
+            MSG_TYPE_SUBMIT_REQUEST => {
+                Self::handle_submit_request(
+                    identity,
+                    payload,
+                    client_socket,
+                    ingestor_socket,
+                )?;
+            }
+            _ => {
+                warn!("Unknown message type from client: 0x{:02x}", msg_type);
+            }
+        }
+
+        Ok(())
+    }
+
+    fn handle_submit_request(
+        client_identity: Vec<u8>,
+        payload: &[u8],
+        client_socket: &zmq::Socket,
+        ingestor_socket: &zmq::Socket,
+    ) -> Result<()> {
+        // Parse protobuf request
+        let request = proto::SubmitHistoricalRequest::decode(payload)
+            .context("Failed to parse SubmitHistoricalRequest")?;
+
+        let request_id = request.request_id.clone();
+        let ticker = request.ticker.clone();
+        let client_id = request.client_id.clone();
+
+        info!("Handling request submission: request_id={}, ticker={}, client_id={:?}",
+            request_id, ticker, client_id);
+
+        // Extract exchange prefix from ticker
+        let exchange_prefix = ticker.split(':').next()
+            .map(|s| format!("{}:", s))
+            .unwrap_or_else(|| String::from(""));
+
+        if exchange_prefix.is_empty() {
+            warn!("Ticker '{}' missing exchange prefix", ticker);
+        }
+
+        // Build DataRequest protobuf for ingestors
+        let data_request = proto::DataRequest {
+            request_id: request_id.clone(),
+            r#type: proto::data_request::RequestType::HistoricalOhlc as i32,
+            ticker: ticker.clone(),
+            historical: Some(proto::HistoricalParams {
+                start_time: request.start_time,
+                end_time: request.end_time,
+                period_seconds: request.period_seconds,
+                limit: request.limit,
+            }),
+            realtime: None,
+            client_id: client_id.clone(),
+        };
+
+        let mut data_request_bytes = Vec::new();
+        data_request.encode(&mut data_request_bytes)?;
+
+        // Publish to ingestors with exchange prefix
+        let version_frame = vec![PROTOCOL_VERSION];
+        let mut message_frame = vec![MSG_TYPE_DATA_REQUEST];
+        message_frame.extend_from_slice(&data_request_bytes);
+
+        ingestor_socket.send(&exchange_prefix, zmq::SNDMORE)?;
+        ingestor_socket.send(&version_frame, zmq::SNDMORE)?;
+        ingestor_socket.send(&message_frame, 0)?;
+
+        info!("Published to ingestors: prefix={}, request_id={}", exchange_prefix, request_id);
+
+        // Build SubmitResponse protobuf
+        // NOTE: This topic is DETERMINISTIC based on client-generated values.
+        // Client should have already subscribed to this topic BEFORE sending the request
+        // to prevent race condition where notification arrives before client subscribes.
+        let notification_topic = if let Some(cid) = &client_id {
+            format!("RESPONSE:{}", cid)
+        } else {
+            format!("HISTORY_READY:{}", request_id)
+        };
+
+        let response = proto::SubmitResponse {
+            request_id: request_id.clone(),
+            status: proto::submit_response::SubmitStatus::Queued as i32,
+            error_message: None,
+            notification_topic: notification_topic.clone(),
+        };
+
+        let mut response_bytes = Vec::new();
+        response.encode(&mut response_bytes)?;
+
+        // Send immediate response to client
+        let version_frame = vec![PROTOCOL_VERSION];
+        let mut message_frame = vec![MSG_TYPE_SUBMIT_RESPONSE];
+        message_frame.extend_from_slice(&response_bytes);
+
+        client_socket.send(&client_identity, zmq::SNDMORE)?;
+        client_socket.send(&[] as &[u8], zmq::SNDMORE)?;
+        client_socket.send(&version_frame, zmq::SNDMORE)?;
+        client_socket.send(&message_frame, 0)?;
+
+        info!("Sent SubmitResponse to client: request_id={}, topic={}", request_id, notification_topic);
+
+        // Relay is now DONE with this request - completely stateless!
+        // Client will receive notification via pub/sub when Flink publishes HistoryReadyNotification
+
+        Ok(())
+    }
+
+    fn proxy_subscription(
+        frontend: &zmq::Socket,
+        backend: &zmq::Socket,
+    ) -> Result<()> {
+        // Forward subscription message from XPUB to XSUB
+        let msg = frontend.recv_bytes(0)?;
+        backend.send(&msg, 0)?;
+
+        if !msg.is_empty() {
+            let action = if msg[0] == 1 { "subscribe" } else { "unsubscribe" };
+            let topic = String::from_utf8_lossy(&msg[1..]);
+            debug!("Client {} to topic: {}", action, topic);
+        }
+
+        Ok(())
+    }
+
+    fn proxy_market_data(
+        backend: &zmq::Socket,
+        frontend: &zmq::Socket,
+    ) -> Result<()> {
+        // Forward all messages from XSUB to XPUB (zero-copy proxy)
+        // This includes:
+        // - Regular market data (ticks, OHLC)
+        // - HistoryReadyNotification from Flink
+        loop {
+            let msg = backend.recv_bytes(0)?;
+            let more = backend.get_rcvmore()?;
+
+            if more {
+                frontend.send(&msg, zmq::SNDMORE)?;
+            } else {
+                frontend.send(&msg, 0)?;
+                break;
+            }
+        }
+
+        Ok(())
+    }
+}