Use signal.NotifyContext and WaitGroup.Go (#466)

cmd/mc-router/main.go

@@ -3,12 +3,14 @@ package main
 import (
 	"context"
 	"fmt"
+	"os"
+	"os/signal"
+	"sync"
+	"syscall"
+
 	"github.com/itzg/go-flagsfiller"
 	"github.com/itzg/mc-router/server"
 	"github.com/sirupsen/logrus"
-	"os"
-	"os/signal"
-	"syscall"
 )
 
 var (
@@ -49,36 +51,37 @@ func main() {
 		logrus.Debug("Debug logs enabled")
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
+	defer stop()
 
 	signals := make(chan os.Signal, 1)
-	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
+	signal.Notify(signals, syscall.SIGHUP)
 
 	s, err := server.NewServer(ctx, &cliConfig.ServerConfig)
 	if err != nil {
 		logrus.WithError(err).Fatal("Could not setup server")
 	}
 
-	go s.Run()
+	var wg sync.WaitGroup
+	wg.Go(s.Run)
 
+signalsLoop:
 	for {
 		select {
-		case <-s.Done():
-			return
+		case <-ctx.Done():
+			break signalsLoop
 
 		case sig := <-signals:
 			switch sig {
 			case syscall.SIGHUP:
 				s.ReloadConfig()
 
-			case syscall.SIGINT, syscall.SIGTERM:
-				cancel()
-				// but wait for the server to be done
-
 			default:
 				logrus.WithField("signal", sig).Warn("Received unexpected signal")
 			}
 		}
 	}
+
+	logrus.Info("Stopping")
+	wg.Wait()
 }