refactor: migrate the redis command keys to scan (#18825)

Refine the cache interface, migrate the Keys to Scan, change the redis underlying keys command to scan. Signed-off-by: chlins <[email protected]>
goharbor · Jun 29, 2023 · 02a1c41 · 02a1c41
1 parent d36ca80
commit 02a1c41
Show file tree

Hide file tree

Showing 23 changed files with 615 additions and 162 deletions.
diff --git a/src/go.mod b/src/go.mod
@@ -47,7 +47,6 @@ require (
 	github.com/opencontainers/go-digest v1.0.0
 	github.com/opencontainers/image-spec v1.1.0-rc2.0.20221005185240-3a7f492d3f1b
 	github.com/pkg/errors v0.9.1
-
 	github.com/prometheus/client_golang v1.14.0
 	github.com/robfig/cron/v3 v3.0.0
 	github.com/spf13/viper v1.8.1

diff --git a/src/lib/cache/cache.go b/src/lib/cache/cache.go
@@ -40,6 +40,14 @@ var (
 	ErrNotFound = errors.New("key not found")
 )
 
+// Iterator returns the ScanIterator
+type Iterator interface {
+	Next(ctx context.Context) bool
+	Val() string
+}
+
+//go:generate mockery --name Cache --output . --outpkg cache --filename mock_cache_test.go --structname mockCache --inpackage
+
 // Cache cache interface
 type Cache interface {
 	// Contains returns true if key exists
@@ -57,8 +65,9 @@ type Cache interface {
 	// Save cache the value by key
 	Save(ctx context.Context, key string, value interface{}, expiration ...time.Duration) error
 
-	// Keys returns the key matched by prefixes
-	Keys(ctx context.Context, prefixes ...string) ([]string, error)
+	// Scan scans the keys matched by match string
+	// NOTICE: memory cache does not support use wildcard, compared by strings.Contains
+	Scan(ctx context.Context, match string) (Iterator, error)
 }
 
 var (

diff --git a/src/lib/cache/cache_test.go b/src/lib/cache/cache_test.go
@@ -18,10 +18,9 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/goharbor/harbor/src/lib/retry"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
-
-	cachetesting "github.com/goharbor/harbor/src/testing/lib/cache"
-	"github.com/goharbor/harbor/src/testing/mock"
 )
 
 type CacheTestSuite struct {
@@ -30,7 +29,7 @@ type CacheTestSuite struct {
 
 func (suite *CacheTestSuite) SetupSuite() {
 	Register("mock", func(opts Options) (Cache, error) {
-		return &cachetesting.Cache{}, nil
+		return &mockCache{}, nil
 	})
 }
 
@@ -62,8 +61,8 @@ func (suite *CacheTestSuite) TestInitialize() {
 
 	{
 		Register("cache", func(opts Options) (Cache, error) {
-			c := &cachetesting.Cache{}
-			c.On("Ping", mock.Anything).Return(fmt.Errorf("oops"))
+			c := &mockCache{}
+			c.On("Ping", mock.Anything).Return(retry.Abort(fmt.Errorf("oops")))
 
 			return c, nil
 		})
@@ -75,7 +74,7 @@ func (suite *CacheTestSuite) TestInitialize() {
 
 	{
 		Register("cache", func(opts Options) (Cache, error) {
-			c := &cachetesting.Cache{}
+			c := &mockCache{}
 			c.On("Ping", mock.Anything).Return(nil)
 
 			return c, nil

diff --git a/src/lib/cache/helper_test.go b/src/lib/cache/helper_test.go
@@ -23,7 +23,6 @@ import (
 
 	"github.com/stretchr/testify/suite"
 
-	cachetesting "github.com/goharbor/harbor/src/testing/lib/cache"
 	"github.com/goharbor/harbor/src/testing/mock"
 )
 
@@ -42,7 +41,7 @@ func (suite *FetchOrSaveTestSuite) SetupSuite() {
 }
 
 func (suite *FetchOrSaveTestSuite) TestFetchInternalError() {
-	c := &cachetesting.Cache{}
+	c := &mockCache{}
 
 	mock.OnAnything(c, "Fetch").Return(fmt.Errorf("oops"))
 
@@ -55,7 +54,7 @@ func (suite *FetchOrSaveTestSuite) TestFetchInternalError() {
 }
 
 func (suite *FetchOrSaveTestSuite) TestBuildError() {
-	c := &cachetesting.Cache{}
+	c := &mockCache{}
 
 	mock.OnAnything(c, "Fetch").Return(ErrNotFound)
 
@@ -68,7 +67,7 @@ func (suite *FetchOrSaveTestSuite) TestBuildError() {
 }
 
 func (suite *FetchOrSaveTestSuite) TestSaveError() {
-	c := &cachetesting.Cache{}
+	c := &mockCache{}
 
 	mock.OnAnything(c, "Fetch").Return(ErrNotFound)
 	mock.OnAnything(c, "Save").Return(fmt.Errorf("oops"))
@@ -83,7 +82,7 @@ func (suite *FetchOrSaveTestSuite) TestSaveError() {
 }
 
 func (suite *FetchOrSaveTestSuite) TestSaveCalledOnlyOneTime() {
-	c := &cachetesting.Cache{}
+	c := &mockCache{}
 
 	var data sync.Map
 

diff --git a/src/lib/cache/memory/memory.go b/src/lib/cache/memory/memory.go
@@ -117,27 +117,55 @@ func (c *Cache) Save(ctx context.Context, key string, value interface{}, expirat
 	return nil
 }
 
-// Keys returns the key matched by prefixes.
-func (c *Cache) Keys(ctx context.Context, prefixes ...string) ([]string, error) {
-	// if no prefix, means match all keys.
-	matchAll := len(prefixes) == 0
-	// range map to get all keys
-	keys := make([]string, 0)
+// Scan scans the keys matched by match string
+func (c *Cache) Scan(ctx context.Context, match string) (cache.Iterator, error) {
+	var keys []string
 	c.storage.Range(func(k, v interface{}) bool {
-		ks := k.(string)
-		if matchAll {
-			keys = append(keys, ks)
-		} else {
-			for _, p := range prefixes {
-				if strings.HasPrefix(ks, c.opts.Key(p)) {
-					keys = append(keys, strings.TrimPrefix(ks, c.opts.Prefix))
-				}
+		matched := true
+		if match != "" {
+			matched = strings.Contains(k.(string), match)
+		}
+
+		if matched {
+			if v.(*entry).isExpirated() {
+				c.storage.Delete(k)
+			} else {
+				keys = append(keys, strings.TrimPrefix(k.(string), c.opts.Prefix))
 			}
 		}
 		return true
 	})
 
-	return keys, nil
+	return &ScanIterator{keys: keys}, nil
+}
+
+// ScanIterator is a ScanIterator for memory cache
+type ScanIterator struct {
+	mu   sync.Mutex
+	pos  int
+	keys []string
+}
+
+// Next checks whether has the next element
+func (i *ScanIterator) Next(ctx context.Context) bool {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
+	i.pos++
+	return i.pos <= len(i.keys)
+}
+
+// Val returns the key
+func (i *ScanIterator) Val() string {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
+	var val string
+	if i.pos <= len(i.keys) {
+		val = i.keys[i.pos-1]
+	}
+
+	return val
 }
 
 // New returns memory cache

diff --git a/src/lib/cache/memory/memory_test.go b/src/lib/cache/memory/memory_test.go
@@ -16,6 +16,7 @@ package memory
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -109,28 +110,54 @@ func (suite *CacheTestSuite) TestPing() {
 	suite.NoError(suite.cache.Ping(suite.ctx))
 }
 
-func (suite *CacheTestSuite) TestKeys() {
-	key1 := "p1"
-	key2 := "p2"
-
-	var err error
-	err = suite.cache.Save(suite.ctx, key1, "hello, p1")
-	suite.Nil(err)
-	err = suite.cache.Save(suite.ctx, key2, "hello, p2")
-	suite.Nil(err)
-
-	// should match all
-	keys, err := suite.cache.Keys(suite.ctx, "p")
-	suite.Nil(err)
-	suite.ElementsMatch([]string{"p1", "p2"}, keys)
-	// only get p1
-	keys, err = suite.cache.Keys(suite.ctx, key1)
-	suite.Nil(err)
-	suite.Equal([]string{"p1"}, keys)
-	// only get p2
-	keys, err = suite.cache.Keys(suite.ctx, key2)
-	suite.Nil(err)
-	suite.Equal([]string{"p2"}, keys)
+func (suite *CacheTestSuite) TestScan() {
+	seed := func(n int) {
+		for i := 0; i < n; i++ {
+			key := fmt.Sprintf("test-scan-%d", i)
+			err := suite.cache.Save(suite.ctx, key, "")
+			suite.NoError(err)
+		}
+	}
+	clean := func(n int) {
+		for i := 0; i < n; i++ {
+			key := fmt.Sprintf("test-scan-%d", i)
+			err := suite.cache.Delete(suite.ctx, key)
+			suite.NoError(err)
+		}
+	}
+	{
+		// no match should return all keys
+		expect := []string{"test-scan-0", "test-scan-1", "test-scan-2"}
+		// seed data
+		seed(3)
+		// test scan
+		iter, err := suite.cache.Scan(suite.ctx, "")
+		suite.NoError(err)
+		got := []string{}
+		for iter.Next(suite.ctx) {
+			got = append(got, iter.Val())
+		}
+		suite.ElementsMatch(expect, got)
+		// clean up
+		clean(3)
+	}
+
+	{
+		// with match should return matched keys
+		expect := []string{"test-scan-1", "test-scan-10"}
+		// seed data
+		seed(11)
+		// test scan
+		iter, err := suite.cache.Scan(suite.ctx, "test-scan-1")
+		suite.NoError(err)
+		got := []string{}
+		for iter.Next(suite.ctx) {
+			got = append(got, iter.Val())
+		}
+		suite.ElementsMatch(expect, got)
+		// clean up
+		clean(11)
+	}
 }
 
 func TestCacheTestSuite(t *testing.T) {