From fdafbc4f0eb4f8030e4db83afc5038d1bc2f081d Mon Sep 17 00:00:00 2001
From: ErickSkrauch <erickskrauch@yandex.ru>
Date: Wed, 7 Feb 2024 18:33:06 +0100
Subject: [PATCH] Simplify stop signal handling

---
 internal/di/handlers.go    |  6 +---
 internal/http/http.go      | 56 +++++++++++---------------------------
 internal/http/http_test.go | 16 -----------
 3 files changed, 17 insertions(+), 61 deletions(-)

diff --git a/internal/di/handlers.go b/internal/di/handlers.go
index 6fe7125..91f42ac 100644
--- a/internal/di/handlers.go
+++ b/internal/di/handlers.go
@@ -39,11 +39,7 @@ func newHandlerFactory(
 	}
 
 	router.StrictSlash(true)
-	requestEventsMiddleware := CreateRequestEventsMiddleware()
-	router.Use(requestEventsMiddleware)
-	// NotFoundHandler doesn't call for registered middlewares, so we must wrap it manually.
-	// See https://github.com/gorilla/mux/issues/416#issuecomment-600079279
-	router.NotFoundHandler = requestEventsMiddleware(http.HandlerFunc(NotFoundHandler))
+	router.NotFoundHandler = http.HandlerFunc(NotFoundHandler)
 
 	if slices.Contains(enabledModules, "api") {
 		var apiRouter *mux.Router
diff --git a/internal/http/http.go b/internal/http/http.go
index 25e30b1..d6f9565 100644
--- a/internal/http/http.go
+++ b/internal/http/http.go
@@ -3,11 +3,11 @@ package http
 import (
 	"context"
 	"encoding/json"
-	"errors"
 	"net/http"
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"github.com/gorilla/mux"
 	"github.com/mono83/slf"
@@ -19,52 +19,28 @@ import (
 func StartServer(server *http.Server, logger slf.Logger) {
 	logger.Debug("Chrly :v (:c)", wd.StringParam("v", version.Version()), wd.StringParam("c", version.Commit()))
 
-	done := make(chan bool, 1)
+	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM, os.Kill)
+	defer cancel()
+
+	srvErr := make(chan error, 1)
 	go func() {
 		logger.Info("Starting the server, HTTP on: :addr", wd.StringParam("addr", server.Addr))
-		if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
-			logger.Emergency("Error in main(): :err", wd.ErrParam(err))
-			close(done)
-		}
+		srvErr <- server.ListenAndServe()
+		close(srvErr)
 	}()
 
-	go func() {
-		s := waitForExitSignal()
-		logger.Info("Got signal: :signal, starting graceful shutdown", wd.StringParam("signal", s.String()))
-		_ = server.Shutdown(context.Background())
-		logger.Info("Graceful shutdown succeed, exiting", wd.StringParam("signal", s.String()))
-		close(done)
-	}()
+	select {
+	case err := <-srvErr:
+		logger.Emergency("Error in main(): :err", wd.ErrParam(err))
+	case <-ctx.Done():
+		logger.Info("Got stop signal, starting graceful shutdown: :ctx")
 
-	<-done
-}
+		stopCtx, cancelFunc := context.WithTimeout(context.Background(), 3*time.Second)
+		defer cancelFunc()
 
-func waitForExitSignal() os.Signal {
-	ch := make(chan os.Signal, 1)
-	signal.Notify(ch, os.Interrupt, syscall.SIGTERM, os.Kill)
+		_ = server.Shutdown(stopCtx)
 
-	return <-ch
-}
-
-type loggingResponseWriter struct {
-	http.ResponseWriter
-	statusCode int
-}
-
-func (lrw *loggingResponseWriter) WriteHeader(code int) {
-	lrw.statusCode = code
-	lrw.ResponseWriter.WriteHeader(code)
-}
-
-func CreateRequestEventsMiddleware() mux.MiddlewareFunc {
-	return func(handler http.Handler) http.Handler {
-		return http.HandlerFunc(func(resp http.ResponseWriter, req *http.Request) {
-			lrw := &loggingResponseWriter{
-				ResponseWriter: resp,
-				statusCode:     http.StatusOK,
-			}
-			handler.ServeHTTP(lrw, req)
-		})
+		logger.Info("Graceful shutdown succeed, exiting")
 	}
 }
 
diff --git a/internal/http/http_test.go b/internal/http/http_test.go
index 1634e3b..623ab2b 100644
--- a/internal/http/http_test.go
+++ b/internal/http/http_test.go
@@ -11,22 +11,6 @@ import (
 	"github.com/stretchr/testify/mock"
 )
 
-func TestCreateRequestEventsMiddleware(t *testing.T) {
-	req := httptest.NewRequest("GET", "http://example.com", nil)
-	resp := httptest.NewRecorder()
-
-	isHandlerCalled := false
-	middlewareFunc := CreateRequestEventsMiddleware()
-	middlewareFunc.Middleware(http.HandlerFunc(func(resp http.ResponseWriter, req *http.Request) {
-		resp.WriteHeader(400)
-		isHandlerCalled = true
-	})).ServeHTTP(resp, req)
-
-	if !isHandlerCalled {
-		t.Fatal("Handler isn't called from the middleware")
-	}
-}
-
 type authCheckerMock struct {
 	mock.Mock
 }