Simplify code around latest state change map updates

juanfont · Oct 27, 2022 · 0e12b66 · 0e12b66
1 parent a647e6a
commit 0e12b66
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 27 deletions.
diff --git a/app.go b/app.go
@@ -24,7 +24,7 @@ import (
 	"github.com/patrickmn/go-cache"
 	zerolog "github.com/philip-bui/grpc-zerolog"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
-	"github.com/puzpuzpuz/xsync"
+	"github.com/puzpuzpuz/xsync/v2"
 	zl "github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 	"golang.org/x/crypto/acme"
@@ -94,7 +94,7 @@ type Headscale struct {
 	aclPolicy *ACLPolicy
 	aclRules  []tailcfg.FilterRule
 
-	lastStateChange *xsync.MapOf[time.Time]
+	lastStateChange *xsync.MapOf[string, time.Time]
 
 	oidcProvider *oidc.Provider
 	oauth2Config *oauth2.Config
@@ -884,7 +884,7 @@ func (h *Headscale) setLastStateChangeToNow() {
 
 	now := time.Now().UTC()
 
-	namespaces, err := h.ListNamespacesStr()
+	namespaces, err := h.ListNamespaces()
 	if err != nil {
 		log.Error().
 			Caller().
@@ -893,22 +893,22 @@ func (h *Headscale) setLastStateChangeToNow() {
 	}
 
 	for _, namespace := range namespaces {
-		lastStateUpdate.WithLabelValues(namespace, "headscale").Set(float64(now.Unix()))
+		lastStateUpdate.WithLabelValues(namespace.Name, "headscale").Set(float64(now.Unix()))
 		if h.lastStateChange == nil {
 			h.lastStateChange = xsync.NewMapOf[time.Time]()
 		}
-		h.lastStateChange.Store(namespace, now)
+		h.lastStateChange.Store(namespace.Name, now)
 	}
 }
 
-func (h *Headscale) getLastStateChange(namespaces ...string) time.Time {
+func (h *Headscale) getLastStateChange(namespaces ...Namespace) time.Time {
 	times := []time.Time{}
 
 	// getLastStateChange takes a list of namespaces as a "filter", if no namespaces
 	// are past, then use the entier list of namespaces and look for the last update
 	if len(namespaces) > 0 {
 		for _, namespace := range namespaces {
-			if lastChange, ok := h.lastStateChange.Load(namespace); ok {
+			if lastChange, ok := h.lastStateChange.Load(namespace.Name); ok {
 				times = append(times, lastChange)
 			}
 		}

diff --git a/go.mod b/go.mod
@@ -22,7 +22,6 @@ require (
 	github.com/prometheus/client_golang v1.13.0
 	github.com/prometheus/common v0.37.0
 	github.com/pterm/pterm v0.12.45
-	github.com/puzpuzpuz/xsync v1.4.3
 	github.com/rs/zerolog v1.28.0
 	github.com/spf13/cobra v1.5.0
 	github.com/spf13/viper v1.12.0
@@ -117,6 +116,7 @@ require (
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/prometheus/client_model v0.2.0 // indirect
 	github.com/prometheus/procfs v0.8.0 // indirect
+	github.com/puzpuzpuz/xsync/v2 v2.0.2 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect
 	github.com/rivo/uniseg v0.2.0 // indirect
 	github.com/rogpeppe/go-internal v1.8.1-0.20211023094830-115ce09fd6b4 // indirect

diff --git a/go.sum b/go.sum
@@ -846,8 +846,8 @@ github.com/pterm/pterm v0.12.36/go.mod h1:NjiL09hFhT/vWjQHSj1athJpx6H8cjpHXNAK5b
 github.com/pterm/pterm v0.12.40/go.mod h1:ffwPLwlbXxP+rxT0GsgDTzS3y3rmpAO1NMjUkGTYf8s=
 github.com/pterm/pterm v0.12.45 h1:5HATKLTDjl9D74b0x7yiHzFI7OADlSXK3yHrJNhRwZE=
 github.com/pterm/pterm v0.12.45/go.mod h1:hJgLlBafm45w/Hr0dKXxY//POD7CgowhePaG1sdPNBg=
-github.com/puzpuzpuz/xsync v1.4.3 h1:nS/Iqc4EnpJ8jm/MzJ+e3MUaP2Ys2mqXeEfoxoU0HaM=
-github.com/puzpuzpuz/xsync v1.4.3/go.mod h1:K98BYhX3k1dQ2M63t1YNVDanbwUPmBCAhNmVrrxfiGg=
+github.com/puzpuzpuz/xsync/v2 v2.0.2 h1:IpXQ8gGkrnZlLGpJLDmq56sYjNhF88n934Yq5BV5fKw=
+github.com/puzpuzpuz/xsync/v2 v2.0.2/go.mod h1:gD2H2krq/w52MfPLE+Uy64TzJDVY7lP2znR9qmR35kU=
 github.com/quasilyte/go-consistent v0.0.0-20190521200055-c6f3937de18c/go.mod h1:5STLWrekHfjyYwxBRVRXNOSewLJ3PWfDJd1VyTS21fI=
 github.com/quasilyte/go-ruleguard v0.3.1-0.20210203134552-1b5a410e1cc8/go.mod h1:KsAh3x0e7Fkpgs+Q9pNLS5XpFSvYCEVl5gP9Pp1xp30=
 github.com/quasilyte/go-ruleguard v0.3.13/go.mod h1:Ul8wwdqR6kBVOCt2dipDBkE+T6vAV/iixkrKuRTN1oQ=

diff --git a/namespaces.go b/namespaces.go
@@ -148,21 +148,6 @@ func (h *Headscale) ListNamespaces() ([]Namespace, error) {
 	return namespaces, nil
 }
 
-func (h *Headscale) ListNamespacesStr() ([]string, error) {
-	namespaces, err := h.ListNamespaces()
-	if err != nil {
-		return []string{}, err
-	}
-
-	namespaceStrs := make([]string, len(namespaces))
-
-	for index, namespace := range namespaces {
-		namespaceStrs[index] = namespace.Name
-	}
-
-	return namespaceStrs, nil
-}
-
 // ListMachinesInNamespace gets all the nodes in a given namespace.
 func (h *Headscale) ListMachinesInNamespace(name string) ([]Machine, error) {
 	err := CheckForFQDNRules(name)

diff --git a/protocol_common_poll.go b/protocol_common_poll.go
@@ -449,7 +449,7 @@ func (h *Headscale) pollNetMapStream(
 					Bool("noise", isNoise).
 					Str("machine", machine.Hostname).
 					Time("last_successful_update", lastUpdate).
-					Time("last_state_change", h.getLastStateChange(machine.Namespace.Name)).
+					Time("last_state_change", h.getLastStateChange(machine.Namespace)).
 					Msgf("There has been updates since the last successful update to %s", machine.Hostname)
 				data, err := h.getMapResponseData(mapRequest, machine, false)
 				if err != nil {
@@ -549,7 +549,7 @@ func (h *Headscale) pollNetMapStream(
 					Bool("noise", isNoise).
 					Str("machine", machine.Hostname).
 					Time("last_successful_update", lastUpdate).
-					Time("last_state_change", h.getLastStateChange(machine.Namespace.Name)).
+					Time("last_state_change", h.getLastStateChange(machine.Namespace)).
 					Msgf("%s is up to date", machine.Hostname)
 			}