Skip to content

Commit b4953a3

Browse files
yeya24yeya24
authored
Add timeout for dynamodb ring kv (#6544)
* add dynamodb kv with timeout enforced Signed-off-by: yeya24 <[email protected]> * add tests Signed-off-by: yeya24 <[email protected]> * docs Signed-off-by: Ben Ye <[email protected]> * update changelog Signed-off-by: Ben Ye <[email protected]> --------- Signed-off-by: yeya24 <[email protected]> Signed-off-by: Ben Ye <[email protected]>
1 parent de5cfe1 commit b4953a3

File tree

6 files changed

+160
-1
lines changed

6 files changed

+160
-1
lines changed

‎docs/blocks-storage/compactor.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,10 @@ compactor:
225225
# CLI flag: -compactor.ring.dynamodb.max-cas-retries
226226
[max_cas_retries: <int> | default = 10]
227227

228+
# Timeout of dynamoDbClient requests. Default is 2m.
229+
# CLI flag: -compactor.ring.dynamodb.timeout
230+
[timeout: <duration> | default = 2m]
231+
228232
# The consul_config configures the consul client.
229233
# The CLI flags prefix for this block config is: compactor.ring
230234
[consul: <consul_config>]

‎docs/blocks-storage/store-gateway.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -240,6 +240,10 @@ store_gateway:
240240
# CLI flag: -store-gateway.sharding-ring.dynamodb.max-cas-retries
241241
[max_cas_retries: <int> | default = 10]
242242

243+
# Timeout of dynamoDbClient requests. Default is 2m.
244+
# CLI flag: -store-gateway.sharding-ring.dynamodb.timeout
245+
[timeout: <duration> | default = 2m]
246+
243247
# The consul_config configures the consul client.
244248
# The CLI flags prefix for this block config is:
245249
# store-gateway.sharding-ring

‎docs/configuration/config-file-reference.md‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -341,6 +341,10 @@ sharding_ring:
341341
# CLI flag: -alertmanager.sharding-ring.dynamodb.max-cas-retries
342342
[max_cas_retries: <int> | default = 10]
343343
344+
# Timeout of dynamoDbClient requests. Default is 2m.
345+
# CLI flag: -alertmanager.sharding-ring.dynamodb.timeout
346+
[timeout: <duration> | default = 2m]
347+
344348
# The consul_config configures the consul client.
345349
# The CLI flags prefix for this block config is: alertmanager.sharding-ring
346350
[consul: <consul_config>]
@@ -2286,6 +2290,10 @@ sharding_ring:
22862290
# CLI flag: -compactor.ring.dynamodb.max-cas-retries
22872291
[max_cas_retries: <int> | default = 10]
22882292
2293+
# Timeout of dynamoDbClient requests. Default is 2m.
2294+
# CLI flag: -compactor.ring.dynamodb.timeout
2295+
[timeout: <duration> | default = 2m]
2296+
22892297
# The consul_config configures the consul client.
22902298
# The CLI flags prefix for this block config is: compactor.ring
22912299
[consul: <consul_config>]
@@ -2595,6 +2603,10 @@ ha_tracker:
25952603
# CLI flag: -distributor.ha-tracker.dynamodb.max-cas-retries
25962604
[max_cas_retries: <int> | default = 10]
25972605
2606+
# Timeout of dynamoDbClient requests. Default is 2m.
2607+
# CLI flag: -distributor.ha-tracker.dynamodb.timeout
2608+
[timeout: <duration> | default = 2m]
2609+
25982610
# The consul_config configures the consul client.
25992611
# The CLI flags prefix for this block config is: distributor.ha-tracker
26002612
[consul: <consul_config>]
@@ -2689,6 +2701,10 @@ ring:
26892701
# CLI flag: -distributor.ring.dynamodb.max-cas-retries
26902702
[max_cas_retries: <int> | default = 10]
26912703
2704+
# Timeout of dynamoDbClient requests. Default is 2m.
2705+
# CLI flag: -distributor.ring.dynamodb.timeout
2706+
[timeout: <duration> | default = 2m]
2707+
26922708
# The consul_config configures the consul client.
26932709
# The CLI flags prefix for this block config is: distributor.ring
26942710
[consul: <consul_config>]
@@ -3017,6 +3033,10 @@ lifecycler:
30173033
# CLI flag: -dynamodb.max-cas-retries
30183034
[max_cas_retries: <int> | default = 10]
30193035
3036+
# Timeout of dynamoDbClient requests. Default is 2m.
3037+
# CLI flag: -dynamodb.timeout
3038+
[timeout: <duration> | default = 2m]
3039+
30203040
# The consul_config configures the consul client.
30213041
[consul: <consul_config>]
30223042
@@ -4674,6 +4694,10 @@ ring:
46744694
# CLI flag: -ruler.ring.dynamodb.max-cas-retries
46754695
[max_cas_retries: <int> | default = 10]
46764696

4697+
# Timeout of dynamoDbClient requests. Default is 2m.
4698+
# CLI flag: -ruler.ring.dynamodb.timeout
4699+
[timeout: <duration> | default = 2m]
4700+
46774701
# The consul_config configures the consul client.
46784702
# The CLI flags prefix for this block config is: ruler.ring
46794703
[consul: <consul_config>]
@@ -5665,6 +5689,10 @@ sharding_ring:
56655689
# CLI flag: -store-gateway.sharding-ring.dynamodb.max-cas-retries
56665690
[max_cas_retries: <int> | default = 10]
56675691
5692+
# Timeout of dynamoDbClient requests. Default is 2m.
5693+
# CLI flag: -store-gateway.sharding-ring.dynamodb.timeout
5694+
[timeout: <duration> | default = 2m]
5695+
56685696
# The consul_config configures the consul client.
56695697
# The CLI flags prefix for this block config is: store-gateway.sharding-ring
56705698
[consul: <consul_config>]

‎pkg/ring/kv/dynamodb/client.go‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ type Config struct {
2626
TTL time.Duration `yaml:"ttl"`
2727
PullerSyncTime time.Duration `yaml:"puller_sync_time"`
2828
MaxCasRetries int `yaml:"max_cas_retries"`
29+
Timeout time.Duration `yaml:"timeout"`
2930
}
3031

3132
type Client struct {
@@ -53,6 +54,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet, prefix string) {
5354
f.DurationVar(&cfg.TTL, prefix+"dynamodb.ttl-time", 0, "Time to expire items on dynamodb.")
5455
f.DurationVar(&cfg.PullerSyncTime, prefix+"dynamodb.puller-sync-time", 60*time.Second, "Time to refresh local ring with information on dynamodb.")
5556
f.IntVar(&cfg.MaxCasRetries, prefix+"dynamodb.max-cas-retries", maxCasRetries, "Maximum number of retries for DDB KV CAS.")
57+
f.DurationVar(&cfg.Timeout, prefix+"dynamodb.timeout", 2*time.Minute, "Timeout of dynamoDbClient requests. Default is 2m.")
5658
}
5759

5860
func NewClient(cfg Config, cc codec.Codec, logger log.Logger, registerer prometheus.Registerer) (*Client, error) {
@@ -69,8 +71,13 @@ func NewClient(cfg Config, cc codec.Codec, logger log.Logger, registerer prometh
6971
MaxRetries: cfg.MaxCasRetries,
7072
}
7173

74+
var kv dynamoDbClient
75+
kv = dynamodbInstrumentation{kv: dynamoDB, ddbMetrics: ddbMetrics}
76+
if cfg.Timeout > 0 {
77+
kv = newDynamodbKVWithTimeout(kv, cfg.Timeout)
78+
}
7279
c := &Client{
73-
kv: dynamodbInstrumentation{kv: dynamoDB, ddbMetrics: ddbMetrics},
80+
kv: kv,
7481
codec: cc,
7582
logger: ddbLog(logger),
7683
ddbMetrics: ddbMetrics,

‎pkg/ring/kv/dynamodb/client_test.go‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -302,6 +302,29 @@ func Test_UpdateStaleData(t *testing.T) {
302302

303303
}
304304

305+
func Test_DynamodbKVWithTimeout(t *testing.T) {
306+
ddbMock := NewDynamodbClientMock()
307+
// Backend has delay of 5s while the client timeout is 1s.
308+
ddbWithDelay := newDynamodbKVWithDelay(ddbMock, time.Second*5)
309+
dbWithTimeout := newDynamodbKVWithTimeout(ddbWithDelay, time.Second)
310+
311+
ctx := context.Background()
312+
_, _, err := dbWithTimeout.List(ctx, dynamodbKey{primaryKey: key})
313+
require.True(t, errors.Is(err, context.DeadlineExceeded))
314+
315+
err = dbWithTimeout.Delete(ctx, dynamodbKey{primaryKey: key})
316+
require.True(t, errors.Is(err, context.DeadlineExceeded))
317+
318+
_, _, err = dbWithTimeout.Query(ctx, dynamodbKey{primaryKey: key}, true)
319+
require.True(t, errors.Is(err, context.DeadlineExceeded))
320+
321+
err = dbWithTimeout.Put(ctx, dynamodbKey{primaryKey: key}, []byte{})
322+
require.True(t, errors.Is(err, context.DeadlineExceeded))
323+
324+
err = dbWithTimeout.Batch(ctx, nil, nil)
325+
require.True(t, errors.Is(err, context.DeadlineExceeded))
326+
}
327+
305328
// NewClientMock makes a new local dynamodb client.
306329
func NewClientMock(ddbClient dynamoDbClient, cc codec.Codec, logger log.Logger, registerer prometheus.Registerer, time time.Duration, config backoff.Config) *Client {
307330
return &Client{
@@ -429,3 +452,57 @@ func (m *DescMock) FindDifference(that codec.MultiKey) (interface{}, []string, e
429452
}
430453
return args.Get(0), args.Get(1).([]string), err
431454
}
455+
456+
type dynamodbKVWithDelayAndContextCheck struct {
457+
ddbClient dynamoDbClient
458+
delay time.Duration
459+
}
460+
461+
func newDynamodbKVWithDelay(client dynamoDbClient, delay time.Duration) *dynamodbKVWithDelayAndContextCheck {
462+
return &dynamodbKVWithDelayAndContextCheck{ddbClient: client, delay: delay}
463+
}
464+
465+
func (d *dynamodbKVWithDelayAndContextCheck) List(ctx context.Context, key dynamodbKey) ([]string, float64, error) {
466+
select {
467+
case <-ctx.Done():
468+
return nil, 0, ctx.Err()
469+
case <-time.After(d.delay):
470+
return d.ddbClient.List(ctx, key)
471+
}
472+
}
473+
474+
func (d *dynamodbKVWithDelayAndContextCheck) Query(ctx context.Context, key dynamodbKey, isPrefix bool) (map[string][]byte, float64, error) {
475+
select {
476+
case <-ctx.Done():
477+
return nil, 0, ctx.Err()
478+
case <-time.After(d.delay):
479+
return d.ddbClient.Query(ctx, key, isPrefix)
480+
}
481+
}
482+
483+
func (d *dynamodbKVWithDelayAndContextCheck) Delete(ctx context.Context, key dynamodbKey) error {
484+
select {
485+
case <-ctx.Done():
486+
return ctx.Err()
487+
case <-time.After(d.delay):
488+
return d.ddbClient.Delete(ctx, key)
489+
}
490+
}
491+
492+
func (d *dynamodbKVWithDelayAndContextCheck) Put(ctx context.Context, key dynamodbKey, data []byte) error {
493+
select {
494+
case <-ctx.Done():
495+
return ctx.Err()
496+
case <-time.After(d.delay):
497+
return d.ddbClient.Put(ctx, key, data)
498+
}
499+
}
500+
501+
func (d *dynamodbKVWithDelayAndContextCheck) Batch(ctx context.Context, put map[dynamodbKey][]byte, delete []dynamodbKey) error {
502+
select {
503+
case <-ctx.Done():
504+
return ctx.Err()
505+
case <-time.After(d.delay):
506+
return d.ddbClient.Batch(ctx, put, delete)
507+
}
508+
}

‎pkg/ring/kv/dynamodb/dynamodb.go‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -259,6 +259,45 @@ func (kv dynamodbKV) generatePutItemRequest(key dynamodbKey, data []byte) map[st
259259
return item
260260
}
261261

262+
type dynamodbKVWithTimeout struct {
263+
ddbClient dynamoDbClient
264+
timeout time.Duration
265+
}
266+
267+
func newDynamodbKVWithTimeout(client dynamoDbClient, timeout time.Duration) *dynamodbKVWithTimeout {
268+
return &dynamodbKVWithTimeout{ddbClient: client, timeout: timeout}
269+
}
270+
271+
func (d *dynamodbKVWithTimeout) List(ctx context.Context, key dynamodbKey) ([]string, float64, error) {
272+
ctx, cancel := context.WithTimeout(ctx, d.timeout)
273+
defer cancel()
274+
return d.ddbClient.List(ctx, key)
275+
}
276+
277+
func (d *dynamodbKVWithTimeout) Query(ctx context.Context, key dynamodbKey, isPrefix bool) (map[string][]byte, float64, error) {
278+
ctx, cancel := context.WithTimeout(ctx, d.timeout)
279+
defer cancel()
280+
return d.ddbClient.Query(ctx, key, isPrefix)
281+
}
282+
283+
func (d *dynamodbKVWithTimeout) Delete(ctx context.Context, key dynamodbKey) error {
284+
ctx, cancel := context.WithTimeout(ctx, d.timeout)
285+
defer cancel()
286+
return d.ddbClient.Delete(ctx, key)
287+
}
288+
289+
func (d *dynamodbKVWithTimeout) Put(ctx context.Context, key dynamodbKey, data []byte) error {
290+
ctx, cancel := context.WithTimeout(ctx, d.timeout)
291+
defer cancel()
292+
return d.ddbClient.Put(ctx, key, data)
293+
}
294+
295+
func (d *dynamodbKVWithTimeout) Batch(ctx context.Context, put map[dynamodbKey][]byte, delete []dynamodbKey) error {
296+
ctx, cancel := context.WithTimeout(ctx, d.timeout)
297+
defer cancel()
298+
return d.ddbClient.Batch(ctx, put, delete)
299+
}
300+
262301
func generateItemKey(key dynamodbKey) map[string]*dynamodb.AttributeValue {
263302
resp := map[string]*dynamodb.AttributeValue{
264303
primaryKey: {

0 commit comments

Comments
 (0)