From cac9326d3d39179a93ac32f9f28f5691e653c060 Mon Sep 17 00:00:00 2001 From: Vlad <4941176+crn4@users.noreply.github.com> Date: Wed, 14 Jan 2026 17:09:17 +0100 Subject: [PATCH] [management] fetch all users data from external cache in one request (#5104) --------- Co-authored-by: pascal --- management/server/cache/idp.go | 25 ++++++++++++++++++++++++ management/server/user.go | 35 ++++++++++++++++++++++++++-------- management/server/user_test.go | 8 ++++++-- 3 files changed, 58 insertions(+), 10 deletions(-) diff --git a/management/server/cache/idp.go b/management/server/cache/idp.go index 19dfc0f38..6ec42e217 100644 --- a/management/server/cache/idp.go +++ b/management/server/cache/idp.go @@ -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) diff --git a/management/server/user.go b/management/server/user.go index 4f9007b61..d12dd4f11 100644 --- a/management/server/user.go +++ b/management/server/user.go @@ -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 } diff --git a/management/server/user_test.go b/management/server/user_test.go index 6d356a8b1..2dd1cea2e 100644 --- a/management/server/user_test.go +++ b/management/server/user_test.go @@ -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)