Added context.Context to stop and correctly stop metrics now

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 {

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")
 }
 

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")
 }

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)
 	}
 }
 

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
 		}
 	}