use concurrency save connectedmap, more metrics

Fixes juanfont#1862

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

hscontrol/app.go

@@ -509,7 +509,7 @@ func (h *Headscale) Serve() error {
 
 	// Fetch an initial DERP Map before we start serving
 	h.DERPMap = derp.GetDERPMap(h.cfg.DERP)
-	h.mapper = mapper.NewMapper(h.db, h.cfg, h.DERPMap, h.nodeNotifier.ConnectedMap())
+	h.mapper = mapper.NewMapper(h.db, h.cfg, h.DERPMap, h.nodeNotifier)
 
 	if h.cfg.DERP.ServerEnabled {
 		// When embedded DERP is enabled we always need a STUN server

hscontrol/auth.go

@@ -546,7 +546,7 @@ func (h *Headscale) handleNodeLogOut(
 	}
 
 	if node.IsEphemeral() {
-		changedNodes, err := h.db.DeleteNode(&node, h.nodeNotifier.ConnectedMap())
+		changedNodes, err := h.db.DeleteNode(&node, h.nodeNotifier.LikelyConnectedMap())
 		if err != nil {
 			log.Error().
 				Err(err).

hscontrol/db/node.go

@@ -10,6 +10,7 @@ import (
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/juanfont/headscale/hscontrol/util"
 	"github.com/patrickmn/go-cache"
+	"github.com/puzpuzpuz/xsync/v3"
 	"github.com/rs/zerolog/log"
 	"gorm.io/gorm"
 	"tailscale.com/tailcfg"
@@ -260,19 +261,19 @@ func NodeSetExpiry(tx *gorm.DB,
 	return tx.Model(&types.Node{}).Where("id = ?", nodeID).Update("expiry", expiry).Error
 }
 
-func (hsdb *HSDatabase) DeleteNode(node *types.Node, isConnected types.NodeConnectedMap) ([]types.NodeID, error) {
+func (hsdb *HSDatabase) DeleteNode(node *types.Node, isLikelyConnected *xsync.MapOf[types.NodeID, bool]) ([]types.NodeID, error) {
 	return Write(hsdb.DB, func(tx *gorm.DB) ([]types.NodeID, error) {
-		return DeleteNode(tx, node, isConnected)
+		return DeleteNode(tx, node, isLikelyConnected)
 	})
 }
 
 // DeleteNode deletes a Node from the database.
 // Caller is responsible for notifying all of change.
 func DeleteNode(tx *gorm.DB,
 	node *types.Node,
-	isConnected types.NodeConnectedMap,
+	isLikelyConnected *xsync.MapOf[types.NodeID, bool],
 ) ([]types.NodeID, error) {
-	changed, err := deleteNodeRoutes(tx, node, isConnected)
+	changed, err := deleteNodeRoutes(tx, node, isLikelyConnected)
 	if err != nil {
 		return changed, err
 	}

hscontrol/db/node_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/juanfont/headscale/hscontrol/policy"
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/juanfont/headscale/hscontrol/util"
+	"github.com/puzpuzpuz/xsync/v3"
 	"gopkg.in/check.v1"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/key"
@@ -120,7 +121,7 @@ func (s *Suite) TestHardDeleteNode(c *check.C) {
 	}
 	db.DB.Save(&node)
 
-	_, err = db.DeleteNode(&node, types.NodeConnectedMap{})
+	_, err = db.DeleteNode(&node, xsync.NewMapOf[types.NodeID, bool]())
 	c.Assert(err, check.IsNil)
 
 	_, err = db.getNode(user.Name, "testnode3")

hscontrol/db/routes.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/juanfont/headscale/hscontrol/policy"
 	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/puzpuzpuz/xsync/v3"
 	"github.com/rs/zerolog/log"
 	"gorm.io/gorm"
 	"tailscale.com/util/set"
@@ -126,7 +127,7 @@ func EnableRoute(tx *gorm.DB, id uint64) (*types.StateUpdate, error) {
 
 func DisableRoute(tx *gorm.DB,
 	id uint64,
-	isConnected types.NodeConnectedMap,
+	isLikelyConnected *xsync.MapOf[types.NodeID, bool],
 ) ([]types.NodeID, error) {
 	route, err := GetRoute(tx, id)
 	if err != nil {
@@ -147,7 +148,7 @@ func DisableRoute(tx *gorm.DB,
 			return nil, err
 		}
 
-		update, err = failoverRouteTx(tx, isConnected, route)
+		update, err = failoverRouteTx(tx, isLikelyConnected, route)
 		if err != nil {
 			return nil, err
 		}
@@ -182,17 +183,17 @@ func DisableRoute(tx *gorm.DB,
 
 func (hsdb *HSDatabase) DeleteRoute(
 	id uint64,
-	isConnected types.NodeConnectedMap,
+	isLikelyConnected *xsync.MapOf[types.NodeID, bool],
 ) ([]types.NodeID, error) {
 	return Write(hsdb.DB, func(tx *gorm.DB) ([]types.NodeID, error) {
-		return DeleteRoute(tx, id, isConnected)
+		return DeleteRoute(tx, id, isLikelyConnected)
 	})
 }
 
 func DeleteRoute(
 	tx *gorm.DB,
 	id uint64,
-	isConnected types.NodeConnectedMap,
+	isLikelyConnected *xsync.MapOf[types.NodeID, bool],
 ) ([]types.NodeID, error) {
 	route, err := GetRoute(tx, id)
 	if err != nil {
@@ -207,7 +208,7 @@ func DeleteRoute(
 	// https://github.com/juanfont/headscale/issues/804#issuecomment-1399314002
 	var update []types.NodeID
 	if !route.IsExitRoute() {
-		update, err = failoverRouteTx(tx, isConnected, route)
+		update, err = failoverRouteTx(tx, isLikelyConnected, route)
 		if err != nil {
 			return nil, nil
 		}
@@ -252,7 +253,7 @@ func DeleteRoute(
 	return update, nil
 }
 
-func deleteNodeRoutes(tx *gorm.DB, node *types.Node, isConnected types.NodeConnectedMap) ([]types.NodeID, error) {
+func deleteNodeRoutes(tx *gorm.DB, node *types.Node, isLikelyConnected *xsync.MapOf[types.NodeID, bool]) ([]types.NodeID, error) {
 	routes, err := GetNodeRoutes(tx, node)
 	if err != nil {
 		return nil, fmt.Errorf("getting node routes: %w", err)
@@ -266,7 +267,7 @@ func deleteNodeRoutes(tx *gorm.DB, node *types.Node, isConnected types.NodeConne
 
 		// TODO(kradalby): This is a bit too aggressive, we could probably
 		// figure out which routes needs to be failed over rather than all.
-		chn, err := failoverRouteTx(tx, isConnected, &routes[i])
+		chn, err := failoverRouteTx(tx, isLikelyConnected, &routes[i])
 		if err != nil {
 			return changed, fmt.Errorf("failing over route after delete: %w", err)
 		}
@@ -409,7 +410,7 @@ func SaveNodeRoutes(tx *gorm.DB, node *types.Node) (bool, error) {
 // If needed, the failover will be attempted.
 func FailoverNodeRoutesIfNeccessary(
 	tx *gorm.DB,
-	isConnected types.NodeConnectedMap,
+	isLikelyConnected *xsync.MapOf[types.NodeID, bool],
 	node *types.Node,
 ) (*types.StateUpdate, error) {
 	nodeRoutes, err := GetNodeRoutes(tx, node)
@@ -430,12 +431,12 @@ nodeRouteLoop:
 			if route.IsPrimary {
 				// if we have a primary route, and the node is connected
 				// nothing needs to be done.
-				if conn, ok := isConnected[route.Node.ID]; conn && ok {
+				if val, ok := isLikelyConnected.Load(route.Node.ID); ok && val {
 					continue nodeRouteLoop
 				}
 
 				// if not, we need to failover the route
-				failover := failoverRoute(isConnected, &route, routes)
+				failover := failoverRoute(isLikelyConnected, &route, routes)
 				if failover != nil {
 					err := failover.save(tx)
 					if err != nil {
@@ -477,7 +478,7 @@ nodeRouteLoop:
 // If the given route was not primary, it returns early.
 func failoverRouteTx(
 	tx *gorm.DB,
-	isConnected types.NodeConnectedMap,
+	isLikelyConnected *xsync.MapOf[types.NodeID, bool],
 	r *types.Route,
 ) ([]types.NodeID, error) {
 	if r == nil {
@@ -500,7 +501,7 @@ func failoverRouteTx(
 		return nil, fmt.Errorf("getting routes by prefix: %w", err)
 	}
 
-	fo := failoverRoute(isConnected, r, routes)
+	fo := failoverRoute(isLikelyConnected, r, routes)
 	if fo == nil {
 		return nil, nil
 	}
@@ -538,7 +539,7 @@ func (f *failover) save(tx *gorm.DB) error {
 }
 
 func failoverRoute(
-	isConnected types.NodeConnectedMap,
+	isLikelyConnected *xsync.MapOf[types.NodeID, bool],
 	routeToReplace *types.Route,
 	altRoutes types.Routes,
 
@@ -570,9 +571,11 @@ func failoverRoute(
 			continue
 		}
 
-		if isConnected != nil && isConnected[route.Node.ID] {
-			newPrimary = &altRoutes[idx]
-			break
+		if isLikelyConnected != nil {
+			if val, ok := isLikelyConnected.Load(route.Node.ID); ok && val {
+				newPrimary = &altRoutes[idx]
+				break
+			}
 		}
 	}