Merge pull request #1053 from gofiber/memory-keys

Add support for Keys() to Memory Driver

memory/README.md

@@ -31,6 +31,7 @@ func (s *Storage) Delete(key string) error
 func (s *Storage) Reset() error
 func (s *Storage) Close() error
 func (s *Storage) Conn() map[string]entry
+func (s *Storage) Keys() ([][]byte, error)
 ```
 
 ### Installation

memory/go.mod

@@ -2,14 +2,10 @@ module github.com/gofiber/storage/memory/v2
 
 go 1.19
 
-require (
-	github.com/gofiber/utils/v2 v2.0.0-beta.3
-	github.com/stretchr/testify v1.8.4
-)
+require github.com/stretchr/testify v1.8.4
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/google/uuid v1.3.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

memory/go.sum

@@ -1,9 +1,5 @@
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/gofiber/utils/v2 v2.0.0-beta.3 h1:pfOhUDDVjBJpkWv6C5jaDyYLvpui7zQ97zpyFFsUOKw=
-github.com/gofiber/utils/v2 v2.0.0-beta.3/go.mod h1:jsl17+MsKfwJjM3ONCE9Rzji/j8XNbwjhUVTjzgfDCo=
-github.com/google/uuid v1.3.1 h1:KjJaJ9iWZ3jOFZIf1Lqf4laDRCasjl0BCmnEGxkdLb4=
-github.com/google/uuid v1.3.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=

memory/memory.go

@@ -141,3 +141,29 @@ func (s *Storage) Conn() map[string]entry {
 	defer s.mux.RUnlock()
 	return s.db
 }
+
+// Return all the keys
+func (s *Storage) Keys() ([][]byte, error) {
+	s.mux.RLock()
+	defer s.mux.RUnlock()
+
+	if len(s.db) == 0 {
+		return nil, nil
+	}
+
+	ts := atomic.LoadUint32(&internal.Timestamp)
+	keys := make([][]byte, 0, len(s.db))
+	for key, v := range s.db {
+		// Filter out the expired keys
+		if v.expiry == 0 || v.expiry > ts {
+			keys = append(keys, []byte(key))
+		}
+	}
+
+	// Double check if no valid keys were found
+	if len(keys) == 0 {
+		return nil, nil
+	}
+
+	return keys, nil
+}

memory/memory_test.go

@@ -7,35 +7,44 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-var testStore = New()
-
 func Test_Storage_Memory_Set(t *testing.T) {
 	var (
-		key = "john"
-		val = []byte("doe")
+		testStore = New()
+		key       = "john"
+		val       = []byte("doe")
 	)
 
 	err := testStore.Set(key, val, 0)
 	require.NoError(t, err)
+
+	keys, err := testStore.Keys()
+	require.NoError(t, err)
+	require.Len(t, keys, 1)
 }
 
 func Test_Storage_Memory_Set_Override(t *testing.T) {
 	var (
-		key = "john"
-		val = []byte("doe")
+		testStore = New()
+		key       = "john"
+		val       = []byte("doe")
 	)
 
 	err := testStore.Set(key, val, 0)
 	require.NoError(t, err)
 
 	err = testStore.Set(key, val, 0)
 	require.NoError(t, err)
+
+	keys, err := testStore.Keys()
+	require.NoError(t, err)
+	require.Len(t, keys, 1)
 }
 
 func Test_Storage_Memory_Get(t *testing.T) {
 	var (
-		key = "john"
-		val = []byte("doe")
+		testStore = New()
+		key       = "john"
+		val       = []byte("doe")
 	)
 
 	err := testStore.Set(key, val, 0)
@@ -44,53 +53,104 @@ func Test_Storage_Memory_Get(t *testing.T) {
 	result, err := testStore.Get(key)
 	require.NoError(t, err)
 	require.Equal(t, val, result)
+
+	keys, err := testStore.Keys()
+	require.NoError(t, err)
+	require.Len(t, keys, 1)
 }
 
 func Test_Storage_Memory_Set_Expiration(t *testing.T) {
 	var (
-		key = "john"
-		val = []byte("doe")
-		exp = 1 * time.Second
+		testStore = New()
+		key       = "john"
+		val       = []byte("doe")
+		exp       = 1 * time.Second
 	)
 
 	err := testStore.Set(key, val, exp)
 	require.NoError(t, err)
 
 	time.Sleep(1100 * time.Millisecond)
+
+	result, err := testStore.Get(key)
+	require.NoError(t, err)
+	require.Zero(t, len(result))
+
+	keys, err := testStore.Keys()
+	require.NoError(t, err)
+	require.Nil(t, keys)
 }
 
-func Test_Storage_Memory_Get_Expired(t *testing.T) {
-	key := "john"
+func Test_Storage_Memory_Set_Long_Expiration_with_Keys(t *testing.T) {
+	var (
+		testStore = New()
+		key       = "john"
+		val       = []byte("doe")
+		exp       = 5 * time.Second
+	)
+
+	keys, err := testStore.Keys()
+	require.NoError(t, err)
+	require.Nil(t, keys)
 
+	err = testStore.Set(key, val, exp)
+	require.NoError(t, err)
+
+	time.Sleep(1100 * time.Millisecond)
+
+	keys, err = testStore.Keys()
+	require.NoError(t, err)
+	require.Len(t, keys, 1)
+
+	time.Sleep(4000 * time.Millisecond)
 	result, err := testStore.Get(key)
 	require.NoError(t, err)
 	require.Zero(t, len(result))
+
+	keys, err = testStore.Keys()
+	require.NoError(t, err)
+	require.Nil(t, keys)
 }
 
 func Test_Storage_Memory_Get_NotExist(t *testing.T) {
+	testStore := New()
 	result, err := testStore.Get("notexist")
 	require.NoError(t, err)
 	require.Zero(t, len(result))
+
+	keys, err := testStore.Keys()
+	require.NoError(t, err)
+	require.Nil(t, keys)
 }
 
 func Test_Storage_Memory_Delete(t *testing.T) {
 	var (
-		key = "john"
-		val = []byte("doe")
+		testStore = New()
+		key       = "john"
+		val       = []byte("doe")
 	)
 
 	err := testStore.Set(key, val, 0)
 	require.NoError(t, err)
 
+	keys, err := testStore.Keys()
+	require.NoError(t, err)
+	require.Len(t, keys, 1)
+
 	err = testStore.Delete(key)
 	require.NoError(t, err)
 
 	result, err := testStore.Get(key)
 	require.NoError(t, err)
 	require.Zero(t, len(result))
+
+	keys, err = testStore.Keys()
+	require.NoError(t, err)
+	require.Nil(t, keys)
 }
 
 func Test_Storage_Memory_Reset(t *testing.T) {
+	testStore := New()
 	val := []byte("doe")
 
 	err := testStore.Set("john1", val, 0)
@@ -99,6 +159,10 @@ func Test_Storage_Memory_Reset(t *testing.T) {
 	err = testStore.Set("john2", val, 0)
 	require.NoError(t, err)
 
+	keys, err := testStore.Keys()
+	require.NoError(t, err)
+	require.Len(t, keys, 2)
+
 	err = testStore.Reset()
 	require.NoError(t, err)
 
@@ -109,17 +173,24 @@ func Test_Storage_Memory_Reset(t *testing.T) {
 	result, err = testStore.Get("john2")
 	require.NoError(t, err)
 	require.Zero(t, len(result))
+
+	keys, err = testStore.Keys()
+	require.NoError(t, err)
+	require.Nil(t, keys)
 }
 
 func Test_Storage_Memory_Close(t *testing.T) {
+	testStore := New()
 	require.Nil(t, testStore.Close())
 }
 
 func Test_Storage_Memory_Conn(t *testing.T) {
+	testStore := New()
 	require.True(t, testStore.Conn() != nil)
 }
 
 func Benchmark_Memory_Set(b *testing.B) {
+	testStore := New()
 	b.ReportAllocs()
 	b.ResetTimer()
 
@@ -132,6 +203,7 @@ func Benchmark_Memory_Set(b *testing.B) {
 }
 
 func Benchmark_Memory_Get(b *testing.B) {
+	testStore := New()
 	err := testStore.Set("john", []byte("doe"), 0)
 	require.NoError(b, err)
 
@@ -146,6 +218,7 @@ func Benchmark_Memory_Get(b *testing.B) {
 }
 
 func Benchmark_Memory_SetAndDelete(b *testing.B) {
+	testStore := New()
 	b.ReportAllocs()
 	b.ResetTimer()