From 13ab50b9012d8618c6a684e82418df8e3c9cb298 Mon Sep 17 00:00:00 2001 From: Bethuel Mmbaga Date: Fri, 4 Sep 2026 11:03:28 +0300 Subject: [PATCH] [management] Add SetNX and GetDel cache store operations (#7084) --- .../service/manager/manager_test.go | 6 +- management/internals/server/boot.go | 6 +- .../internals/shared/grpc/pkce_verifier.go | 19 +-- .../shared/grpc/pkce_verifier_test.go | 85 ++++++++++ .../internals/shared/grpc/proxy_test.go | 3 +- management/server/auth/session.go | 32 ++-- management/server/auth/session_test.go | 51 ++++++ management/server/cache/memory.go | 57 +++++++ management/server/cache/memory_test.go | 76 +++++++++ management/server/cache/redis.go | 63 ++++++++ management/server/cache/redis_test.go | 153 ++++++++++++++++++ management/server/cache/store.go | 47 ++---- management/server/cache/store_test.go | 126 +++++---------- 13 files changed, 565 insertions(+), 159 deletions(-) create mode 100644 management/internals/shared/grpc/pkce_verifier_test.go create mode 100644 management/server/cache/memory.go create mode 100644 management/server/cache/memory_test.go create mode 100644 management/server/cache/redis.go create mode 100644 management/server/cache/redis_test.go diff --git a/management/internals/modules/reverseproxy/service/manager/manager_test.go b/management/internals/modules/reverseproxy/service/manager/manager_test.go index 10893673e..dd0edec60 100644 --- a/management/internals/modules/reverseproxy/service/manager/manager_test.go +++ b/management/internals/modules/reverseproxy/service/manager/manager_test.go @@ -7,11 +7,10 @@ import ( "testing" "time" - cachestore "github.com/eko/gocache/lib/v4/store" - "go.uber.org/mock/gomock" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.opentelemetry.io/otel/metric/noop" + "go.uber.org/mock/gomock" "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy" proxymanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy/manager" @@ -31,7 +30,7 @@ import ( "github.com/netbirdio/netbird/shared/management/status" ) -func testCacheStore(t *testing.T) cachestore.StoreInterface { +func testCacheStore(t *testing.T) nbcache.Store { t.Helper() s, err := nbcache.NewStore(context.Background(), 30*time.Minute, 10*time.Minute, 100) require.NoError(t, err) @@ -295,6 +294,7 @@ func TestPersistNewService(t *testing.T) { assert.Equal(t, status.AlreadyExists, sErr.Type()) }) } + func TestPreserveExistingAuthSecrets(t *testing.T) { mgr := &Manager{} diff --git a/management/internals/server/boot.go b/management/internals/server/boot.go index 0a4df3924..87a36a93b 100644 --- a/management/internals/server/boot.go +++ b/management/internals/server/boot.go @@ -21,8 +21,6 @@ import ( "google.golang.org/grpc/credentials" "google.golang.org/grpc/keepalive" - cachestore "github.com/eko/gocache/lib/v4/store" - "github.com/netbirdio/netbird/encryption" "github.com/netbirdio/netbird/formatter/hook" "github.com/netbirdio/netbird/management/internals/modules/agentnetwork" @@ -75,8 +73,8 @@ func (s *BaseServer) Metrics() telemetry.AppMetrics { // CacheStore returns a shared cache store backed by Redis or in-memory depending on the environment. // All consumers should reuse this store to avoid creating multiple Redis connections. -func (s *BaseServer) CacheStore() cachestore.StoreInterface { - return Create(s, func() cachestore.StoreInterface { +func (s *BaseServer) CacheStore() nbcache.Store { + return Create(s, func() nbcache.Store { cs, err := nbcache.NewStore(context.Background(), nbcache.DefaultStoreMaxTimeout, nbcache.DefaultStoreCleanupInterval, nbcache.DefaultStoreMaxConn) if err != nil { log.Fatalf("failed to create shared cache store: %v", err) diff --git a/management/internals/shared/grpc/pkce_verifier.go b/management/internals/shared/grpc/pkce_verifier.go index a1325256c..18155dc1d 100644 --- a/management/internals/shared/grpc/pkce_verifier.go +++ b/management/internals/shared/grpc/pkce_verifier.go @@ -5,22 +5,23 @@ import ( "fmt" "time" - "github.com/eko/gocache/lib/v4/cache" "github.com/eko/gocache/lib/v4/store" log "github.com/sirupsen/logrus" + + nbcache "github.com/netbirdio/netbird/management/server/cache" ) // PKCEVerifierStore manages PKCE verifiers for OAuth flows. // Supports both in-memory and Redis storage via NB_IDP_CACHE_REDIS_ADDRESS env var. type PKCEVerifierStore struct { - cache *cache.Cache[string] + cache nbcache.Store ctx context.Context } // NewPKCEVerifierStore creates a PKCE verifier store using the provided shared cache store. -func NewPKCEVerifierStore(ctx context.Context, cacheStore store.StoreInterface) *PKCEVerifierStore { +func NewPKCEVerifierStore(ctx context.Context, cacheStore nbcache.Store) *PKCEVerifierStore { return &PKCEVerifierStore{ - cache: cache.New[string](cacheStore), + cache: cacheStore, ctx: ctx, } } @@ -40,14 +41,14 @@ func (s *PKCEVerifierStore) Store(state, verifier string, ttl time.Duration) err // Returns the verifier and true if found, or empty string and false if not found. // This enforces single-use semantics for PKCE verifiers. func (s *PKCEVerifierStore) LoadAndDelete(state string) (string, bool) { - verifier, err := s.cache.Get(s.ctx, state) + verifier, found, err := s.cache.GetDel(s.ctx, state) if err != nil { - log.Debugf("PKCE verifier not found for state") + log.Warnf("Failed to consume PKCE verifier: %v", err) return "", false } - - if err := s.cache.Delete(s.ctx, state); err != nil { - log.Warnf("Failed to delete PKCE verifier for state: %v", err) + if !found { + log.Debug("PKCE verifier not found for state") + return "", false } return verifier, true diff --git a/management/internals/shared/grpc/pkce_verifier_test.go b/management/internals/shared/grpc/pkce_verifier_test.go new file mode 100644 index 000000000..e7175b6c5 --- /dev/null +++ b/management/internals/shared/grpc/pkce_verifier_test.go @@ -0,0 +1,85 @@ +package grpc + +import ( + "context" + "testing" + "time" +) + +func TestPKCEVerifierStoreLoadAndDelete(t *testing.T) { + const ( + state = "state" + verifier = "verifier" + attempts = 64 + ) + + t.Run("exactly one concurrent caller consumes the verifier", func(t *testing.T) { + store := NewPKCEVerifierStore(context.Background(), testCacheStore(t)) + if err := store.Store(state, verifier, time.Minute); err != nil { + t.Fatalf("couldn't store PKCE verifier: %s", err) + } + + start := make(chan struct{}) + type result struct { + verifier string + found bool + } + results := make(chan result, attempts) + for range attempts { + go func() { + <-start + verifier, found := store.LoadAndDelete(state) + results <- result{verifier: verifier, found: found} + }() + } + close(start) + + winners := 0 + for range attempts { + result := <-results + if result.found { + winners++ + if result.verifier != verifier { + t.Fatalf("unexpected verifier: got %q, expected %q", result.verifier, verifier) + } + } + } + if winners != 1 { + t.Fatalf("expected exactly one PKCE verifier consumer, got %d", winners) + } + }) + + t.Run("replayed state is rejected", func(t *testing.T) { + store := NewPKCEVerifierStore(context.Background(), testCacheStore(t)) + if err := store.Store(state, verifier, time.Minute); err != nil { + t.Fatalf("couldn't store PKCE verifier: %s", err) + } + + if got, found := store.LoadAndDelete(state); !found || got != verifier { + t.Fatalf("first load should return the verifier, got %q, found %t", got, found) + } + if got, found := store.LoadAndDelete(state); found { + t.Fatalf("replayed state should not resolve, got %q", got) + } + }) + + t.Run("unknown state is rejected", func(t *testing.T) { + store := NewPKCEVerifierStore(context.Background(), testCacheStore(t)) + + if got, found := store.LoadAndDelete("never-stored"); found { + t.Fatalf("unknown state should not resolve, got %q", got) + } + }) + + t.Run("expired verifier is rejected", func(t *testing.T) { + store := NewPKCEVerifierStore(context.Background(), testCacheStore(t)) + if err := store.Store(state, verifier, 50*time.Millisecond); err != nil { + t.Fatalf("couldn't store PKCE verifier: %s", err) + } + + time.Sleep(100 * time.Millisecond) + if got, found := store.LoadAndDelete(state); found { + t.Fatalf("expired verifier should not resolve, got %q", got) + } + }) +} diff --git a/management/internals/shared/grpc/proxy_test.go b/management/internals/shared/grpc/proxy_test.go index 0379edc6d..29b7c9523 100644 --- a/management/internals/shared/grpc/proxy_test.go +++ b/management/internals/shared/grpc/proxy_test.go @@ -9,7 +9,6 @@ import ( "testing" "time" - cachestore "github.com/eko/gocache/lib/v4/store" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "google.golang.org/grpc/codes" @@ -21,7 +20,7 @@ import ( "github.com/netbirdio/netbird/shared/management/proto" ) -func testCacheStore(t *testing.T) cachestore.StoreInterface { +func testCacheStore(t *testing.T) nbcache.Store { t.Helper() s, err := nbcache.NewStore(context.Background(), 30*time.Minute, 10*time.Minute, 100) require.NoError(t, err) diff --git a/management/server/auth/session.go b/management/server/auth/session.go index 7621a1c10..778146589 100644 --- a/management/server/auth/session.go +++ b/management/server/auth/session.go @@ -7,9 +7,6 @@ import ( "errors" "fmt" "time" - - "github.com/eko/gocache/lib/v4/cache" - "github.com/eko/gocache/lib/v4/store" ) const ( @@ -22,12 +19,17 @@ var ( ErrTokenExpired = errors.New("JWT expired") ) -type SessionStore struct { - cache *cache.Cache[string] +// TokenCache atomically records used JWTs until their expiration. +type TokenCache interface { + SetNX(ctx context.Context, key, value string, ttl time.Duration) (bool, error) } -func NewSessionStore(cacheStore store.StoreInterface) *SessionStore { - return &SessionStore{cache: cache.New[string](cacheStore)} +type SessionStore struct { + cache TokenCache +} + +func NewSessionStore(cacheStore TokenCache) *SessionStore { + return &SessionStore{cache: cacheStore} } // RegisterToken records a JWT until its exp time and rejects reuse. @@ -38,20 +40,14 @@ func (s *SessionStore) RegisterToken(ctx context.Context, token string, expiresA } key := usedTokenKeyPrefix + hashToken(token) - _, err := s.cache.Get(ctx, key) - if err == nil { + created, err := s.cache.SetNX(ctx, key, usedTokenMarker, ttl) + if err != nil { + return fmt.Errorf("store used token entry: %w", err) + } + if !created { return ErrTokenAlreadyUsed } - var notFound *store.NotFound - if !errors.As(err, ¬Found) { - return fmt.Errorf("failed to lookup used token entry: %w", err) - } - - if err := s.cache.Set(ctx, key, usedTokenMarker, store.WithExpiration(ttl)); err != nil { - return fmt.Errorf("failed to store used token entry: %w", err) - } - return nil } diff --git a/management/server/auth/session_test.go b/management/server/auth/session_test.go index 3a7d85f4c..7c82dfc43 100644 --- a/management/server/auth/session_test.go +++ b/management/server/auth/session_test.go @@ -2,6 +2,7 @@ package auth import ( "context" + "errors" "testing" "time" @@ -38,6 +39,39 @@ func TestSessionStore_RegisterSameTokenTwiceIsRejected(t *testing.T) { assert.ErrorIs(t, err, ErrTokenAlreadyUsed) } +func TestSessionStore_ConcurrentRegistrationAllowsOneCaller(t *testing.T) { + s := newTestSessionStore(t) + ctx := context.Background() + const attempts = 100 + + start := make(chan struct{}) + results := make(chan error, attempts) + for range attempts { + go func() { + <-start + results <- s.RegisterToken(ctx, "token", time.Now().Add(time.Hour)) + }() + } + close(start) + + succeeded := 0 + alreadyUsed := 0 + for range attempts { + err := <-results + switch { + case err == nil: + succeeded++ + case errors.Is(err, ErrTokenAlreadyUsed): + alreadyUsed++ + default: + require.NoError(t, err, "concurrent registration returned an unexpected error") + } + } + + assert.Equal(t, 1, succeeded, "exactly one concurrent caller should register the token") + assert.Equal(t, attempts-1, alreadyUsed, "every other caller should be rejected as already used") +} + func TestSessionStore_RegisterDifferentTokensAreIndependent(t *testing.T) { s := newTestSessionStore(t) ctx := context.Background() @@ -72,6 +106,23 @@ func TestSessionStore_EntryEvictsAtTTLAndAllowsReRegistration(t *testing.T) { require.NoError(t, s.RegisterToken(ctx, token, time.Now().Add(time.Hour))) } +type failingTokenCache struct { + err error +} + +func (f failingTokenCache) SetNX(context.Context, string, string, time.Duration) (bool, error) { + return false, f.err +} + +func TestSessionStore_CacheErrorIsReturned(t *testing.T) { + cacheErr := errors.New("cache unavailable") + s := NewSessionStore(failingTokenCache{err: cacheErr}) + + err := s.RegisterToken(context.Background(), "token", time.Now().Add(time.Hour)) + require.Error(t, err, "cache failure should be surfaced to the caller") + assert.ErrorIs(t, err, cacheErr, "cache error should be wrapped, not replaced") +} + func TestHashToken_StableAndDoesNotLeak(t *testing.T) { a := hashToken("tokenA") b := hashToken("tokenB") diff --git a/management/server/cache/memory.go b/management/server/cache/memory.go new file mode 100644 index 000000000..f140f3ec8 --- /dev/null +++ b/management/server/cache/memory.go @@ -0,0 +1,57 @@ +package cache + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/eko/gocache/lib/v4/store" + gocachestore "github.com/eko/gocache/store/go_cache/v4" + gocache "github.com/patrickmn/go-cache" +) + +type goCacheStore struct { + store.StoreInterface + client *gocache.Cache + mu sync.Mutex +} + +func newMemoryStore(maxTimeout, cleanupInterval time.Duration) Store { + client := gocache.New(maxTimeout, cleanupInterval) + return &goCacheStore{ + StoreInterface: gocachestore.NewGoCache(client), + client: client, + } +} + +func (s *goCacheStore) SetNX(_ context.Context, key, value string, ttl time.Duration) (bool, error) { + // Add only returns an error when a non-expired entry already exists. + if err := s.client.Add(key, value, ttl); err != nil { + return false, nil //nolint:nilerr + } + return true, nil +} + +// GetDel reads the value under key and removes it. go-cache has no native read-and-delete +// and releases its own lock between the two calls, so mu holds the pair together and no +// value is consumed twice. +// +// Writes do not take mu: a Set landing mid-pair is lost, since GetDel returns the prior +// value and deletes the new one. Callers must write a consumed key only once. +func (s *goCacheStore) GetDel(_ context.Context, key string) (string, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + + value, found := s.client.Get(key) + if !found { + return "", false, nil + } + s.client.Delete(key) + + str, ok := value.(string) + if !ok { + return "", false, fmt.Errorf("cached value is %T, not a string", value) + } + return str, true, nil +} diff --git a/management/server/cache/memory_test.go b/management/server/cache/memory_test.go new file mode 100644 index 000000000..363504921 --- /dev/null +++ b/management/server/cache/memory_test.go @@ -0,0 +1,76 @@ +package cache_test + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/management/server/cache" +) + +func TestMemoryStore(t *testing.T) { + memStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100) + require.NoError(t, err, "couldn't create memory store") + + ctx := context.Background() + key, value := "testing", "tested" + err = memStore.Set(ctx, key, value) + assert.NoError(t, err, "couldn't set testing data") + + result, err := memStore.Get(ctx, key) + assert.NoError(t, err, "couldn't get testing data") + assert.Equal(t, value, result, "value returned doesn't match testing data") + + created, err := memStore.SetNX(ctx, "conditional", value, 100*time.Millisecond) + require.NoError(t, err, "couldn't conditionally set testing data") + require.True(t, created, "first conditional set should create the entry") + + created, err = memStore.SetNX(ctx, "conditional", value, 100*time.Millisecond) + require.NoError(t, err, "couldn't conditionally check testing data") + require.False(t, created, "second conditional set should not replace the entry") + + // test expiration + time.Sleep(300 * time.Millisecond) + _, err = memStore.Get(ctx, key) + assert.Error(t, err, "value should not be found") +} + +func TestMemoryStoreGetDel(t *testing.T) { + ctx := context.Background() + newStore := func(t *testing.T) cache.Store { + t.Helper() + memStore, err := cache.NewStore(ctx, time.Minute, time.Minute, 100) + require.NoError(t, err, "couldn't create memory store") + + return memStore + } + + const ( + key = "consume" + value = "verifier" + ) + + t.Run("exactly one concurrent caller consumes the key", func(t *testing.T) { + memStore := newStore(t) + require.NoError(t, memStore.Set(ctx, key, value), "couldn't set testing data") + + assertGetDelConsumedOnce(ctx, t, []cache.Store{memStore}, key, value) + assertGetDelMisses(ctx, t, memStore, key) + }) + + t.Run("missing key is not an error", func(t *testing.T) { + assertGetDelMisses(ctx, t, newStore(t), "never-set") + }) + + t.Run("expired key is not found", func(t *testing.T) { + memStore := newStore(t) + _, err := memStore.SetNX(ctx, key, value, 50*time.Millisecond) + require.NoError(t, err, "couldn't set testing data") + + time.Sleep(100 * time.Millisecond) + assertGetDelMisses(ctx, t, memStore, key) + }) +} diff --git a/management/server/cache/redis.go b/management/server/cache/redis.go new file mode 100644 index 000000000..0cd921c92 --- /dev/null +++ b/management/server/cache/redis.go @@ -0,0 +1,63 @@ +package cache + +import ( + "context" + "errors" + "fmt" + "math" + "time" + + "github.com/eko/gocache/lib/v4/store" + redisstore "github.com/eko/gocache/store/redis/v4" + "github.com/redis/go-redis/v9" + log "github.com/sirupsen/logrus" +) + +type redisStore struct { + store.StoreInterface + client *redis.Client +} + +func getRedisStore(ctx context.Context, redisEnvAddr string, maxConn int) (Store, error) { + options, err := redis.ParseURL(redisEnvAddr) + if err != nil { + return nil, fmt.Errorf("parsing redis cache url: %s", err) + } + + options.MaxIdleConns = int(math.Ceil(float64(maxConn) * 0.5)) // 50% of max conns + options.MinIdleConns = int(math.Ceil(float64(maxConn) * 0.1)) // 10% of max conns + options.MaxActiveConns = maxConn + options.ConnMaxIdleTime = 30 * time.Minute + options.ConnMaxLifetime = 0 + options.PoolTimeout = 10 * time.Second + redisClient := redis.NewClient(options) + subCtx, cancel := context.WithTimeout(ctx, 2*time.Second) + defer cancel() + + _, err = redisClient.Ping(subCtx).Result() + if err != nil { + return nil, err + } + + log.WithContext(subCtx).Infof("using redis cache at %s", redisEnvAddr) + + return &redisStore{ + StoreInterface: redisstore.NewRedis(redisClient), + client: redisClient, + }, nil +} + +func (s *redisStore) SetNX(ctx context.Context, key, value string, ttl time.Duration) (bool, error) { + return s.client.SetNX(ctx, key, value, ttl).Result() +} + +func (s *redisStore) GetDel(ctx context.Context, key string) (string, bool, error) { + value, err := s.client.GetDel(ctx, key).Result() + if errors.Is(err, redis.Nil) { + return "", false, nil + } + if err != nil { + return "", false, err + } + return value, true, nil +} diff --git a/management/server/cache/redis_test.go b/management/server/cache/redis_test.go new file mode 100644 index 000000000..994ec7490 --- /dev/null +++ b/management/server/cache/redis_test.go @@ -0,0 +1,153 @@ +package cache_test + +import ( + "context" + "testing" + "time" + + "github.com/eko/gocache/lib/v4/store" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + testcontainersredis "github.com/testcontainers/testcontainers-go/modules/redis" + + "github.com/netbirdio/netbird/management/server/cache" +) + +func startRedis(t *testing.T) string { + t.Helper() + + ctx := context.Background() + redisContainer, err := testcontainersredis.Run(ctx, "redis:7") + require.NoError(t, err, "couldn't start redis container") + + t.Cleanup(func() { + if err := redisContainer.Terminate(ctx); err != nil { + t.Logf("failed to terminate container: %s", err) + } + }) + + redisURL, err := redisContainer.ConnectionString(ctx) + require.NoError(t, err, "couldn't get connection string") + + t.Setenv(cache.RedisStoreEnvVar, redisURL) + return redisURL +} + +func newRedisStore(t *testing.T) cache.Store { + t.Helper() + + redisStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100) + require.NoError(t, err) + + return redisStore +} + +func TestRedisStoreConnectionFailure(t *testing.T) { + t.Setenv(cache.RedisStoreEnvVar, "redis://127.0.0.1:6379") + _, err := cache.NewStore(context.Background(), 10*time.Millisecond, 30*time.Millisecond, 100) + require.Error(t, err, "getting redis cache store should return error") +} + +func TestRedisStoreConnectionSuccess(t *testing.T) { + ctx := context.Background() + redisURL := startRedis(t) + redisStore := newRedisStore(t) + + key, value := "testing", "tested" + err := redisStore.Set(ctx, key, value, store.WithExpiration(100*time.Millisecond)) + assert.NoError(t, err, "couldn't set testing data") + + result, err := redisStore.Get(ctx, key) + assert.NoError(t, err, "couldn't get testing data") + assert.Equal(t, value, result, "value returned doesn't match testing data") + + options, err := redis.ParseURL(redisURL) + require.NoError(t, err, "parsing redis cache url") + + redisClient := redis.NewClient(options) + r, err := redisClient.Get(ctx, key).Result() + assert.NoError(t, err, "couldn't get testing data from redis") + assert.Equal(t, value, r, "value returned from redis doesn't match testing data") + + // test expiration + time.Sleep(300 * time.Millisecond) + _, err = redisStore.Get(ctx, key) + assert.Error(t, err, "value should not be found") +} + +func TestRedisStoreSetNX(t *testing.T) { + ctx := context.Background() + redisURL := startRedis(t) + redisStore, secondRedisStore := newRedisStore(t), newRedisStore(t) + + const ( + key = "conditional" + value = "tested" + ) + + start := make(chan struct{}) + type setResult struct { + created bool + err error + } + results := make(chan setResult, 2) + for _, cacheStore := range []cache.Store{redisStore, secondRedisStore} { + go func() { + <-start + created, err := cacheStore.SetNX(ctx, key, value, time.Minute) + results <- setResult{created: created, err: err} + }() + } + close(start) + + created := 0 + for range 2 { + result := <-results + require.NoError(t, result.err, "conditional redis set failed") + if result.created { + created++ + } + } + require.Equal(t, 1, created, "expected exactly one redis client to create the entry") + + options, err := redis.ParseURL(redisURL) + require.NoError(t, err, "parsing redis cache url") + + ttl, err := redis.NewClient(options).PTTL(ctx, key).Result() + require.NoError(t, err, "couldn't read entry TTL") + require.Positive(t, ttl, "created entry should have a positive TTL") +} + +func TestRedisStoreGetDel(t *testing.T) { + ctx := context.Background() + startRedis(t) + redisStore, secondRedisStore := newRedisStore(t), newRedisStore(t) + + const ( + key = "consume" + value = "verifier" + ) + + t.Run("exactly one caller across independent clients consumes the key", func(t *testing.T) { + // A generous TTL: the key is consumed explicitly, so expiry racing the + // concurrent callers would only make the test flaky on a loaded runner. + err := redisStore.Set(ctx, key, value, store.WithExpiration(time.Minute)) + require.NoError(t, err, "couldn't set value to consume") + + assertGetDelConsumedOnce(ctx, t, []cache.Store{redisStore, secondRedisStore}, key, value) + assertGetDelMisses(ctx, t, secondRedisStore, key) + }) + + t.Run("missing key is not an error", func(t *testing.T) { + assertGetDelMisses(ctx, t, redisStore, "never-set") + }) + + t.Run("expired key is not found", func(t *testing.T) { + err := redisStore.Set(ctx, key, value, store.WithExpiration(50*time.Millisecond)) + require.NoError(t, err, "couldn't set value to consume") + + time.Sleep(100 * time.Millisecond) + assertGetDelMisses(ctx, t, redisStore, key) + }) +} diff --git a/management/server/cache/store.go b/management/server/cache/store.go index 2ca8e8603..a0c093e5d 100644 --- a/management/server/cache/store.go +++ b/management/server/cache/store.go @@ -2,17 +2,10 @@ package cache import ( "context" - "fmt" - "math" "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" - log "github.com/sirupsen/logrus" ) // RedisStoreEnvVar is the environment variable that determines if a redis store should be used. @@ -31,15 +24,23 @@ const ( DefaultStoreMaxConn = 1000 ) +// Store extends the shared cache interface with conditional and consuming operations. +type Store interface { + store.StoreInterface + // SetNX stores a value with a TTL only when the key does not exist. + SetNX(ctx context.Context, key, value string, ttl time.Duration) (bool, error) + // GetDel reads a value and removes it, so only one caller can consume a key. + GetDel(ctx context.Context, key string) (value string, found bool, err error) +} + // 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(ctx context.Context, maxTimeout, cleanupInterval time.Duration, maxConn int) (store.StoreInterface, error) { +func NewStore(ctx context.Context, maxTimeout, cleanupInterval time.Duration, maxConn int) (Store, error) { redisAddr := GetAddrFromEnv() if redisAddr != "" { return getRedisStore(ctx, redisAddr, maxConn) } - goc := gocache.New(maxTimeout, cleanupInterval) - return gocache_store.NewGoCache(goc), nil + return newMemoryStore(maxTimeout, cleanupInterval), nil } // GetAddrFromEnv returns the redis address from the environment variable RedisStoreEnvVar or its legacy counterpart. @@ -50,29 +51,3 @@ func GetAddrFromEnv() string { } return addr } - -func getRedisStore(ctx context.Context, redisEnvAddr string, maxConn int) (store.StoreInterface, error) { - options, err := redis.ParseURL(redisEnvAddr) - if err != nil { - return nil, fmt.Errorf("parsing redis cache url: %s", err) - } - - options.MaxIdleConns = int(math.Ceil(float64(maxConn) * 0.5)) // 50% of max conns - options.MinIdleConns = int(math.Ceil(float64(maxConn) * 0.1)) // 10% of max conns - options.MaxActiveConns = maxConn - options.ConnMaxIdleTime = 30 * time.Minute - options.ConnMaxLifetime = 0 - options.PoolTimeout = 10 * time.Second - redisClient := redis.NewClient(options) - subCtx, cancel := context.WithTimeout(ctx, 2*time.Second) - defer cancel() - - _, err = redisClient.Ping(subCtx).Result() - if err != nil { - return nil, err - } - - log.WithContext(subCtx).Infof("using redis cache at %s", redisEnvAddr) - - return redis_store.NewRedis(redisClient), nil -} diff --git a/management/server/cache/store_test.go b/management/server/cache/store_test.go index b869170f0..a59be8393 100644 --- a/management/server/cache/store_test.go +++ b/management/server/cache/store_test.go @@ -3,101 +3,53 @@ package cache_test import ( "context" "testing" - "time" - "github.com/eko/gocache/lib/v4/store" - "github.com/redis/go-redis/v9" - testcontainersredis "github.com/testcontainers/testcontainers-go/modules/redis" + "github.com/stretchr/testify/require" "github.com/netbirdio/netbird/management/server/cache" ) -func TestMemoryStore(t *testing.T) { - memStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100) - 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 assertGetDelConsumedOnce(ctx context.Context, t *testing.T, stores []cache.Store, key, value string) { + t.Helper() -func TestRedisStoreConnectionFailure(t *testing.T) { - t.Setenv(cache.RedisStoreEnvVar, "redis://127.0.0.1:6379") - _, err := cache.NewStore(context.Background(), 10*time.Millisecond, 30*time.Millisecond, 100) - if err == nil { - t.Fatal("getting redis cache store should return error") - } -} + const getDelAttempts = 64 -func TestRedisStoreConnectionSuccess(t *testing.T) { - ctx := context.Background() - redisContainer, err := testcontainersredis.Run(ctx, "redis:7") - if err != nil { - t.Fatalf("couldn't start redis container: %s", err) + type getDelResult struct { + value string + found bool + err error } - defer func() { - if err := redisContainer.Terminate(ctx); err != nil { - t.Logf("failed to terminate container: %s", err) + + start := make(chan struct{}) + results := make(chan getDelResult, getDelAttempts) + for i := range getDelAttempts { + cacheStore := stores[i%len(stores)] + go func() { + <-start + value, found, err := cacheStore.GetDel(ctx, key) + results <- getDelResult{value: value, found: found, err: err} + }() + } + close(start) + + consumers := 0 + for range getDelAttempts { + result := <-results + require.NoError(t, result.err, "concurrent GetDel failed") + if !result.found { + continue } - }() - 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(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100) - 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") + consumers++ + require.Equal(t, value, result.value, "consumed value doesn't match testing data") } + require.Equal(t, 1, consumers, "expected exactly one consumer") +} + +func assertGetDelMisses(ctx context.Context, t *testing.T, cacheStore cache.Store, key string) { + t.Helper() + + value, found, err := cacheStore.GetDel(ctx, key) + require.NoError(t, err, "GetDel on a missing key should not error") + require.False(t, found, "GetDel should not find key %q, got value %q", key, value) + require.Empty(t, value, "GetDel should return an empty value when not found") }