Gracefully Shutdown Gorilla Server

I'm building a server in go using gorilla multiplexer library found in https://github.com/gorilla/mux.
The problem is, I want it to gracefully shutdown when I'm using Ctrl+C, or when there is a specific API call, for example "/shutdown".

I already know that in Go 1.8, graceful shutdown is already implemented. But how to combine it with gorilla multiplexer? Also, how to combine it with SIGINT signal?

Can anyone show me how to do it?

Channel can be used to capture shutdown request through API call (/shutdown) or interrupt signal (Ctrl+C).

1. Embed http.Server into a custom struct, so we can call http Server.Shutdown later
2. Add channel field (shutdownReq) for passing shutdown request from API call /shutdown
3. Register http handlers including /shutdown in gorilla/mux's router, then assign the router to http.Server.Handler
4. Register os.Interrupt/syscall.SIGINT, syscall.SIGTERM handler
5. Use select to capture shutdown request through API call or interrupt signal
6. Perform clean shutdown by calling Server.Shutdown

Below is the example code:

package main
import (
"context"
"log"
"net/http"
"sync/atomic"
"syscall"
"time"
"os"
"os/signal"
"github.com/gorilla/mux"
)
type myServer struct {
http.Server
shutdownReq chan bool
reqCount    uint32
}
func NewServer() *myServer {
//create server
s := &myServer{
Server: http.Server{
Addr:         ":8080",
ReadTimeout:  10 * time.Second,
WriteTimeout: 10 * time.Second,
},
shutdownReq: make(chan bool),
}
router := mux.NewRouter()
//register handlers
router.HandleFunc("/", s.RootHandler)
router.HandleFunc("/shutdown", s.ShutdownHandler)
//set http server handler
s.Handler = router
return s
}
func (s *myServer) WaitShutdown() {
irqSig := make(chan os.Signal, 1)
signal.Notify(irqSig, syscall.SIGINT, syscall.SIGTERM)
//Wait interrupt or shutdown request through /shutdown
select {
case sig := <-irqSig:
log.Printf("Shutdown request (signal: %v)", sig)
case sig := <-s.shutdownReq:
log.Printf("Shutdown request (/shutdown %v)", sig)
}
log.Printf("Stoping http server ...")
//Create shutdown context with 10 second timeout
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
//shutdown the server
err := s.Shutdown(ctx)
if err != nil {
log.Printf("Shutdown request error: %v", err)
}
}
func (s *myServer) RootHandler(w http.ResponseWriter, r *http.Request) {
w.Write([]byte("Hello Gorilla MUX!\n"))
}
func (s *myServer) ShutdownHandler(w http.ResponseWriter, r *http.Request) {
w.Write([]byte("Shutdown server"))
//Do nothing if shutdown request already issued
//if s.reqCount == 0 then set to 1, return true otherwise false
if !atomic.CompareAndSwapUint32(&s.reqCount, 0, 1) {
log.Printf("Shutdown through API call in progress...")
return
}
go func() {
s.shutdownReq <- true
}()
}
func main() {
//Start the server
server := NewServer()
done := make(chan bool)
go func() {
err := server.ListenAndServe()
if err != nil {
log.Printf("Listen and serve: %v", err)
}
done <- true
}()
//wait shutdown
server.WaitShutdown()
<-done
log.Printf("DONE!")
}

Note: Please watch this issue which is related to gracefull shutdown.

A variation of the original question - how to shut down in a unit test? Watching for an OS signal is not needed, timeouts aren't needed, and so there's less to synchronize. Also, there's another issue - a unit test usually has to wait for the server to start listening before doing its test.

Here's @putu's answer simplified for this purpose.

package main
import (
"context"
"fmt"
"net"
"net/http"
"sync"
"sync/atomic"
"testing"
"github.com/gorilla/mux"
)
type (
myServer struct {
ctx      context.Context
srv      http.Server
stopping uint32
err      error
wg       sync.WaitGroup
}
)
func newServer(ctx context.Context, port int) (s *myServer) {
s = &myServer{
ctx: ctx,
srv: http.Server{
Addr: fmt.Sprintf(":%d", port),
},
}
// make the routes
router := mux.NewRouter()
router.HandleFunc("/", s.getRoot)
s.srv.Handler = router
return
}
func (s *myServer) getRoot(w http.ResponseWriter, r *http.Request) {
// example route
w.WriteHeader(http.StatusOK)
}
func (s *myServer) start() (err error) {
addr := s.srv.Addr
ln, err := net.Listen("tcp", addr)
if err != nil {
return
}
s.wg.Add(1)
go func() {
fmt.Println("Server started")
err = s.srv.Serve(ln)
if err != http.ErrServerClosed {
s.err = err
}
s.wg.Done()
}()
return
}
func (s *myServer) stop() error {
if atomic.CompareAndSwapUint32(&s.stopping, 0, 1) {
s.srv.Shutdown(s.ctx)
fmt.Println("Server stopped")
}
s.wg.Wait()
return s.err
}
func TestMockServer(t *testing.T) {
s := newServer(context.Background(), 8000)
err := s.start()
if err != nil {
t.Fatal(err)
}
// ...do a test that makes use of the mock server...
_, err = http.Get("http://localhost:8000")
if err != nil {
t.Fatal(err)
}
err = s.stop()
if err != nil {
t.Fatal(err)
}
}