fix(Raw) Fixed count option and ensure records are flushed before exit.

cmd/kinesis-tail/main.go

@@ -152,6 +152,10 @@ func processRawData(svc kinesisiface.KinesisAPI, stream string, timeout int64, c
 
 	timer1 := time.NewTimer(time.Duration(timeout) * time.Millisecond)
 
+	if count > 0 {
+		logger.WithField("count", count).Debug("waiting for records")
+	}
+
 LOOP:
 	for {
 
@@ -172,16 +176,19 @@ LOOP:
 				msgResults = append(msgResults, msg)
 			}
 
+			messageSorter.PushBatch(msgResults)
+
+			recordCount = recordCount + len(result.Records)
+
 			if count != 0 {
-				if recordCount == count {
+				if recordCount >= count {
+					messageSorter.Flush()
+
 					logger.WithField("recordCount", recordCount).Info("reached count exit")
 					break LOOP
 				}
 			}
 
-			recordCount++
-			messageSorter.PushBatch(msgResults)
-
 		case <-timer1.C:
 			logger.Info("timer expired exit")
 			break LOOP

pkg/sorter/sorter.go

@@ -34,12 +34,8 @@ func (lms *MessageSorter) PushBatch(logMessageBatch []*ktail.LogMessage) bool {
 	return lms.flushCheck()
 }
 
-func (lms *MessageSorter) flushCheck() bool {
-	lms.current++
-
-	if lms.current != lms.batchSize {
-		return false
-	}
+// Flush force a flush of messages
+func (lms *MessageSorter) Flush() {
 	sort.Sort(ktail.ByTimestamp(lms.cache))
 
 	for _, msg := range lms.cache {
@@ -48,6 +44,16 @@ func (lms *MessageSorter) flushCheck() bool {
 
 	lms.cache = []*ktail.LogMessage{}
 	lms.current = 0
+}
+
+func (lms *MessageSorter) flushCheck() bool {
+	lms.current++
+
+	if lms.current != lms.batchSize {
+		return false
+	}
+
+	lms.Flush()
 
 	return true
 }