diff --git a/app.go b/app.go index 695f18eb..e6a5121a 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 e214952a..5549afd8 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 } } }