Rework map session

This commit restructures the map session in to a struct
holding the state of what is needed during its lifetime.

For streaming sessions, the event loop is structured a
bit differently not hammering the clients with updates
but rather batching them over a short, configurable time
which should significantly improve cpu usage, and potentially
flakyness.

The use of Patch updates has been dialed back a little as
it does not look like its a 100% ready for prime time. Nodes
are now updated with full changes, except for a few things
like online status.

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>

hscontrol/types/common.go

@@ -90,6 +90,25 @@ func (i StringList) Value() (driver.Value, error) {
 
 type StateUpdateType int
 
+func (su StateUpdateType) String() string {
+	switch su {
+	case StateFullUpdate:
+		return "StateFullUpdate"
+	case StatePeerChanged:
+		return "StatePeerChanged"
+	case StatePeerChangedPatch:
+		return "StatePeerChangedPatch"
+	case StatePeerRemoved:
+		return "StatePeerRemoved"
+	case StateSelfUpdate:
+		return "StateSelfUpdate"
+	case StateDERPUpdated:
+		return "StateDERPUpdated"
+	}
+
+	return "unknown state update type"
+}
+
 const (
 	StateFullUpdate StateUpdateType = iota
 	// StatePeerChanged is used for updates that needs
@@ -118,7 +137,7 @@ type StateUpdate struct {
 	// ChangeNodes must be set when Type is StatePeerAdded
 	// and StatePeerChanged and contains the full node
 	// object for added nodes.
-	ChangeNodes Nodes
+	ChangeNodes []NodeID
 
 	// ChangePatches must be set when Type is StatePeerChangedPatch
 	// and contains a populated PeerChange object.
@@ -127,7 +146,7 @@ type StateUpdate struct {
 	// Removed must be set when Type is StatePeerRemoved and
 	// contain a list of the nodes that has been removed from
 	// the network.
-	Removed []tailcfg.NodeID
+	Removed []NodeID
 
 	// DERPMap must be set when Type is StateDERPUpdated and
 	// contain the new DERP Map.
@@ -138,39 +157,6 @@ type StateUpdate struct {
 	Message string
 }
 
-// Valid reports if a StateUpdate is correctly filled and
-// panics if the mandatory fields for a type is not
-// filled.
-// Reports true if valid.
-func (su *StateUpdate) Valid() bool {
-	switch su.Type {
-	case StatePeerChanged:
-		if su.ChangeNodes == nil {
-			panic("Mandatory field ChangeNodes is not set on StatePeerChanged update")
-		}
-	case StatePeerChangedPatch:
-		if su.ChangePatches == nil {
-			panic("Mandatory field ChangePatches is not set on StatePeerChangedPatch update")
-		}
-	case StatePeerRemoved:
-		if su.Removed == nil {
-			panic("Mandatory field Removed is not set on StatePeerRemove update")
-		}
-	case StateSelfUpdate:
-		if su.ChangeNodes == nil || len(su.ChangeNodes) != 1 {
-			panic(
-				"Mandatory field ChangeNodes is not set for StateSelfUpdate or has more than one node",
-			)
-		}
-	case StateDERPUpdated:
-		if su.DERPMap == nil {
-			panic("Mandatory field DERPMap is not set on StateDERPUpdated update")
-		}
-	}
-
-	return true
-}
-
 // Empty reports if there are any updates in the StateUpdate.
 func (su *StateUpdate) Empty() bool {
 	switch su.Type {
@@ -185,12 +171,12 @@ func (su *StateUpdate) Empty() bool {
 	return false
 }
 
-func StateUpdateExpire(nodeID uint64, expiry time.Time) StateUpdate {
+func StateUpdateExpire(nodeID NodeID, expiry time.Time) StateUpdate {
 	return StateUpdate{
 		Type: StatePeerChangedPatch,
 		ChangePatches: []*tailcfg.PeerChange{
 			{
-				NodeID:    tailcfg.NodeID(nodeID),
+				NodeID:    nodeID.NodeID(),
 				KeyExpiry: &expiry,
 			},
 		},

hscontrol/types/config.go

@@ -69,6 +69,8 @@ type Config struct {
 	CLI CLIConfig
 
 	ACL ACLConfig
+
+	Tuning Tuning
 }
 
 type SqliteConfig struct {
@@ -161,6 +163,11 @@ type LogConfig struct {
 	Level  zerolog.Level
 }
 
+type Tuning struct {
+	BatchChangeDelay               time.Duration
+	NodeMapSessionBufferedChanSize int
+}
+
 func LoadConfig(path string, isFile bool) error {
 	if isFile {
 		viper.SetConfigFile(path)
@@ -220,6 +227,9 @@ func LoadConfig(path string, isFile bool) error {
 
 	viper.SetDefault("node_update_check_interval", "10s")
 
+	viper.SetDefault("tuning.batch_change_delay", "800ms")
+	viper.SetDefault("tuning.node_mapsession_buffered_chan_size", 30)
+
 	if IsCLIConfigured() {
 		return nil
 	}
@@ -719,6 +729,12 @@ func GetHeadscaleConfig() (*Config, error) {
 		},
 
 		Log: GetLogConfig(),
+
+		// TODO(kradalby): Document these settings when more stable
+		Tuning: Tuning{
+			BatchChangeDelay:               viper.GetDuration("tuning.batch_change_delay"),
+			NodeMapSessionBufferedChanSize: viper.GetInt("tuning.node_mapsession_buffered_chan_size"),
+		},
 	}, nil
 }
 

hscontrol/types/node.go

@@ -7,11 +7,13 @@ import (
 	"fmt"
 	"net/netip"
 	"sort"
+	"strconv"
 	"strings"
 	"time"
 
 	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
 	"github.com/juanfont/headscale/hscontrol/policy/matcher"
+	"github.com/juanfont/headscale/hscontrol/util"
 	"github.com/rs/zerolog/log"
 	"go4.org/netipx"
 	"google.golang.org/protobuf/types/known/timestamppb"
@@ -27,9 +29,24 @@ var (
 	ErrNodeUserHasNoName    = errors.New("node user has no name")
 )
 
+type NodeID uint64
+type NodeConnectedMap map[NodeID]bool
+
+func (id NodeID) StableID() tailcfg.StableNodeID {
+	return tailcfg.StableNodeID(strconv.FormatUint(uint64(id), util.Base10))
+}
+
+func (id NodeID) NodeID() tailcfg.NodeID {
+	return tailcfg.NodeID(id)
+}
+
+func (id NodeID) Uint64() uint64 {
+	return uint64(id)
+}
+
 // Node is a Headscale client.
 type Node struct {
-	ID uint64 `gorm:"primary_key"`
+	ID NodeID `gorm:"primary_key"`
 
 	// MachineKeyDatabaseField is the string representation of MachineKey
 	// it is _only_ used for reading and writing the key to the
@@ -198,7 +215,7 @@ func (node Node) IsExpired() bool {
 		return false
 	}
 
-	return time.Now().UTC().After(*node.Expiry)
+	return time.Since(*node.Expiry) > 0
 }
 
 // IsEphemeral returns if the node is registered as an Ephemeral node.
@@ -319,7 +336,7 @@ func (node *Node) AfterFind(tx *gorm.DB) error {
 
 func (node *Node) Proto() *v1.Node {
 	nodeProto := &v1.Node{
-		Id:         node.ID,
+		Id:         uint64(node.ID),
 		MachineKey: node.MachineKey.String(),
 
 		NodeKey:     node.NodeKey.String(),
@@ -486,8 +503,8 @@ func (nodes Nodes) String() string {
 	return fmt.Sprintf("[ %s ](%d)", strings.Join(temp, ", "), len(temp))
 }
 
-func (nodes Nodes) IDMap() map[uint64]*Node {
-	ret := map[uint64]*Node{}
+func (nodes Nodes) IDMap() map[NodeID]*Node {
+	ret := map[NodeID]*Node{}
 
 	for _, node := range nodes {
 		ret[node.ID] = node