Merge remote-tracking branch 'origin/main' into dmitri-catch-disconnected-peer

management/internals/shared/db/conn.go

@@ -11,6 +11,8 @@ import (
 	"github.com/jackc/pgx/v5/pgxpool"
 	log "github.com/sirupsen/logrus"
 	"gorm.io/gorm"
+
+	"github.com/netbirdio/netbird/util"
 )
 
 const (
@@ -64,13 +66,11 @@ func NewConn(ctx context.Context, gormDB *gorm.DB, engine Engine, pool *pgxpool.
 		opt(&options)
 	}
 
-	txTimeout := defaultTransactionTimeout
-	if v := os.Getenv("NB_STORE_TRANSACTION_TIMEOUT"); v != "" {
-		if parsed, err := time.ParseDuration(v); err == nil {
-			txTimeout = parsed
-		}
-	}
+	txTimeout := util.DurationFromEnv(transactionTimeoutEnv, defaultTransactionTimeout)
 	log.WithContext(ctx).Infof("Setting transaction timeout to %v", txTimeout)
+	if d := util.DurationFromEnv(statementTimeoutEnv, 0); engine == PostgresStoreEngine && d > 0 {
+		log.WithContext(ctx).Infof("Setting statement timeout to %v", d)
+	}
 
 	conns := runtime.NumCPU()
 	configuredConns, err := strconv.Atoi(os.Getenv("NB_SQL_MAX_OPEN_CONNS"))

management/internals/shared/db/open.go

@@ -10,7 +10,9 @@ import (
 	"strings"
 	"time"
 
+	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/jackc/pgx/v5/stdlib"
 	"gorm.io/driver/postgres"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
@@ -100,9 +102,10 @@ func OpenSqliteFile(ctx context.Context, dataDir, storeFile string) (*Conn, erro
 	return conn, nil
 }
 
-// OpenPostgres opens a Postgres database through gorm and a pgx pool sized by pool.
+// OpenPostgres opens a Postgres database through gorm and a pgx pool sized by
+// pool, applying the configured timeouts to every session.
 func OpenPostgres(ctx context.Context, dsn string, pool PoolConfig, opts ...Option) (*Conn, error) {
-	gormDB, err := gorm.Open(postgres.Open(dsn), GormConfig())
+	gormDB, err := openGormPostgres(ctx, dsn)
 	if err != nil {
 		return nil, err
 	}
@@ -123,7 +126,7 @@ func OpenPostgres(ctx context.Context, dsn string, pool PoolConfig, opts ...Opti
 // OpenPostgresWithoutPool opens a Postgres database through gorm alone, for
 // repositories that never take the pgx read paths.
 func OpenPostgresWithoutPool(ctx context.Context, dsn string, opts ...Option) (*Conn, error) {
-	gormDB, err := gorm.Open(postgres.Open(dsn), GormConfig())
+	gormDB, err := openGormPostgres(ctx, dsn)
 	if err != nil {
 		return nil, err
 	}
@@ -135,13 +138,29 @@ func OpenPostgresWithoutPool(ctx context.Context, dsn string, opts ...Option) (*
 	return conn, nil
 }
 
-// OpenPgxPool opens a pgx pool to dsn sized by cfg and verifies it with a ping.
+func openGormPostgres(ctx context.Context, dsn string) (*gorm.DB, error) {
+	config, err := pgx.ParseConfig(dsn)
+	if err != nil {
+		return nil, fmt.Errorf("parse database config: %w", err)
+	}
+	sqlDB := stdlib.OpenDB(*config, stdlib.OptionAfterConnect(sessionTimeouts()))
+	gormDB, err := gorm.Open(postgres.New(postgres.Config{Conn: sqlDB}), GormConfig())
+	if err != nil {
+		_ = sqlDB.Close()
+		return nil, err
+	}
+	return gormDB, nil
+}
+
+// OpenPgxPool opens a pgx pool to dsn sized by cfg, applying the configured
+// timeouts to every session, and verifies it with a ping.
 func OpenPgxPool(ctx context.Context, dsn string, cfg PoolConfig) (*pgxpool.Pool, error) {
 	config, err := pgxpool.ParseConfig(dsn)
 	if err != nil {
 		return nil, fmt.Errorf("unable to parse database config: %w", err)
 	}
 
+	config.AfterConnect = sessionTimeouts()
 	config.MaxConns = cfg.MaxConns
 	config.MinConns = cfg.MinConns
 	config.MaxConnLifetime = cfg.MaxConnLifetime

management/internals/shared/db/timeout.go

@@ -0,0 +1,52 @@
+package db
+
+import (
+	"context"
+	"fmt"
+	"math"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+
+	"github.com/netbirdio/netbird/util"
+)
+
+const (
+	transactionTimeoutEnv = "NB_STORE_TRANSACTION_TIMEOUT"
+	statementTimeoutEnv   = "NB_STORE_STATEMENT_TIMEOUT"
+)
+
+// sessionTimeouts sets the configured timeouts on each new Postgres session with
+// SET, since poolers such as PgBouncer reject them as startup parameters.
+func sessionTimeouts() func(context.Context, *pgx.Conn) error {
+	var settings []string
+	if d := util.DurationFromEnv(statementTimeoutEnv, 0); d > 0 {
+		settings = append(settings, "SET statement_timeout = "+postgresMillis(d))
+	}
+	if d := util.DurationFromEnv(transactionTimeoutEnv, defaultTransactionTimeout); d > 0 {
+		settings = append(settings, "SET idle_in_transaction_session_timeout = "+postgresMillis(d))
+	}
+	query := strings.Join(settings, "; ")
+
+	return func(ctx context.Context, conn *pgx.Conn) error {
+		if query == "" {
+			return nil
+		}
+		if _, err := conn.Exec(ctx, query); err != nil {
+			return fmt.Errorf("set session timeouts: %w", err)
+		}
+		return nil
+	}
+}
+
+// postgresMillis converts d to milliseconds for a Postgres timeout, rounding up
+// so a positive d never becomes 0 (disabled) and capping at the server maximum.
+func postgresMillis(d time.Duration) string {
+	ms := d.Milliseconds()
+	if d%time.Millisecond != 0 {
+		ms++
+	}
+	return strconv.FormatInt(min(ms, math.MaxInt32), 10)
+}

management/internals/shared/db/timeout_test.go

@@ -0,0 +1,146 @@
+package db
+
+import (
+	"context"
+	"errors"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+
+	"github.com/netbirdio/netbird/management/server/testutil"
+)
+
+// queryCanceledCode is the SQLSTATE Postgres reports when statement_timeout fires.
+const queryCanceledCode = "57014"
+
+func TestPostgresMillis(t *testing.T) {
+	cases := map[time.Duration]string{
+		time.Minute:             "60000",
+		1500 * time.Microsecond: "2",
+		500 * time.Microsecond:  "1",
+		720 * time.Hour:         "2147483647",
+	}
+	for d, want := range cases {
+		assert.Equal(t, want, postgresMillis(d), "milliseconds for %v", d)
+	}
+}
+
+func TestSessionTimeouts_NothingToApply(t *testing.T) {
+	t.Setenv(statementTimeoutEnv, "0")
+	t.Setenv(transactionTimeoutEnv, "0")
+	ctx := context.Background()
+
+	assert.NoError(t, sessionTimeouts()(ctx, nil), "the hook should not touch the session when both timeouts are disabled")
+}
+
+func TestOpenPostgres_BoundsEverySession(t *testing.T) {
+	t.Setenv(statementTimeoutEnv, "500ms")
+	t.Setenv(transactionTimeoutEnv, "3m")
+	conn := openPostgresTestConn(t)
+	ctx := context.Background()
+
+	requireStatementTimeout(t, conn.DB(nil).Exec("SELECT pg_sleep(1)").Error)
+
+	_, err := conn.Pool(nil).Exec(ctx, "SELECT pg_sleep(1)")
+	requireStatementTimeout(t, err)
+
+	var inTx string
+	require.NoError(t, conn.RunInTx(ctx, func(tx *Tx) error {
+		return conn.DB(tx).Raw("SHOW statement_timeout").Scan(&inTx).Error
+	}))
+	assert.Equal(t, "1min", inTx, "transactions should keep their own 1 minute bound")
+
+	assert.Equal(t, "3min", showSetting(t, conn.DB(nil), "idle_in_transaction_session_timeout"),
+		"the idle transaction timeout should follow the transaction timeout")
+}
+
+func TestOpenPostgres_NoStatementTimeoutKeepsServerSetting(t *testing.T) {
+	for name, value := range map[string]string{"unset": "", "zero": "0"} {
+		t.Run(name, func(t *testing.T) {
+			t.Setenv(statementTimeoutEnv, value)
+			if value == "" {
+				require.NoError(t, os.Unsetenv(statementTimeoutEnv))
+			}
+			ctx := context.Background()
+			dsn := postgresTestDSN(t)
+			admin, err := pgx.Connect(ctx, dsn)
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = admin.Close(ctx) })
+			_, err = admin.Exec(ctx, "ALTER ROLE CURRENT_USER SET statement_timeout = '7s'")
+			require.NoError(t, err)
+
+			conn, err := OpenPostgres(ctx, dsn, DefaultPoolConfig)
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = conn.Close() })
+
+			assert.Equal(t, "7s", showSetting(t, conn.DB(nil), "statement_timeout"), "sessions should keep the server setting")
+		})
+	}
+}
+
+func TestOpenPostgres_CapsTimeoutsAtServerMaximum(t *testing.T) {
+	t.Setenv(transactionTimeoutEnv, "720h")
+	conn := openPostgresTestConn(t)
+
+	assert.Equal(t, "2147483647ms", showSetting(t, conn.DB(nil), "idle_in_transaction_session_timeout"),
+		"a timeout above the server maximum should be capped instead of failing the connection")
+}
+
+func TestOpenPostgresWithoutPool_BoundsSessions(t *testing.T) {
+	t.Setenv(statementTimeoutEnv, "500ms")
+	conn, err := OpenPostgresWithoutPool(context.Background(), postgresTestDSN(t))
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+
+	requireStatementTimeout(t, conn.DB(nil).Exec("SELECT pg_sleep(1)").Error)
+}
+
+func TestOpenPgxPool_BoundsSessions(t *testing.T) {
+	t.Setenv(statementTimeoutEnv, "500ms")
+	ctx := context.Background()
+	pool, err := OpenPgxPool(ctx, postgresTestDSN(t), DefaultPoolConfig)
+	require.NoError(t, err)
+	t.Cleanup(pool.Close)
+
+	_, err = pool.Exec(ctx, "SELECT pg_sleep(1)")
+	requireStatementTimeout(t, err)
+}
+
+func postgresTestDSN(t *testing.T) string {
+	t.Helper()
+	if os.Getenv("NETBIRD_STORE_ENGINE") != string(PostgresStoreEngine) {
+		t.Skip("requires NETBIRD_STORE_ENGINE=postgres")
+	}
+	cleanup, dsn, err := testutil.CreatePostgresTestContainer()
+	require.NoError(t, err)
+	t.Cleanup(cleanup)
+	return dsn
+}
+
+func openPostgresTestConn(t *testing.T) *Conn {
+	t.Helper()
+	conn, err := OpenPostgres(context.Background(), postgresTestDSN(t), DefaultPoolConfig)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+	return conn
+}
+
+func requireStatementTimeout(t *testing.T, err error) {
+	t.Helper()
+	var pgErr *pgconn.PgError
+	require.True(t, errors.As(err, &pgErr), "expected a Postgres error, got %v", err)
+	assert.Equal(t, queryCanceledCode, pgErr.Code, "the statement should be cancelled by statement_timeout")
+}
+
+func showSetting(t *testing.T, db *gorm.DB, name string) string {
+	t.Helper()
+	var value string
+	require.NoError(t, db.Raw("SHOW "+name).Scan(&value).Error)
+	return value
+}

management/internals/shared/grpc/server.go

@@ -91,6 +91,12 @@ type Server struct {
 
 	reverseProxyManager rpservice.Manager
 	reverseProxyMu      sync.RWMutex
+
+	// drainMu guards draining and pendingDisconnects. While draining, ending sync streams
+	// queue their disconnect for FlushDisconnects instead of writing it.
+	drainMu            sync.Mutex
+	draining           bool
+	pendingDisconnects []types.PeerSession
 }
 
 // NewServer creates a new Management server
@@ -313,7 +319,14 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
 	metahash := metaHash(peerMeta)
 	s.loginFilter.addLogin(peerKey.String(), metahash)
 
-	peer, netMap, postureChecks, dnsFwdPort, err := s.accountManager.SyncAndMarkPeer(ctx, accountID, peerKey.String(), peerMeta, realIP, syncStart)
+	settings, err := s.settingsManager.GetSettings(ctx, accountID, activity.SystemInitiator)
+	if err != nil {
+		log.WithContext(ctx).Debugf("error while getting settings for peer %s: %v", peerKey.String(), err)
+		s.syncSem.Add(-1)
+		return status.Errorf(codes.Internal, "error handling request")
+	}
+
+	peer, netMap, postureChecks, dnsFwdPort, err := s.accountManager.SyncAndMarkPeer(ctx, accountID, peerKey.String(), peerMeta, realIP, syncStart, settings)
 	if err != nil {
 		log.WithContext(ctx).Debugf("error while syncing peer %s: %v", peerKey.String(), err)
 		s.syncSem.Add(-1)
@@ -321,11 +334,11 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
 	}
 
 	trackChallenge := func() { s.accountManager.TrackCertificateChallenges(ctx, accountID, peer.ID, syncStart) }
-	err = s.sendInitialSync(ctx, peerKey, peer, netMap, postureChecks, srv, dnsFwdPort, trackChallenge)
+	err = s.sendInitialSync(ctx, peerKey, peer, netMap, postureChecks, srv, dnsFwdPort, settings, trackChallenge)
 	if err != nil {
 		log.WithContext(ctx).Debugf("error while sending initial sync for %s: %v", peerKey.String(), err)
 		s.syncSem.Add(-1)
-		s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil)
+		s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil, false)
 		return err
 	}
 
