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

management/internals/network_map_db/pgsql/peer.go

@@ -11,7 +11,7 @@ import (
 const (
 	GetPeersQuery = `
 	select id, key, ssh_key, dns_label, extra_dns_labels, user_id, ssh_enabled, login_expiration_enabled, last_login, ip, ipv6,
-	peer_status_requires_approval, peer_status_connected, proxy_meta_embedded, proxy_meta_cluster,
+	peer_status_requires_approval, peer_status_connected, peer_status_last_seen, proxy_meta_embedded, proxy_meta_cluster,
 	meta_wt_version, meta_go_os, meta_os_version, meta_kernel_version, meta_network_addresses, meta_files, meta_certificates, meta_capabilities, meta_flags, meta_sync_message_version,
 	location_country_code, location_city_name, location_connection_ip
 	from peers

management/internals/network_map_db/shared_types.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/miekg/dns"
 	"github.com/netbirdio/netbird/management/server/integrations/integrated_validator"
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/settings"
 	"github.com/netbirdio/netbird/shared/management/networkmap"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
@@ -186,6 +187,7 @@ type Peer struct {
 	IPv6                       []byte         `nmap:"json"`
 	PeerStatusRequiresApproval sql.NullBool   `nmap:"map_to:RequiresApproval"`
 	PeerStatusConnected        sql.NullBool   `nmap:"skip"`
+	PeerStatusLastSeen         sql.NullTime   `nmap:"skip"`
 	ProxyMetaEmbedded          sql.NullBool   `nmap:"skip"`
 	ProxyMetaCluster           sql.NullString `nmap:"skip"`
 	MetaWtVersion              sql.NullString `nmap:"skip"`
@@ -318,8 +320,8 @@ func ConvertToNmdataPeers(peers []Peer) ([]nmdata.Peer, map[string][]*nmdata.Pee
 		}
 		dp.ProxyMeta.Cluster = p.ProxyMetaCluster.String
 		// This is only used to build private service candidates, not connected peers are skipped
-		dp.Connected = p.PeerStatusConnected.Bool
-		if dp.ProxyMeta.Embedded && p.PeerStatusConnected.Bool {
+		dp.Connected = nbpeer.IsConnected(p.PeerStatusConnected.Bool, p.PeerStatusLastSeen.Time)
+		if dp.ProxyMeta.Embedded && dp.Connected {
 			clusterToPeerIdx[p.ProxyMetaCluster.String] = append(clusterToPeerIdx[p.ProxyMetaCluster.String], &dp)
 		}
 		if p.MetaWtVersion.Valid {

management/internals/network_map_db/shared_types_test.go

@@ -1,10 +1,14 @@
 package networkmapdb
 
 import (
+	"database/sql"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 )
 
 func TestRecordTypeAndRdata(t *testing.T) {
@@ -58,3 +62,31 @@ func TestConvertToNmdataPeers_RejectsMalformedCertificates(t *testing.T) {
 
 	assert.Error(t, err, "a stored value that is not a list of chains must not load as no certificates")
 }
+
+// TestConvertToNmdataPeers_StaleEmbeddedProxyIsNotACandidate: a management instance that
+// stops without marking its peers disconnected leaves the connected flag set. An embedded
+// proxy peer in that state must not become a private service DNS target, or clients
+// resolve the service to a tunnel IP that no longer answers.
+func TestConvertToNmdataPeers_StaleEmbeddedProxyIsNotACandidate(t *testing.T) {
+	embedded := func(id string, lastSeen time.Time) Peer {
+		return Peer{
+			ID:                  id,
+			PeerStatusConnected: sql.NullBool{Bool: true, Valid: true},
+			PeerStatusLastSeen:  sql.NullTime{Time: lastSeen, Valid: true},
+			ProxyMetaEmbedded:   sql.NullBool{Bool: true, Valid: true},
+			ProxyMetaCluster:    sql.NullString{String: "cluster", Valid: true},
+		}
+	}
+
+	peers, clusterToPeers, err := ConvertToNmdataPeers([]Peer{
+		embedded("live", time.Now().Add(-time.Minute)),
+		embedded("stale", time.Now().Add(-nbpeer.StaleAfter-time.Minute)),
+	})
+
+	require.NoError(t, err)
+	require.Len(t, peers, 2)
+	assert.True(t, peers[0].Connected)
+	assert.False(t, peers[1].Connected, "a peer not seen for longer than StaleAfter must read as disconnected")
+	require.Len(t, clusterToPeers["cluster"], 1)
+	assert.Equal(t, "live", clusterToPeers["cluster"][0].ID)
+}

management/internals/network_map_db/sqlite/peer.go

@@ -10,7 +10,7 @@ import (
 const (
 	GetPeersQuery = `
 	select id, key, ssh_key, dns_label, extra_dns_labels, user_id, ssh_enabled, login_expiration_enabled, last_login, ip, ipv6,
-	peer_status_requires_approval, peer_status_connected, proxy_meta_embedded, proxy_meta_cluster,
+	peer_status_requires_approval, peer_status_connected, peer_status_last_seen, proxy_meta_embedded, proxy_meta_cluster,
 	meta_wt_version, meta_go_os, meta_os_version, meta_kernel_version, meta_network_addresses, meta_files, meta_certificates, meta_capabilities, meta_flags, meta_sync_message_version,
 	location_country_code, location_city_name, location_connection_ip
 	from peers

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)
 	}
 }
 

