Merge pull request #57 from rewardStyle/feature/limit-reset-timer

add reset timer

Author: fjordan

kinesis.go

@@ -20,11 +20,13 @@ const (
 	statusActive
 	staticUpdating
 
+	kinesisWritesPerSec int = 1000
+	kinesisReadsPerSec  int = 5
+
 	// http://docs.aws.amazon.com/kinesis/latest/APIReference/API_GetRecords.html#API_GetRecords_RequestSyntax
 	defaultLimit int = 10000
 
-	kinesisWritesPerSec int = 1000
-	kinesisReadsPerSec  int = 5
+	limitResetDuration time.Duration = 1 * time.Minute
 
 	// Timeout TODO
 	Timeout = 60
@@ -61,17 +63,19 @@ type kinesis struct {
 	sequenceNumber    string
 	sequenceNumberMu  sync.Mutex
 
-	limit     int
-	origLimit int
-	limitMu   sync.Mutex
-	client    gokinesis.KinesisClient
+	limit        int
+	origLimit    int
+	limitMu      sync.Mutex
+	limitReset   *time.Timer
+	limitResetMu sync.Mutex
+
+	client gokinesis.KinesisClient
 
 	msgCount int64
 	errCount int64
 }
 
 func (k *kinesis) init(stream, shard, shardIteratorType, accessKey, secretKey, region string) (*kinesis, error) {
-
 	auth, err := authenticate(accessKey, secretKey)
 	k = &kinesis{
 		stream:            stream,
@@ -98,7 +102,10 @@ func (k *kinesis) args() *gokinesis.RequestArgs {
 	args.Add("StreamName", k.stream)
 	args.Add("ShardId", k.shard)
 	args.Add("ShardIterator", k.shardIterator)
+
+	k.limitMu.Lock()
 	args.Add("Limit", k.limit)
+	k.limitMu.Unlock()
 
 	if k.sequenceNumber != "" {
 		args.Add("StartingSequenceNumber", k.sequenceNumber)
@@ -200,12 +207,41 @@ func getLock(sem chan bool) {
 
 func (k *kinesis) decreaseRequestLimit() {
 	k.limitMu.Lock()
+
 	k.limit = k.limit >> 1
+
+	if k.limit <= 0 {
+		k.limit = 1
+	}
+
 	k.limitMu.Unlock()
 }
 
+func (k *kinesis) limitResetter() {
+	<-k.limitReset.C
+	k.resetRequestLimit()
+}
+
 func (k *kinesis) resetRequestLimit() {
 	k.limitMu.Lock()
 	k.limit = k.origLimit
 	k.limitMu.Unlock()
 }
+
+func (k *kinesis) startRequestLimitReset(dur time.Duration) {
+	k.limitResetMu.Lock()
+
+	// Init or reset the limitReset ticker
+	if k.limitReset == nil {
+		k.limitReset = time.NewTimer(dur)
+	} else {
+		if !k.limitReset.Stop() {
+			<-k.limitReset.C
+		}
+		k.limitReset.Reset(dur)
+	}
+
+	go k.limitResetter()
+
+	k.limitResetMu.Unlock()
+}

listener.go

@@ -43,6 +43,9 @@ type Listener struct {
 	consumingMu sync.Mutex
 	messageMu   sync.Mutex
 
+	paused  bool
+	pauseMu sync.Mutex
+
 	errors     chan error
 	messages   chan *Message
 	interrupts chan os.Signal
@@ -208,6 +211,28 @@ retry:
 	l.messageMu.Unlock()
 }
 
+func (l *Listener) pause(duration time.Duration) {
+	l.pauseMu.Lock()
+
+	if l.paused == false {
+		go func() {
+			<-time.After(duration)
+			l.pauseMu.Lock()
+			l.paused = false
+			l.pauseMu.Unlock()
+		}()
+	}
+
+	l.paused = true
+	l.pauseMu.Unlock()
+}
+
+func (l *Listener) isPaused() bool {
+	l.pauseMu.Lock()
+	defer l.pauseMu.Unlock()
+	return l.paused
+}
+
 func (l *Listener) getRecords(GsiCounter int, GsiTimer time.Time) (*gokinesis.GetRecordsResp, error) {
 	// args() will give us the shard iterator and type as well as the shard id
 	response, err := l.client.GetRecords(l.args())
@@ -220,9 +245,19 @@ func (l *Listener) getRecords(GsiCounter int, GsiTimer time.Time) (*gokinesis.Ge
 		// the AWS General Reference.
 		if strings.Contains(err.Error(), `ProvisionedThroughputExceededException`) {
 			log.Println("Received ProvisionedThroughputExceededException. Temporarily decreasing request limit and retrying.")
+
 			log.Println("Previous Request Limit: ", l.limit)
 			l.decreaseRequestLimit()
 			log.Println("New Request Limit: ", l.limit)
+
+			// The maximum size of data that GetRecords can return is 10 MB.
+			// If a call returns this amount of data, subsequent calls made
+			// within the next 5 seconds throw ProvisionedThroughputExceededException.
+			// If there is insufficient provisioned throughput on the stream,
+			/// subsequent calls made within the next 1 second throw ProvisionedThroughputExceededException.
+			l.pause(5 * time.Second)
+
+			go l.startRequestLimitReset(limitResetDuration)
 		}
 	}
 
@@ -282,7 +317,6 @@ func (l *Listener) consume() {
 
 		if response != nil {
 			l.setShardIterator(response.NextShardIterator)
-			l.resetRequestLimit()
 
 			if len(response.Records) > 0 {
 				for _, record := range response.Records {
@@ -427,6 +461,9 @@ func (l *Listener) handleError(err error) {
 // Kinesis allows five read ops per second per shard
 // http://docs.aws.amazon.com/kinesis/latest/dev/service-sizes-and-limits.html
 func (l *Listener) throttle(counter *int, timer *time.Time) {
+	for l.isPaused() {
+		runtime.Gosched()
+	}
 	// If a second has passed since the last timer start, reset the timer
 	if time.Now().After(timer.Add(1 * time.Second)) {
 		*timer = time.Now()

listener_test.go

@@ -135,6 +135,64 @@ func TestRetrieveMessage(t *testing.T) {
 	listener.Close()
 }
 
+func TestLimitReset(t *testing.T) {
+	l, _ := new(Listener).InitC("your-stream", "0", ShardIterTypes[3], "accesskey", "secretkey", "us-east-1", 10)
+	l.NewEndpoint(testEndpoint, "your-stream")
+
+	testLimitResetDuration := 2 * time.Second
+
+	Convey("If the request limit has not been decreased", t, func() {
+		Convey("It should be equal to the defaultLimit", func() {
+			So(l.limit, ShouldEqual, defaultLimit)
+
+			Convey("If the request limit has been decreased", func() {
+				l.decreaseRequestLimit()
+				Convey("It should not be reset until the timer has elapsed", func() {
+
+					go l.startRequestLimitReset(testLimitResetDuration)
+					<-time.After(1 * time.Second)
+					l.limitMu.Lock()
+					So(l.limit, ShouldBeLessThan, defaultLimit)
+					l.limitMu.Unlock()
+
+					go l.startRequestLimitReset(testLimitResetDuration)
+					<-time.After(1 * time.Second)
+					l.limitMu.Lock()
+					So(l.limit, ShouldBeLessThan, defaultLimit)
+					l.limitMu.Unlock()
+
+					go l.startRequestLimitReset(testLimitResetDuration)
+					<-time.After(1 * time.Second)
+					l.limitMu.Lock()
+					So(l.limit, ShouldBeLessThan, defaultLimit)
+					l.limitMu.Unlock()
+
+					<-time.After(testLimitResetDuration + 1*time.Second)
+					l.limitMu.Lock()
+					So(l.limit, ShouldEqual, defaultLimit)
+					l.limitMu.Unlock()
+				})
+			})
+		})
+	})
+}
+
+func TestLimitGreaterThanZero(t *testing.T) {
+	l, _ := new(Listener).InitC("your-stream", "0", ShardIterTypes[3], "accesskey", "secretkey", "us-east-1", 10)
+	l.NewEndpoint(testEndpoint, "your-stream")
+
+	Convey("If the listener limit is decreased", t, func() {
+
+		for idx := 0; idx < 100; idx++ {
+			l.decreaseRequestLimit()
+		}
+
+		Convey("it should always be greater than zero", func() {
+			So(l.limit, ShouldBeGreaterThan, 0)
+		})
+	})
+}
+
 var cases = []struct {
 	message []byte
 }{