fix: adding monitoring

v1/server/server.go

@@ -1,12 +1,14 @@
 package server
 
 import (
+	"context"
 	"crypto/subtle"
 	"fmt"
 	"log"
 	"net/http"
 	"strings"
 	"syscall"
+	"time"
 
 	"git.neurocipta.com/rogerferdinan/safe-web-socket/internal"
 	"github.com/gorilla/websocket"
@@ -29,6 +31,7 @@ type SafeWebsocketServerBuilder struct {
 	apiKey   *string             `nil_checker:"required"`
 	upgrader *websocket.Upgrader `nil_checker:"required"`
 	mux      *http.ServeMux      `nil_checker:"required"`
+	ctx      context.Context
 }
 
 func NewSafeWebsocketServerBuilder() *SafeWebsocketServerBuilder {
@@ -41,6 +44,7 @@ func NewSafeWebsocketServerBuilder() *SafeWebsocketServerBuilder {
 			},
 		},
 		mux: http.NewServeMux(),
+		ctx: context.Background(),
 	}
 }
 
@@ -59,16 +63,26 @@ func (b *SafeWebsocketServerBuilder) ApiKey(apiKey string) *SafeWebsocketServerB
 	return b
 }
 
+// Context sets the lifecycle context for all hub and writeFunc goroutines.
+// When ctx is cancelled, hubs stop dispatching and all active connections
+// are unblocked so their goroutines can exit cleanly.
+// Call this before HandleFuncWebsocket. Defaults to context.Background().
+func (b *SafeWebsocketServerBuilder) Context(ctx context.Context) *SafeWebsocketServerBuilder {
+	b.ctx = ctx
+	return b
+}
+
 func (b *SafeWebsocketServerBuilder) HandleFunc(pattern string, fn func(http.ResponseWriter, *http.Request)) *SafeWebsocketServerBuilder {
 	b.mux.HandleFunc(pattern, fn)
 	return b
 }
 
-func (b *SafeWebsocketServerBuilder) HandleFuncWebsocket(pattern string, subscribedPath string, writeFunc func(writeChannel chan []byte)) *SafeWebsocketServerBuilder {
-	h := internal.NewHub()
-	h.Run()
+// HandleFuncWebsocket registers a WebSocket endpoint.
+func (b *SafeWebsocketServerBuilder) HandleFuncWebsocket(pattern string, subscribedPath string, maxClients int, writeFunc func(ctx context.Context, writeChannel chan []byte)) *SafeWebsocketServerBuilder {
+	h := internal.NewHub(pattern+subscribedPath, maxClients)
+	h.Run(b.ctx)
 
-	go writeFunc(h.Broadcast)
+	go writeFunc(b.ctx, h.Broadcast)
 
 	b.mux.HandleFunc(pattern+subscribedPath, func(w http.ResponseWriter, r *http.Request) {
 		conn, err := b.upgrader.Upgrade(w, r, nil)
@@ -87,7 +101,6 @@ func (b *SafeWebsocketServerBuilder) HandleFuncWebsocket(pattern string, subscri
 
 		go internal.WritePump(c, h)
 		go internal.ReadPump(c, h)
-
 	})
 	return b
 }
@@ -102,6 +115,7 @@ func (b *SafeWebsocketServerBuilder) Build() (*SafeWebsocketServer, error) {
 		mux:    b.mux,
 		url:    fmt.Sprintf("%s:%d", *b.baseHost, *b.basePort),
 		apiKey: *b.apiKey,
+		ctx:    b.ctx,
 	}
 	return &safeServer, nil
 }
@@ -110,6 +124,7 @@ type SafeWebsocketServer struct {
 	mux    *http.ServeMux
 	url    string
 	apiKey string
+	ctx    context.Context
 }
 
 func (s *SafeWebsocketServer) AuthMiddleware(next http.Handler) http.Handler {
@@ -130,9 +145,30 @@ func (s *SafeWebsocketServer) AuthMiddleware(next http.Handler) http.Handler {
 }
 
 func (s *SafeWebsocketServer) ListenAndServe() error {
-	log.Printf("HTTP serve on %s\n", s.url)
-	if err := http.ListenAndServe(s.url, s.AuthMiddleware(s.mux)); err != nil {
-		return fmt.Errorf("failed to serve websocket: %w", err)
+	srv := &http.Server{
+		Addr:    s.url,
+		Handler: s.AuthMiddleware(s.mux),
+	}
+
+	errCh := make(chan error, 1)
+	go func() {
+		log.Printf("HTTP serve on %s\n", s.url)
+		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+			errCh <- fmt.Errorf("failed to serve websocket: %w", err)
+		}
+		close(errCh)
+	}()
+
+	select {
+	case err := <-errCh:
+		return err
+	case <-s.ctx.Done():
+		log.Println("Server shutting down...")
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		if err := srv.Shutdown(shutdownCtx); err != nil {
+			return fmt.Errorf("server shutdown: %w", err)
+		}
+		return <-errCh
 	}
-	return nil
 }