inmemory return ttlExpiryTime in GetResponse (#2870)

Signed-off-by: joshvanl <me@joshvanl.dev>
Signed-off-by: Alessandro (Ale) Segala <43508+ItalyPaleAle@users.noreply.github.com>
Co-authored-by: Alessandro (Ale) Segala <43508+ItalyPaleAle@users.noreply.github.com>
This commit is contained in:
Josh van Leeuwen 2023-06-06 16:28:08 +01:00 committed by GitHub
parent ff06ae6b33
commit 638b106bdb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 90 additions and 11 deletions

View File

@ -24,6 +24,7 @@ import (
"time"
"github.com/google/uuid"
"k8s.io/utils/clock"
"github.com/dapr/components-contrib/state"
"github.com/dapr/components-contrib/state/utils"
@ -37,6 +38,7 @@ type inMemoryStore struct {
items map[string]*inMemStateStoreItem
lock sync.RWMutex
log logger.Logger
clock clock.Clock
closeCh chan struct{}
closed atomic.Bool
wg sync.WaitGroup
@ -51,6 +53,7 @@ func newStateStore(log logger.Logger) *inMemoryStore {
items: map[string]*inMemStateStoreItem{},
log: log,
closeCh: make(chan struct{}),
clock: clock.RealClock{},
}
s.BulkStore = state.NewDefaultBulkStore(s)
return s
@ -146,7 +149,7 @@ func (store *inMemoryStore) Get(ctx context.Context, req *state.GetRequest) (*st
store.lock.RLock()
item := store.items[req.Key]
store.lock.RUnlock()
if item != nil && item.isExpired() {
if item != nil && item.isExpired(store.clock.Now()) {
store.lock.Lock()
item = store.getAndExpire(req.Key)
store.lock.Unlock()
@ -156,7 +159,14 @@ func (store *inMemoryStore) Get(ctx context.Context, req *state.GetRequest) (*st
return &state.GetResponse{}, nil
}
return &state.GetResponse{Data: item.data, ETag: item.etag}, nil
var metadata map[string]string
if item.expire != nil {
metadata = map[string]string{
state.GetRespMetaKeyTTLExpireTime: item.expire.UTC().Format(time.RFC3339),
}
}
return &state.GetResponse{Data: item.data, ETag: item.etag, Metadata: metadata}, nil
}
func (store *inMemoryStore) BulkGet(ctx context.Context, req []state.GetRequest, _ state.BulkGetOpts) ([]state.BulkGetResponse, error) {
@ -171,12 +181,18 @@ func (store *inMemoryStore) BulkGet(ctx context.Context, req []state.GetRequest,
for i, r := range req {
item := store.items[r.Key]
if item != nil && !item.isExpired() {
if item != nil && !item.isExpired(store.clock.Now()) {
res[i] = state.BulkGetResponse{
Key: r.Key,
Data: item.data,
ETag: item.etag,
}
if item.expire != nil {
res[i].Metadata = map[string]string{
state.GetRespMetaKeyTTLExpireTime: item.expire.UTC().Format(time.RFC3339),
}
}
} else {
res[i] = state.BulkGetResponse{
Key: r.Key,
@ -193,7 +209,7 @@ func (store *inMemoryStore) getAndExpire(key string) *inMemStateStoreItem {
if item == nil {
return nil
}
if item.isExpired() {
if item.isExpired(store.clock.Now()) {
delete(store.items, key)
return nil
}
@ -280,7 +296,7 @@ func (store *inMemoryStore) doSet(ctx context.Context, key string, data []byte,
etag: &etag,
}
if ttlInSeconds > 0 {
el.expire = ptr.Of(time.Now().UnixMilli() + int64(ttlInSeconds)*1000)
el.expire = ptr.Of(store.clock.Now().Add(time.Duration(ttlInSeconds) * time.Second))
}
store.items[key] = el
@ -388,7 +404,7 @@ func (store *inMemoryStore) doCleanExpiredItems() {
defer store.lock.Unlock()
for key, item := range store.items {
if item.expire != nil && item.isExpired() {
if item.expire != nil && item.isExpired(store.clock.Now()) {
store.doDelete(context.Background(), key)
}
}
@ -402,12 +418,12 @@ func (store *inMemoryStore) GetComponentMetadata() map[string]string {
type inMemStateStoreItem struct {
data []byte
etag *string
expire *int64
expire *time.Time
}
func (item *inMemStateStoreItem) isExpired() bool {
func (item *inMemStateStoreItem) isExpired(now time.Time) bool {
if item == nil || item.expire == nil {
return false
}
return time.Now().UnixMilli() > *item.expire
return now.After(*item.expire)
}

View File

@ -15,11 +15,14 @@ package inmemory
import (
"context"
"sort"
"testing"
"time"
"github.com/golang/mock/gomock"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
clocktesting "k8s.io/utils/clock/testing"
"github.com/dapr/components-contrib/state"
"github.com/dapr/kit/logger"
@ -30,7 +33,9 @@ func TestReadAndWrite(t *testing.T) {
defer ctl.Finish()
store := NewInMemoryStateStore(logger.NewLogger("test"))
store := NewInMemoryStateStore(logger.NewLogger("test")).(*inMemoryStore)
fakeClock := clocktesting.NewFakeClock(time.Now())
store.clock = fakeClock
store.Init(context.Background(), state.Metadata{})
keyA := "theFirstKey"
@ -65,7 +70,7 @@ func TestReadAndWrite(t *testing.T) {
err := store.Set(context.Background(), setReq)
assert.NoError(t, err)
// simulate expiration
time.Sleep(2 * time.Second)
fakeClock.Step(2 * time.Second)
// get
getReq := &state.GetRequest{
Key: keyA,
@ -77,6 +82,64 @@ func TestReadAndWrite(t *testing.T) {
assert.Nil(t, resp.ETag)
})
t.Run("return expire time when ttlInSeconds set with Get", func(t *testing.T) {
now := fakeClock.Now()
// set with LWW
setReq := &state.SetRequest{
Key: keyA,
Value: valueA,
Metadata: map[string]string{"ttlInSeconds": "1000"},
}
err := store.Set(context.Background(), setReq)
assert.NoError(t, err)
// get
getReq := &state.GetRequest{
Key: keyA,
}
resp, err := store.Get(context.Background(), getReq)
assert.NoError(t, err)
assert.NotNil(t, resp)
assert.Equal(t, `"value of key"`, string(resp.Data))
assert.Len(t, resp.Metadata, 1)
require.Contains(t, resp.Metadata, "ttlExpireTime")
assert.Equal(t, now.Add(time.Second*1000).UTC().Format(time.RFC3339), resp.Metadata["ttlExpireTime"])
})
t.Run("return expire time when ttlInSeconds set with GetBulk", func(t *testing.T) {
assert.NoError(t, store.Set(context.Background(), &state.SetRequest{
Key: "a",
Value: "123",
Metadata: map[string]string{"ttlInSeconds": "1000"},
}))
assert.NoError(t, store.Set(context.Background(), &state.SetRequest{
Key: "b",
Value: "456",
Metadata: map[string]string{"ttlInSeconds": "2001"},
}))
resp, err := store.BulkGet(context.Background(), []state.GetRequest{
{Key: "a"},
{Key: "b"},
}, state.BulkGetOpts{})
assert.NoError(t, err)
assert.NotNil(t, resp)
require.Len(t, resp, 2)
sort.Slice(resp, func(i, j int) bool {
return resp[i].Key < resp[j].Key
})
assert.Equal(t, `"123"`, string(resp[0].Data))
assert.Equal(t, `"456"`, string(resp[1].Data))
assert.Len(t, resp[0].Metadata, 1)
require.Contains(t, resp[0].Metadata, "ttlExpireTime")
assert.Len(t, resp[1].Metadata, 1)
require.Contains(t, resp[1].Metadata, "ttlExpireTime")
assert.Equal(t, fakeClock.Now().Add(time.Second*1000).UTC().Format(time.RFC3339), resp[0].Metadata["ttlExpireTime"])
assert.Equal(t, fakeClock.Now().Add(time.Second*2001).UTC().Format(time.RFC3339), resp[1].Metadata["ttlExpireTime"])
})
t.Run("set and get the second key successfully", func(t *testing.T) {
// set
setReq := &state.SetRequest{