management/server/peer/peer.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
 
 	"github.com/netbirdio/netbird/management/server/util"
 	"github.com/netbirdio/netbird/shared/management/http/api"
@@ -101,6 +102,32 @@ type PeerStatus struct { //nolint:revive
 	RequiresApproval bool
 }
 
+const (
+	// HeartbeatInterval is how often the management instance serving a peer's sync
+	// stream refreshes the peer's LastSeen.
+	HeartbeatInterval = time.Minute
+	// StaleAfter is how long a connected peer may go without a LastSeen refresh before
+	// it reads as disconnected. It spans several heartbeats so one missed refresh does
+	// not disconnect a live peer.
+	StaleAfter = 5 * time.Minute
+)
+
+// IsConnected reports whether a peer stored with the given connected flag and LastSeen
+// counts as connected. A management instance that stops without marking its peers
+// disconnected leaves the flag set, so a peer whose LastSeen is older than StaleAfter
+// reads as disconnected.
+func IsConnected(connected bool, lastSeen time.Time) bool {
+	return connected && time.Since(lastSeen) < StaleAfter
+}
+
+// AfterFind applies IsConnected to every peer loaded through gorm.
+func (p *Peer) AfterFind(*gorm.DB) error {
+	if p.Status != nil {
+		p.Status.Connected = IsConnected(p.Status.Connected, p.Status.LastSeen)
+	}
+	return nil
+}
+
 // Location is a geo location information of a Peer based on public connection IP
 type Location struct {
 	ConnectionIP net.IP `gorm:"serializer:json"` // from grpc peer or reverse proxy headers depends on setup

management/server/peer/peer_test.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net/netip"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -164,3 +165,22 @@ func TestPeerCapabilities(t *testing.T) {
 		})
 	}
 }
+
+func TestIsConnected(t *testing.T) {
+	tests := []struct {
+		name      string
+		connected bool
+		lastSeen  time.Time
+		want      bool
+	}{
+		{name: "connected and recently seen", connected: true, lastSeen: time.Now().Add(-time.Minute), want: true},
+		{name: "connected but not seen for longer than StaleAfter", connected: true, lastSeen: time.Now().Add(-StaleAfter - time.Minute)},
+		{name: "connected but never seen", connected: true},
+		{name: "disconnected and recently seen", lastSeen: time.Now()},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, IsConnected(tt.connected, tt.lastSeen))
+		})
+	}
+}

management/server/store/sql_store_activity_test.go

@@ -109,6 +109,88 @@ func TestRefreshPeerLastSeenLeavesSessionStateAlone(t *testing.T) {
 	assert.Equal(t, int64(1234567890), peer.Status.SessionStartedAt, "session token must survive an activity write")
 }
 
