In this episode, we'll extend our server to serve static files. We'll also refactor our code to support connection reuse, and 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:
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.
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>,
}
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 Responsestruct by using a trait object instead of a generic type parameter. This will allow us 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.
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:
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
}
}
#[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:
We'll use the two methods provided by the CancellationToken:
cancel: request cancellation from the main thread when the user presses Ctrl+C
cancelled: return a future that resolves when the cancellation is requested. Used in conjunction with select!, 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!
