From f0c54490ed9f8deec835cf59a59c16780a4374fe Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Wed, 6 Oct 2021 22:06:07 +0000
Subject: [PATCH] Allow multiple namespaces to be checked for state at the same
 time

---
 app.go     | 29 ++++++++++++++++++++++-------
 go.mod     |  1 +
 go.sum     |  2 ++
 machine.go | 23 +++++++++++------------
 4 files changed, 36 insertions(+), 19 deletions(-)

diff --git a/app.go b/app.go
index 78412a24..acaa7f1b 100644
--- a/app.go
+++ b/app.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"sort"
 	"strings"
 	"sync"
 	"time"
@@ -247,14 +248,28 @@ func (h *Headscale) setLastStateChangeToNow(namespace string) {
 	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)
+func (h *Headscale) getLastStateChange(namespaces ...string) time.Time {
+	times := []time.Time{}
+
+	for _, namespace := range namespaces {
+		if wrapped, ok := h.lastStateChange.Load(namespace); ok {
+			lastChange, _ := wrapped.(time.Time)
+
+			times = append(times, lastChange)
+		}
 
-		return lastChange
 	}
 
-	now := time.Now().UTC()
-	h.lastStateChange.Store(namespace, now)
-	return now
+	sort.Slice(times, func(i, j int) bool {
+		return times[i].After(times[j])
+	})
+
+	log.Trace().Msgf("Latest times %#v", times)
+
+	if len(times) == 0 {
+		return time.Now().UTC()
+
+	} else {
+		return times[0]
+	}
 }
diff --git a/go.mod b/go.mod
index 70428315..1fadd6b7 100644
--- a/go.mod
+++ b/go.mod
@@ -10,6 +10,7 @@ require (
 	github.com/docker/cli v20.10.8+incompatible // indirect
 	github.com/docker/docker v20.10.8+incompatible // indirect
 	github.com/efekarakus/termcolor v1.0.1
+	github.com/fatih/set v0.2.1 // indirect
 	github.com/gin-gonic/gin v1.7.4
 	github.com/gofrs/uuid v4.0.0+incompatible
 	github.com/google/go-github v17.0.0+incompatible // indirect
diff --git a/go.sum b/go.sum
index 6dc92b6d..b429ca95 100644
--- a/go.sum
+++ b/go.sum
@@ -204,6 +204,8 @@ github.com/fanliao/go-promise v0.0.0-20141029170127-1890db352a72/go.mod h1:Pjfxu
 github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
 github.com/fatih/color v1.9.0/go.mod h1:eQcE1qtQxscV5RaZvpXrrb8Drkc3/DdQ+uUYCNjL+zU=
 github.com/fatih/color v1.10.0/go.mod h1:ELkj/draVOlAH/xkhN6mQ50Qd0MPOk5AAr3maGEBuJM=
+github.com/fatih/set v0.2.1 h1:nn2CaJyknWE/6txyUDGwysr3G5QC6xWB/PtVjPBbeaA=
+github.com/fatih/set v0.2.1/go.mod h1:+RKtMCH+favT2+3YecHGxcc0b4KyVWA1QWWJUs4E0CI=
 github.com/flynn/go-shlex v0.0.0-20150515145356-3f9db97f8568/go.mod h1:xEzjJPgXI435gkrCt3MPfRiAkVrwSbHsst4LCFVfpJc=
 github.com/franela/goblin v0.0.0-20200105215937-c9ffbefa60db/go.mod h1:7dvUGVsVBjqR7JHJk0brhHOZYGmfBYOrK0ZhYMEtBr4=
 github.com/franela/goreq v0.0.0-20171204163338-bcd34c9993f8/go.mod h1:ZhphrRTfi2rbfLwlschooIH4+wKKDR4Pdxhh+TRoA20=
diff --git a/machine.go b/machine.go
index ebe6f84b..326c2fca 100644
--- a/machine.go
+++ b/machine.go
@@ -8,6 +8,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/fatih/set"
 	"github.com/rs/zerolog/log"
 
 	"gorm.io/datatypes"
@@ -221,23 +222,21 @@ func (h *Headscale) isOutdated(m *Machine) bool {
 
 	sharedMachines, _ := h.getShared(m)
 
+	namespaceSet := set.New(set.ThreadSafe)
+	namespaceSet.Add(m.Namespace.Name)
+
 	// Check if any of our shared namespaces has updates that we have
 	// not propagated.
 	for _, sharedMachine := range sharedMachines {
-		lastChange := h.getLastStateChange(sharedMachine.Namespace.Name)
-		log.Trace().
-			Str("func", "keepAlive").
-			Str("machine", m.Name).
-			Time("last_successful_update", *m.LastSuccessfulUpdate).
-			Time("last_state_change", lastChange).
-			Msgf("Checking if %s is missing updates", m.Name)
-		// Only return if we have a shared node with a newer update.
-		if m.LastSuccessfulUpdate.Before(lastChange) {
-			return true
-		}
+		namespaceSet.Add(sharedMachine.Namespace.Name)
 	}
 
-	lastChange := h.getLastStateChange(m.Namespace.Name)
+	namespaces := make([]string, namespaceSet.Size())
+	for index, namespace := range namespaceSet.List() {
+		namespaces[index] = namespace.(string)
+	}
+
+	lastChange := h.getLastStateChange(namespaces...)
 	log.Trace().
 		Str("func", "keepAlive").
 		Str("machine", m.Name).