[management] batch disconnect and lock free reconnect (#8373)

management/server/store/sql_store_peer.go

@@ -106,7 +106,6 @@ func (s *SqlStore) MarkPeerConnectedIfNewerSession(ctx context.Context, accountI
 			"peer_status_connected":          true,
 			"peer_status_last_seen":          gorm.Expr("CURRENT_TIMESTAMP"),
 			"peer_status_session_started_at": newSessionStartedAt,
-			"peer_status_login_expired":      false,
 		})
 	if result.Error != nil {
 		return false, status.Errorf(status.Internal, "mark peer connected: %v", result.Error)
@@ -142,6 +141,51 @@ func (s *SqlStore) MarkPeerDisconnectedIfSameSession(ctx context.Context, accoun
 	return result.RowsAffected > 0, nil
 }
 
+// markPeersDisconnectedChunkSize bounds the sessions per statement so the
+// bound parameter count stays below the database limits.
+const markPeersDisconnectedChunkSize = 1000
+
+// MarkPeersDisconnectedIfSameSession is the batched form of
+// MarkPeerDisconnectedIfSameSession: one UPDATE per chunk of sessions, each
+// peer guarded by its own (account, peer, session) tuple. A failed chunk does
+// not stop the others, only a done ctx does; the count then covers the chunks
+// that succeeded and the error joins the failures.
+func (s *SqlStore) MarkPeersDisconnectedIfSameSession(ctx context.Context, sessions []types.PeerSession) (int64, error) {
+	var total int64
+	var errs []error
+	for start := 0; start < len(sessions); start += markPeersDisconnectedChunkSize {
+		if err := ctx.Err(); err != nil {
+			errs = append(errs, status.Errorf(status.Internal, "mark peers disconnected: %v", err))
+			break
+		}
+		end := min(start+markPeersDisconnectedChunkSize, len(sessions))
+		tuples := make([][]any, 0, end-start)
+		for _, session := range sessions[start:end] {
+			if session.SessionStartedAt == 0 {
+				continue
+			}
+			tuples = append(tuples, []any{session.AccountID, session.PeerID, session.SessionStartedAt})
+		}
+		if len(tuples) == 0 {
+			continue
+		}
+		result := s.db.WithContext(ctx).
+			Model(&nbpeer.Peer{}).
+			Where("(account_id, id, peer_status_session_started_at) IN ?", tuples).
+			Updates(map[string]any{
+				"peer_status_connected":          false,
+				"peer_status_last_seen":          gorm.Expr("CURRENT_TIMESTAMP"),
+				"peer_status_session_started_at": int64(0),
+			})
+		if result.Error != nil {
+			errs = append(errs, status.Errorf(status.Internal, "mark peers disconnected: %v", result.Error))
+			continue
+		}
+		total += result.RowsAffected
+	}
+	return total, errors.Join(errs...)
+}
+
 // ApproveAccountPeers marks all peers that currently require approval in the given account as approved.
 func (s *SqlStore) ApproveAccountPeers(ctx context.Context, accountID string) (int, error) {
 	result := s.db.Model(&nbpeer.Peer{}).

management/server/store/sql_store_peer_test.go

@@ -3,16 +3,20 @@ package store
 import (
 	"context"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"net"
 	"net/netip"
 	"reflect"
 	"sort"
+	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
 
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/types"
@@ -941,3 +945,146 @@ func TestSqlStore_ApproveAccountPeers(t *testing.T) {
 		})
 	})
 }
