From 908245010e75756972b5b639cac1573f4fece5eb Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Wed, 27 Sep 2023 15:56:15 -0700
Subject: [PATCH] update lastseen in db and mapper

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 hscontrol/mapper/mapper.go | 25 ++++++++++++++++++-
 hscontrol/poll.go          | 50 ++++++++++++++++++--------------------
 hscontrol/types/common.go  |  1 +
 3 files changed, 49 insertions(+), 27 deletions(-)

diff --git a/hscontrol/mapper/mapper.go b/hscontrol/mapper/mapper.go
index 80985b95f42..72832f543c3 100644
--- a/hscontrol/mapper/mapper.go
+++ b/hscontrol/mapper/mapper.go
@@ -46,6 +46,7 @@ var debugDumpMapResponsePath = envknob.String("HEADSCALE_DEBUG_DUMP_MAPRESPONSE_
 // - Keep information about the previous mapresponse so we can send a diff
 // - Store hashes
 // - Create a "minifier" that removes info not needed for the node
+// - some sort of batching, wait for 5 or 60 seconds before sending
 
 type Mapper struct {
 	privateKey2019 *key.MachinePrivate
@@ -329,11 +330,28 @@ func (m *Mapper) PeerChangedResponse(
 		return nil, err
 	}
 
-	// resp.PeerSeenChange = lastSeen
+	resp.PeerSeenChange = lastSeen
 
 	return m.marshalMapResponse(mapRequest, &resp, node, mapRequest.Compress)
 }
 
+// PeerOnlineChanged is a no response function that internally updates
+// the mappers internal node map with LastSeen and Online information.
+func (m *Mapper) PeerOnlineChanged(
+	mapRequest tailcfg.MapRequest,
+	node *types.Node,
+	changed types.Nodes,
+	pol *policy.ACLPolicy,
+) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	// Update our internal map.
+	for _, node := range changed {
+		m.peers[node.ID].LastSeen = node.LastSeen
+	}
+}
+
 func (m *Mapper) PeerRemovedResponse(
 	mapRequest tailcfg.MapRequest,
 	node *types.Node,
@@ -617,5 +635,10 @@ func appendPeerChanges(
 	// TODO(kradalby): This currently does not take last seen in keepalives into account
 	resp.OnlineChange = peers.OnlineNodeMap()
 
+	log.Trace().
+		Interface("OnlineMap", resp.OnlineChange).
+		Str("hostname", node.Hostname).
+		Msg("node online map")
+
 	return nil
 }
diff --git a/hscontrol/poll.go b/hscontrol/poll.go
index 5d0d40e95b6..4cb39f3e45d 100644
--- a/hscontrol/poll.go
+++ b/hscontrol/poll.go
@@ -270,14 +270,7 @@ func (h *Headscale) handlePoll(
 			// One alternative is to split these different channels into
 			// goroutines, but then you might have a problem without a lock
 			// if a keepalive is written at the same time as an update.
-			go func() {
-				err = h.db.UpdateLastSeen(node)
-				if err != nil {
-					logErr(err, "Cannot update node LastSeen")
-
-					return
-				}
-			}()
+			go h.updateOnline(node)
 
 		case update := <-updateChan:
 			logInfo("Received update")
@@ -299,6 +292,10 @@ func (h *Headscale) handlePoll(
 			case types.StateFullUpdate:
 				logInfo("Sending Full MapResponse")
 				data, err = mapp.FullMapResponse(mapRequest, node, h.ACLPolicy)
+			case types.StatePeerOnlineChanged:
+				logInfo("Sending PeerOnlineChanged MapResponse")
+				mapp.PeerOnlineChanged(mapRequest, node, update.Changed, h.ACLPolicy)
+				continue
 			}
 
 			if err != nil {
@@ -325,16 +322,6 @@ func (h *Headscale) handlePoll(
 				return
 			}
 
-			// See comment in keepAliveTicker
-			go func() {
-				err = h.db.UpdateLastSeen(node)
-				if err != nil {
-					logErr(err, "Cannot update node LastSeen")
-
-					return
-				}
-			}()
-
 			log.Info().
 				Caller().
 				Bool("noise", isNoise).
@@ -348,14 +335,7 @@ func (h *Headscale) handlePoll(
 		case <-ctx.Done():
 			logInfo("The client has closed the connection")
 
-			go func() {
-				err = h.db.UpdateLastSeen(node)
-				if err != nil {
-					logErr(err, "Cannot update node LastSeen")
-
-					return
-				}
-			}()
+			go h.updateOnline(node)
 
 			// The connection has been closed, so we can stop polling.
 			return
@@ -368,6 +348,24 @@ func (h *Headscale) handlePoll(
 	}
 }
 
+func (h *Headscale) updateOnline(node *types.Node) {
+	now := time.Now()
+
+	node.LastSeen = &now
+
+	h.nodeNotifier.NotifyWithIgnore(types.StateUpdate{
+		Type:    types.StatePeerOnlineChanged,
+		Changed: types.Nodes{node},
+	}, node.MachineKey)
+
+	err := h.db.UpdateLastSeen(node)
+	if err != nil {
+		log.Error().Err(err).Msg("Cannot update node LastSeen")
+
+		return
+	}
+}
+
 func closeChanWithLog[C chan []byte | chan struct{} | chan types.StateUpdate](channel C, node, name string) {
 	log.Trace().
 		Str("handler", "PollNetMap").
diff --git a/hscontrol/types/common.go b/hscontrol/types/common.go
index b275fa4ad22..e76817493ae 100644
--- a/hscontrol/types/common.go
+++ b/hscontrol/types/common.go
@@ -114,6 +114,7 @@ const (
 	StatePeerChanged
 	StatePeerRemoved
 	StateDERPUpdated
+	StatePeerOnlineChanged
 )
 
 // StateUpdate is an internal message containing information about