Merge pull request #5 from Versent/fix_capture_ts_before_iterators

fix(iterators) Remove difference between start time of iterators.

Author: wolfeidau

.gitignore

@@ -4,3 +4,4 @@ console.log
 trace.out
 dist
 node_modules
+/bin

Makefile

@@ -6,13 +6,12 @@ GO ?= go
 
 # Install all the build and lint dependencies
 setup:
-	go get -u github.com/alecthomas/gometalinter
-	go get -u github.com/golang/dep/cmd/dep
-	go get -u github.com/pierrre/gotestcover
-	go get -u golang.org/x/tools/cmd/cover
-	go get github.com/vektra/mockery/...
-	gometalinter --install
-	dep ensure
+	@$(GO) get -u github.com/alecthomas/gometalinter
+	@$(GO) get -u github.com/golang/dep/cmd/dep
+	@$(GO) get -u github.com/axw/gocov/...
+	@$(GO) get github.com/vektra/mockery/...
+	@gometalinter --install
+	@dep ensure
 .PHONY: setup
 
 # Install from source.
@@ -23,17 +22,12 @@ install:
 
 # Run all the tests
 test:
-	@gotestcover $(TEST_OPTIONS) -covermode=atomic -coverprofile=coverage.txt $(SOURCE_FILES) -run $(TEST_PATTERN) -timeout=2m
+	@gocov test -timeout=2m ./... | gocov report
 .PHONY: test
 
-# Run all the tests and opens the coverage report
-cover: test
-	@$(GO) tool cover -html=coverage.txt
-.PHONY: cover
-
 # Run all the linters
 lint:
-	gometalinter --vendor ./...
+	gometalinter --deadline 300s --vendor ./...
 .PHONY: lint
 
 # Run all the tests and code checks

cmd/kinesis-tail/main.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/alecthomas/kingpin"
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/credentials"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/kinesis"
 	"github.com/aws/aws-sdk-go/service/kinesis/kinesisiface"
@@ -30,6 +31,8 @@ var (
 	tracing       = kingpin.Flag("trace", "Enable trace mode.").Short('t').Bool()
 	debug         = kingpin.Flag("debug", "Enable debug logging.").Short('d').Bool()
 	region        = kingpin.Flag("region", "Configure the aws region.").Short('r').String()
+	profile       = kingpin.Flag("profile", "Configure the aws profile.").Short('p').String()
+	timestamp     = kingpin.Flag("timestamp", "Start time in epoch milliseconds.").Short('T').Int64()
 	cwlogsCommand = kingpin.Command("cwlogs", "Process cloudwatch logs data from kinesis.")
 	includes      = cwlogsCommand.Flag("include", "Include anything in log group names which match the supplied string.").Strings()
 	excludes      = cwlogsCommand.Flag("exclude", "Exclude anything in log group names which match the supplied string.").Strings()
@@ -58,7 +61,6 @@ func main() {
 		}
 
 		defer trace.Stop()
-
 	}
 
 	if *debug {
@@ -68,44 +70,35 @@ func main() {
 		logger.SetLevel(logrus.DebugLevel)
 	}
 
-	sess := session.Must(session.NewSession())
-
-	var svc kinesisiface.KinesisAPI
-
-	if *region == "" {
-		svc = kinesis.New(sess)
-	} else {
-		// Create a Kinesis client with additional configuration
-		svc = kinesis.New(sess, aws.NewConfig().WithRegion(*region))
-	}
+	svc := newKinesis(region, profile)
 
 	logger.Debug("built kinesis service")
 
 	switch subCommand {
 	case "cwlogs":
-		err := processLogData(svc, *cwlogsStream, *includes, *excludes)
+		err := processLogData(svc, *cwlogsStream, *timestamp, *includes, *excludes)
 		if err != nil {
 			logger.WithError(err).Fatal("failed to process log data")
 		}
 	case "raw":
-		err := processRawData(svc, *rawStream, *timeout, *count)
+		err := processRawData(svc, *rawStream, *timeout, *timestamp, *count)
 		if err != nil {
 			logger.WithError(err).Fatal("failed to process log data")
 		}
 	}
 
 }
 
