Replace the timestamp based state system

This commit replaces the timestamp based state system with a new
one that has update channels directly to the connected nodes. It
will send an update to all listening clients via the polling
mechanism.

It introduces a new package notifier, which has a concurrency safe
manager for all our channels to the connected nodes.

Signed-off-by: Kristoffer Dalby <[email protected]>

hscontrol/app.go

@@ -10,7 +10,6 @@ import (
 	"net/http"
 	"os"
 	"os/signal"
-	"sort"
 	"strconv"
 	"strings"
 	"sync"
@@ -26,13 +25,13 @@ import (
 	"github.com/juanfont/headscale/hscontrol/db"
 	"github.com/juanfont/headscale/hscontrol/derp"
 	derpServer "github.com/juanfont/headscale/hscontrol/derp/server"
+	"github.com/juanfont/headscale/hscontrol/notifier"
 	"github.com/juanfont/headscale/hscontrol/policy"
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/juanfont/headscale/hscontrol/util"
 	"github.com/patrickmn/go-cache"
 	zerolog "github.com/philip-bui/grpc-zerolog"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
-	"github.com/puzpuzpuz/xsync/v2"
 	zl "github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 	"golang.org/x/crypto/acme"
@@ -84,7 +83,7 @@ type Headscale struct {
 
 	ACLPolicy *policy.ACLPolicy
 
-	lastStateChange *xsync.MapOf[string, time.Time]
+	nodeNotifier *notifier.Notifier
 
 	oidcProvider *oidc.Provider
 	oauth2Config *oauth2.Config
@@ -93,9 +92,6 @@ type Headscale struct {
 
 	shutdownChan       chan struct{}
 	pollNetMapStreamWG sync.WaitGroup
-
-	stateUpdateChan       chan struct{}
-	cancelStateUpdateChan chan struct{}
 }
 
 func NewHeadscale(cfg *types.Config) (*Headscale, error) {
@@ -158,19 +154,14 @@ func NewHeadscale(cfg *types.Config) (*Headscale, error) {
 		noisePrivateKey:    noisePrivateKey,
 		registrationCache:  registrationCache,
 		pollNetMapStreamWG: sync.WaitGroup{},
-		lastStateChange:    xsync.NewMapOf[time.Time](),
-
-		stateUpdateChan:       make(chan struct{}),
-		cancelStateUpdateChan: make(chan struct{}),
+		nodeNotifier:       notifier.NewNotifier(),
 	}
 
-	go app.watchStateChannel()
-
 	database, err := db.NewHeadscaleDatabase(
 		cfg.DBtype,
 		dbString,
 		app.dbDebug,
-		app.stateUpdateChan,
+		app.nodeNotifier,
 		cfg.IPPrefixes,
 		cfg.BaseDomain)
 	if err != nil {
@@ -203,7 +194,11 @@ func NewHeadscale(cfg *types.Config) (*Headscale, error) {
 
 	if cfg.DERP.ServerEnabled {
 		// TODO(kradalby): replace this key with a dedicated DERP key.
-		embeddedDERPServer, err := derpServer.NewDERPServer(cfg.ServerURL, key.NodePrivate(*privateKey), &cfg.DERP)
+		embeddedDERPServer, err := derpServer.NewDERPServer(
+			cfg.ServerURL,
+			key.NodePrivate(*privateKey),
+			&cfg.DERP,
+		)
 		if err != nil {
 			return nil, err
 		}
@@ -230,10 +225,14 @@ func (h *Headscale) expireEphemeralNodes(milliSeconds int64) {
 
 // expireExpiredMachines expires machines that have an explicit expiry set
 // after that expiry time has passed.
-func (h *Headscale) expireExpiredMachines(milliSeconds int64) {
-	ticker := time.NewTicker(time.Duration(milliSeconds) * time.Millisecond)
+func (h *Headscale) expireExpiredMachines(intervalMs int64) {
+	interval := time.Duration(intervalMs) * time.Millisecond
+	ticker := time.NewTicker(interval)
+
+	lastCheck := time.Unix(0, 0)
+
 	for range ticker.C {
-		h.db.ExpireExpiredMachines(h.getLastStateChange())
+		lastCheck = h.db.ExpireExpiredMachines(lastCheck)
 	}
 }
 
@@ -258,7 +257,7 @@ func (h *Headscale) scheduledDERPMapUpdateWorker(cancelChan <-chan struct{}) {
 				h.DERPMap.Regions[region.RegionID] = &region
 			}
 
-			h.setLastStateChangeToNow()
+			h.nodeNotifier.NotifyAll()
 		}
 	}
 }
@@ -722,7 +721,7 @@ func (h *Headscale) Serve() error {
 						Str("path", aclPath).
 						Msg("ACL policy successfully reloaded, notifying nodes of change")
 
-					h.setLastStateChangeToNow()
+					h.nodeNotifier.NotifyAll()
 				}
 
 			default:
@@ -760,10 +759,6 @@ func (h *Headscale) Serve() error {
 				// Stop listening (and unlink the socket if unix type):
 				socketListener.Close()
 
-				<-h.cancelStateUpdateChan
-				close(h.stateUpdateChan)
-				close(h.cancelStateUpdateChan)
-
 				// Close db connections
 				err = h.db.Close()
 				if err != nil {
@@ -859,73 +854,6 @@ func (h *Headscale) getTLSSettings() (*tls.Config, error) {
 	}
 }
 
-// TODO(kradalby): baby steps, make this more robust.
-func (h *Headscale) watchStateChannel() {
-	for {
-		select {
-		case <-h.stateUpdateChan:
-			h.setLastStateChangeToNow()
-
-		case <-h.cancelStateUpdateChan:
-			return
-		}
-	}
-}
-
-func (h *Headscale) setLastStateChangeToNow() {
-	var err error
-
-	now := time.Now().UTC()
-
-	users, err := h.db.ListUsers()
-	if err != nil {
-		log.Error().
-			Caller().
-			Err(err).
-			Msg("failed to fetch all users, failing to update last changed state.")
-	}
-
-	for _, user := range users {
-		lastStateUpdate.WithLabelValues(user.Name, "headscale").Set(float64(now.Unix()))
-		if h.lastStateChange == nil {
-			h.lastStateChange = xsync.NewMapOf[time.Time]()
-		}
-		h.lastStateChange.Store(user.Name, now)
-	}
-}
-
-func (h *Headscale) getLastStateChange(users ...types.User) time.Time {
-	times := []time.Time{}
-
-	// getLastStateChange takes a list of users as a "filter", if no users
-	// are past, then use the entier list of users and look for the last update
-	if len(users) > 0 {
-		for _, user := range users {
-			if lastChange, ok := h.lastStateChange.Load(user.Name); ok {
-				times = append(times, lastChange)
-			}
-		}
-	} else {
-		h.lastStateChange.Range(func(key string, value time.Time) bool {
-			times = append(times, value)
-
-			return true
-		})
-	}
-
-	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]
-	}
-}
-
 func notFoundHandler(
 	writer http.ResponseWriter,
 	req *http.Request,

hscontrol/db/addresses_test.go

@@ -63,8 +63,6 @@ func (s *Suite) TestGetUsedIps(c *check.C) {
 
 	c.Assert(len(machine1.IPAddresses), check.Equals, 1)
 	c.Assert(machine1.IPAddresses[0], check.Equals, expected)
-
-	c.Assert(channelUpdates, check.Equals, int32(0))
 }
 
 func (s *Suite) TestGetMultiIp(c *check.C) {
@@ -153,8 +151,6 @@ func (s *Suite) TestGetMultiIp(c *check.C) {
 
 	c.Assert(len(nextIP2), check.Equals, 1)
 	c.Assert(nextIP2[0].String(), check.Equals, expectedNextIP.String())
-
-	c.Assert(channelUpdates, check.Equals, int32(0))
 }
 
 func (s *Suite) TestGetAvailableIpMachineWithoutIP(c *check.C) {
@@ -192,6 +188,4 @@ func (s *Suite) TestGetAvailableIpMachineWithoutIP(c *check.C) {
 
 	c.Assert(len(ips2), check.Equals, 1)
 	c.Assert(ips2[0].String(), check.Equals, expected.String())
-
-	c.Assert(channelUpdates, check.Equals, int32(0))
 }

hscontrol/db/api_key_test.go

@@ -22,8 +22,6 @@ func (*Suite) TestCreateAPIKey(c *check.C) {
 	keys, err := db.ListAPIKeys()
 	c.Assert(err, check.IsNil)
 	c.Assert(len(keys), check.Equals, 1)
-
-	c.Assert(channelUpdates, check.Equals, int32(0))
 }
 
 func (*Suite) TestAPIKeyDoesNotExist(c *check.C) {
@@ -41,8 +39,6 @@ func (*Suite) TestValidateAPIKeyOk(c *check.C) {
 	valid, err := db.ValidateAPIKey(apiKeyStr)
 	c.Assert(err, check.IsNil)
 	c.Assert(valid, check.Equals, true)
-
-	c.Assert(channelUpdates, check.Equals, int32(0))
 }
 
 func (*Suite) TestValidateAPIKeyNotOk(c *check.C) {
@@ -71,8 +67,6 @@ func (*Suite) TestValidateAPIKeyNotOk(c *check.C) {
 	validWithErr, err := db.ValidateAPIKey("produceerrorkey")
 	c.Assert(err, check.NotNil)
 	c.Assert(validWithErr, check.Equals, false)
-
-	c.Assert(channelUpdates, check.Equals, int32(0))
 }
 
 func (*Suite) TestExpireAPIKey(c *check.C) {
@@ -92,6 +86,4 @@ func (*Suite) TestExpireAPIKey(c *check.C) {
 	notValid, err := db.ValidateAPIKey(apiKeyStr)
 	c.Assert(err, check.IsNil)
 	c.Assert(notValid, check.Equals, false)
-
-	c.Assert(channelUpdates, check.Equals, int32(0))
 }

hscontrol/db/db.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/glebarez/sqlite"
+	"github.com/juanfont/headscale/hscontrol/notifier"
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/juanfont/headscale/hscontrol/util"
 	"github.com/rs/zerolog/log"
@@ -36,8 +37,8 @@ type KV struct {
 }
 
 type HSDatabase struct {
-	db              *gorm.DB
-	notifyStateChan chan<- struct{}
+	db       *gorm.DB
+	notifier *notifier.Notifier
 
 	ipAllocationMutex sync.Mutex
 
@@ -50,7 +51,7 @@ type HSDatabase struct {
 func NewHeadscaleDatabase(
 	dbType, connectionAddr string,
 	debug bool,
-	notifyStateChan chan<- struct{},
+	notifier *notifier.Notifier,
 	ipPrefixes []netip.Prefix,
 	baseDomain string,
 ) (*HSDatabase, error) {
@@ -60,8 +61,8 @@ func NewHeadscaleDatabase(
 	}
 
 	db := HSDatabase{
-		db:              dbConn,
-		notifyStateChan: notifyStateChan,
+		db:       dbConn,
+		notifier: notifier,
 
 		ipPrefixes: ipPrefixes,
 		baseDomain: baseDomain,
@@ -297,10 +298,6 @@ func openDB(dbType, connectionAddr string, debug bool) (*gorm.DB, error) {
 	)
 }
 
-func (hsdb *HSDatabase) notifyStateChange() {
-	hsdb.notifyStateChan <- struct{}{}
-}
-
 // getValue returns the value for the given key in KV.
 func (hsdb *HSDatabase) getValue(key string) (string, error) {
 	var row KV