Skip to content

Instantly share code, notes, and snippets.

@dcoles
Created October 12, 2021 04:29
Show Gist options
  • Save dcoles/fdfef0d89ec03555b3ef6eb6e5dad259 to your computer and use it in GitHub Desktop.
Save dcoles/fdfef0d89ec03555b3ef6eb6e5dad259 to your computer and use it in GitHub Desktop.
Tokio-based TCP server
// Tokio-based server
use std::io;
use std::iter::FromIterator;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
// Listening server.
struct Listener {
listener: TcpListener,
}
impl Listener {
fn new(listener: TcpListener) -> Self {
Listener { listener }
}
async fn run(&self) {
loop {
let (stream, peer) = match self.listener.accept().await {
Err(err) => {
println!("ERROR: {}", err);
continue;
},
Ok(c) => c,
};
println!("New connection from {}", peer);
let mut handler = Handler::new(stream);
tokio::spawn(async move {
if let Err(e) = handler.run().await {
println!("ERROR: Failed to handle connection from {}: {}", peer, e);
}
println!("Disconnecting from {}", peer);
});
}
}
}
// Per connection handler
struct Handler {
connection: Connection,
heartbeat: tokio::time::Interval,
}
impl Handler {
fn new(stream: TcpStream) -> Self {
let connection = Connection::new(stream);
let heartbeat = tokio::time::interval(Duration::from_secs(1));
Handler { connection, heartbeat }
}
async fn run(&mut self) -> io::Result<()> {
loop {
tokio::select! {
// Heartbeat
_ = self.heartbeat.tick() => {
self.connection.write_frame(b"Heartbeat\n").await?;
},
// Received message
res = self.connection.read_frame() => {
let frame = res?;
if frame.is_empty() {
// Client disconnected
return Ok(());
}
self.connection.write_frame(b"Not implemented!\n").await?;
}
}
}
}
}
// Low-level connection
struct Connection {
stream: TcpStream,
}
impl Connection {
fn new(stream: TcpStream) -> Self {
Connection { stream }
}
async fn read_frame(&mut self) -> io::Result<Vec<u8>> {
let mut buf = [0u8; 1024];
let n = self.stream.read(&mut buf).await?;
Ok(Vec::from_iter(buf[..n].iter().copied()))
}
async fn write_frame(&mut self, frame: &[u8]) -> io::Result<()> {
self.stream.write_all(&frame).await?;
Ok(())
}
}
#[tokio::main]
async fn main() {
let tcp_listener = TcpListener::bind("0.0.0.0:5000").await.unwrap();
let listener = Listener::new(tcp_listener);
listener.run().await;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment