[Rust Guide] 20.3 The Final Project: Graceful Shutdown and Cleanup of the Web Server

rust dev.to

20.3.0 Review

In the previous article, we completed the multithreaded web server, but there are still some improvements we can make. In this article, we will refine the code.

Note: This article continues from 20.2. The Final Project - Multithreaded Web Server. If you want to understand the web server construction process from scratch in detail, please read all the articles in Chapter 20.

20.3.1 Implementing the Drop Trait for ThreadPool

When we want to shut down the server (using the less graceful Ctrl + C method to stop the main thread), all other threads stop immediately too, even if they are still processing requests.

The trait used to manage cleanup is the Drop trait. We only need to write a local drop function to override the default implementation, allowing the threads to finish the work they are currently processing before shutting down. We also need some way to prevent the threads from receiving new requests and prepare for shutdown.

Let’s implement the Drop trait for ThreadPool:

impl Drop for ThreadPool {
    fn drop(&mut self) {
        for worker in &mut self.workers {
            println!("Shutting down worker {}", worker.id);

            worker.thread.join().unwrap();
        }
    }
}
Enter fullscreen mode Exit fullscreen mode

The logic is simply to iterate over each worker and call the join method on the thread field inside worker (see 16.1. Running Code Concurrently with Threads).

Run cargo check:

error[E0507]: cannot move out of `worker.thread` which is behind a mutable reference
   --> src/lib.rs:42:13
    |
 42 |             worker.thread.join().unwrap();
    |             ^^^^^^^^^^^^^ ------ `worker.thread` moved due to this method call
    |             |
    |             move occurs because `worker.thread` has type `JoinHandle<()>`, which does not implement the `Copy` trait
    |
note: `JoinHandle::<T>::join` takes ownership of the receiver `self`, which moves `worker.thread`
   --> /Users/stanyin/.rustup/toolchains/stable-aarch64-apple-darwin/lib/rustlib/src/rust/library/std/src/thread/join_handle.rs:149:17
    |
149 |     pub fn join(self) -> Result<T> {
    |                 ^^^^

For more information about this error, try `rustc --explain E0507`.
error: could not compile `web_server` (lib) due to 1 previous error
Enter fullscreen mode Exit fullscreen mode

The error message shows that we cannot move the thread field out of worker, because we only have a mutable reference to each worker, but join requires ownership of the JoinHandle (that is, ownership of worker.thread).

To satisfy the ownership requirement, we need to change the type of the thread field in Worker and wrap thread::JoinHandle<()> in Option<T>. That way we can call Option<T>::take to obtain ownership:

struct Worker {
    id: usize,
    thread: Option<thread::JoinHandle<()>>,
}
Enter fullscreen mode Exit fullscreen mode

Anywhere the thread field is used must also be updated because of Option<T>:

impl Worker {
    fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Job>>>) -> Worker {
        let thread = thread::spawn(move || loop {
            let job = receiver.lock().unwrap().recv().unwrap();
            println!("Worker {} got a job; executing.", id);
            job();
        });

        Worker {
            id,
            thread: Some(thread),
        }
    }
}
Enter fullscreen mode Exit fullscreen mode

Change the value of the thread field from thread to Some(thread).

impl Drop for ThreadPool {
    fn drop(&mut self) {
        for worker in &mut self.workers {
            println!("Shutting down worker {}", worker.id);

            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap();
            }
        }
    }
}
Enter fullscreen mode Exit fullscreen mode

Use if let pattern matching to extract the value when worker.thread is Some (using take gives us ownership instead of a mutable reference).

20.3.2 Signaling Threads to Exit

This change compiles, but it still does not achieve the desired result. Calling drop does not actually shut down the threads, because the threads are still stuck in the loop waiting for work.

If we drop ThreadPool with this drop method, the main thread will block forever while waiting for the first thread to finish (because each thread keeps looping and looking for work, and never breaks out of the loop).

We need the sender field in ThreadPool to have two states: a live state with work attached to the sender, and a terminated state:

pub struct ThreadPool {
    workers: Vec<Worker>,
    sender: Option<mpsc::Sender<Job>>,
}
Enter fullscreen mode Exit fullscreen mode

Using Option<T> lets it represent both states.

Anywhere the sender field is used must also be changed:

impl Drop for ThreadPool {
    fn drop(&mut self) {
        drop(self.sender.take());

        for worker in &mut self.workers {
            println!("Shutting down worker {}", worker.id);

            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap();
            }
        }
    }
}
Enter fullscreen mode Exit fullscreen mode

Add drop(self.sender.take()); to explicitly drop the sender, which closes the channel. When that happens, all the recv calls executed inside the workers’ infinite loop will return an error, and the workers will stop running.

pub fn new(size: usize) -> ThreadPool {
    assert!(size > 0);
    let (sender, receiver) = mpsc::channel();
    let mut workers = Vec::with_capacity(size);

    let receiver = Arc::new(Mutex::new(receiver));
    for id in 0..size {
        workers.push(Worker::new(id, Arc::clone(&receiver)));
    }

    ThreadPool {
        workers,
        sender: Some(sender),
    }
}
Enter fullscreen mode Exit fullscreen mode

Wrap the sender field in the return value with Some.

pub fn execute<F>(&self, f: F)
where
    F: FnOnce() + Send + 'static,
{
    let job = Box::new(f);

    self.sender.as_ref().unwrap().send(job).unwrap();
}
Enter fullscreen mode Exit fullscreen mode

Because sender is now an Option, we call as_ref to get an Option<&Sender> (a reference to the sender inside) without moving it out of self. Then we unwrap and call send, which takes &self on the sender and moves the job into the channel.