@@ -333,10 +346,17 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
 	if err != nil {
 		log.WithContext(ctx).Debugf("error while notify peer connected for %s: %v", peerKey.String(), err)
 		s.syncSem.Add(-1)
-		s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil)
+		s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil, false)
 		return err
 	}
 
+	if err = s.accountManager.CheckPeerMaySync(ctx, accountID, peerKey.String(), settings); err != nil {
+		log.WithContext(ctx).Debugf("peer %s may no longer sync after connecting: %v", peerKey.String(), err)
+		s.syncSem.Add(-1)
+		s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, updates, false)
+		return mapError(ctx, err)
+	}
+
 	s.secretsManager.SetupRefresh(ctx, accountID, peer.ID)
 
 	unlock()
@@ -455,21 +475,28 @@ func (s *Server) heartbeatPeerSession(ctx context.Context, accountID, peerID str
 	}
 }
 
+// cancelPeerRoutines tears down the stream under the peer lock. While draining, the
+// disconnect is queued before waiting for the lock, so the shutdown flush covers the
+// stream even when the lock is held past the gRPC stop timeout.
 func (s *Server) cancelPeerRoutines(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time, session chan *network_map.UpdateMessage) {
 	uncanceledCTX := context.WithoutCancel(ctx)
+	disconnectQueued := s.queueDisconnect(accountID, peer, streamStartTime)
 	unlock := s.acquirePeerLockByUID(uncanceledCTX, peer.Key)
 	defer unlock()
 
-	s.cancelPeerRoutinesWithoutLock(uncanceledCTX, accountID, peer, streamStartTime, session)
+	s.cancelPeerRoutinesWithoutLock(uncanceledCTX, accountID, peer, streamStartTime, session, disconnectQueued)
 }
 
 // cancelPeerRoutinesWithoutLock tears down the stream of the session identified by streamStartTime
 // and its updates channel. A nil session means the stream failed before it registered a channel;
