Merge branch 'main' into add-atomic-cache-ops

Split the memory and Redis cache implementations into separate files and
provide backend-native atomic set-if-absent support.

go.mod

@@ -24,7 +24,7 @@ require (
 	golang.zx2c4.com/wireguard v0.0.0-20231211153847-12269c276173
 	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20241231184526-a9ab2273dd10
 	golang.zx2c4.com/wireguard/windows v0.5.3
-	google.golang.org/grpc v1.81.1
+	google.golang.org/grpc v1.80.0
 	google.golang.org/protobuf v1.36.11
 )
 
@@ -116,11 +116,11 @@ require (
 	github.com/wailsapp/wails/v3 v3.0.0-alpha2.117
 	github.com/yusufpapurcu/wmi v1.2.4
 	github.com/zcalusic/sysinfo v1.1.3
-	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.69.0
-	go.opentelemetry.io/otel v1.44.0
-	go.opentelemetry.io/otel/exporters/prometheus v0.66.0
-	go.opentelemetry.io/otel/metric v1.44.0
-	go.opentelemetry.io/otel/sdk/metric v1.44.0
+	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.67.0
+	go.opentelemetry.io/otel v1.43.0
+	go.opentelemetry.io/otel/exporters/prometheus v0.64.0
+	go.opentelemetry.io/otel/metric v1.43.0
+	go.opentelemetry.io/otel/sdk/metric v1.43.0
 	go.uber.org/mock v0.6.0
 	go.uber.org/zap v1.27.0
 	goauthentik.io/api/v3 v3.2023051.3
@@ -292,7 +292,7 @@ require (
 	github.com/prometheus/client_model v0.6.2 // indirect
 	github.com/prometheus/common v0.67.5 // indirect
 	github.com/prometheus/otlptranslator v1.0.0 // indirect
-	github.com/prometheus/procfs v0.20.1 // indirect
+	github.com/prometheus/procfs v0.19.2 // indirect
 	github.com/russellhaering/goxmldsig v1.6.0 // indirect
 	github.com/ryanuber/go-glob v1.0.0 // indirect
 	github.com/shopspring/decimal v1.4.0 // indirect
@@ -309,15 +309,15 @@ require (
 	go.mongodb.org/mongo-driver v1.17.9 // indirect
 	go.opentelemetry.io/auto/sdk v1.2.1 // indirect
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.67.0 // indirect
-	go.opentelemetry.io/otel/sdk v1.44.0 // indirect
-	go.opentelemetry.io/otel/trace v1.44.0 // indirect
+	go.opentelemetry.io/otel/sdk v1.43.0 // indirect
+	go.opentelemetry.io/otel/trace v1.43.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
-	go.yaml.in/yaml/v2 v2.4.4 // indirect
+	go.yaml.in/yaml/v2 v2.4.3 // indirect
 	golang.org/x/text v0.40.0 // indirect
 	golang.org/x/tools v0.47.0 // indirect
 	golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20260319201613-d00831a3d3e7 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20260401024825-9d38bb4040a9 // indirect
 	gopkg.in/square/go-jose.v2 v2.6.0 // indirect
 	gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect

go.sum

@@ -578,8 +578,8 @@ github.com/prometheus/common v0.67.5 h1:pIgK94WWlQt1WLwAC5j2ynLaBRDiinoAb86HZHTU
 github.com/prometheus/common v0.67.5/go.mod h1:SjE/0MzDEEAyrdr5Gqc6G+sXI67maCxzaT3A2+HqjUw=
 github.com/prometheus/otlptranslator v1.0.0 h1:s0LJW/iN9dkIH+EnhiD3BlkkP5QVIUVEoIwkU+A6qos=
 github.com/prometheus/otlptranslator v1.0.0/go.mod h1:vRYWnXvI6aWGpsdY/mOT/cbeVRBlPWtBNDb7kGR3uKM=
-github.com/prometheus/procfs v0.20.1 h1:XwbrGOIplXW/AU3YhIhLODXMJYyC1isLFfYCsTEycfc=
-github.com/prometheus/procfs v0.20.1/go.mod h1:o9EMBZGRyvDrSPH1RqdxhojkuXstoe4UlK79eF5TGGo=
+github.com/prometheus/procfs v0.19.2 h1:zUMhqEW66Ex7OXIiDkll3tl9a1ZdilUOd/F6ZXw4Vws=
+github.com/prometheus/procfs v0.19.2/go.mod h1:M0aotyiemPhBCM0z5w87kL22CxfcH05ZpYlu+b4J7mw=
 github.com/quic-go/quic-go v0.55.0 h1:zccPQIqYCXDt5NmcEabyYvOnomjs8Tlwl7tISjJh9Mk=
 github.com/quic-go/quic-go v0.55.0/go.mod h1:DR51ilwU1uE164KuWXhinFcKWGlEjzys2l8zUl5Ss1U=
 github.com/redis/go-redis/v9 v9.7.3 h1:YpPyAayJV+XErNsatSElgRZZVCwXX9QzkKYNvO7x0wM=
@@ -684,28 +684,26 @@ go.mongodb.org/mongo-driver v1.17.9 h1:IexDdCuuNJ3BHrELgBlyaH9p60JXAvdzWR128q+U5
 go.mongodb.org/mongo-driver v1.17.9/go.mod h1:LlOhpH5NUEfhxcAwG0UEkMqwYcc4JU18gtCdGudk/tQ=
 go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
 go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
-go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.69.0 h1:2yEATaop1/a1I4psnSLgWVPLWwCzkqWakgJy7xTDVy0=
-go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.69.0/go.mod h1:D7J12YRapIekYyPWgGPlA/23pRmpSEZC5xJC/TTLI9U=
+go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.67.0 h1:yI1/OhfEPy7J9eoa6Sj051C7n5dvpj0QX8g4sRchg04=
+go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.67.0/go.mod h1:NoUCKYWK+3ecatC4HjkRktREheMeEtrXoQxrqYFeHSc=
 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.67.0 h1:OyrsyzuttWTSur2qN/Lm0m2a8yqyIjUVBZcxFPuXq2o=
 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.67.0/go.mod h1:C2NGBr+kAB4bk3xtMXfZ94gqFDtg/GkI7e9zqGh5Beg=
-go.opentelemetry.io/otel v1.44.0 h1:JjwHmHpA4iZ3wBxluu2fbbE7j4kqlE8jXyAyPXH7HqU=
-go.opentelemetry.io/otel v1.44.0/go.mod h1:BMgjTHL9WPRlRjL2oZCBTL4whCGtXch2H4BhOPIAyYc=
+go.opentelemetry.io/otel v1.43.0 h1:mYIM03dnh5zfN7HautFE4ieIig9amkNANT+xcVxAj9I=
+go.opentelemetry.io/otel v1.43.0/go.mod h1:JuG+u74mvjvcm8vj8pI5XiHy1zDeoCS2LB1spIq7Ay0=
 go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.19.0 h1:Mne5On7VWdx7omSrSSZvM4Kw7cS7NQkOOmLcgscI51U=
 go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.19.0/go.mod h1:IPtUMKL4O3tH5y+iXVyAXqpAwMuzC1IrxVS81rummfE=
 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.19.0 h1:IeMeyr1aBvBiPVYihXIaeIZba6b8E1bYp7lbdxK8CQg=
 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.19.0/go.mod h1:oVdCUtjq9MK9BlS7TtucsQwUcXcymNiEDjgDD2jMtZU=
-go.opentelemetry.io/otel/exporters/prometheus v0.66.0 h1:vkrK8PAznv2NKt2r+kdu252ccGzkEqLc2aSXbQIALYQ=
-go.opentelemetry.io/otel/exporters/prometheus v0.66.0/go.mod h1:V/UB6D3vMF/UBOL5igAsAYnk1nG/bzYYTzvsB16cy7o=
-go.opentelemetry.io/otel/metric v1.44.0 h1:1w0gILTcHdr3YI+ixLyjemwrVnsMURbTZFrSYCdDdmc=
-go.opentelemetry.io/otel/metric v1.44.0/go.mod h1:8O7hanEPBNgEMmybD3s2VBKcgWOCsA6tzHBPODAiquo=
-go.opentelemetry.io/otel/metric/x v0.66.0 h1:YkCrx1zLOChi9ZcZ6euupOcsgzbVlec7D/xoEU1+cTA=
-go.opentelemetry.io/otel/metric/x v0.66.0/go.mod h1:d1+BDj9t96do0/1LoU1ayfCv79ZgNE41qbhBvnMOBZk=
-go.opentelemetry.io/otel/sdk v1.44.0 h1:nHYwb9lK+fJPU/dnT6s7W7Z8itMWyqrnVfbheVYrZ58=
-go.opentelemetry.io/otel/sdk v1.44.0/go.mod h1:Osuydd3Se74nqjAKxid74N5eC+jfEqfTegHRnq58oK0=
-go.opentelemetry.io/otel/sdk/metric v1.44.0 h1:3LlKgI+VjbVsjNRFZJZAJ30WjXC5VkNRks6si09iEfI=
-go.opentelemetry.io/otel/sdk/metric v1.44.0/go.mod h1:5B5pMARnXxKhltooO4xUuCBorl65a4EpnTalObqOigA=
-go.opentelemetry.io/otel/trace v1.44.0 h1:jxF5CsGYCe74MCRx2X4g7WsY/VBKRqqpNvXlX/6gtIk=
-go.opentelemetry.io/otel/trace v1.44.0/go.mod h1:oLl1jrMQAVo6v3GAggN+1VH9VIz9iUSvW53sW1Q8PIE=
+go.opentelemetry.io/otel/exporters/prometheus v0.64.0 h1:g0LRDXMX/G1SEZtK8zl8Chm4K6GBwRkjPKE36LxiTYs=
+go.opentelemetry.io/otel/exporters/prometheus v0.64.0/go.mod h1:UrgcjnarfdlBDP3GjDIJWe6HTprwSazNjwsI+Ru6hro=
+go.opentelemetry.io/otel/metric v1.43.0 h1:d7638QeInOnuwOONPp4JAOGfbCEpYb+K6DVWvdxGzgM=
+go.opentelemetry.io/otel/metric v1.43.0/go.mod h1:RDnPtIxvqlgO8GRW18W6Z/4P462ldprJtfxHxyKd2PY=
+go.opentelemetry.io/otel/sdk v1.43.0 h1:pi5mE86i5rTeLXqoF/hhiBtUNcrAGHLKQdhg4h4V9Dg=
+go.opentelemetry.io/otel/sdk v1.43.0/go.mod h1:P+IkVU3iWukmiit/Yf9AWvpyRDlUeBaRg6Y+C58QHzg=
+go.opentelemetry.io/otel/sdk/metric v1.43.0 h1:S88dyqXjJkuBNLeMcVPRFXpRw2fuwdvfCGLEo89fDkw=
+go.opentelemetry.io/otel/sdk/metric v1.43.0/go.mod h1:C/RJtwSEJ5hzTiUz5pXF1kILHStzb9zFlIEe85bhj6A=
+go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09nk+3A=
+go.opentelemetry.io/otel/trace v1.43.0/go.mod h1:/QJhyVBUUswCphDVxq+8mld+AvhXZLhe+8WVFxiFff0=
 go.opentelemetry.io/proto/otlp v1.0.0 h1:T0TX0tmXU8a3CbNXzEKGeU5mIVOdf0oykP+u2lIVU/I=
 go.opentelemetry.io/proto/otlp v1.0.0/go.mod h1:Sy6pihPLfYHkr3NkUbEhGHFhINUSI/v80hjKIs5JXpM=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
@@ -716,8 +714,8 @@ go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
 go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
 go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8=
 go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E=
-go.yaml.in/yaml/v2 v2.4.4 h1:tuyd0P+2Ont/d6e2rl3be67goVK4R6deVxCUX5vyPaQ=
-go.yaml.in/yaml/v2 v2.4.4/go.mod h1:gMZqIpDtDqOfM0uNfy0SkpRhvUryYH0Z6wdMYcacYXQ=
+go.yaml.in/yaml/v2 v2.4.3 h1:6gvOSjQoTB3vt1l+CU+tSyi/HOjfOjRLJ4YwYZGwRO0=
+go.yaml.in/yaml/v2 v2.4.3/go.mod h1:zSxWcmIDjOzPXpjlTTbAsKokqkDNAVtZO0WOMiT90s8=
 go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
 goauthentik.io/api/v3 v3.2023051.3 h1:NebAhD/TeTWNo/9X3/Uj+rM5fG1HaiLOlKTNLQv9Qq4=
 goauthentik.io/api/v3 v3.2023051.3/go.mod h1:nYECml4jGbp/541hj8GcylKQG1gVBsKppHy4+7G8u4U=
@@ -884,10 +882,10 @@ google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7 h1:XzmzkmB14QhVhgn
 google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7/go.mod h1:L43LFes82YgSonw6iTXTxXUX1OlULt4AQtkik4ULL/I=
 google.golang.org/genproto/googleapis/api v0.0.0-20260319201613-d00831a3d3e7 h1:41r6JMbpzBMen0R/4TZeeAmGXSJC7DftGINUodzTkPI=
 google.golang.org/genproto/googleapis/api v0.0.0-20260319201613-d00831a3d3e7/go.mod h1:EIQZ5bFCfRQDV4MhRle7+OgjNtZ6P1PiZBgAKuxXu/Y=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa h1:mZHHdPZl0dbGHCflZgAq/Q468DWVFcU2whhB2KAo8fk=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8=
-google.golang.org/grpc v1.81.1 h1:VnnIIZ88UzOOKLukQi+ImGz8O1Wdp8nAGGnvOfEIWQQ=
-google.golang.org/grpc v1.81.1/go.mod h1:xGH9GfzOyMTGIOXBJmXt+BX/V0kcdQbdcuwQ/zNw42I=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20260401024825-9d38bb4040a9 h1:m8qni9SQFH0tJc1X0vmnpw/0t+AImlSvp30sEupozUg=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20260401024825-9d38bb4040a9/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8=
+google.golang.org/grpc v1.80.0 h1:Xr6m2WmWZLETvUNvIUmeD5OAagMw3FiKmMlTdViWsHM=
+google.golang.org/grpc v1.80.0/go.mod h1:ho/dLnxwi3EDJA4Zghp7k2Ec1+c2jqup0bFkw07bwF4=
 google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8=
 google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0=
 google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM=

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)

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,19 +40,13 @@ func (s *SessionStore) RegisterToken(ctx context.Context, token string, expiresA
 	}
 
 	key := usedTokenKeyPrefix + hashToken(token)
-	_, err := s.cache.Get(ctx, key)
-	if err == nil {
-		return ErrTokenAlreadyUsed
-	}
-
-	var notFound *store.NotFound
-	if !errors.As(err, &notFound) {
-		return fmt.Errorf("failed to lookup used token entry: %w", err)
-	}
-
-	if err := s.cache.Set(ctx, key, usedTokenMarker, store.WithExpiration(ttl)); err != nil {
+	created, err := s.cache.SetNX(ctx, key, usedTokenMarker, ttl)
+	if err != nil {
 		return fmt.Errorf("failed to store used token entry: %w", err)
 	}
+	if !created {
+		return ErrTokenAlreadyUsed
+	}
 
 	return nil
 }

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)
+		}
+	}
+
+	assert.Equal(t, 1, succeeded)
+	assert.Equal(t, attempts-1, alreadyUsed)
+}
+
 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)
+	assert.ErrorIs(t, err, cacheErr)
+}
+
 func TestHashToken_StableAndDoesNotLeak(t *testing.T) {
 	a := hashToken("tokenA")
 	b := hashToken("tokenB")

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
+}

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")
+	}
+}

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()
+}

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)

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
-}

proxy/Dockerfile.multistage

@@ -10,7 +10,6 @@ COPY encryption ./encryption
 COPY flow ./flow
 COPY formatter ./formatter
 COPY monotime ./monotime
-COPY management ./management
 COPY proxy ./proxy
 COPY route ./route
 COPY shared ./shared