From febdf04bceb525512554c46b44d8e155e6c22593 Mon Sep 17 00:00:00 2001
From: Robert Volkmann <Robert.Volkmann@x-cellent.com>
Date: Tue, 26 Sep 2023 15:46:02 +0200
Subject: [PATCH] wait until go routines finish on SIGTERM

---
 cmd/server.go | 44 +++++++++++++++++++++++++++++++++++++-------
 1 file changed, 37 insertions(+), 7 deletions(-)

diff --git a/cmd/server.go b/cmd/server.go
index ba973d0e..625e817a 100644
--- a/cmd/server.go
+++ b/cmd/server.go
@@ -11,7 +11,10 @@ import (
 	"net/http"
 	httppprof "net/http/pprof"
 	"os"
+	"os/signal"
 	"strings"
+	"sync"
+	"syscall"
 	"time"
 
 	"github.com/kelseyhightower/envconfig"
@@ -112,10 +115,21 @@ func Run() {
 		os.Exit(1)
 	}
 
-	ctx := context.Background()
+	var wg sync.WaitGroup
 
-	go c.ConstantlyReconfigureSwitch(ctx, cfg.ReconfigureSwitchInterval)
-	go c.ConstantlyPhoneHome(ctx, phonedHomeInterval)
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		c.ConstantlyReconfigureSwitch(ctx, cfg.ReconfigureSwitchInterval)
+	}()
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		c.ConstantlyPhoneHome(ctx, phonedHomeInterval)
+	}()
 
 	// Start metrics
 	metricsAddr := fmt.Sprintf("%v:%d", cfg.MetricsServerBindAddress, cfg.MetricsServerPort)
@@ -137,8 +151,24 @@ func Run() {
 		ReadHeaderTimeout: 3 * time.Second,
 	}
 
-	if err = srv.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
-		log.Error("unable to start metrics listener", "error", err)
-		os.Exit(1)
-	}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		if err = srv.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
+			log.Error("unable to start metrics listener", "error", err)
+			os.Exit(1)
+		}
+	}()
+
+	<-ctx.Done()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		if err = srv.Shutdown(context.Background()); err != nil {
+			log.Error("unable to shutdown metrics listener", "error", err)
+		}
+	}()
+
+	wg.Wait()
 }