From b0ec945dbb59196b7542386a2a543ddf5ba987b9 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Thu, 19 Aug 2021 18:19:26 +0100
Subject: [PATCH] Make lastStateChange namespaced

---
 app.go     | 40 ++++++++++++++++++++--------------------
 machine.go |  2 +-
 poll.go    |  6 +++---
 3 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/app.go b/app.go
index 255a7df1..76cf92ea 100644
--- a/app.go
+++ b/app.go
@@ -60,8 +60,7 @@ type Headscale struct {
 
 	clientsUpdateChannels sync.Map
 
-	lastStateChangeMutex sync.RWMutex
-	lastStateChange      time.Time
+	lastStateChange sync.Map
 }
 
 // NewHeadscale returns the Headscale app
@@ -88,13 +87,12 @@ func NewHeadscale(cfg Config) (*Headscale, error) {
 	}
 
 	h := Headscale{
-		cfg:             cfg,
-		dbType:          cfg.DBtype,
-		dbString:        dbString,
-		privateKey:      privKey,
-		publicKey:       &pubKey,
-		aclRules:        &tailcfg.FilterAllowAll, // default allowall
-		lastStateChange: time.Now().UTC(),
+		cfg:        cfg,
+		dbType:     cfg.DBtype,
+		dbString:   dbString,
+		privateKey: privKey,
+		publicKey:  &pubKey,
+		aclRules:   &tailcfg.FilterAllowAll, // default allowall
 	}
 
 	err = h.initDB()
@@ -229,17 +227,19 @@ func (h *Headscale) Serve() error {
 	return err
 }
 
-func (h *Headscale) setLastStateChangeToNow() {
-	h.lastStateChangeMutex.Lock()
+func (h *Headscale) setLastStateChangeToNow(namespace string) {
+	now := time.Now().UTC()
+	h.lastStateChange.Store(namespace, now)
+}
+
+func (h *Headscale) getLastStateChange(namespace string) time.Time {
+	if wrapped, ok := h.lastStateChange.Load(namespace); ok {
+		lastChange, _ := wrapped.(time.Time)
+		return lastChange
+
+	}
 
 	now := time.Now().UTC()
-	h.lastStateChange = now
-
-	h.lastStateChangeMutex.Unlock()
-}
-
-func (h *Headscale) getLastStateChange() time.Time {
-	h.lastStateChangeMutex.RLock()
-	defer h.lastStateChangeMutex.RUnlock()
-	return h.lastStateChange
+	h.lastStateChange.Store(namespace, now)
+	return now
 }
diff --git a/machine.go b/machine.go
index 13e3529a..5352f741 100644
--- a/machine.go
+++ b/machine.go
@@ -320,7 +320,7 @@ func (h *Headscale) isOutdated(m *Machine) bool {
 		return true
 	}
 
-	lastChange := h.getLastStateChange()
+	lastChange := h.getLastStateChange(m.Namespace.Name)
 	log.Trace().
 		Str("func", "keepAlive").
 		Str("machine", m.Name).
diff --git a/poll.go b/poll.go
index 522c529e..e85c7a9f 100644
--- a/poll.go
+++ b/poll.go
@@ -123,7 +123,7 @@ func (h *Headscale) PollNetMapHandler(c *gin.Context) {
 
 	// There has been an update to _any_ of the nodes that the other nodes would
 	// need to know about
-	h.setLastStateChangeToNow()
+	h.setLastStateChangeToNow(m.Namespace.Name)
 
 	// The request is not ReadOnly, so we need to set up channels for updating
 	// peers via longpoll
@@ -310,7 +310,7 @@ func (h *Headscale) PollNetMapStream(
 					Str("handler", "PollNetMapStream").
 					Str("machine", m.Name).
 					Time("last_successful_update", *m.LastSuccessfulUpdate).
-					Time("last_state_change", h.getLastStateChange()).
+					Time("last_state_change", h.getLastStateChange(m.Namespace.Name)).
 					Msgf("There has been updates since the last successful update to %s", m.Name)
 				data, err := h.getMapResponse(mKey, req, m)
 				if err != nil {
@@ -348,7 +348,7 @@ func (h *Headscale) PollNetMapStream(
 					Str("handler", "PollNetMapStream").
 					Str("machine", m.Name).
 					Time("last_successful_update", *m.LastSuccessfulUpdate).
-					Time("last_state_change", h.getLastStateChange()).
+					Time("last_state_change", h.getLastStateChange(m.Namespace.Name)).
 					Msgf("%s is up to date", m.Name)
 			}
 			return true