From a9c568c801a514855396c7dcec031b3598457f20 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@dalby.cc>
Date: Thu, 2 May 2024 13:39:19 +0200
Subject: [PATCH] trace log and notifier shutdown (#1922)

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 hscontrol/app.go               | 25 +++++++++++++++++++++++++
 hscontrol/notifier/notifier.go |  5 +++++
 2 files changed, 30 insertions(+)

diff --git a/hscontrol/app.go b/hscontrol/app.go
index ce2fd1d..e72aca2 100644
--- a/hscontrol/app.go
+++ b/hscontrol/app.go
@@ -800,10 +800,23 @@ func (h *Headscale) Serve() error {
 				}
 
 			default:
+				trace := log.Trace().Msgf
 				log.Info().
 					Str("signal", sig.String()).
 					Msg("Received signal to stop, shutting down gracefully")
 
+				trace("closing map sessions")
+				wg := sync.WaitGroup{}
+				for _, mapSess := range h.mapSessions {
+					wg.Add(1)
+					go func() {
+						mapSess.close()
+						wg.Done()
+					}()
+				}
+				wg.Wait()
+
+				trace("waiting for netmap stream to close")
 				h.pollNetMapStreamWG.Wait()
 
 				// Gracefully shut down servers
@@ -811,32 +824,44 @@ func (h *Headscale) Serve() error {
 					context.Background(),
 					types.HTTPShutdownTimeout,
 				)
+				trace("shutting down debug http server")
 				if err := debugHTTPServer.Shutdown(ctx); err != nil {
 					log.Error().Err(err).Msg("Failed to shutdown prometheus http")
 				}
+				trace("shutting down main http server")
 				if err := httpServer.Shutdown(ctx); err != nil {
 					log.Error().Err(err).Msg("Failed to shutdown http")
 				}
+
+				trace("shutting down grpc server (socket)")
 				grpcSocket.GracefulStop()
 
 				if grpcServer != nil {
+					trace("shutting down grpc server (external)")
 					grpcServer.GracefulStop()
 					grpcListener.Close()
 				}
 
 				if tailsqlContext != nil {
+					trace("shutting down tailsql")
 					tailsqlContext.Done()
 				}
 
+				trace("closing node notifier")
+				h.nodeNotifier.Close()
+
 				// Close network listeners
+				trace("closing network listeners")
 				debugHTTPListener.Close()
 				httpListener.Close()
 				grpcGatewayConn.Close()
 
 				// Stop listening (and unlink the socket if unix type):
+				trace("closing socket listener")
 				socketListener.Close()
 
 				// Close db connections
+				trace("closing database connection")
 				err = h.db.Close()
 				if err != nil {
 					log.Error().Err(err).Msg("Failed to close db")
diff --git a/hscontrol/notifier/notifier.go b/hscontrol/notifier/notifier.go
index 74b6645..d5ef89f 100644
--- a/hscontrol/notifier/notifier.go
+++ b/hscontrol/notifier/notifier.go
@@ -34,6 +34,11 @@ func NewNotifier(cfg *types.Config) *Notifier {
 	return n
 }
 
+// Close stops the batcher inside the notifier.
+func (n *Notifier) Close() {
+	n.b.close()
+}
+
 func (n *Notifier) AddNode(nodeID types.NodeID, c chan<- types.StateUpdate) {
 	log.Trace().Caller().Uint64("node.id", nodeID.Uint64()).Msg("acquiring lock to add node")
 	defer log.Trace().