review comments: skip update on nmap, fix other readers and add timeout to avoid staleness

management/internals/shared/grpc/peer_update_handler.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/netbirdio/netbird/encryption"
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/telemetry"
 	"github.com/netbirdio/netbird/shared/management/certposture"
 	"github.com/netbirdio/netbird/shared/management/proto"
@@ -21,6 +22,7 @@ func PeerUpdateHandlerFactory(
 	secretsManager SecretsManager,
 	challenger *certposture.Challenger,
 	onChallengeStamped func(),
+	heartbeat func(),
 	srv proto.ManagementService_SyncServer,
 	cleanupfunc func()) *PeerUpdateHandler {
 	return &PeerUpdateHandler{
@@ -29,6 +31,7 @@ func PeerUpdateHandlerFactory(
 		secretsManager:     secretsManager,
 		challenger:         challenger,
 		onChallengeStamped: onChallengeStamped,
+		heartbeat:          heartbeat,
 		srv:                srv,
 		encrypter:          encryption.DefaultEncrypter{},
 		debouncer:          NewUpdateDebouncer(1000 * time.Millisecond),
@@ -54,6 +57,7 @@ type PeerUpdateHandler struct {
 	encrypter          encryption.Encrypter
 	debouncer          Debouncer
 	cleanupFunc        func()
+	heartbeat          func()
 }
 
 func (pu *PeerUpdateHandler) WithMetrics(appMetrics telemetry.AppMetrics) *PeerUpdateHandler {
@@ -72,6 +76,13 @@ func (pu *PeerUpdateHandler) HandleUpdates(ctx context.Context) error {
 
 	defer pu.debouncer.Stop()
 
+	var heartbeat <-chan time.Time
+	if pu.heartbeat != nil {
+		ticker := time.NewTicker(nbpeer.HeartbeatInterval)
+		defer ticker.Stop()
+		heartbeat = ticker.C
+	}
+
 	for {
 		select {
 		// condition when there are some updates
@@ -110,6 +121,9 @@ func (pu *PeerUpdateHandler) HandleUpdates(ctx context.Context) error {
 				}
 			}
 
+		case <-heartbeat:
+			pu.heartbeat()
+
 		// condition when client <-> server connection has been terminated
 		case <-pu.srv.Context().Done():
 			// happens when connection drops, e.g. client disconnects

management/internals/shared/grpc/peer_update_handler_test.go

@@ -5,14 +5,17 @@ import (
 	"fmt"
 	"sync"
 	"testing"
+	"testing/synctest"
 	"time"
 
 	pb "github.com/golang/protobuf/proto" //nolint
-	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
-	"github.com/netbirdio/netbird/shared/management/proto"
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/mock/gomock"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+
+	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/shared/management/proto"
 )
 
 func TestSendPeerUpdates_FirstUpdate(t *testing.T) {
@@ -118,6 +121,61 @@ func TestSendPeerUpdates_ServerContextDone(t *testing.T) {
 	wg.Wait()
 }
 
+func TestSendPeerUpdates_HeartbeatOnlyOnInterval(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		secretsManager := NewMockSecretsManager(ctrl)
+		updateDebouncer := NewMockDebouncer(ctrl)
+		syncSender := NewMocksyncSender(ctrl)
+
+		var heartbeats int
+		pu := PeerUpdateHandler{
+			peerKey:        mustGenerateKey(t),
+			updates:        make(chan *network_map.UpdateMessage),
+			secretsManager: secretsManager,
+			encrypter:      testEncrypter{},
+			debouncer:      updateDebouncer,
+			srv:            syncSender,
+			cleanupFunc:    func() {},
+			heartbeat:      func() { heartbeats++ },
+		}
+
+		msg := network_map.UpdateMessage{
+			Update: &proto.SyncResponse{Version: 1},
+		}
+
+		srvCtx, cancel := context.WithCancel(context.Background())
+		srvKey := mustGenerateKey(t)
+		updateDebouncer.EXPECT().ProcessUpdate(gomock.Eq(&msg)).Return(true)
+		updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(make(chan time.Time))
+		syncSender.EXPECT().Context().AnyTimes().Return(srvCtx)
+		secretsManager.EXPECT().GetWGKey().Return(srvKey, nil)
+		syncSender.EXPECT().Send(gomock.Any())
+		updateDebouncer.EXPECT().Stop()
+
+		done := make(chan struct{})
+		go func() {
+			defer close(done)
+			pu.HandleUpdates(context.Background()) //nolint:errcheck
+		}()
+
+		pu.updates <- &msg
+		synctest.Wait()
+		assert.Zero(t, heartbeats, "an update must not write a heartbeat")
+
+		time.Sleep(nbpeer.HeartbeatInterval)
+		synctest.Wait()
+		assert.Equal(t, 1, heartbeats, "one heartbeat per interval")
+
+		time.Sleep(nbpeer.HeartbeatInterval)
+		synctest.Wait()
+		assert.Equal(t, 2, heartbeats, "one heartbeat per interval")
+
+		cancel()
+		<-done
+	})
+}
+
 func mustGenerateKey(t *testing.T) wgtypes.Key {
 	t.Helper()
 	k, err := wgtypes.GenerateKey()

management/internals/shared/grpc/server.go

@@ -57,6 +57,8 @@ const (
 	envConcurrentSyncs = "NB_MAX_CONCURRENT_SYNCS"
 
 	defaultSyncLim = 1000
+
+	peerHeartbeatTimeout = 5 * time.Second
 )
 
 // Server an instance of a Management gRPC API server
@@ -348,7 +350,7 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
 	s.syncSem.Add(-1)
 
 	return PeerUpdateHandlerFactory(peerKey, updates, s.secretsManager, s.challenger,
-		trackChallenge, func() { s.refreshPeerLastSeen(ctx, accountID, peer.ID) },
+		trackChallenge, func() { s.heartbeatPeerSession(ctx, accountID, peer.ID, syncStart) },
 		srv, func() { s.cancelPeerRoutines(ctx, accountID, peer, syncStart, updates) }).
 		WithMetrics(s.appMetrics).HandleUpdates(ctx)
 }
@@ -442,13 +444,14 @@ func (s *Server) sendJob(ctx context.Context, peerKey wgtypes.Key, job *job.Even
 	return nil
 }
 
-// refreshPeerLastSeen records that the peer of a live sync stream is still connected,
-// so the stale peer sweep leaves it alone. A failed refresh keeps the stream open: the
-// sweep only acts after several missed heartbeats.
-func (s *Server) refreshPeerLastSeen(ctx context.Context, accountID, peerID string) {
-	staleBefore := time.Now().UTC().Add(-nbpeer.HeartbeatInterval / 2)
-	if _, err := s.accountManager.GetStore().RefreshPeerLastSeen(ctx, accountID, peerID, staleBefore); err != nil && ctx.Err() == nil {
-		log.WithContext(ctx).Warnf("refresh last seen of peer %s: %v", peerID, err)
+// heartbeatPeerSession records that the peer of a live sync stream is still connected.
+// A failed heartbeat keeps the stream open: the peer only reads as disconnected after
+// nbpeer.StaleAfter without one, and reads as connected again once a heartbeat lands.
+func (s *Server) heartbeatPeerSession(ctx context.Context, accountID, peerID string, streamStartTime time.Time) {
+	ctx, cancel := context.WithTimeout(ctx, peerHeartbeatTimeout)
+	defer cancel()
+	if _, err := s.accountManager.GetStore().HeartbeatPeerSession(ctx, accountID, peerID, streamStartTime.UnixNano()); err != nil && !errors.Is(ctx.Err(), context.Canceled) {
+		log.WithContext(ctx).Warnf("heartbeat of peer %s: %v", peerID, err)
 	}
 }