Refactors

2026-10-03 15:56:29 +05:30
parent a3385e7702
commit c9438d00a6
14 changed files with 358 additions and 401 deletions
@@ -5,5 +5,4 @@ edition = "2024"

 [dependencies]
 tracing = { workspace = true }
-tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] }
-hd-lib = { path = "../hd-lib" }
+hd-lib = {path = "../hd-lib"}
@@ -1,30 +0,0 @@
-use std::{sync::Arc, thread, time::Duration};
-
-use hd_lib::{bus::EventData, error::HError};
-use hd_server::TcpEventBus;
-use tracing::info_span;
-
-fn main() -> Result<(), HError> {
-    tracing_subscriber::fmt()
-        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
-        .init();
-
-    let addr = "0.0.0.0:21368";
-    let server = Arc::new(TcpEventBus::new());
-    server.publish(EventData::new(1, "Hello"))?;
-    server.publish(EventData::new(1, "World"))?;
-
-    let server_clone = server.clone();
-    thread::spawn(move || {
-        let _scope = info_span!("tcp_server").entered();
-        server.start(addr)
-    });
-
-    let _scope = info_span!("health_checks").entered();
-    for i in 0..255 {
-        thread::sleep(Duration::from_secs(5));
-        server_clone.publish(EventData::new(i, "asd"))?;
-    }
-
-    Ok(())
-}
@@ -0,0 +1,185 @@
+use std::{net::TcpStream, time::Duration};
+
+use hd_lib::{
+    bus::{packet::*, *},
+    error::HdError,
+    hd_tcp,
+    thread_pool::ThreadPool,
+};
+use tracing::{debug, error, info, info_span};
+
+#[derive(Copy, Clone, PartialEq)]
+enum HdClientConnectionState {
+    Connect,
+    Deliver,
+    Wait,
+    Disconnect,
+}
+
+struct HdClientConnection {
+    pub addr: String,
+    stream: TcpStream,
+    state: HdClientConnectionState,
+    subscription: Option<SubscriptionInfo>,
+    bus: EventBus,
+}
+
+impl std::fmt::Display for HdClientConnectionState {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let to_write = match self {
+            HdClientConnectionState::Connect => "CONNECT",
+            HdClientConnectionState::Deliver => "DELIVER",
+            HdClientConnectionState::Wait => "WAIT",
+            HdClientConnectionState::Disconnect => "DISCONNECT",
+        };
+
+        write!(f, "{}", to_write)
+    }
+}
+
+impl HdClientConnection {
+    fn new(bus: EventBus, stream: TcpStream) -> Result<Self, HdError> {
+        let addr = stream
+            .peer_addr()
+            .map(|a| a.to_string())
+            .unwrap_or("unknown".to_string());
+
+        stream.set_read_timeout(Some(Duration::from_secs(5)))?;
+
+        let state = HdClientConnectionState::Connect;
+
+        Ok(Self {
+            stream,
+            addr,
+            state,
+            bus,
+            subscription: None,
+        })
+    }
+}
+
+pub(crate) fn handle_client_connection(
+    thread_pool: &ThreadPool,
+    bus: EventBus,
+    stream: TcpStream,
+) -> Result<(), HdError> {
+    thread_pool.execute(move || {
+        let mut client = HdClientConnection::new(bus, stream)?;
+
+        let _scope = info_span!("handle_client", addr = client.addr,).entered();
+
+        loop {
+            if let HdClientConnectionState::Disconnect = client.state {
+                unsubscribe_client(&mut client)?;
+                info!("Closing client connection");
+                return Ok(());
+            }
+
+            let prev_state = client.state;
+            let next_state = match handle_event_bus_client_state(&mut client) {
+                Ok(s) => s,
+                Err(e) => {
+                    error!("Error while handling client {}", e);
+                    HdClientConnectionState::Disconnect
+                }
+            };
+
+            if prev_state != next_state {
+                info!("State change {} -> {}", prev_state, next_state);
+            }
+
+            client.state = next_state;
+        }
+    })
+}
+
+fn handle_event_bus_client_state(client: &mut HdClientConnection) -> Result<HdClientConnectionState, HdError> {
+    match client.state {
+        HdClientConnectionState::Connect => handle_state_connect(client),
+        HdClientConnectionState::Deliver => handle_state_deliver_event(client),
+        HdClientConnectionState::Wait => handle_state_wait_event(client),
+        HdClientConnectionState::Disconnect => Ok(HdClientConnectionState::Disconnect),
+    }
+}
+
+fn handle_state_connect(client: &mut HdClientConnection) -> Result<HdClientConnectionState, HdError> {
+    match hd_tcp::read_packet(&mut client.stream)? {
+        Packet::Connect(connect_packet) => {
+            info!("Received connect");
+            let next = match connect_packet.mode {
+                ConnectMode::Publish => HdClientConnectionState::Wait,
+                ConnectMode::Subscribe => {
+                    let subscription = client.bus.subscribe()?;
+                    client.subscription = Some(subscription);
+                    HdClientConnectionState::Deliver
+                }
+            };
+
+            Ok(next)
+        }
+        p => {
+            error!(
+                "Received wrong packet of type {}, expected {}. Will disconnect",
+                get_packet_type(&p),
+                PacketType::Connect
+            );
+            Ok(HdClientConnectionState::Disconnect)
+        }
+    }
+}
+
+fn handle_state_deliver_event(client: &mut HdClientConnection) -> Result<HdClientConnectionState, HdError> {
+    if let Some(subscription) = &mut client.subscription {
+        // TODO: When client sends unsubscribe, server shouldn't take forever
+        // to ack that.
+        let Ok(evt) = subscription.event_recv.recv() else {
+            error!("Failed to receive a message for this subscription, disconnecting");
+            return Ok(HdClientConnectionState::Disconnect);
+        };
+
+        let packet = Packet::create_deliver_packet(evt);
+        hd_tcp::write_packet(&mut client.stream, &packet)?;
+        debug!("Published event to client, waiting for ack");
+
+        let response_packet = hd_tcp::read_packet(&mut client.stream)?;
+        if let Packet::Disconnect = &response_packet {
+            info!("Client disconnected");
+            // TODO: Allow unsubscribe, and actually disconnect
+            return Ok(HdClientConnectionState::Disconnect);
+        }
+
+        debug!("Publish ack received, message is settled");
+    }
+
+    Ok(HdClientConnectionState::Deliver)
+}
+
+fn unsubscribe_client(client: &mut HdClientConnection) -> Result<(), HdError> {
+    if let Some(subscription) = &client.subscription {
+        client.bus.unsubscribe(subscription.id)?;
+        client.subscription = None;
+    };
+
+    Ok(())
+}
+
+fn handle_state_wait_event(client: &mut HdClientConnection) -> Result<HdClientConnectionState, HdError> {
+    match hd_tcp::read_packet(&mut client.stream)? {
+        Packet::Publish(publish_packet) => {
+            debug!("Received event from client, publishing to bus");
+            let data = String::from_utf8(publish_packet.data_utf8)?;
+            let evt = EventData::new(publish_packet.event_type, &data);
+            client.bus.publish(evt)?;
+
+            Ok(HdClientConnectionState::Wait)
+        }
+        p => {
+            error!(
+                "Received wrong packet of type {}, expected {}. Will disconnect",
+                get_packet_type(&p),
+                PacketType::Connect
+            );
+            Ok(HdClientConnectionState::Disconnect)
+        }
+    }
+}
@@ -1,29 +1,21 @@
-use std::{
-    net::TcpStream,
-    sync::{Arc, mpsc::Receiver},
-    time::Duration,
-};
+mod connection;

 use hd_lib::{
-    bus::{
-        Event, EventBus, EventData,
-        packet::{self, Packet, PacketType, PublishPacket, get_packet_type},
-    },
-    error::HError,
-    hd_tcp,
+    bus::{EventBus, EventData},
+    error::HdError,
    thread_pool::ThreadPool,
 };

