54 lines
1.1 KiB
Go
54 lines
1.1 KiB
Go
|
package utility
|
||
|
|
|
||
|
|
import (
|
||
|
|
"math/rand"
|
||
|
|
"sync/atomic"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/aws/aws-sdk-go/aws/awserr"
|
||
|
|
"github.com/aws/aws-sdk-go/service/dynamodb"
|
||
|
|
|
||
|
|
checkpoint "github.com/harlow/kinesis-consumer/checkpoint/ddb"
|
||
|
|
)
|
||
|
|
|
||
|
|
const defaultSleepInterval = 30 * time.Second
|
||
|
|
|
||
|
|
// Retryer used for checkpointing
|
||
|
|
type Retryer struct {
|
||
|
|
checkpoint.Retryer
|
||
|
|
count uint64
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewRetryer() *Retryer {
|
||
|
|
return &Retryer{count: 0}
|
||
|
|
}
|
||
|
|
|
||
|
|
// ShouldRetry implements custom logic for when a checkpont should retry
|
||
|
|
func (r *Retryer) ShouldRetry(err error) bool {
|
||
|
|
if awsErr, ok := err.(awserr.Error); ok {
|
||
|
|
switch awsErr.Code() {
|
||
|
|
case dynamodb.ErrCodeProvisionedThroughputExceededException, dynamodb.ErrCodeLimitExceededException:
|
||
|
|
jitter := rand.New(rand.NewSource(0))
|
||
|
|
atomic.AddUint64(&r.count, 1)
|
||
|
|
// You can have more sophisticated sleep mechanism
|
||
|
|
time.Sleep(30 * time.Second)
|
||
|
|
randomSleep(defaultSleepInterval, jitter)
|
||
|
|
return true
|
||
|
|
default:
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r Retryer) Count() uint {
|
||
|
|
return uint(r.count)
|
||
|
|
}
|
||
|
|
|
||
|
|
func randomSleep(d time.Duration, r *rand.Rand) time.Duration {
|
||
|
|
if d == 0 {
|
||
|
|
return 0
|
||
|
|
}
|
||
|
|
return d + time.Duration(r.Int63n(2*int64(d)))
|
||
|
|
}
|