FEAT: server graceful shutdown

main.go

@@ -1,13 +1,39 @@
 package main
 
 import (
+	"context"
 	"log"
+	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
 
 	"github.com/xray-web/web-check-api/config"
 	"github.com/xray-web/web-check-api/server"
 )
 
 func main() {
-	s := server.New(config.New())
-	log.Println(s.Run())
+	srv := server.New(config.New())
+
+	done := make(chan os.Signal, 1)
+	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		if err := srv.Run(); err != nil && err != http.ErrServerClosed {
+			log.Fatalf("listen: %v\n", err)
+		}
+	}()
+
+	<-done
+
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer func() {
+		// extra handling here, databases etc
+		cancel()
+	}()
+
+	if err := srv.Shutdown(ctx); err != nil {
+		log.Fatalf("Server Shutdown Failed:%+v", err)
+	}
 }

server/server.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/http"
@@ -14,10 +15,12 @@
 	conf   config.Config
 	mux    *http.ServeMux
 	checks *checks.Checks
+	srv    *http.Server
 }
 
 func New(conf config.Config) *Server {
 	return &Server{
+		srv:    &http.Server{},
 		conf:   conf,
 		mux:    http.NewServeMux(),
 		checks: checks.NewChecks(),
@@ -49,12 +52,19 @@
 	s.mux.Handle("GET /api/social-tags", handlers.HandleGetSocialTags(s.checks.SocialTags))
 	s.mux.Handle("GET /api/tls", handlers.HandleTLS(s.checks.Tls))
 	s.mux.Handle("GET /api/trace-route", handlers.HandleTraceRoute())
+
+	s.srv.Handler = s.CORS(s.mux)
 }
 
 func (s *Server) Run() error {
 	s.routes()
 
 	addr := fmt.Sprintf("%s:%s", s.conf.Host, s.conf.Port)
 	log.Printf("Server started, listening on: %v\n", addr)
-	return http.ListenAndServe(addr, s.CORS(s.mux))
+	s.srv.Addr = addr
+	return s.srv.ListenAndServe()
+}
+
+func (s *Server) Shutdown(ctx context.Context) error {
+	return s.srv.Shutdown(ctx)
 }

server/server_test.go

@@ -0,0 +1,41 @@
+package server
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/xray-web/web-check-api/config"
+	"golang.org/x/net/context"
+)
+
+func TestServer(t *testing.T) {
+	t.Parallel()
+
+	t.Run("start server", func(t *testing.T) {
+		t.Parallel()
+
+		srv := New(config.New())
+		srv.routes()
+		ts := httptest.NewServer(srv.CORS(srv.mux))
+		defer ts.Close()
+
+		// wait up tot 10 seconds for health check to return 200
+		ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(+10*time.Second))
+		defer cancel()
+		for {
+			req, err := http.NewRequestWithContext(ctx, http.MethodGet, ts.URL+"/health", nil)
+			assert.NoError(t, err)
+			resp, err := http.DefaultClient.Do(req)
+			if err == nil && resp.StatusCode == http.StatusOK {
+				break
+			}
+		}
+		ctx, cancel = context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		err := srv.Shutdown(ctx)
+		assert.NoError(t, err)
+	})
+}