+// TestPeerReadsAsDisconnectedWhenStale covers a management instance that stopped
+// without marking its peers disconnected: the stored flag stays set, so reads must
+// fall back to LastSeen, both when loading the peer and when filtering on it.
+func TestPeerReadsAsDisconnectedWhenStale(t *testing.T) {
+	ctx := context.Background()
+
+	tests := []struct {
+		name          string
+		lastSeen      time.Time
+		wantConnected bool
+	}{
+		{name: "seen within StaleAfter", lastSeen: time.Now().UTC().Add(-time.Minute), wantConnected: true},
+		{name: "not seen for longer than StaleAfter", lastSeen: time.Now().UTC().Add(-nbpeer.StaleAfter - time.Minute)},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			store := newActivityTestStore(t)
+			stored := activityPeer(tt.lastSeen)
+			stored.Status.Connected = true
+			require.NoError(t, store.AddPeerToAccount(ctx, stored))
+
+			peer, err := store.GetPeerByID(ctx, LockingStrengthNone, activityAccountID, "activityPeer")
+			require.NoError(t, err)
+			assert.Equal(t, tt.wantConnected, peer.Status.Connected)
+
+			for _, connected := range []bool{true, false} {
+				peers, _, err := store.GetAccountPeersPaginated(ctx, LockingStrengthNone, activityAccountID, PaginationState{}, PeerFilters{Connected: &connected}, PeerSorting{})
+				require.NoError(t, err)
+				if connected == tt.wantConnected {
+					assert.Len(t, peers, 1, "filter connected=%t", connected)
+				} else {
+					assert.Empty(t, peers, "filter connected=%t", connected)
+				}
+			}
+		})
+	}
+}
+
+func TestHeartbeatPeerSession(t *testing.T) {
+	ctx := context.Background()
+	const session = int64(1234567890)
+	stale := time.Now().UTC().Add(-time.Hour)
+
+	tests := []struct {
+		name          string
+		storedSession int64
+		loginExpired  bool
+		session       int64
+		wantUpdated   bool
+	}{
+		{name: "stream that owns the session", storedSession: session, session: session, wantUpdated: true},
+		{name: "login expired while the stream stayed open", storedSession: session, loginExpired: true, session: session},
+		{name: "session replaced by a newer stream", storedSession: session + 1, session: session},
+		{name: "session already ended", session: session},
+		{name: "zero token never matches", session: 0},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			store := newActivityTestStore(t)
+			stored := activityPeer(stale)
+			stored.Status.SessionStartedAt = tt.storedSession
+			stored.Status.LoginExpired = tt.loginExpired
+			require.NoError(t, store.AddPeerToAccount(ctx, stored))
+
+			updated, err := store.HeartbeatPeerSession(ctx, activityAccountID, "activityPeer", tt.session)
+			require.NoError(t, err)
+			assert.Equal(t, tt.wantUpdated, updated)
+
+			peer, err := store.GetPeerByID(ctx, LockingStrengthNone, activityAccountID, "activityPeer")
+			require.NoError(t, err)
+			assert.Equal(t, tt.storedSession, peer.Status.SessionStartedAt, "the session token must never change")
+			if tt.wantUpdated {
+				assert.True(t, peer.Status.Connected, "a heartbeat must undo a connected flag that was saved back as false")
+				assert.WithinDuration(t, time.Now().UTC(), peer.Status.LastSeen.UTC(), time.Minute)
+				return
+			}
+			assert.False(t, peer.Status.Connected)
+			assert.WithinDuration(t, stale, peer.Status.LastSeen.UTC(), time.Second)
+		})
+	}
+}
+
 func activityPeer(lastSeen time.Time) *nbpeer.Peer {
 	return &nbpeer.Peer{
 		ID:        "activityPeer",

management/server/store/sql_store_get_account_test.go

@@ -233,7 +233,7 @@ func TestGetAccount_ComprehensiveFieldValidation(t *testing.T) {
 			SystemManufacturer: "Dell Inc.",
 		},
 		Status: &nbpeer.PeerStatus{
-			LastSeen:         now.Add(-5 * time.Minute),
+			LastSeen:         now.Add(-1 * time.Minute),
 			Connected:        true,
 			LoginExpired:     false,
 			RequiresApproval: false,

management/server/store/sql_store_peer.go

@@ -182,6 +182,30 @@ func (s *SqlStore) RefreshPeerLastSeen(ctx context.Context, accountID, peerID st
 	return result.RowsAffected > 0, nil
 }
 
+// HeartbeatPeerSession stamps LastSeen with the database clock and marks the peer
+// connected, but only while sessionStartedAt still owns the peer and its login has not
+// expired. Setting Connected again undoes a save of a peer that was read as disconnected
+// while its heartbeats were delayed; the session match keeps an ended or replaced
+// session from coming back, and the login check keeps an expired peer disconnected
+// until it logs in again.
+func (s *SqlStore) HeartbeatPeerSession(ctx context.Context, accountID, peerID string, sessionStartedAt int64) (bool, error) {
+	if sessionStartedAt == 0 {
+		return false, nil
+	}
+	result := s.db.WithContext(ctx).
+		Model(&nbpeer.Peer{}).
+		Where(accountAndIDQueryCondition, accountID, peerID).
+		Where("peer_status_session_started_at = ? AND peer_status_login_expired = ?", sessionStartedAt, false).
+		Updates(map[string]any{
+			"peer_status_connected": true,
+			"peer_status_last_seen": gorm.Expr("CURRENT_TIMESTAMP"),
+		})
+	if result.Error != nil {
+		return false, status.Errorf(status.Internal, "heartbeat peer session: %v", result.Error)
+	}
+	return result.RowsAffected > 0, nil
+}
+
 func (s *SqlStore) getPeers(ctx context.Context, accountID string) ([]nbpeer.Peer, error) {
 	const query = `SELECT id, account_id, key, ip, name, dns_label, user_id, ssh_key, ssh_enabled, login_expiration_enabled,
 	inactivity_expiration_enabled, last_login, created_at, ephemeral, extra_dns_labels, allow_extra_dns_labels, meta_hostname,
@@ -252,7 +276,7 @@ func (s *SqlStore) getPeers(ctx context.Context, accountID string) ([]nbpeer.Pee
 				p.Status.SessionStartedAt = peerStatusSessionStartedAt.Int64
 			}
 			if peerStatusConnected.Valid {
-				p.Status.Connected = peerStatusConnected.Bool
+				p.Status.Connected = nbpeer.IsConnected(peerStatusConnected.Bool, p.Status.LastSeen)
 			}
 			if peerStatusLoginExpired.Valid {
 				p.Status.LoginExpired = peerStatusLoginExpired.Bool
@@ -599,7 +623,12 @@ func (f *PeerFilters) AsConditions(db *gorm.DB) *gorm.DB {
 		db = db.Where("user_id=?", f.UserId)
 	}
 	if f.Connected != nil {
-		db = db.Where("peer_status_connected=?", *f.Connected)
+		lastSeenAfter := time.Now().UTC().Add(-nbpeer.StaleAfter)
+		if *f.Connected {
+			db = db.Where("peer_status_connected = ? AND peer_status_last_seen >= ?", true, lastSeenAfter)
+		} else {
+			db = db.Where("(peer_status_connected = ? OR peer_status_last_seen IS NULL OR peer_status_last_seen < ?)", false, lastSeenAfter)
+		}
 	}
 	if f.ApprovalRequired != nil {
 		db = db.Where("peer_status_requires_approval=?", *f.ApprovalRequired)

management/server/store/sql_store_peers_filter_test.go

@@ -3,12 +3,21 @@ package store
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 )
 
 func TestSqlStore_FilterPeers(t *testing.T) {
 	runTestForAllEnginesNoSqliteSeed(t, "../testdata/peers_pagination.sql", func(t *testing.T, engine string, store Store) {
+		for _, peerID := range []string{"peer-id-one", "peer-id-two"} {
+			peer, err := store.GetPeerByID(context.Background(), LockingStrengthNone, "account-1", peerID)
+			require.NoError(t, err)
+			peer.Status.Connected = true
+			peer.Status.LastSeen = time.Now().UTC()
+			require.NoError(t, store.SavePeerStatus(context.Background(), "account-1", peerID, *peer.Status))
+		}
+
 		var tests = []struct {
 			description    string
 			filters        PeerFilters

management/server/store/store.go

@@ -196,6 +196,11 @@ type Store interface {
 	// recorded by the database. Returns true when the update happened,
 	// false when a newer session has taken over.
 	MarkPeerDisconnectedIfSameSession(ctx context.Context, accountID, peerID string, sessionStartedAt int64) (bool, error)
+	// HeartbeatPeerSession stamps LastSeen with the database clock and marks the
+	// peer connected, but only while the stored SessionStartedAt equals
+	// sessionStartedAt and the peer's login has not expired. Returns true when
+	// the update happened.
+	HeartbeatPeerSession(ctx context.Context, accountID, peerID string, sessionStartedAt int64) (bool, error)
 	ApproveAccountPeers(ctx context.Context, accountID string) (int, error)
 	DeletePeer(ctx context.Context, accountID string, peerID string) error
 

management/server/store/store_mock.go

@@ -3185,6 +3185,21 @@ func (mr *MockStoreMockRecorder) HasGatewayEndpointByOtherAccount(ctx, host, acc
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HasGatewayEndpointByOtherAccount", reflect.TypeOf((*MockStore)(nil).HasGatewayEndpointByOtherAccount), ctx, host, accountID)
 }
 
+// HeartbeatPeerSession mocks base method.
+func (m *MockStore) HeartbeatPeerSession(ctx context.Context, accountID, peerID string, sessionStartedAt int64) (bool, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "HeartbeatPeerSession", ctx, accountID, peerID, sessionStartedAt)
+	ret0, _ := ret[0].(bool)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// HeartbeatPeerSession indicates an expected call of HeartbeatPeerSession.
+func (mr *MockStoreMockRecorder) HeartbeatPeerSession(ctx, accountID, peerID, sessionStartedAt any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HeartbeatPeerSession", reflect.TypeOf((*MockStore)(nil).HeartbeatPeerSession), ctx, accountID, peerID, sessionStartedAt)
+}
+
 // IncrementAgentNetworkConsumption mocks base method.
 func (m *MockStore) IncrementAgentNetworkConsumption(ctx context.Context, accountID string, kind types.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time, tokensIn, tokensOut int64, costUSD float64) error {
 	m.ctrl.T.Helper()