+
+func TestSqlStore_MarkPeersDisconnectedIfSameSession(t *testing.T) {
+	ctx := context.Background()
+	store, cleanUp, err := NewTestStoreFromSQL(ctx, "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanUp)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	sqlStore := store.(*SqlStore)
+
+	sessionOf := map[string]int64{"peer-a": 101, "peer-b": 102, "peer-c": 103, "peer-d": 104}
+	for id, session := range sessionOf {
+		require.NoError(t, store.AddPeerToAccount(ctx, &nbpeer.Peer{
+			ID:        id,
+			AccountID: accountID,
+			Key:       "key-" + id,
+			DNSLabel:  id,
+			IP:        netip.AddrFrom4([4]byte{10, 9, 0, byte(session)}),
+			IPv6:      netip.MustParseAddr(fmt.Sprintf("fd00::9:%d", session)),
+		}))
+		require.NoError(t, sqlStore.db.Model(&nbpeer.Peer{}).Where("id = ?", id).Updates(map[string]any{
+			"peer_status_connected":          true,
+			"peer_status_session_started_at": session,
+		}).Error)
+	}
+
+	sessions := []types.PeerSession{{AccountID: accountID, PeerID: "peer-a", SessionStartedAt: 101}}
+	for i := 0; i < markPeersDisconnectedChunkSize+5; i++ {
+		sessions = append(sessions, types.PeerSession{AccountID: accountID, PeerID: fmt.Sprintf("missing-%d", i), SessionStartedAt: 1})
+	}
+	sessions = append(sessions,
+		types.PeerSession{AccountID: accountID, PeerID: "peer-b", SessionStartedAt: 102},
+		types.PeerSession{AccountID: accountID, PeerID: "peer-c", SessionStartedAt: 103},
+		types.PeerSession{AccountID: accountID, PeerID: "peer-d", SessionStartedAt: 999},
+		types.PeerSession{AccountID: accountID, PeerID: "peer-d", SessionStartedAt: 0},
+	)
+
+	marked, err := store.MarkPeersDisconnectedIfSameSession(ctx, sessions)
+	require.NoError(t, err)
+	assert.Equal(t, int64(3), marked)
+
+	for id, wantConnected := range map[string]bool{"peer-a": false, "peer-b": false, "peer-c": false, "peer-d": true} {
+		peer, err := store.GetPeerByID(ctx, LockingStrengthNone, accountID, id)
+		require.NoError(t, err)
+		assert.Equal(t, wantConnected, peer.Status.Connected, id)
+		wantSession := int64(0)
+		if wantConnected {
+			wantSession = sessionOf[id]
+		}
+		assert.Equal(t, wantSession, peer.Status.SessionStartedAt, id)
+	}
+}
+
+func TestSqlStore_MarkPeersDisconnectedIfSameSessionContinuesAfterFailedChunk(t *testing.T) {
+	ctx := context.Background()
+	store, cleanUp, err := NewTestStoreFromSQL(ctx, "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanUp)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	sqlStore := store.(*SqlStore)
+
+	sessionOf := map[string]int64{"peer-a": 101, "peer-b": 102, "peer-c": 103}
+	for id, session := range sessionOf {
+		require.NoError(t, store.AddPeerToAccount(ctx, &nbpeer.Peer{
+			ID:        id,
+			AccountID: accountID,
+			Key:       "key-" + id,
+			DNSLabel:  id,
+			IP:        netip.AddrFrom4([4]byte{10, 9, 0, byte(session)}),
+			IPv6:      netip.MustParseAddr(fmt.Sprintf("fd00::9:%d", session)),
+		}))
+		require.NoError(t, sqlStore.db.Model(&nbpeer.Peer{}).Where("id = ?", id).Updates(map[string]any{
+			"peer_status_connected":          true,
+			"peer_status_session_started_at": session,
+		}).Error)
+	}
+
+	sessions := []types.PeerSession{{AccountID: accountID, PeerID: "peer-a", SessionStartedAt: 101}}
+	for i := 0; i < markPeersDisconnectedChunkSize-1; i++ {
+		sessions = append(sessions, types.PeerSession{AccountID: accountID, PeerID: fmt.Sprintf("missing-%d", i), SessionStartedAt: 1})
+	}
+	sessions = append(sessions,
+		types.PeerSession{AccountID: accountID, PeerID: "peer-b", SessionStartedAt: 102},
+		types.PeerSession{AccountID: accountID, PeerID: "peer-c", SessionStartedAt: 103},
+	)
+
+	chunkErr := errors.New("deadlock detected")
+	var updates atomic.Int32
+	require.NoError(t, sqlStore.db.Callback().Update().Before("gorm:update").Register("test:fail_first_chunk", func(db *gorm.DB) {
+		if updates.Add(1) == 1 {
+			_ = db.AddError(chunkErr)
+		}
+	}))
+
+	marked, err := store.MarkPeersDisconnectedIfSameSession(ctx, sessions)
+	require.Error(t, err)
+	assert.ErrorContains(t, err, chunkErr.Error())
+	assert.Equal(t, int64(2), marked, "only the second chunk is counted")
+	assert.Equal(t, int32(2), updates.Load(), "the second chunk must still run")
+
+	for id, wantConnected := range map[string]bool{"peer-a": true, "peer-b": false, "peer-c": false} {
+		peer, err := store.GetPeerByID(ctx, LockingStrengthNone, accountID, id)
+		require.NoError(t, err)
+		assert.Equal(t, wantConnected, peer.Status.Connected, id)
+	}
+
+	cancelled, cancel := context.WithCancel(ctx)
+	cancel()
+	updates.Store(1)
+	marked, err = store.MarkPeersDisconnectedIfSameSession(cancelled, sessions)
+	assert.Zero(t, marked)
+	assert.ErrorContains(t, err, context.Canceled.Error())
+	assert.Equal(t, int32(1), updates.Load(), "no chunk runs once the context is done")
+	assert.Equal(t, 1, strings.Count(err.Error(), context.Canceled.Error()), "a done context is reported once")
+}
+
+func TestSqlStore_MarkPeerConnectedKeepsLoginExpired(t *testing.T) {
+	ctx := context.Background()
+	store, cleanUp, err := NewTestStoreFromSQL(ctx, "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanUp)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	require.NoError(t, store.AddPeerToAccount(ctx, &nbpeer.Peer{
+		ID:        "expired-peer",
+		AccountID: accountID,
+		Key:       "key-expired-peer",
+		DNSLabel:  "expired-peer",
+		IP:        netip.AddrFrom4([4]byte{10, 9, 1, 1}),
+		IPv6:      netip.MustParseAddr("fd00::9:101"),
+		Status:    &nbpeer.PeerStatus{LoginExpired: true},
+	}))
+
+	updated, err := store.MarkPeerConnectedIfNewerSession(ctx, accountID, "expired-peer", 1)
+	require.NoError(t, err)
+	require.True(t, updated)
+
+	peer, err := store.GetPeerByID(ctx, LockingStrengthNone, accountID, "expired-peer")
+	require.NoError(t, err)
+	assert.True(t, peer.Status.Connected)
+	assert.True(t, peer.Status.LoginExpired, "connecting must not clear a login expiry")
+}

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)
+	// MarkPeersDisconnectedIfSameSession applies the semantics of
+	// MarkPeerDisconnectedIfSameSession to many peers in batched statements.
+	// Sessions with a zero SessionStartedAt are skipped. Returns the number
+	// of peers that were marked disconnected.
+	MarkPeersDisconnectedIfSameSession(ctx context.Context, sessions []types.PeerSession) (int64, error)
 	ApproveAccountPeers(ctx context.Context, accountID string) (int, error)
 	DeletePeer(ctx context.Context, accountID string, peerID string) error
 

