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

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

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 {