From 39b58f7d4cb68ae22c1442fc9f899a81990a1228 Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Thu, 23 Jun 2022 19:40:07 +0200
Subject: [PATCH] Use a signal to close the longpolls on shutdown

---
 app.go  | 5 +++++
 poll.go | 7 +++++++
 2 files changed, 12 insertions(+)

diff --git a/app.go b/app.go
index 695f18e..e6a5121 100644
--- a/app.go
+++ b/app.go
@@ -93,6 +93,8 @@ type Headscale struct {
 	registrationCache *cache.Cache
 
 	ipAllocationMutex sync.Mutex
+
+	shutdownChan chan struct{}
 }
 
 // Look up the TLS constant relative to user-supplied TLS client
@@ -642,6 +644,7 @@ func (h *Headscale) Serve() error {
 		Msgf("listening and serving metrics on: %s", h.cfg.MetricsAddr)
 
 	// Handle common process-killing signals so we can gracefully shut down:
+	h.shutdownChan = make(chan struct{})
 	sigc := make(chan os.Signal, 1)
 	signal.Notify(sigc,
 		syscall.SIGHUP,
@@ -679,6 +682,8 @@ func (h *Headscale) Serve() error {
 					Str("signal", sig.String()).
 					Msg("Received signal to stop, shutting down gracefully")
 
+				h.shutdownChan <- struct{}{}
+
 				// Gracefully shut down servers
 				promHTTPServer.Shutdown(ctx)
 				httpServer.Shutdown(ctx)
diff --git a/poll.go b/poll.go
index e214952..5549afd 100644
--- a/poll.go
+++ b/poll.go
@@ -572,6 +572,13 @@ func (h *Headscale) PollNetMapStream(
 
 			// The connection has been closed, so we can stop polling.
 			return
+
+		case <-h.shutdownChan:
+			log.Info().
+				Str("handler", "PollNetMapStream").
+				Str("machine", machine.Hostname).
+				Msg("The long-poll handler is shutting down")
+			return
 		}
 	}
 }