[management] Add redis cache (#3516)

Refactor IdP user data caching by introducing a Redis cache implementation alongside an in-memory fallback, adding a Marshaler interface for flexible serialization, and updating related tests and account management code.

- Added a new cache store implementation with support for Redis and in-memory backends.
- Introduced Marshaler and wrapper types for handling serialization with msgpack and JSON.
- Updated account and user management modules to integrate and test the new caching strategy.

management/server/account.go

@@ -13,16 +13,16 @@ import (
 	"sync"
 	"time"
 
-	"github.com/eko/gocache/v3/cache"
-	cacheStore "github.com/eko/gocache/v3/store"
-	gocache "github.com/patrickmn/go-cache"
+	cacheStore "github.com/eko/gocache/lib/v4/store"
 	"github.com/rs/xid"
 	log "github.com/sirupsen/logrus"
+	"github.com/vmihailenco/msgpack/v5"
 	"golang.org/x/exp/maps"
 
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/server/account"
 	"github.com/netbirdio/netbird/management/server/activity"
+	nbcache "github.com/netbirdio/netbird/management/server/cache"
 	nbcontext "github.com/netbirdio/netbird/management/server/context"
 	"github.com/netbirdio/netbird/management/server/geolocation"
 	"github.com/netbirdio/netbird/management/server/idp"
@@ -40,8 +40,6 @@ import (
 )
 
 const (
-	CacheExpirationMax         = 7 * 24 * 3600 * time.Second // 7 days
-	CacheExpirationMin         = 3 * 24 * 3600 * time.Second // 3 days
 	peerSchedulerRetryInterval = 3 * time.Second
 	emptyUserID                = "empty user ID in claims"
 	errorGettingDomainAccIDFmt = "error getting account ID by private domain: %v"
@@ -50,7 +48,7 @@ const (
 type userLoggedInOnce bool
 
 func cacheEntryExpiration() time.Duration {
-	r := rand.Intn(int(CacheExpirationMax.Milliseconds()-CacheExpirationMin.Milliseconds())) + int(CacheExpirationMin.Milliseconds())
+	r := rand.Intn(int(nbcache.DefaultIDPCacheExpirationMax.Milliseconds()-nbcache.DefaultIDPCacheExpirationMin.Milliseconds())) + int(nbcache.DefaultIDPCacheExpirationMin.Milliseconds())
 	return time.Duration(r) * time.Millisecond
 }
 
@@ -62,8 +60,8 @@ type DefaultAccountManager struct {
 	cacheLoading         map[string]chan struct{}
 	peersUpdateManager   *PeersUpdateManager
 	idpManager           idp.Manager
-	cacheManager         cache.CacheInterface[[]*idp.UserData]
-	externalCacheManager account.ExternalCacheManager
+	cacheManager         *nbcache.AccountUserDataCache
+	externalCacheManager nbcache.UserDataCache
 	ctx                  context.Context
 	eventStore           activity.Store
 	geo                  geolocation.Geolocation
@@ -200,14 +198,12 @@ func BuildManager(
 		log.WithContext(ctx).Infof("single account mode disabled, accounts number %d", accountsCounter)
 	}
 
-	goCacheClient := gocache.New(CacheExpirationMax, 30*time.Minute)
-	goCacheStore := cacheStore.NewGoCache(goCacheClient)
-	am.cacheManager = cache.NewLoadable[[]*idp.UserData](am.loadAccount, cache.New[[]*idp.UserData](goCacheStore))
-
-	// TODO: what is max expiration time? Should be quite long
-	am.externalCacheManager = cache.New[*idp.UserData](
-		cacheStore.NewGoCache(goCacheClient),
-	)
+	cacheStore, err := nbcache.NewStore(nbcache.DefaultIDPCacheExpirationMax, nbcache.DefaultIDPCacheCleanupInterval)
+	if err != nil {
+		return nil, fmt.Errorf("getting cache store: %s", err)
+	}
+	am.externalCacheManager = nbcache.NewUserDataCache(cacheStore)
+	am.cacheManager = nbcache.NewAccountUserDataCache(am.loadAccount, cacheStore)
 
 	if !isNil(am.idpManager) {
 		go func() {
@@ -489,7 +485,7 @@ func (am *DefaultAccountManager) warmupIDPCache(ctx context.Context) error {
 	rcvdUsers := 0
 	for accountID, users := range userData {
 		rcvdUsers += len(users)
-		err = am.cacheManager.Set(am.ctx, accountID, users, cacheStore.WithExpiration(cacheEntryExpiration()))
+		err = am.cacheManager.Set(am.ctx, accountID, users, cacheEntryExpiration())
 		if err != nil {
 			return err
 		}
@@ -633,18 +629,18 @@ func (am *DefaultAccountManager) addAccountIDToIDPAppMeta(ctx context.Context, u
 	return nil
 }
 
-func (am *DefaultAccountManager) loadAccount(ctx context.Context, accountID interface{}) ([]*idp.UserData, error) {
+func (am *DefaultAccountManager) loadAccount(ctx context.Context, accountID any) (any, []cacheStore.Option, error) {
 	log.WithContext(ctx).Debugf("account %s not found in cache, reloading", accountID)
 	accountIDString := fmt.Sprintf("%v", accountID)
 
 	account, err := am.Store.GetAccount(ctx, accountIDString)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	userData, err := am.idpManager.GetAccount(ctx, accountIDString)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	log.WithContext(ctx).Debugf("%d entries received from IdP management", len(userData))
 
@@ -665,7 +661,13 @@ func (am *DefaultAccountManager) loadAccount(ctx context.Context, accountID inte
 		}
 		matchedUserData = append(matchedUserData, datum)
 	}
-	return matchedUserData, nil
+
+	data, err := msgpack.Marshal(matchedUserData)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return data, []cacheStore.Option{cacheStore.WithExpiration(cacheEntryExpiration())}, nil
 }
 
 func (am *DefaultAccountManager) lookupUserInCacheByEmail(ctx context.Context, email string, accountID string) (*idp.UserData, error) {
@@ -851,7 +853,7 @@ func (am *DefaultAccountManager) removeUserFromCache(ctx context.Context, accoun
 		}
 	}
 
-	return am.cacheManager.Set(am.ctx, accountID, data, cacheStore.WithExpiration(cacheEntryExpiration()))
+	return am.cacheManager.Set(am.ctx, accountID, data, cacheEntryExpiration())
 }
 
 // updateAccountDomainAttributesIfNotUpToDate updates the account domain attributes if they are not up to date and then, saves the account changes

management/server/account/manager.go

@@ -6,11 +6,10 @@ import (
 	"net/netip"
 	"time"
 
-	"github.com/eko/gocache/v3/cache"
-
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/domain"
 	"github.com/netbirdio/netbird/management/server/activity"
+	nbcache "github.com/netbirdio/netbird/management/server/cache"
 	nbcontext "github.com/netbirdio/netbird/management/server/context"
 	"github.com/netbirdio/netbird/management/server/idp"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
@@ -20,7 +19,7 @@ import (
 	"github.com/netbirdio/netbird/route"
 )
 
-type ExternalCacheManager cache.CacheInterface[*idp.UserData]
+type ExternalCacheManager nbcache.UserDataCache
 
 type Manager interface {
 	GetOrCreateAccountByUser(ctx context.Context, userId, domain string) (*types.Account, error)

management/server/cache/idp.go

@@ -0,0 +1,113 @@
+package cache
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/eko/gocache/lib/v4/cache"
+	"github.com/eko/gocache/lib/v4/marshaler"
+	"github.com/eko/gocache/lib/v4/store"
+	"github.com/eko/gocache/store/redis/v4"
+	"github.com/vmihailenco/msgpack/v5"
+
+	"github.com/netbirdio/netbird/management/server/idp"
+)
+
+const (
+	DefaultIDPCacheExpirationMax   = 7 * 24 * time.Hour // 7 days
+	DefaultIDPCacheExpirationMin   = 3 * 24 * time.Hour // 3 days
+	DefaultIDPCacheCleanupInterval = 30 * time.Minute
+)
+
+// UserDataCache is an interface that wraps the basic Get, Set and Delete methods for idp.UserData objects.
+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
+}
+
+// UserDataCacheImpl is a struct that implements the UserDataCache interface.
+type UserDataCacheImpl struct {
+	cache Marshaler
+}
+
+func (u *UserDataCacheImpl) Get(ctx context.Context, key string) (*idp.UserData, error) {
+	v, err := u.cache.Get(ctx, key, new(idp.UserData))
+	if err != nil {
+		return nil, err
+	}
+
+	data := v.(*idp.UserData)
+	return data, nil
+}
+
+func (u *UserDataCacheImpl) Set(ctx context.Context, key string, value *idp.UserData, expiration time.Duration) error {
+	return u.cache.Set(ctx, key, value, store.WithExpiration(expiration))
+}
+
+func (u *UserDataCacheImpl) Delete(ctx context.Context, key string) error {
+	return u.cache.Delete(ctx, key)
+}
+
+// NewUserDataCache creates a new UserDataCacheImpl object.
+func NewUserDataCache(store store.StoreInterface) *UserDataCacheImpl {
+	simpleCache := cache.New[any](store)
+	if store.GetType() == redis.RedisType {
+		m := marshaler.New(simpleCache)
+		return &UserDataCacheImpl{cache: m}
+	}
+	return &UserDataCacheImpl{cache: &marshalerWraper{simpleCache}}
+}
+
+// AccountUserDataCache wraps the basic Get, Set and Delete methods for []*idp.UserData objects.
+type AccountUserDataCache struct {
+	cache Marshaler
+}
+
+func (a *AccountUserDataCache) Get(ctx context.Context, key string) ([]*idp.UserData, error) {
+	var m []*idp.UserData
+	v, err := a.cache.Get(ctx, key, &m)
+	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 unmarshalUserData(data []byte) ([]*idp.UserData, error) {
+	returnObj := &[]*idp.UserData{}
+	err := msgpack.Unmarshal(data, returnObj)
+	if err != nil {
+		return nil, err
+	}
+	return *returnObj, nil
+}
+
+func (a *AccountUserDataCache) Set(ctx context.Context, key string, value []*idp.UserData, expiration time.Duration) error {
+	return a.cache.Set(ctx, key, value, store.WithExpiration(expiration))
+}
+
+func (a *AccountUserDataCache) Delete(ctx context.Context, key string) error {
+	return a.cache.Delete(ctx, key)
+}
+
+// NewAccountUserDataCache creates a new AccountUserDataCache object.
+func NewAccountUserDataCache(loadableFunc cache.LoadFunction[any], store store.StoreInterface) *AccountUserDataCache {
+	simpleCache := cache.New[any](store)
+	loadable := cache.NewLoadable[any](loadableFunc, simpleCache)
+	if store.GetType() == redis.RedisType {
+		m := marshaler.New(loadable)
+		return &AccountUserDataCache{cache: m}
+	}
+	return &AccountUserDataCache{cache: &marshalerWraper{loadable}}
+}

management/server/cache/idp_test.go

@@ -0,0 +1,135 @@
+package cache_test
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/eko/gocache/lib/v4/store"
+	"github.com/redis/go-redis/v9"
+	"github.com/testcontainers/testcontainers-go"
+	testcontainersredis "github.com/testcontainers/testcontainers-go/modules/redis"
+	"github.com/vmihailenco/msgpack/v5"
+
+	"github.com/netbirdio/netbird/management/server/cache"
+	"github.com/netbirdio/netbird/management/server/idp"
+)
+
+func TestNewIDPCacheManagers(t *testing.T) {
+	tt := []struct {
+		name  string
+		redis bool
+	}{
+		{"memory", false},
+		{"redis", true},
+	}
+	for _, tc := range tt {
+		t.Run(tc.name, func(t *testing.T) {
+			if tc.redis {
+				ctx := context.Background()
+				redisContainer, err := testcontainersredis.RunContainer(ctx, testcontainers.WithImage("redis:7"))
+				if err != nil {
+					t.Fatalf("couldn't start redis container: %s", err)
+				}
+				defer func() {
+					if err := redisContainer.Terminate(ctx); err != nil {
+						t.Logf("failed to terminate container: %s", err)
+					}
+				}()
+				redisURL, err := redisContainer.ConnectionString(ctx)
+				if err != nil {
+					t.Fatalf("couldn't get connection string: %s", err)
+				}
+
+				t.Setenv(cache.RedisStoreEnvVar, redisURL)
+			}
+			cacheStore, err := cache.NewStore(cache.DefaultIDPCacheExpirationMax, cache.DefaultIDPCacheCleanupInterval)
+			if err != nil {
+				t.Fatalf("couldn't create cache store: %s", err)
+			}
+
+			simple := cache.NewUserDataCache(cacheStore)
+			loadable := cache.NewAccountUserDataCache(loader, cacheStore)
+
+			ctx := context.Background()
+			value := &idp.UserData{ID: "v", Name: "vv"}
+			err = simple.Set(ctx, "key1", value, time.Minute)
+			if err != nil {
+				t.Errorf("couldn't set testing data: %s", err)
+			}
+
+			result, err := simple.Get(ctx, "key1")
+			if err != nil {
+				t.Errorf("couldn't get testing data: %s", err)
+			}
+			if value.ID != result.ID || value.Name != result.Name {
+				t.Errorf("value returned doesn't match testing data, got %v, expected %v", result, "value1")
+			}
+			values := []*idp.UserData{
+				{ID: "v2", Name: "v2v2"},
+				{ID: "v3", Name: "v3v3"},
+				{ID: "v4", Name: "v4v4"},
+			}
+			err = loadable.Set(ctx, "key2", values, time.Minute)
+
+			if err != nil {
+				t.Errorf("couldn't set testing data: %s", err)
+			}
+			result2, err := loadable.Get(ctx, "key2")
+			if err != nil {
+				t.Errorf("couldn't get testing data: %s", err)
+			}
+
+			if values[0].ID != result2[0].ID || values[0].Name != result2[0].Name {
+				t.Errorf("value returned doesn't match testing data, got %v, expected %v", result2[0], values[0])
+			}
+			if values[1].ID != result2[1].ID || values[1].Name != result2[1].Name {
+				t.Errorf("value returned doesn't match testing data, got %v, expected %v", result2[1], values[1])
+			}
+
+			// checking with direct store client
+			if tc.redis {
+				// wait for redis to sync
+				options, err := redis.ParseURL(os.Getenv(cache.RedisStoreEnvVar))
+				if err != nil {
+					t.Fatalf("parsing redis cache url: %s", err)
+				}
+
+				redisClient := redis.NewClient(options)
+				_, err = redisClient.Get(ctx, "loadKey").Result()
+				if err == nil {
+					t.Errorf("shouldn't find testing data from redis")
+				}
+			}
+
+			// testing loadable capability
+			result2, err = loadable.Get(ctx, "loadKey")
+			if err != nil {
+				t.Errorf("couldn't get testing data: %s", err)
+			}
+
+			if loadData[0].ID != result2[0].ID || loadData[0].Name != result2[0].Name {
+				t.Errorf("value returned doesn't match testing data, got %v, expected %v", result2[0], loadData[0])
+			}
+			if loadData[1].ID != result2[1].ID || loadData[1].Name != result2[1].Name {
+				t.Errorf("value returned doesn't match testing data, got %v, expected %v", result2[1], loadData[1])
+			}
+		})
+	}
+
+}
+
+var loadData = []*idp.UserData{
+	{ID: "a", Name: "aa"},
+	{ID: "b", Name: "bb"},
+	{ID: "c", Name: "cc"},
+}
+
+func loader(ctx context.Context, key any) (any, []store.Option, error) {
+	bytes, err := msgpack.Marshal(loadData)
+	if err != nil {
+		return nil, nil, err
+	}
+	return bytes, nil, nil
+}

management/server/cache/marshaler.go

@@ -0,0 +1,35 @@
+package cache
+
+import (
+	"context"
+
+	"github.com/eko/gocache/lib/v4/store"
+)
+
+type Marshaler interface {
+	Get(ctx context.Context, key any, returnObj any) (any, error)
+	Set(ctx context.Context, key, object any, options ...store.Option) error
+	Delete(ctx context.Context, key any) error
+}
+
+type cacher[T any] interface {
+	Get(ctx context.Context, key any) (T, error)
+	Set(ctx context.Context, key any, object T, options ...store.Option) error
+	Delete(ctx context.Context, key any) error
+}
+
+type marshalerWraper struct {
+	cache cacher[any]
+}
+
+func (m marshalerWraper) Get(ctx context.Context, key any, _ any) (any, error) {
+	return m.cache.Get(ctx, key)
+}
+
+func (m marshalerWraper) Set(ctx context.Context, key, object any, options ...store.Option) error {
+	return m.cache.Set(ctx, key, object, options...)
+}
+
+func (m marshalerWraper) Delete(ctx context.Context, key any) error {
+	return m.cache.Delete(ctx, key)
+}

management/server/cache/store.go

@@ -0,0 +1,50 @@
+package cache
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/eko/gocache/lib/v4/store"
+	gocache_store "github.com/eko/gocache/store/go_cache/v4"
+	redis_store "github.com/eko/gocache/store/redis/v4"
+	gocache "github.com/patrickmn/go-cache"
+	"github.com/redis/go-redis/v9"
+)
+
+// RedisStoreEnvVar is the environment variable that determines if a redis store should be used.
+// The value should follow redis URL format. https://github.com/redis/redis-specifications/blob/master/uri/redis.txt
+const RedisStoreEnvVar = "NB_IDP_CACHE_REDIS_ADDRESS"
+
+// NewStore creates a new cache store with the given max timeout and cleanup interval. It checks for the environment Variable RedisStoreEnvVar
+// to determine if a redis store should be used. If the environment variable is set, it will attempt to connect to the redis store.
+func NewStore(maxTimeout, cleanupInterval time.Duration) (store.StoreInterface, error) {
+	redisAddr := os.Getenv(RedisStoreEnvVar)
+	if redisAddr != "" {
+		return getRedisStore(redisAddr)
+	}
+	goc := gocache.New(maxTimeout, cleanupInterval)
+	return gocache_store.NewGoCache(goc), nil
+}
+
+func getRedisStore(redisEnvAddr string) (store.StoreInterface, error) {
+	options, err := redis.ParseURL(redisEnvAddr)
+	if err != nil {
+		return nil, fmt.Errorf("parsing redis cache url: %s", err)
+	}
+
+	options.MaxIdleConns = 6
+	options.MinIdleConns = 3
+	options.MaxActiveConns = 100
+	redisClient := redis.NewClient(options)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+
+	_, err = redisClient.Ping(ctx).Result()
+	if err != nil {
+		return nil, err
+	}
+
+	return redis_store.NewRedis(redisClient), nil
+}

management/server/cache/store_test.go

@@ -0,0 +1,105 @@
+package cache_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/eko/gocache/lib/v4/store"
+	"github.com/redis/go-redis/v9"
+	"github.com/testcontainers/testcontainers-go"
+
+	testcontainersredis "github.com/testcontainers/testcontainers-go/modules/redis"
+
+	"github.com/netbirdio/netbird/management/server/cache"
+)
+
+func TestMemoryStore(t *testing.T) {
+	memStore, err := cache.NewStore(100*time.Millisecond, 300*time.Millisecond)
+	if err != nil {
+		t.Fatalf("couldn't create memory store: %s", err)
+	}
+	ctx := context.Background()
+	key, value := "testing", "tested"
+	err = memStore.Set(ctx, key, value)
+	if err != nil {
+		t.Errorf("couldn't set testing data: %s", err)
+	}
+	result, err := memStore.Get(ctx, key)
+	if err != nil {
+		t.Errorf("couldn't get testing data: %s", err)
+	}
+	if value != result.(string) {
+		t.Errorf("value returned doesn't match testing data, got %s, expected %s", result, value)
+	}
+	// test expiration
+	time.Sleep(300 * time.Millisecond)
+	_, err = memStore.Get(ctx, key)
+	if err == nil {
+		t.Error("value should not be found")
+	}
+}
+
+func TestRedisStoreConnectionFailure(t *testing.T) {
+	t.Setenv(cache.RedisStoreEnvVar, "redis://127.0.0.1:6379")
+	_, err := cache.NewStore(10*time.Millisecond, 30*time.Millisecond)
+	if err == nil {
+		t.Fatal("getting redis cache store should return error")
+	}
+}
+
+func TestRedisStoreConnectionSuccess(t *testing.T) {
+	ctx := context.Background()
+	redisContainer, err := testcontainersredis.RunContainer(ctx, testcontainers.WithImage("redis:7"))
+	if err != nil {
+		t.Fatalf("couldn't start redis container: %s", err)
+	}
+	defer func() {
+		if err := redisContainer.Terminate(ctx); err != nil {
+			t.Logf("failed to terminate container: %s", err)
+		}
+	}()
+	redisURL, err := redisContainer.ConnectionString(ctx)
+	if err != nil {
+		t.Fatalf("couldn't get connection string: %s", err)
+	}
+
+	t.Setenv(cache.RedisStoreEnvVar, redisURL)
+	redisStore, err := cache.NewStore(100*time.Millisecond, 300*time.Millisecond)
+	if err != nil {
+		t.Fatalf("couldn't create redis store: %s", err)
+	}
+
+	key, value := "testing", "tested"
+	err = redisStore.Set(ctx, key, value, store.WithExpiration(100*time.Millisecond))
+	if err != nil {
+		t.Errorf("couldn't set testing data: %s", err)
+	}
+	result, err := redisStore.Get(ctx, key)
+	if err != nil {
+		t.Errorf("couldn't get testing data: %s", err)
+	}
+	if value != result.(string) {
+		t.Errorf("value returned doesn't match testing data, got %s, expected %s", result, value)
+	}
+
+	options, err := redis.ParseURL(redisURL)
+	if err != nil {
+		t.Errorf("parsing redis cache url: %s", err)
+	}
+
+	redisClient := redis.NewClient(options)
+	r, e := redisClient.Get(ctx, key).Result()
+	if e != nil {
+		t.Errorf("couldn't get testing data from redis: %s", e)
+	}
+	if value != r {
+		t.Errorf("value returned from redis doesn't match testing data, got %s, expected %s", r, value)
+	}
+	// test expiration
+	time.Sleep(300 * time.Millisecond)
+	_, err = redisStore.Get(ctx, key)
+	if err == nil {
+		t.Error("value should not be found")
+	}
+}

management/server/idp/idp.go

@@ -2,6 +2,7 @@ package idp
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"net/http"
 	"strings"
@@ -73,6 +74,23 @@ type UserData struct {
 	AppMetadata AppMetadata `json:"app_metadata"`
 }
 
+func (u *UserData) MarshalBinary() (data []byte, err error) {
+	return json.Marshal(u)
+}
+
+func (u *UserData) UnmarshalBinary(data []byte) (err error) {
+	return json.Unmarshal(data, &u)
+}
+
+func (u *UserData) Marshal() (data string, err error) {
+	d, err := json.Marshal(u)
+	return string(d), err
+}
+
+func (u *UserData) Unmarshal(data []byte) (err error) {
+	return json.Unmarshal(data, &u)
+}
+
 // AppMetadata user app metadata to associate with a profile
 type AppMetadata struct {
 	// WTAccountID is a NetBird (previously Wiretrustee) account id to update in the IDP

management/server/user_test.go

@@ -7,19 +7,18 @@ import (
 	"testing"
 	"time"
 
-	"github.com/eko/gocache/v3/cache"
-	cacheStore "github.com/eko/gocache/v3/store"
 	"github.com/google/go-cmp/cmp"
 
+	"golang.org/x/exp/maps"
+
+	nbcache "github.com/netbirdio/netbird/management/server/cache"
 	nbcontext "github.com/netbirdio/netbird/management/server/context"
 	"github.com/netbirdio/netbird/management/server/util"
-	"golang.org/x/exp/maps"
 
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/types"
 
-	gocache "github.com/patrickmn/go-cache"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
@@ -517,9 +516,10 @@ func TestUser_InviteNewUser(t *testing.T) {
 		cacheLoading: map[string]chan struct{}{},
 	}
 
-	goCacheClient := gocache.New(CacheExpirationMax, 30*time.Minute)
-	goCacheStore := cacheStore.NewGoCache(goCacheClient)
-	am.cacheManager = cache.NewLoadable[[]*idp.UserData](am.loadAccount, cache.New[[]*idp.UserData](goCacheStore))
+	cs, err := nbcache.NewStore(nbcache.DefaultIDPCacheExpirationMax, nbcache.DefaultIDPCacheCleanupInterval)
+	require.NoError(t, err)
+
+	am.cacheManager = nbcache.NewAccountUserDataCache(am.loadAccount, cs)
 
 	mockData := []*idp.UserData{
 		{
@@ -1092,21 +1092,19 @@ func TestDefaultAccountManager_ExternalCache(t *testing.T) {
 		eventStore:   &activity.InMemoryEventStore{},
 		idpManager:   &idp.GoogleWorkspaceManager{}, // empty manager
 		cacheLoading: map[string]chan struct{}{},
-		cacheManager: cache.New[[]*idp.UserData](
-			cacheStore.NewGoCache(gocache.New(CacheExpirationMax, 30*time.Minute)),
-		),
-		externalCacheManager: cache.New[*idp.UserData](
-			cacheStore.NewGoCache(gocache.New(CacheExpirationMax, 30*time.Minute)),
-		),
 	}
 
+	cacheStore, err := nbcache.NewStore(nbcache.DefaultIDPCacheExpirationMax, nbcache.DefaultIDPCacheCleanupInterval)
+	assert.NoError(t, err)
+	am.externalCacheManager = nbcache.NewUserDataCache(cacheStore)
+	am.cacheManager = nbcache.NewAccountUserDataCache(am.loadAccount, cacheStore)
 	// pretend that we receive mockUserID from IDP
-	err = am.cacheManager.Set(am.ctx, mockAccountID, []*idp.UserData{{Name: mockUserID, ID: mockUserID}})
+	err = am.cacheManager.Set(am.ctx, mockAccountID, []*idp.UserData{{Name: mockUserID, ID: mockUserID}}, time.Minute)
 	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"})
+	err = cacheManager.Set(context.Background(), cacheKey, &idp.UserData{ID: externalUser.Id, Name: "Test User", Email: "user@example.com"}, time.Minute)
 	assert.NoError(t, err)
 
 	infos, err := am.GetUsersFromAccount(context.Background(), mockAccountID, mockUserID)