diff --git a/Cargo.lock b/Cargo.lock index f1bc9f8..0b8e2d2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -23,6 +23,7 @@ version = "0.1.0" dependencies = [ "hd-lib", "tracing", + "tracing-subscriber", ] [[package]] diff --git a/crates/hd-client/Cargo.toml b/crates/hd-client/Cargo.toml index e4b56ed..875005f 100644 --- a/crates/hd-client/Cargo.toml +++ b/crates/hd-client/Cargo.toml @@ -6,3 +6,4 @@ edition = "2024" [dependencies] tracing = { workspace = true } hd-lib = { path = "../hd-lib" } +tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] } diff --git a/crates/hd-client/src/bin/client.rs b/crates/hd-client/src/bin/client.rs index 12e92ed..228baf1 100644 --- a/crates/hd-client/src/bin/client.rs +++ b/crates/hd-client/src/bin/client.rs @@ -1,6 +1,10 @@ use hd_client::{EventBusClient, EventBusClientError}; fn main() -> Result<(), EventBusClientError> { + tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .init(); + let addr = "0.0.0.0:21368"; let mut client = EventBusClient::start(addr)?; client.subscribe() diff --git a/crates/hd-client/src/lib.rs b/crates/hd-client/src/lib.rs index b0dcf5e..e270a2b 100644 --- a/crates/hd-client/src/lib.rs +++ b/crates/hd-client/src/lib.rs @@ -1,11 +1,22 @@ -use std::{ - io::{Read, Write}, - net::TcpStream, +use std::net::TcpStream; + +use hd_lib::{ + bus::packet::{self, Packet, PacketType}, + error::HError, + hd_tcp, }; +use tracing::{debug, info}; #[derive(Debug)] pub enum EventBusClientError { TcpStreamError(std::io::Error), + HError(HError), +} + +impl From for EventBusClientError { + fn from(value: HError) -> Self { + Self::HError(value) + } } pub struct EventBusClient { @@ -14,43 +25,28 @@ pub struct EventBusClient { impl EventBusClient { pub fn start(addr: &'static str) -> Result { - let stream = TcpStream::connect(addr).map_err(EventBusClientError::TcpStreamError)?; - println!("Event bus client connected to {}", addr); + let mut stream = TcpStream::connect(addr).map_err(EventBusClientError::TcpStreamError)?; + hd_tcp::write_packet(&mut stream, &Packet::create_connect_packet())?; + hd_tcp::read_ack(&mut stream, &[PacketType::Connect])?; + info!("Event bus client connected to {}", addr); Ok(Self { tcp_stream: stream }) } pub fn subscribe(&mut self) -> Result<(), EventBusClientError> { - let bytes = [1]; - self.tcp_stream - .write_all(&bytes) - .map_err(EventBusClientError::TcpStreamError)?; + debug!("Going to subscribe"); + hd_tcp::write_packet(&mut self.tcp_stream, &Packet::create_subscribe_packet(0))?; + hd_tcp::read_ack(&mut self.tcp_stream, &[PacketType::Subscribe])?; - self.tcp_stream.flush().map_err(EventBusClientError::TcpStreamError)?; - - let mut buf = vec![0; 1024]; loop { - let Ok(bytes_read) = self.tcp_stream.read(&mut buf) else { - return Ok(()); - }; - - if bytes_read == 0 { - eprintln!("Connection closed"); + let packet = hd_tcp::read_packet(&mut self.tcp_stream)?; + if let Packet::Disconnect = packet { + info!("Disconnecting from the server"); return Ok(()); } - println!("MESSAGE BEGIN"); - for b in &buf[0..bytes_read] { - print!("[{}] ", b); - } - println!("MESSAGE END"); - - // ACK - self.tcp_stream - .write_all(&bytes) - .map_err(EventBusClientError::TcpStreamError)?; - - self.tcp_stream.flush().map_err(EventBusClientError::TcpStreamError)?; + debug!("Received packet ... {}", packet::get_packet_type(&packet)); + hd_tcp::write_ack(&mut self.tcp_stream, PacketType::SendEvent)?; } } } diff --git a/crates/hd-lib/src/bus/mod.rs b/crates/hd-lib/src/bus/mod.rs index 4163293..a8e70b0 100644 --- a/crates/hd-lib/src/bus/mod.rs +++ b/crates/hd-lib/src/bus/mod.rs @@ -1,6 +1,6 @@ pub mod packet; -use tracing::{info, info_span}; +use tracing::{debug_span, info, info_span}; use crate::error::HError; use std::{ @@ -95,6 +95,7 @@ impl EventBus { for s in &bus.subscription_handles { if let Err(e) = s.new_event_signal.send(1) { + // TODO: Multiple failures should automatically unsubscribe tracing::error!("Failed to notify {}", e); } } @@ -195,7 +196,7 @@ impl Subscription { fn start_delivery(mut subscription: Subscription, bus: EventBus) -> JoinHandle> { thread::spawn(move || { loop { - let _scope = info_span!("subscription_delivery", subscription.id).entered(); + let _scope = debug_span!("subscription_delivery", subscription.id).entered(); loop { if let Ok(evt) = bus.get_next_event(subscription.cursor) { info!("Delivering event {} ", evt.id); diff --git a/crates/hd-lib/src/bus/packet.rs b/crates/hd-lib/src/bus/packet.rs index 15bfc7a..5f47f7f 100644 --- a/crates/hd-lib/src/bus/packet.rs +++ b/crates/hd-lib/src/bus/packet.rs @@ -18,6 +18,23 @@ pub enum PacketType { Disconnect, } +impl Display for PacketType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let to_write = match self { + PacketType::Connect => "CONNECT", + PacketType::Subscribe => "SUBSCRIBE", + PacketType::Unsubscribe => "UNSUBSCRIBE", + PacketType::Publish => "PUBLISH", + PacketType::SendEvent => "SEND_EVENT", + PacketType::Settle => "SETTLE_EVENT", + PacketType::Ack => "ACK", + PacketType::Disconnect => "DISCONNECT", + }; + + write!(f, "{}", to_write) + } +} + pub enum Packet { Connect, Subscribe(SubscribePacket), @@ -49,11 +66,15 @@ pub struct SettlePacket { } pub struct AckPacket { - packet_type: PacketType, + pub packet_type: PacketType, } impl Packet { - pub fn create_connect_packet(cursor: u64) -> Packet { + pub fn create_connect_packet() -> Packet { + Packet::Connect + } + + pub fn create_subscribe_packet(cursor: u64) -> Packet { Packet::Subscribe(SubscribePacket { cursor }) } @@ -115,7 +136,7 @@ pub struct ReadPacketState { pub fn start_read_packet(buf: &mut [u8; packet_constants::HEADER_SIZE]) -> Result { let packet_type = get_packet_type_from_byte(buf[0])?; - let required_buffer_size = read_u16(&buf[1..3]); + let required_buffer_size = read_u16(&buf[1..3])?; Ok(ReadPacketState { required_buffer_size, @@ -133,31 +154,29 @@ pub fn read_packet(buf: &[u8], read_state: ReadPacketState) -> Result Ok(Packet::Unsubscribe), PacketType::Publish => { let event_type = buf[0]; - let len = read_u16(&buf[1..]); + let event_len = read_u16(&buf[1..3])?; // TODO: Share underlying bytes without copying let publish_packet = PublishPacket { event_type, - data_utf8: buf[3..len].to_vec(), + data_utf8: buf[3..3 + event_len].to_vec(), }; Ok(Packet::Publish(publish_packet)) } PacketType::SendEvent => { - let event_id = read_u64(&buf)?; + let event_id = read_u64(&buf[0..8])?; let event_type = buf[8]; - let len = usize::from_le_bytes([buf[9], buf[10], 0, 0, 0, 0, 0, 0]); + let event_len = read_u16(&buf[9..11])?; let send_event_packet = SendEventPacket { event_id, event_type, - data_utf8: buf[11..len].to_vec(), + data_utf8: buf[11..11 + event_len].to_vec(), }; Ok(Packet::SendEvent(send_event_packet)) } PacketType::Settle => { - let mut le_bytes = [0u8; 8]; - le_bytes.copy_from_slice(&buf[0..8]); - let cursor = u64::from_le_bytes(le_bytes); + let cursor = read_u64(&buf[0..8])?; Ok(Packet::Settle(SettlePacket { cursor })) } PacketType::Ack => { @@ -214,7 +233,7 @@ fn write_data_frame(data_frame: &mut [u8], packet: &Packet) { write_event_data( send_event_packet.event_type, &send_event_packet.data_utf8, - &mut data_frame[1..], + &mut data_frame[8..], ); } Packet::Settle(settle_packet) => { @@ -231,12 +250,12 @@ fn write_event_data(event_type: u8, data_utf8: &[u8], buf: &mut [u8]) { buf[0] = event_type; write_u16(data_utf8.len(), &mut buf[1..3]); let available_length = buf.len() - 3; - buf[3..].copy_from_slice(&data_utf8[..available_length]); + buf[3..].copy_from_slice(&data_utf8[0..available_length]); } fn write_u64(v64: u64, data_frame: &mut [u8]) { let bytes = v64.to_le_bytes(); - data_frame[0..8].copy_from_slice(&bytes); + data_frame[0..8].copy_from_slice(&bytes[0..8]); } fn write_u16(mut length: usize, buf: &mut [u8]) { @@ -245,17 +264,25 @@ fn write_u16(mut length: usize, buf: &mut [u8]) { } let bytes = length.to_le_bytes(); - buf[0..2].copy_from_slice(&bytes); + buf[0..2].copy_from_slice(&bytes[0..2]); } fn read_u64(buf: &[u8]) -> Result { let mut le_bytes = [0u8; 8]; + if le_bytes.len() != buf.len() { + return Err(PacketError::PacketProtocolError); + } + le_bytes.copy_from_slice(buf); Ok(u64::from_le_bytes(le_bytes)) } -fn read_u16(buf: &[u8]) -> usize { - usize::from_le_bytes([buf[0], buf[1], 0, 0, 0, 0, 0, 0]) +fn read_u16(buf: &[u8]) -> Result { + if buf.len() != 2 { + return Err(PacketError::PacketProtocolError); + } + + Ok(usize::from_le_bytes([buf[0], buf[1], 0, 0, 0, 0, 0, 0])) } pub fn get_packet_type(packet: &Packet) -> PacketType { @@ -286,13 +313,14 @@ fn get_byte_from_packet_type(packet_type: &PacketType) -> u8 { fn get_packet_type_from_byte(byte: u8) -> Result { let packet_type = match byte { - 0 => Some(PacketType::Subscribe), - 1 => Some(PacketType::Unsubscribe), - 2 => Some(PacketType::Publish), - 3 => Some(PacketType::SendEvent), - 4 => Some(PacketType::Settle), - 5 => Some(PacketType::Ack), - 6 => Some(PacketType::Disconnect), + 0 => Some(PacketType::Connect), + 1 => Some(PacketType::Subscribe), + 2 => Some(PacketType::Unsubscribe), + 3 => Some(PacketType::Publish), + 4 => Some(PacketType::SendEvent), + 5 => Some(PacketType::Settle), + 6 => Some(PacketType::Ack), + 7 => Some(PacketType::Disconnect), _ => None, }; diff --git a/crates/hd-lib/src/lib.rs b/crates/hd-lib/src/lib.rs index a53ebec..1662807 100644 --- a/crates/hd-lib/src/lib.rs +++ b/crates/hd-lib/src/lib.rs @@ -1,3 +1,62 @@ +pub mod bus; pub mod error; pub mod thread_pool; -pub mod bus; + +pub mod hd_tcp { + use std::{ + io::{Read, Write}, + net::TcpStream, + }; + + use crate::{ + bus::packet::{self, Packet, PacketType}, + error::HError, + }; + + pub fn read_packet(stream: &mut TcpStream) -> Result { + let mut header_buf = [0u8; 3]; + stream.read_exact(&mut header_buf)?; + let read_state = packet::start_read_packet(&mut header_buf)?; + + let mut buf = vec![0u8; read_state.required_buffer_size]; + if read_state.required_buffer_size > 0 { + stream.read_exact(&mut buf)?; + } + let packet = packet::read_packet(&buf[0..], read_state)?; + + if let Packet::Disconnect = packet { + return Err(HError::PeerDisconnect); + }; + + Ok(packet) + } + + pub fn read_ack(stream: &mut TcpStream, packet_types: &[PacketType]) -> Result { + let packet = read_packet(stream)?; + if let Packet::Ack(ack_packet) = &packet { + for packet_type in packet_types.iter() { + if *packet_type == ack_packet.packet_type { + return Ok(packet); + } + } + } + + Err(HError::ProtocolError) + } + + pub fn write_ack(stream: &mut TcpStream, packet_type: PacketType) -> Result<(), HError> { + let packet = Packet::create_ack_packet(packet_type); + write_packet(stream, &packet) + } + + pub fn write_packet(stream: &mut TcpStream, packet: &Packet) -> Result<(), HError> { + let write_state = packet::start_write_packet(packet); + let mut buf = vec![0u8; write_state.required_buffer_size]; + packet::write_packet(&mut buf, packet, write_state)?; + + stream.write_all(&mut buf)?; + stream.flush()?; + + Ok(()) + } +} diff --git a/crates/hd-server/src/bin/server.rs b/crates/hd-server/src/bin/server.rs index af41f64..08a2801 100644 --- a/crates/hd-server/src/bin/server.rs +++ b/crates/hd-server/src/bin/server.rs @@ -10,14 +10,17 @@ fn main() -> Result<(), HError> { .init(); let addr = "0.0.0.0:21368"; - let _scope = info_span!("example_server").entered(); 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 || server.start(addr)); + 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"))?; diff --git a/crates/hd-server/src/lib.rs b/crates/hd-server/src/lib.rs index d6a69e6..60ee5f1 100644 --- a/crates/hd-server/src/lib.rs +++ b/crates/hd-server/src/lib.rs @@ -1,5 +1,4 @@ use std::{ - io::{Read, Write}, net::TcpStream, sync::{Arc, mpsc::Receiver}, time::Duration, @@ -8,13 +7,14 @@ use std::{ use hd_lib::{ bus::{ Event, EventBus, EventData, - packet::{self, Packet, PacketType}, + packet::{self, Packet, PacketType, PublishPacket, get_packet_type}, }, error::HError, + hd_tcp, thread_pool::ThreadPool, }; -use tracing::{Span, info, info_span}; +use tracing::{debug, error, info, info_span}; use std::net::TcpListener; @@ -40,11 +40,11 @@ impl TcpEventBus { impl TcpEventBus { pub fn start(&self, addr: &'static str) -> Result<(), HError> { let listener = TcpListener::bind(addr)?; - info!("Heimdall listening at {}", 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)?, - Err(e) => tracing::error!("TCP connection failed {}", e), + Err(e) => error!("TCP connection failed {}", e), } } @@ -90,57 +90,11 @@ impl EventBusClientConnection { } } -fn read_ack(stream: &mut TcpStream, packet_type: PacketType) -> Result<(), HError> { - let packet = read_packet(stream)?; - if packet_type == packet::get_packet_type(&packet) { - return Ok(()) - } - - Err(HError::ProtocolError) -} - -fn read_packet(stream: &mut TcpStream) -> Result { - let mut header_buf = [0u8; 3]; - stream.read_exact(&mut header_buf)?; - let read_state = packet::start_read_packet(&mut header_buf)?; - - let mut buf = vec![0u8; read_state.required_buffer_size]; - stream.read_exact(&mut buf)?; - let packet = packet::read_packet(&buf[0..], read_state)?; - - if let Packet::Disconnect = packet { - return Err(HError::PeerDisconnect); - }; - - Ok(packet) -} - -fn write_ack(stream: &mut TcpStream, packet_type: PacketType) -> Result<(), HError> { - let packet = Packet::create_ack_packet(packet_type); - write_packet(stream, &packet) -} - -fn write_packet(stream: &mut TcpStream, packet: &Packet) -> Result<(), HError> { - let write_state = packet::start_write_packet(packet); - let mut buf = vec![0u8; write_state.required_buffer_size]; - packet::write_packet(&mut buf, packet, write_state)?; - - stream.write_all(&mut buf)?; - stream.flush()?; - - Ok(()) -} - 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_event_bus_client", - addr = client.addr, - state = %client.state - ) - .entered(); + let _scope = info_span!("handle_client", addr = client.addr,).entered(); loop { if let EventBusClientConnectionState::Disconnect = client.state { @@ -154,8 +108,6 @@ pub fn handle_event_bus_client(thread_pool: &ThreadPool, bus: EventBus, stream: if prev_state != client.state { info!("State change {} -> {}", prev_state, client.state); } - - Span::current().record("state", format!("{}", client.state)); } }) } @@ -163,7 +115,7 @@ pub fn handle_event_bus_client(thread_pool: &ThreadPool, bus: EventBus, stream: 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_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), @@ -174,66 +126,112 @@ fn handle_event_bus_client_state(client: &mut EventBusClientConnection) -> Resul } fn handle_state_connect(client: &mut EventBusClientConnection) -> Result { - let Packet::Connect = read_packet(&mut client.stream)? else { - return Ok(EventBusClientConnectionState::Disconnect); - }; + match hd_tcp::read_packet(&mut client.stream)? { + Packet::Connect => { + info!("Received connect, acking"); - write_ack(&mut client.stream, PacketType::Connect)?; - Ok(EventBusClientConnectionState::Command) + 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 { - match read_packet(&mut client.stream)? { + match hd_tcp::read_packet(&mut client.stream)? { // TODO: Read cursor Packet::Subscribe(_) => { - write_ack(&mut client.stream, PacketType::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) => { - write_ack(&mut client.stream, PacketType::Publish)?; + info!("Received publish packet"); + publish_packet_to_bus(client, publish_packet)?; - // TODO: Possible without copying data_utf8? - let data = String::from_utf8(publish_packet.data_utf8)?; - client.bus.publish(EventData::new(publish_packet.event_type, &data))?; + // TODO: Return error if any + hd_tcp::write_ack(&mut client.stream, PacketType::Publish)?; + debug!("Publish ack complete"); Ok(EventBusClientConnectionState::WaitEvent) } - _ => Ok(EventBusClientConnectionState::Disconnect), + 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 { if let Some(subscription) = &mut client.subscription { - loop { - let Ok(evt) = subscription.recv() else { - return Ok(EventBusClientConnectionState::Disconnect); - }; + // 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); + }; - info!( - "SEND => ID: {} TYPE: {} LEN: {} <=", - evt.id, - evt.data.event_type, - evt.data.data_utf8.len() - ); + 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 packet = Packet::create_send_event_packet(evt); - write_packet(&mut client.stream, &packet)?; - read_ack(&mut client.stream, PacketType::SendEvent)?; + 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::Disconnect) + Ok(EventBusClientConnectionState::SendEvent) } fn handle_state_wait_event(client: &mut EventBusClientConnection) -> Result { - todo!() + 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 write_le_bytes(v64: u64, buf: &mut [u8]) { - let le_bytes = v64.to_le_bytes(); - buf[0..8].copy_from_slice(&le_bytes); +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 {