Created
October 12, 2021 04:29
-
-
Save dcoles/fdfef0d89ec03555b3ef6eb6e5dad259 to your computer and use it in GitHub Desktop.
Tokio-based TCP server
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
// 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