If we shut down our program (currently via sigint, and crtl+c), it will cut off any current connections. Go's http.Server actually handles graceful shutdowns! We just need to orchestrate it properly with our whacky setup.

We'll do just that, giving connections up to 10 seconds to finish their request before shutting down forcefully.

Let's start a little backwards this time, here's an updated main.go file.

Read the previous article for more context. What we do there is call a new Stop() method on my ReverseProxy after we send an interrupt signal:

package main

import (
    "github.com/fideloper/someproxy/reverseproxy"
    "github.com/gorilla/mux"
)

func main() {
    proxy := &reverseproxy.ReverseProxy{}

    // Let's assume we have 2 backends to proxy to
    // localhost:8000 (for our fun API requests)
    // and localhost:8001 (for everything else)

    // Match requests to "localhost/api" and "localhost/api/*"
    r := mux.NewRouter()
    r.Host("localhost").PathPrefix("/api")
    proxy.AddTarget("http://localhost:8000", r)

    // Catch-all for all other requests
    proxy.AddTarget("http://localhost:8001", nil)

    // Listen for http://
    proxy.AddListener(":80")

    // Listen for https://
    proxy.AddListenerTLS(":443", "keys/fid.dev.pem", "keys/fid.dev-key.pem")

    if err := proxy.Start(); err != nil {
        log.Fatal(err)
    }

    // Shutdown when we receive Ctrl+c (interrupt)
    c := make(chan os.Signal, 1)
    
    // We'll accept graceful shutdowns when quit via SIGINT (Ctrl+C)
    // SIGKILL, SIGQUIT or SIGTERM (Ctrl+/) will not be caught.
    signal.Notify(c, os.Interrupt)

    // Block until we receive our signal.
    <-c

    proxy.Stop() // This is the only new thing here
}

The addition of proxy.Stop() is all we changed.

Adding the Stop Method

Next we'll edit reverseproxy.go and add that Stop() method. Let's first talk about what it will do!

It turns out that http.Server has a Shutdown() method. This will handle gracefully closing listeners. If any are active, it will wait indefinitely for them to disconnect.

To make sure "indefinitely" isn't forever, we can pass it a context.Context with a timeout. We'll set that timeout to 10 seconds (I randomly chose 10 seconds).

After 10 seconds, any connections that refuse to finish are forcefully cut and the server will shut down.

It looks a bit like this:

// For a given HTTP server listening for connections
srv := &http.Server{}

srv.Serve(someListener)

// ...

// We can later shut it down gracefully, with a 10 second deadline
context, close := context.WithTimeout(context.Background(), time.Second * 10)
srv.Shutdown(context)

The Shutdown() method is blocking, so it will take up to 10 seconds to complete.

But ... multiple servers!

We need to handle shutting down multiple servers, so our Stop() function is just a tad more complex than the above logic. Here's the updated Stop() method in reverseproxy.go:

// Stop will gracefully shut down all listening servers
func (r *ReverseProxy) Stop() {
    // A context that times out in 10 seconds
    ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
    defer cancel()

    // A waitgroup allows us to block until all
    // goroutines are finished - when all servers
    // have finished shutting down.
    var wg sync.WaitGroup

    for _, srv := range r.servers {
        // This prevents a nasty and extremely common bug
        // Google for "golang range loop variable re-use"
        srv := srv

        // Tell the WaitGroup to wait for
        // +1 things to finish
        wg.Add(1)
        go func() {
            // Tell the WaitGroup we finished
            // one of the things
            defer wg.Done()

            // Wait up to 10 seconds for the
            // server to shutdown
            if err := srv.Shutdown(ctx); err != nil {
                log.Println(err)
                return
            }
            log.Println("A listener was shutdown successfully")
        }()
    }

    // Block until all servers are shut down
    wg.Wait()
    log.Println("Server shut down")
}

Let's cover what's going on there.

First we create a context. We just tell the cancel() method to run at the end of the function.

We have multiple servers to shutdown, and we want them to shutdown in parallel (not in serial!). Therefore, we'll run each Shutdown method in a goroutine.

This creates a race condition - the function will just finish immediately if we run those in a goroutine, and it's possible our entire app shuts down before the shutdown calls are complete.

To ensure we wait for all servers to finish shutting down, we'll use a sync.Waitgroup, which helps us orchestrate that. The call to wg.Wait() blocks further execution until wg.Done() is called for each wg.Add(1). If we setup 3 servers, we need to make sure all 3 are done shutting down.

This will now gracefully shutdown each server! If a connection is doing something that takes longer than 10 seconds, the context will timeout and the server will forcefully shutdown. In that case, you'll see log output similar to this:

2022/10/08 13:42:04 A listener was shutdown successfully
2022/10/08 13:42:14 context deadline exceeded
2022/10/08 13:42:14 A listener was shutdown successfully
2022/10/08 13:42:14 Server shut down

One server shutdown immediately (it didn't have any active connections), but the other had a long-running connection that exceeded the 10 second limit. It was shutdown forcefully.

And voilà, our reverse proxy will now shutdown gracefully!