This change is still not elegant, because all the recv calls executed in the workers’ infinite loop will return errors. It would be better not to exit because of an error, so we need one more change:

fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Job>>>) -> Worker {
    let thread = thread::spawn(move || loop {
        let job = receiver.lock().unwrap().recv();

        match job {
            Ok(job) => {
                println!("Worker {} got a job; executing.", id);
                job();
            }
            Err(_) => {
                println!("Worker {} disconnected; shutting down.", id);
                break;
            }
        }
    });
Enter fullscreen mode Exit fullscreen mode

Remove the last unwrap from job, and instead use match branches: execute job for the Ok variant, and for the Err variant print that the worker is disconnecting and then break.

20.3.3 Trial Run

To test the modified behavior, change main.rs so that the server only accepts two requests (by limiting the number of iterations with take):

fn main() {
    let listener = TcpListener::bind("127.0.0.1:7878").unwrap();
    let pool = ThreadPool::new(4);

    for stream in listener.incoming().take(2) {
        let stream = stream.unwrap();

        pool.execute(|| {
            handle_connection(stream);
        });
    }

    println!("Shutting down.");
}
Enter fullscreen mode Exit fullscreen mode


Console output (one possible run; job assignment and interleaving are nondeterministic):

Shutting down.
Shutting down worker 0
Worker 0 got a job; executing.
Worker 3 got a job; executing.
Worker 1 disconnected; shutting down.
Worker 2 disconnected; shutting down.
Worker 3 disconnected; shutting down.
Worker 0 disconnected; shutting down.
Shutting down worker 1
Shutting down worker 2
Shutting down worker 3
Enter fullscreen mode Exit fullscreen mode

You may see different worker ids and a different interleaving of “got a job”, “disconnected”, and “Shutting down worker” lines, but the overall pattern should be similar.

20.3.4 Summary

main.rs:

use std::{
    fs,
    io::{prelude::*, BufReader},
    net::{TcpListener, TcpStream},
    thread,
    time::Duration,
};
use web_server::ThreadPool;

fn main() {
    let listener = TcpListener::bind("127.0.0.1:7878").unwrap();
    let pool = ThreadPool::new(4);

    for stream in listener.incoming().take(2) {
        let stream = stream.unwrap();

        pool.execute(|| {
            handle_connection(stream);
        });
    }

    println!("Shutting down.");
}

fn handle_connection(mut stream: TcpStream) {
    let buf_reader = BufReader::new(&stream);
    let request_line = buf_reader.lines().next().unwrap().unwrap();

    let (status_line, filename) = match &request_line[..] {
        "GET / HTTP/1.1" => ("HTTP/1.1 200 OK", "hello.html"),
        "GET /sleep HTTP/1.1" => {
            thread::sleep(Duration::from_secs(5));
            ("HTTP/1.1 200 OK", "hello.html")
        }
        _ => ("HTTP/1.1 404 NOT FOUND", "404.html"),
    };

    let contents = fs::read_to_string(filename).unwrap();
    let length = contents.len();

    let response =
        format!("{status_line}\r\nContent-Length: {length}\r\n\r\n{contents}");

    stream.write_all(response.as_bytes()).unwrap();
}
Enter fullscreen mode Exit fullscreen mode

lib.rs:

use std::{
    sync::{mpsc, Arc, Mutex},
    thread,
};

pub struct ThreadPool {
    workers: Vec<Worker>,
    sender: Option<mpsc::Sender<Job>>,
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        drop(self.sender.take());

        for worker in &mut self.workers {
            println!("Shutting down worker {}", worker.id);

            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap();
            }
        }
    }
}

type Job = Box<dyn FnOnce() + Send + 'static>;

impl ThreadPool {
    /// Create a new ThreadPool.
    ///
    /// The size is the number of threads in the pool.
    ///
    /// # Panics
    ///
    /// The `new` function will panic if the size is zero.
    pub fn new(size: usize) -> ThreadPool {
        assert!(size > 0);
        let (sender, receiver) = mpsc::channel();
        let mut workers = Vec::with_capacity(size);

        let receiver = Arc::new(Mutex::new(receiver));
        for id in 0..size {
            workers.push(Worker::new(id, Arc::clone(&receiver)));
        }

        ThreadPool {
            workers,
            sender: Some(sender),
        }
    }

    pub fn execute<F>(&self, f: F)
    where
        F: FnOnce() + Send + 'static,
    {
        let job = Box::new(f);

        self.sender.as_ref().unwrap().send(job).unwrap();
    }
}

struct Worker {
    id: usize,
    thread: Option<thread::JoinHandle<()>>,
}

impl Worker {
    fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Job>>>) -> Worker {
        let thread = thread::spawn(move || loop {
            let job = receiver.lock().unwrap().recv();

            match job {
                Ok(job) => {
                    println!("Worker {} got a job; executing.", id);
                    job();
                }
                Err(_) => {
                    println!("Worker {} disconnected; shutting down.", id);
                    break;
                }
            }
        });

        Worker {
            id,
            thread: Some(thread),
        }
    }
}
Enter fullscreen mode Exit fullscreen mode

hello.html:

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="utf-8">
    <title>Hello!</title>
</head>
<body>
<h1>Hello!</h1>
<p>Hi from Rust</p>
</body>
</html>
Enter fullscreen mode Exit fullscreen mode

404.html:

<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <title>Hello!</title>
</head>
<body>
<h1>Oops!</h1>
<p>Sorry, I don't know what you're asking for.</p>
</body>
</html>
Enter fullscreen mode Exit fullscreen mode

Source: dev.to

arrow_back Back to Tutorials