-func processLogData(svc kinesisiface.KinesisAPI, stream string, includes []string, excludes []string) error {
+func processLogData(svc kinesisiface.KinesisAPI, stream string, timestamp int64, includes []string, excludes []string) error {
 
 	helper := ktail.New(svc, logger)
 
-	iterators, err := helper.GetStreamIterators(stream)
+	iterators, err := helper.GetStreamIterators(stream, timestamp)
 	if err != nil {
 		return errors.Wrap(err, "get iterators failed")
 	}
 
-	kstream := streamer.New(svc, iterators, 5000)
+	kstream := streamer.New(svc, iterators, 5000, logger)
 	ch := kstream.StartGetRecords()
 
 	messageSorter := sorter.New(os.Stdout, len(iterators), formatLogsMsg)
@@ -136,16 +129,16 @@ func processLogData(svc kinesisiface.KinesisAPI, stream string, includes []strin
 	return nil
 }
 
-func processRawData(svc kinesisiface.KinesisAPI, stream string, timeout int64, count int) error {
+func processRawData(svc kinesisiface.KinesisAPI, stream string, timeout int64, timestamp int64, count int) error {
 
 	helper := ktail.New(svc, logger)
 
-	iterators, err := helper.GetStreamIterators(stream)
+	iterators, err := helper.GetStreamIterators(stream, timestamp)
 	if err != nil {
 		return errors.Wrap(err, "get iterators failed")
 	}
 
-	kstream := streamer.New(svc, iterators, 5000)
+	kstream := streamer.New(svc, iterators, 5000, logger)
 	ch := kstream.StartGetRecords()
 
 	messageSorter := sorter.New(os.Stdout, len(iterators), formatRawMsg)
@@ -217,3 +210,19 @@ func formatLogsMsg(wr io.Writer, msg *ktail.LogMessage) {
 		logger.WithError(err).Fatal("failed to create trace file")
 	}
 }
+
+func newKinesis(region, profile *string) kinesisiface.KinesisAPI {
+	sess := session.Must(session.NewSession())
+
+	cfg := aws.NewConfig()
+
+	if aws.StringValue(region) != "" {
+		cfg = cfg.WithRegion(*region)
+	}
+
+	if aws.StringValue(profile) != "" {
+		cfg = cfg.WithCredentials(credentials.NewSharedCredentials("", *profile))
+	}
+
+	return kinesis.New(sess, cfg)
+}

pkg/ktail/kinesis.go

@@ -1,6 +1,8 @@
 package ktail
 
 import (
+	"time"
+
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/kinesis"
 	"github.com/aws/aws-sdk-go/service/kinesis/kinesisiface"
@@ -28,7 +30,19 @@ func New(svc kinesisiface.KinesisAPI, logger *logrus.Logger) *KinesisHelper {
 }
 
 // GetStreamIterators build a list of iterators for the stream
-func (kh *KinesisHelper) GetStreamIterators(streamName string) (map[string]*string, error) {
+func (kh *KinesisHelper) GetStreamIterators(streamName string, timestamp int64) (map[string]*string, error) {
+
+	var ts time.Time
+
+	if timestamp > 0 {
+		tsec := timestamp / 1000
+		tnano := (timestamp % 1000) * 1000
+		ts = time.Unix(tsec, tnano)
+	} else {
+		ts = time.Now()
+	}
+
+	kh.logger.WithField("ts", ts.Unix()).Info("starting stream")
 
 	respDesc, err := kh.svc.DescribeStream(&kinesis.DescribeStreamInput{
 		StreamName: aws.String(streamName),
@@ -42,7 +56,7 @@ func (kh *KinesisHelper) GetStreamIterators(streamName string) (map[string]*stri
 	iterators := map[string]*string{}
 
 	for _, shard := range respDesc.StreamDescription.Shards {
-		go kh.asyncGetShardIterator(streamName, aws.StringValue(shard.ShardId), ch)
+		go kh.asyncGetShardIterator(streamName, aws.StringValue(shard.ShardId), ts, ch)
 	}
 
 	for range respDesc.StreamDescription.Shards {
@@ -54,13 +68,14 @@ func (kh *KinesisHelper) GetStreamIterators(streamName string) (map[string]*stri
 	return iterators, nil
 }
 
-func (kh *KinesisHelper) asyncGetShardIterator(streamName, shardID string, ch chan *iteratorResult) {
+func (kh *KinesisHelper) asyncGetShardIterator(streamName, shardID string, ts time.Time, ch chan *iteratorResult) {
 	kh.logger.WithField("shard", shardID).Debug("get shard iterator")
 
 	respShard, err := kh.svc.GetShardIterator(&kinesis.GetShardIteratorInput{
 		StreamName:        aws.String(streamName),
 		ShardIteratorType: aws.String(kinesis.ShardIteratorTypeLatest),
 		ShardId:           aws.String(shardID),
+		Timestamp:         aws.Time(ts),
 	})
 	if err != nil {
 		kh.logger.WithError(err).Fatal("get shard iterator failed")

pkg/ktail/kinesis_test.go

@@ -0,0 +1 @@
+package ktail

pkg/streamer/kinesis.go

@@ -4,6 +4,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/sirupsen/logrus"
+
 	"github.com/aws/aws-sdk-go/service/kinesis"
 	"github.com/aws/aws-sdk-go/service/kinesis/kinesisiface"
 	"github.com/pkg/errors"
@@ -15,6 +17,7 @@ type KinesisStreamer struct {
 	iterators     map[string]*string
 	iteratorMutex *sync.Mutex
 	pollFreqMs    int64
+	logger        *logrus.Logger
 }
 
 // GetRecordsEntry returns the results of the last get records request
@@ -26,12 +29,13 @@ type GetRecordsEntry struct {
 }
 
 // New return a new configured streamer
-func New(svc kinesisiface.KinesisAPI, iterators map[string]*string, pollFreqMs int64) *KinesisStreamer {
+func New(svc kinesisiface.KinesisAPI, iterators map[string]*string, pollFreqMs int64, logger *logrus.Logger) *KinesisStreamer {
 	return &KinesisStreamer{
 		svc:           svc,
 		iterators:     iterators,
 		pollFreqMs:    pollFreqMs,
 		iteratorMutex: &sync.Mutex{},
+		logger:        logger,
 	}
 }
 
@@ -53,14 +57,19 @@ func (ks *KinesisStreamer) asyncGetRecords(shard string, ch chan *GetRecordsEntr
 
 	for now := range c {
 
+		if ks.iterators[shard] == nil {
+			ks.logger.Debug("nil iterator for shard as it is CLOSED: %s", shard)
+			continue
+		}
+
 		resp, err := ks.svc.GetRecords(&kinesis.GetRecordsInput{
 			ShardIterator: ks.iterators[shard],
 		})
 		if err != nil {
 			ch <- &GetRecordsEntry{Created: now, Shard: shard, Err: errors.Wrap(err, "get records failed")}
 		}
 
-		//logger.WithField("iterator", aws.StringValue(iterator)).Info("get records shard")
+		ks.logger.WithField("iterator", resp).Debug("get records shard")
 
 		ch <- &GetRecordsEntry{Created: now, Shard: shard, Records: resp.Records}