From b64bee35fa4cbe1b3b03620bc5ad84768652f920 Mon Sep 17 00:00:00 2001 From: Pedro Maia Costa <550684+pnmcosta@users.noreply.github.com> Date: Sat, 22 Feb 2025 10:31:39 +0000 Subject: [PATCH] [management] faster server bootstrap (#3365) Faster server bootstrap by counting accounts rather than fetching all from storage in the account manager instantiation. This change moved the deprecated need to ensure accounts have an All group to tests instead. --- management/server/account.go | 81 ++++------------------- management/server/store/sql_store.go | 20 +++++- management/server/store/sql_store_test.go | 43 +----------- management/server/store/store.go | 29 ++++++++ management/server/testdata/storev1.sql | 1 - management/server/types/account.go | 41 ++++++++++++ 6 files changed, 104 insertions(+), 111 deletions(-) diff --git a/management/server/account.go b/management/server/account.go index 76c984286..332d356e2 100644 --- a/management/server/account.go +++ b/management/server/account.go @@ -246,6 +246,11 @@ func BuildManager( integratedPeerValidator integrated_validator.IntegratedValidator, metrics telemetry.AppMetrics, ) (*DefaultAccountManager, error) { + start := time.Now() + defer func() { + log.WithContext(ctx).Debugf("took %v to instantiate account manager", time.Since(start)) + }() + am := &DefaultAccountManager{ Store: store, geo: geo, @@ -263,39 +268,21 @@ func BuildManager( metrics: metrics, requestBuffer: NewAccountRequestBuffer(ctx, store), } - allAccounts := store.GetAllAccounts(ctx) + accountsCounter, err := store.GetAccountsCounter(ctx) + if err != nil { + log.WithContext(ctx).Error(err) + } + // enable single account mode only if configured by user and number of existing accounts is not grater than 1 - am.singleAccountMode = singleAccountModeDomain != "" && len(allAccounts) <= 1 + am.singleAccountMode = singleAccountModeDomain != "" && accountsCounter <= 1 if am.singleAccountMode { if !isDomainValid(singleAccountModeDomain) { return nil, status.Errorf(status.InvalidArgument, "invalid domain \"%s\" provided for a single account mode. Please review your input for --single-account-mode-domain", singleAccountModeDomain) } am.singleAccountModeDomain = singleAccountModeDomain - log.WithContext(ctx).Infof("single account mode enabled, accounts number %d", len(allAccounts)) + log.WithContext(ctx).Infof("single account mode enabled, accounts number %d", accountsCounter) } else { - log.WithContext(ctx).Infof("single account mode disabled, accounts number %d", len(allAccounts)) - } - - // if account doesn't have a default group - // we create 'all' group and add all peers into it - // also we create default rule with source as destination - for _, account := range allAccounts { - shouldSave := false - - _, err := account.GetGroupAll() - if err != nil { - if err := addAllGroup(account); err != nil { - return nil, err - } - shouldSave = true - } - - if shouldSave { - err = store.SaveAccount(ctx, account) - if err != nil { - return nil, err - } - } + log.WithContext(ctx).Infof("single account mode disabled, accounts number %d", accountsCounter) } goCacheClient := gocache.New(CacheExpirationMax, 30*time.Minute) @@ -1619,46 +1606,6 @@ func (am *DefaultAccountManager) GetAccountSettings(ctx context.Context, account return am.Store.GetAccountSettings(ctx, store.LockingStrengthShare, accountID) } -// addAllGroup to account object if it doesn't exist -func addAllGroup(account *types.Account) error { - if len(account.Groups) == 0 { - allGroup := &types.Group{ - ID: xid.New().String(), - Name: "All", - Issued: types.GroupIssuedAPI, - } - for _, peer := range account.Peers { - allGroup.Peers = append(allGroup.Peers, peer.ID) - } - account.Groups = map[string]*types.Group{allGroup.ID: allGroup} - - id := xid.New().String() - - defaultPolicy := &types.Policy{ - ID: id, - Name: types.DefaultRuleName, - Description: types.DefaultRuleDescription, - Enabled: true, - Rules: []*types.PolicyRule{ - { - ID: id, - Name: types.DefaultRuleName, - Description: types.DefaultRuleDescription, - Enabled: true, - Sources: []string{allGroup.ID}, - Destinations: []string{allGroup.ID}, - Bidirectional: true, - Protocol: types.PolicyRuleProtocolALL, - Action: types.PolicyTrafficActionAccept, - }, - }, - } - - account.Policies = []*types.Policy{defaultPolicy} - } - return nil -} - // newAccountWithId creates a new Account with a default SetupKey (doesn't store in a Store) and provided id func newAccountWithId(ctx context.Context, accountID, userID, domain string) *types.Account { log.WithContext(ctx).Debugf("creating new account") @@ -1703,7 +1650,7 @@ func newAccountWithId(ctx context.Context, accountID, userID, domain string) *ty }, } - if err := addAllGroup(acc); err != nil { + if err := acc.AddAllGroup(); err != nil { log.WithContext(ctx).Errorf("error adding all group to account %s: %v", acc.Id, err) } return acc diff --git a/management/server/store/sql_store.go b/management/server/store/sql_store.go index 947694420..efc2539ff 100644 --- a/management/server/store/sql_store.go +++ b/management/server/store/sql_store.go @@ -15,7 +15,6 @@ import ( "sync" "time" - "github.com/netbirdio/netbird/management/server/util" log "github.com/sirupsen/logrus" "gorm.io/driver/mysql" "gorm.io/driver/postgres" @@ -24,6 +23,8 @@ import ( "gorm.io/gorm/clause" "gorm.io/gorm/logger" + "github.com/netbirdio/netbird/management/server/util" + nbdns "github.com/netbirdio/netbird/dns" "github.com/netbirdio/netbird/management/server/account" resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types" @@ -615,6 +616,16 @@ func (s *SqlStore) GetResourceGroups(ctx context.Context, lockStrength LockingSt return groups, nil } +func (s *SqlStore) GetAccountsCounter(ctx context.Context) (int64, error) { + var count int64 + result := s.db.Model(&types.Account{}).Count(&count) + if result.Error != nil { + return 0, fmt.Errorf("failed to get all accounts counter: %w", result.Error) + } + + return count, nil +} + func (s *SqlStore) GetAllAccounts(ctx context.Context) (all []*types.Account) { var accounts []types.Account result := s.db.Find(&accounts) @@ -1035,6 +1046,13 @@ func NewSqliteStoreFromFileStore(ctx context.Context, fileStore *FileStore, data } for _, account := range fileStore.GetAllAccounts(ctx) { + _, err = account.GetGroupAll() + if err != nil { + if err := account.AddAllGroup(); err != nil { + return nil, err + } + } + err := store.SaveAccount(ctx, account) if err != nil { return nil, err diff --git a/management/server/store/sql_store_test.go b/management/server/store/sql_store_test.go index dd240ce6c..5cb092190 100644 --- a/management/server/store/sql_store_test.go +++ b/management/server/store/sql_store_test.go @@ -15,7 +15,6 @@ import ( "time" "github.com/google/uuid" - "github.com/rs/xid" log "github.com/sirupsen/logrus" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -2045,52 +2044,12 @@ func newAccountWithId(ctx context.Context, accountID, userID, domain string) *ty }, } - if err := addAllGroup(acc); err != nil { + if err := acc.AddAllGroup(); err != nil { log.WithContext(ctx).Errorf("error adding all group to account %s: %v", acc.Id, err) } return acc } -// addAllGroup to account object if it doesn't exist -func addAllGroup(account *types.Account) error { - if len(account.Groups) == 0 { - allGroup := &types.Group{ - ID: xid.New().String(), - Name: "All", - Issued: types.GroupIssuedAPI, - } - for _, peer := range account.Peers { - allGroup.Peers = append(allGroup.Peers, peer.ID) - } - account.Groups = map[string]*types.Group{allGroup.ID: allGroup} - - id := xid.New().String() - - defaultPolicy := &types.Policy{ - ID: id, - Name: types.DefaultRuleName, - Description: types.DefaultRuleDescription, - Enabled: true, - Rules: []*types.PolicyRule{ - { - ID: id, - Name: types.DefaultRuleName, - Description: types.DefaultRuleDescription, - Enabled: true, - Sources: []string{allGroup.ID}, - Destinations: []string{allGroup.ID}, - Bidirectional: true, - Protocol: types.PolicyRuleProtocolALL, - Action: types.PolicyTrafficActionAccept, - }, - }, - } - - account.Policies = []*types.Policy{defaultPolicy} - } - return nil -} - func TestSqlStore_GetAccountNetworks(t *testing.T) { store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir()) t.Cleanup(cleanup) diff --git a/management/server/store/store.go b/management/server/store/store.go index e074c4c60..2686c3597 100644 --- a/management/server/store/store.go +++ b/management/server/store/store.go @@ -48,6 +48,7 @@ const ( ) type Store interface { + GetAccountsCounter(ctx context.Context) (int64, error) GetAllAccounts(ctx context.Context) []*types.Account GetAccount(ctx context.Context, accountID string) (*types.Account, error) AccountExists(ctx context.Context, lockStrength LockingStrength, id string) (bool, error) @@ -352,9 +353,37 @@ func NewTestStoreFromSQL(ctx context.Context, filename string, dataDir string) ( return nil, nil, fmt.Errorf("failed to create test store: %v", err) } + err = addAllGroupToAccount(ctx, store) + if err != nil { + return nil, nil, fmt.Errorf("failed to add all group to account: %v", err) + } + return getSqlStoreEngine(ctx, store, kind) } +func addAllGroupToAccount(ctx context.Context, store Store) error { + allAccounts := store.GetAllAccounts(ctx) + for _, account := range allAccounts { + shouldSave := false + + _, err := account.GetGroupAll() + if err != nil { + if err := account.AddAllGroup(); err != nil { + return err + } + shouldSave = true + } + + if shouldSave { + err = store.SaveAccount(ctx, account) + if err != nil { + return err + } + } + } + return nil +} + func getSqlStoreEngine(ctx context.Context, store *SqlStore, kind Engine) (Store, func(), error) { var cleanup func() var err error diff --git a/management/server/testdata/storev1.sql b/management/server/testdata/storev1.sql index cda333d4f..8b09ec2be 100644 --- a/management/server/testdata/storev1.sql +++ b/management/server/testdata/storev1.sql @@ -36,4 +36,3 @@ INSERT INTO peers VALUES('xlx9/9D8+ibnRiIIB8nHGMxGOzxV17r8ShPHgi4aYSM=','auth0|6 INSERT INTO peers VALUES('6kjbmVq1hmucVzvBXo5OucY5OYv+jSsB1jUTLq291Dw=','google-oauth2|103201118415301331038','6kjbmVq1hmucVzvBXo5OucY5OYv+jSsB1jUTLq291Dw=','5AFB60DB-61F2-4251-8E11-494847EE88E9','"100.64.0.2"','braginini','linux','Linux','21.04','x86_64','Ubuntu','','','','',NULL,'','','','{"Cloud":"","Platform":""}',NULL,'braginini','braginini','2021-12-24 16:12:05.994305438+01:00',0,0,0,'','',0,0,NULL,'2024-10-02 17:00:54.228182+02:00',0,'""','','',0); INSERT INTO peers VALUES('Ok+5QMdt/UjoktNOvicGYj+IX2g98p+0N2PJ3vJ45RI=','google-oauth2|103201118415301331038','Ok+5QMdt/UjoktNOvicGYj+IX2g98p+0N2PJ3vJ45RI=','A72E4DC2-00DE-4542-8A24-62945438104E','"100.64.0.1"','braginini','linux','Linux','21.04','x86_64','Ubuntu','','','','',NULL,'','','','{"Cloud":"","Platform":""}',NULL,'braginini','braginini-1','2021-12-24 16:11:27.015739803+01:00',0,0,0,'','',0,0,NULL,'2024-10-02 17:00:54.228182+02:00',1,'""','','',0); INSERT INTO installations VALUES(1,''); - diff --git a/management/server/types/account.go b/management/server/types/account.go index 4c68b9523..c890a7730 100644 --- a/management/server/types/account.go +++ b/management/server/types/account.go @@ -12,6 +12,7 @@ import ( "github.com/hashicorp/go-multierror" "github.com/miekg/dns" + "github.com/rs/xid" log "github.com/sirupsen/logrus" nbdns "github.com/netbirdio/netbird/dns" @@ -1525,3 +1526,43 @@ func getPoliciesSourcePeers(policies []*Policy, groups map[string]*Group) map[st return sourcePeers } + +// AddAllGroup to account object if it doesn't exist +func (a *Account) AddAllGroup() error { + if len(a.Groups) == 0 { + allGroup := &Group{ + ID: xid.New().String(), + Name: "All", + Issued: GroupIssuedAPI, + } + for _, peer := range a.Peers { + allGroup.Peers = append(allGroup.Peers, peer.ID) + } + a.Groups = map[string]*Group{allGroup.ID: allGroup} + + id := xid.New().String() + + defaultPolicy := &Policy{ + ID: id, + Name: DefaultRuleName, + Description: DefaultRuleDescription, + Enabled: true, + Rules: []*PolicyRule{ + { + ID: id, + Name: DefaultRuleName, + Description: DefaultRuleDescription, + Enabled: true, + Sources: []string{allGroup.ID}, + Destinations: []string{allGroup.ID}, + Bidirectional: true, + Protocol: PolicyRuleProtocolALL, + Action: PolicyTrafficActionAccept, + }, + }, + } + + a.Policies = []*Policy{defaultPolicy} + } + return nil +}