[management] add peer semaphore

management/server/account.go

@@ -44,6 +44,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/types"
 	"github.com/netbirdio/netbird/management/server/util"
 	"github.com/netbirdio/netbird/route"
+	semaphoregroup "github.com/netbirdio/netbird/util/semaphore-group"
 )
 
 const (
@@ -101,6 +102,10 @@ type DefaultAccountManager struct {
 
 	accountUpdateLocks               sync.Map
 	updateAccountPeersBufferInterval atomic.Int64
+
+	addPeerMap         sync.Map
+	addPeerSemaphore   *semaphoregroup.SemaphoreGroup
+	addPeerSemaphoreMu sync.Mutex
 }
 
 // getJWTGroupsChanges calculates the changes needed to sync a user's JWT groups.
@@ -603,11 +608,15 @@ func (am *DefaultAccountManager) DeleteAccount(ctx context.Context, accountID, u
 	}
 
 	for _, otherUser := range account.Users {
-		if otherUser.IsServiceUser {
+		if otherUser.Id == userID {
 			continue
 		}
 
-		if otherUser.Id == userID {
+		if otherUser.IsServiceUser {
+			err = am.deleteServiceUser(ctx, accountID, userID, otherUser)
+			if err != nil {
+				return err
+			}
 			continue
 		}
 

management/server/account_test.go

@@ -853,6 +853,42 @@ func TestAccountManager_DeleteAccount(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	account.Users["service-user-1"] = &types.User{
+		Id:            "service-user-1",
+		Role:          types.UserRoleAdmin,
+		IsServiceUser: true,
+		Issued:        types.UserIssuedAPI,
+		PATs: map[string]*types.PersonalAccessToken{
+			"pat-1": {
+				ID:          "pat-1",
+				UserID:      "service-user-1",
+				Name:        "service-user-1",
+				HashedToken: "hashedToken",
+				CreatedAt:   time.Now(),
+			},
+		},
+	}
+	account.Users[userId] = &types.User{
+		Id:            "service-user-2",
+		Role:          types.UserRoleUser,
+		IsServiceUser: true,
+		Issued:        types.UserIssuedAPI,
+		PATs: map[string]*types.PersonalAccessToken{
+			"pat-2": {
+				ID:          "pat-2",
+				UserID:      userId,
+				Name:        userId,
+				HashedToken: "hashedToken",
+				CreatedAt:   time.Now(),
+			},
+		},
+	}
+
+	err = manager.Store.SaveAccount(context.Background(), account)
+	if err != nil {
+		t.Fatal(err)
+	}
+
 	err = manager.DeleteAccount(context.Background(), account.Id, userId)
 	if err != nil {
 		t.Fatal(err)
@@ -862,6 +898,14 @@ func TestAccountManager_DeleteAccount(t *testing.T) {
 	if err == nil {
 		t.Fatal(fmt.Errorf("expected to get an error when trying to get deleted account, got %v", getAccount))
 	}
+
+	pats, err := manager.Store.GetUserPATs(context.Background(), store.LockingStrengthShare, "service-user-1")
+	require.NoError(t, err)
+	assert.Len(t, pats, 0)
+
+	pats, err = manager.Store.GetUserPATs(context.Background(), store.LockingStrengthShare, userId)
+	require.NoError(t, err)
+	assert.Len(t, pats, 0)
 }
 
 func BenchmarkTest_GetAccountWithclaims(b *testing.B) {

management/server/peer.go

@@ -19,6 +19,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/geolocation"
 	"github.com/netbirdio/netbird/management/server/permissions/modules"
 	"github.com/netbirdio/netbird/management/server/permissions/operations"
+	semaphoregroup "github.com/netbirdio/netbird/util/semaphore-group"
 
 	"github.com/netbirdio/netbird/management/server/idp"
 	"github.com/netbirdio/netbird/management/server/posture"
@@ -762,6 +763,13 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
 	return am.getValidatedPeerWithMap(ctx, peerNotValid, accountID, peer)
 }
 
+type addPeerResult struct {
+	Peer       *nbpeer.Peer
+	NetworkMap *types.NetworkMap
+	Checks     []*posture.Checks
+	Error      error
+}
+
 func (am *DefaultAccountManager) handlePeerLoginNotFound(ctx context.Context, login types.PeerLogin, err error) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, error) {
 	if errStatus, ok := status.FromError(err); ok && errStatus.Type() == status.NotFound {
 		// we couldn't find this peer by its public key which can mean that peer hasn't been registered yet.
@@ -773,7 +781,43 @@ func (am *DefaultAccountManager) handlePeerLoginNotFound(ctx context.Context, lo
 			Location: nbpeer.Location{ConnectionIP: login.ConnectionIP},
 		}
 
-		return am.AddPeer(ctx, login.SetupKey, login.UserID, newPeer)
+		if v, ok := am.addPeerMap.Load(login.WireGuardPubKey); ok {
+			if r, ok := v.(addPeerResult); ok {
+				return r.Peer, r.NetworkMap, r.Checks, r.Error
+			}
+		}
+
+		am.addPeerSemaphoreMu.Lock()
+		if am.addPeerSemaphore == nil {
+			am.addPeerSemaphore = semaphoregroup.NewSemaphoreGroup(10)
+			// ensures cleanup go routine waits for at least one add peer to be done
+			am.addPeerSemaphore.Add(context.Background())
+			go func() {
+				// cleanup routine, keep recreating the semaphore
+				am.addPeerSemaphore.Wait()
+				am.addPeerSemaphore = nil
+			}()
+		} else {
+			am.addPeerSemaphore.Add(context.Background())
+		}
+		am.addPeerSemaphoreMu.Unlock()
+
+		go func() {
+			defer am.addPeerSemaphore.Done(context.Background())
+
+			peer, nMap, checks, err := am.AddPeer(context.Background(), login.SetupKey, login.UserID, newPeer)
+
+			am.addPeerMap.Store(login.WireGuardPubKey, addPeerResult{peer, nMap, checks, err})
+		}()
+
+		// check again!
+		if v, ok := am.addPeerMap.Load(login.WireGuardPubKey); ok {
+			if r, ok := v.(addPeerResult); ok {
+				return r.Peer, r.NetworkMap, r.Checks, r.Error
+			}
+		}
+
+		return nil, nil, nil, status.Errorf(status.Internal, "failed while logging in peer")
 	}
 
 	log.WithContext(ctx).Errorf("failed while logging in peer %s: %v", login.WireGuardPubKey, err)

management/server/store/sql_store.go

@@ -1683,18 +1683,26 @@ func (s *SqlStore) SavePolicy(ctx context.Context, lockStrength LockingStrength,
 }
 
 func (s *SqlStore) DeletePolicy(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) error {
-	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
-		Delete(&types.Policy{}, accountAndIDQueryCondition, accountID, policyID)
-	if err := result.Error; err != nil {
-		log.WithContext(ctx).Errorf("failed to delete policy from store: %s", err)
-		return status.Errorf(status.Internal, "failed to delete policy from store")
-	}
+	return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		if err := tx.Where("policy_id = ?", policyID).Delete(&types.PolicyRule{}).Error; err != nil {
+			return fmt.Errorf("delete policy rules: %w", err)
+		}
 
-	if result.RowsAffected == 0 {
-		return status.NewPolicyNotFoundError(policyID)
-	}
+		result := tx.Clauses(clause.Locking{Strength: string(lockStrength)}).
+			Where(accountAndIDQueryCondition, accountID, policyID).
+			Delete(&types.Policy{})
 
-	return nil
+		if err := result.Error; err != nil {
+			log.WithContext(ctx).Errorf("failed to delete policy from store: %s", err)
+			return status.Errorf(status.Internal, "failed to delete policy from store")
+		}
+
+		if result.RowsAffected == 0 {
+			return status.NewPolicyNotFoundError(policyID)
+		}
+
+		return nil
+	})
 }
 
 // GetAccountPostureChecks retrieves posture checks for an account.