benchmarks: add a version agnostic simple router

Cargo.toml

@@ -3,5 +3,6 @@
 members = [
   "rumqttc",
   "rumqttd",
-  "benchmarks"
+  "benchmarks",
+  "benchmarks/simplerouter",
 ]

benchmarks/simplerouter/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "simplerouter"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+bytes = "1.1.0"
+log = "0.4.14"
+pretty_env_logger = "0.4.0"
+thiserror = "1.0.30"
+tokio = { version = "1.17.0", features = ["net", "sync", "rt-multi-thread", "io-util", "macros"] }

benchmarks/simplerouter/src/bin/simplerouter.rs

@@ -0,0 +1,11 @@
+use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
+
+#[tokio::main]
+async fn main() {
+    pretty_env_logger::init();
+    simplerouter::run(simplerouter::Config {
+        addr: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 1883)),
+    })
+    .await
+    .unwrap();
+}

benchmarks/simplerouter/src/lib.rs

@@ -0,0 +1,122 @@
+use std::{io, net::SocketAddr};
+
+use bytes::BytesMut;
+use log::*;
+use tokio::net::TcpListener;
+
+mod network;
+mod protocol;
+use network::Network;
+use protocol::{v4, v5};
+
+pub struct Config {
+    pub addr: SocketAddr,
+}
+
+pub async fn run(config: Config) -> Result<(), Error> {
+    let listener = TcpListener::bind(config.addr).await?;
+    info!("router: listening on {}", config.addr);
+
+    loop {
+        let (stream, addr) = listener.accept().await?;
+        info!("router: accepted connection from {}", addr);
+        let (network, _) = match Network::read_connect(stream).await {
+            Ok(v) => v,
+            Err(e) => {
+                error!("router: unable to read connect : {}", e);
+                continue;
+            }
+        };
+        info!("connection: sent connack");
+        tokio::spawn(publisher_handle(network));
+    }
+}
+
+async fn publisher_handle(mut network: Network) {
+    let mut payload = BytesMut::with_capacity(2);
+    v4::pingresp::write(&mut payload).unwrap();
+    let pingresp_bytes = payload.split().freeze();
+
+    loop {
+        let packet = match network.poll().await {
+            Ok(packet) => packet,
+            Err(e) => {
+                error!("connection: unable to read packet: {}", e);
+                return;
+            }
+        };
+        match packet {
+            protocol::Packet::V4(packet) => match packet {
+                v4::Packet::Disconnect => {
+                    info!("connection: received disconnect, exiting");
+                    return;
+                }
+                v4::Packet::PingReq => {
+                    if let Err(e) = network.send_data(&pingresp_bytes).await {
+                        error!("unable to send pingresp, exiting : {}", e);
+                        return;
+                    };
+                }
+                v4::Packet::Publish(publish) => {
+                    let pkid = match publish.view_meta() {
+                        Ok(v) => v.2,
+                        Err(e) => {
+                            error!("connection: malformed publish packet : {}", e);
+                            continue;
+                        }
+                    };
+                    payload.reserve(2);
+                    v4::puback::write(pkid, &mut payload).unwrap();
+                    if let Err(e) = network.send_data(&payload.split().freeze()).await {
+                        error!("unable to send puback pkid = {}, exiting : {}", pkid, e);
+                        return;
+                    };
+                }
+                p => {
+                    error!("connection: invalid packet {:?}", p);
+                    continue;
+                }
+            },
+            protocol::Packet::V5(packet) => match packet {
+                v5::Packet::Disconnect => {
+                    info!("connection: received disconnect, exiting");
+                    return;
+                }
+                v5::Packet::PingReq => {
+                    if let Err(e) = network.send_data(&pingresp_bytes).await {
+                        error!("unable to send pingresp, exiting : {}", e);
+                        return;
+                    };
+                }
+                v5::Packet::Publish(publish) => {
+                    let pkid = match publish.view_meta() {
+                        Ok(v) => v.2,
+                        Err(e) => {
+                            error!("connection: malformed publish packet : {}", e);
+                            continue;
+                        }
+                    };
+                    payload.reserve(8);
+                    v5::puback::write(pkid, v5::puback::PubAckReason::Success, None, &mut payload)
+                        .unwrap();
+                    if let Err(e) = network.send_data(&payload.split().freeze()).await {
+                        error!("unable to send puback pkid = {}, exiting : {}", pkid, e);
+                        return;
+                    };
+                }
+                p => {
+                    error!("connection: invalid packet {:?}", p);
+                    continue;
+                }
+            },
+        }
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("MQTT : {0}")]
+    MQTT(#[from] crate::protocol::Error),
+    #[error("i/O : {0}")]
+    IO(#[from] io::Error),
+}

benchmarks/simplerouter/src/network.rs

@@ -0,0 +1,102 @@
+use std::io;
+
+use bytes::{Bytes, BytesMut};
+use log::*;
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    net::TcpStream,
+};
+
+use crate::{
+    protocol::{self, v4, v5, Connect, Packet},
+    Error,
+};
+
+pub(crate) struct Network {
+    stream: TcpStream,
+    buf: BytesMut,
+    protocol_level: u8,
+}
+
+impl Network {
+    pub(crate) async fn read_connect(stream: TcpStream) -> Result<(Self, Connect), Error> {
+        let mut network = Self {
+            stream,
+            buf: BytesMut::with_capacity(4096),
+            protocol_level: 0,
+        };
+        debug!("network: reading from stream");
+        network.stream.read_buf(&mut network.buf).await?;
+        let connect_packet = loop {
+            match protocol::read_first_connect(&mut network.buf, 4096) {
+                Err(protocol::Error::InsufficientBytes(count)) => {
+                    network.read_atleast(count).await?
+                }
+                res => break res?,
+            }
+        };
+        debug!("network: read connect");
+        match &connect_packet {
+            Connect::V4(_) => {
+                network.protocol_level = 4;
+                let mut payload = BytesMut::with_capacity(10);
+                v4::connack::write(v4::connack::ConnectReturnCode::Success, false, &mut payload)?;
+                network.send_data(&payload.split().freeze()).await?;
+            }
+            Connect::V5(_) => {
+                network.protocol_level = 5;
+                let mut payload = BytesMut::with_capacity(10);
+                v5::connack::write(
+                    v5::connack::ConnectReturnCode::Success,
+                    false,
+                    None,
+                    &mut payload,
+                )?;
+                network.send_data(&payload.split().freeze()).await?;
+            }
+        }
+        debug!("network: sent connack");
+        Ok((network, connect_packet))
+    }
+
+    async fn read_atleast(&mut self, count: usize) -> io::Result<()> {
+        let mut len = 0;
+        while len < count {
+            len += self.stream.read_buf(&mut self.buf).await?;
+        }
+        debug!("network: read {} bytes", len);
+
+        Ok(())
+    }
+
+    pub(crate) async fn poll(&mut self) -> Result<protocol::Packet, Error> {
+        loop {
+            match self.protocol_level {
+                4 => match v4::read_mut(&mut self.buf, 4096) {
+                    Err(protocol::Error::InsufficientBytes(count)) => {
+                        self.read_atleast(count).await?;
+                        continue;
+                    }
+                    res => return Ok(Packet::V4(res?)),
+                },
+                5 => match v5::read_mut(&mut self.buf, 4096) {
+                    Err(protocol::Error::InsufficientBytes(count)) => {
+                        self.read_atleast(count).await?;
+                        continue;
+                    }
+                    res => return Ok(Packet::V5(res?)),
+                },
+                // SAFETY: we don't allow changing protocol_level
+                _ => unsafe { std::hint::unreachable_unchecked() },
+            }
+        }
+    }
+
+    pub(crate) async fn send_data(&mut self, data: &Bytes) -> Result<(), Error> {
+        debug!(
+            "network: sent {} bytes",
+            self.stream.write(data.as_ref()).await?
+        );
+        Ok(())
+    }
+}