From 64a89fb89530853990db709b5813a42ae4c28f0b Mon Sep 17 00:00:00 2001 From: Flo <53355483+Flo4604@users.noreply.github.com> Date: Mon, 21 Jul 2025 17:32:19 +0200 Subject: [PATCH 1/3] Add cache hit status to cache and add logic to return null correctly --- .../cache_validation_test.go | 5 +- .../api/routes/v2_ratelimit_limit/handler.go | 11 +- go/internal/services/keys/get.go | 14 +- go/pkg/cache/cache.go | 24 ++- go/pkg/cache/cache_test.go | 8 +- go/pkg/cache/interface.go | 2 +- go/pkg/cache/middleware/tracing.go | 7 +- go/pkg/cache/noop.go | 4 +- go/pkg/cache/simulation_test.go | 2 +- go/pkg/cache/swr_test.go | 157 ++++++++++++++++++ 10 files changed, 214 insertions(+), 20 deletions(-) create mode 100644 go/pkg/cache/swr_test.go diff --git a/go/apps/api/routes/v2_apis_delete_api/cache_validation_test.go b/go/apps/api/routes/v2_apis_delete_api/cache_validation_test.go index 1ae0abe4010..202a6457512 100644 --- a/go/apps/api/routes/v2_apis_delete_api/cache_validation_test.go +++ b/go/apps/api/routes/v2_apis_delete_api/cache_validation_test.go @@ -46,10 +46,11 @@ func TestCacheInvalidation(t *testing.T) { api := h.CreateApi(seed.CreateApiRequest{WorkspaceID: h.Resources().UserWorkspace.ID}) // Get API to ensure it's in the cache - _, err := h.Caches.ApiByID.SWR(ctx, api.ID, func(ctx context.Context) (db.Api, error) { + _, hit, err := h.Caches.ApiByID.SWR(ctx, api.ID, func(ctx context.Context) (db.Api, error) { return db.Query.FindApiByID(ctx, h.DB.RO(), api.ID) }, caches.DefaultFindFirstOp) require.NoError(t, err) + require.Equal(t, cache.Hit, hit) // Delete the API req := handler.Request{ ApiId: api.ID, @@ -70,7 +71,7 @@ func TestCacheInvalidation(t *testing.T) { require.True(t, apiAfterDelete.DeletedAtM.Valid) // Verify the API is deleted in the cache - _, hit := h.Caches.ApiByID.Get(ctx, api.ID) + _, hit = h.Caches.ApiByID.Get(ctx, api.ID) require.Equal(t, cache.Null, hit) }) } diff --git a/go/apps/api/routes/v2_ratelimit_limit/handler.go b/go/apps/api/routes/v2_ratelimit_limit/handler.go index e042d1f00c5..e45a706c827 100644 --- a/go/apps/api/routes/v2_ratelimit_limit/handler.go +++ b/go/apps/api/routes/v2_ratelimit_limit/handler.go @@ -70,7 +70,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { } ctx, span := tracing.Start(ctx, "FindRatelimitNamespace") - namespace, err := h.RatelimitNamespaceByNameCache.SWR(ctx, req.Namespace, func(ctx context.Context) (db.FindRatelimitNamespace, error) { + namespace, hit, err := h.RatelimitNamespaceByNameCache.SWR(ctx, req.Namespace, func(ctx context.Context) (db.FindRatelimitNamespace, error) { response, err := db.Query.FindRatelimitNamespace(ctx, h.DB.RO(), db.FindRatelimitNamespaceParams{ WorkspaceID: auth.AuthorizedWorkspaceID, Name: sql.NullString{String: req.Namespace, Valid: true}, @@ -120,6 +120,15 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } + if hit == cache.Null { + if db.IsNotFound(err) { + return fault.New("namespace cache null", + fault.Code(codes.Data.RatelimitNamespace.NotFound.URN()), + fault.Public("This namespace does not exist."), + ) + } + } + if namespace.DeletedAtM.Valid { return fault.New("namespace was deleted", fault.Code(codes.Data.RatelimitNamespace.NotFound.URN()), diff --git a/go/internal/services/keys/get.go b/go/internal/services/keys/get.go index 06dcb0e5f9c..a07c4d6fb66 100644 --- a/go/internal/services/keys/get.go +++ b/go/internal/services/keys/get.go @@ -4,10 +4,12 @@ import ( "context" "encoding/json" "fmt" + "log" "time" "github.com/unkeyed/unkey/go/internal/services/caches" "github.com/unkeyed/unkey/go/pkg/assert" + "github.com/unkeyed/unkey/go/pkg/cache" "github.com/unkeyed/unkey/go/pkg/db" "github.com/unkeyed/unkey/go/pkg/fault" "github.com/unkeyed/unkey/go/pkg/hash" @@ -60,7 +62,7 @@ func (s *service) Get(ctx context.Context, sess *zen.Session, rawKey string) (*K } h := hash.Sha256(rawKey) - key, err := s.keyCache.SWR(ctx, h, func(ctx context.Context) (db.FindKeyForVerificationRow, error) { + key, hit, err := s.keyCache.SWR(ctx, h, func(ctx context.Context) (db.FindKeyForVerificationRow, error) { return db.Query.FindKeyForVerification(ctx, s.db.RO(), h) }, caches.DefaultFindFirstOp) if err != nil { @@ -79,6 +81,16 @@ func (s *service) Get(ctx context.Context, sess *zen.Session, rawKey string) (*K ) } + if hit == cache.Null { + // nolint:exhaustruct + return &KeyVerifier{ + Status: StatusNotFound, + message: "key does not exist", + }, nil + } + + log.Printf("Found key %#v", key) + // ForWorkspace set but that doesn't exist if key.ForWorkspaceID.Valid && !key.ForWorkspaceEnabled.Valid { // nolint:exhaustruct diff --git a/go/pkg/cache/cache.go b/go/pkg/cache/cache.go index 013fa30ea25..c2c8f1fcd34 100644 --- a/go/pkg/cache/cache.go +++ b/go/pkg/cache/cache.go @@ -239,7 +239,7 @@ func (c *cache[K, V]) SWR( key K, refreshFromOrigin func(context.Context) (V, error), op func(error) Op, -) (V, error) { +) (V, CacheHit, error) { now := c.clock.Now() e, ok := c.get(ctx, key) if ok { @@ -247,7 +247,7 @@ func (c *cache[K, V]) SWR( if now.Before(e.Fresh) { // We have data and it's fresh, so we return it - return e.Value, nil + return e.Value, e.Hit, nil } if now.Before(e.Stale) { @@ -260,7 +260,7 @@ func (c *cache[K, V]) SWR( c.revalidate(context.WithoutCancel(ctx), key, refreshFromOrigin, op) } - return e.Value, nil + return e.Value, e.Hit, nil } // We have old data, that we should not serve anymore @@ -281,5 +281,21 @@ func (c *cache[K, V]) SWR( break } - return v, err + if err != nil { + // Error occurred, return Miss as the cache hit status + return v, Miss, err + } + + // Determine cache hit status based on the operation + var hit CacheHit + switch op(err) { + case WriteValue: + hit = Hit + case WriteNull: + hit = Null + default: + hit = Miss + } + + return v, hit, err } diff --git a/go/pkg/cache/cache_test.go b/go/pkg/cache/cache_test.go index b0f57b5e976..f2e3c05bb2d 100644 --- a/go/pkg/cache/cache_test.go +++ b/go/pkg/cache/cache_test.go @@ -15,7 +15,7 @@ import ( func TestWriteRead(t *testing.T) { - c, err := cache.New[string, string](cache.Config[string, string]{ + c, err := cache.New(cache.Config[string, string]{ MaxSize: 10_000, Fresh: time.Minute, @@ -33,7 +33,7 @@ func TestWriteRead(t *testing.T) { func TestEviction(t *testing.T) { clk := clock.NewTestClock() - c, err := cache.New[string, string](cache.Config[string, string]{ + c, err := cache.New(cache.Config[string, string]{ MaxSize: 10_000, Fresh: time.Second, @@ -56,7 +56,7 @@ func TestRefresh(t *testing.T) { // count how many times we refreshed from origin refreshedFromOrigin := atomic.Int32{} - c, err := cache.New[string, string](cache.Config[string, string]{ + c, err := cache.New(cache.Config[string, string]{ MaxSize: 10_000, Fresh: time.Second * 2, @@ -80,7 +80,7 @@ func TestRefresh(t *testing.T) { func TestNull(t *testing.T) { - c, err := cache.New[string, string](cache.Config[string, string]{ + c, err := cache.New(cache.Config[string, string]{ MaxSize: 10_000, Fresh: time.Second * 1, Stale: time.Minute * 5, diff --git a/go/pkg/cache/interface.go b/go/pkg/cache/interface.go index 0978698e06a..e5831774334 100644 --- a/go/pkg/cache/interface.go +++ b/go/pkg/cache/interface.go @@ -18,7 +18,7 @@ type Cache[K comparable, V any] interface { // Removes the key from the cache. Remove(ctx context.Context, key K) - SWR(ctx context.Context, key K, refreshFromOrigin func(ctx context.Context) (V, error), op func(error) Op) (value V, err error) + SWR(ctx context.Context, key K, refreshFromOrigin func(ctx context.Context) (V, error), op func(error) Op) (value V, hit CacheHit, err error) // Dump returns a serialized representation of the cache. Dump(ctx context.Context) ([]byte, error) diff --git a/go/pkg/cache/middleware/tracing.go b/go/pkg/cache/middleware/tracing.go index e96360ea5db..a1ba00dc23e 100644 --- a/go/pkg/cache/middleware/tracing.go +++ b/go/pkg/cache/middleware/tracing.go @@ -84,20 +84,19 @@ func (mw *tracingMiddleware[K, V]) Clear(ctx context.Context) { mw.next.Clear(ctx) } -func (mw *tracingMiddleware[K, V]) SWR(ctx context.Context, key K, refreshFromOrigin func(ctx context.Context) (V, error), op func(err error) cache.Op) (V, error) { +func (mw *tracingMiddleware[K, V]) SWR(ctx context.Context, key K, refreshFromOrigin func(ctx context.Context) (V, error), op func(err error) cache.Op) (V, cache.CacheHit, error) { ctx, span := tracing.Start(ctx, "cache.SWR") defer span.End() span.SetAttributes(attribute.String("key", fmt.Sprintf("%v", key))) - value, err := mw.next.SWR(ctx, key, func(innerCtx context.Context) (V, error) { + value, hit, err := mw.next.SWR(ctx, key, func(innerCtx context.Context) (V, error) { innerCtx, innerSpan := tracing.Start(innerCtx, "refreshFromOrigin") defer innerSpan.End() - return refreshFromOrigin(innerCtx) }, op) if err != nil { tracing.RecordError(span, err) } - return value, err + return value, hit, err } diff --git a/go/pkg/cache/noop.go b/go/pkg/cache/noop.go index e81fb6db2fe..43fb5711201 100644 --- a/go/pkg/cache/noop.go +++ b/go/pkg/cache/noop.go @@ -22,9 +22,9 @@ func (c *noopCache[K, V]) Restore(ctx context.Context, data []byte) error { return nil } func (c *noopCache[K, V]) Clear(ctx context.Context) {} -func (c *noopCache[K, V]) SWR(ctx context.Context, key K, refreshFromOrigin func(context.Context) (V, error), op func(err error) Op) (V, error) { +func (c *noopCache[K, V]) SWR(ctx context.Context, key K, refreshFromOrigin func(context.Context) (V, error), op func(err error) Op) (V, CacheHit, error) { var v V - return v, nil + return v, Miss, nil } func NewNoopCache[K comparable, V any]() Cache[K, V] { diff --git a/go/pkg/cache/simulation_test.go b/go/pkg/cache/simulation_test.go index ddd40caed7b..7b9db03b040 100644 --- a/go/pkg/cache/simulation_test.go +++ b/go/pkg/cache/simulation_test.go @@ -106,7 +106,7 @@ func TestSimulation(t *testing.T) { fresh := time.Second + time.Duration(rng.IntN(60*60*1000))*time.Millisecond stale := fresh + time.Duration(rng.IntN(24*60*60*1000))*time.Millisecond - c, err := cache.New[uint64, uint64](cache.Config[uint64, uint64]{ + c, err := cache.New(cache.Config[uint64, uint64]{ Clock: clk, Fresh: fresh, Stale: stale, diff --git a/go/pkg/cache/swr_test.go b/go/pkg/cache/swr_test.go new file mode 100644 index 00000000000..69c19a77bb9 --- /dev/null +++ b/go/pkg/cache/swr_test.go @@ -0,0 +1,157 @@ +package cache_test + +import ( + "context" + "database/sql" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/require" + "github.com/unkeyed/unkey/go/pkg/cache" + "github.com/unkeyed/unkey/go/pkg/clock" + "github.com/unkeyed/unkey/go/pkg/db" + "github.com/unkeyed/unkey/go/pkg/otel/logging" +) + +func TestSWR_CacheHit(t *testing.T) { + ctx := context.Background() + mockClock := clock.NewTestClock() + logger := logging.New() + + c, err := cache.New(cache.Config[string, string]{ + Fresh: 1 * time.Minute, + Stale: 5 * time.Minute, + Logger: logger, + MaxSize: 100, + Resource: "test", + Clock: mockClock, + }) + require.NoError(t, err) + + t.Run("miss on first call", func(t *testing.T) { + value, hit, err := c.SWR(ctx, "key1", func(ctx context.Context) (string, error) { + return "value1", nil + }, func(err error) cache.Op { + if err != nil { + return cache.Noop + } + return cache.WriteValue + }) + + require.NoError(t, err) + require.Equal(t, "value1", value) + require.Equal(t, cache.Hit, hit) + }) + + t.Run("hit on subsequent call within fresh time", func(t *testing.T) { + // First call to populate cache + _, _, err := c.SWR(ctx, "key2", func(ctx context.Context) (string, error) { + return "value2", nil + }, func(err error) cache.Op { + return cache.WriteValue + }) + require.NoError(t, err) + + // Second call should hit cache + value, hit, err := c.SWR(ctx, "key2", func(ctx context.Context) (string, error) { + t.Fatal("should not call refresh function") + return "", nil + }, func(err error) cache.Op { + return cache.WriteValue + }) + + require.NoError(t, err) + require.Equal(t, "value2", value) + require.Equal(t, cache.Hit, hit) + }) + + t.Run("null cache hit", func(t *testing.T) { + // First call returns not found error + _, _, err := c.SWR(ctx, "key3", func(ctx context.Context) (string, error) { + return "", sql.ErrNoRows + }, func(err error) cache.Op { + if db.IsNotFound(err) { + return cache.WriteNull + } + return cache.Noop + }) + require.Error(t, err) + require.True(t, db.IsNotFound(err)) + + // Second call should return null hit + value, hit, err := c.SWR(ctx, "key3", func(ctx context.Context) (string, error) { + t.Fatal("should not call refresh function") + return "", nil + }, func(err error) cache.Op { + return cache.WriteValue + }) + + require.NoError(t, err) + require.Equal(t, "", value) + require.Equal(t, cache.Null, hit) + }) + + t.Run("stale hit returns cached value", func(t *testing.T) { + // First call to populate cache + _, _, err := c.SWR(ctx, "key4", func(ctx context.Context) (string, error) { + return "value4", nil + }, func(err error) cache.Op { + return cache.WriteValue + }) + require.NoError(t, err) + + // Move time forward past fresh but within stale + mockClock.Tick(2 * time.Minute) + + // Should return cached value with hit status + value, hit, err := c.SWR(ctx, "key4", func(ctx context.Context) (string, error) { + // This will be called in background + return "updated_value4", nil + }, func(err error) cache.Op { + return cache.WriteValue + }) + + require.NoError(t, err) + require.Equal(t, "value4", value) + require.Equal(t, cache.Hit, hit) + }) + + t.Run("miss after stale time", func(t *testing.T) { + // First call to populate cache + _, _, err := c.SWR(ctx, "key5", func(ctx context.Context) (string, error) { + return "value5", nil + }, func(err error) cache.Op { + return cache.WriteValue + }) + require.NoError(t, err) + + // Move time forward past stale + mockClock.Tick(6 * time.Minute) + + // Should call refresh and return new value + value, hit, err := c.SWR(ctx, "key5", func(ctx context.Context) (string, error) { + return "new_value5", nil + }, func(err error) cache.Op { + return cache.WriteValue + }) + + require.NoError(t, err) + require.Equal(t, "new_value5", value) + require.Equal(t, cache.Hit, hit) + }) + + t.Run("error returns miss", func(t *testing.T) { + expectedErr := errors.New("refresh error") + value, hit, err := c.SWR(ctx, "key6", func(ctx context.Context) (string, error) { + return "", expectedErr + }, func(err error) cache.Op { + return cache.Noop + }) + + require.Error(t, err) + require.Equal(t, expectedErr, err) + require.Equal(t, "", value) + require.Equal(t, cache.Miss, hit) + }) +} From e52df0b73db7c36aec77d638149cd7acfbf0792b Mon Sep 17 00:00:00 2001 From: Flo <53355483+Flo4604@users.noreply.github.com> Date: Mon, 21 Jul 2025 17:35:24 +0200 Subject: [PATCH 2/3] remove log --- go/internal/services/keys/get.go | 3 --- 1 file changed, 3 deletions(-) diff --git a/go/internal/services/keys/get.go b/go/internal/services/keys/get.go index a07c4d6fb66..92bc1ec0b22 100644 --- a/go/internal/services/keys/get.go +++ b/go/internal/services/keys/get.go @@ -4,7 +4,6 @@ import ( "context" "encoding/json" "fmt" - "log" "time" "github.com/unkeyed/unkey/go/internal/services/caches" @@ -89,8 +88,6 @@ func (s *service) Get(ctx context.Context, sess *zen.Session, rawKey string) (*K }, nil } - log.Printf("Found key %#v", key) - // ForWorkspace set but that doesn't exist if key.ForWorkspaceID.Valid && !key.ForWorkspaceEnabled.Valid { // nolint:exhaustruct From f82797c31495e6402815617ccefb27359e5292b2 Mon Sep 17 00:00:00 2001 From: Flo <53355483+Flo4604@users.noreply.github.com> Date: Mon, 21 Jul 2025 18:59:20 +0200 Subject: [PATCH 3/3] fix logic --- go/apps/api/routes/v2_ratelimit_limit/handler.go | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/go/apps/api/routes/v2_ratelimit_limit/handler.go b/go/apps/api/routes/v2_ratelimit_limit/handler.go index e45a706c827..92926de5187 100644 --- a/go/apps/api/routes/v2_ratelimit_limit/handler.go +++ b/go/apps/api/routes/v2_ratelimit_limit/handler.go @@ -121,12 +121,10 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { } if hit == cache.Null { - if db.IsNotFound(err) { - return fault.New("namespace cache null", - fault.Code(codes.Data.RatelimitNamespace.NotFound.URN()), - fault.Public("This namespace does not exist."), - ) - } + return fault.New("namespace cache null", + fault.Code(codes.Data.RatelimitNamespace.NotFound.URN()), + fault.Public("This namespace does not exist."), + ) } if namespace.DeletedAtM.Valid {