Skip to content

Commit 02668ec

Browse files
berndverstberndverst
committed
Adds support for TTL in Redis State Store
1 parent b1f584d commit 02668ec

File tree

2 files changed

+189
-6
lines changed

2 files changed

+189
-6
lines changed

‎state/redis/redis.go‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ const (
3030
infoReplicationDelimiter = "\r\n"
3131
maxRetries = "maxRetries"
3232
maxRetryBackoff = "maxRetryBackoff"
33+
ttlInSeconds = "ttlInSeconds"
3334
defaultBase = 10
3435
defaultBitSize = 0
3536
defaultDB = 0
@@ -230,6 +231,7 @@ func (r *StateStore) setValue(req *state.SetRequest) error {
230231
if err != nil {
231232
return err
232233
}
234+
ttl := r.parseTTL(req)
233235

234236
bt, _ := utils.Marshal(req.Value, r.json.Marshal)
235237

@@ -241,6 +243,19 @@ func (r *StateStore) setValue(req *state.SetRequest) error {
241243

242244
return fmt.Errorf("failed to set key %s: %s", req.Key, err)
243245
}
246+
switch {
247+
case ttl == -1:
248+
_, err = r.client.Do(r.ctx, "PERSIST", req.Key).Result()
249+
if err != nil {
250+
return fmt.Errorf("failed to persist key %s: %s", req.Key, err)
251+
}
252+
253+
case ttl > 0:
254+
_, err = r.client.Do(r.ctx, "EXPIRE", req.Key, ttl).Result()
255+
if err != nil {
256+
return fmt.Errorf("failed to set key %s ttl: %s", req.Key, err)
257+
}
258+
}
244259

245260
if req.Options.Consistency == state.Strong && r.replicas > 0 {
246261
_, err = r.client.Do(r.ctx, "WAIT", r.replicas, 1000).Result()
@@ -264,11 +279,18 @@ func (r *StateStore) Multi(request *state.TransactionalStateRequest) error {
264279
if o.Operation == state.Upsert {
265280
req := o.Request.(state.SetRequest)
266281
ver, err := r.parseETag(&req)
282+
ttl := r.parseTTL(&req)
267283
if err != nil {
268284
return err
269285
}
270286
bt, _ := utils.Marshal(req.Value, r.json.Marshal)
271287
pipe.Do(r.ctx, "EVAL", setQuery, 1, req.Key, ver, bt)
288+
switch {
289+
case ttl == -1:
290+
pipe.Do(r.ctx, "PERSIST", req.Key)
291+
case ttl > 0:
292+
pipe.Do(r.ctx, "EXPIRE", req.Key, ttl)
293+
}
272294
} else if o.Operation == state.Delete {
273295
req := o.Request.(state.DeleteRequest)
274296
if req.ETag == nil {
@@ -318,6 +340,19 @@ func (r *StateStore) parseETag(req *state.SetRequest) (int, error) {
318340
return ver, nil
319341
}
320342

343+
func (r *StateStore) parseTTL(req *state.SetRequest) int {
344+
if val, ok := req.Metadata[ttlInSeconds]; ok && val != "" {
345+
parsedVal, err := strconv.ParseInt(val, defaultBase, defaultBitSize)
346+
if err != nil {
347+
return 0
348+
}
349+
350+
return int(parsedVal)
351+
}
352+
353+
return 0
354+
}
355+
321356
func (r *StateStore) Close() error {
322357
r.cancel()
323358

‎state/redis/redis_test.go‎

Lines changed: 154 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,9 @@ package redis
77

88
import (
99
"context"
10+
"strconv"
1011
"testing"
12+
"time"
1113

1214
"github.com/agrea/ptr"
1315
miniredis "github.com/alicebob/miniredis/v2"
@@ -90,6 +92,48 @@ func TestParseEtag(t *testing.T) {
9092
})
9193
}
9294

95+
func TestParseTTL(t *testing.T) {
96+
store := NewRedisStateStore(logger.NewLogger("test"))
97+
t.Run("TTL Not an integer", func(t *testing.T) {
98+
ttlInSeconds := "not an integer"
99+
ttl := store.parseTTL(&state.SetRequest{
100+
Metadata: map[string]string{
101+
"ttlInSeconds": ttlInSeconds,
102+
},
103+
})
104+
assert.Equal(t, ttl, 0)
105+
})
106+
t.Run("TTL specified with wrong key", func(t *testing.T) {
107+
ttlInSeconds := 12345
108+
ttl := store.parseTTL(&state.SetRequest{
109+
Metadata: map[string]string{
110+
"expirationTime": strconv.Itoa(ttlInSeconds),
111+
},
112+
})
113+
assert.Equal(t, ttl, 0)
114+
})
115+
t.Run("TTL is a number", func(t *testing.T) {
116+
ttlInSeconds := 12345
117+
ttl := store.parseTTL(&state.SetRequest{
118+
Metadata: map[string]string{
119+
"ttlInSeconds": strconv.Itoa(ttlInSeconds),
120+
},
121+
})
122+
assert.Equal(t, ttl, ttlInSeconds)
123+
})
124+
125+
t.Run("TTL never expires", func(t *testing.T) {
126+
ttlInSeconds := -1
127+
ttl := store.parseTTL(&state.SetRequest{
128+
Metadata: map[string]string{
129+
"ttlInSeconds": strconv.Itoa(ttlInSeconds),
130+
},
131+
})
132+
assert.Equal(t, ttl, ttlInSeconds)
133+
})
134+
135+
}
136+
93137
func TestParseConnectedSlavs(t *testing.T) {
94138
store := NewRedisStateStore(logger.NewLogger("test"))
95139

@@ -126,13 +170,35 @@ func TestTransactionalUpsert(t *testing.T) {
126170
ss.ctx, ss.cancel = context.WithCancel(context.Background())
127171

128172
err := ss.Multi(&state.TransactionalStateRequest{
129-
Operations: []state.TransactionalStateOperation{{
130-
Operation: state.Upsert,
131-
Request: state.SetRequest{
132-
Key: "weapon",
133-
Value: "deathstar",
173+
Operations: []state.TransactionalStateOperation{
174+
{
175+
Operation: state.Upsert,
176+
Request: state.SetRequest{
177+
Key: "weapon",
178+
Value: "deathstar",
179+
},
134180
},
135-
}},
181+
{
182+
Operation: state.Upsert,
183+
Request: state.SetRequest{
184+
Key: "weapon2",
185+
Value: "deathstar2",
186+
Metadata: map[string]string{
187+
"ttlInSeconds": "123",
188+
},
189+
},
190+
},
191+
{
192+
Operation: state.Upsert,
193+
Request: state.SetRequest{
194+
Key: "weapon3",
195+
Value: "deathstar3",
196+
Metadata: map[string]string{
197+
"ttlInSeconds": "-1",
198+
},
199+
},
200+
},
201+
},
136202
})
137203
assert.Equal(t, nil, err)
138204

@@ -144,6 +210,18 @@ func TestTransactionalUpsert(t *testing.T) {
144210
assert.Equal(t, nil, err)
145211
assert.Equal(t, ptr.String("1"), version)
146212
assert.Equal(t, `"deathstar"`, data)
213+
214+
res, err = c.Do(context.Background(), "TTL", "weapon").Result()
215+
assert.Equal(t, nil, err)
216+
assert.Equal(t, int64(-1), res)
217+
218+
res, err = c.Do(context.Background(), "TTL", "weapon2").Result()
219+
assert.Equal(t, nil, err)
220+
assert.Equal(t, int64(123), res)
221+
222+
res, err = c.Do(context.Background(), "TTL", "weapon3").Result()
223+
assert.Equal(t, nil, err)
224+
assert.Equal(t, int64(-1), res)
147225
}
148226

149227
func TestTransactionalDelete(t *testing.T) {
@@ -201,6 +279,76 @@ func TestPing(t *testing.T) {
201279
assert.Error(t, err)
202280
}
203281

282+
func TestSetRequestWithTTL(t *testing.T) {
283+
s, c := setupMiniredis()
284+
defer s.Close()
285+
286+
ss := &StateStore{
287+
client: c,
288+
json: jsoniter.ConfigFastest,
289+
logger: logger.NewLogger("test"),
290+
}
291+
ss.ctx, ss.cancel = context.WithCancel(context.Background())
292+
293+
t.Run("TTL specified", func(t *testing.T) {
294+
ttlInSeconds := 100
295+
ss.Set(&state.SetRequest{
296+
Key: "weapon100",
297+
Value: "deathstar100",
298+
Metadata: map[string]string{
299+
"ttlInSeconds": strconv.Itoa(ttlInSeconds),
300+
},
301+
})
302+
303+
ttl, _ := ss.client.TTL(ss.ctx, "weapon100").Result()
304+
305+
assert.Equal(t, time.Duration(ttlInSeconds)*time.Second, ttl)
306+
})
307+
308+
t.Run("TTL not specified", func(t *testing.T) {
309+
ss.Set(&state.SetRequest{
310+
Key: "weapon200",
311+
Value: "deathstar200",
312+
})
313+
314+
ttl, _ := ss.client.TTL(ss.ctx, "weapon200").Result()
315+
316+
assert.Equal(t, time.Duration(-1), ttl)
317+
})
318+
319+
t.Run("TTL Changed for Existing Key", func(t *testing.T) {
320+
ss.Set(&state.SetRequest{
321+
Key: "weapon300",
322+
Value: "deathstar300",
323+
})
324+
ttl, _ := ss.client.TTL(ss.ctx, "weapon300").Result()
325+
assert.Equal(t, time.Duration(-1), ttl)
326+
327+
// make the key no longer persistent
328+
ttlInSeconds := 123
329+
ss.Set(&state.SetRequest{
330+
Key: "weapon300",
331+
Value: "deathstar300",
332+
Metadata: map[string]string{
333+
"ttlInSeconds": strconv.Itoa(ttlInSeconds),
334+
},
335+
})
336+
ttl, _ = ss.client.TTL(ss.ctx, "weapon300").Result()
337+
assert.Equal(t, time.Duration(ttlInSeconds)*time.Second, ttl)
338+
339+
// make the key persistent again
340+
ss.Set(&state.SetRequest{
341+
Key: "weapon300",
342+
Value: "deathstar301",
343+
Metadata: map[string]string{
344+
"ttlInSeconds": strconv.Itoa(-1),
345+
},
346+
})
347+
ttl, _ = ss.client.TTL(ss.ctx, "weapon300").Result()
348+
assert.Equal(t, time.Duration(-1), ttl)
349+
})
350+
}
351+
204352
func TestTransactionalDeleteNoEtag(t *testing.T) {
205353
s, c := setupMiniredis()
206354
defer s.Close()

0 commit comments

Comments
 (0)