From 86a0b0cf4fa4a6f6d5391a45be7daec14b854754 Mon Sep 17 00:00:00 2001
From: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>
Date: Sun, 24 Aug 2025 08:35:17 +0000
Subject: [PATCH] Small fixes

* Shut down the web server first to prevent errors caused by clients trying
to use functionality that has already been shut down, causing errors and
potentially delaying the shutdown process.
* remove write timeout from the websocket Write() function.

Signed-off-by: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>
---
 cmd/garm/main.go       | 14 +++++++-------
 websocket/websocket.go | 18 ++++++------------
 2 files changed, 13 insertions(+), 19 deletions(-)

diff --git a/cmd/garm/main.go b/cmd/garm/main.go
index bfc23d50..28879322 100644
--- a/cmd/garm/main.go
+++ b/cmd/garm/main.go
@@ -372,6 +372,13 @@ func main() {
 
 	<-ctx.Done()
 
+	slog.InfoContext(ctx, "shutting down http server")
+	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer shutdownCancel()
+	if err := srv.Shutdown(shutdownCtx); err != nil {
+		slog.With(slog.Any("error", err)).ErrorContext(ctx, "graceful api server shutdown failed")
+	}
+
 	if err := cacheWorker.Stop(); err != nil {
 		slog.With(slog.Any("error", err)).ErrorContext(ctx, "failed to stop credentials worker")
 	}
@@ -386,13 +393,6 @@ func main() {
 		slog.With(slog.Any("error", err)).ErrorContext(ctx, "failed to stop provider worker")
 	}
 
-	slog.InfoContext(ctx, "shutting down http server")
-	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 60*time.Second)
-	defer shutdownCancel()
-	if err := srv.Shutdown(shutdownCtx); err != nil {
-		slog.With(slog.Any("error", err)).ErrorContext(ctx, "graceful api server shutdown failed")
-	}
-
 	slog.With(slog.Any("error", err)).InfoContext(ctx, "waiting for runner to stop")
 	if err := runner.Wait(); err != nil {
 		slog.With(slog.Any("error", err)).ErrorContext(ctx, "failed to shutdown workers")
diff --git a/websocket/websocket.go b/websocket/websocket.go
index e59eb70e..d1f626e6 100644
--- a/websocket/websocket.go
+++ b/websocket/websocket.go
@@ -118,23 +118,17 @@ func (h *Hub) Unregister(client *Client) error {
 }
 
 func (h *Hub) Write(msg []byte) (int, error) {
-	h.mux.Lock()
-	if !h.running {
-		h.mux.Unlock()
-		return 0, fmt.Errorf("websocket writer is not running")
-	}
-	h.mux.Unlock()
-
 	tmp := make([]byte, len(msg))
 	copy(tmp, msg)
-	timer := time.NewTimer(5 * time.Second)
-	defer timer.Stop()
+
 	select {
-	case <-timer.C:
-		return 0, fmt.Errorf("timed out sending message to client")
 	case h.broadcast <- tmp:
+		return len(tmp), nil
+	case <-h.quit:
+		return 0, fmt.Errorf("websocket hub is shutting down")
+	default:
+		return 0, fmt.Errorf("failed to broadcast over websocket")
 	}
-	return len(tmp), nil
 }
 
 func (h *Hub) Start() error {