Fix slow shutdown

cmd/headscale/cli/serve.go

@@ -1,6 +1,9 @@
 package cli
 
 import (
+	"errors"
+	"net/http"
+
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
 )
@@ -22,8 +25,8 @@ var serveCmd = &cobra.Command{
 		}
 
 		err = app.Serve()
-		if err != nil {
-			log.Fatal().Caller().Err(err).Msg("Error starting server")
+		if err != nil && !errors.Is(err, http.ErrServerClosed) {
+			log.Fatal().Caller().Err(err).Msg("Headscale ran into an error and had to shut down.")
 		}
 	},
 }

hscontrol/app.go

@@ -780,9 +780,6 @@ func (h *Headscale) Serve() error {
 				expireNodeCancel()
 				h.ephemeralGC.Close()
 
-				trace("waiting for netmap stream to close")
-				h.pollNetMapStreamWG.Wait()
-
 				// Gracefully shut down servers
 				ctx, cancel := context.WithTimeout(
 					context.Background(),
@@ -797,6 +794,12 @@ func (h *Headscale) Serve() error {
 					log.Error().Err(err).Msg("Failed to shutdown http")
 				}
 
+				trace("closing node notifier")
+				h.nodeNotifier.Close()
+
+				trace("waiting for netmap stream to close")
+				h.pollNetMapStreamWG.Wait()
+
 				trace("shutting down grpc server (socket)")
 				grpcSocket.GracefulStop()
 
@@ -811,9 +814,6 @@ func (h *Headscale) Serve() error {
 					tailsqlContext.Done()
 				}
 
-				trace("closing node notifier")
-				h.nodeNotifier.Close()
-
 				// Close network listeners
 				trace("closing network listeners")
 				debugHTTPListener.Close()

hscontrol/notifier/notifier.go

@@ -36,13 +36,15 @@ type Notifier struct {
 	connected *xsync.MapOf[types.NodeID, bool]
 	b         *batcher
 	cfg       *types.Config
+	closed    bool
 }
 
 func NewNotifier(cfg *types.Config) *Notifier {
 	n := &Notifier{
 		nodes:     make(map[types.NodeID]chan<- types.StateUpdate),
 		connected: xsync.NewMapOf[types.NodeID, bool](),
 		cfg:       cfg,
+		closed:    false,
 	}
 	b := newBatcher(cfg.Tuning.BatchChangeDelay, n)
 	n.b = b
@@ -51,9 +53,19 @@ func NewNotifier(cfg *types.Config) *Notifier {
 	return n
 }
 
-// Close stops the batcher inside the notifier.
+// Close stops the batcher and closes all channels.
 func (n *Notifier) Close() {
+	notifierWaitersForLock.WithLabelValues("lock", "close").Inc()
+	n.l.Lock()
+	defer n.l.Unlock()
+	notifierWaitersForLock.WithLabelValues("lock", "close").Dec()
+
+	n.closed = true
 	n.b.close()
+
+	for _, c := range n.nodes {
+		close(c)
+	}
 }
 
 func (n *Notifier) tracef(nID types.NodeID, msg string, args ...any) {
@@ -70,6 +82,10 @@ func (n *Notifier) AddNode(nodeID types.NodeID, c chan<- types.StateUpdate) {
 	notifierWaitersForLock.WithLabelValues("lock", "add").Dec()
 	notifierWaitForLock.WithLabelValues("add").Observe(time.Since(start).Seconds())
 
+	if n.closed {
+		return
+	}
+
 	// If a channel exists, it means the node has opened a new
 	// connection. Close the old channel and replace it.
 	if curr, ok := n.nodes[nodeID]; ok {
@@ -96,6 +112,10 @@ func (n *Notifier) RemoveNode(nodeID types.NodeID, c chan<- types.StateUpdate) b
 	notifierWaitersForLock.WithLabelValues("lock", "remove").Dec()
 	notifierWaitForLock.WithLabelValues("remove").Observe(time.Since(start).Seconds())
 
+	if n.closed {
+		return true
+	}
+
 	if len(n.nodes) == 0 {
 		return true
 	}
@@ -154,6 +174,10 @@ func (n *Notifier) NotifyWithIgnore(
 	update types.StateUpdate,
 	ignoreNodeIDs ...types.NodeID,
 ) {
+	if n.closed {
+		return
+	}
+
 	notifierUpdateReceived.WithLabelValues(update.Type.String(), types.NotifyOriginKey.Value(ctx)).Inc()
 	n.b.addOrPassthrough(update)
 }
@@ -170,6 +194,10 @@ func (n *Notifier) NotifyByNodeID(
 	notifierWaitersForLock.WithLabelValues("lock", "notify").Dec()
 	notifierWaitForLock.WithLabelValues("notify").Observe(time.Since(start).Seconds())
 
+	if n.closed {
+		return
+	}
+
 	if c, ok := n.nodes[nodeID]; ok {
 		select {
 		case <-ctx.Done():
@@ -205,6 +233,10 @@ func (n *Notifier) sendAll(update types.StateUpdate) {
 	notifierWaitersForLock.WithLabelValues("lock", "send-all").Dec()
 	notifierWaitForLock.WithLabelValues("send-all").Observe(time.Since(start).Seconds())
 
+	if n.closed {
+		return
+	}
+
 	for id, c := range n.nodes {
 		// Whenever an update is sent to all nodes, there is a chance that the node
 		// has disconnected and the goroutine that was supposed to consume the update