starred/netbird
1
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2026-03-31 06:34:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

[management] fetch all users data from external cache in one request (#5104)

Browse Source 
---------

Co-authored-by: pascal <pascal@netbird.io>
This commit is contained in:
Vlad
2026-01-14 17:09:17 +01:00
committed by GitHub
parent 520d9c66cf
commit cac9326d3d
3 changed files with 58 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
management/server/cache/idp.go vendored
View File

@@ -26,6 +26,8 @@ type UserDataCache interface {
Get(ctx context.Context, key string) (*idp.UserData, error)
Set(ctx context.Context, key string, value *idp.UserData, expiration time.Duration) error
Delete(ctx context.Context, key string) error
GetUsers(ctx context.Context, key string) ([]*idp.UserData, error)
SetUsers(ctx context.Context, key string, users []*idp.UserData, expiration time.Duration) error
}
// UserDataCacheImpl is a struct that implements the UserDataCache interface.
@@ -51,6 +53,29 @@ func (u *UserDataCacheImpl) Delete(ctx context.Context, key string) error {
return u.cache.Delete(ctx, key)
}
func (u *UserDataCacheImpl) GetUsers(ctx context.Context, key string) ([]*idp.UserData, error) {
var users []*idp.UserData
v, err := u.cache.Get(ctx, key, &users)
if err != nil {
return nil, err
}
switch v := v.(type) {
case []*idp.UserData:
return v, nil
case *[]*idp.UserData:
return *v, nil
case []byte:
return unmarshalUserData(v)
}
return nil, fmt.Errorf("unexpected type: %T", v)
}
func (u *UserDataCacheImpl) SetUsers(ctx context.Context, key string, users []*idp.UserData, expiration time.Duration) error {
return u.cache.Set(ctx, key, users, store.WithExpiration(expiration))
}
// NewUserDataCache creates a new UserDataCacheImpl object.
func NewUserDataCache(store store.StoreInterface) *UserDataCacheImpl {
simpleCache := cache.New[any](store)

35
management/server/user.go
View File

@@ -911,10 +911,12 @@ func (am *DefaultAccountManager) GetUsersFromAccount(ctx context.Context, accoun
accountUsers := []*types.User{}
switch {
case allowed:
start := time.Now()
accountUsers, err = am.Store.GetAccountUsers(ctx, store.LockingStrengthNone, accountID)
if err != nil {
return nil, err
}
log.WithContext(ctx).Tracef("Got %d users from account %s after %s", len(accountUsers), accountID, time.Since(start))
case user != nil && user.AccountID == accountID:
accountUsers = append(accountUsers, user)
default:
@@ -933,23 +935,40 @@ func (am *DefaultAccountManager) BuildUserInfosForAccount(ctx context.Context, a
if !isNil(am.idpManager) && !IsEmbeddedIdp(am.idpManager) {
users := make(map[string]userLoggedInOnce, len(accountUsers))
usersFromIntegration := make([]*idp.UserData, 0)
filtered := make(map[string]*idp.UserData, len(accountUsers))
log.WithContext(ctx).Tracef("Querying users from IDP for account %s", accountID)
start := time.Now()
integrationKeys := make(map[string]struct{})
for _, user := range accountUsers {
if user.Issued == types.UserIssuedIntegration {
key := user.IntegrationReference.CacheKey(accountID, user.Id)
info, err := am.externalCacheManager.Get(am.ctx, key)
if err != nil {
log.WithContext(ctx).Infof("Get ExternalCache for key: %s, error: %s", key, err)
users[user.Id] = true
continue
}
usersFromIntegration = append(usersFromIntegration, info)
integrationKeys[user.IntegrationReference.CacheKey(accountID)] = struct{}{}
continue
}
if !user.IsServiceUser {
users[user.Id] = userLoggedInOnce(!user.GetLastLogin().IsZero())
}
}
for key := range integrationKeys {
usersData, err := am.externalCacheManager.GetUsers(am.ctx, key)
if err != nil {
log.WithContext(ctx).Debugf("GetUsers from ExternalCache for key: %s, error: %s", key, err)
continue
}
for _, ud := range usersData {
filtered[ud.ID] = ud
}
}
for _, ud := range filtered {
usersFromIntegration = append(usersFromIntegration, ud)
}
log.WithContext(ctx).Tracef("Got user info from external cache after %s", time.Since(start))
start = time.Now()
queriedUsers, err = am.lookupCache(ctx, users, accountID)
log.WithContext(ctx).Tracef("Got user info from cache for %d users after %s", len(queriedUsers), time.Since(start))
if err != nil {
return nil, err
}

8
management/server/user_test.go
View File

@@ -1086,8 +1086,12 @@ func TestDefaultAccountManager_ExternalCache(t *testing.T) {
assert.NoError(t, err)
cacheManager := am.GetExternalCacheManager()
cacheKey := externalUser.IntegrationReference.CacheKey(mockAccountID, externalUser.Id)
err = cacheManager.Set(context.Background(), cacheKey, &idp.UserData{ID: externalUser.Id, Name: "Test User", Email: "user@example.com"}, time.Minute)
tud := &idp.UserData{ID: externalUser.Id, Name: "Test User", Email: "user@example.com"}
cacheKeyUser := externalUser.IntegrationReference.CacheKey(mockAccountID, externalUser.Id)
err = cacheManager.Set(context.Background(), cacheKeyUser, tud, time.Minute)
assert.NoError(t, err)
cacheKeyAccount := externalUser.IntegrationReference.CacheKey(mockAccountID)
err = cacheManager.SetUsers(context.Background(), cacheKeyAccount, []*idp.UserData{tud}, time.Minute)
assert.NoError(t, err)
infos, err := am.GetUsersFromAccount(context.Background(), mockAccountID, mockUserID)