management/server/store/store_mock.go

@@ -3419,6 +3419,21 @@ func (mr *MockStoreMockRecorder) MarkPeerDisconnectedIfSameSession(ctx, accountI
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeerDisconnectedIfSameSession", reflect.TypeOf((*MockStore)(nil).MarkPeerDisconnectedIfSameSession), ctx, accountID, peerID, sessionStartedAt)
 }
 
+// MarkPeersDisconnectedIfSameSession mocks base method.
+func (m *MockStore) MarkPeersDisconnectedIfSameSession(ctx context.Context, sessions []types3.PeerSession) (int64, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "MarkPeersDisconnectedIfSameSession", ctx, sessions)
+	ret0, _ := ret[0].(int64)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// MarkPeersDisconnectedIfSameSession indicates an expected call of MarkPeersDisconnectedIfSameSession.
+func (mr *MockStoreMockRecorder) MarkPeersDisconnectedIfSameSession(ctx, sessions any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeersDisconnectedIfSameSession", reflect.TypeOf((*MockStore)(nil).MarkPeersDisconnectedIfSameSession), ctx, sessions)
+}
+
 // MarkPendingJobsAsFailed mocks base method.
 func (m *MockStore) MarkPendingJobsAsFailed(ctx context.Context, accountID, peerID, jobID, reason string) error {
 	m.ctrl.T.Helper()