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

.github/workflows/test.yml

@@ -3,10 +3,10 @@ name: Test
 on:
   push:
     branches:
-      - master
+      - main
   pull_request:
     branches:
-      - master
+      - main
 
 jobs:
   build:

README.md

@@ -312,7 +312,7 @@ For the port it will look in `spec.ports` for a port named `mc-router`, if not p
   `"mc-router.itzg.me/externalServerName"` annotation that declares their external server name(s)
 
 ```bash
-kubectl apply -f https://raw.githubusercontent.com/itzg/mc-router/master/docs/k8s-example-auto.yaml
+kubectl apply -f https://raw.githubusercontent.com/itzg/mc-router/main/docs/k8s-example-auto.yaml
 ```
 
 ![](docs/example-deployment-auto.drawio.png)

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())
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
-	defer cancel()
+	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():
+		case <-ctx.Done():
-			return
+			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()
 }

go.mod

@@ -1,6 +1,6 @@
 module github.com/itzg/mc-router
 
-go 1.24.4
+go 1.25
 
 require (
 	github.com/fsnotify/fsnotify v1.9.0

server/server.go

@@ -3,12 +3,13 @@ package server
 import (
 	"context"
 	"fmt"
-	"github.com/sirupsen/logrus"
 	"net"
 	"os"
 	"runtime/pprof"
 	"strconv"
 	"time"
+
+	"github.com/sirupsen/logrus"
 )
 
 type Server struct {
@@ -179,15 +180,6 @@ func NewServer(ctx context.Context, config *Config) (*Server, error) {
 	}, nil
 }
 
-// Done provides a channel notified when the server has closed all connections, etc
-func (s *Server) Done() <-chan struct{} {
-	return s.doneChan
-}
-
-func (s *Server) notifyDone() {
-	s.doneChan <- struct{}{}
-}
-
 // ReloadConfig indicates that an external request, such as a SIGHUP,
 // is requesting the routes config file to be reloaded, if enabled
 func (s *Server) ReloadConfig() {
@@ -209,7 +201,6 @@ func (s *Server) Run() {
 	)
 	if err != nil {
 		logrus.WithError(err).Error("Could not start accepting connections")
-		s.notifyDone()
 		return
 	}
 
@@ -222,10 +213,9 @@ func (s *Server) Run() {
 			}
 
 		case <-s.ctx.Done():
-			logrus.Info("Stopping. Waiting for connections to complete...")
+			logrus.Info("Server Stopping. Waiting for connections to complete...")
 			s.connector.WaitForConnections()
 			logrus.Info("Stopped")
-			s.notifyDone()
 			return
 		}
 	}