-// the controller then closes any channel still registered.
-func (s *Server) cancelPeerRoutinesWithoutLock(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time, session chan *network_map.UpdateMessage) {
-	err := s.accountManager.OnPeerDisconnected(ctx, accountID, peer.Key, streamStartTime)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to disconnect peer %s properly: %v", peer.Key, err)
+// the controller then closes any channel still registered. disconnectQueued tells that the
+// caller already queued the disconnect for the shutdown flush.
+func (s *Server) cancelPeerRoutinesWithoutLock(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time, session chan *network_map.UpdateMessage, disconnectQueued bool) {
+	if !disconnectQueued && !s.queueDisconnect(accountID, peer, streamStartTime) {
+		err := s.accountManager.OnPeerDisconnected(ctx, accountID, peer.Key, streamStartTime)
+		if err != nil {
+			log.WithContext(ctx).Errorf("failed to disconnect peer %s properly: %v", peer.Key, err)
+		}
 	}
 	if !s.networkMapController.OnPeerDisconnected(ctx, accountID, peer.ID, session) {
 		log.WithContext(ctx).Debugf("skipped peer routines teardown for %s: a newer session owns the peer", peer.Key)
@@ -481,6 +508,60 @@ func (s *Server) cancelPeerRoutinesWithoutLock(ctx context.Context, accountID st
 	log.WithContext(ctx).Debugf("peer %s has been disconnected", peer.Key)
 }
 
+// queueDisconnect queues the disconnect of the stream for FlushDisconnects while the
+// server drains and reports whether it did. Embedded proxy peers are never queued
+// because their disconnect refreshes the network maps of the peers using them.
+func (s *Server) queueDisconnect(accountID string, peer *nbpeer.Peer, streamStartTime time.Time) bool {
+	if peer.ProxyMeta.Embedded {
+		return false
+	}
+
+	s.drainMu.Lock()
+	defer s.drainMu.Unlock()
+	if !s.draining {
+		return false
+	}
+	s.pendingDisconnects = append(s.pendingDisconnects, types.PeerSession{
+		AccountID:        accountID,
+		PeerID:           peer.ID,
+		SessionStartedAt: streamStartTime.UnixNano(),
+	})
+	return true
+}
+
+// StartDraining makes ending sync streams queue their disconnect for FlushDisconnects
+// instead of writing one each.
+func (s *Server) StartDraining() {
+	s.drainMu.Lock()
+	defer s.drainMu.Unlock()
+	s.draining = true
+}
+
+// FlushDisconnects stops draining and marks the queued sessions disconnected in batched
+// writes. Streams ending afterwards write their own disconnect again. It returns an error
+// when not every queued session could be written.
+func (s *Server) FlushDisconnects(ctx context.Context) error {
+	s.drainMu.Lock()
+	sessions := s.pendingDisconnects
+	s.pendingDisconnects = nil
+	s.draining = false
+	s.drainMu.Unlock()
+
+	if len(sessions) == 0 {
+		return nil
+	}
+	err := s.accountManager.MarkPeersDisconnected(ctx, sessions)
+	if err != nil && ctx.Err() == nil {
+		// Retrying the whole batch is safe: rows already marked no longer match their session token.
+		log.WithContext(ctx).Warnf("retrying to mark %d peers disconnected on shutdown: %v", len(sessions), err)
+		err = s.accountManager.MarkPeersDisconnected(ctx, sessions)
+	}
+	if err != nil {
+		return fmt.Errorf("mark %d peers disconnected on shutdown: %w", len(sessions), err)
+	}
+	return nil
+}
+
 func (s *Server) validateToken(ctx context.Context, peerKey, jwtToken string) (string, error) {
 	if s.authManager == nil {
 		return "", status.Errorf(codes.Internal, "missing auth manager")
@@ -701,6 +782,7 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
 	//nolint
 	ctx = context.WithValue(ctx, nbContext.PeerIDKey, peerKey.String())
 	accountID, err := s.accountManager.GetAccountIDForPeerKey(ctx, peerKey.String())
+	accountKnown := err == nil
 	if err != nil {
 		// this case should not happen and already indicates an issue but we don't want the system to fail due to being unable to log in detail
 		accountID = "UNKNOWN"
@@ -727,6 +809,15 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
 		sshKey = loginReq.GetPeerKeys().GetSshPubKey()
 	}
 
+	var settings *types.Settings
+	if accountKnown {
+		settings, err = s.settingsManager.GetSettings(ctx, accountID, activity.SystemInitiator)
+		if err != nil {
+			log.WithContext(ctx).Warnf("failed getting settings for peer %s: %s", peerKey, err)
+			return nil, status.Errorf(codes.Internal, "failed logging in peer")
+		}
+	}
+
 	peer, network, postureChecks, enableSSH, err := s.accountManager.LoginPeer(ctx, types.PeerLogin{
 		WireGuardPubKey: peerKey.String(),
 		SSHKey:          cleanClientString(string(sshKey)),
@@ -735,6 +826,7 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
 		SetupKey:        loginReq.GetSetupKey(),
 		ConnectionIP:    realIP,
 		ExtraDNSLabels:  loginReq.GetDnsLabels(),
+		Settings:        settings,
 	})
 	if err != nil {
 		if errors.Is(err, internalStatus.ErrNoAuthMethodProvided) {
@@ -745,7 +837,7 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
 		return nil, mapError(ctx, err)
 	}
 
-	loginResp, err := s.prepareLoginResponse(ctx, peer, network, postureChecks, enableSSH)
+	loginResp, err := s.prepareLoginResponse(ctx, peer, network, postureChecks, enableSSH, settings)
 	if err != nil {
 		log.WithContext(ctx).Warnf("failed preparing login response for peer %s: %s", peerKey, err)
 		return nil, status.Errorf(codes.Internal, "failed logging in peer")
@@ -850,7 +942,7 @@ func (s *Server) ExtendAuthSession(ctx context.Context, req *proto.EncryptedMess
 	}, nil
 }
 
-func (s *Server) prepareLoginResponse(ctx context.Context, peer *nbpeer.Peer, network *types.Network, postureChecks []*nmdata.PostureChecks, enableSSH bool) (*proto.LoginResponse, error) {
+func (s *Server) prepareLoginResponse(ctx context.Context, peer *nbpeer.Peer, network *types.Network, postureChecks []*nmdata.PostureChecks, enableSSH bool, settings *types.Settings) (*proto.LoginResponse, error) {
 	var relayToken *Token
 	var err error
 	if s.config.Relay != nil && len(s.config.Relay.Addresses) > 0 {
@@ -860,10 +952,12 @@ func (s *Server) prepareLoginResponse(ctx context.Context, peer *nbpeer.Peer, ne
 		}
 	}
 
-	settings, err := s.settingsManager.GetSettings(ctx, peer.AccountID, activity.SystemInitiator)
-	if err != nil {
-		log.WithContext(ctx).Warnf("failed getting settings for peer %s: %s", peer.Key, err)
-		return nil, status.Errorf(codes.Internal, "failed getting settings")
+	if settings == nil {
+		settings, err = s.settingsManager.GetSettings(ctx, peer.AccountID, activity.SystemInitiator)
+		if err != nil {
+			log.WithContext(ctx).Warnf("failed getting settings for peer %s: %s", peer.Key, err)
+			return nil, status.Errorf(codes.Internal, "failed getting settings")
+		}
 	}
 
 	// if peer has reached this point then it has logged in
@@ -938,7 +1032,7 @@ func (s *Server) IsHealthy(ctx context.Context, req *proto.Empty) (*proto.Empty,
 }
 
 // sendInitialSync sends initial proto.SyncResponse to the peer requesting synchronization
-func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer *nbpeer.Peer, networkMap *types.NetworkMap, postureChecks []*nmdata.PostureChecks, srv proto.ManagementService_SyncServer, dnsFwdPort int64, onChallengeStamped func()) error {
+func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer *nbpeer.Peer, networkMap *types.NetworkMap, postureChecks []*nmdata.PostureChecks, srv proto.ManagementService_SyncServer, dnsFwdPort int64, settings *types.Settings, onChallengeStamped func()) error {
 	var err error
 	var turnToken *Token
 
@@ -957,11 +1051,6 @@ func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer
 		}
 	}
 
-	settings, err := s.settingsManager.GetSettings(ctx, peer.AccountID, activity.SystemInitiator)
-	if err != nil {
-		return status.Errorf(codes.Internal, "error handling request")
-	}
-
 	peerGroups, err := s.accountManager.GetStore().GetPeerGroupIDs(ctx, store.LockingStrengthNone, peer.AccountID, peer.ID)
 	if err != nil {
 		return status.Errorf(codes.Internal, "failed to get peer groups %s", err)

management/internals/shared/grpc/server_disconnect_test.go

@@ -2,14 +2,18 @@ package grpc
 
 import (
 	"context"
+	"errors"
+	"sync"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
 	"github.com/netbirdio/netbird/management/server/account"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/types"
 )
 
 func TestCancelPeerRoutines_SessionOwnership(t *testing.T) {
@@ -49,3 +53,144 @@ func TestCancelPeerRoutines_SessionOwnership(t *testing.T) {
 		})
 	}
 }
+
+func TestCancelPeerRoutines_DrainingQueuesDisconnect(t *testing.T) {
+	streamStart := time.Unix(1700000000, 0)
+
+	tests := []struct {
+		name         string
+		embedded     bool
+		draining     bool
+		flushBefore  bool
+		wantQueued   bool
+		wantOwnWrite bool
+	}{
+		{name: "regular peer while draining is queued", draining: true, wantQueued: true},
+		{name: "embedded proxy peer while draining writes its own disconnect", embedded: true, draining: true, wantOwnWrite: true},
+		{name: "regular peer without draining writes its own disconnect", wantOwnWrite: true},
+		{name: "regular peer after the flush writes its own disconnect", draining: true, flushBefore: true, wantOwnWrite: true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			accountManager := account.NewMockManager(ctrl)
+			controller := network_map.NewMockController(ctrl)
+			secretsManager := NewMockSecretsManager(ctrl)
+			s := &Server{accountManager: accountManager, networkMapController: controller, secretsManager: secretsManager}
+
+			peer := &nbpeer.Peer{ID: "peer-1", Key: "peer-key"}
+			peer.ProxyMeta.Embedded = tt.embedded
+
+			if tt.draining {
+				s.StartDraining()
+			}
+			if tt.flushBefore {
+				require.NoError(t, s.FlushDisconnects(context.Background()))
+			}
+
+			if tt.wantOwnWrite {
+				accountManager.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.Key, streamStart).Return(nil)
+			}
+			controller.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.ID, gomock.Any()).Return(true)
+			secretsManager.EXPECT().CancelRefresh(peer.ID)
+			accountManager.EXPECT().UntrackCertificateChallenges("account-1", peer.ID, streamStart)
+
+			s.cancelPeerRoutines(context.Background(), "account-1", peer, streamStart, nil)
+
+			if tt.wantQueued {
+				accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), []types.PeerSession{
+					{AccountID: "account-1", PeerID: peer.ID, SessionStartedAt: streamStart.UnixNano()},
+				}).Return(nil)
+			}
+			require.NoError(t, s.FlushDisconnects(context.Background()))
+		})
+	}
+}
+
+func TestFlushDisconnects_FailedBatchStopsDraining(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	accountManager := account.NewMockManager(ctrl)
+	controller := network_map.NewMockController(ctrl)
+	secretsManager := NewMockSecretsManager(ctrl)
+	s := &Server{accountManager: accountManager, networkMapController: controller, secretsManager: secretsManager}
+
+	peer := &nbpeer.Peer{ID: "peer-1", Key: "peer-key"}
+	first := time.Unix(1700000000, 0)
+	second := first.Add(time.Second)
+	controller.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.ID, gomock.Any()).Return(true).Times(2)
+	secretsManager.EXPECT().CancelRefresh(peer.ID).Times(2)
+	accountManager.EXPECT().UntrackCertificateChallenges("account-1", peer.ID, gomock.Any()).Times(2)
+
+	s.StartDraining()
+	s.cancelPeerRoutines(context.Background(), "account-1", peer, first, nil)
+
+	accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), gomock.Any()).Return(errors.New("db down")).Times(2)
+	require.Error(t, s.FlushDisconnects(context.Background()))
+
+	accountManager.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.Key, second).Return(nil)
+	s.cancelPeerRoutines(context.Background(), "account-1", peer, second, nil)
+}
+
+func TestFlushDisconnects_RetriesFailedBatch(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	accountManager := account.NewMockManager(ctrl)
+	controller := network_map.NewMockController(ctrl)
+	secretsManager := NewMockSecretsManager(ctrl)
+	s := &Server{accountManager: accountManager, networkMapController: controller, secretsManager: secretsManager}
+
+	peer := &nbpeer.Peer{ID: "peer-1", Key: "peer-key"}
+	streamStart := time.Unix(1700000000, 0)
+	controller.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.ID, gomock.Any()).Return(true)
+	secretsManager.EXPECT().CancelRefresh(peer.ID)
+	accountManager.EXPECT().UntrackCertificateChallenges("account-1", peer.ID, streamStart)
+
+	s.StartDraining()
+	s.cancelPeerRoutines(context.Background(), "account-1", peer, streamStart, nil)
+
+	want := []types.PeerSession{{AccountID: "account-1", PeerID: peer.ID, SessionStartedAt: streamStart.UnixNano()}}
+	gomock.InOrder(
+		accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), want).Return(errors.New("deadlock detected")),
+		accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), want).Return(nil),
+	)
+	require.NoError(t, s.FlushDisconnects(context.Background()))
+}
+
+func TestCancelPeerRoutines_QueuesDisconnectBeforeWaitingForPeerLock(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	accountManager := account.NewMockManager(ctrl)
+	controller := network_map.NewMockController(ctrl)
+	secretsManager := NewMockSecretsManager(ctrl)
+	s := &Server{accountManager: accountManager, networkMapController: controller, secretsManager: secretsManager}
+
+	peer := &nbpeer.Peer{ID: "peer-1", Key: "peer-key"}
+	streamStart := time.Unix(1700000000, 0)
+	controller.EXPECT().OnPeerDisconnected(gomock.Any(), "account-1", peer.ID, gomock.Any()).Return(true)
+	secretsManager.EXPECT().CancelRefresh(peer.ID)
+	accountManager.EXPECT().UntrackCertificateChallenges("account-1", peer.ID, streamStart)
+
+	unlock := s.acquirePeerLockByUID(context.Background(), peer.Key)
+	var unlockOnce sync.Once
+	t.Cleanup(func() { unlockOnce.Do(unlock) })
+	s.StartDraining()
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		s.cancelPeerRoutines(context.Background(), "account-1", peer, streamStart, nil)
+	}()
+
+	require.Eventually(t, func() bool {
+		s.drainMu.Lock()
+		defer s.drainMu.Unlock()
+		return len(s.pendingDisconnects) == 1
+	}, time.Second, 5*time.Millisecond, "the disconnect must be queued while the teardown waits for the peer lock")
+
+	accountManager.EXPECT().MarkPeersDisconnected(gomock.Any(), []types.PeerSession{
+		{AccountID: "account-1", PeerID: peer.ID, SessionStartedAt: streamStart.UnixNano()},
+	}).Return(nil)
+	require.NoError(t, s.FlushDisconnects(context.Background()))
+
+	unlockOnce.Do(unlock)
+	<-done
+}