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

management/internals/server/boot.go

@@ -269,6 +269,9 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
 			grpc.KeepaliveParams(kasp),
 			grpc.ChainUnaryInterceptor(realip.UnaryServerInterceptorOpts(realipOpts...), unaryInterceptor, proxyUnary),
 			grpc.ChainStreamInterceptor(realip.StreamServerInterceptorOpts(realipOpts...), streamInterceptor, proxyStream),
+			// Stop waits for the sync handlers, so their teardown has queued every
+			// disconnect before FlushDisconnects writes them.
+			grpc.WaitForHandlers(true),
 		}
 
 		// Append interceptors contributed by registered gRPC extensions. These
@@ -292,6 +295,21 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
 		}
 
 		gRPCAPIHandler := grpc.NewServer(gRPCOpts...)
+		mgmtProto.RegisterManagementServiceServer(gRPCAPIHandler, s.ManagementGRPCServer())
+
+		mgmtProto.RegisterProxyServiceServer(gRPCAPIHandler, s.ReverseProxyGRPCServer())
+		log.Info("ProxyService registered on gRPC server")
+
+		// Register services contributed by external modules via the extension seam.
+		registerExtensions(gRPCAPIHandler, s.grpcExtensions)
+
+		return gRPCAPIHandler
+	})
+}
+
+// ManagementGRPCServer returns the management gRPC service, created once and served by GRPCServer.
+func (s *BaseServer) ManagementGRPCServer() *nbgrpc.Server {
+	return Create(s, func() *nbgrpc.Server {
 		srv, err := nbgrpc.NewServer(s.Config, s.AccountManager(), s.SettingsManager(), s.JobManager(), s.SecretsManager(), s.Metrics(), s.AuthManager(), s.IntegratedValidator(), s.NetworkMapController(), s.OAuthConfigProvider(), s.SessionStore())
 		if err != nil {
 			log.Fatalf("failed to create management server: %v", err)
@@ -301,15 +319,7 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
 		if serviceMgr != nil {
 			serviceMgr.StartExposeReaper(context.Background())
 		}
-		mgmtProto.RegisterManagementServiceServer(gRPCAPIHandler, srv)
-
-		mgmtProto.RegisterProxyServiceServer(gRPCAPIHandler, s.ReverseProxyGRPCServer())
-		log.Info("ProxyService registered on gRPC server")
-
-		// Register services contributed by external modules via the extension seam.
-		registerExtensions(gRPCAPIHandler, s.grpcExtensions)
-
-		return gRPCAPIHandler
+		return srv
 	})
 }
 

management/internals/server/container.go

@@ -1,6 +1,9 @@
 package server
 
-import "fmt"
+import (
+	"fmt"
+	"reflect"
+)
 
 // Create a dependency and add it to the BaseServer's container. A string key identifier will be based on its type definition.
 func Create[T any](s Server, createFunc func() T) T {
@@ -33,15 +36,41 @@ func InjectNamed[T any](c Server, name string, thing T) {
 }
 
 func maybeCreate[T any](s Server, createFunc func() T) (result T, isNew bool) {
-	key := fmt.Sprintf("%T", (*T)(nil))[1:]
-	return maybeCreateKeyed(s, key, createFunc)
+	return maybeCreateKeyed(s, typeKey(reflect.TypeFor[T]()), createFunc)
 }
 
 func maybeCreateNamed[T any](s Server, name string, createFunc func() T) (result T, isNew bool) {
-	key := fmt.Sprintf("%T:%s", (*T)(nil), name)[1:]
+	key := fmt.Sprintf("%s:%s", typeKey(reflect.TypeFor[T]()), name)
 	return maybeCreateKeyed(s, key, createFunc)
 }
 
+// typeKey names t by the full package path of every named type it is built from, so
+// same-named types of different packages, such as google.golang.org/grpc.Server and a
+// local grpc.Server, never share a container key.
+func typeKey(t reflect.Type) string {
+	if t.Name() != "" {
+		if t.PkgPath() == "" {
+			return t.Name()
+		}
+		return t.PkgPath() + "." + t.Name()
+	}
+
+	switch t.Kind() {
+	case reflect.Pointer:
+		return "*" + typeKey(t.Elem())
+	case reflect.Slice:
+		return "[]" + typeKey(t.Elem())
+	case reflect.Array:
+		return fmt.Sprintf("[%d]%s", t.Len(), typeKey(t.Elem()))
+	case reflect.Map:
+		return "map[" + typeKey(t.Key()) + "]" + typeKey(t.Elem())
+	case reflect.Chan:
+		return t.ChanDir().String() + " " + typeKey(t.Elem())
+	default:
+		return t.String()
+	}
+}
+
 func maybeCreateKeyed[T any](s Server, key string, createFunc func() T) (result T, isNew bool) {
 	if t, ok := s.GetContainer(key); ok {
 		if t == nil {

management/internals/server/server.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"crypto/tls"
+	"errors"
 	"fmt"
 	"net"
 	"net/http"
@@ -40,6 +41,15 @@ const (
 	ContainerKeyBaseServer = "baseServer"
 
 	applicationName = "management"
+
+	// shutdownTimeout stays below the 10s that systemd (TimeoutStopSec) and Docker grant
+	// before killing the process.
+	shutdownTimeout = 9 * time.Second
+	grpcStopTimeout = 3 * time.Second
+	// shutdownCloseReserve is kept from the flush for closing the stores.
+	shutdownCloseReserve = 2 * time.Second
+	// shutdownIdPReserve is kept from the extension shutdown hooks for stopping the embedded IdP.
+	shutdownIdPReserve = 1 * time.Second
 )
 
 type Server interface {
@@ -297,9 +307,13 @@ func (s *BaseServer) setupTLS(ctx context.Context) (bool, error) {
 	}
 }
 
-// Stop attempts a graceful shutdown, waiting up to 5 seconds for active connections to finish
+// Stop attempts a graceful shutdown within shutdownTimeout. Writing the disconnects of the
+// drained streams may use the budget left after stopping the services and gRPC streams,
+// except for shutdownCloseReserve, which stays for closing the stores; the extension hooks
+// leave shutdownIdPReserve for stopping the embedded IdP. It returns an error when the
+// disconnects could not all be written or the shutdown ran out of time.
 func (s *BaseServer) Stop() error {
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
 	defer cancel()
 	defer s.RunStopHandlers()
 	if s.domainCleanupStop != nil {
@@ -318,12 +332,20 @@ func (s *BaseServer) Stop() error {
 	if s.certListener != nil {
 		_ = s.certListener.Close()
 	}
-	s.GRPCServer().Stop()
+	s.ManagementGRPCServer().StartDraining()
+	s.stopGRPCServer()
+	deadline, _ := ctx.Deadline()
+	flushCtx, cancelFlush := context.WithDeadline(ctx, deadline.Add(-shutdownCloseReserve))
+	flushErr := s.ManagementGRPCServer().FlushDisconnects(flushCtx)
+	cancelFlush()
+
 	if s.proxyAuthClose != nil {
 		s.proxyAuthClose()
 		s.proxyAuthClose = nil
 	}
-	runExtensionShutdownHooks(ctx, s.grpcExtensions)
+	hooksCtx, cancelHooks := context.WithDeadline(ctx, deadline.Add(-shutdownIdPReserve))
+	runExtensionShutdownHooks(hooksCtx, s.grpcExtensions)
+	cancelHooks()
 	_ = s.Store().Close(ctx)
 	_ = s.EventStore().Close(ctx)
 	if s.update != nil {
@@ -338,12 +360,29 @@ func (s *BaseServer) Stop() error {
 	case err := <-s.Errors():
 		if err != nil {
 			log.WithContext(ctx).Infof("stopped Management Service: %v", err)
-			return nil
+			return flushErr
 		}
 		log.WithContext(ctx).Infof("stopped Management Service")
-		return nil
+		return flushErr
 	case <-ctx.Done():
-		return ctx.Err()
+		return errors.Join(flushErr, ctx.Err())
+	}
+}
+
+// stopGRPCServer stops the gRPC server, which waits for the stream handlers to finish
+// their teardown, but gives up waiting after grpcStopTimeout so a stuck handler cannot
+// block the shutdown.
+func (s *BaseServer) stopGRPCServer() {
+	stopped := make(chan struct{})
+	go func() {
+		s.GRPCServer().Stop()
+		close(stopped)
+	}()
+
+	select {
+	case <-stopped:
+	case <-time.After(grpcStopTimeout):
+		log.Warnf("gRPC server did not stop within %s, continuing shutdown", grpcStopTimeout)
 	}
 }
 

management/internals/shared/grpc/server.go

@@ -89,6 +89,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
@@ -311,7 +317,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)
@@ -319,11 +332,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
 	}
 
@@ -331,10 +344,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()
@@ -442,21 +462,28 @@ func (s *Server) sendJob(ctx context.Context, peerKey wgtypes.Key, job *job.Even
 	return nil
 }
 
+// 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)
@@ -468,6 +495,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")
@@ -688,6 +769,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"
@@ -714,6 +796,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)),
@@ -722,6 +813,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) {
@@ -732,7 +824,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")
@@ -837,7 +929,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 {
@@ -847,10 +939,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
@@ -925,7 +1019,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
 
@@ -944,11 +1038,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
+}