Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)
})
}
11 changes: 10 additions & 1 deletion go/apps/api/routes/v2_ratelimit_limit/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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."),
)
}
}
Comment thread
Flo4604 marked this conversation as resolved.

if namespace.DeletedAtM.Valid {
return fault.New("namespace was deleted",
fault.Code(codes.Data.RatelimitNamespace.NotFound.URN()),
Expand Down
11 changes: 10 additions & 1 deletion go/internal/services/keys/get.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (

"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"
Expand Down Expand Up @@ -60,7 +61,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 {
Expand All @@ -79,6 +80,14 @@ 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
}

// ForWorkspace set but that doesn't exist
if key.ForWorkspaceID.Valid && !key.ForWorkspaceEnabled.Valid {
// nolint:exhaustruct
Expand Down
24 changes: 20 additions & 4 deletions go/pkg/cache/cache.go
Original file line number Diff line number Diff line change
Expand Up @@ -239,15 +239,15 @@ 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 {
// Cache Hit

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) {
Expand All @@ -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
Expand All @@ -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
}
8 changes: 4 additions & 4 deletions go/pkg/cache/cache_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion go/pkg/cache/interface.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
7 changes: 3 additions & 4 deletions go/pkg/cache/middleware/tracing.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
4 changes: 2 additions & 2 deletions go/pkg/cache/noop.go
Original file line number Diff line number Diff line change
Expand Up @@ -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] {
Expand Down
2 changes: 1 addition & 1 deletion go/pkg/cache/simulation_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
157 changes: 157 additions & 0 deletions go/pkg/cache/swr_test.go
Original file line number Diff line number Diff line change
@@ -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)
})
}