Merge pull request #6 from postmates/rhettg-retry

Detect ProvisionedThroughputExceededException and retry

Author: rhettg

triton/stream.go

@@ -2,6 +2,7 @@ package triton
 
 import (
 	"fmt"
+	"log"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -28,11 +29,13 @@ type ShardStreamReader struct {
 	service     KinesisService
 	records     []*kinesis.Record
 	lastRequest *time.Time
+	retries     int64
 }
 
 // Recommended minimum polling interval to keep from overloading a Kinesis
 // shard.
-const MIN_POLL_INTERVAL = 1.0 * time.Second
+const minPollInterval = 1.0 * time.Second
+const maxRetries = 3
 
 func (s *ShardStreamReader) initIterator() (err error) {
 	gsi := kinesis.GetShardIteratorInput{
@@ -56,7 +59,9 @@ func (s *ShardStreamReader) initIterator() (err error) {
 
 func (s *ShardStreamReader) wait(minInterval time.Duration) {
 	if s.lastRequest != nil {
-		sleepTime := minInterval - time.Since(*s.lastRequest)
+		retryDelay := time.Duration(s.retries*250) * time.Millisecond
+
+		sleepTime := (minInterval - time.Since(*s.lastRequest)) + retryDelay
 		if sleepTime >= time.Duration(0) {
 			time.Sleep(sleepTime)
 		}
@@ -66,8 +71,44 @@ func (s *ShardStreamReader) wait(minInterval time.Duration) {
 	s.lastRequest = &n
 }
 
+// Documented Kinesis specific errors as well as common errors we
+// should probably just retry on.
+// http://docs.aws.amazon.com/kinesis/latest/APIReference/CommonErrors.html
+var retryErrorCodes = [...]string{
+	"ProvisionedThroughputExceededException",
+	"ServiceUnavailable",
+	"InternalFailure",
+	"Throttling",
+}
+
+func (s *ShardStreamReader) isRetryError(err error) bool {
+	if awsErr, ok := err.(awserr.Error); ok {
+		retry := false
+
+		for _, code := range retryErrorCodes {
+			if awsErr.Code() == code {
+				retry = true
+				break
+			}
+		}
+
+		if retry {
+			s.retries += 1
+			if s.retries <= maxRetries {
+				log.Printf("%s: %s. Retrying", awsErr.Code(), awsErr.Message())
+				return true
+			} else {
+				log.Printf("%s: %s. Max retries attempted", awsErr.Code(), awsErr.Message())
+				return false
+			}
+		}
+	}
+
+	return false
+}
+
 func (s *ShardStreamReader) fetchMoreRecords() (err error) {
-	s.wait(MIN_POLL_INTERVAL)
+	s.wait(minPollInterval)
 
 	if s.NextIteratorValue == nil {
 		err := s.initIterator()
@@ -83,9 +124,15 @@ func (s *ShardStreamReader) fetchMoreRecords() (err error) {
 
 	gro, err := s.service.GetRecords(gri)
 	if err != nil {
-		return err
+		if s.isRetryError(err) {
+			return nil
+		} else {
+			return err
+		}
 	}
 
+	s.retries = 0
+
 	s.records = gro.Records
 	s.NextIteratorValue = gro.NextShardIterator
 

triton/stream_test.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/service/kinesis"
 )
 
@@ -30,6 +31,22 @@ func (s *NullKinesisService) DescribeStream(input *kinesis.DescribeStreamInput)
 	return nil, fmt.Errorf("Not Implemented")
 }
 
+type FailingKinesisService struct{}
+
+func (s *FailingKinesisService) DescribeStream(input *kinesis.DescribeStreamInput) (*kinesis.DescribeStreamOutput, error) {
+	return nil, fmt.Errorf("Not Implemented")
+}
+
+func (s *FailingKinesisService) GetRecords(*kinesis.GetRecordsInput) (*kinesis.GetRecordsOutput, error) {
+	err := awserr.New("ProvisionedThroughputExceededException", "slow down dummy", fmt.Errorf("error"))
+	return nil, err
+}
+
+func (s *FailingKinesisService) GetShardIterator(*kinesis.GetShardIteratorInput) (*kinesis.GetShardIteratorOutput, error) {
+	gso := &kinesis.GetShardIteratorOutput{ShardIterator: aws.String("123")}
+	return gso, nil
+}
+
 func TestNewShardStreamReader(t *testing.T) {
 	svc := NullKinesisService{}
 
@@ -108,6 +125,26 @@ func TestFetchMoreRecords(t *testing.T) {
 	}
 }
 
+func TestRetryShardStreamReader(t *testing.T) {
+	svc := FailingKinesisService{}
+
+	s := NewShardStreamReader(&svc, "test-stream", "shard-0001")
+
+	err := s.fetchMoreRecords()
+	if err != nil {
+		t.Errorf("Received error %v", err)
+		return
+	}
+
+	if len(s.records) != 0 {
+		t.Errorf("Should have no records")
+	}
+
+	if s.retries != 1 {
+		t.Errorf("Should have a retry")
+	}
+}
+
 func TestRead(t *testing.T) {
 	svc := newTestKinesisService()
 	st := newTestKinesisStream("test-stream")