What context to pass to just wait until something is completed?

Question

For various reasons, I feel that http.ListenAndServe does not suit my needs.

I needed to be able to determine the bound address and port (i.e. when using ":0"), so I introduced a net.Listener, read listener.Addr() and then passed to http.Serve(listener, nil).

Then I needed to be able to run two HTTP servers with different URL handlers, so I introduced an http.NewServeMux(), added the necessary mux.HandleFunc("/path", fn) handlers, and passed as http.Serve(listener, mux).

Then I needed to be able to stop these servers cleanly, and shut down any connections, independently of the main program itself, so now I have introduced &http.Server{Handler: mux} which I can go func() { server.Serve(listener) }().

In theory, I can stop this by calling server.Shutdown(ctx), but now none of the available contexts in import "context" seem to offer what I want either. I want to be able to wait until the clean shutdown has finished, then continue with my code.

My understanding is that I should be able to <- ctx.Done() to achieve this, but I've tried both context.Background() and context.TODO() and neither seem to "trigger" ctx.Done(), and I end up blocking forever. The other context options appear to be time-based.

If I don't wait on something, or pass nil, server.Shutdown(ctx) seems to finish too quickly and I can see nothing is actually closed (runtime.Numgoroutine() != 1)

I can time.Sleep(duration) for some arbitrary duration, but I don't want an arbitrary duration. I want to know that server.Shutdown has completed cleanly.

package main

import (
    "fmt"
    "net"
    "net/http"
    "runtime"
    "time"
)

func main() {
    var err error

    listener, err := net.Listen("tcp", "localhost:0")
    fmt.Printf("Listening on http://%v\n", listener.Addr())

    mux := http.NewServeMux()
    mux.HandleFunc("/", handleIndex)

    stop, err := startHTTPServer(listener, mux)

    d, _ := time.ParseDuration("5s")

    time.Sleep(d)   // delay here just for example of "long-running" server
    close(stop)     // closing the channel returned by my helper should trigger shutdown
    time.Sleep(d)   // if this delay is here, I see the "Stopped" message

    if err != nil {
        panic(err)
    }

    fmt.Printf("End of program, active goroutines: %v", runtime.NumGoroutine())
}

// startHTTPServer is a helper function to start a server and return a channel that can trigger shutdown
func startHTTPServer(listener net.Listener, handler http.Handler) (stop chan struct{}, err error) {
    stop = make(chan struct{})
    server := &http.Server{Handler: handler}

    go func() {
        fmt.Println("Starting server...")
        err = server.Serve(listener)
    }()
    go func() {
        select {
        case <-stop:
            fmt.Println("Stop channel closed; stopping server...")
            err = server.Shutdown(nil)    // what is passed instead of nil here?
            fmt.Println("Stopped.")
            return
        }
    }()
    return
}

func handleIndex(w http.ResponseWriter, r *http.Request) {
    fmt.Fprintln(w, "Hello world")
}

I've tried both context.Background() and context.TODO(). I've tried new(context.Context) but that threw a SIGSEGV. I've tried nil and that doesn't wait at all.

I tried adding a sync.WaitGroup and calling wg.Wait() instead of the second time.Sleep(d), but I still need to wait until server.Shutdown() is finished before calling wg.Done() (and defer wg.Done() called it too early).

I feel like, with Contexts, WaitGroups, etc., I'm just adding cruft to the code without really understanding why any of it is necessary.

What is the correct, clean, idiomatic way to wait for the server.Shutdown to complete?

Answer 1

To wait for Shutdown to complete in the main goroutine, call Shutdown from that goroutine. Eliminate the channel and extra goroutine.

listener, _ := net.Listen("tcp", "localhost:0")
fmt.Printf("Listening on http://%v\n", listener.Addr())

mux := http.NewServeMux()
mux.HandleFunc("/", handleIndex)

server := &http.Server{Handler: mux}
go func() {
    fmt.Println("Starting server...")
    err := server.Serve(listener)
    if err != nil && err != http.ErrServerClosed {
        log.Fatal(err)
    }
}()

time.Sleep(5 * time.Second) // delay here just for example of "long-running" server

// Shutdown and wait for server to complete.
server.Shutdown(context.Background())

If you want to limit how long Shutdown waits for the server to shutdown, replace context.Background() with a context created with a deadline.

Answer 2

In theory, I can stop this by calling server.Shutdown(ctx), but now none of the available contexts in import "context" seem to offer what I want either. I want to be able to wait until the clean shutdown has finished, then continue with my code.

This means you should pass a context which doesn't time out. Neither context.Background() nor context.TODO() time out and thus are suitable (really, see below). Whether you use the one or the other depends on whether you plan to time out the shutdown (you should to prevent rough client from stopping you shutting down your server) or not.

My understanding is that I should be able to <- ctx.Done() to achieve this, but I've tried both context.Background() and context.TODO() and neither seem to "trigger" ctx.Done(), and I end up blocking forever. The other context options appear to be time-based.

Well, this one is wrong. Neither context.Background() nor context.TODO() will close their Done as they do not time out. But there is no need to wait for done: server.Shutdown is a normal function which returns once the server actually is shut down properly (that's what you seem to want) or the context times out. In any case server.Shutdown just returns.

A simple

server.Shutdown(context.TODO())

is what you want. (For now, in the long run: Pass a context which times out after a long ,but finit, time.)

But your code looks fishy anyway: Your func startHTTPServer doesn't properly handle errors: Not the error from starting and not the one from stopping. If your server didn't start you cannot stop it and your code just swallows the error. It is also racy on err. Your problems probably do not stem from the context passed to server.Shutdown but from somewhere else.

The problem is that your code does not wait for server.Shutdown to return as this function is started in a goroutine without any syncronisation back to the caller of startHTTPServer: Don't do that. Really: the Context of server.Shutdown is not the problem.

The following is untested code which overcomes the problems a bit. It is not production ready as you can see at all the TOODs.

// startHTTPServer is a helper function to start a server and returns
// a channel to stop the server and a channel reporting errors during
// starting/stopping the server or nil if the server was shut down
// properly
func startHTTPServer(listener net.Listener, handler http.Handler) (stop chan bool, problems chan error) {
    stop, problems = make(chan bool), make(chan error)
    server := &http.Server{Handler: handler} // TODO: set timeouts

    go func() {
        fmt.Println("Starting server...")
        err := server.Serve(listener)
        if err != http.ErrServerClosed {
            problems <- err // TODO: tag/classify as startup error
        }
    }()
    go func() {
        select {
        case <-stop:
            fmt.Println("Stop channel closed; stopping server...")
            err := server.Shutdown(context.TODO())
            fmt.Println("Stopped.")
            problems <- err // TODO: tag/classify as shutdown error
        case e := <-problems:
            problems <- e  // resend, this error  is not for us
            return // stop waiting for stop as server did not start anyway.
        }
    }()
    return stop, problems

}

Other solutions would be possible too, e.g. returning separate startup and shutdown error channels, etc.