Attempt to shutdown server gracefully

go.mod

@@ -3,11 +3,18 @@ module github.com/CoderCookE/goaround
 go 1.13
 
 require (
-	github.com/OneOfOne/xxhash v1.2.7 // indirect
-	github.com/dgraph-io/ristretto v0.0.2-0.20200130184421-ff325ada232b
+	github.com/OneOfOne/xxhash v1.2.8 // indirect
+	github.com/dgraph-io/ristretto v0.0.2
 	github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13 // indirect
 	github.com/kr/pretty v0.2.0 // indirect
-	github.com/prometheus/client_golang v1.4.1
+	github.com/kr/text v0.2.0 // indirect
+	github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
+	github.com/prometheus/client_golang v1.7.1
 	github.com/spaolacci/murmur3 v1.1.0 // indirect
-	gopkg.in/yaml.v2 v2.2.8 // indirect
+	github.com/stretchr/testify v1.6.1 // indirect
+	golang.org/x/sys v0.0.0-20200625212154-ddb9806d33ae // indirect
+	google.golang.org/protobuf v1.25.0 // indirect
+	gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f // indirect
+	gopkg.in/yaml.v2 v2.3.0 // indirect
+	gopkg.in/yaml.v3 v3.0.0-20200615113413-eeeca48fe776 // indirect
 )

internal/assert/assert.go

@@ -25,6 +25,14 @@ func (a *Asserter) Equal(actual, expected interface{}) {
 	}
 }
 
+func (a *Asserter) StringContains(actual, expected string) {
+	a.T.Helper()
+	if !strings.Contains(actual, expected) {
+		message := fmt.Sprintf("expected: %s was not contained in actual: %s", expected, actual)
+		a.T.Error(message)
+	}
+}
+
 func (a *Asserter) NotEqual(actual, expected interface{}) {
 	a.T.Helper()
 

internal/gracefulserver/main.go

@@ -0,0 +1,71 @@
+package gracefulserver
+
+import (
+	"context"
+	"log"
+	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+)
+
+type GracefulServer struct {
+	*http.Server
+	errChan      chan error
+	shutdownChan chan os.Signal
+}
+
+func New(server *http.Server) *GracefulServer {
+	gs := &GracefulServer{
+		Server:       server,
+		errChan:      make(chan error, 1),
+		shutdownChan: make(chan os.Signal, 1),
+	}
+
+	signal.Notify(gs.shutdownChan, syscall.SIGTERM)
+
+	return gs
+}
+
+func (gs *GracefulServer) ListenAndServe() error {
+	go func() {
+		err := gs.Server.ListenAndServe()
+		if err == http.ErrServerClosed {
+			err = nil
+		}
+
+		gs.errChan <- err
+	}()
+
+	return gs.listenForSignals()
+}
+
+func (gs *GracefulServer) ListenAndServeTLS(certFile, keyFile string) error {
+	go func() {
+		err := gs.Server.ListenAndServeTLS(certFile, keyFile)
+		if err == http.ErrServerClosed {
+			err = nil
+		}
+
+		gs.errChan <- err
+	}()
+
+	return gs.listenForSignals()
+}
+
+func (gs *GracefulServer) listenForSignals() error {
+	select {
+	case err := <-gs.errChan:
+		return err
+	case <-gs.shutdownChan:
+		signal.Stop(gs.shutdownChan)
+		log.Printf("Gracefully shutting down")
+		defer log.Printf("Graceful shutdown complete")
+
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		return gs.Server.Shutdown(ctx)
+	}
+}

internal/gracefulserver/main_test.go

@@ -0,0 +1,59 @@
+package gracefulserver
+
+import (
+	"bytes"
+	"log"
+	"net/http"
+	"os"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/CoderCookE/goaround/internal/assert"
+)
+
+func TestMain(m *testing.M) {
+	os.Exit(m.Run())
+}
+
+func testHandler(rw http.ResponseWriter, req *http.Request) {
+	rw.Write([]byte("ok"))
+}
+
+func TestListenAndServe(t *testing.T) {
+	assertion := &assert.Asserter{T: t}
+
+	t.Run("gracefully shutsdown", func(t *testing.T) {
+		var buf bytes.Buffer
+		log.SetOutput(&buf)
+		defer func() {
+			log.SetOutput(os.Stderr)
+		}()
+
+		done := make(chan struct{})
+		server := New(&http.Server{
+			Addr:        ":9999",
+			Handler:     http.HandlerFunc(testHandler),
+			ReadTimeout: 1 * time.Second,
+		})
+
+		go func() {
+			err := server.ListenAndServe()
+			assertion.Equal(err, nil)
+
+			<-time.After(500 * time.Millisecond)
+
+			close(done)
+		}()
+
+		<-time.After(500 * time.Millisecond)
+		process, err := os.FindProcess(os.Getpid())
+		assertion.Equal(err, nil)
+
+		err = process.Signal(syscall.SIGTERM)
+		assertion.Equal(err, nil)
+		<-done
+
+		assertion.StringContains(buf.String(), "Graceful shutdown complete")
+	})
+}

internal/stats/main.go

@@ -3,7 +3,9 @@ package stats
 import (
 	"log"
 	"net/http"
+	"time"
 
+	"github.com/CoderCookE/goaround/internal/gracefulserver"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
@@ -51,10 +53,22 @@ func init() {
 }
 
 func StartUp(addr string) {
-	http.Handle("/metrics", promhttp.HandlerFor(
+	handler := promhttp.HandlerFor(
 		prometheus.DefaultGatherer,
 		promhttp.HandlerOpts{},
-	))
+	)
+
+	server := &http.Server{
+		Addr:         addr,
+		Handler:      handler,
+		ReadTimeout:  10 * time.Second,
+		WriteTimeout: 30 * time.Second,
+	}
 
-	log.Fatal(http.ListenAndServe(addr, nil))
+	graceful := gracefulserver.New(server)
+	log.Printf("Starting Prometheus server on port %s", addr)
+	err := graceful.ListenAndServe()
+	if err != nil {
+		log.Printf("Error starting Prometheus server: %s", err.Error())
+	}
 }

main.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/CoderCookE/goaround/internal/customflags"
+	"github.com/CoderCookE/goaround/internal/gracefulserver"
 	"github.com/CoderCookE/goaround/internal/pool"
 	"github.com/CoderCookE/goaround/internal/stats"
 )
@@ -46,11 +47,12 @@ func main() {
 
 	defer server.Close()
 
+	graceful := gracefulserver.New(server)
 	var err error
 	if *cacert != "" && *privkey != "" {
-		err = server.ListenAndServeTLS(*cacert, *privkey)
+		err = graceful.ListenAndServeTLS(*cacert, *privkey)
 	} else {
-		err = server.ListenAndServe()
+		err = graceful.ListenAndServe()
 	}
 
 	if err != nil {