geoffreycopin
Posted on May 20, 2023
In this episode, we'll extend our server to serve static files.
We'll also refactor our code to support connection reuse, and we'll implement a graceful shutdown mechanism.
If your didn't follow the previous episode, you can find the code on GitHub.
As we will use new dependencies, we'll need to update our Cargo.toml
file:
cargo add clap tokio-util futures
Connection reuse
Under HTTP/1.0, a separate TCP connection is established for each request/response pair.
This is inefficient, as it requires a new TCP handshake for each request. HTTP/1.1 introduced connection reuse,
which allows multiple requests to be sent over the same TCP connection. This mechanism is also necessary
to support request pipelining, which we'll see in a future episode.
We'll keep waiting for new requests on the same connection until the client closes it, unless the client sets
the Connection: close
header. In that case, we'll close the connection after sending the response.
All we have to do to implement this change is to wrap our client handling code in a loop:
// main.rs
// [...]
info!(?addr, "new connection");
loop {
let req = match req::parse_request(&mut stream).await {
Ok(req) => {
info!(?req, "incoming request");
req
}
Err(e) => {
error!(?e, "failed to parse request");
break;
}
};
let close_connection =
req.headers.get("Connection") == Some(&"close".to_string());
let resp = resp::Response::from_html(
resp::Status::NotFound,
include_str!("../static/404.html"),
);
resp.write(&mut stream).await.unwrap();
if close_connection {
break;
}
}
// [...]
Serving static files
Answering every request with the same 'Not found' page was ok for a start, but the time has come to serve
some real content. In this section, we'll implement a handler that serves static files from a directory.
This directory will be either the current working directory, or a directory specified by the user.
As our CLI is getting more complex, we'll use the clap crate to parse the
command line arguments.
// args.rs
use std::path::PathBuf;
use clap::Parser;
#[derive(Parser, Debug)]
pub struct Args {
#[arg(short, long, default_value_t = 8080)]
pub port: u16,
#[arg(short, long)]
pub root: Option<PathBuf>,
}
// main.rs
use clap::Parser;
// [...]
#[tokio::main]
async fn main() -> anyhow::Result<()> {
// [...]
let args = args::Args::parse();
let port = args.port;
let listener = TcpListener::bind(format!("0.0.0.0:{port}")).await.unwrap();
// [...]
}
In order to simplify our handling code, we'll also refactor our Response
struct
by
using a trait object instead of a generic type parameter. This will allow use to use the same
Response
type for all our handlers, even if they don't return the same type of response.
We'll also add a new from_file
constructor to our Response
type, which will allow us to
create a response struct
from a file on disk. At this stage, we will not implement any
kind of sophisticated content negotiation, so we'll just set the Content-Type
header to
a mime type based on the file extension.
// resp.rs
pub struct Response {
pub status: Status,
pub headers: HashMap<String, String>,
pub data: Box<dyn AsyncRead + Unpin + Send>,
}
impl Response {
// [...]
pub async fn from_file(path: &Path, file: File) -> anyhow::Result<Response> {
let headers = hashmap! {
"Content-Length".to_string() => file.metadata().await?.len().to_string(),
"Content-Type".to_string() => mime_type(path).to_string(),
};
Ok(Response {
headers,
status: Status::Ok,
data: Box::new(file),
})
}
// [...]
}
fn mime_type(path: &Path) -> &str {
match path.extension().and_then(|ext| ext.to_str()) {
Some("html") => "text/html",
Some("css") => "text/css",
Some("js") => "text/javascript",
Some("png") => "image/png",
Some("jpg") => "image/jpeg",
Some("gif") => "image/gif",
_ => "application/octet-stream",
}
}
These preparatory steps allow us to implement our static file handler in a few lines of code:
// handler.rs
use std::{env::current_dir, io, path::PathBuf};
use crate::{
req::Request,
resp::{Response, Status},
};
#[derive(Debug, Clone)]
pub struct StaticFileHandler {
root: PathBuf,
}
impl StaticFileHandler {
pub fn in_current_dir() -> io::Result<StaticFileHandler> {
current_dir().map(StaticFileHandler::with_root)
}
pub fn with_root(root: PathBuf) -> StaticFileHandler {
StaticFileHandler { root }
}
pub async fn handle(&self, request: Request) -> anyhow::Result<Response> {
let path = self.root.join(request.path.strip_prefix('/').unwrap());
if !path.is_file() {
return Ok(Response::from_html(
Status::NotFound,
include_str!("../static/404.html"),
));
}
let file = tokio::fs::File::open(&path).await?;
Response::from_file(&path, file).await
}
}
// main.rs
#[tokio::main]
async fn main() -> anyhow::Result<()> {
// [...]
let handler = args
.root
.map(handler::StaticFileHandler::with_root)
.unwrap_or_else(|| {
handler::StaticFileHandler::in_current_dir().expect("failed to get current dir")
});
// [...]
match handler.handle(req).await {
Ok(resp) => {
resp.write(stream).await.unwrap();
}
Err(e) => {
error!(?e, "failed to handle request");
return Ok(false);
}
};
// [...]
}
After copying these files to the static
directory, we can now serve them with our server:
cargo run -- --root static
Pressing Ctrl+C
will stop the server, but in a rather abrupt way: the server will not log any
message regarding the shutdown, and will not wait for the pending requests to be processed before
closing the connection. This is not a big deal for a toy server, but in a real world application,
we would want to handle this more gracefully.
Graceful shutdown
The most straightforward way to stop our server would be to collect the client handling tasks
into a JoinSet
and abort them when we receive a SIGINT
signal. However, by doing so, we would
have no way to wait for the pending requests to be processed before exiting the program.
Instead, we'll use a CancellationToken
to signal that the server should:
- stop accepting new connection,
- stop processing requests from the current connections
We'll use the two methods provided by the CancellationToken
:
-
cancel
: request cancellation from the main thread when the user pressesCtrl+C
-
cancelled
: return a future that resolves when the cancellation is requested. Used in conjunction withselect!
, this will allow us to stop the client handling tasks when the cancellation is requested.
Our main
function now looks like this:
// [...]
async fn main() -> anyhow::Result<()> {
// [...]
let cancel_token = CancellationToken::new();
tokio::spawn({
let cancel_token = cancel_token.clone();
async move {
if let Ok(()) = signal::ctrl_c().await {
info!("received Ctrl-C, shutting down");
cancel_token.cancel();
}
}
});
// [...]
let mut tasks = Vec::new();
loop {
let cancel_token = cancel_token.clone();
tokio::select! {
Ok((stream, addr)) = listener.accept() => {
let handler = handler.clone();
let client_task = tokio::spawn(async move {
if let Err(e) = handle_client(cancel_token, stream, addr, &handler).await {
error!(?e, "failed to handle client");
}
});
tasks.push(client_task);
},
_ = cancel_token.cancelled() => {
info!("stop listening");
break;
}
}
}
futures::future::join_all(tasks).await;
Ok(())
}
async fn handle_client(
cancel_token: CancellationToken,
stream: TcpStream,
addr: SocketAddr,
handler: &handler::StaticFileHandler,
) -> anyhow::Result<()> {
let mut stream = BufStream::new(stream);
info!(?addr, "new connection");
loop {
tokio::select! {
req = req::parse_request(&mut stream) => {
match req {
Ok(req) => {
info!(?req, "incoming request");
let close_conn = handle_req(req, &handler, &mut stream).await?;
if close_conn {
break;
}
}
Err(e) => {
error!(?e, "failed to parse request");
break;
}
}
}
_ = cancel_token.cancelled() => {
info!(?addr, "closing connection");
break;
}
}
}
Ok(())
}
async fn handle_req<S: AsyncWrite + Unpin>(
req: req::Request,
handler: &handler::StaticFileHandler,
stream: &mut S,
) -> anyhow::Result<bool> {
let close_connection = req.headers.get("Connection") == Some(&"close".to_string());
match handler.handle(req).await {
Ok(resp) => {
resp.write(stream).await.unwrap();
}
Err(e) => {
error!(?e, "failed to handle request");
return Ok(false);
}
};
Ok(close_connection)
}
We can now stop the server by pressing Ctrl+C
, and the server will wait for the pending requests
to be processed before exiting.
The full source code for this part is available here.
Looking for a Rust dev? Let’s get in touch!
Posted on May 20, 2023
Join Our Newsletter. No Spam, Only the good stuff.
Sign up to receive the latest update from our blog.