From e90c211d0fc096098bbb2f0b1701cc0050a9e665 Mon Sep 17 00:00:00 2001
From: Pablu23 <pablu23@gmx.de>
Date: Wed, 15 Oct 2025 13:27:09 +0200
Subject: [PATCH] Added context.Context to stop and correctly stop metrics now

---
 cmd/domain-router/main.go |  7 +++++--
 middleware/logging.go     |  3 ++-
 middleware/metrics.go     | 16 ++++++++++++----
 middleware/pipeline.go    |  7 ++++---
 middleware/rate-limit.go  |  5 +++--
 5 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/cmd/domain-router/main.go b/cmd/domain-router/main.go
index 25fed55..2e53b13 100644
--- a/cmd/domain-router/main.go
+++ b/cmd/domain-router/main.go
@@ -64,8 +64,11 @@ func main() {
 		defer wg.Done()
 		<-sigs
 		log.Info().Msg("Stopping server")
-		server.Shutdown(context.Background())
-		pipeline.Stop()
+		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+		defer cancel()
+
+		server.Shutdown(ctx)
+		pipeline.Stop(ctx)
 	}()
 
 	if config.Server.Ssl.Enabled {
diff --git a/middleware/logging.go b/middleware/logging.go
index 8b20be6..f3f9b88 100644
--- a/middleware/logging.go
+++ b/middleware/logging.go
@@ -1,6 +1,7 @@
 package middleware
 
 import (
+	"context"
 	"net/http"
 	"time"
 
@@ -11,7 +12,7 @@ import (
 
 type RequestLogger struct{}
 
-func (_ *RequestLogger) Stop() {
+func (_ *RequestLogger) Stop(ctx context.Context) {
 	log.Info().Msg("Stopped Logging")
 }
 
diff --git a/middleware/metrics.go b/middleware/metrics.go
index 8ed0a0e..ed61883 100644
--- a/middleware/metrics.go
+++ b/middleware/metrics.go
@@ -2,6 +2,7 @@ package middleware
 
 import (
 	"cmp"
+	"context"
 	"encoding/json"
 	"net/http"
 	"os"
@@ -17,7 +18,7 @@ type Metrics struct {
 	endpointMetrics []EndpointMetrics
 	ticker          *time.Ticker
 	file            string
-	stop            chan bool
+	stop            chan struct{}
 }
 
 type EndpointMetrics struct {
@@ -133,12 +134,19 @@ func (m *Metrics) Flush() {
 	log.Debug().Str("file", m.file).Int("count", len(a)).Msg("Completed Metrics flush")
 }
 
-func (m *Metrics) Stop() {
+func (m *Metrics) Stop(ctx context.Context) {
 	log.Info().Msg("Stopping Request Metrics")
 	for len(m.c) > 0 {
-		rm := <- m.c
-		m.calculateDuration(rm)
+		select {
+		case rm := <-m.c:
+			m.calculateDuration(rm)
+		case <-ctx.Done():
+			m.stop <- struct{}{}
+			log.Warn().Msg("Hard Stopped Request Metrics")
+			return
+		}
 	}
 	m.Flush()
+	m.stop <- struct{}{}
 	log.Info().Msg("Stopped Request Metrics")
 }
diff --git a/middleware/pipeline.go b/middleware/pipeline.go
index 2513114..3f6e663 100644
--- a/middleware/pipeline.go
+++ b/middleware/pipeline.go
@@ -1,6 +1,7 @@
 package middleware
 
 import (
+	"context"
 	"net/http"
 	"slices"
 )
@@ -8,7 +9,7 @@ import (
 type Middleware interface {
 	Use(http.Handler) http.Handler
 	Manage()
-	Stop()
+	Stop(context.Context)
 }
 
 type Pipeline struct {
@@ -33,9 +34,9 @@ func (p *Pipeline) Use() func(http.Handler) http.Handler {
 	}
 }
 
-func (p *Pipeline) Stop() {
+func (p *Pipeline) Stop(ctx context.Context) {
 	for _, m := range p.middleware {
-		m.Stop()
+		m.Stop(ctx)
 	}
 }
 
diff --git a/middleware/rate-limit.go b/middleware/rate-limit.go
index c5dfcfa..3ccf778 100644
--- a/middleware/rate-limit.go
+++ b/middleware/rate-limit.go
@@ -1,6 +1,7 @@
 package middleware
 
 import (
+	"context"
 	"net/http"
 	"strings"
 	"sync"
@@ -37,7 +38,7 @@ func (l *Limiter) UpdateCleanupTime(new time.Duration) {
 	l.cleanupTicker.Reset(new)
 }
 
-func (l *Limiter) Stop() {
+func (l *Limiter) Stop(ctx context.Context) {
 	l.stop <- struct{}{}
 	log.Info().Msg("Stopped Ratelimits")
 }
@@ -79,7 +80,7 @@ func (l *Limiter) Manage() {
 			l.rwLock.Unlock()
 			duration := time.Since(start)
 			log.Debug().Str("duration", duration.String()).Int("deleted_buckets", deletedBuckets).Msg("Cleaned up Buckets")
-		case <- l.stop:
+		case <-l.stop:
 			return
 		}
 	}