diff --git a/management/internals/server/boot.go b/management/internals/server/boot.go index 1c78af9d0..b06407e15 100644 --- a/management/internals/server/boot.go +++ b/management/internals/server/boot.go @@ -20,17 +20,15 @@ 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" "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs" accesslogsmanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs/manager" rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc" "github.com/netbirdio/netbird/management/server/activity" activitystore "github.com/netbirdio/netbird/management/server/activity/store" - "github.com/netbirdio/netbird/management/internals/modules/agentnetwork" nbcache "github.com/netbirdio/netbird/management/server/cache" nbContext "github.com/netbirdio/netbird/management/server/context" nbhttp "github.com/netbirdio/netbird/management/server/http" @@ -70,8 +68,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/server/cache/memory.go b/management/server/cache/memory.go new file mode 100644 index 000000000..ba3b4cc89 --- /dev/null +++ b/management/server/cache/memory.go @@ -0,0 +1,31 @@ +package cache + +import ( + "context" + "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 +} + +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 + } + return true, nil +} diff --git a/management/server/cache/memory_test.go b/management/server/cache/memory_test.go new file mode 100644 index 000000000..226da8a45 --- /dev/null +++ b/management/server/cache/memory_test.go @@ -0,0 +1,49 @@ +package cache_test + +import ( + "context" + "testing" + "time" + + "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) + } + created, err := memStore.SetNX(ctx, "atomic", value, 100*time.Millisecond) + if err != nil { + t.Fatalf("couldn't atomically set testing data: %s", err) + } + if !created { + t.Fatal("first atomic set should create the entry") + } + created, err = memStore.SetNX(ctx, "atomic", value, 100*time.Millisecond) + if err != nil { + t.Fatalf("couldn't atomically check testing data: %s", err) + } + if created { + t.Fatal("second atomic set should not replace the entry") + } + // test expiration + time.Sleep(300 * time.Millisecond) + _, err = memStore.Get(ctx, key) + if err == nil { + t.Error("value should not be found") + } +} diff --git a/management/server/cache/redis.go b/management/server/cache/redis.go new file mode 100644 index 000000000..15b936d06 --- /dev/null +++ b/management/server/cache/redis.go @@ -0,0 +1,51 @@ +package cache + +import ( + "context" + "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() +} diff --git a/management/server/cache/store_test.go b/management/server/cache/redis_test.go similarity index 69% rename from management/server/cache/store_test.go rename to management/server/cache/redis_test.go index b869170f0..6e663e790 100644 --- a/management/server/cache/store_test.go +++ b/management/server/cache/redis_test.go @@ -12,32 +12,6 @@ import ( "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 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) @@ -94,6 +68,47 @@ func TestRedisStoreConnectionSuccess(t *testing.T) { if value != r { t.Errorf("value returned from redis doesn't match testing data, got %s, expected %s", r, value) } + + secondRedisStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100) + if err != nil { + t.Fatalf("couldn't create second redis store: %s", err) + } + 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, "atomic", value, time.Second) + results <- setResult{created: created, err: err} + }() + } + close(start) + + created := 0 + for range 2 { + result := <-results + if result.err != nil { + t.Fatalf("atomic redis set failed: %s", result.err) + } + if result.created { + created++ + } + } + if created != 1 { + t.Fatalf("expected exactly one redis client to create the entry, got %d", created) + } + ttl, err := redisClient.PTTL(ctx, "atomic").Result() + if err != nil { + t.Fatalf("couldn't read atomic entry TTL: %s", err) + } + if ttl <= 0 { + t.Fatalf("atomic entry should have a positive TTL, got %s", ttl) + } + // test expiration time.Sleep(300 * time.Millisecond) _, err = redisStore.Get(ctx, key) diff --git a/management/server/cache/store.go b/management/server/cache/store.go index 2ca8e8603..63c3e0bb7 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,21 @@ const ( DefaultStoreMaxConn = 1000 ) +// Store extends the shared cache interface with atomic insertion support. +type Store interface { + store.StoreInterface + // SetNX atomically 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) +} + // 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 +49,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 -}