-use tracing::{debug, error, info, info_span};
+use tracing::{error, info};

 use std::net::TcpListener;

-pub struct TcpEventBus {
+pub struct HdBus {
    bus: EventBus,
    pool: ThreadPool,
 }

-impl TcpEventBus {
+impl HdBus {
    pub fn new() -> Self {
        let pool = ThreadPool::new(10);
        Self {
@@ -32,18 +24,18 @@ impl TcpEventBus {
        }
    }

-    pub fn publish(&self, evt: EventData) -> Result<(), HError> {
+    pub fn publish(&self, evt: EventData) -> Result<(), HdError> {
        self.bus.publish(evt)
    }
 }

-impl TcpEventBus {
-    pub fn start(&self, addr: &'static str) -> Result<(), HError> {
+impl HdBus {
+    pub fn start(&self, addr: &'static str) -> Result<(), HdError> {
        let listener = TcpListener::bind(addr)?;
        info!("Heimdall running at {}", addr);
        for stream in listener.incoming() {
            match stream {
-                Ok(stream) => handle_event_bus_client(&self.pool, self.bus.clone(), stream)?,
+                Ok(stream) => connection::handle_client_connection(&self.pool, self.bus.clone(), stream)?,
                Err(e) => error!("TCP connection failed {}", e),
            }
        }
@@ -51,199 +43,3 @@ impl TcpEventBus {
        Ok(())
    }
 }
-
-#[derive(Copy, Clone, PartialEq)]
-enum EventBusClientConnectionState {
-    Connect,
-    Command,
-    SendEvent,
-    WaitEvent,
-    Disconnect,
-}
-
-struct EventBusClientConnection {
-    pub addr: String,
-    stream: TcpStream,
-    state: EventBusClientConnectionState,
-    subscription: Option<Receiver<Arc<Event>>>,
-    bus: EventBus,
-}
-
-impl EventBusClientConnection {
-    fn new(bus: EventBus, stream: TcpStream) -> Result<Self, HError> {
-        let addr = stream
-            .peer_addr()
-            .map(|a| a.to_string())
-            .unwrap_or("unknown".to_string());
-
-        stream.set_read_timeout(Some(Duration::from_secs(5)))?;
-
-        let state = EventBusClientConnectionState::Connect;
-
-        Ok(Self {
-            stream,
-            addr,
-            state,
-            bus,
-            subscription: None,
-        })
-    }
-}
-
-pub fn handle_event_bus_client(thread_pool: &ThreadPool, bus: EventBus, stream: TcpStream) -> Result<(), HError> {
-    thread_pool.execute(move || {
-        let mut client = EventBusClientConnection::new(bus, stream)?;
-
-        let _scope = info_span!("handle_client", addr = client.addr,).entered();
-
-        loop {
-            if let EventBusClientConnectionState::Disconnect = client.state {
-                info!("Closing connection due to invalid input");
-                return Ok(());
-            }
-
-            let prev_state = client.state;
-            handle_event_bus_client_state(&mut client)
-                .inspect_err(|e| info!("Closing connection due to error {}", e))?;
-            if prev_state != client.state {
-                info!("State change {} -> {}", prev_state, client.state);
-            }
-        }
-    })
-}
-
-fn handle_event_bus_client_state(client: &mut EventBusClientConnection) -> Result<(), HError> {
-    let next_state = match client.state {
-        EventBusClientConnectionState::Connect => handle_state_connect(client),
-        EventBusClientConnectionState::Command => handle_state_command(client),
-        EventBusClientConnectionState::SendEvent => handle_state_send_event(client),
-        EventBusClientConnectionState::WaitEvent => handle_state_wait_event(client),
-        EventBusClientConnectionState::Disconnect => Ok(EventBusClientConnectionState::Disconnect),
-    }?;
-
-    client.state = next_state;
-    Ok(())
-}
-
-fn handle_state_connect(client: &mut EventBusClientConnection) -> Result<EventBusClientConnectionState, HError> {
-    match hd_tcp::read_packet(&mut client.stream)? {
-        Packet::Connect => {
-            info!("Received connect, acking");
-
-            hd_tcp::write_ack(&mut client.stream, PacketType::Connect)?;
-            info!("Connect ack complete");
-
-            Ok(EventBusClientConnectionState::Command)
-        }
-        p => {
-            error!(
-                "Received wrong packet of type {}, expected {}. Will disconnect",
-                get_packet_type(&p),
-                PacketType::Connect
-            );
-            Ok(EventBusClientConnectionState::Disconnect)
-        }
-    }
-}
-
-fn handle_state_command(client: &mut EventBusClientConnection) -> Result<EventBusClientConnectionState, HError> {
-    match hd_tcp::read_packet(&mut client.stream)? {
-        // TODO: Read cursor
-        Packet::Subscribe(_) => {
-            info!("Received subscribe");
-            let subscription = client.bus.subscribe()?;
-
-            info!("Subscribed to the event bus");
-            client.subscription = Some(subscription);
-
-            // TODO: Return error if any
-            hd_tcp::write_ack(&mut client.stream, PacketType::Subscribe)?;
-            info!("Subscribe ack complete");
-            Ok(EventBusClientConnectionState::SendEvent)
-        }
-        Packet::Publish(publish_packet) => {
-            info!("Received publish packet");
-            publish_packet_to_bus(client, publish_packet)?;
-
-            // TODO: Return error if any
-            hd_tcp::write_ack(&mut client.stream, PacketType::Publish)?;
-            debug!("Publish ack complete");
-            Ok(EventBusClientConnectionState::WaitEvent)
-        }
-        p => {
-            error!(
-                "Received wrong packet of type {}, expected {}. Will disconnect",
-                get_packet_type(&p),
-                PacketType::Connect
-            );
-            Ok(EventBusClientConnectionState::Disconnect)
-        }
-    }
-}
-
-fn handle_state_send_event(client: &mut EventBusClientConnection) -> Result<EventBusClientConnectionState, HError> {
-    if let Some(subscription) = &mut client.subscription {
-        // TODO: When client sends unsubscribe, server shouldn't take forever
-        // to ack that.
-        let Ok(evt) = subscription.recv() else {
-            error!("Failed to receive a message for this subscription, disconnecting");
-            return Ok(EventBusClientConnectionState::Disconnect);
-        };
-
-        let packet = Packet::create_send_event_packet(evt);
-        hd_tcp::write_packet(&mut client.stream, &packet)?;
-        debug!("Published event to client, waiting for ack");
-
-        let ack_packet = hd_tcp::read_ack(&mut client.stream, &[PacketType::SendEvent, PacketType::Unsubscribe])?;
-        if packet::get_packet_type(&ack_packet) == PacketType::Unsubscribe {
-            info!("Subscriber unsubscribed");
-            return Ok(EventBusClientConnectionState::Command);
-        }
-
-        debug!("Publish ack received, message is settled");
-    }
-
-    Ok(EventBusClientConnectionState::SendEvent)
-}
-
-fn handle_state_wait_event(client: &mut EventBusClientConnection) -> Result<EventBusClientConnectionState, HError> {
-    match hd_tcp::read_packet(&mut client.stream)? {
-        Packet::Publish(publish_packet) => {
-            debug!("Received event from client, publishing to bus");
-            publish_packet_to_bus(client, publish_packet)?;
-
-            hd_tcp::write_ack(&mut client.stream, PacketType::Publish)?;
-            debug!("Ack receive event complete");
-            Ok(EventBusClientConnectionState::WaitEvent)
-        }
-        p => {
-            error!(
-                "Received wrong packet of type {}, expected {}. Will disconnect",
-                get_packet_type(&p),
-                PacketType::Connect
-            );
-            Ok(EventBusClientConnectionState::Disconnect)
-        }
-    }
-}
-
-fn publish_packet_to_bus(client: &mut EventBusClientConnection, packet: PublishPacket) -> Result<(), HError> {
-    let data = String::from_utf8(packet.data_utf8)?;
-    let evt = EventData::new(packet.event_type, &data);
-    client.bus.publish(evt)?;
-    Ok(())
-}
-
-impl std::fmt::Display for EventBusClientConnectionState {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let to_write = match self {
-            EventBusClientConnectionState::Connect => "CONNECT",
-            EventBusClientConnectionState::Command => "COMMAND",
-            EventBusClientConnectionState::SendEvent => "SEND_EVENT",
-            EventBusClientConnectionState::WaitEvent => "WAIT_EVENT",
-            EventBusClientConnectionState::Disconnect => "DISCONNECT",
-        };
-
-        write!(f, "{}", to_write)
-    }
-}