From f6f2f66c7893165067ceba989ac73d82a28b3bd4 Mon Sep 17 00:00:00 2001 From: Tommy Wang Date: Mon, 27 Mar 2017 17:57:13 -0500 Subject: [PATCH 001/109] new listener --- .gitignore | 24 ++ glide.lock | 6 +- glide.yaml | 2 +- listener/config.go | 112 ++++++ listener/listener.go | 752 ++++++++++++++++++++++++++++++++++++++ listener/listener_test.go | 527 ++++++++++++++++++++++++++ listener/stats.go | 37 ++ utils.go | 20 + 8 files changed, 1476 insertions(+), 4 deletions(-) create mode 100644 listener/config.go create mode 100644 listener/listener.go create mode 100644 listener/listener_test.go create mode 100644 listener/stats.go create mode 100644 utils.go diff --git a/.gitignore b/.gitignore index c6d755d..3204325 100644 --- a/.gitignore +++ b/.gitignore @@ -37,3 +37,27 @@ Session.vim *~ # auto-generated tag files tags + +# Created by https://www.gitignore.io/api/go + +### Go ### +# Binaries for programs and plugins +*.exe +*.dll +*.so +*.dylib + +# Test binary, build with `go test -c` +*.test + +# Output of the go coverage tool, specifically when used with LiteIDE +*.out + +# Project-local glide cache, RE: https://github.com/Masterminds/glide/issues/736 +.glide/ + +# End of https://www.gitignore.io/api/go + +# custom +# external packages folder +vendor/ diff --git a/glide.lock b/glide.lock index 219ba3f..c884078 100644 --- a/glide.lock +++ b/glide.lock @@ -1,8 +1,8 @@ -hash: c178dbd27e8a90c66854415d5bc65d76be516a5033ea92216cd1d0f32f709373 -updated: 2017-03-20T14:16:53.247369041-05:00 +hash: 9f381497df21cd2b2e97a96a4526bb2f48443a7d63baaf81e306f4b6c3cdad98 +updated: 2017-03-23T00:24:46.690555736-05:00 imports: - name: github.com/aws/aws-sdk-go - version: 5b99715ae2945a2434a2371f4e6c5542e839a32d + version: 72df63b404d3f9820db08c73176c1b277d9f614f subpackages: - aws - aws/awserr diff --git a/glide.yaml b/glide.yaml index f865d0d..afcdda7 100644 --- a/glide.yaml +++ b/glide.yaml @@ -1,7 +1,7 @@ package: github.com/rewardStyle/kinetic import: - package: github.com/aws/aws-sdk-go - version: 5b99715ae2945a2434a2371f4e6c5542e839a32d + version: v1.8.1 subpackages: - aws - aws/awserr diff --git a/listener/config.go b/listener/config.go new file mode 100644 index 0000000..d574256 --- /dev/null +++ b/listener/config.go @@ -0,0 +1,112 @@ +package listener + +import ( + "net/http" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/credentials" + "github.com/aws/aws-sdk-go/aws/session" +) + +const () + +type Config struct { + awsConfig *aws.Config + + stream string + + logLevel aws.LogLevelType + stats StatsListener + + getRecordsReadTimeout time.Duration + + shard string + batchSize int + concurrency int + shardIterator *ShardIterator +} + +func NewConfig(stream string) *Config { + return &Config{ + awsConfig: aws.NewConfig().WithHTTPClient(&http.Client{ + Timeout: 5 * time.Minute, + }), + stream: stream, + logLevel: aws.LogOff, + stats: &NilStatsListener{}, + batchSize: 10000, + concurrency: 10000, + shardIterator: NewShardIterator(), + getRecordsReadTimeout: 1 * time.Second, + } +} + +func (c *Config) WithCredentials(accessKey, secretKey, securityToken string) *Config { + c.awsConfig.WithCredentials( + credentials.NewStaticCredentials(accessKey, secretKey, securityToken), + ) + return c +} + +func (c *Config) WithRegion(region string) *Config { + c.awsConfig.WithRegion(region) + return c +} + +func (c *Config) WithEndpoint(endpoint string) *Config { + c.awsConfig.WithEndpoint(endpoint) + return c +} + +func (c *Config) WithLogger(logger aws.Logger) *Config { + c.awsConfig.WithLogger(logger) + return c +} + +func (c *Config) WithLogLevel(logLevel aws.LogLevelType) *Config { + c.awsConfig.WithLogLevel(logLevel & 0xffff) + c.logLevel = logLevel >> 16 + return c +} + +func (c *Config) WithHttpClientTimeout(timeout time.Duration) *Config { + c.awsConfig.WithHTTPClient(&http.Client{ + Timeout: timeout, + }) + return c +} + +func (c *Config) WithStatsListener(stats StatsListener) *Config { + c.stats = stats + return c +} + +func (c *Config) WithShardId(shard string) *Config { + c.shard = shard + return c +} + +func (c *Config) WithBatchSize(batchSize int) *Config { + c.batchSize = batchSize + return c +} + +func (c *Config) WithConcurrency(concurrency int) *Config { + c.concurrency = concurrency + return c +} + +func (c *Config) WithInitialShardIterator(shardIterator *ShardIterator) *Config { + c.shardIterator = shardIterator + return c +} + +func (c *Config) WithGetRecordsReadTimeout(timouet time.Duration) *Config { + c.getRecordsReadTimeout = timouet + return c +} + +func (c *Config) GetAwsSession() (*session.Session, error) { + return session.NewSession(c.awsConfig) +} diff --git a/listener/listener.go b/listener/listener.go new file mode 100644 index 0000000..aa3e131 --- /dev/null +++ b/listener/listener.go @@ -0,0 +1,752 @@ +package listener + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "sync" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/awserr" + "github.com/aws/aws-sdk-go/aws/request" + "github.com/aws/aws-sdk-go/service/kinesis" + "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" + + "github.com/rewardStyle/kinetic" +) + +var ( + // errors tha can occur in Retrieve / RetrieveFn (and their WithContext variants) + ErrAlreadyConsuming = errors.New("Listener already consuming. Only one Listen, Retrieve, or RetrieveFn may be active at a time") + + // errors that can occur in GetShards + ErrNilDescribeStreamResponse = errors.New("DescribeStream returned a nil response") + ErrNilStreamDescription = errors.New("DescribeStream returned a nil StreamDescription") + + // errors that can occur in SetShard + ErrCannotSetShard = errors.New("Cannot set shard while consuming") + + // errors that can occur in fetchBatch + ErrEmptySequenceNumber = errors.New("Attempted to set sequence number with empty value") + ErrEmptyShardIterator = errors.New("Attempted to set shard iterator with empty value") + ErrNilGetShardIteratorResponse = errors.New("GetShardIteratore returned a nil response") + ErrNilShardIterator = errors.New("GetShardIterator returned a nil ShardIterator") + ErrNilGetRecordsResponse = errors.New("GetRecords returned an nil response") + ErrTimeoutReadResponseBody = errors.New("Timeout while reading response body") + + ErrPipeOfDeath = errors.New("Received pipe of death") +) + +type Empty struct{} + +type MessageFn func([]byte, *sync.WaitGroup) + +type ShardIterator struct { + shardIteratorType string + sequenceNumber string + timestamp time.Time +} + +func NewShardIterator() *ShardIterator { + return &ShardIterator{ + shardIteratorType: "TRIM_HORIZON", + } +} + +func (it *ShardIterator) TrimHorizon() *ShardIterator { + it.shardIteratorType = "TRIM_HORIZON" + return it +} + +func (it *ShardIterator) Latest() *ShardIterator { + it.shardIteratorType = "LATEST" + return it +} + +func (it *ShardIterator) AtSequenceNumber(sequenceNumber string) *ShardIterator { + it.shardIteratorType = "AT_SEQUENCE_NUMBER" + it.sequenceNumber = sequenceNumber + return it +} + +func (it *ShardIterator) AfterSequenceNumber(sequenceNumber string) *ShardIterator { + it.shardIteratorType = "AFTER_SEQUENCE_NUMBER" + it.sequenceNumber = sequenceNumber + return it +} + +func (it *ShardIterator) AtTimestamp(timestamp time.Time) *ShardIterator { + it.shardIteratorType = "AT_TIMESTAMP" + it.timestamp = timestamp + return it +} + +func (it *ShardIterator) getStartingSequenceNumber() *string { + if it.sequenceNumber == "" { + return nil + } + return aws.String(it.sequenceNumber) +} + +func (it *ShardIterator) getTimestamp() *time.Time { + if it.timestamp.IsZero() { + return nil + } + return aws.Time(it.timestamp) +} + +type Listener struct { + Config *Config + + nextShardIterator string + + messages chan *kinetic.Message + concurrencySem chan Empty + throttleSem chan Empty + pipeOfDeath chan Empty + wg sync.WaitGroup + + consuming bool + consumingMu sync.Mutex + + client kinesisiface.KinesisAPI + clientMu sync.Mutex +} + +// NewListener creates a new listener for listening to message on a Kinesis +// stream. +func NewListener(config *Config) (*Listener, error) { + l := &Listener{ + Config: config, + concurrencySem: make(chan Empty, config.concurrency), + throttleSem: make(chan Empty, 5), + pipeOfDeath: make(chan Empty), + } + return l, nil +} + +// Logs a debug message using the AWS SDK logger. +func (l *Listener) Log(args ...interface{}) { + l.ensureClient() + if l.client != nil && l.Config.logLevel.AtLeast(aws.LogDebug) { + l.Config.awsConfig.Logger.Log(args...) + } +} + +// CreateStream creates the stream. +func (l *Listener) CreateStream(shards int) error { + if err := l.ensureClient(); err != nil { + return err + } + _, err := l.client.CreateStream(&kinesis.CreateStreamInput{ + StreamName: aws.String(l.Config.stream), + ShardCount: aws.Int64(int64(shards)), + }) + if err != nil { + l.Log("Error creating stream:", err) + } + return err +} + +// WaitUntilActive waits until the stream is active. Timeouts can be set via +// context. Defaults to 18 retries with a 10s delay (180s or 3 minutes total). +func (l *Listener) WaitUntilActive(ctx context.Context, opts ...request.WaiterOption) error { + if err := l.ensureClient(); err != nil { + return err + } + return l.client.WaitUntilStreamExistsWithContext(ctx, &kinesis.DescribeStreamInput{ + StreamName: aws.String(l.Config.stream), // Required + }, opts...) +} + +// DeleteStream deletes the stream. +func (l *Listener) DeleteStream() error { + if err := l.ensureClient(); err != nil { + return err + } + _, err := l.client.DeleteStream(&kinesis.DeleteStreamInput{ + StreamName: aws.String(l.Config.stream), + }) + if err != nil { + l.Log("Error deleting stream:", err) + } + return err +} + +// WaitUntilDeleted waits until the stream is does not exist. Timeouts can be +// set via context. Defaults to 18 retries with a 10s delay (180s or 3 minutes +// total). +func (l *Listener) WaitUntilDeleted(ctx context.Context, opts ...request.WaiterOption) error { + if err := l.ensureClient(); err != nil { + return err + } + w := request.Waiter{ + Name: "WaitUntilStreamIsDeleted", + MaxAttempts: 18, + Delay: request.ConstantWaiterDelay(10 * time.Second), + Acceptors: []request.WaiterAcceptor{ + { + State: request.SuccessWaiterState, + Matcher: request.ErrorWaiterMatch, + Expected: kinesis.ErrCodeResourceNotFoundException, + }, + }, + Logger: l.Config.awsConfig.Logger, + NewRequest: func(opts []request.Option) (*request.Request, error) { + req, _ := l.client.DescribeStreamRequest(&kinesis.DescribeStreamInput{ + StreamName: aws.String(l.Config.stream), // Required + }) + req.SetContext(ctx) + req.ApplyOptions(opts...) + return req, nil + }, + } + w.ApplyOptions(opts...) + return w.WaitWithContext(ctx) +} + +// GetShards gets a list of shards in a stream. +func (l *Listener) GetShards() ([]string, error) { + if err := l.ensureClient(); err != nil { + return nil, err + } + resp, err := l.client.DescribeStream(&kinesis.DescribeStreamInput{ + StreamName: aws.String(l.Config.stream), + }) + if err != nil { + l.Log("Error describing stream", err) + return nil, err + } + if resp == nil { + return nil, ErrNilDescribeStreamResponse + } + if resp.StreamDescription == nil { + return nil, ErrNilStreamDescription + } + var shards []string + for _, shard := range resp.StreamDescription.Shards { + if shard.ShardId != nil { + shards = append(shards, aws.StringValue(shard.ShardId)) + } + } + return shards, nil +} + +// SetShard sets the shard for the listener. +func (l *Listener) SetShard(shard string) error { + if !l.blockConsumers() { + return ErrCannotSetShard + } + defer l.allowConsumers() + l.Config.shard = shard + return nil +} + +// setNextShardIterator sets the nextShardIterator to use when calling +// GetRecords. +// +// Not thread-safe. Only called from fetchBatch (and ensureShardIterator, +// which is called from fetchBatch). Care must be taken to ensure that only +// one call to Listen and Retrieve/RetrieveFn can be running at a time. +func (l *Listener) setNextShardIterator(shardIterator string) error { + if len(shardIterator) == 0 { + return ErrEmptyShardIterator + } + l.nextShardIterator = shardIterator + return nil +} + +// setSequenceNumber sets the sequenceNumber of shardIterator to the last +// delivered message and updates the shardIteratorType to AT_SEQUENCE_NUMBER. +// This is only used when we need to call getShardIterator (say, to refresh the +// shard iterator). +// +// Not thread-safe. Only called from fetchBatch. Care must be taken to ensure +// that only one call to Listen and Retrieve/RetrieveFn can be running at a +// time. +func (l *Listener) setSequenceNumber(sequenceNumber string) error { + if len(sequenceNumber) == 0 { + return ErrEmptySequenceNumber + } + l.Config.shardIterator.AtSequenceNumber(sequenceNumber) + return nil +} + +// ensureClient will lazily make sure we have an AWS Kinesis client. +func (l *Listener) ensureClient() error { + // From the aws-go-sdk documentation: + // http://docs.aws.amazon.com/sdk-for-go/api/aws/session/ + // + // Concurrency: + // Sessions are safe to use concurrently as long as the Session is not + // being modified. The SDK will not modify the Session once the Session + // has been created. Creating service clients concurrently from a + // shared Session is safe. + // + // We need to think through the impact of creating a new client (for + // example, after receiving an error from Kinesis) while there may be + // outstanding goroutines still processing messages. My cursory thought + // is that this is safe to do, as any outstanding messages will likely + // not interact with the Kinesis stream. At worst, we would need a lock + // around the ensureClient method to make sure that no two goroutines + // are trying to ensure the client at the same time. + // + // As we don't expose any methods (or in fact, even the Listener object + // itself) to the client through the API, I don't forsee needing to add + // this lock unless something dramatically changes about the design of + // this library. + l.clientMu.Lock() + defer l.clientMu.Unlock() + if l.client != nil { + return nil + } + + session, err := l.Config.GetAwsSession() + if err != nil { + return err + } + l.client = kinesis.New(session) + return nil +} + +// ensureShardIterator will lazily make sure that we have a valid ShardIterator, +// calling the GetShardIterator API with the configured ShardIteratorType (with +// any applicable StartingSequenceNumber or Timestamp) if necessary. +// +// Not thread-safe. Only called from fetchBatch Care must be taken to ensure +// that only one call to Listen and Retrieve/RetrieveFn can be running at a +// time. +func (l *Listener) ensureShardIterator() error { + if l.nextShardIterator != "" { + return nil + } + + resp, err := l.client.GetShardIterator(&kinesis.GetShardIteratorInput{ + ShardId: aws.String(l.Config.shard), // Required + ShardIteratorType: aws.String(l.Config.shardIterator.shardIteratorType), // Required + StreamName: aws.String(l.Config.stream), // Required + StartingSequenceNumber: l.Config.shardIterator.getStartingSequenceNumber(), + Timestamp: l.Config.shardIterator.getTimestamp(), + }) + if err != nil { + l.Log(err) + return err + } + if resp == nil { + return ErrNilGetShardIteratorResponse + } + if resp.ShardIterator == nil { + return ErrNilShardIterator + } + return l.setNextShardIterator(*resp.ShardIterator) +} + +// 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(sem chan Empty) { + sem <- Empty{} + time.AfterFunc(1*time.Second, func() { + <-sem + }) +} + +// fetchBatch calls GetRecords and delivers each record into the messages +// channel. +// TODO: Convert timeout implementation to use context.Context +// FIXME: Need to investigate that the timeout implementation doesn't result in +// an fd leak. Since we call Read on the HTTPResonse.Body in a select with a +// timeout channel, we do prevent ourself from blocking. Once we timeout, we +// return an error to the outer ioutil.ReadAll, which should result in a call +// to our io.ReadCloser's Close function. This will in turn call Close on the +// underlying HTTPResponse.Body. The question is whether this actually shuts +// down the TCP connection. Worst case scenario is that our client Timeout +// eventually fires and closes the socket, but this can be susceptible to FD +// exhaustion. +func (l *Listener) fetchBatch(size int) (int, error) { + if err := l.ensureClient(); err != nil { + return 0, err + } + + if err := l.ensureShardIterator(); err != nil { + return 0, err + } + + l.throttle(l.throttleSem) + + // We use the GetRecordsRequest method of creating requests to allow for + // registering custom handlers for better control over the API request. + var startReadTime time.Time + var startUnmarshalTime time.Time + start := time.Now() + req, resp := l.client.GetRecordsRequest(&kinesis.GetRecordsInput{ + Limit: aws.Int64(int64(size)), + ShardIterator: aws.String(l.nextShardIterator), + }) + + // If debug is turned on, add some handlers for GetRecords logging + if l.Config.logLevel.AtLeast(aws.LogDebug) { + req.Handlers.Send.PushBack(func(r *request.Request) { + l.Log("Finished GetRecords Send, took", time.Since(start)) + }) + } + + // Here, we insert a handler to be called after the Send handler and + // before the the Unmarshal handler in the aws-go-sdk library. + // + // The Send handler will call http.Client.Do() on the request, which + // blocks until the response headers have been read before returning an + // HTTPResponse. + // + // The Unmarshal handler will ultimately call ioutil.ReadAll() on the + // HTTPResponse.Body stream. + // + // Our handler wraps the HTTPResponse.Body with our own ReadCloser so + // that we can implement a timeout mechanism on the Read() call (which + // is called by the ioutil.ReadAll() function) + req.Handlers.Unmarshal.PushFront(func(r *request.Request) { + l.Log("Started GetRecords Unmarshal, took", time.Since(start)) + // Here, we set a timer that the initial Read() call on + // HTTPResponse.Body must return by. Note that the normal + // http.Client Timeout is still in effect. + startReadTime = time.Now() + timer := time.NewTimer(l.Config.getRecordsReadTimeout) + + r.HTTPResponse.Body = &kinetic.TimeoutReadCloser{ + ReadCloser: r.HTTPResponse.Body, + OnReadFn: func(stream io.ReadCloser, b []byte) (n int, err error) { + // The OnReadFn will be called each time + // ioutil.ReadAll calls Read on the + // TimeoutReadCloser. + + // First, we set up a struct that to hold the + // results of the Read() call that can go + // through a channel + type Result struct { + n int + err error + } + + // Next, we build a channel with which to pass + // the Read() results + c := make(chan Result, 1) + + // Now, we call the Read() on the + // HTTPResponse.Body in a goroutine and feed the + // results into the channel + readStart := time.Now() + go func() { + var result Result + result.n, result.err = stream.Read(b) + c <- result + }() + + // Finally, we poll for the Read() to complete + // or the timer to elapse. + select { + case result := <-c: + // If we sucessfully Read() from the + // HTTPResponse.Body, we reset our + // timeout and return the results from + // the Read() + timer.Reset(l.Config.getRecordsReadTimeout) + n, err = result.n, result.err + l.Log(fmt.Sprintf("DEBUG: read %d bytes, took %v", n, time.Since(readStart))) + case <-timer.C: + // If we timeout, we return an error + // that will unblock ioutil.ReadAll(). + // This will cause the Unmarshal handler + // to return an error. This error will + // propogate to the original req.Send() + // call (below) + l.Log(fmt.Sprintf("DEBUG: read timed out after %v", time.Since(readStart))) + err = ErrTimeoutReadResponseBody + } + return + }, + OnCloseFn: func() { + l.Config.stats.AddGetRecordsReadResponseTime(time.Since(startReadTime)) + l.Log("Finished GetRecords body read, took", time.Since(start)) + startUnmarshalTime = time.Now() + }, + } + }) + + req.Handlers.Unmarshal.PushBack(func(r *request.Request) { + l.Config.stats.AddGetRecordsUnmarshalTime(time.Since(startUnmarshalTime)) + l.Log("Finished GetRecords Unmarshal, took", time.Since(start)) + }) + + // Send the GetRecords request + l.Log("Starting GetRecords build/sign request, took", time.Since(start)) + l.Config.stats.AddGetRecordsCalled(1) + if err := req.Send(); err != nil { + l.Log("Error getting records:", err) + return 0, err + } + + // Process Records + l.Log(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), l.Config.shard, time.Since(start))) + if resp == nil { + return 0, ErrNilGetRecordsResponse + } + delivered := 0 + l.Config.stats.AddBatchSizeSample(len(resp.Records)) + for _, record := range resp.Records { + if record != nil { + delivered++ + l.messages <- &kinetic.Message{*record} + l.Config.stats.AddConsumedSample(1) + } + if record.SequenceNumber != nil { + // We can safely ignore if this call returns + // error, as if we somehow receive an empty + // sequence number from AWS, we will simply not + // set it. At worst, this causes us to + // reprocess this record if we happen to refresh + // the iterator. + l.setSequenceNumber(*record.SequenceNumber) + } + } + if resp.NextShardIterator != nil { + // TODO: According to AWS docs: + // http://docs.aws.amazon.com/sdk-for-go/api/service/kinesis/#GetRecordsOutput + // + // NextShardIterator: The next position in the shard + // from which to start sequentially reading data + // records. If set to null, the shard has been closed + // and the requested iterator will not return any more + // data. + // + // When dealing with streams that will merge or split, + // we need to detect that the shard has closed and + // notify the client library. + // + // TODO: I don't know if we should be ignoring an error returned + // by setShardIterator in case of an empty shard iterator in the + // response. There isn't much we can do, and the best path for + // recovery may be simply to reprocess the batch and see if we + // get a valid NextShardIterator from AWS the next time around. + l.setNextShardIterator(*resp.NextShardIterator) + } + return delivered, nil +} + +// blockConsumers will set consuming to true if there is not already another +// consume loop running. +func (l *Listener) blockConsumers() bool { + l.consumingMu.Lock() + defer l.consumingMu.Unlock() + if !l.consuming { + l.consuming = true + return true + } + return false +} + +// startConsuming handles any initialization needed in preparation to start +// consuming. +func (l *Listener) startConsuming() { + l.messages = make(chan *kinetic.Message, l.Config.batchSize) +} + +// shouldConsume is a convenience function that allows functions to break their +// loops if the context receives a cancellation. +func (l *Listener) shouldConsume(ctx context.Context) (bool, error) { + select { + case <-l.pipeOfDeath: + return false, ErrPipeOfDeath + case <-ctx.Done(): + return false, ctx.Err() + default: + return true, nil + } +} + +// stopConsuming handles any cleanup after a consuming has stopped. +func (l *Listener) stopConsuming() { + close(l.messages) +} + +// allowConsumers allows consuming. Called after blockConsumers to release the +// lock on consuming. +func (l *Listener) allowConsumers() { + l.consumingMu.Lock() + defer l.consumingMu.Unlock() + l.consuming = false +} + +// IsConsuming returns true while consuming. +func (l *Listener) IsConsuming() bool { + l.consumingMu.Lock() + defer l.consumingMu.Unlock() + return l.consuming +} + +// Retrieve waits for a message from the stream and return the value. +// Cancellation supported through contexts. +func (l *Listener) RetrieveWithContext(ctx context.Context) (*kinetic.Message, error) { + if !l.blockConsumers() { + return nil, ErrAlreadyConsuming + } + l.startConsuming() + defer func() { + l.stopConsuming() + l.allowConsumers() + }() + for { + ok, err := l.shouldConsume(ctx) + if !ok { + return nil, err + } + n, err := l.fetchBatch(1) + if err != nil { + return nil, err + } + if n > 0 { + l.Config.stats.AddDeliveredSample(1) + return <-l.messages, nil + } + } +} + +// Retrieve waits for a message from the stream and return the value. +func (l *Listener) Retrieve() (*kinetic.Message, error) { + return l.RetrieveWithContext(context.TODO()) +} + +// RetrieveFn retrieves a message from the stream and dispatches it to the +// supplied function. RetrieveFn will wait until the function completes. +// Cancellation supported through context. +func (l *Listener) RetrieveFnWithContext(ctx context.Context, fn MessageFn) error { + msg, err := l.RetrieveWithContext(ctx) + if err != nil { + return err + } + var wg sync.WaitGroup + wg.Add(1) + go fn(msg.Value(), &wg) + wg.Wait() + l.Config.stats.AddProcessedSample(1) + return nil +} + +// RetrieveFn retrieves a message from the stream and dispatches it to the +// supplied function. RetrieveFn will wait until the function completes. +func (l *Listener) RetrieveFn(fn MessageFn) error { + return l.RetrieveFnWithContext(context.TODO(), fn) +} + +// consume calls fetchBatch with configured batch size in a loop until the +// listener is stopped. +func (l *Listener) consume(ctx context.Context) { + // We need to run blockConsumers & startConsuming to make sure that we + // are okay and ready to start consuming. This is mainly to avoid a + // race condition where Listen() will attempt to read the messages + // channel prior to consume() initializing it. We can then launch a + // goroutine to handle the actual consume operation. + if !l.blockConsumers() { + return + } + l.startConsuming() + go func() { + defer func() { + l.stopConsuming() + l.allowConsumers() + }() + stop: + for { + ok, err := l.shouldConsume(ctx) + if !ok { + break stop + } + _, err = l.fetchBatch(l.Config.batchSize) + + if err != nil { + switch err := err.(type) { + case net.Error: + if err.Timeout() { + l.Config.stats.AddGetRecordsTimeout(1) + l.Log("Received net error:", err.Error()) + } else { + l.Log("Received unknown net error:", err.Error()) + } + case error: + switch err { + case ErrTimeoutReadResponseBody: + l.Config.stats.AddGetRecordsReadTimeout(1) + l.Log("Received error:", err.Error()) + case ErrEmptySequenceNumber: + fallthrough + case ErrEmptyShardIterator: + fallthrough + case ErrNilGetShardIteratorResponse: + fallthrough + case ErrNilShardIterator: + fallthrough + case ErrNilGetRecordsResponse: + fallthrough + default: + l.Log("Received error:", err.Error()) + } + case awserr.Error: + switch err.Code() { + case kinesis.ErrCodeProvisionedThroughputExceededException: + l.Config.stats.AddProvisionedThroughputExceeded(1) + case kinesis.ErrCodeResourceNotFoundException: + fallthrough + case kinesis.ErrCodeInvalidArgumentException: + fallthrough + case kinesis.ErrCodeExpiredIteratorException: + fallthrough + default: + l.Log("Received AWS error:", err.Error()) + } + default: + l.Log("Received unknown error:", err.Error()) + } + } + } + }() +} + +// Listen listens and delivers message to the supplied function. Upon +// cancellation, Listen will stop the consumer loop and wait until the messages +// channel is closed and all messages are delivered. +func (l *Listener) ListenWithContext(ctx context.Context, fn MessageFn) { + l.consume(ctx) + var wg sync.WaitGroup + defer wg.Wait() +stop: + for { + select { + case msg, ok := <-l.messages: + // Listen should always run until it the messages + // channel closes. + if !ok { + break stop + } + l.Config.stats.AddDeliveredSample(1) + l.concurrencySem <- Empty{} + wg.Add(1) + go func() { + defer func() { + <-l.concurrencySem + }() + var fnWg sync.WaitGroup + fnWg.Add(1) + fn(msg.Value(), &fnWg) + fnWg.Wait() + l.Config.stats.AddProcessedSample(1) + wg.Done() + }() + } + } +} + +// Listen listens and delivers message to the supplied function. +func (l *Listener) Listen(fn MessageFn) { + l.ListenWithContext(context.TODO(), fn) +} diff --git a/listener/listener_test.go b/listener/listener_test.go new file mode 100644 index 0000000..aff75fe --- /dev/null +++ b/listener/listener_test.go @@ -0,0 +1,527 @@ +package listener + +import ( + . "github.com/smartystreets/goconvey/convey" + + "context" + "fmt" + "math/rand" + "strconv" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/awserr" + "github.com/aws/aws-sdk-go/aws/request" + "github.com/aws/aws-sdk-go/aws/session" + "github.com/aws/aws-sdk-go/service/kinesis" +) + +type DebugLogger struct { + logs []string +} + +func (l *DebugLogger) Log(args ...interface{}) { + l.logs = append(l.logs, fmt.Sprint(args...)) +} + +func (l *DebugLogger) Reset() { + l.logs = nil +} + +type DebugStatsListener struct{} + +func (l *DebugStatsListener) AddConsumedSample(int) {} +func (l *DebugStatsListener) AddDeliveredSample(int) {} +func (l *DebugStatsListener) AddProcessedSample(int) {} +func (l *DebugStatsListener) AddBatchSizeSample(int) {} +func (l *DebugStatsListener) AddGetRecordsCalled(int) {} +func (l *DebugStatsListener) AddProvisionedThroughputExceeded(int) {} +func (l *DebugStatsListener) AddGetRecordsTimeout(int) {} +func (l *DebugStatsListener) AddGetRecordsReadTimeout(int) {} +func (l *DebugStatsListener) AddGetRecordsReadResponseTime(time.Duration) {} +func (l *DebugStatsListener) AddGetRecordsUnmarshalTime(time.Duration) {} + +func getSession(config *Config) *session.Session { + sess, err := config.GetAwsSession() + So(err, ShouldBeNil) + So(sess, ShouldNotBeNil) + return sess +} + +func putRecord(l *Listener, b []byte) (*string, error) { + resp, err := l.client.PutRecord(&kinesis.PutRecordInput{ + Data: b, + PartitionKey: aws.String("dummy"), + StreamName: aws.String(l.Config.stream), + }) + if err != nil { + return nil, err + } + return resp.SequenceNumber, nil +} + +func TestNewConfig(t *testing.T) { + Convey("given a new listener config", t, func() { + config := NewConfig("some-stream") + + Convey("check the default values for its non-zero config", func() { + So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 5*time.Minute) + So(config.stream, ShouldEqual, "some-stream") + So(config.logLevel, ShouldEqual, aws.LogOff) + So(config.stats, ShouldHaveSameTypeAs, &NilStatsListener{}) + So(config.batchSize, ShouldEqual, 10000) + So(config.concurrency, ShouldEqual, 10000) + So(config.shardIterator.shardIteratorType, ShouldEqual, "TRIM_HORIZON") + So(config.getRecordsReadTimeout, ShouldEqual, 1*time.Second) + }) + + Convey("check that we can retrieve an aws.Session from it ", func() { + getSession(config) + }) + + Convey("check that we can set credentials", func() { + config = config.WithCredentials("access-key", "secret-key", "security-token") + sess := getSession(config) + creds, err := sess.Config.Credentials.Get() + So(err, ShouldBeNil) + So(creds.AccessKeyID, ShouldEqual, "access-key") + So(creds.SecretAccessKey, ShouldEqual, "secret-key") + So(creds.SessionToken, ShouldEqual, "security-token") + }) + + Convey("check that we can set the region", func() { + config = config.WithRegion("my-region") + sess := getSession(config) + So(aws.StringValue(sess.Config.Region), ShouldEqual, "my-region") + }) + + Convey("check that we can set the endpoint", func() { + config = config.WithEndpoint("my-endpoint") + sess := getSession(config) + So(aws.StringValue(sess.Config.Endpoint), ShouldEqual, "my-endpoint") + }) + + Convey("check that we can configure a logger", func() { + l := &DebugLogger{} + config = config.WithLogger(l) + sess := getSession(config) + So(sess.Config.Logger, ShouldHaveSameTypeAs, l) + + Convey("check that basic logging should work", func() { + sess.Config.Logger.Log("one") + sess.Config.Logger.Log("two") + sess.Config.Logger.Log("three") + So(len(l.logs), ShouldEqual, 3) + So(l.logs, ShouldContain, "one") + So(l.logs, ShouldContain, "two") + So(l.logs, ShouldContain, "three") + Reset(func() { + l.Reset() + }) + }) + }) + + Convey("check that the default log level is off for both the sdk and kinetic", func() { + sess := getSession(config) + So(sess.Config.LogLevel.Value(), ShouldEqual, aws.LogOff) + So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeFalse) + So(config.logLevel.Value(), ShouldEqual, aws.LogOff) + So(config.logLevel.AtLeast(aws.LogDebug), ShouldBeFalse) + }) + + Convey("check that we can set the sdk log level", func() { + config = config.WithLogLevel(aws.LogDebug | aws.LogDebugWithSigning) + sess := getSession(config) + So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) + So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) + So(config.logLevel.Value(), ShouldEqual, aws.LogOff) + So(config.logLevel.AtLeast(aws.LogDebug), ShouldBeFalse) + }) + + Convey("check that we can set the kinetic log level", func() { + config = config.WithLogLevel((aws.LogDebug | 1) << 16) + sess := getSession(config) + So(sess.Config.LogLevel.Value(), ShouldEqual, aws.LogOff) + So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeFalse) + So(config.logLevel.AtLeast(aws.LogDebug), ShouldBeTrue) + So(config.logLevel.Matches(1), ShouldBeTrue) + }) + + Convey("check that we can set both the sdk and kinetic log level", func() { + config = config.WithLogLevel(aws.LogDebug | aws.LogDebugWithSigning | ((aws.LogDebug | 1) << 16)) + sess := getSession(config) + So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) + So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) + So(config.logLevel.AtLeast(aws.LogDebug), ShouldBeTrue) + So(config.logLevel.Matches(1), ShouldBeTrue) + }) + + Convey("check that we can set the http.Client Timeout", func() { + config = config.WithHttpClientTimeout(10 * time.Minute) + So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 10*time.Minute) + }) + + Convey("check that we can configure a stats listener", func() { + config = config.WithStatsListener(&DebugStatsListener{}) + So(config.stats, ShouldHaveSameTypeAs, &DebugStatsListener{}) + }) + + Convey("check that we can set the shard", func() { + config = config.WithShardId("some-shard") + So(config.shard, ShouldEqual, "some-shard") + }) + + Convey("check that we can set the batch size", func() { + config = config.WithBatchSize(1000) + So(config.batchSize, ShouldEqual, 1000) + }) + + Convey("check that we can set the concurrency limit", func() { + config = config.WithConcurrency(50) + So(config.concurrency, ShouldEqual, 50) + }) + + Convey("check that the default shard iterator is TRIM_HORIZON", func() { + config = config.WithInitialShardIterator(NewShardIterator()) + So(config.shardIterator.shardIteratorType, ShouldEqual, "TRIM_HORIZON") + So(config.shardIterator.getStartingSequenceNumber(), ShouldBeNil) + So(config.shardIterator.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can set the initial shard iterator (to LATEST)", func() { + config = config.WithInitialShardIterator(NewShardIterator().Latest()) + So(config.shardIterator.shardIteratorType, ShouldEqual, "LATEST") + So(config.shardIterator.getStartingSequenceNumber(), ShouldBeNil) + So(config.shardIterator.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can set the read timeout for the GetRecords request", func() { + config = config.WithGetRecordsReadTimeout(10 * time.Second) + So(config.getRecordsReadTimeout, ShouldEqual, 10*time.Second) + }) + }) +} + +func TestShardIterator(t *testing.T) { + Convey("given a new shard iterator", t, func() { + it := NewShardIterator() + + Convey("check that the default shard iterator type is TRIM_HORIZON", func() { + So(it.shardIteratorType, ShouldEqual, "TRIM_HORIZON") + So(it.getStartingSequenceNumber(), ShouldBeNil) + So(it.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can explicitly set it to TRIM_HORIZON", func() { + it = it.TrimHorizon() + So(it.shardIteratorType, ShouldEqual, "TRIM_HORIZON") + So(it.getStartingSequenceNumber(), ShouldBeNil) + So(it.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can explicitly set it to LATEST", func() { + it = it.Latest() + So(it.shardIteratorType, ShouldEqual, "LATEST") + So(it.getStartingSequenceNumber(), ShouldBeNil) + So(it.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can explicitly set it to AT_SEQEUENCE_NUMBER", func() { + it = it.AtSequenceNumber("some-sequence") + So(it.shardIteratorType, ShouldEqual, "AT_SEQUENCE_NUMBER") + So(aws.StringValue(it.getStartingSequenceNumber()), ShouldEqual, "some-sequence") + So(it.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can explicitly set it to AFTER_SEQEUENCE_NUMBER", func() { + it = it.AfterSequenceNumber("some-sequence") + So(it.shardIteratorType, ShouldEqual, "AFTER_SEQUENCE_NUMBER") + So(aws.StringValue(it.getStartingSequenceNumber()), ShouldEqual, "some-sequence") + So(it.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can explicitly set it to AT_TIMESTAMP", func() { + n := time.Now() + it = it.AtTimestamp(n) + So(it.shardIteratorType, ShouldEqual, "AT_TIMESTAMP") + So(aws.TimeValue(it.getTimestamp()).Equal(n), ShouldBeTrue) + So(it.getStartingSequenceNumber(), ShouldBeNil) + }) + }) +} + +func TestListener(t *testing.T) { + Convey("given a listener", t, func() { + logger := &DebugLogger{} + l, err := NewListener(NewConfig("some-stream"). + WithRegion("some-region"). + WithEndpoint("http://127.0.0.1:4567"). + WithLogger(logger). + WithLogLevel(aws.LogDebug << 16). + WithConcurrency(10)) + So(l, ShouldNotBeNil) + So(err, ShouldBeNil) + + Convey("check that logging works", func() { + l.Log("foo") + So(logger.logs, ShouldContain, "foo") + So(len(logger.logs), ShouldEqual, 1) + Reset(func() { logger.Reset() }) + }) + + Convey("check that calling ensureClient twice doesn't overwrite existing client", func() { + So(l.client, ShouldBeNil) + l.ensureClient() + So(l.client, ShouldNotBeNil) + client := l.client + l.ensureClient() + So(l.client, ShouldEqual, client) + }) + + Convey("check deleting a non-existent stream returns an error", func() { + err := l.DeleteStream() + So(err, ShouldNotBeNil) + e := err.(awserr.Error) + So(e.Code(), ShouldEqual, kinesis.ErrCodeResourceNotFoundException) + }) + + Convey("check that setting an empty shard iterator returns an error", func() { + err := l.setNextShardIterator("") + So(err, ShouldEqual, ErrEmptyShardIterator) + }) + + Convey("check that setting an empty sequence number returns an error", func() { + err := l.setSequenceNumber("") + So(err, ShouldEqual, ErrEmptySequenceNumber) + }) + + Convey("check getting shards on a non-existent stream returns an error", func() { + shards, err := l.GetShards() + So(shards, ShouldBeNil) + So(err, ShouldNotBeNil) + e := err.(awserr.Error) + So(e.Code(), ShouldEqual, kinesis.ErrCodeResourceNotFoundException) + }) + + Convey("check that we can create a stream", func() { + err := l.CreateStream(1) + So(err, ShouldBeNil) + err = l.WaitUntilActive(context.TODO(), request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + So(err, ShouldBeNil) + + Convey("check that creating an existing stream returns an error", func() { + err := l.CreateStream(1) + So(err, ShouldNotBeNil) + e := err.(awserr.Error) + So(e.Code(), ShouldEqual, kinesis.ErrCodeResourceInUseException) + }) + + Convey("check that getting the shard iterator with an empty shard returns an error", func() { + So(l.Config.shard, ShouldBeEmpty) + err := l.ensureShardIterator() + So(err, ShouldNotBeNil) + e := err.(awserr.Error) + So(e.Code(), ShouldEqual, request.InvalidParameterErrCode) + }) + + Convey("check that getting the shard iterator without an invalid shard returns an error", func() { + l.SetShard("some-shard") + err := l.ensureShardIterator() + So(err, ShouldNotBeNil) + e := err.(awserr.Error) + So(e.Code(), ShouldEqual, kinesis.ErrCodeResourceNotFoundException) + }) + + Convey("check that we can obtain a list of shards", func() { + shards, err := l.GetShards() + So(err, ShouldBeNil) + So(len(shards), ShouldEqual, 1) + + Convey("check that we can set the shard id", func() { + l.SetShard(shards[0]) + So(l.Config.shard, ShouldEqual, shards[0]) + + Convey("check that we can get the TRIM_HORIZON shard iterator", func() { + err := l.ensureShardIterator() + So(err, ShouldBeNil) + So(l.nextShardIterator, ShouldNotBeEmpty) + }) + + Convey("check that we can retrieve records one by one", func() { + data := []string{"foo", "bar"} + for n, datum := range data { + seq, err := putRecord(l, []byte(datum)) + So(err, ShouldBeNil) + So(seq, ShouldNotBeNil) + msg, err := l.Retrieve() + So(err, ShouldBeNil) + So(string(msg.Value()), ShouldEqual, datum) + Convey(fmt.Sprintf("check that iteration %d properly advanced the shard iterator", n), func() { + So(l.Config.shardIterator.shardIteratorType, ShouldEqual, "AT_SEQUENCE_NUMBER") + So(l.Config.shardIterator.sequenceNumber, ShouldEqual, *seq) + }) + } + }) + + Convey("check that retrieve will block until record comes", func(c C) { + start := time.Now() + data := "hello" + go func() { + <-time.After(1 * time.Second) + _, err := putRecord(l, []byte(data)) + c.So(err, ShouldBeNil) + }() + msg, err := l.Retrieve() + elapsed := time.Since(start) + Printf("(it blocked %f seconds)\n", elapsed.Seconds()) + So(err, ShouldBeNil) + So(string(msg.Value()), ShouldEqual, data) + So(elapsed.Seconds(), ShouldBeGreaterThan, 1) + }) + + Convey("check that we can use a context to cancel the retrieve", func() { + start := time.Now() + ctx, _ := context.WithTimeout(context.TODO(), 1*time.Second) + _, err := l.RetrieveWithContext(ctx) + elapsed := time.Since(start) + Printf("(it blocked %f seconds)\n", elapsed.Seconds()) + So(err, ShouldNotBeNil) + So(err, ShouldHaveSameTypeAs, context.DeadlineExceeded) + So(elapsed.Seconds(), ShouldBeGreaterThan, 1) + }) + + Convey("check that we can use a context to cancel the retrieve (again)", func() { + start := time.Now() + ctx, _ := context.WithTimeout(context.TODO(), 10*time.Millisecond) + _, err := l.RetrieveWithContext(ctx) + elapsed := time.Since(start) + Printf("(it blocked %f seconds)\n", elapsed.Seconds()) + So(err, ShouldNotBeNil) + So(err, ShouldHaveSameTypeAs, context.DeadlineExceeded) + So(elapsed.Seconds(), ShouldBeGreaterThan, 0.01) + }) + + Convey("check that retrieve still works with a canceller if a message comes before the deadline", func(c C) { + ctx, _ := context.WithTimeout(context.TODO(), 5*time.Second) + + data := "goodbye" + go func() { + <-time.After(1 * time.Second) + _, err := putRecord(l, []byte(data)) + c.So(err, ShouldBeNil) + }() + msg, err := l.RetrieveWithContext(ctx) + So(err, ShouldBeNil) + So(string(msg.Value()), ShouldEqual, data) + }) + + Convey("check that retrieve properly blocks other retrieves and attempts to set the shard id", func(c C) { + var wg sync.WaitGroup + wg.Add(1) + go func() { + ctx, _ := context.WithTimeout(context.TODO(), 1*time.Second) + _, err := l.RetrieveWithContext(ctx) + c.So(err, ShouldNotBeNil) + c.So(err, ShouldHaveSameTypeAs, context.DeadlineExceeded) + wg.Done() + }() + <-time.After(10 * time.Millisecond) + So(l.IsConsuming(), ShouldBeTrue) + _, err := l.Retrieve() + So(err, ShouldEqual, ErrAlreadyConsuming) + err = l.SetShard("bogus") + So(err, ShouldEqual, ErrCannotSetShard) + wg.Wait() + }) + + Convey("check that throttle mechanism prevents more than 5 calls to get records", func() { + start := time.Now() + secs := []float64{} + for i := 1; i <= 6; i++ { + start := time.Now() + l.fetchBatch(1) + secs = append(secs, time.Since(start).Seconds()) + } + elapsed := time.Since(start).Seconds() + So(elapsed, ShouldBeGreaterThan, 1) + Printf("%f seconds total, (%v)", elapsed, secs) + }) + + Convey("check that retrievefn can deliver messages to the fn", func(c C) { + called := false + data := "retrieved" + _, err := putRecord(l, []byte(data)) + So(err, ShouldBeNil) + err = l.RetrieveFn(func(b []byte, wg *sync.WaitGroup) { + called = true + // Note that because this is called in a goroutine, we have to use + // the goconvey context + c.So(string(b), ShouldEqual, data) + wg.Done() + }) + So(err, ShouldBeNil) + So(called, ShouldBeTrue) + }) + + Convey("check that listen can deliver messages to fn", func(c C) { + planets := []string{"mercury", "venus", "earth", "mars", "jupiter", "saturn", "neptune", "uranus"} + var count int64 = 0 + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + l.Listen(func(b []byte, wg *sync.WaitGroup) { + defer wg.Done() + atomic.AddInt64(&count, 1) + }) + }() + for _, planet := range planets { + _, err := putRecord(l, []byte(planet)) + So(err, ShouldBeNil) + } + close(l.pipeOfDeath) + wg.Wait() + So(count, ShouldEqual, len(planets)) + }) + + Convey("check that listen can be cancelled by context", func(c C) { + for i := 0; i < 20; i++ { + _, err := putRecord(l, []byte(fmt.Sprintf("%d", i))) + So(err, ShouldBeNil) + } + var count int64 = 0 + var wg sync.WaitGroup + wg.Add(1) + go func() { + ctx, cancel := context.WithCancel(context.TODO()) + defer wg.Done() + l.ListenWithContext(ctx, func(b []byte, wg *sync.WaitGroup) { + defer wg.Done() + time.AfterFunc(time.Duration(rand.Intn(10))*time.Second, func() { + n, err := strconv.Atoi(string(b)) + c.So(err, ShouldBeNil) + atomic.AddInt64(&count, 1) + if n == 15 { + cancel() + } + }) + }) + }() + wg.Wait() + So(count, ShouldBeBetweenOrEqual, 1, 20) + Printf("(count was %d)", count) + }) + }) + + }) + + Reset(func() { + l.DeleteStream() + l.WaitUntilDeleted(context.TODO(), request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + }) + }) + }) +} diff --git a/listener/stats.go b/listener/stats.go new file mode 100644 index 0000000..b6da386 --- /dev/null +++ b/listener/stats.go @@ -0,0 +1,37 @@ +package listener + +import ( + "time" +) + +// This was really built with rcrowley/go-metrics in mind. +type StatsListener interface { + // histograms + AddConsumedSample(int) + AddDeliveredSample(int) + AddProcessedSample(int) + AddBatchSizeSample(int) + + // meters + AddGetRecordsCalled(int) + AddProvisionedThroughputExceeded(int) + AddGetRecordsTimeout(int) + AddGetRecordsReadTimeout(int) + + // timers + AddGetRecordsReadResponseTime(time.Duration) + AddGetRecordsUnmarshalTime(time.Duration) +} + +type NilStatsListener struct{} + +func (l *NilStatsListener) AddConsumedSample(int) {} +func (l *NilStatsListener) AddDeliveredSample(int) {} +func (l *NilStatsListener) AddProcessedSample(int) {} +func (l *NilStatsListener) AddBatchSizeSample(int) {} +func (l *NilStatsListener) AddGetRecordsCalled(int) {} +func (l *NilStatsListener) AddProvisionedThroughputExceeded(int) {} +func (l *NilStatsListener) AddGetRecordsTimeout(int) {} +func (l *NilStatsListener) AddGetRecordsReadTimeout(int) {} +func (l *NilStatsListener) AddGetRecordsReadResponseTime(time.Duration) {} +func (l *NilStatsListener) AddGetRecordsUnmarshalTime(time.Duration) {} diff --git a/utils.go b/utils.go new file mode 100644 index 0000000..99fe9c2 --- /dev/null +++ b/utils.go @@ -0,0 +1,20 @@ +package kinetic + +import ( + "io" +) + +type TimeoutReadCloser struct { + io.ReadCloser + OnReadFn func(io.ReadCloser, []byte) (int, error) + OnCloseFn func() +} + +func (r *TimeoutReadCloser) Read(b []byte) (int, error) { + return r.OnReadFn(r.ReadCloser, b) +} + +func (r *TimeoutReadCloser) Close() error { + r.OnCloseFn() + return r.ReadCloser.Close() +} From dc17c6fbb67257f139eb649d0174120df160385e Mon Sep 17 00:00:00 2001 From: Tommy Wang Date: Wed, 29 Mar 2017 01:23:04 -0500 Subject: [PATCH 002/109] refactored listener, added kinetic General: * Factored out a logging and utils package in attempt to reduce circular dependencies. Listener: * Pulled out non-shard operations from listener (CreateStream, DeleteStream, etc). This is primarily because it is a much safer to have the shard set and immutable at listener creation. Operations like SetShard in previous implementation were dangerous if the shard iterators have already been retrieved. Non-shard operations are now moved to the Kinetic object. * Use an embedded struct (listenerConfig) to share common fields in Listener and its Config. * Added a FromKinetic to Listener's Config to import any AWS/relevant config from an existing Kinetic object. * Fixed the logLevel implementationt to be a bit more sane. * Separated out testing of Listener's Config to config_test.go Added Kinetic: * Added a Kinetic object that can be used to create Kinesis streams and Firehose Delivery Streams. Note that since the APIs for Kinesis and Firehose differ, they can co-exist in Kinetic comfortably. General Notes: The new Kinetic struct is completely optional, but may be useful in situations where you do not know what your ShardId is. For example, you may be spawning off multiple Listeners to handle all shards, or you may have an external mechanism to determine which ShardId to listen to. The Kinetic object can be instantiated and the GetShards method can be used to obtain a list of all shards. Once you have figured out the ShardId, you can construct a Listener object using: import "github.com/rewardStyle/kinetic" import "github.com/rewardStyle/kinetic/listener" k := kinetic.New(kinetic.NewConfig(). WithCredentials(accessKey, secretKe, securityToken). WithRegion("us-east-1")) l := listener.NewListener( listener.NewConfig(stream, shard). FromKinetic(k)) This allows you to use the same config from your Kinetic object to instantiate the Listener. Alternatively, you can opt to create a Listener directly: l := listener.NewListener( listener.NewConfig(stream, shard). WithCredentials(accessKey, secretKey, securityToken). WithRegion("us-east")) I'm not very satisfied with the amount of code duplication between the Listener's Config and Kinetic's Config. But, with the prototypal pattern With*, I could not figure out an easy way to maintain the chaining of the Config using embedded structs. Furthermore, I'm not to thrilled at the dependency implementation. I think it would be easier to allow something like this: import "github.com/rewardStyle/kinetic" k := kinetic.New(kinetic.NewConfig(). WithCredentials(accessKey, secretKe, securityToken). WithRegion("us-east-1")) l := k.NewListener(stream, shard) Where: func (k *Kinetic) NewListener(stream, shard string) (*Listener, error) { return listener.NewListener( listener.NewConfig(stream, shard). FromKinetic(k)) } But, this would require the kinetic package to import the listener package. Unfortunately, the listener package also needs to import the kinetic package. Moving logging and utils out to their separate packages has helped minimize the circular dependencies, but the listener test cases still rely directly on the kinetic package. Finally, the tests do require unexported fields. Despite that being a smell, I am torn between several forces: * Several of the unexported variables carry important state that really should not be tampered with. * I don't like the idea of adding a bunch of accessors for internal state variables simply for the purpose of testing. Furthermore, this actually just makes the public API more unstable, as it couples the API to the internal implementation of the packages. * The tests really need access to the unexported fields. TODOS: * We need to implement checkpointing such that we don't always start at TRIM_HORIZON. This would probably be done via dependency injection of an interface, similar to the StatsListener. * We need to add support for merging and splitting straem. At minimum, the Kinetic library should provide a mechanism for the calling application to decide what to do when a stream closes. --- .gitignore | 26 -- auth.go | 28 -- config.go | 120 +++---- config_test.go | 119 +++++-- firehose.go | 387 --------------------- firehose_test.go | 87 ----- kinesis.go | 209 ----------- kinesis_test.go | 25 -- kinetic.go | 151 ++++++++ kinetic.sample.conf | 12 - kinetic_test.go | 81 +++++ listener.go | 400 --------------------- listener/config.go | 53 ++- listener/config_test.go | 167 +++++++++ listener/listener.go | 225 ++++-------- listener/listener_test.go | 580 ++++++++++--------------------- listener_test.go | 166 --------- logging/logging.go | 14 + message.go => message/message.go | 8 +- producer.go | 398 --------------------- producer_test.go | 138 -------- utils.go => utils/readcloser.go | 2 +- 22 files changed, 810 insertions(+), 2586 deletions(-) delete mode 100644 auth.go delete mode 100644 firehose.go delete mode 100644 firehose_test.go delete mode 100644 kinesis.go delete mode 100644 kinesis_test.go create mode 100644 kinetic.go delete mode 100644 kinetic.sample.conf create mode 100644 kinetic_test.go delete mode 100644 listener.go create mode 100644 listener/config_test.go delete mode 100644 listener_test.go create mode 100644 logging/logging.go rename message.go => message/message.go (85%) delete mode 100644 producer.go delete mode 100644 producer_test.go rename utils.go => utils/readcloser.go (95%) diff --git a/.gitignore b/.gitignore index 3204325..838c953 100644 --- a/.gitignore +++ b/.gitignore @@ -1,29 +1,3 @@ -# Compiled Object files, Static and Dynamic libs (Shared Objects) -kinetic.conf -*.o -*.a -*.so - -# Folders -_obj -_test - -# Architecture specific extensions/prefixes -*.[568vq] -[568vq].out - -*.cgo1.go -*.cgo2.c -_cgo_defun.c -_cgo_gotypes.go -_cgo_export.* - -_testmain.go - -*.exe -*.test -*.prof - # Created by https://www.gitignore.io/api/vim ### Vim ### diff --git a/auth.go b/auth.go deleted file mode 100644 index b36030a..0000000 --- a/auth.go +++ /dev/null @@ -1,28 +0,0 @@ -package kinetic - -import ( - "errors" - - "github.com/aws/aws-sdk-go/aws" - - "github.com/aws/aws-sdk-go/aws/credentials" - "github.com/aws/aws-sdk-go/aws/session" - //gokinesis "github.com/rewardStyle/go-kinesis" -) - -// ErrMetaAuthentication represents an error that occurred on authentication from meta -var ErrMetaAuthentication = errors.New("Authentication error: failed to auth from meta. Your IAM roles are bad, or you need to specify an AccessKey and SecretKey") - -func authenticate(accessKey, secretKey string) (sess *session.Session, err error) { - if accessKey == "" || secretKey == "" { - if sess, err = session.NewSession(); err != nil { - return nil, ErrMetaAuthentication - } - } else { - conf := &aws.Config{} - conf = conf.WithCredentials(credentials.NewStaticCredentials(accessKey, secretKey, "")) - sess, err = session.NewSessionWithOptions(session.Options{Config: *conf}) - } - return - -} diff --git a/config.go b/config.go index 6a884af..e96bc38 100644 --- a/config.go +++ b/config.go @@ -1,89 +1,65 @@ package kinetic import ( - "io/ioutil" - "log" - "os" + "net/http" + "time" - "gopkg.in/gcfg.v1" + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/credentials" + "github.com/aws/aws-sdk-go/aws/session" ) -var ( - configPath = "/etc/kinetic.conf" - - defaultConfig = ` -[kinesis] -stream = stream-name -shard = 0 -sharditeratortype = 3 - -[firehose] -stream = firehose-stream-name - -[aws] -accesskey = accesskey -secretkey = secretkey -region = us-east-1 - -[debug] -verbose = true - -[concurrency] -listener = 100 -producer = 100 -` -) - -type config struct { - Kinesis struct { - Stream string - Shard string - ShardIteratorType int - } +type Config struct { + awsConfig *aws.Config + *kineticConfig +} - Firehose struct { - Stream string +func NewConfig() *Config { + return &Config{ + awsConfig: aws.NewConfig().WithHTTPClient(&http.Client{ + Timeout: 5 * time.Minute, + }), + kineticConfig: &kineticConfig{ + logLevel: aws.LogOff, + }, } +} - AWS struct { - AccessKey string - SecretKey string - Region string - } +func (c *Config) WithCredentials(accessKey, secretKey, securityToken string) *Config { + c.awsConfig.WithCredentials( + credentials.NewStaticCredentials(accessKey, secretKey, securityToken), + ) + return c +} - Debug struct { - Verbose bool - } +func (c *Config) WithRegion(region string) *Config { + c.awsConfig.WithRegion(region) + return c +} - Concurrency struct { - Listener int - Producer int - } +func (c *Config) WithEndpoint(endpoint string) *Config { + c.awsConfig.WithEndpoint(endpoint) + return c } -func getConfig() *config { - con := new(config) +func (c *Config) WithLogger(logger aws.Logger) *Config { + c.awsConfig.WithLogger(logger) + return c +} - file, err := ioutil.ReadFile(configPath) - if err != nil { - switch err.(type) { - case *os.PathError: - log.Println("Failed to parse config. Loading default configuration.") - file = []byte(defaultConfig) - default: - log.Println("Missing config: " + configPath + ". Loading default configuration.") - file = []byte(defaultConfig) - } - } +func (c *Config) WithLogLevel(logLevel aws.LogLevelType) *Config { + c.awsConfig.WithLogLevel(logLevel & 0xffff) + c.logLevel = logLevel & 0xffff0000 + return c +} - err = gcfg.ReadStringInto(con, string(file)) - if err != nil { - log.Println("Failed to parse config. Loading default configuration.") - err = gcfg.ReadStringInto(con, string(defaultConfig)) - if err != nil { - panic(err) - } - } +func (c *Config) WithHttpClientTimeout(timeout time.Duration) *Config { + c.awsConfig.WithHTTPClient(&http.Client{ + Timeout: timeout, + }) + return c +} - return con +func (c *Config) GetSession() (*session.Session, error) { + return session.NewSession(c.awsConfig) } diff --git a/config_test.go b/config_test.go index c40d364..da8d622 100644 --- a/config_test.go +++ b/config_test.go @@ -1,58 +1,101 @@ package kinetic import ( - "io/ioutil" - "os/exec" + . "github.com/smartystreets/goconvey/convey" + + "fmt" "testing" + "time" - . "github.com/smartystreets/goconvey/convey" + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/session" + + "github.com/rewardStyle/kinetic/logging" ) -const testConfigPath = "kinetic.conf" +func getSession(config *Config) *session.Session { + sess, err := config.GetSession() + So(err, ShouldBeNil) + So(sess, ShouldNotBeNil) + return sess +} -func TestBadConfig(t *testing.T) { - Convey("Given an incorrectly formatted config file", t, func() { - moveConfig(t) +func TestNewConfig(t *testing.T) { + Convey("given a new kinetic config", t, func() { + config := NewConfig() - Convey("The default configuration should be loaded", func() { - makeBadConfig(t, testConfigPath) + Convey("check the default values for its non-zero config", func() { + So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 5*time.Minute) + }) - config := getConfig() + Convey("check that we can retrieve an aws.Session from it ", func() { + getSession(config) + }) - So(config.Kinesis.Stream, ShouldNotResemble, nil) - So(config.AWS.AccessKey, ShouldNotResemble, nil) - So(config.AWS.SecretKey, ShouldNotResemble, nil) - restoreConfig(t) + Convey("check that we can set credentials", func() { + config = config.WithCredentials("access-key", "secret-key", "security-token") + sess := getSession(config) + creds, err := sess.Config.Credentials.Get() + So(err, ShouldBeNil) + So(creds.AccessKeyID, ShouldEqual, "access-key") + So(creds.SecretAccessKey, ShouldEqual, "secret-key") + So(creds.SessionToken, ShouldEqual, "security-token") + }) + + Convey("check that we can set the region", func() { + config = config.WithRegion("my-region") + sess := getSession(config) + So(aws.StringValue(sess.Config.Region), ShouldEqual, "my-region") }) - }) -} -func TestMissingConfig(t *testing.T) { - Convey("Given a missing config file", t, func() { - moveConfig(t) + Convey("check that we can set the endpoint", func() { + config = config.WithEndpoint("my-endpoint") + sess := getSession(config) + So(aws.StringValue(sess.Config.Endpoint), ShouldEqual, "my-endpoint") + }) - Convey("The default configuration should be loaded", func() { - config := getConfig() + Convey("check that we can configure a logger", func() { + var logs []string + loggerFn := func(args ...interface{}) { + logs = append(logs, fmt.Sprint(args...)) + } + config = config.WithLogger(aws.LoggerFunc(loggerFn)) + sess := getSession(config) - So(config.Kinesis.Stream, ShouldNotResemble, nil) - So(config.AWS.AccessKey, ShouldNotResemble, nil) - So(config.AWS.SecretKey, ShouldNotResemble, nil) - restoreConfig(t) + Convey("check that basic logging should work", func() { + sess.Config.Logger.Log("one") + sess.Config.Logger.Log("two") + sess.Config.Logger.Log("three") + So(len(logs), ShouldEqual, 3) + So(logs, ShouldContain, "one") + So(logs, ShouldContain, "two") + So(logs, ShouldContain, "three") + Reset(func() { + logs = nil + }) + }) }) - }) -} -func moveConfig(t *testing.T) { - exec.Command("mv", testConfigPath, testConfigPath+".missing").Run() -} + Convey("check that the default log level is off for both the sdk and kinetic", func() { + sess := getSession(config) + So(sess.Config.LogLevel.Value(), ShouldEqual, aws.LogOff) + So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeFalse) + So(config.logLevel.Value(), ShouldEqual, aws.LogOff) + So(config.logLevel.AtLeast(aws.LogDebug), ShouldBeFalse) + }) -func restoreConfig(t *testing.T) { - exec.Command("mv", testConfigPath+".missing", testConfigPath).Run() -} + Convey("check that we can set both the sdk and kinetic log level", func() { + ll := aws.LogDebug | aws.LogDebugWithSigning | logging.LogDebug + config = config.WithLogLevel(ll) + sess := getSession(config) + So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) + So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) + So(config.logLevel.AtLeast(logging.LogDebug), ShouldBeTrue) + }) -func makeBadConfig(t *testing.T, path string) { - err := ioutil.WriteFile(path, []byte("bad=config"), 0644) - if err != nil { - t.Fatalf(err.Error()) - } + Convey("check that we can set the http.Client Timeout", func() { + config = config.WithHttpClientTimeout(10 * time.Minute) + So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 10*time.Minute) + }) + }) } diff --git a/firehose.go b/firehose.go deleted file mode 100644 index 73be0be..0000000 --- a/firehose.go +++ /dev/null @@ -1,387 +0,0 @@ -package kinetic - -//gokinesis "github.com/rewardStyle/go-kinesis" -import ( - "errors" - "log" - "os" - "os/signal" - "runtime" - "sync" - "sync/atomic" - "syscall" - "time" - - "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/aws/credentials" - "github.com/aws/aws-sdk-go/aws/session" - awsFirehose "github.com/aws/aws-sdk-go/service/firehose" - awsFirehoseIface "github.com/aws/aws-sdk-go/service/firehose/firehoseiface" -) - -const ( - firehoseWritesPerSec int64 = 2000 - truncatedRecordTerminator string = `` -) - -// Firehose is a Producer -type Firehose struct { - stream string - client awsFirehoseIface.FirehoseAPI - msgCount int64 - errCount int64 - concurrency int - concurrencyMu sync.Mutex - sem chan Empty - wg sync.WaitGroup - producing bool - producingMu sync.Mutex - errors chan error - messages chan *Message - interrupts chan os.Signal -} - -func (p *Firehose) activate() (Producer, error) { - active, err := p.checkActive() - if err != nil || !active { - if err != nil { - return p, err - } - return p, ErrNotActive - } - - // go start feeder consumer and let listen processes them - go p.produce() - - return p, err -} - -// Close stops queuing and producing and waits for all tasks to finish -func (p *Firehose) Close() error { - if conf.Debug.Verbose { - log.Println("Producer is waiting for all tasks to finish...") - } - - p.wg.Wait() - - // Stop producing - go func() { - p.interrupts <- syscall.SIGINT - }() - - if conf.Debug.Verbose { - log.Println("Producer is shutting down.") - } - runtime.Gosched() - return nil -} - -// CloseSync closes the Firehose producer in a syncronous manner. -func (p *Firehose) CloseSync() error { - if conf.Debug.Verbose { - log.Println("Producer is waiting for all tasks to finish...") - } - - select { - case p.interrupts <- syscall.SIGINT: - break - default: - if conf.Debug.Verbose { - log.Println("Already closing listener.") - } - runtime.Gosched() - return nil - } - p.wg.Wait() - for p.IsProducing() { - runtime.Gosched() - } - if conf.Debug.Verbose { - log.Println("Listener is shutting down.") - } - runtime.Gosched() - return nil -} - -func (p *Firehose) initChannels() { - p.sem = make(chan Empty, p.getConcurrency()) - p.errors = make(chan error, p.getConcurrency()) - p.messages = make(chan *Message, p.msgBufSize()) - - p.interrupts = make(chan os.Signal, 1) - signal.Notify(p.interrupts, os.Interrupt) -} - -func (p *Firehose) checkActive() (bool, error) { - status, err := p.client.DescribeDeliveryStream( - &awsFirehose.DescribeDeliveryStreamInput{ - DeliveryStreamName: aws.String(p.stream)}) - if err != nil { - return false, err - } - - if streamStatuses[statusActive] == aws.StringValue(status.DeliveryStreamDescription.DeliveryStreamStatus) { - return true, nil - } - return false, nil -} - -func (p *Firehose) setConcurrency(concurrency int) { - p.concurrencyMu.Lock() - p.concurrency = concurrency - p.concurrencyMu.Unlock() -} - -func (p *Firehose) getConcurrency() (concurrency int) { - p.concurrencyMu.Lock() - concurrency = p.concurrency - p.concurrencyMu.Unlock() - return -} - -func (p *Firehose) msgBufSize() int { - return p.getConcurrency() * 1000 -} - -// Send a message to Firehose asyncronously -func (p *Firehose) Send(msg *Message) { - msg.SetValue(append(msg.Value(), truncatedRecordTerminator...)) - p.wg.Add(1) - go func() { - p.messages <- msg - p.wg.Done() - }() -} - -// TryToSend tries to send the message, but if the channel is full it drops the message, and returns an error. -func (p *Firehose) TryToSend(msg *Message) error { - msg.SetValue(append(msg.Value(), truncatedRecordTerminator...)) - p.wg.Add(1) - select { - case p.messages <- msg: - return nil - default: - return ErrDroppedMessage - } -} - -// Init initalizes a firehose producer with the config file defaults -func (p *Firehose) Init() (Producer, error) { - if conf.Concurrency.Producer < 1 { - return nil, ErrBadConcurrency - } - p.setConcurrency(conf.Concurrency.Producer) - p.initChannels() - sess, err := authenticate(conf.AWS.AccessKey, conf.AWS.SecretKey) - p.stream = conf.Firehose.Stream - p.client = awsFirehose.New(sess) - //gokinesis.NewWithEndpoint(auth, conf.AWS.Region, fmt.Sprintf(firehoseURL, conf.AWS.Region)), - - if err != nil { - return p, err - } - - return p.activate() -} - -// InitC initializes a producer for Kinesis Firehose with the specified params -func (p *Firehose) InitC(stream, _, _, accessKey, secretKey, region string, concurrency int) (Producer, error) { - if concurrency < 1 { - return nil, ErrBadConcurrency - } - if stream == "" { - return nil, ErrNullStream - } - - p.setConcurrency(concurrency) - p.initChannels() - sess, err := authenticate(accessKey, secretKey) - conf := &aws.Config{Region: aws.String(region)} - p.stream = stream - p.client = awsFirehose.New(sess, conf) - if err != nil { - return p, err - } - - return p.activate() -} - -// IsProducing returns true if Firehose is producing otherwise false -func (p *Firehose) IsProducing() (isProducing bool) { - p.producingMu.Lock() - isProducing = p.producing - p.producingMu.Unlock() - return -} - -func (p *Firehose) setProducing(producing bool) { - p.producingMu.Lock() - p.producing = producing - p.producingMu.Unlock() -} - -// NewEndpoint switches the endpoint of the firehose stream. This is useful for testing. -func (p *Firehose) NewEndpoint(endpoint, stream string) (err error) { - conf := &aws.Config{} - conf = conf.WithCredentials( - credentials.NewStaticCredentials("BAD_ACCESS_KEY", "BAD_SECRET_KEY", "BAD_TOKEN"), - ).WithEndpoint(endpoint).WithRegion("us-east-1") - sess, err := session.NewSessionWithOptions(session.Options{Config: *conf}) - p.client = awsFirehose.New(sess) - return err -} - -// Messages gets the current message channel from the producer -func (p *Firehose) Messages() <-chan *Message { - return p.messages -} - -// Errors gets the current number of errors on the Producer -func (p *Firehose) Errors() <-chan error { - return p.errors -} - -func (p *Firehose) produce() { - p.setProducing(true) - var counter int64 -stop: - for { - getLock(p.sem) - select { - case msg := <-p.Messages(): - timer := time.Now() - if p.firehoseFlush(&counter, &timer) { - p.wg.Add(1) - go func() { - records := []*awsFirehose.Record{&awsFirehose.Record{Data: msg.Value()}} - if conf.Debug.Verbose && atomic.LoadInt64(&counter)%100 == 0 { - log.Println("Attempting to send firehose messages") - } - p.sendFirehoseRecords(records...) - p.wg.Done() - }() - } - <-p.sem - case <-p.interrupts: - if conf.Debug.Verbose { - log.Println("Producer received interrupt signal") - } - <-p.sem - p.Close() - break stop - case err := <-p.Errors(): - if err != nil && conf.Debug.Verbose { - log.Println("Received error: ", err.Error()) - } - p.incErrCount() - <-p.sem - } - } -} - -// Each firehose stream can support up to 2,000 transactions per second for writes, -// 5,000 records a second, up to a maximum total data write rate of 5 MB per second by default. -// TODO: payload inspection & throttling -// PutRecordBatch can take up to 500 records per call or 4 MB per call, whichever is smaller. -// http://docs.aws.amazon.com/firehose/latest/dev/limits.html -func (p *Firehose) firehoseFlush(counter *int64, timer *time.Time) bool { - // 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() - atomic.StoreInt64(counter, 0) - } - - atomic.AddInt64(counter, 1) - - // If we have attempted 5000 times and it has been less than one second - // since we started sending then we need to wait for the second to finish - if atomic.LoadInt64(counter) >= firehoseWritesPerSec && !(time.Now().After(timer.Add(1 * time.Second))) { - // Wait for the remainder of the second - timer and counter - // will be reset on next pass - time.Sleep(time.Since(*timer)) - } - - return true -} - -// Queue the messages sent to firehose for POSTing. -// msgNum is number of the messge till now. It is just for only logging every 100 messages. -// if you want it to always -func (p *Firehose) sendFirehoseRecords(records ...*awsFirehose.Record) { - putResp, err := p.client.PutRecordBatch(&awsFirehose.PutRecordBatchInput{ - DeliveryStreamName: aws.String(p.stream), - Records: records, - }) - if err != nil && conf.Debug.Verbose { - p.errors <- err - } - - // Because we do not know which of the records was successful or failed - // we need to put them all back on the queue - if putResp != nil { - if aws.Int64Value(putResp.FailedPutCount) > 0 { - if conf.Debug.Verbose { - log.Printf("Failed firehose records: %d\n", aws.Int64Value(putResp.FailedPutCount)) - } - - for idx, resp := range putResp.RequestResponses { - // Put failed records back on the queue - if aws.StringValue(resp.ErrorCode) != "" || aws.StringValue(resp.ErrorMessage) != "" { - p.errors <- errors.New(aws.StringValue(resp.ErrorMessage)) - p.Send(new(Message).Init(records[idx].Data, "")) - - if conf.Debug.Verbose { - log.Println("Messages in failed PutRecords put back on the queue: " + string(records[idx].Data)) - } - } else { - p.incMsgCount() - } - } - } else { - p.incMsgCountBy(len(putResp.RequestResponses)) - } - } else if putResp == nil { - //resend records when we get a nil response - for _, record := range records { - p.Send(new(Message).Init(record.Data, "")) - if conf.Debug.Verbose { - log.Println("Message in nil send response put back on the queue: " + string(record.Data)) - } - } - } - - if conf.Debug.Verbose && p.getMsgCount()%100 == 0 { - log.Printf("Messages sent so far: %d\n", p.getMsgCount()) - } -} - -// ReInit re-initializes the shard iterator. Used with conjucntion with NewEndpoint -func (p *Firehose) ReInit() { - if !p.IsProducing() { - go p.produce() - } -} - -func (p *Firehose) incMsgCount() { - atomic.AddInt64(&p.msgCount, 1) -} - -func (p *Firehose) incMsgCountBy(n int) { - atomic.AddInt64(&p.msgCount, int64(n)) -} - -func (p *Firehose) getMsgCount() int64 { - return atomic.LoadInt64(&p.msgCount) -} - -func (p *Firehose) decErrCount() { - atomic.AddInt64(&p.errCount, -1) -} - -func (p *Firehose) incErrCount() { - atomic.AddInt64(&p.errCount, 1) -} - -func (p *Firehose) getErrCount() int64 { - return atomic.LoadInt64(&p.errCount) -} diff --git a/firehose_test.go b/firehose_test.go deleted file mode 100644 index 8a991c7..0000000 --- a/firehose_test.go +++ /dev/null @@ -1,87 +0,0 @@ -package kinetic - -import ( - "runtime" - "sync/atomic" - "testing" - "time" - - "github.com/aws/aws-sdk-go/aws/request" - "github.com/aws/aws-sdk-go/service/firehose" - . "github.com/smartystreets/goconvey/convey" -) - -func TestFireHose(t *testing.T) { - producer, _ := new(Firehose).InitC("your-stream", "", "", "accesskey", "secretkey", "us-east-1", 4) - producer.NewEndpoint("localhost", "your-stream") - producer.(*Firehose).client = new(fakefirehose) - producer.ReInit() - Convey("Given a running firehose producer", t, func() { - Convey("it should send data to the firehose stream", func() { - for i := 0; i < 100; i++ { - producer.Send(new(Message).Init([]byte("this is a message"), "")) - runtime.Gosched() - } - time.Sleep(10 * time.Second) - So(atomic.LoadInt64(&(producer.(*Firehose).client.(*fakefirehose).count)), ShouldEqual, 100) - So(producer.(*Firehose).getMsgCount(), ShouldEqual, 100) - }) - }) -} - -// Mocks for aws Firehose. -// This implements github.com/aws/aws-sdk-go/service/firehose/firehoseiface.FirehoseAPI -type fakefirehose struct { - count int64 -} - -func (f *fakefirehose) CreateDeliveryStreamRequest(*firehose.CreateDeliveryStreamInput) (*request.Request, *firehose.CreateDeliveryStreamOutput) { - return nil, nil -} - -func (f *fakefirehose) CreateDeliveryStream(*firehose.CreateDeliveryStreamInput) (*firehose.CreateDeliveryStreamOutput, error) { - return nil, nil -} - -func (f *fakefirehose) DeleteDeliveryStreamRequest(*firehose.DeleteDeliveryStreamInput) (*request.Request, *firehose.DeleteDeliveryStreamOutput) { - return nil, nil -} - -func (f *fakefirehose) DeleteDeliveryStream(*firehose.DeleteDeliveryStreamInput) (*firehose.DeleteDeliveryStreamOutput, error) { - return nil, nil -} -func (f *fakefirehose) DescribeDeliveryStreamRequest(*firehose.DescribeDeliveryStreamInput) (*request.Request, *firehose.DescribeDeliveryStreamOutput) { - return nil, nil -} -func (f *fakefirehose) DescribeDeliveryStream(*firehose.DescribeDeliveryStreamInput) (*firehose.DescribeDeliveryStreamOutput, error) { - return nil, nil -} -func (f *fakefirehose) ListDeliveryStreamsRequest(*firehose.ListDeliveryStreamsInput) (*request.Request, *firehose.ListDeliveryStreamsOutput) { - return nil, nil -} -func (f *fakefirehose) ListDeliveryStreams(*firehose.ListDeliveryStreamsInput) (*firehose.ListDeliveryStreamsOutput, error) { - return nil, nil -} -func (f *fakefirehose) PutRecordRequest(*firehose.PutRecordInput) (*request.Request, *firehose.PutRecordOutput) { - return nil, nil -} -func (f *fakefirehose) PutRecord(*firehose.PutRecordInput) (*firehose.PutRecordOutput, error) { - return nil, nil -} -func (f *fakefirehose) PutRecordBatchRequest(*firehose.PutRecordBatchInput) (*request.Request, *firehose.PutRecordBatchOutput) { - return nil, nil -} -func (f *fakefirehose) PutRecordBatch(*firehose.PutRecordBatchInput) (*firehose.PutRecordBatchOutput, error) { - atomic.AddInt64(&(f.count), 1) - return &firehose.PutRecordBatchOutput{ - RequestResponses: []*firehose.PutRecordBatchResponseEntry{ - &firehose.PutRecordBatchResponseEntry{}, - }, - }, nil -} -func (f *fakefirehose) UpdateDestinationRequest(*firehose.UpdateDestinationInput) (*request.Request, *firehose.UpdateDestinationOutput) { - return nil, nil -} -func (f *fakefirehose) UpdateDestination(*firehose.UpdateDestinationInput) (*firehose.UpdateDestinationOutput, error) { - return nil, nil -} diff --git a/kinesis.go b/kinesis.go deleted file mode 100644 index 2b9c3bd..0000000 --- a/kinesis.go +++ /dev/null @@ -1,209 +0,0 @@ -package kinetic - -import ( - "errors" - "sync" - "sync/atomic" - - "github.com/aws/aws-sdk-go/aws" - - "github.com/aws/aws-sdk-go/aws/credentials" - "github.com/aws/aws-sdk-go/aws/session" - awsKinesis "github.com/aws/aws-sdk-go/service/kinesis" - "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" -) - -const ( - atSequenceNumber = iota - afterSequenceNumber - trimHorizon - latest - - statusCreating = iota - statusDeleting - statusActive - staticUpdating - - kinesisWritesPerSec int = 1000 - kinesisReadsPerSec int = 5 - // Timeout TODO - Timeout = 60 -) - -// ErrNilShardIterator is an error for when we get back a nil shard iterator -var ErrNilShardIterator = errors.New("Nil shard iterator") - -// ErrNilShardStatus is an error for when we get back a nil stream status -var ErrNilShardStatus = errors.New("Nil stream status") - -// Empty is an empty struct. It is mostly used for counting semaphore purposes -type Empty struct{} - -var ( - conf = getConfig() - - // ShardIterTypes are the types of iterators to use within Kinesis - ShardIterTypes shardIteratorTypes = map[int]string{ - atSequenceNumber: "AT_SEQUENCE_NUMBER", - afterSequenceNumber: "AFTER_SEQUENCE_NUMBER", - trimHorizon: "TRIM_HORIZON", - latest: "LATEST", - } - - streamStatuses streamStatusTypes = map[int]string{ - statusCreating: "CREATING", - statusDeleting: "DELETING", - statusActive: "ACTIVE", - staticUpdating: "UPDATING", - } -) - -type msgFn func([]byte, *sync.WaitGroup) -type shardIteratorTypes map[int]string -type streamStatusTypes map[int]string - -type kinesis struct { - stream string - shard string - endPoint string - shardIteratorType string - shardIterator string - sequenceNumber string - sequenceNumberMu sync.Mutex - - client kinesisiface.KinesisAPI - - msgCount int64 - errCount int64 -} - -func (k *kinesis) init(stream, shard, shardIteratorType, accessKey, secretKey, region string) (*kinesis, error) { - - sess, err := authenticate(accessKey, secretKey) - conf := aws.NewConfig().WithRegion(region) - if k.endPoint != "" { - conf = conf.WithEndpoint(k.endPoint) - } - client := awsKinesis.New(sess, conf) - - k = &kinesis{ - stream: stream, - shard: shard, - shardIteratorType: shardIteratorType, - client: client, - } - if err != nil { - return k, err - } - - err = k.initShardIterator() - return k, err -} - -func (k *kinesis) initShardIterator() error { - var awsSeqNumber *string - if k.sequenceNumber != "" { - awsSeqNumber = aws.String(k.sequenceNumber) - k.shardIteratorType = ShardIterTypes[atSequenceNumber] - } - resp, err := k.client.GetShardIterator(&awsKinesis.GetShardIteratorInput{ - ShardId: aws.String(k.shard), // Required - ShardIteratorType: aws.String(k.shardIteratorType), // Required - StreamName: aws.String(k.stream), // Required - StartingSequenceNumber: awsSeqNumber, - }) - if err != nil { - return err - } - if resp.ShardIterator != nil { - return k.setShardIterator(*resp.ShardIterator) - } - - return ErrNilShardIterator -} - -func (k *kinesis) setSequenceNumber(sequenceNum string) { - if sequenceNum == "" || len(sequenceNum) == 0 { - return - } - - k.sequenceNumberMu.Lock() - k.sequenceNumber = sequenceNum - k.sequenceNumberMu.Unlock() -} - -func (k *kinesis) setShardIterator(shardIter string) error { - if shardIter == "" || len(shardIter) == 0 { - return errors.New("Attempted to set shard iterator with empty value") - } - - k.shardIterator = shardIter - - return nil -} - -func (k *kinesis) checkActive() (bool, error) { - status, err := k.client.DescribeStream(&awsKinesis.DescribeStreamInput{ - StreamName: aws.String(k.stream), // Required - }) - if err != nil { - return false, err - } - if status.StreamDescription.StreamStatus == nil { - return false, ErrNilShardStatus - } - if streamStatuses[statusActive] == *status.StreamDescription.StreamStatus { - return true, nil - } - return false, nil -} - -func (k *kinesis) newClient(endpoint, stream string) (*awsKinesis.Kinesis, error) { - k.endPoint = endpoint - conf := &aws.Config{} - conf = conf.WithCredentials( - credentials.NewStaticCredentials("BAD_ACCESS_KEY", "BAD_SECRET_KEY", "BAD_TOKEN"), - ).WithEndpoint(endpoint).WithRegion("us-east-1").WithDisableSSL(true).WithMaxRetries(3) - - // fake region - sess, err := session.NewSessionWithOptions(session.Options{Config: *conf}) - return awsKinesis.New(sess, conf), err -} - -func (k *kinesis) refreshClient(accessKey, secretKey, region string) error { - sess, err := authenticate(accessKey, secretKey) - conf := aws.NewConfig().WithRegion(region).WithEndpoint(k.endPoint) - if err != nil { - return err - } - k.client = awsKinesis.New(sess, conf) - return nil -} - -func (k *kinesis) decMsgCount() { - atomic.AddInt64(&k.msgCount, -1) -} - -func (k *kinesis) incMsgCount() { - atomic.AddInt64(&k.msgCount, 1) -} - -func (k *kinesis) getMsgCount() int64 { - return atomic.LoadInt64(&k.msgCount) -} - -func (k *kinesis) decErrCount() { - atomic.AddInt64(&k.errCount, -1) -} - -func (k *kinesis) incErrCount() { - atomic.AddInt64(&k.errCount, 1) -} - -func (k *kinesis) getErrCount() int64 { - return atomic.LoadInt64(&k.errCount) -} - -func getLock(sem chan Empty) { - sem <- Empty{} -} diff --git a/kinesis_test.go b/kinesis_test.go deleted file mode 100644 index 6fc521f..0000000 --- a/kinesis_test.go +++ /dev/null @@ -1,25 +0,0 @@ -package kinetic - -import ( - "testing" - - . "github.com/smartystreets/goconvey/convey" -) - -// TestKineticCreation tests to make sure that the Kinetic creation doesn't return early or return a nil. -func TestKineticCreation(t *testing.T) { - kinesisKinetic, err := new(kinesis).init("fake", "ShardId-00000001", "TRIM_HORIZON", "BADaccessKey", "BADsecretKey", "region") - - Convey("Given an badly configured init-ed kinetic", t, func() { - Convey("the error returned should not be nil", func() { - So(err, ShouldNotBeNil) - }) - Convey("the returned kinesis struct pointer should not be nil", func() { - So(kinesisKinetic, ShouldNotBeNil) - }) - Convey("it should also have some data in it", func() { - So(kinesisKinetic.stream, ShouldEqual, "fake") - So(kinesisKinetic.shard, ShouldEqual, "ShardId-00000001") - }) - }) -} diff --git a/kinetic.go b/kinetic.go new file mode 100644 index 0000000..aa88af1 --- /dev/null +++ b/kinetic.go @@ -0,0 +1,151 @@ +package kinetic + +import ( + "context" + "errors" + "sync" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/request" + "github.com/aws/aws-sdk-go/aws/session" + //"github.com/aws/aws-sdk-go/service/firehose" + "github.com/aws/aws-sdk-go/service/firehose/firehoseiface" + "github.com/aws/aws-sdk-go/service/kinesis" + "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" + + "github.com/rewardStyle/kinetic/logging" +) + +var ( + // GetShards errors + ErrNilDescribeStreamResponse = errors.New("DescribeStream returned a nil response") + ErrNilStreamDescription = errors.New("DescribeStream returned a nil StreamDescription") +) + +type kineticConfig struct { + logLevel aws.LogLevelType +} + +type Kinetic struct { + *kineticConfig + session *session.Session + + fclient firehoseiface.FirehoseAPI + kclient kinesisiface.KinesisAPI + clientMu sync.Mutex +} + +func New(config *Config) (*Kinetic, error) { + session, err := config.GetSession() + if err != nil { + return nil, err + } + return &Kinetic{ + kineticConfig: config.kineticConfig, + session: session, + }, nil +} + +func (k *Kinetic) Log(args ...interface{}) { + if k.logLevel.Matches(logging.LogDebug) { + k.session.Config.Logger.Log(args...) + } +} + +func (k *Kinetic) ensureKinesisClient() { + k.clientMu.Lock() + defer k.clientMu.Unlock() + if k.kclient == nil { + k.kclient = kinesis.New(k.session) + } +} + +func (k *Kinetic) CreateStream(stream string, shards int) error { + k.ensureKinesisClient() + _, err := k.kclient.CreateStream(&kinesis.CreateStreamInput{ + StreamName: aws.String(stream), + ShardCount: aws.Int64(int64(shards)), + }) + if err != nil { + k.Log("Error creating kinesis stream:", err) + } + return err +} + +func (k *Kinetic) WaitUntilStreamExists(ctx context.Context, stream string, opts ...request.WaiterOption) error { + k.ensureKinesisClient() + return k.kclient.WaitUntilStreamExistsWithContext(ctx, &kinesis.DescribeStreamInput{ + StreamName: aws.String(stream), // Required + }, opts...) +} + +func (k *Kinetic) DeleteStream(stream string) error { + k.ensureKinesisClient() + _, err := k.kclient.DeleteStream(&kinesis.DeleteStreamInput{ + StreamName: aws.String(stream), + }) + if err != nil { + k.Log("Error deleting kinesis stream:", err) + } + return err +} + +func (k *Kinetic) WaitUntilStreamDeleted(ctx context.Context, stream string, opts ...request.WaiterOption) error { + k.ensureKinesisClient() + w := request.Waiter{ + Name: "WaitUntilStreamIsDeleted", + MaxAttempts: 18, + Delay: request.ConstantWaiterDelay(10 * time.Second), + Acceptors: []request.WaiterAcceptor{ + { + State: request.SuccessWaiterState, + Matcher: request.ErrorWaiterMatch, + Expected: kinesis.ErrCodeResourceNotFoundException, + }, + }, + Logger: k.session.Config.Logger, + NewRequest: func(opts []request.Option) (*request.Request, error) { + req, _ := k.kclient.DescribeStreamRequest(&kinesis.DescribeStreamInput{ + StreamName: aws.String(stream), // Required + }) + req.SetContext(ctx) + req.ApplyOptions(opts...) + return req, nil + }, + } + w.ApplyOptions(opts...) + return w.WaitWithContext(ctx) +} + +func (k *Kinetic) GetShards(stream string) ([]string, error) { + k.ensureKinesisClient() + resp, err := k.kclient.DescribeStream(&kinesis.DescribeStreamInput{ + StreamName: aws.String(stream), + }) + if err != nil { + k.Log("Error describing kinesis stream:", err) + return nil, err + } + if resp == nil { + return nil, ErrNilDescribeStreamResponse + } + if resp.StreamDescription == nil { + return nil, ErrNilStreamDescription + } + var shards []string + for _, shard := range resp.StreamDescription.Shards { + if shard.ShardId != nil { + shards = append(shards, aws.StringValue(shard.ShardId)) + } + } + return shards, nil +} + +func (k *Kinetic) GetSession() *session.Session { + return k.session +} + +// func (k *Kinetic) NewListener(config *listener.Config) (*listener.Listener, error) { +// return listener.NewListener(config, k, k.session, k.kclient) +// } diff --git a/kinetic.sample.conf b/kinetic.sample.conf deleted file mode 100644 index 602a8cb..0000000 --- a/kinetic.sample.conf +++ /dev/null @@ -1,12 +0,0 @@ -[kinesis] -stream = stream-name -shard = 0 -sharditeratortype = 3 - -[aws] -accesskey = accesskey -secretkey = secretkey -region = us-east-1 - -[debug] -verbose = true diff --git a/kinetic_test.go b/kinetic_test.go new file mode 100644 index 0000000..dd0b4e3 --- /dev/null +++ b/kinetic_test.go @@ -0,0 +1,81 @@ +package kinetic + +import ( + . "github.com/smartystreets/goconvey/convey" + + "context" + "testing" + "time" + + "github.com/aws/aws-sdk-go/aws/awserr" + "github.com/aws/aws-sdk-go/aws/request" + "github.com/aws/aws-sdk-go/service/kinesis" +) + +func TestKinetic(t *testing.T) { + Convey("given a kinetic object", t, func() { + k, err := New(NewConfig(). + WithRegion("some-region"). + WithEndpoint("http://127.0.0.1:4567")) + So(k, ShouldNotBeNil) + So(err, ShouldBeNil) + + stream := "some-kinetic-stream" + + Convey("check that calling ensureKinesisClient twice doesn't overwrite existing client", func() { + So(k.kclient, ShouldBeNil) + k.ensureKinesisClient() + So(k.kclient, ShouldNotBeNil) + kclient := k.kclient + k.ensureKinesisClient() + So(k.kclient, ShouldEqual, kclient) + }) + + Convey("check deleting a non-existent stream returns an error", func() { + err := k.DeleteStream(stream) + So(err, ShouldNotBeNil) + e := err.(awserr.Error) + So(e.Code(), ShouldEqual, kinesis.ErrCodeResourceNotFoundException) + }) + + Convey("check getting shards on a non-existent stream returns an error", func() { + shards, err := k.GetShards(stream) + So(shards, ShouldBeNil) + So(err, ShouldNotBeNil) + e := err.(awserr.Error) + So(e.Code(), ShouldEqual, kinesis.ErrCodeResourceNotFoundException) + }) + + Convey("check that we can create a stream", func() { + err := k.CreateStream(stream, 1) + So(err, ShouldBeNil) + err = k.WaitUntilStreamExists(context.TODO(), stream, request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + So(err, ShouldBeNil) + + Convey("check that creating an existing stream returns an error", func() { + err := k.CreateStream(stream, 1) + So(err, ShouldNotBeNil) + e := err.(awserr.Error) + So(e.Code(), ShouldEqual, kinesis.ErrCodeResourceInUseException) + }) + + Convey("check that we can delete an existing stream", func() { + err := k.DeleteStream(stream) + So(err, ShouldBeNil) + err = k.WaitUntilStreamDeleted(context.TODO(), stream, request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + So(err, ShouldBeNil) + }) + + Convey("check that we can obtain a list of shards", func() { + shards, err := k.GetShards(stream) + So(err, ShouldBeNil) + So(len(shards), ShouldEqual, 1) + }) + + Reset(func() { + k.DeleteStream(stream) + k.WaitUntilStreamDeleted(context.TODO(), stream, request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + }) + }) + }) +} diff --git a/listener.go b/listener.go deleted file mode 100644 index 9fe0881..0000000 --- a/listener.go +++ /dev/null @@ -1,400 +0,0 @@ -package kinetic - -import ( - "errors" - "log" - "net" - "os" - "os/signal" - "runtime" - "sync" - "syscall" - "time" - - "github.com/aws/aws-sdk-go/aws" - awsKinesis "github.com/aws/aws-sdk-go/service/kinesis" -) - -var ( - // ErrNullStream represents an error where the stream was not specified - ErrNullStream = errors.New("A stream must be specified") - // ErrNotActive represents an error where the stream is not ready for processing - ErrNotActive = errors.New("The Stream is not yet active") -) - -// Listener represents a kinesis listener -type Listener struct { - *kinesis - - accessKey string - secretKey string - region string - - concurrency int - concurrencyMu sync.Mutex - sem chan Empty - - wg sync.WaitGroup - - listening bool - listeningMu sync.Mutex - consuming bool - consumingMu sync.Mutex - - errors chan error - messages chan *Message - interrupts chan os.Signal -} - -func (l *Listener) init(stream, shard, shardIterType, accessKey, secretKey, region string, concurrency int) (*Listener, error) { - var err error - if concurrency < 1 { - return nil, ErrBadConcurrency - } - if stream == "" { - return nil, ErrNullStream - } - - l.setConcurrency(concurrency) - - l.accessKey = accessKey - l.secretKey = secretKey - l.region = region - - l.sem = make(chan Empty, l.getConcurrency()) - l.errors = make(chan error, l.getConcurrency()) - l.messages = make(chan *Message, l.msgBufSize()) - - l.interrupts = make(chan os.Signal, 1) - signal.Notify(l.interrupts, os.Interrupt) - - l.kinesis, err = new(kinesis).init(stream, shard, shardIterType, accessKey, secretKey, region) - if err != nil { - return l, err - } - - // Is the stream ready? - active, err := l.checkActive() - if err != nil || !active { - if err != nil { - return l, err - } - return l, ErrNotActive - } - - // Start feeder consumer - go l.consume() - - return l, err -} - -// Init initializes a listener with the params specified in the configuration file -func (l *Listener) Init() (*Listener, error) { - return l.init(conf.Kinesis.Stream, conf.Kinesis.Shard, ShardIterTypes[conf.Kinesis.ShardIteratorType], conf.AWS.AccessKey, conf.AWS.SecretKey, conf.AWS.Region, conf.Concurrency.Listener) -} - -// InitC initialize a listener with the supplied params -func (l *Listener) InitC(stream, shard, shardIterType, accessKey, secretKey, region string, concurrency int) (*Listener, error) { - return l.init(stream, shard, shardIterType, accessKey, secretKey, region, concurrency) -} - -// NewEndpoint re-initializes kinesis client with new endpoint. Used for testing with kinesalite -func (l *Listener) NewEndpoint(endpoint, stream string) (err error) { - l.kinesis.client, err = l.kinesis.newClient(endpoint, stream) - return -} - -// ReInit re-initializes the shard iterator. Used with conjucntion with NewEndpoint -func (l *Listener) ReInit() { - l.initShardIterator() - - if !l.IsConsuming() { - go l.consume() - } -} - -func (l *Listener) setConcurrency(concurrency int) { - l.concurrencyMu.Lock() - l.concurrency = concurrency - l.concurrencyMu.Unlock() -} - -func (l *Listener) getConcurrency() int { - l.concurrencyMu.Lock() - defer l.concurrencyMu.Unlock() - return l.concurrency -} - -func (l *Listener) msgBufSize() int { - l.concurrencyMu.Lock() - defer l.concurrencyMu.Unlock() - return l.concurrency * 1000 -} - -func (l *Listener) setListening(listening bool) { - l.listeningMu.Lock() - l.listening = listening - l.listeningMu.Unlock() -} - -// IsListening identifies whether or not messages and errors are being handled after consumption -func (l *Listener) IsListening() bool { - l.listeningMu.Lock() - defer l.listeningMu.Unlock() - return l.listening -} - -func (l *Listener) setConsuming(consuming bool) { - l.consumingMu.Lock() - l.consuming = consuming - l.consumingMu.Unlock() -} - -// IsConsuming identifies whether or not the kinesis stream is being polled -func (l *Listener) IsConsuming() bool { - l.consumingMu.Lock() - defer l.consumingMu.Unlock() - return l.consuming -} - -func (l *Listener) shouldConsume() bool { - select { - case <-l.interrupts: - return false - default: - return true - } - -} - -// Listen handles the consumed messages, errors and interrupts -func (l *Listener) Listen(fn msgFn) { - l.setListening(true) -stop: - for { - getLock(l.sem) //counting semaphore - - select { - case err := <-l.errors: - l.incErrCount() - l.wg.Add(1) - go l.handleError(err) - case msg := <-l.messages: - l.incMsgCount() - l.wg.Add(1) - go l.handleMsg(msg, fn) - case sig := <-l.interrupts: - l.handleInterrupt(sig) - break stop - } - } - l.setListening(false) -} - -// Continually poll the Kinesis stream -func (l *Listener) consume() { - l.setConsuming(true) - - readCounter := 0 - readTimer := time.Now() - - GsiCounter := 0 - GsiTimer := time.Now() - - for { - if !l.shouldConsume() { - l.setConsuming(false) - break - } - - l.throttle(&readCounter, &readTimer) - - // args() will give us the shard iterator and type as well as the shard id - response, err := l.client.GetRecords( - &awsKinesis.GetRecordsInput{ - Limit: aws.Int64(10000), - ShardIterator: aws.String(l.shardIterator), - }, - ) - if err != nil { - go func() { - l.errors <- err - }() - - // We receive net.OpError if kinesis terminates the socket. - // It will contain a message resembling: - // - // Received error: Post https://kinesis.us-east-1.amazonaws.com: - // read tcp 172.16.0.38:37680->54.239.28.39:443: read: connection reset by peer - // - // If this happens we need to refresh the kinesis client to - // reestablish the connection - if _, ok := err.(*net.OpError); ok { - log.Println("Received net.OpError. Recreating kinesis client and connection.") - l.refreshClient(l.accessKey, l.secretKey, l.region) - } - - refresh_iterator: - - // Refresh the shard iterator - err := l.initShardIterator() - if err != nil { - log.Println("Failed to refresh iterator: " + err.Error()) - // If we received an error we should wait and attempt to - // refresh the shard iterator again - l.throttle(&GsiCounter, &GsiTimer) - - goto refresh_iterator - } - } - - if response != nil && response.NextShardIterator != nil { - l.setShardIterator(*response.NextShardIterator) - - if len(response.Records) > 0 { - for _, record := range response.Records { - if record != nil { - l.messages <- &Message{*record} - } - l.setSequenceNumber(*record.SequenceNumber) - } - } - } - } -} - -// Retrieve a message from the stream and return the value -func (l *Listener) Retrieve() (*Message, error) { - select { - case msg := <-l.messages: - return msg, nil - case err := <-l.Errors(): - return nil, err - case sig := <-l.interrupts: - l.handleInterrupt(sig) - return nil, nil - } -} - -// RetrieveFn retrieves a message from the queue and apply the supplied function to the message -func (l *Listener) RetrieveFn(fn msgFn) { - select { - case err := <-l.Errors(): - l.wg.Add(1) - go l.handleError(err) - case msg := <-l.messages: - l.wg.Add(1) - go fn(msg.Value(), &l.wg) - case sig := <-l.interrupts: - l.handleInterrupt(sig) - } -} - -// Close stops consuming and listening and waits for all tasks to finish -func (l *Listener) Close() error { - if conf.Debug.Verbose { - log.Println("Listener is waiting for all tasks to finish...") - } - // Stop consuming - go func() { - l.interrupts <- syscall.SIGINT - }() - - l.wg.Wait() - - if conf.Debug.Verbose { - log.Println("Listener is shutting down.") - } - runtime.Gosched() - return nil -} - -// CloseSync closes the Listener in a syncronous manner. -func (l *Listener) CloseSync() error { - if conf.Debug.Verbose { - log.Println("Listener is waiting for all tasks to finish...") - } - var err error - // Stop consuming - select { - case l.interrupts <- syscall.SIGINT: - break - default: - if conf.Debug.Verbose { - log.Println("Already closing listener.") - } - runtime.Gosched() - return err - } - l.wg.Wait() - for l.IsConsuming() { - runtime.Gosched() - } - if conf.Debug.Verbose { - log.Println("Listener is shutting down.") - } - runtime.Gosched() - return nil -} - -// Errors gets the current number of errors on the Listener -func (l *Listener) Errors() <-chan error { - return l.errors -} - -func (l *Listener) handleMsg(msg *Message, fn msgFn) { - if conf.Debug.Verbose { - if l.getMsgCount()%100 == 0 { - log.Printf("Messages received: %d", l.getMsgCount()) - } - } - - defer func() { - <-l.sem - }() - - fn(msg.Value(), &l.wg) -} - -func (l *Listener) handleInterrupt(signal os.Signal) { - if conf.Debug.Verbose { - log.Println("Listener received interrupt signal") - } - - defer func() { - <-l.sem - }() - - l.Close() -} - -func (l *Listener) handleError(err error) { - if err != nil && conf.Debug.Verbose { - log.Println("Received error: ", err.Error()) - } - - defer func() { - <-l.sem - }() - - l.wg.Done() -} - -// 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) { - // 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() - *counter = 0 - } - - *counter++ - - // If we have attempted five times and it has been less than one second - // since we started reading then we need to wait for the second to finish - if *counter >= kinesisReadsPerSec && !(time.Now().After(timer.Add(1 * time.Second))) { - // Wait for the remainder of the second - timer and counter will be reset on next pass - time.Sleep(1*time.Second - time.Since(*timer)) - } -} diff --git a/listener/config.go b/listener/config.go index d574256..66d4d0a 100644 --- a/listener/config.go +++ b/listener/config.go @@ -7,38 +7,31 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/credentials" "github.com/aws/aws-sdk-go/aws/session" -) -const () + "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/logging" +) type Config struct { awsConfig *aws.Config - - stream string - - logLevel aws.LogLevelType - stats StatsListener - - getRecordsReadTimeout time.Duration - - shard string - batchSize int - concurrency int - shardIterator *ShardIterator + *listenerConfig } -func NewConfig(stream string) *Config { +func NewConfig(stream, shard string) *Config { return &Config{ awsConfig: aws.NewConfig().WithHTTPClient(&http.Client{ Timeout: 5 * time.Minute, }), - stream: stream, - logLevel: aws.LogOff, - stats: &NilStatsListener{}, - batchSize: 10000, - concurrency: 10000, - shardIterator: NewShardIterator(), - getRecordsReadTimeout: 1 * time.Second, + listenerConfig: &listenerConfig{ + stream: stream, + shard: shard, + batchSize: 10000, + concurrency: 10000, + shardIterator: NewShardIterator(), + getRecordsReadTimeout: 1 * time.Second, + stats: &NilStatsListener{}, + logLevel: logging.LogOff, + }, } } @@ -66,7 +59,7 @@ func (c *Config) WithLogger(logger aws.Logger) *Config { func (c *Config) WithLogLevel(logLevel aws.LogLevelType) *Config { c.awsConfig.WithLogLevel(logLevel & 0xffff) - c.logLevel = logLevel >> 16 + c.logLevel = logLevel & 0xffff0000 return c } @@ -77,13 +70,12 @@ func (c *Config) WithHttpClientTimeout(timeout time.Duration) *Config { return c } -func (c *Config) WithStatsListener(stats StatsListener) *Config { - c.stats = stats - return c +func (c *Config) GetSession() (*session.Session, error) { + return session.NewSession(c.awsConfig) } -func (c *Config) WithShardId(shard string) *Config { - c.shard = shard +func (c *Config) FromKinetic(k *kinetic.Kinetic) *Config { + c.awsConfig = k.GetSession().Config return c } @@ -107,6 +99,7 @@ func (c *Config) WithGetRecordsReadTimeout(timouet time.Duration) *Config { return c } -func (c *Config) GetAwsSession() (*session.Session, error) { - return session.NewSession(c.awsConfig) +func (c *Config) WithStatsListener(stats StatsListener) *Config { + c.stats = stats + return c } diff --git a/listener/config_test.go b/listener/config_test.go new file mode 100644 index 0000000..5305a75 --- /dev/null +++ b/listener/config_test.go @@ -0,0 +1,167 @@ +package listener + +import ( + . "github.com/smartystreets/goconvey/convey" + + "fmt" + "testing" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/session" + + "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/logging" +) + +type DebugStatsListener struct{} + +func (l *DebugStatsListener) AddConsumedSample(int) {} +func (l *DebugStatsListener) AddDeliveredSample(int) {} +func (l *DebugStatsListener) AddProcessedSample(int) {} +func (l *DebugStatsListener) AddBatchSizeSample(int) {} +func (l *DebugStatsListener) AddGetRecordsCalled(int) {} +func (l *DebugStatsListener) AddProvisionedThroughputExceeded(int) {} +func (l *DebugStatsListener) AddGetRecordsTimeout(int) {} +func (l *DebugStatsListener) AddGetRecordsReadTimeout(int) {} +func (l *DebugStatsListener) AddGetRecordsReadResponseTime(time.Duration) {} +func (l *DebugStatsListener) AddGetRecordsUnmarshalTime(time.Duration) {} + +func getSession(config *Config) *session.Session { + sess, err := config.GetSession() + So(err, ShouldBeNil) + So(sess, ShouldNotBeNil) + return sess +} + +func TestNewConfig(t *testing.T) { + Convey("given a new listener config", t, func() { + stream := "some-stream" + shard := "some-shard" + config := NewConfig(stream, shard) + + Convey("check the default values for its non-zero config", func() { + So(config.stream, ShouldEqual, stream) + So(config.shard, ShouldEqual, shard) + So(config.batchSize, ShouldEqual, 10000) + So(config.concurrency, ShouldEqual, 10000) + So(config.shardIterator.shardIteratorType, ShouldEqual, "TRIM_HORIZON") + So(config.getRecordsReadTimeout, ShouldEqual, 1*time.Second) + So(config.stats, ShouldHaveSameTypeAs, &NilStatsListener{}) + So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 5*time.Minute) + }) + + Convey("check that we can retrieve an aws.Session from it ", func() { + getSession(config) + }) + + Convey("check that we can set credentials", func() { + config = config.WithCredentials("access-key", "secret-key", "security-token") + sess := getSession(config) + creds, err := sess.Config.Credentials.Get() + So(err, ShouldBeNil) + So(creds.AccessKeyID, ShouldEqual, "access-key") + So(creds.SecretAccessKey, ShouldEqual, "secret-key") + So(creds.SessionToken, ShouldEqual, "security-token") + }) + + Convey("check that we can set the region", func() { + config = config.WithRegion("my-region") + sess := getSession(config) + So(aws.StringValue(sess.Config.Region), ShouldEqual, "my-region") + }) + + Convey("check that we can set the endpoint", func() { + config = config.WithEndpoint("my-endpoint") + sess := getSession(config) + So(aws.StringValue(sess.Config.Endpoint), ShouldEqual, "my-endpoint") + }) + + Convey("check that we can configure a logger", func() { + var logs []string + loggerFn := func(args ...interface{}) { + logs = append(logs, fmt.Sprint(args...)) + } + config = config.WithLogger(aws.LoggerFunc(loggerFn)) + sess := getSession(config) + + Convey("check that basic logging should work", func() { + sess.Config.Logger.Log("one") + sess.Config.Logger.Log("two") + sess.Config.Logger.Log("three") + So(len(logs), ShouldEqual, 3) + So(logs, ShouldContain, "one") + So(logs, ShouldContain, "two") + So(logs, ShouldContain, "three") + Reset(func() { + logs = nil + }) + }) + }) + + Convey("check that the default log level is off for both the sdk and kinetic", func() { + sess := getSession(config) + So(sess.Config.LogLevel.Value(), ShouldEqual, aws.LogOff) + So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeFalse) + So(config.logLevel.Value(), ShouldEqual, aws.LogOff) + So(config.logLevel.AtLeast(aws.LogDebug), ShouldBeFalse) + }) + + Convey("check that we can set both the sdk and kinetic log level", func() { + ll := aws.LogDebug | aws.LogDebugWithSigning | logging.LogDebug + config = config.WithLogLevel(ll) + sess := getSession(config) + So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) + So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) + So(config.logLevel.AtLeast(logging.LogDebug), ShouldBeTrue) + }) + + Convey("check that we can set the http.Client Timeout", func() { + config = config.WithHttpClientTimeout(10 * time.Minute) + So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 10*time.Minute) + }) + + Convey("check that we can import configuration from kinetic", func() { + k, err := kinetic.New(kinetic.NewConfig(). + WithEndpoint("bogus-endpoint")) + So(err, ShouldBeNil) + config = config.FromKinetic(k) + sess := getSession(config) + So(aws.StringValue(sess.Config.Endpoint), ShouldEqual, "bogus-endpoint") + }) + + Convey("check that we can set the batch size", func() { + config = config.WithBatchSize(1000) + So(config.batchSize, ShouldEqual, 1000) + }) + + Convey("check that we can set the concurrency limit", func() { + config = config.WithConcurrency(50) + So(config.concurrency, ShouldEqual, 50) + }) + + Convey("check that the default shard iterator is TRIM_HORIZON", func() { + config = config.WithInitialShardIterator(NewShardIterator()) + So(config.shardIterator.shardIteratorType, ShouldEqual, "TRIM_HORIZON") + So(config.shardIterator.getStartingSequenceNumber(), ShouldBeNil) + So(config.shardIterator.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can set the initial shard iterator (to LATEST)", func() { + config = config.WithInitialShardIterator(NewShardIterator().Latest()) + So(config.shardIterator.shardIteratorType, ShouldEqual, "LATEST") + So(config.shardIterator.getStartingSequenceNumber(), ShouldBeNil) + So(config.shardIterator.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can set the read timeout for the GetRecords request", func() { + config = config.WithGetRecordsReadTimeout(10 * time.Second) + So(config.getRecordsReadTimeout, ShouldEqual, 10*time.Second) + }) + + Convey("check that we can configure a stats listener", func() { + config = config.WithStatsListener(&DebugStatsListener{}) + So(config.stats, ShouldHaveSameTypeAs, &DebugStatsListener{}) + }) + }) +} diff --git a/listener/listener.go b/listener/listener.go index aa3e131..baa2475 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -12,10 +12,13 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/awserr" "github.com/aws/aws-sdk-go/aws/request" + "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" - "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/logging" + "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic/utils" ) var ( @@ -98,20 +101,33 @@ func (it *ShardIterator) getTimestamp() *time.Time { return aws.Time(it.timestamp) } +type listenerConfig struct { + stream string + shard string + + batchSize int + concurrency int + shardIterator *ShardIterator + getRecordsReadTimeout time.Duration + stats StatsListener + + logLevel aws.LogLevelType +} + type Listener struct { - Config *Config + *listenerConfig nextShardIterator string - messages chan *kinetic.Message + messages chan *message.Message concurrencySem chan Empty throttleSem chan Empty pipeOfDeath chan Empty - wg sync.WaitGroup consuming bool consumingMu sync.Mutex + session *session.Session client kinesisiface.KinesisAPI clientMu sync.Mutex } @@ -119,130 +135,24 @@ type Listener struct { // NewListener creates a new listener for listening to message on a Kinesis // stream. func NewListener(config *Config) (*Listener, error) { - l := &Listener{ - Config: config, + session, err := config.GetSession() + if err != nil { + return nil, err + } + return &Listener{ + listenerConfig: config.listenerConfig, concurrencySem: make(chan Empty, config.concurrency), throttleSem: make(chan Empty, 5), pipeOfDeath: make(chan Empty), - } - return l, nil + session: session, + }, nil } // Logs a debug message using the AWS SDK logger. func (l *Listener) Log(args ...interface{}) { - l.ensureClient() - if l.client != nil && l.Config.logLevel.AtLeast(aws.LogDebug) { - l.Config.awsConfig.Logger.Log(args...) - } -} - -// CreateStream creates the stream. -func (l *Listener) CreateStream(shards int) error { - if err := l.ensureClient(); err != nil { - return err - } - _, err := l.client.CreateStream(&kinesis.CreateStreamInput{ - StreamName: aws.String(l.Config.stream), - ShardCount: aws.Int64(int64(shards)), - }) - if err != nil { - l.Log("Error creating stream:", err) - } - return err -} - -// WaitUntilActive waits until the stream is active. Timeouts can be set via -// context. Defaults to 18 retries with a 10s delay (180s or 3 minutes total). -func (l *Listener) WaitUntilActive(ctx context.Context, opts ...request.WaiterOption) error { - if err := l.ensureClient(); err != nil { - return err - } - return l.client.WaitUntilStreamExistsWithContext(ctx, &kinesis.DescribeStreamInput{ - StreamName: aws.String(l.Config.stream), // Required - }, opts...) -} - -// DeleteStream deletes the stream. -func (l *Listener) DeleteStream() error { - if err := l.ensureClient(); err != nil { - return err - } - _, err := l.client.DeleteStream(&kinesis.DeleteStreamInput{ - StreamName: aws.String(l.Config.stream), - }) - if err != nil { - l.Log("Error deleting stream:", err) - } - return err -} - -// WaitUntilDeleted waits until the stream is does not exist. Timeouts can be -// set via context. Defaults to 18 retries with a 10s delay (180s or 3 minutes -// total). -func (l *Listener) WaitUntilDeleted(ctx context.Context, opts ...request.WaiterOption) error { - if err := l.ensureClient(); err != nil { - return err + if l.session.Config.LogLevel.Matches(logging.LogDebug) { + l.session.Config.Logger.Log(args...) } - w := request.Waiter{ - Name: "WaitUntilStreamIsDeleted", - MaxAttempts: 18, - Delay: request.ConstantWaiterDelay(10 * time.Second), - Acceptors: []request.WaiterAcceptor{ - { - State: request.SuccessWaiterState, - Matcher: request.ErrorWaiterMatch, - Expected: kinesis.ErrCodeResourceNotFoundException, - }, - }, - Logger: l.Config.awsConfig.Logger, - NewRequest: func(opts []request.Option) (*request.Request, error) { - req, _ := l.client.DescribeStreamRequest(&kinesis.DescribeStreamInput{ - StreamName: aws.String(l.Config.stream), // Required - }) - req.SetContext(ctx) - req.ApplyOptions(opts...) - return req, nil - }, - } - w.ApplyOptions(opts...) - return w.WaitWithContext(ctx) -} - -// GetShards gets a list of shards in a stream. -func (l *Listener) GetShards() ([]string, error) { - if err := l.ensureClient(); err != nil { - return nil, err - } - resp, err := l.client.DescribeStream(&kinesis.DescribeStreamInput{ - StreamName: aws.String(l.Config.stream), - }) - if err != nil { - l.Log("Error describing stream", err) - return nil, err - } - if resp == nil { - return nil, ErrNilDescribeStreamResponse - } - if resp.StreamDescription == nil { - return nil, ErrNilStreamDescription - } - var shards []string - for _, shard := range resp.StreamDescription.Shards { - if shard.ShardId != nil { - shards = append(shards, aws.StringValue(shard.ShardId)) - } - } - return shards, nil -} - -// SetShard sets the shard for the listener. -func (l *Listener) SetShard(shard string) error { - if !l.blockConsumers() { - return ErrCannotSetShard - } - defer l.allowConsumers() - l.Config.shard = shard - return nil } // setNextShardIterator sets the nextShardIterator to use when calling @@ -271,12 +181,12 @@ func (l *Listener) setSequenceNumber(sequenceNumber string) error { if len(sequenceNumber) == 0 { return ErrEmptySequenceNumber } - l.Config.shardIterator.AtSequenceNumber(sequenceNumber) + l.shardIterator.AtSequenceNumber(sequenceNumber) return nil } // ensureClient will lazily make sure we have an AWS Kinesis client. -func (l *Listener) ensureClient() error { +func (l *Listener) ensureClient() { // From the aws-go-sdk documentation: // http://docs.aws.amazon.com/sdk-for-go/api/aws/session/ // @@ -300,16 +210,9 @@ func (l *Listener) ensureClient() error { // this library. l.clientMu.Lock() defer l.clientMu.Unlock() - if l.client != nil { - return nil + if l.client == nil { + l.client = kinesis.New(l.session) } - - session, err := l.Config.GetAwsSession() - if err != nil { - return err - } - l.client = kinesis.New(session) - return nil } // ensureShardIterator will lazily make sure that we have a valid ShardIterator, @@ -320,16 +223,17 @@ func (l *Listener) ensureClient() error { // that only one call to Listen and Retrieve/RetrieveFn can be running at a // time. func (l *Listener) ensureShardIterator() error { + l.ensureClient() if l.nextShardIterator != "" { return nil } resp, err := l.client.GetShardIterator(&kinesis.GetShardIteratorInput{ - ShardId: aws.String(l.Config.shard), // Required - ShardIteratorType: aws.String(l.Config.shardIterator.shardIteratorType), // Required - StreamName: aws.String(l.Config.stream), // Required - StartingSequenceNumber: l.Config.shardIterator.getStartingSequenceNumber(), - Timestamp: l.Config.shardIterator.getTimestamp(), + ShardId: aws.String(l.shard), // Required + ShardIteratorType: aws.String(l.shardIterator.shardIteratorType), // Required + StreamName: aws.String(l.stream), // Required + StartingSequenceNumber: l.shardIterator.getStartingSequenceNumber(), + Timestamp: l.shardIterator.getTimestamp(), }) if err != nil { l.Log(err) @@ -366,10 +270,7 @@ func (l *Listener) throttle(sem chan Empty) { // eventually fires and closes the socket, but this can be susceptible to FD // exhaustion. func (l *Listener) fetchBatch(size int) (int, error) { - if err := l.ensureClient(); err != nil { - return 0, err - } - + l.ensureClient() if err := l.ensureShardIterator(); err != nil { return 0, err } @@ -387,7 +288,7 @@ func (l *Listener) fetchBatch(size int) (int, error) { }) // If debug is turned on, add some handlers for GetRecords logging - if l.Config.logLevel.AtLeast(aws.LogDebug) { + if l.session.Config.LogLevel.AtLeast(aws.LogDebug) { req.Handlers.Send.PushBack(func(r *request.Request) { l.Log("Finished GetRecords Send, took", time.Since(start)) }) @@ -412,9 +313,9 @@ func (l *Listener) fetchBatch(size int) (int, error) { // HTTPResponse.Body must return by. Note that the normal // http.Client Timeout is still in effect. startReadTime = time.Now() - timer := time.NewTimer(l.Config.getRecordsReadTimeout) + timer := time.NewTimer(l.getRecordsReadTimeout) - r.HTTPResponse.Body = &kinetic.TimeoutReadCloser{ + r.HTTPResponse.Body = &utils.TimeoutReadCloser{ ReadCloser: r.HTTPResponse.Body, OnReadFn: func(stream io.ReadCloser, b []byte) (n int, err error) { // The OnReadFn will be called each time @@ -451,7 +352,7 @@ func (l *Listener) fetchBatch(size int) (int, error) { // HTTPResponse.Body, we reset our // timeout and return the results from // the Read() - timer.Reset(l.Config.getRecordsReadTimeout) + timer.Reset(l.getRecordsReadTimeout) n, err = result.n, result.err l.Log(fmt.Sprintf("DEBUG: read %d bytes, took %v", n, time.Since(readStart))) case <-timer.C: @@ -467,7 +368,7 @@ func (l *Listener) fetchBatch(size int) (int, error) { return }, OnCloseFn: func() { - l.Config.stats.AddGetRecordsReadResponseTime(time.Since(startReadTime)) + l.stats.AddGetRecordsReadResponseTime(time.Since(startReadTime)) l.Log("Finished GetRecords body read, took", time.Since(start)) startUnmarshalTime = time.Now() }, @@ -475,30 +376,30 @@ func (l *Listener) fetchBatch(size int) (int, error) { }) req.Handlers.Unmarshal.PushBack(func(r *request.Request) { - l.Config.stats.AddGetRecordsUnmarshalTime(time.Since(startUnmarshalTime)) + l.stats.AddGetRecordsUnmarshalTime(time.Since(startUnmarshalTime)) l.Log("Finished GetRecords Unmarshal, took", time.Since(start)) }) // Send the GetRecords request l.Log("Starting GetRecords build/sign request, took", time.Since(start)) - l.Config.stats.AddGetRecordsCalled(1) + l.stats.AddGetRecordsCalled(1) if err := req.Send(); err != nil { l.Log("Error getting records:", err) return 0, err } // Process Records - l.Log(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), l.Config.shard, time.Since(start))) + l.Log(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), l.shard, time.Since(start))) if resp == nil { return 0, ErrNilGetRecordsResponse } delivered := 0 - l.Config.stats.AddBatchSizeSample(len(resp.Records)) + l.stats.AddBatchSizeSample(len(resp.Records)) for _, record := range resp.Records { if record != nil { delivered++ - l.messages <- &kinetic.Message{*record} - l.Config.stats.AddConsumedSample(1) + l.messages <- &message.Message{*record} + l.stats.AddConsumedSample(1) } if record.SequenceNumber != nil { // We can safely ignore if this call returns @@ -549,7 +450,7 @@ func (l *Listener) blockConsumers() bool { // startConsuming handles any initialization needed in preparation to start // consuming. func (l *Listener) startConsuming() { - l.messages = make(chan *kinetic.Message, l.Config.batchSize) + l.messages = make(chan *message.Message, l.batchSize) } // shouldConsume is a convenience function that allows functions to break their @@ -587,7 +488,7 @@ func (l *Listener) IsConsuming() bool { // Retrieve waits for a message from the stream and return the value. // Cancellation supported through contexts. -func (l *Listener) RetrieveWithContext(ctx context.Context) (*kinetic.Message, error) { +func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, error) { if !l.blockConsumers() { return nil, ErrAlreadyConsuming } @@ -606,14 +507,14 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*kinetic.Message, e return nil, err } if n > 0 { - l.Config.stats.AddDeliveredSample(1) + l.stats.AddDeliveredSample(1) return <-l.messages, nil } } } // Retrieve waits for a message from the stream and return the value. -func (l *Listener) Retrieve() (*kinetic.Message, error) { +func (l *Listener) Retrieve() (*message.Message, error) { return l.RetrieveWithContext(context.TODO()) } @@ -629,7 +530,7 @@ func (l *Listener) RetrieveFnWithContext(ctx context.Context, fn MessageFn) erro wg.Add(1) go fn(msg.Value(), &wg) wg.Wait() - l.Config.stats.AddProcessedSample(1) + l.stats.AddProcessedSample(1) return nil } @@ -662,13 +563,13 @@ func (l *Listener) consume(ctx context.Context) { if !ok { break stop } - _, err = l.fetchBatch(l.Config.batchSize) + _, err = l.fetchBatch(l.batchSize) if err != nil { switch err := err.(type) { case net.Error: if err.Timeout() { - l.Config.stats.AddGetRecordsTimeout(1) + l.stats.AddGetRecordsTimeout(1) l.Log("Received net error:", err.Error()) } else { l.Log("Received unknown net error:", err.Error()) @@ -676,7 +577,7 @@ func (l *Listener) consume(ctx context.Context) { case error: switch err { case ErrTimeoutReadResponseBody: - l.Config.stats.AddGetRecordsReadTimeout(1) + l.stats.AddGetRecordsReadTimeout(1) l.Log("Received error:", err.Error()) case ErrEmptySequenceNumber: fallthrough @@ -694,7 +595,7 @@ func (l *Listener) consume(ctx context.Context) { case awserr.Error: switch err.Code() { case kinesis.ErrCodeProvisionedThroughputExceededException: - l.Config.stats.AddProvisionedThroughputExceeded(1) + l.stats.AddProvisionedThroughputExceeded(1) case kinesis.ErrCodeResourceNotFoundException: fallthrough case kinesis.ErrCodeInvalidArgumentException: @@ -728,7 +629,7 @@ stop: if !ok { break stop } - l.Config.stats.AddDeliveredSample(1) + l.stats.AddDeliveredSample(1) l.concurrencySem <- Empty{} wg.Add(1) go func() { @@ -739,7 +640,7 @@ stop: fnWg.Add(1) fn(msg.Value(), &fnWg) fnWg.Wait() - l.Config.stats.AddProcessedSample(1) + l.stats.AddProcessedSample(1) wg.Done() }() } diff --git a/listener/listener_test.go b/listener/listener_test.go index aff75fe..a874d3d 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -13,49 +13,18 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/aws/awserr" "github.com/aws/aws-sdk-go/aws/request" - "github.com/aws/aws-sdk-go/aws/session" "github.com/aws/aws-sdk-go/service/kinesis" -) - -type DebugLogger struct { - logs []string -} - -func (l *DebugLogger) Log(args ...interface{}) { - l.logs = append(l.logs, fmt.Sprint(args...)) -} - -func (l *DebugLogger) Reset() { - l.logs = nil -} -type DebugStatsListener struct{} - -func (l *DebugStatsListener) AddConsumedSample(int) {} -func (l *DebugStatsListener) AddDeliveredSample(int) {} -func (l *DebugStatsListener) AddProcessedSample(int) {} -func (l *DebugStatsListener) AddBatchSizeSample(int) {} -func (l *DebugStatsListener) AddGetRecordsCalled(int) {} -func (l *DebugStatsListener) AddProvisionedThroughputExceeded(int) {} -func (l *DebugStatsListener) AddGetRecordsTimeout(int) {} -func (l *DebugStatsListener) AddGetRecordsReadTimeout(int) {} -func (l *DebugStatsListener) AddGetRecordsReadResponseTime(time.Duration) {} -func (l *DebugStatsListener) AddGetRecordsUnmarshalTime(time.Duration) {} - -func getSession(config *Config) *session.Session { - sess, err := config.GetAwsSession() - So(err, ShouldBeNil) - So(sess, ShouldNotBeNil) - return sess -} + "github.com/rewardStyle/kinetic" +) func putRecord(l *Listener, b []byte) (*string, error) { + l.ensureClient() resp, err := l.client.PutRecord(&kinesis.PutRecordInput{ Data: b, PartitionKey: aws.String("dummy"), - StreamName: aws.String(l.Config.stream), + StreamName: aws.String(l.stream), }) if err != nil { return nil, err @@ -63,148 +32,6 @@ func putRecord(l *Listener, b []byte) (*string, error) { return resp.SequenceNumber, nil } -func TestNewConfig(t *testing.T) { - Convey("given a new listener config", t, func() { - config := NewConfig("some-stream") - - Convey("check the default values for its non-zero config", func() { - So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 5*time.Minute) - So(config.stream, ShouldEqual, "some-stream") - So(config.logLevel, ShouldEqual, aws.LogOff) - So(config.stats, ShouldHaveSameTypeAs, &NilStatsListener{}) - So(config.batchSize, ShouldEqual, 10000) - So(config.concurrency, ShouldEqual, 10000) - So(config.shardIterator.shardIteratorType, ShouldEqual, "TRIM_HORIZON") - So(config.getRecordsReadTimeout, ShouldEqual, 1*time.Second) - }) - - Convey("check that we can retrieve an aws.Session from it ", func() { - getSession(config) - }) - - Convey("check that we can set credentials", func() { - config = config.WithCredentials("access-key", "secret-key", "security-token") - sess := getSession(config) - creds, err := sess.Config.Credentials.Get() - So(err, ShouldBeNil) - So(creds.AccessKeyID, ShouldEqual, "access-key") - So(creds.SecretAccessKey, ShouldEqual, "secret-key") - So(creds.SessionToken, ShouldEqual, "security-token") - }) - - Convey("check that we can set the region", func() { - config = config.WithRegion("my-region") - sess := getSession(config) - So(aws.StringValue(sess.Config.Region), ShouldEqual, "my-region") - }) - - Convey("check that we can set the endpoint", func() { - config = config.WithEndpoint("my-endpoint") - sess := getSession(config) - So(aws.StringValue(sess.Config.Endpoint), ShouldEqual, "my-endpoint") - }) - - Convey("check that we can configure a logger", func() { - l := &DebugLogger{} - config = config.WithLogger(l) - sess := getSession(config) - So(sess.Config.Logger, ShouldHaveSameTypeAs, l) - - Convey("check that basic logging should work", func() { - sess.Config.Logger.Log("one") - sess.Config.Logger.Log("two") - sess.Config.Logger.Log("three") - So(len(l.logs), ShouldEqual, 3) - So(l.logs, ShouldContain, "one") - So(l.logs, ShouldContain, "two") - So(l.logs, ShouldContain, "three") - Reset(func() { - l.Reset() - }) - }) - }) - - Convey("check that the default log level is off for both the sdk and kinetic", func() { - sess := getSession(config) - So(sess.Config.LogLevel.Value(), ShouldEqual, aws.LogOff) - So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeFalse) - So(config.logLevel.Value(), ShouldEqual, aws.LogOff) - So(config.logLevel.AtLeast(aws.LogDebug), ShouldBeFalse) - }) - - Convey("check that we can set the sdk log level", func() { - config = config.WithLogLevel(aws.LogDebug | aws.LogDebugWithSigning) - sess := getSession(config) - So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) - So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) - So(config.logLevel.Value(), ShouldEqual, aws.LogOff) - So(config.logLevel.AtLeast(aws.LogDebug), ShouldBeFalse) - }) - - Convey("check that we can set the kinetic log level", func() { - config = config.WithLogLevel((aws.LogDebug | 1) << 16) - sess := getSession(config) - So(sess.Config.LogLevel.Value(), ShouldEqual, aws.LogOff) - So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeFalse) - So(config.logLevel.AtLeast(aws.LogDebug), ShouldBeTrue) - So(config.logLevel.Matches(1), ShouldBeTrue) - }) - - Convey("check that we can set both the sdk and kinetic log level", func() { - config = config.WithLogLevel(aws.LogDebug | aws.LogDebugWithSigning | ((aws.LogDebug | 1) << 16)) - sess := getSession(config) - So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) - So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) - So(config.logLevel.AtLeast(aws.LogDebug), ShouldBeTrue) - So(config.logLevel.Matches(1), ShouldBeTrue) - }) - - Convey("check that we can set the http.Client Timeout", func() { - config = config.WithHttpClientTimeout(10 * time.Minute) - So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 10*time.Minute) - }) - - Convey("check that we can configure a stats listener", func() { - config = config.WithStatsListener(&DebugStatsListener{}) - So(config.stats, ShouldHaveSameTypeAs, &DebugStatsListener{}) - }) - - Convey("check that we can set the shard", func() { - config = config.WithShardId("some-shard") - So(config.shard, ShouldEqual, "some-shard") - }) - - Convey("check that we can set the batch size", func() { - config = config.WithBatchSize(1000) - So(config.batchSize, ShouldEqual, 1000) - }) - - Convey("check that we can set the concurrency limit", func() { - config = config.WithConcurrency(50) - So(config.concurrency, ShouldEqual, 50) - }) - - Convey("check that the default shard iterator is TRIM_HORIZON", func() { - config = config.WithInitialShardIterator(NewShardIterator()) - So(config.shardIterator.shardIteratorType, ShouldEqual, "TRIM_HORIZON") - So(config.shardIterator.getStartingSequenceNumber(), ShouldBeNil) - So(config.shardIterator.getTimestamp(), ShouldBeNil) - }) - - Convey("check that we can set the initial shard iterator (to LATEST)", func() { - config = config.WithInitialShardIterator(NewShardIterator().Latest()) - So(config.shardIterator.shardIteratorType, ShouldEqual, "LATEST") - So(config.shardIterator.getStartingSequenceNumber(), ShouldBeNil) - So(config.shardIterator.getTimestamp(), ShouldBeNil) - }) - - Convey("check that we can set the read timeout for the GetRecords request", func() { - config = config.WithGetRecordsReadTimeout(10 * time.Second) - So(config.getRecordsReadTimeout, ShouldEqual, 10*time.Second) - }) - }) -} - func TestShardIterator(t *testing.T) { Convey("given a new shard iterator", t, func() { it := NewShardIterator() @@ -255,23 +82,28 @@ func TestShardIterator(t *testing.T) { func TestListener(t *testing.T) { Convey("given a listener", t, func() { - logger := &DebugLogger{} - l, err := NewListener(NewConfig("some-stream"). + k, err := kinetic.New(kinetic.NewConfig(). WithRegion("some-region"). - WithEndpoint("http://127.0.0.1:4567"). - WithLogger(logger). - WithLogLevel(aws.LogDebug << 16). + WithEndpoint("http://127.0.0.1:4567")) + + stream := "some-listener-stream" + + err = k.CreateStream(stream, 1) + So(err, ShouldBeNil) + + err = k.WaitUntilStreamExists(context.TODO(), stream, request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + So(err, ShouldBeNil) + + shards, err := k.GetShards(stream) + So(err, ShouldBeNil) + So(len(shards), ShouldEqual, 1) + + l, err := NewListener(NewConfig(stream, shards[0]). + FromKinetic(k). WithConcurrency(10)) So(l, ShouldNotBeNil) So(err, ShouldBeNil) - Convey("check that logging works", func() { - l.Log("foo") - So(logger.logs, ShouldContain, "foo") - So(len(logger.logs), ShouldEqual, 1) - Reset(func() { logger.Reset() }) - }) - Convey("check that calling ensureClient twice doesn't overwrite existing client", func() { So(l.client, ShouldBeNil) l.ensureClient() @@ -281,13 +113,6 @@ func TestListener(t *testing.T) { So(l.client, ShouldEqual, client) }) - Convey("check deleting a non-existent stream returns an error", func() { - err := l.DeleteStream() - So(err, ShouldNotBeNil) - e := err.(awserr.Error) - So(e.Code(), ShouldEqual, kinesis.ErrCodeResourceNotFoundException) - }) - Convey("check that setting an empty shard iterator returns an error", func() { err := l.setNextShardIterator("") So(err, ShouldEqual, ErrEmptyShardIterator) @@ -298,230 +123,179 @@ func TestListener(t *testing.T) { So(err, ShouldEqual, ErrEmptySequenceNumber) }) - Convey("check getting shards on a non-existent stream returns an error", func() { - shards, err := l.GetShards() - So(shards, ShouldBeNil) - So(err, ShouldNotBeNil) - e := err.(awserr.Error) - So(e.Code(), ShouldEqual, kinesis.ErrCodeResourceNotFoundException) - }) - - Convey("check that we can create a stream", func() { - err := l.CreateStream(1) + Convey("check that we can get the TRIM_HORIZON shard iterator", func() { + err := l.ensureShardIterator() So(err, ShouldBeNil) - err = l.WaitUntilActive(context.TODO(), request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) - So(err, ShouldBeNil) - - Convey("check that creating an existing stream returns an error", func() { - err := l.CreateStream(1) - So(err, ShouldNotBeNil) - e := err.(awserr.Error) - So(e.Code(), ShouldEqual, kinesis.ErrCodeResourceInUseException) - }) - - Convey("check that getting the shard iterator with an empty shard returns an error", func() { - So(l.Config.shard, ShouldBeEmpty) - err := l.ensureShardIterator() - So(err, ShouldNotBeNil) - e := err.(awserr.Error) - So(e.Code(), ShouldEqual, request.InvalidParameterErrCode) - }) - - Convey("check that getting the shard iterator without an invalid shard returns an error", func() { - l.SetShard("some-shard") - err := l.ensureShardIterator() - So(err, ShouldNotBeNil) - e := err.(awserr.Error) - So(e.Code(), ShouldEqual, kinesis.ErrCodeResourceNotFoundException) - }) + So(l.nextShardIterator, ShouldNotBeEmpty) + }) - Convey("check that we can obtain a list of shards", func() { - shards, err := l.GetShards() + Convey("check that we can retrieve records one by one", func() { + data := []string{"foo", "bar"} + for n, datum := range data { + seq, err := putRecord(l, []byte(datum)) So(err, ShouldBeNil) - So(len(shards), ShouldEqual, 1) - - Convey("check that we can set the shard id", func() { - l.SetShard(shards[0]) - So(l.Config.shard, ShouldEqual, shards[0]) - - Convey("check that we can get the TRIM_HORIZON shard iterator", func() { - err := l.ensureShardIterator() - So(err, ShouldBeNil) - So(l.nextShardIterator, ShouldNotBeEmpty) - }) + So(seq, ShouldNotBeNil) + msg, err := l.Retrieve() + So(err, ShouldBeNil) + So(string(msg.Value()), ShouldEqual, datum) + Convey(fmt.Sprintf("check that iteration %d properly advanced the shard iterator", n), func() { + So(l.shardIterator.shardIteratorType, ShouldEqual, "AT_SEQUENCE_NUMBER") + So(l.shardIterator.sequenceNumber, ShouldEqual, *seq) + }) + } + }) - Convey("check that we can retrieve records one by one", func() { - data := []string{"foo", "bar"} - for n, datum := range data { - seq, err := putRecord(l, []byte(datum)) - So(err, ShouldBeNil) - So(seq, ShouldNotBeNil) - msg, err := l.Retrieve() - So(err, ShouldBeNil) - So(string(msg.Value()), ShouldEqual, datum) - Convey(fmt.Sprintf("check that iteration %d properly advanced the shard iterator", n), func() { - So(l.Config.shardIterator.shardIteratorType, ShouldEqual, "AT_SEQUENCE_NUMBER") - So(l.Config.shardIterator.sequenceNumber, ShouldEqual, *seq) - }) - } - }) + Convey("check that retrieve will block until record comes", func(c C) { + start := time.Now() + data := "hello" + go func() { + <-time.After(1 * time.Second) + _, err := putRecord(l, []byte(data)) + c.So(err, ShouldBeNil) + }() + msg, err := l.Retrieve() + elapsed := time.Since(start) + Printf("(it blocked %f seconds)\n", elapsed.Seconds()) + So(err, ShouldBeNil) + So(string(msg.Value()), ShouldEqual, data) + So(elapsed.Seconds(), ShouldBeGreaterThan, 1) + }) - Convey("check that retrieve will block until record comes", func(c C) { - start := time.Now() - data := "hello" - go func() { - <-time.After(1 * time.Second) - _, err := putRecord(l, []byte(data)) - c.So(err, ShouldBeNil) - }() - msg, err := l.Retrieve() - elapsed := time.Since(start) - Printf("(it blocked %f seconds)\n", elapsed.Seconds()) - So(err, ShouldBeNil) - So(string(msg.Value()), ShouldEqual, data) - So(elapsed.Seconds(), ShouldBeGreaterThan, 1) - }) + Convey("check that we can use a context to cancel the retrieve", func() { + start := time.Now() + ctx, _ := context.WithTimeout(context.TODO(), 1*time.Second) + _, err := l.RetrieveWithContext(ctx) + elapsed := time.Since(start) + Printf("(it blocked %f seconds)\n", elapsed.Seconds()) + So(err, ShouldNotBeNil) + So(err, ShouldHaveSameTypeAs, context.DeadlineExceeded) + So(elapsed.Seconds(), ShouldBeGreaterThan, 1) + }) - Convey("check that we can use a context to cancel the retrieve", func() { - start := time.Now() - ctx, _ := context.WithTimeout(context.TODO(), 1*time.Second) - _, err := l.RetrieveWithContext(ctx) - elapsed := time.Since(start) - Printf("(it blocked %f seconds)\n", elapsed.Seconds()) - So(err, ShouldNotBeNil) - So(err, ShouldHaveSameTypeAs, context.DeadlineExceeded) - So(elapsed.Seconds(), ShouldBeGreaterThan, 1) - }) + Convey("check that we can use a context to cancel the retrieve (again)", func() { + start := time.Now() + ctx, _ := context.WithTimeout(context.TODO(), 10*time.Millisecond) + _, err := l.RetrieveWithContext(ctx) + elapsed := time.Since(start) + Printf("(it blocked %f seconds)\n", elapsed.Seconds()) + So(err, ShouldNotBeNil) + So(err, ShouldHaveSameTypeAs, context.DeadlineExceeded) + So(elapsed.Seconds(), ShouldBeGreaterThan, 0.01) + }) - Convey("check that we can use a context to cancel the retrieve (again)", func() { - start := time.Now() - ctx, _ := context.WithTimeout(context.TODO(), 10*time.Millisecond) - _, err := l.RetrieveWithContext(ctx) - elapsed := time.Since(start) - Printf("(it blocked %f seconds)\n", elapsed.Seconds()) - So(err, ShouldNotBeNil) - So(err, ShouldHaveSameTypeAs, context.DeadlineExceeded) - So(elapsed.Seconds(), ShouldBeGreaterThan, 0.01) - }) + Convey("check that retrieve still works with a canceller if a message comes before the deadline", func(c C) { + ctx, _ := context.WithTimeout(context.TODO(), 5*time.Second) - Convey("check that retrieve still works with a canceller if a message comes before the deadline", func(c C) { - ctx, _ := context.WithTimeout(context.TODO(), 5*time.Second) - - data := "goodbye" - go func() { - <-time.After(1 * time.Second) - _, err := putRecord(l, []byte(data)) - c.So(err, ShouldBeNil) - }() - msg, err := l.RetrieveWithContext(ctx) - So(err, ShouldBeNil) - So(string(msg.Value()), ShouldEqual, data) - }) + data := "goodbye" + go func() { + <-time.After(1 * time.Second) + _, err := putRecord(l, []byte(data)) + c.So(err, ShouldBeNil) + }() + msg, err := l.RetrieveWithContext(ctx) + So(err, ShouldBeNil) + So(string(msg.Value()), ShouldEqual, data) + }) - Convey("check that retrieve properly blocks other retrieves and attempts to set the shard id", func(c C) { - var wg sync.WaitGroup - wg.Add(1) - go func() { - ctx, _ := context.WithTimeout(context.TODO(), 1*time.Second) - _, err := l.RetrieveWithContext(ctx) - c.So(err, ShouldNotBeNil) - c.So(err, ShouldHaveSameTypeAs, context.DeadlineExceeded) - wg.Done() - }() - <-time.After(10 * time.Millisecond) - So(l.IsConsuming(), ShouldBeTrue) - _, err := l.Retrieve() - So(err, ShouldEqual, ErrAlreadyConsuming) - err = l.SetShard("bogus") - So(err, ShouldEqual, ErrCannotSetShard) - wg.Wait() - }) + Convey("check that retrieve properly blocks other retrieves and attempts to set the shard id", func(c C) { + var wg sync.WaitGroup + wg.Add(1) + go func() { + ctx, _ := context.WithTimeout(context.TODO(), 1*time.Second) + _, err := l.RetrieveWithContext(ctx) + c.So(err, ShouldNotBeNil) + c.So(err, ShouldHaveSameTypeAs, context.DeadlineExceeded) + wg.Done() + }() + <-time.After(10 * time.Millisecond) + So(l.IsConsuming(), ShouldBeTrue) + _, err := l.Retrieve() + So(err, ShouldEqual, ErrAlreadyConsuming) + wg.Wait() + }) - Convey("check that throttle mechanism prevents more than 5 calls to get records", func() { - start := time.Now() - secs := []float64{} - for i := 1; i <= 6; i++ { - start := time.Now() - l.fetchBatch(1) - secs = append(secs, time.Since(start).Seconds()) - } - elapsed := time.Since(start).Seconds() - So(elapsed, ShouldBeGreaterThan, 1) - Printf("%f seconds total, (%v)", elapsed, secs) - }) + Convey("check that throttle mechanism prevents more than 5 calls to get records", func() { + start := time.Now() + secs := []float64{} + for i := 1; i <= 6; i++ { + start := time.Now() + l.fetchBatch(1) + secs = append(secs, time.Since(start).Seconds()) + } + elapsed := time.Since(start).Seconds() + So(elapsed, ShouldBeGreaterThan, 1) + Printf("%f seconds total, (%v)", elapsed, secs) + }) - Convey("check that retrievefn can deliver messages to the fn", func(c C) { - called := false - data := "retrieved" - _, err := putRecord(l, []byte(data)) - So(err, ShouldBeNil) - err = l.RetrieveFn(func(b []byte, wg *sync.WaitGroup) { - called = true - // Note that because this is called in a goroutine, we have to use - // the goconvey context - c.So(string(b), ShouldEqual, data) - wg.Done() - }) - So(err, ShouldBeNil) - So(called, ShouldBeTrue) - }) + Convey("check that retrievefn can deliver messages to the fn", func(c C) { + called := false + data := "retrieved" + _, err := putRecord(l, []byte(data)) + So(err, ShouldBeNil) + err = l.RetrieveFn(func(b []byte, wg *sync.WaitGroup) { + called = true + // Note that because this is called in a goroutine, we have to use + // the goconvey context + c.So(string(b), ShouldEqual, data) + wg.Done() + }) + So(err, ShouldBeNil) + So(called, ShouldBeTrue) + }) - Convey("check that listen can deliver messages to fn", func(c C) { - planets := []string{"mercury", "venus", "earth", "mars", "jupiter", "saturn", "neptune", "uranus"} - var count int64 = 0 - var wg sync.WaitGroup - wg.Add(1) - go func() { - defer wg.Done() - l.Listen(func(b []byte, wg *sync.WaitGroup) { - defer wg.Done() - atomic.AddInt64(&count, 1) - }) - }() - for _, planet := range planets { - _, err := putRecord(l, []byte(planet)) - So(err, ShouldBeNil) - } - close(l.pipeOfDeath) - wg.Wait() - So(count, ShouldEqual, len(planets)) - }) + Convey("check that listen can deliver messages to fn", func(c C) { + planets := []string{"mercury", "venus", "earth", "mars", "jupiter", "saturn", "neptune", "uranus"} + var count int64 = 0 + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + l.Listen(func(b []byte, wg *sync.WaitGroup) { + defer wg.Done() + atomic.AddInt64(&count, 1) + }) + }() + for _, planet := range planets { + _, err := putRecord(l, []byte(planet)) + So(err, ShouldBeNil) + } + close(l.pipeOfDeath) + wg.Wait() + So(count, ShouldEqual, len(planets)) + }) - Convey("check that listen can be cancelled by context", func(c C) { - for i := 0; i < 20; i++ { - _, err := putRecord(l, []byte(fmt.Sprintf("%d", i))) - So(err, ShouldBeNil) + Convey("check that listen can be cancelled by context", func(c C) { + for i := 0; i < 20; i++ { + _, err := putRecord(l, []byte(fmt.Sprintf("%d", i))) + So(err, ShouldBeNil) + } + var count int64 = 0 + var wg sync.WaitGroup + wg.Add(1) + go func() { + ctx, cancel := context.WithCancel(context.TODO()) + defer wg.Done() + l.ListenWithContext(ctx, func(b []byte, wg *sync.WaitGroup) { + defer wg.Done() + time.AfterFunc(time.Duration(rand.Intn(10))*time.Second, func() { + n, err := strconv.Atoi(string(b)) + c.So(err, ShouldBeNil) + atomic.AddInt64(&count, 1) + if n == 15 { + cancel() } - var count int64 = 0 - var wg sync.WaitGroup - wg.Add(1) - go func() { - ctx, cancel := context.WithCancel(context.TODO()) - defer wg.Done() - l.ListenWithContext(ctx, func(b []byte, wg *sync.WaitGroup) { - defer wg.Done() - time.AfterFunc(time.Duration(rand.Intn(10))*time.Second, func() { - n, err := strconv.Atoi(string(b)) - c.So(err, ShouldBeNil) - atomic.AddInt64(&count, 1) - if n == 15 { - cancel() - } - }) - }) - }() - wg.Wait() - So(count, ShouldBeBetweenOrEqual, 1, 20) - Printf("(count was %d)", count) }) }) + }() + wg.Wait() + So(count, ShouldBeBetweenOrEqual, 1, 20) + Printf("(count was %d)", count) + }) - }) - - Reset(func() { - l.DeleteStream() - l.WaitUntilDeleted(context.TODO(), request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) - }) + Reset(func() { + k.DeleteStream(stream) + k.WaitUntilStreamDeleted(context.TODO(), stream, request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) }) + }) } diff --git a/listener_test.go b/listener_test.go deleted file mode 100644 index 9dd77e8..0000000 --- a/listener_test.go +++ /dev/null @@ -1,166 +0,0 @@ -package kinetic - -import ( - "errors" - "runtime" - "sync" - "syscall" - "testing" - "time" - - "github.com/aws/aws-sdk-go/aws" - awsKinesis "github.com/aws/aws-sdk-go/service/kinesis" - awsKinesisIface "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" - . "github.com/smartystreets/goconvey/convey" -) - -const testEndpoint = "http://127.0.0.1:4567" - -func CreateAndWaitForStream(client awsKinesisIface.KinesisAPI, name string) { - client.CreateStream(&awsKinesis.CreateStreamInput{ - StreamName: aws.String(name), - ShardCount: aws.Int64(1), - }) - stream := &awsKinesis.DescribeStreamInput{StreamName: aws.String(name), Limit: aws.Int64(1)} - client.WaitUntilStreamExists(stream) -} - -func TestListenerStop(t *testing.T) { - listener, _ := new(Listener).Init() - listener.NewEndpoint(testEndpoint, "stream-name") - CreateAndWaitForStream(listener.client, "stream-name") - listener.ReInit() - - Convey("Given a running listener", t, func() { - go listener.Listen(func(msg []byte, wg *sync.WaitGroup) { - wg.Done() - }) - - Convey("It should stop listening if sent an interrupt signal", func() { - listener.interrupts <- syscall.SIGINT - runtime.Gosched() - // Let it finish stopping - time.Sleep(3 * time.Second) - - So(listener.IsListening(), ShouldEqual, false) - }) - }) - - listener.Close() -} - -func TestListenerSyncStop(t *testing.T) { - listener, _ := new(Listener).Init() - listener.NewEndpoint(testEndpoint, "stream-name") - CreateAndWaitForStream(listener.client, "stream-name") - listener.ReInit() - - Convey("Given a running listener", t, func() { - go listener.Listen(func(msg []byte, wg *sync.WaitGroup) { - wg.Done() - }) - - Convey("It should stop listening if sent an interrupt signal", func() { - err := listener.CloseSync() - So(err, ShouldBeNil) - So(listener.IsListening(), ShouldEqual, false) - }) - }) - - listener.Close() -} - -func TestListenerError(t *testing.T) { - listener, _ := new(Listener).Init() - listener.NewEndpoint(testEndpoint, "stream-name") - CreateAndWaitForStream(listener.client, "stream-name") - listener.ReInit() - - Convey("Given a running listener", t, func() { - go listener.Listen(func(msg []byte, wg *sync.WaitGroup) { - wg.Done() - }) - - Convey("It should handle errors successfully", func() { - listener.errors <- errors.New("All your base are belong to us") - - // Let the error propagate - <-time.After(3 * time.Second) - - So(listener.getErrCount(), ShouldNotEqual, 0) - So(listener.IsListening(), ShouldEqual, true) - }) - }) - - listener.Close() -} - -func TestListenerMessage(t *testing.T) { - listener, _ := new(Listener).Init() - listener.NewEndpoint(testEndpoint, "stream-name") - CreateAndWaitForStream(listener.client, "stream-name") - listener.ReInit() - - go listener.Listen(func(msg []byte, wg *sync.WaitGroup) { - wg.Done() - }) - - time.Sleep(3 * time.Second) - - for _, c := range cases { - Convey("Given a running listener", t, func() { - listener.messages <- new(Message).Init(c.message, "test") - - Convey("It should handle messages successfully", func() { - So(listener.IsListening(), ShouldEqual, true) - So(listener.Errors(), ShouldNotResemble, nil) - }) - }) - } - time.Sleep(2 * time.Second) - listener.CloseSync() -} - -func TestRetrieveMessage(t *testing.T) { - listener, _ := new(Listener).InitC("your-stream", "0", ShardIterTypes[3], "accesskey", "secretkey", "us-east-1", 10) - producer, _ := new(KinesisProducer).InitC("your-stream", "0", ShardIterTypes[3], "accesskey", "secretkey", "us-east-1", 10) - - listener.NewEndpoint(testEndpoint, "your-stream") - producer.NewEndpoint(testEndpoint, "your-stream") - CreateAndWaitForStream(listener.client, "your-stream") - listener.ReInit() - producer.ReInit() - - time.Sleep(10 * time.Millisecond) - for _, c := range cases { - Convey("Given a valid message", t, func() { - producer.Send(new(Message).Init(c.message, "test")) - time.Sleep(3 * time.Millisecond) - - Convey("It should be passed on the queue without error", func() { - msg, err := listener.Retrieve() - So(err, ShouldBeNil) - - So(string(msg.Value()), ShouldResemble, string(c.message)) - }) - }) - } - time.Sleep(1 * time.Second) - producer.CloseSync() - time.Sleep(1 * time.Second) - listener.CloseSync() -} - -var cases = []struct { - message []byte -}{ - { - message: []byte(`{"foo":"bar"}`), - }, - { - message: []byte(`{"bar":"baz"}`), - }, - { - message: []byte(`{"baz":"qux"}`), - }, -} diff --git a/logging/logging.go b/logging/logging.go new file mode 100644 index 0000000..0d2631e --- /dev/null +++ b/logging/logging.go @@ -0,0 +1,14 @@ +package logging + +import ( + "github.com/aws/aws-sdk-go/aws" +) + +const ( + LogOff aws.LogLevelType = (iota * 0x1000) << 16 + LogDebug +) + +const ( + LogPlaceholder aws.LogLevelType = LogDebug | (1 << (16 + iota)) +) diff --git a/message.go b/message/message.go similarity index 85% rename from message.go rename to message/message.go index f3dd50a..333b7b9 100644 --- a/message.go +++ b/message/message.go @@ -1,17 +1,17 @@ -package kinetic +package message -import awsKinesis "github.com/aws/aws-sdk-go/service/kinesis" +import kinesis "github.com/aws/aws-sdk-go/service/kinesis" // Message represents an item on the Kinesis stream type Message struct { - awsKinesis.Record + kinesis.Record } // Init initializes a Message. // Currently we are ignoring sequenceNumber. func (k *Message) Init(msg []byte, key string) *Message { return &Message{ - awsKinesis.Record{ + kinesis.Record{ Data: msg, PartitionKey: &key, }, diff --git a/producer.go b/producer.go deleted file mode 100644 index 82b5d81..0000000 --- a/producer.go +++ /dev/null @@ -1,398 +0,0 @@ -package kinetic - -import ( - "errors" - "log" - "os" - "os/signal" - "runtime" - "strconv" - "sync" - "syscall" - "time" - - "github.com/aws/aws-sdk-go/aws" - awsKinesis "github.com/aws/aws-sdk-go/service/kinesis" -) - -var _ = awsKinesis.EndpointsID - -var ( - // ErrThroughputExceeded represents an error when the Kinesis throughput has been exceeded - ErrThroughputExceeded = errors.New("Configured AWS Kinesis throughput has been exceeded") - // ErrKinesisFailure represents a generic internal AWS Kinesis error - ErrKinesisFailure = errors.New("AWS Kinesis internal failure") - // ErrBadConcurrency represents an error when the provided concurrency value is invalid - ErrBadConcurrency = errors.New("Concurrency must be greater than zero") - // ErrDroppedMessage represents an error the channel is full and messages are being dropped - ErrDroppedMessage = errors.New("Channel is full, dropped message") -) - -// Producer is an interface for sending messages to a stream of some sort. -type Producer interface { - Init() (Producer, error) - InitC(stream, shard, shardIterType, accessKey, secretKey, region string, concurrency int) (Producer, error) - NewEndpoint(endpoint, stream string) (err error) - ReInit() - IsProducing() bool - Send(msg *Message) - TryToSend(msg *Message) error - Close() error - CloseSync() error -} - -// KinesisProducer keeps a queue of messages on a channel and continually attempts -// to POST the records using the PutRecords method. If the messages were -// not sent successfully they are placed back on the queue to retry -type KinesisProducer struct { - *kinesis - - concurrency int - concurrencyMu sync.Mutex - sem chan Empty - - wg sync.WaitGroup - - producing bool - producingMu sync.Mutex - - errors chan error - - // We need to ensure that the sent messages were successfully processed - // before removing them from this local queue - messages chan *Message - interrupts chan os.Signal -} - -func (p *KinesisProducer) init(stream, shard, shardIterType, accessKey, secretKey, region string, concurrency int) (Producer, error) { - var err error - if concurrency < 1 { - return nil, ErrBadConcurrency - } - if stream == "" { - return nil, ErrNullStream - } - - p.setConcurrency(concurrency) - - p.initChannels() - - p.kinesis, err = new(kinesis).init(stream, shard, shardIterType, accessKey, secretKey, region) - if err != nil { - return p, err - } - - return p.activate() -} - -func (p *KinesisProducer) initChannels() { - p.sem = make(chan Empty, p.getConcurrency()) - p.errors = make(chan error, p.getConcurrency()) - p.messages = make(chan *Message, p.msgBufSize()) - - p.interrupts = make(chan os.Signal, 1) - signal.Notify(p.interrupts, os.Interrupt) -} - -func (p *KinesisProducer) setConcurrency(concurrency int) { - p.concurrencyMu.Lock() - p.concurrency = concurrency - p.concurrencyMu.Unlock() -} - -func (p *KinesisProducer) getConcurrency() int { - p.concurrencyMu.Lock() - defer p.concurrencyMu.Unlock() - return p.concurrency -} - -func (p *KinesisProducer) msgBufSize() int { - p.concurrencyMu.Lock() - defer p.concurrencyMu.Unlock() - return p.concurrency * 1000 -} - -func (p *KinesisProducer) activate() (Producer, error) { - // Is the stream ready? - active, err := p.checkActive() - if err != nil || !active { - if err != nil { - return p, err - } - return p, ErrNotActive - } - - // go start feeder consumer and let listen processes them - go p.produce() - - return p, err -} - -// Init initializes a producer with the params specified in the configuration file -func (p *KinesisProducer) Init() (Producer, error) { - return p.init(conf.Kinesis.Stream, conf.Kinesis.Shard, ShardIterTypes[conf.Kinesis.ShardIteratorType], conf.AWS.AccessKey, conf.AWS.SecretKey, conf.AWS.Region, conf.Concurrency.Producer) -} - -// InitC initializes a producer with the specified configuration: stream, shard, shard-iter-type, access-key, secret-key, and region -func (p *KinesisProducer) InitC(stream, shard, shardIterType, accessKey, secretKey, region string, concurrency int) (Producer, error) { - return p.init(stream, shard, shardIterType, accessKey, secretKey, region, concurrency) -} - -// NewEndpoint re-initializes kinesis client with new endpoint. Used for testing with kinesalite -func (p *KinesisProducer) NewEndpoint(endpoint, stream string) (err error) { - // Re-initialize kinesis client for testing - p.kinesis.client, err = p.kinesis.newClient(endpoint, stream) - return -} - -// ReInit re-initializes the shard iterator. Used with conjucntion with NewEndpoint -func (p *KinesisProducer) ReInit() { - if !p.IsProducing() { - go p.produce() - } - return -} - -// Each shard can support up to 1,000 records per second for writes, up to a maximum total -// data write rate of 1 MB per second (including partition keys). This write limit applies -// to operations such as PutRecord and PutRecords. -// TODO: payload inspection & throttling -// http://docs.aws.amazon.com/kinesis/latest/dev/service-sizes-and-limits.html -func (p *KinesisProducer) kinesisFlush(counter *int, timer *time.Time) bool { - // 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() - *counter = 0 - } - - *counter++ - - // If we have attempted 1000 times and it has been less than one second - // since we started sending then we need to wait for the second to finish - if *counter >= kinesisWritesPerSec && !(time.Now().After(timer.Add(1 * time.Second))) { - // Wait for the remainder of the second - timer and counter - // will be reset on next pass - time.Sleep(1*time.Second - time.Since(*timer)) - } - - return true -} - -func (p *KinesisProducer) setProducing(producing bool) { - p.producingMu.Lock() - p.producing = producing - p.producingMu.Unlock() - -} - -// IsProducing identifies whether or not the messages are queued for POSTing to the stream -func (p *KinesisProducer) IsProducing() bool { - p.producingMu.Lock() - defer p.producingMu.Unlock() - return p.producing -} - -// http://docs.aws.amazon.com/kinesis/latest/APIReference/API_PutRecords.html -// -// Maximum of 1000 requests a second for a single shard. Each PutRecords can -// accept a maximum of 500 records per request and each record can be as large -// as 1MB per record OR 5MB per request -func (p *KinesisProducer) produce() { - p.setProducing(true) - - counter := 0 - timer := time.Now() - -stop: - for { - getLock(p.sem) - - select { - case msg := <-p.Messages(): - p.incMsgCount() - - if conf.Debug.Verbose && p.getMsgCount()%100 == 0 { - log.Println("Received message to send. Total messages received: " + strconv.FormatInt(p.getMsgCount(), 10)) - } - kargs := &awsKinesis.PutRecordsInput{StreamName: aws.String(p.stream)} - kargs.Records = append( - kargs.Records, - &awsKinesis.PutRecordsRequestEntry{ - Data: msg.Value(), - PartitionKey: aws.String(string(msg.Key()))}) - if p.kinesisFlush(&counter, &timer) { - p.wg.Add(1) - go func() { - p.sendRecords(kargs) - p.wg.Done() - }() - } - - <-p.sem - case sig := <-p.interrupts: - go p.handleInterrupt(sig) - break stop - case err := <-p.Errors(): - p.incErrCount() - p.wg.Add(1) - go p.handleError(err) - } - } - - p.setProducing(false) -} - -// Messages gets the current message channel from the producer -func (p *KinesisProducer) Messages() <-chan *Message { - return p.messages -} - -// Errors gets the current number of errors on the Producer -func (p *KinesisProducer) Errors() <-chan error { - return p.errors -} - -// Send a message to the queue for POSTing -func (p *KinesisProducer) Send(msg *Message) { - // Add the terminating record indicator - p.wg.Add(1) - go func() { - p.messages <- msg - p.wg.Done() - }() -} - -// TryToSend tries to send the message, but if the channel is full it drops the message, and returns an error. -func (p *KinesisProducer) TryToSend(msg *Message) error { - // Add the terminating record indicator - select { - case p.messages <- msg: - return nil - default: - return ErrDroppedMessage - } -} - -// If our payload is larger than allowed Kinesis will write as much as -// possible and fail the rest. We can then put them back on the queue -// to re-send -func (p *KinesisProducer) sendRecords(args *awsKinesis.PutRecordsInput) { - - putResp, err := p.client.PutRecords(args) - if err != nil && conf.Debug.Verbose { - p.errors <- err - } - - // Because we do not know which of the records was successful or failed - // we need to put them all back on the queue - failedRecordCount := aws.Int64Value(putResp.FailedRecordCount) - if putResp != nil && failedRecordCount > 0 { - if conf.Debug.Verbose { - log.Printf("Failed records: %d", failedRecordCount) - } - - for idx, resp := range putResp.Records { - // Put failed records back on the queue - errorCode := aws.StringValue(resp.ErrorCode) - errorMessage := aws.StringValue(resp.ErrorMessage) - if errorCode != "" || errorMessage != "" { - p.decMsgCount() - p.errors <- errors.New(errorMessage) - p.Send(new(Message).Init(args.Records[idx].Data, aws.StringValue(args.Records[idx].PartitionKey))) - - if conf.Debug.Verbose { - log.Println("Message in failed PutRecords put back on the queue: " + string(args.Records[idx].Data)) - } - } - } - } else if putResp == nil { - // Retry posting these records as they most likely were not posted successfully - p.retryRecords(args.Records) - } - - if conf.Debug.Verbose && p.getMsgCount()%100 == 0 { - log.Println("Messages sent so far: " + strconv.FormatInt(p.getMsgCount(), 10)) - } -} - -func (p *KinesisProducer) retryRecords(records []*awsKinesis.PutRecordsRequestEntry) { - for _, record := range records { - p.Send(new(Message).Init(record.Data, aws.StringValue(record.PartitionKey))) - - if conf.Debug.Verbose { - log.Println("Message in nil send response put back on the queue: " + string(record.Data)) - } - } -} - -// Close stops queuing and producing and waits for all tasks to finish -func (p *KinesisProducer) Close() error { - if conf.Debug.Verbose { - log.Println("Producer is waiting for all tasks to finish...") - } - - p.wg.Wait() - - // Stop producing - go func() { - p.interrupts <- syscall.SIGINT - }() - - if conf.Debug.Verbose { - log.Println("Producer is shutting down.") - } - runtime.Gosched() - return nil -} - -// CloseSync closes the Producer in a syncronous manner. -func (p *KinesisProducer) CloseSync() error { - if conf.Debug.Verbose { - log.Println("Listener is waiting for all tasks to finish...") - } - var err error - // Stop consuming - select { - case p.interrupts <- syscall.SIGINT: - break - default: - if conf.Debug.Verbose { - log.Println("Already closing listener.") - } - runtime.Gosched() - return err - } - p.wg.Wait() - for p.IsProducing() { - runtime.Gosched() - } - if conf.Debug.Verbose { - log.Println("Listener is shutting down.") - } - runtime.Gosched() - return nil -} - -func (p *KinesisProducer) handleInterrupt(signal os.Signal) { - if conf.Debug.Verbose { - log.Println("Producer received interrupt signal") - } - - defer func() { - <-p.sem - }() - - p.Close() -} - -func (p *KinesisProducer) handleError(err error) { - if err != nil && conf.Debug.Verbose { - log.Println("Received error: ", err.Error()) - } - - defer func() { - <-p.sem - }() - - p.wg.Done() -} diff --git a/producer_test.go b/producer_test.go deleted file mode 100644 index 1d2bdf3..0000000 --- a/producer_test.go +++ /dev/null @@ -1,138 +0,0 @@ -package kinetic - -import ( - "encoding/binary" - "errors" - "runtime" - "syscall" - "testing" - "time" - - . "github.com/smartystreets/goconvey/convey" -) - -func TestProducerStop(t *testing.T) { - producerInterface, _ := new(KinesisProducer).Init() - producerInterface.NewEndpoint(testEndpoint, "stream-name") - producer := producerInterface.(*KinesisProducer) - CreateAndWaitForStream(producer.client, "stream-name") - producerInterface.ReInit() - - Convey("Given a running producer", t, func() { - go producer.produce() - - Convey("It should stop producing if sent an interrupt signal", func() { - producer.interrupts <- syscall.SIGINT - runtime.Gosched() - - // Wait for it to stop - for { - if !producer.IsProducing() { - break - } - } - - So(producer.IsProducing(), ShouldEqual, false) - }) - }) - - producer.Close() -} - -func TestSyncStop(t *testing.T) { - producerInterface, _ := new(KinesisProducer).Init() - producerInterface.NewEndpoint(testEndpoint, "stream-name") - producer := producerInterface.(*KinesisProducer) - CreateAndWaitForStream(producer.client, "stream-name") - producerInterface.ReInit() - - Convey("Given a running producer", t, func() { - go producer.produce() - runtime.Gosched() - Convey("It should stop producing if sent an interrupt signal", func() { - err := producer.CloseSync() - So(err, ShouldBeNil) - // Wait for it to stop - So(producer.IsProducing(), ShouldEqual, false) - }) - }) - - producer.Close() -} - -func TestProducerError(t *testing.T) { - producerInterface, _ := new(KinesisProducer).Init() - producerInterface.NewEndpoint(testEndpoint, "stream-name") - producer := producerInterface.(*KinesisProducer) - CreateAndWaitForStream(producer.client, "stream-name") - producerInterface.ReInit() - - Convey("Given a running producer", t, func() { - go producer.produce() - - Convey("It should handle errors successfully", func() { - producer.errors <- errors.New("All your base are belong to us") - // Let the error propagate - <-time.After(3 * time.Second) - So(producer.getErrCount(), ShouldEqual, 1) - So(producer.IsProducing(), ShouldEqual, true) - }) - }) - - producer.Close() -} - -func TestProducerMessage(t *testing.T) { - listener, _ := new(Listener).InitC("your-stream", "0", "LATEST", "accesskey", "secretkey", "us-east-1", 4) - producer, _ := new(KinesisProducer).InitC("your-stream", "0", "LATEST", "accesskey", "secretkey", "us-east-1", 4) - - listener.NewEndpoint(testEndpoint, "your-stream") - producer.NewEndpoint(testEndpoint, "your-stream") - - CreateAndWaitForStream(producer.(*KinesisProducer).client, "your-stream") - listener.ReInit() - producer.ReInit() - - for _, c := range cases { - Convey("Given a valid message", t, func() { - producer.Send(new(Message).Init(c.message, "test")) - - Convey("It should be passed on the queue without error", func() { - msg, err := listener.Retrieve() - if err != nil { - t.Fatalf(err.Error()) - } - - So(string(msg.Value()), ShouldResemble, string(c.message)) - }) - }) - } - - listener.Close() - producer.Close() -} - -func TestProducerTryToSend(t *testing.T) { - producer, _ := new(KinesisProducer).InitC("your-stream", "0", "LATEST", "accesskey", "secretkey", "us-east-1", 4) - producer.NewEndpoint(testEndpoint, "your-stream") - CreateAndWaitForStream(producer.(*KinesisProducer).client, "stream-name") - producer.ReInit() - - closeError := producer.CloseSync() // This is to make the test deterministic. It stops producer from sending messages. - runtime.Gosched() - var totDropped int - for i := 0; i < 5000; i++ { - b := make([]byte, 2) - binary.LittleEndian.PutUint16(b, uint16(i)) - if err := producer.TryToSend(new(Message).Init(b, "foo")); nil != err { - totDropped++ - } - } - Convey("Given a producer", t, func() { - Convey("TryToSend should drop messages when the queue is full", func() { - So(closeError, ShouldBeNil) - So(totDropped, ShouldEqual, 1000) - So(len(producer.(*KinesisProducer).messages), ShouldEqual, 4000) - }) - }) -} diff --git a/utils.go b/utils/readcloser.go similarity index 95% rename from utils.go rename to utils/readcloser.go index 99fe9c2..c15d89c 100644 --- a/utils.go +++ b/utils/readcloser.go @@ -1,4 +1,4 @@ -package kinetic +package utils import ( "io" From f94ded4231eaa41fed93fa41f6be2e2559c90c24 Mon Sep 17 00:00:00 2001 From: Tommy Wang Date: Wed, 29 Mar 2017 08:31:10 -0500 Subject: [PATCH 003/109] go vet & golint fixes --- config.go | 17 ++++++++- config_test.go | 2 +- kinetic.go | 21 +++++++++-- kinetic_test.go | 6 ++++ listener/config.go | 26 +++++++++++++- listener/config_test.go | 4 ++- listener/listener.go | 74 +++++++++++++++++++++++++++------------ listener/listener_test.go | 17 +++++---- listener/stats.go | 66 ++++++++++++++++++++++++++++------ logging/logging.go | 4 +++ utils/readcloser.go | 13 +++++-- 11 files changed, 203 insertions(+), 47 deletions(-) diff --git a/config.go b/config.go index e96bc38..70271a5 100644 --- a/config.go +++ b/config.go @@ -9,11 +9,13 @@ import ( "github.com/aws/aws-sdk-go/aws/session" ) +// Config is used to configure a Kinetic instance type Config struct { awsConfig *aws.Config *kineticConfig } +// NewConfig creates a new instance of Config func NewConfig() *Config { return &Config{ awsConfig: aws.NewConfig().WithHTTPClient(&http.Client{ @@ -25,6 +27,9 @@ func NewConfig() *Config { } } +// WithCredentials configures AWS credentials. Leave blank to use environment +// variables, IAM roles or ~/.aws/credentials. See aws-sdk-go docs for more +// details. func (c *Config) WithCredentials(accessKey, secretKey, securityToken string) *Config { c.awsConfig.WithCredentials( credentials.NewStaticCredentials(accessKey, secretKey, securityToken), @@ -32,34 +37,44 @@ func (c *Config) WithCredentials(accessKey, secretKey, securityToken string) *Co return c } +// WithRegion configures the AWS region. Leave blank to use environment +// variables. See aws-sdk-go for more details. func (c *Config) WithRegion(region string) *Config { c.awsConfig.WithRegion(region) return c } +// WithEndpoint sets the endpoint to be used by aws-sdk-go. func (c *Config) WithEndpoint(endpoint string) *Config { c.awsConfig.WithEndpoint(endpoint) return c } +// WithLogger configures the logger for Kinetic and the aws-sdk-go. func (c *Config) WithLogger(logger aws.Logger) *Config { c.awsConfig.WithLogger(logger) return c } +// WithLogLevel configures the log levels for Kinetic and the aws-sdk-go. Note +// that log levels for the SDK can be found in aws-sdk-go/aws package. Kinetic +// log levels are found in the kinetic/logging package. func (c *Config) WithLogLevel(logLevel aws.LogLevelType) *Config { c.awsConfig.WithLogLevel(logLevel & 0xffff) c.logLevel = logLevel & 0xffff0000 return c } -func (c *Config) WithHttpClientTimeout(timeout time.Duration) *Config { +// WithHTTPClientTimeout configures the HTTP timeout for the SDK. +func (c *Config) WithHTTPClientTimeout(timeout time.Duration) *Config { c.awsConfig.WithHTTPClient(&http.Client{ Timeout: timeout, }) return c } +// GetSession creates an instance of the session.Session to be used when creating service +// clients in aws-sdk-go. func (c *Config) GetSession() (*session.Session, error) { return session.NewSession(c.awsConfig) } diff --git a/config_test.go b/config_test.go index da8d622..99e2483 100644 --- a/config_test.go +++ b/config_test.go @@ -94,7 +94,7 @@ func TestNewConfig(t *testing.T) { }) Convey("check that we can set the http.Client Timeout", func() { - config = config.WithHttpClientTimeout(10 * time.Minute) + config = config.WithHTTPClientTimeout(10 * time.Minute) So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 10*time.Minute) }) }) diff --git a/kinetic.go b/kinetic.go index aa88af1..f84069a 100644 --- a/kinetic.go +++ b/kinetic.go @@ -18,15 +18,22 @@ import ( ) var ( - // GetShards errors + // ErrNilDescribeStreamResponse is an error returned by GetShards when + // the DescribeStream request returns a nil response ErrNilDescribeStreamResponse = errors.New("DescribeStream returned a nil response") - ErrNilStreamDescription = errors.New("DescribeStream returned a nil StreamDescription") + + // ErrNilStreamDescription is an error returned by GetShards when the + // DescribeStream request returns a response with a nil + // StreamDescription + ErrNilStreamDescription = errors.New("DescribeStream returned a nil StreamDescription") ) type kineticConfig struct { logLevel aws.LogLevelType } +// Kinetic represents a kinesis and firehose client and provides some utility +// methods for interacting with the AWS services. type Kinetic struct { *kineticConfig session *session.Session @@ -36,6 +43,7 @@ type Kinetic struct { clientMu sync.Mutex } +// New creates a new instance of Kientic. func New(config *Config) (*Kinetic, error) { session, err := config.GetSession() if err != nil { @@ -47,6 +55,7 @@ func New(config *Config) (*Kinetic, error) { }, nil } +// Log logs a message if LogDebug is set. func (k *Kinetic) Log(args ...interface{}) { if k.logLevel.Matches(logging.LogDebug) { k.session.Config.Logger.Log(args...) @@ -61,6 +70,7 @@ func (k *Kinetic) ensureKinesisClient() { } } +// CreateStream creates a new Kinesis stream. func (k *Kinetic) CreateStream(stream string, shards int) error { k.ensureKinesisClient() _, err := k.kclient.CreateStream(&kinesis.CreateStreamInput{ @@ -73,6 +83,8 @@ func (k *Kinetic) CreateStream(stream string, shards int) error { return err } +// WaitUntilStreamExists is meant to be used after CreateStream to wait until a +// Kinesis stream is ACTIVE. func (k *Kinetic) WaitUntilStreamExists(ctx context.Context, stream string, opts ...request.WaiterOption) error { k.ensureKinesisClient() return k.kclient.WaitUntilStreamExistsWithContext(ctx, &kinesis.DescribeStreamInput{ @@ -80,6 +92,7 @@ func (k *Kinetic) WaitUntilStreamExists(ctx context.Context, stream string, opts }, opts...) } +// DeleteStream deletes an existing Kinesis stream. func (k *Kinetic) DeleteStream(stream string) error { k.ensureKinesisClient() _, err := k.kclient.DeleteStream(&kinesis.DeleteStreamInput{ @@ -91,6 +104,8 @@ func (k *Kinetic) DeleteStream(stream string) error { return err } +// WaitUntilStreamDeleted is meant to be used after DeleteStream to wait until a +// Kinesis stream no longer exists. func (k *Kinetic) WaitUntilStreamDeleted(ctx context.Context, stream string, opts ...request.WaiterOption) error { k.ensureKinesisClient() w := request.Waiter{ @@ -118,6 +133,7 @@ func (k *Kinetic) WaitUntilStreamDeleted(ctx context.Context, stream string, opt return w.WaitWithContext(ctx) } +// GetShards returns a list of the shards in a Kinesis stream. func (k *Kinetic) GetShards(stream string) ([]string, error) { k.ensureKinesisClient() resp, err := k.kclient.DescribeStream(&kinesis.DescribeStreamInput{ @@ -142,6 +158,7 @@ func (k *Kinetic) GetShards(stream string) ([]string, error) { return shards, nil } +// GetSession returns the aws-sdk-go session.Session object. func (k *Kinetic) GetSession() *session.Session { return k.session } diff --git a/kinetic_test.go b/kinetic_test.go index dd0b4e3..befb61b 100644 --- a/kinetic_test.go +++ b/kinetic_test.go @@ -15,11 +15,17 @@ import ( func TestKinetic(t *testing.T) { Convey("given a kinetic object", t, func() { k, err := New(NewConfig(). + WithCredentials("some-access-key", "some-secret-key", "some-security-token"). WithRegion("some-region"). WithEndpoint("http://127.0.0.1:4567")) So(k, ShouldNotBeNil) So(err, ShouldBeNil) + Convey("check that we can get the session", func() { + sess := k.GetSession() + So(sess, ShouldNotBeNil) + }) + stream := "some-kinetic-stream" Convey("check that calling ensureKinesisClient twice doesn't overwrite existing client", func() { diff --git a/listener/config.go b/listener/config.go index 66d4d0a..a309b12 100644 --- a/listener/config.go +++ b/listener/config.go @@ -12,11 +12,13 @@ import ( "github.com/rewardStyle/kinetic/logging" ) +// Config is used to configure a Listener instance type Config struct { awsConfig *aws.Config *listenerConfig } +// NewConfig creates a new instance of Config func NewConfig(stream, shard string) *Config { return &Config{ awsConfig: aws.NewConfig().WithHTTPClient(&http.Client{ @@ -35,6 +37,9 @@ func NewConfig(stream, shard string) *Config { } } +// WithCredentials configures AWS credentials. Leave blank to use environment +// variables, IAM roles or ~/.aws/credentials. See aws-sdk-go docs for more +// details. func (c *Config) WithCredentials(accessKey, secretKey, securityToken string) *Config { c.awsConfig.WithCredentials( credentials.NewStaticCredentials(accessKey, secretKey, securityToken), @@ -42,63 +47,82 @@ func (c *Config) WithCredentials(accessKey, secretKey, securityToken string) *Co return c } +// WithRegion configures the AWS region. Leave blank to use environment +// variables. See aws-sdk-go for more details. func (c *Config) WithRegion(region string) *Config { c.awsConfig.WithRegion(region) return c } +// WithEndpoint sets the endpoint to be used by aws-sdk-go. func (c *Config) WithEndpoint(endpoint string) *Config { c.awsConfig.WithEndpoint(endpoint) return c } +// WithLogger configures the logger for Kinetic and the aws-sdk-go. func (c *Config) WithLogger(logger aws.Logger) *Config { c.awsConfig.WithLogger(logger) return c } +// WithLogLevel configures the log levels for Kinetic and the aws-sdk-go. Note +// that log levels for the SDK can be found in aws-sdk-go/aws package. Kinetic +// log levels are found in the kinetic/logging package. func (c *Config) WithLogLevel(logLevel aws.LogLevelType) *Config { c.awsConfig.WithLogLevel(logLevel & 0xffff) c.logLevel = logLevel & 0xffff0000 return c } -func (c *Config) WithHttpClientTimeout(timeout time.Duration) *Config { +// WithHTTPClientTimeout configures the HTTP timeout for the SDK. +func (c *Config) WithHTTPClientTimeout(timeout time.Duration) *Config { c.awsConfig.WithHTTPClient(&http.Client{ Timeout: timeout, }) return c } +// GetSession creates an instance of the session.Session to be used when creating service +// clients in aws-sdk-go. func (c *Config) GetSession() (*session.Session, error) { return session.NewSession(c.awsConfig) } +// FromKinetic configures the session from Kinetic. func (c *Config) FromKinetic(k *kinetic.Kinetic) *Config { c.awsConfig = k.GetSession().Config return c } +// WithBatchSize configures the batch size of the GetRecords call. func (c *Config) WithBatchSize(batchSize int) *Config { c.batchSize = batchSize return c } +// WithConcurrency controls the number of goroutines the Listener will spawn to +// process messages. func (c *Config) WithConcurrency(concurrency int) *Config { c.concurrency = concurrency return c } +// WithInitialShardIterator configures the settings used to retrieve initial +// shard iterator via the GetShardIterator call. func (c *Config) WithInitialShardIterator(shardIterator *ShardIterator) *Config { c.shardIterator = shardIterator return c } +// WithGetRecordsReadTimeout configures the time to wait for each successive +// Read operation on the GetRecords response payload. func (c *Config) WithGetRecordsReadTimeout(timouet time.Duration) *Config { c.getRecordsReadTimeout = timouet return c } +// WithStatsListener configures a listener to handle metrics. func (c *Config) WithStatsListener(stats StatsListener) *Config { c.stats = stats return c diff --git a/listener/config_test.go b/listener/config_test.go index 5305a75..76b72b0 100644 --- a/listener/config_test.go +++ b/listener/config_test.go @@ -24,6 +24,8 @@ func (l *DebugStatsListener) AddGetRecordsCalled(int) {} func (l *DebugStatsListener) AddProvisionedThroughputExceeded(int) {} func (l *DebugStatsListener) AddGetRecordsTimeout(int) {} func (l *DebugStatsListener) AddGetRecordsReadTimeout(int) {} +func (l *DebugStatsListener) AddProcessedTime(time.Duration) {} +func (l *DebugStatsListener) AddGetRecordsTime(time.Duration) {} func (l *DebugStatsListener) AddGetRecordsReadResponseTime(time.Duration) {} func (l *DebugStatsListener) AddGetRecordsUnmarshalTime(time.Duration) {} @@ -117,7 +119,7 @@ func TestNewConfig(t *testing.T) { }) Convey("check that we can set the http.Client Timeout", func() { - config = config.WithHttpClientTimeout(10 * time.Minute) + config = config.WithHTTPClientTimeout(10 * time.Minute) So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 10*time.Minute) }) diff --git a/listener/listener.go b/listener/listener.go index baa2475..ec10ca3 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -22,65 +22,89 @@ import ( ) var ( - // errors tha can occur in Retrieve / RetrieveFn (and their WithContext variants) + // ErrAlreadyConsuming is returned when attempting to consume when the + // Listener is already consuming. May be returned by + // Retrieve/RetrieveFn. ErrAlreadyConsuming = errors.New("Listener already consuming. Only one Listen, Retrieve, or RetrieveFn may be active at a time") - // errors that can occur in GetShards - ErrNilDescribeStreamResponse = errors.New("DescribeStream returned a nil response") - ErrNilStreamDescription = errors.New("DescribeStream returned a nil StreamDescription") + // ErrEmptySequenceNumber is returned when attempting to set an empty + // sequence number. + ErrEmptySequenceNumber = errors.New("Attempted to set sequence number with empty value") - // errors that can occur in SetShard - ErrCannotSetShard = errors.New("Cannot set shard while consuming") + // ErrEmptyShardIterator is returned when attempting to set an empty + // sequence number. + ErrEmptyShardIterator = errors.New("Attempted to set shard iterator with empty value") - // errors that can occur in fetchBatch - ErrEmptySequenceNumber = errors.New("Attempted to set sequence number with empty value") - ErrEmptyShardIterator = errors.New("Attempted to set shard iterator with empty value") + // ErrNilGetShardIteratorResponse is returned when the GetShardIterator + // call returns a nil response. ErrNilGetShardIteratorResponse = errors.New("GetShardIteratore returned a nil response") - ErrNilShardIterator = errors.New("GetShardIterator returned a nil ShardIterator") - ErrNilGetRecordsResponse = errors.New("GetRecords returned an nil response") - ErrTimeoutReadResponseBody = errors.New("Timeout while reading response body") + // ErrNilShardIterator is returned when the GetShardIterator call + // returns a nil shard iterator. + ErrNilShardIterator = errors.New("GetShardIterator returned a nil ShardIterator") + + // ErrNilGetRecordsResponse is returned when the GetRecords calls + // returns a nil response. + ErrNilGetRecordsResponse = errors.New("GetRecords returned an nil response") + + // ErrTimeoutReadResponseBody is returned when a timeout occurs while + // reading the GetRecords response body. + ErrTimeoutReadResponseBody = errors.New("Timeout while reading response body") + + // ErrPipeOfDeath returns when the pipe of death is closed. ErrPipeOfDeath = errors.New("Received pipe of death") ) +// Empty is used a as a dummy type for counting semaphore channels. type Empty struct{} +// MessageFn defines the signature of a message handler used by Listen and +// RetrieveFn. type MessageFn func([]byte, *sync.WaitGroup) +// ShardIterator represents the settings used to retrieve a shard iterator from +// the GetShardIterator API. type ShardIterator struct { shardIteratorType string sequenceNumber string timestamp time.Time } +// NewShardIterator creates a new ShardIterator. The default shard iterator +// type is TRIM_HORIZON. func NewShardIterator() *ShardIterator { return &ShardIterator{ shardIteratorType: "TRIM_HORIZON", } } +// TrimHorizon sets the shard iterator to TRIM_HORIZON. func (it *ShardIterator) TrimHorizon() *ShardIterator { it.shardIteratorType = "TRIM_HORIZON" return it } +// Latest sets the shard iterator to LATEST. func (it *ShardIterator) Latest() *ShardIterator { it.shardIteratorType = "LATEST" return it } +// AtSequenceNumber sets the shard iterator to AT_SEQUENCE_NUMBER. func (it *ShardIterator) AtSequenceNumber(sequenceNumber string) *ShardIterator { it.shardIteratorType = "AT_SEQUENCE_NUMBER" it.sequenceNumber = sequenceNumber return it } +// AfterSequenceNumber sets the shard iterator to AFTER_SEQUENCE_NUMBER. func (it *ShardIterator) AfterSequenceNumber(sequenceNumber string) *ShardIterator { it.shardIteratorType = "AFTER_SEQUENCE_NUMBER" it.sequenceNumber = sequenceNumber return it } +// AtTimestamp sets the shard iterator to AT_TIMESTAMP. func (it *ShardIterator) AtTimestamp(timestamp time.Time) *ShardIterator { it.shardIteratorType = "AT_TIMESTAMP" it.timestamp = timestamp @@ -114,6 +138,7 @@ type listenerConfig struct { logLevel aws.LogLevelType } +// Listener polls the Kinesis stream for messages. type Listener struct { *listenerConfig @@ -148,7 +173,7 @@ func NewListener(config *Config) (*Listener, error) { }, nil } -// Logs a debug message using the AWS SDK logger. +// Log a debug message using the AWS SDK logger. func (l *Listener) Log(args ...interface{}) { if l.session.Config.LogLevel.Matches(logging.LogDebug) { l.session.Config.Logger.Log(args...) @@ -315,12 +340,12 @@ func (l *Listener) fetchBatch(size int) (int, error) { startReadTime = time.Now() timer := time.NewTimer(l.getRecordsReadTimeout) - r.HTTPResponse.Body = &utils.TimeoutReadCloser{ + r.HTTPResponse.Body = &utils.ReadCloserWrapper{ ReadCloser: r.HTTPResponse.Body, OnReadFn: func(stream io.ReadCloser, b []byte) (n int, err error) { // The OnReadFn will be called each time // ioutil.ReadAll calls Read on the - // TimeoutReadCloser. + // ReadCloserWrapper. // First, we set up a struct that to hold the // results of the Read() call that can go @@ -387,6 +412,7 @@ func (l *Listener) fetchBatch(size int) (int, error) { l.Log("Error getting records:", err) return 0, err } + l.stats.AddGetRecordsTime(time.Since(start)) // Process Records l.Log(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), l.shard, time.Since(start))) @@ -398,7 +424,7 @@ func (l *Listener) fetchBatch(size int) (int, error) { for _, record := range resp.Records { if record != nil { delivered++ - l.messages <- &message.Message{*record} + l.messages <- &message.Message{Record: *record} l.stats.AddConsumedSample(1) } if record.SequenceNumber != nil { @@ -486,7 +512,7 @@ func (l *Listener) IsConsuming() bool { return l.consuming } -// Retrieve waits for a message from the stream and return the value. +// RetrieveWithContext waits for a message from the stream and return the value. // Cancellation supported through contexts. func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, error) { if !l.blockConsumers() { @@ -518,8 +544,8 @@ func (l *Listener) Retrieve() (*message.Message, error) { return l.RetrieveWithContext(context.TODO()) } -// RetrieveFn retrieves a message from the stream and dispatches it to the -// supplied function. RetrieveFn will wait until the function completes. +// RetrieveFnWithContext retrieves a message from the stream and dispatches it +// to the supplied function. RetrieveFn will wait until the function completes. // Cancellation supported through context. func (l *Listener) RetrieveFnWithContext(ctx context.Context, fn MessageFn) error { msg, err := l.RetrieveWithContext(ctx) @@ -528,8 +554,10 @@ func (l *Listener) RetrieveFnWithContext(ctx context.Context, fn MessageFn) erro } var wg sync.WaitGroup wg.Add(1) + start := time.Now() go fn(msg.Value(), &wg) wg.Wait() + l.stats.AddProcessedTime(time.Since(start)) l.stats.AddProcessedSample(1) return nil } @@ -613,9 +641,9 @@ func (l *Listener) consume(ctx context.Context) { }() } -// Listen listens and delivers message to the supplied function. Upon -// cancellation, Listen will stop the consumer loop and wait until the messages -// channel is closed and all messages are delivered. +// ListenWithContext listens and delivers message to the supplied function. +// Upon cancellation, Listen will stop the consumer loop and wait until the +// messages channel is closed and all messages are delivered. func (l *Listener) ListenWithContext(ctx context.Context, fn MessageFn) { l.consume(ctx) var wg sync.WaitGroup @@ -638,8 +666,10 @@ stop: }() var fnWg sync.WaitGroup fnWg.Add(1) + start := time.Now() fn(msg.Value(), &fnWg) fnWg.Wait() + l.stats.AddProcessedTime(time.Since(start)) l.stats.AddProcessedSample(1) wg.Done() }() diff --git a/listener/listener_test.go b/listener/listener_test.go index a874d3d..7672e93 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -83,6 +83,7 @@ func TestShardIterator(t *testing.T) { func TestListener(t *testing.T) { Convey("given a listener", t, func() { k, err := kinetic.New(kinetic.NewConfig(). + WithCredentials("some-access-key", "some-secret-key", "some-security-token"). WithRegion("some-region"). WithEndpoint("http://127.0.0.1:4567")) @@ -163,7 +164,8 @@ func TestListener(t *testing.T) { Convey("check that we can use a context to cancel the retrieve", func() { start := time.Now() - ctx, _ := context.WithTimeout(context.TODO(), 1*time.Second) + ctx, cancel := context.WithTimeout(context.TODO(), 1*time.Second) + defer cancel() _, err := l.RetrieveWithContext(ctx) elapsed := time.Since(start) Printf("(it blocked %f seconds)\n", elapsed.Seconds()) @@ -174,7 +176,8 @@ func TestListener(t *testing.T) { Convey("check that we can use a context to cancel the retrieve (again)", func() { start := time.Now() - ctx, _ := context.WithTimeout(context.TODO(), 10*time.Millisecond) + ctx, cancel := context.WithTimeout(context.TODO(), 10*time.Millisecond) + defer cancel() _, err := l.RetrieveWithContext(ctx) elapsed := time.Since(start) Printf("(it blocked %f seconds)\n", elapsed.Seconds()) @@ -184,7 +187,8 @@ func TestListener(t *testing.T) { }) Convey("check that retrieve still works with a canceller if a message comes before the deadline", func(c C) { - ctx, _ := context.WithTimeout(context.TODO(), 5*time.Second) + ctx, cancel := context.WithTimeout(context.TODO(), 5*time.Second) + defer cancel() data := "goodbye" go func() { @@ -201,7 +205,8 @@ func TestListener(t *testing.T) { var wg sync.WaitGroup wg.Add(1) go func() { - ctx, _ := context.WithTimeout(context.TODO(), 1*time.Second) + ctx, cancel := context.WithTimeout(context.TODO(), 1*time.Second) + defer cancel() _, err := l.RetrieveWithContext(ctx) c.So(err, ShouldNotBeNil) c.So(err, ShouldHaveSameTypeAs, context.DeadlineExceeded) @@ -245,7 +250,7 @@ func TestListener(t *testing.T) { Convey("check that listen can deliver messages to fn", func(c C) { planets := []string{"mercury", "venus", "earth", "mars", "jupiter", "saturn", "neptune", "uranus"} - var count int64 = 0 + var count int64 var wg sync.WaitGroup wg.Add(1) go func() { @@ -269,7 +274,7 @@ func TestListener(t *testing.T) { _, err := putRecord(l, []byte(fmt.Sprintf("%d", i))) So(err, ShouldBeNil) } - var count int64 = 0 + var count int64 var wg sync.WaitGroup wg.Add(1) go func() { diff --git a/listener/stats.go b/listener/stats.go index b6da386..c7eeb73 100644 --- a/listener/stats.go +++ b/listener/stats.go @@ -4,7 +4,9 @@ import ( "time" ) -// This was really built with rcrowley/go-metrics in mind. +// StatsListener allows for a listener to "listen" for various metrics produced +// by the Kinetic library. This was really built with rcrowley/go-metrics in +// mind. type StatsListener interface { // histograms AddConsumedSample(int) @@ -19,19 +21,63 @@ type StatsListener interface { AddGetRecordsReadTimeout(int) // timers + AddProcessedTime(time.Duration) + AddGetRecordsTime(time.Duration) AddGetRecordsReadResponseTime(time.Duration) AddGetRecordsUnmarshalTime(time.Duration) } +// NilStatsListener is a stats listener that ignores all metrics. type NilStatsListener struct{} -func (l *NilStatsListener) AddConsumedSample(int) {} -func (l *NilStatsListener) AddDeliveredSample(int) {} -func (l *NilStatsListener) AddProcessedSample(int) {} -func (l *NilStatsListener) AddBatchSizeSample(int) {} -func (l *NilStatsListener) AddGetRecordsCalled(int) {} -func (l *NilStatsListener) AddProvisionedThroughputExceeded(int) {} -func (l *NilStatsListener) AddGetRecordsTimeout(int) {} -func (l *NilStatsListener) AddGetRecordsReadTimeout(int) {} +// AddConsumedSample records a count of the number of messages received from AWS +// Kinesis by the listener. +func (l *NilStatsListener) AddConsumedSample(int) {} + +// AddDeliveredSample records a count of the number of messages delivered to the +// application by the listener. +func (l *NilStatsListener) AddDeliveredSample(int) {} + +// AddProcessedSample records a count of the number of messages processed by the +// application by the listener. This is based on a WaitGroup that is sent to +// the RetrieveFn and Listen functions. Retrieve does not count processed +// messages. +func (l *NilStatsListener) AddProcessedSample(int) {} + +// AddBatchSizeSample records a count of the number of messages returned by +// GetRecords in the listener. +func (l *NilStatsListener) AddBatchSizeSample(int) {} + +// AddGetRecordsCalled records the number of times the GetRecords API was called +// by the listener. +func (l *NilStatsListener) AddGetRecordsCalled(int) {} + +// AddProvisionedThroughputExceeded records the number of times the GetRecords +// API returned a ErrCodeProvisionedThroughputExceededException by the listener. +func (l *NilStatsListener) AddProvisionedThroughputExceeded(int) {} + +// AddGetRecordsTimeout records the number of times the GetRecords API timed out +// on the HTTP level. This is influenced by the WithHTTPClientTimeout +// configuration. +func (l *NilStatsListener) AddGetRecordsTimeout(int) {} + +// AddGetRecordsReadTimeout records the number of times the GetRecords API timed out +// while reading the response body. This is influenced by the +// WithGetRecordsReadTimeout configuration. +func (l *NilStatsListener) AddGetRecordsReadTimeout(int) {} + +// AddProcessedTime records the duration to process a record. See notes on +// AddProcessedSample. +func (l *NilStatsListener) AddProcessedTime(time.Duration) {} + +// AddGetRecordsTime records the duration that the GetRecords API request took. +// Only the times of successful calls are measured. +func (l *NilStatsListener) AddGetRecordsTime(time.Duration) {} + +// AddGetRecordsReadResponseTime records the duration that it took to read the +// response body of a GetRecords API request. func (l *NilStatsListener) AddGetRecordsReadResponseTime(time.Duration) {} -func (l *NilStatsListener) AddGetRecordsUnmarshalTime(time.Duration) {} + +// AddGetRecordsUnmarshalTime records the duration that it took to unmarshal the +// response body of a GetRecords API request. +func (l *NilStatsListener) AddGetRecordsUnmarshalTime(time.Duration) {} diff --git a/logging/logging.go b/logging/logging.go index 0d2631e..16c67f9 100644 --- a/logging/logging.go +++ b/logging/logging.go @@ -5,10 +5,14 @@ import ( ) const ( + // LogOff disables all logging. LogOff aws.LogLevelType = (iota * 0x1000) << 16 + + // LogDebug enables debug logging. LogDebug ) const ( + // LogPlaceholder is a placeholder debug sub-level. LogPlaceholder aws.LogLevelType = LogDebug | (1 << (16 + iota)) ) diff --git a/utils/readcloser.go b/utils/readcloser.go index c15d89c..2365f3f 100644 --- a/utils/readcloser.go +++ b/utils/readcloser.go @@ -4,17 +4,24 @@ import ( "io" ) -type TimeoutReadCloser struct { +// ReadCloserWrapper wraps a stream (like HTTPResponse.Body) and allows us to +// intercept Read() and Close() calls. Used to implement a timeout mechanism in +// the aws-sdk-go. +type ReadCloserWrapper struct { io.ReadCloser OnReadFn func(io.ReadCloser, []byte) (int, error) OnCloseFn func() } -func (r *TimeoutReadCloser) Read(b []byte) (int, error) { +// Read is called to read the wrapped stream. The supplied OnReadFn is +// responsible for making the read to the wrapped stream. +func (r *ReadCloserWrapper) Read(b []byte) (int, error) { return r.OnReadFn(r.ReadCloser, b) } -func (r *TimeoutReadCloser) Close() error { +// Close is called to close the wrapped stream. The supplied OnCloseFn is not +// responsible for closing the stream. +func (r *ReadCloserWrapper) Close() error { r.OnCloseFn() return r.ReadCloser.Close() } From 41e8fb30b6c44f11fb07d42dedda7e725add0bee Mon Sep 17 00:00:00 2001 From: Tommy Wang Date: Wed, 29 Mar 2017 08:49:16 -0500 Subject: [PATCH 004/109] gosimple fixes --- listener/listener.go | 42 +++++++++++++++++------------------------- 1 file changed, 17 insertions(+), 25 deletions(-) diff --git a/listener/listener.go b/listener/listener.go index ec10ca3..33c524b 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -648,32 +648,24 @@ func (l *Listener) ListenWithContext(ctx context.Context, fn MessageFn) { l.consume(ctx) var wg sync.WaitGroup defer wg.Wait() -stop: - for { - select { - case msg, ok := <-l.messages: - // Listen should always run until it the messages - // channel closes. - if !ok { - break stop - } - l.stats.AddDeliveredSample(1) - l.concurrencySem <- Empty{} - wg.Add(1) - go func() { - defer func() { - <-l.concurrencySem - }() - var fnWg sync.WaitGroup - fnWg.Add(1) - start := time.Now() - fn(msg.Value(), &fnWg) - fnWg.Wait() - l.stats.AddProcessedTime(time.Since(start)) - l.stats.AddProcessedSample(1) - wg.Done() + + for msg := range l.messages { + l.stats.AddDeliveredSample(1) + l.concurrencySem <- Empty{} + wg.Add(1) + go func() { + defer func() { + <-l.concurrencySem }() - } + var fnWg sync.WaitGroup + fnWg.Add(1) + start := time.Now() + fn(msg.Value(), &fnWg) + fnWg.Wait() + l.stats.AddProcessedTime(time.Since(start)) + l.stats.AddProcessedSample(1) + wg.Done() + }() } } From 9481b6a2e9e4e92e92305bba0992587eeeb68144 Mon Sep 17 00:00:00 2001 From: Tommy Wang Date: Wed, 29 Mar 2017 09:03:10 -0500 Subject: [PATCH 005/109] more go vet & go race fixes --- listener/listener.go | 4 ++-- listener/listener_test.go | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/listener/listener.go b/listener/listener.go index 33c524b..0ba0d75 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -653,7 +653,7 @@ func (l *Listener) ListenWithContext(ctx context.Context, fn MessageFn) { l.stats.AddDeliveredSample(1) l.concurrencySem <- Empty{} wg.Add(1) - go func() { + go func(msg *message.Message) { defer func() { <-l.concurrencySem }() @@ -665,7 +665,7 @@ func (l *Listener) ListenWithContext(ctx context.Context, fn MessageFn) { l.stats.AddProcessedTime(time.Since(start)) l.stats.AddProcessedSample(1) wg.Done() - }() + }(msg) } } diff --git a/listener/listener_test.go b/listener/listener_test.go index 7672e93..f2ece02 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -293,8 +293,8 @@ func TestListener(t *testing.T) { }) }() wg.Wait() - So(count, ShouldBeBetweenOrEqual, 1, 20) - Printf("(count was %d)", count) + So(atomic.LoadInt64(&count), ShouldBeBetweenOrEqual, 1, 20) + Printf("(count was %d)", atomic.LoadInt64(&count)) }) Reset(func() { From a1d14e41749534b533fbf74eea1025a868f5e536 Mon Sep 17 00:00:00 2001 From: Tommy Wang Date: Fri, 31 Mar 2017 08:53:35 -0500 Subject: [PATCH 006/109] Simplify configuration implementation * Use a simpler (more DRY) method for configuring Kinetic and Listener objects. TODO: * There should probably be a recover() inside the New functions to handle any panics inside the supplied configuration function. * Should re-evaluate whether embedding should be done by pointer. Could the caller retain a pointer to the structs and therefore gain access to internal state? Do we care? * Should re-evaluate whether much of the config should be exported or unexported. I've been told recently that exported/unexported isn't really about encapsulation, but rather more about stable vs unstable API. I'm not 100% convinced. --- config.go | 98 ++++++++++++++++++---------------- config_test.go | 26 ++++----- kinetic.go | 12 +++-- kinetic_test.go | 9 ++-- listener/config.go | 107 +++++++++++--------------------------- listener/config_test.go | 79 ++++------------------------ listener/listener.go | 18 ++++--- listener/listener_test.go | 16 +++--- 8 files changed, 134 insertions(+), 231 deletions(-) diff --git a/config.go b/config.go index 70271a5..e97a2d1 100644 --- a/config.go +++ b/config.go @@ -9,72 +9,80 @@ import ( "github.com/aws/aws-sdk-go/aws/session" ) -// Config is used to configure a Kinetic instance -type Config struct { - awsConfig *aws.Config - *kineticConfig +// AwsOptions helps configure an aws.Config and session.Session +type AwsOptions struct { + AwsConfig *aws.Config } -// NewConfig creates a new instance of Config -func NewConfig() *Config { - return &Config{ - awsConfig: aws.NewConfig().WithHTTPClient(&http.Client{ - Timeout: 5 * time.Minute, - }), - kineticConfig: &kineticConfig{ - logLevel: aws.LogOff, - }, +// DefaultAwsOptions initializes the default aws.Config struct +func DefaultAwsOptions() *AwsOptions { + return &AwsOptions{ + AwsConfig: aws.NewConfig().WithHTTPClient( + &http.Client{ + Timeout: 10 * time.Minute, + }, + ), } } -// WithCredentials configures AWS credentials. Leave blank to use environment -// variables, IAM roles or ~/.aws/credentials. See aws-sdk-go docs for more -// details. -func (c *Config) WithCredentials(accessKey, secretKey, securityToken string) *Config { - c.awsConfig.WithCredentials( - credentials.NewStaticCredentials(accessKey, secretKey, securityToken), +// SetCredentials configures AWS credentials. +func (c *AwsOptions) SetCredentials(accessKey, secretKey, sessionToken string) { + c.AwsConfig.WithCredentials( + credentials.NewStaticCredentials(accessKey, secretKey, sessionToken), ) - return c } -// WithRegion configures the AWS region. Leave blank to use environment -// variables. See aws-sdk-go for more details. -func (c *Config) WithRegion(region string) *Config { - c.awsConfig.WithRegion(region) - return c +// SetRegion configures the AWS region. +func (c *AwsOptions) SetRegion(region string) { + c.AwsConfig.WithRegion(region) } -// WithEndpoint sets the endpoint to be used by aws-sdk-go. -func (c *Config) WithEndpoint(endpoint string) *Config { - c.awsConfig.WithEndpoint(endpoint) - return c +// SetEndpoint sets the endpoint to be used by aws-sdk-go. +func (c *AwsOptions) SetEndpoint(endpoint string) { + c.AwsConfig.WithEndpoint(endpoint) } -// WithLogger configures the logger for Kinetic and the aws-sdk-go. -func (c *Config) WithLogger(logger aws.Logger) *Config { - c.awsConfig.WithLogger(logger) - return c +// SetLogger configures the logger for Kinetic and the aws-sdk-go. +func (c *AwsOptions) SetLogger(logger aws.Logger) { + c.AwsConfig.WithLogger(logger) } -// WithLogLevel configures the log levels for Kinetic and the aws-sdk-go. Note -// that log levels for the SDK can be found in aws-sdk-go/aws package. Kinetic -// log levels are found in the kinetic/logging package. -func (c *Config) WithLogLevel(logLevel aws.LogLevelType) *Config { - c.awsConfig.WithLogLevel(logLevel & 0xffff) - c.logLevel = logLevel & 0xffff0000 - return c +// SetLogLevel configures the log levels for the SDK. +func (c *AwsOptions) SetLogLevel(logLevel aws.LogLevelType) { + c.AwsConfig.WithLogLevel(logLevel & 0xffff) } -// WithHTTPClientTimeout configures the HTTP timeout for the SDK. -func (c *Config) WithHTTPClientTimeout(timeout time.Duration) *Config { - c.awsConfig.WithHTTPClient(&http.Client{ +// SetHTTPClientTimeout configures the HTTP timeout for the SDK. +func (c *AwsOptions) SetHTTPClientTimeout(timeout time.Duration) { + c.AwsConfig.WithHTTPClient(&http.Client{ Timeout: timeout, }) - return c +} + +// Config is used to configure a Kinetic instance +type Config struct { + *AwsOptions + *kineticOptions +} + +// NewConfig creates a new instance of Config +func NewConfig() *Config { + return &Config{ + AwsOptions: DefaultAwsOptions(), + kineticOptions: &kineticOptions{ + logLevel: aws.LogOff, + }, + } +} + +// SetLogLevel configures both the SDK and Kinetic log levels. +func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { + c.AwsOptions.SetLogLevel(logLevel) + c.kineticOptions.logLevel = logLevel & 0xffff0000 } // GetSession creates an instance of the session.Session to be used when creating service // clients in aws-sdk-go. func (c *Config) GetSession() (*session.Session, error) { - return session.NewSession(c.awsConfig) + return session.NewSession(c.AwsConfig) } diff --git a/config_test.go b/config_test.go index 99e2483..c95abba 100644 --- a/config_test.go +++ b/config_test.go @@ -25,7 +25,9 @@ func TestNewConfig(t *testing.T) { config := NewConfig() Convey("check the default values for its non-zero config", func() { - So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 5*time.Minute) + So(config.AwsConfig.HTTPClient.Timeout, ShouldEqual, 10*time.Minute) + So(config.AwsConfig.LogLevel.Value(), ShouldEqual, aws.LogOff) + So(config.logLevel.Value(), ShouldEqual, logging.LogOff) }) Convey("check that we can retrieve an aws.Session from it ", func() { @@ -33,7 +35,7 @@ func TestNewConfig(t *testing.T) { }) Convey("check that we can set credentials", func() { - config = config.WithCredentials("access-key", "secret-key", "security-token") + config.SetCredentials("access-key", "secret-key", "security-token") sess := getSession(config) creds, err := sess.Config.Credentials.Get() So(err, ShouldBeNil) @@ -43,13 +45,13 @@ func TestNewConfig(t *testing.T) { }) Convey("check that we can set the region", func() { - config = config.WithRegion("my-region") + config.SetRegion("my-region") sess := getSession(config) So(aws.StringValue(sess.Config.Region), ShouldEqual, "my-region") }) Convey("check that we can set the endpoint", func() { - config = config.WithEndpoint("my-endpoint") + config.SetEndpoint("my-endpoint") sess := getSession(config) So(aws.StringValue(sess.Config.Endpoint), ShouldEqual, "my-endpoint") }) @@ -59,7 +61,7 @@ func TestNewConfig(t *testing.T) { loggerFn := func(args ...interface{}) { logs = append(logs, fmt.Sprint(args...)) } - config = config.WithLogger(aws.LoggerFunc(loggerFn)) + config.SetLogger(aws.LoggerFunc(loggerFn)) sess := getSession(config) Convey("check that basic logging should work", func() { @@ -76,17 +78,9 @@ func TestNewConfig(t *testing.T) { }) }) - Convey("check that the default log level is off for both the sdk and kinetic", func() { - sess := getSession(config) - So(sess.Config.LogLevel.Value(), ShouldEqual, aws.LogOff) - So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeFalse) - So(config.logLevel.Value(), ShouldEqual, aws.LogOff) - So(config.logLevel.AtLeast(aws.LogDebug), ShouldBeFalse) - }) - Convey("check that we can set both the sdk and kinetic log level", func() { ll := aws.LogDebug | aws.LogDebugWithSigning | logging.LogDebug - config = config.WithLogLevel(ll) + config.SetLogLevel(ll) sess := getSession(config) So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) @@ -94,8 +88,8 @@ func TestNewConfig(t *testing.T) { }) Convey("check that we can set the http.Client Timeout", func() { - config = config.WithHTTPClientTimeout(10 * time.Minute) - So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 10*time.Minute) + config.SetHTTPClientTimeout(10 * time.Minute) + So(config.AwsConfig.HTTPClient.Timeout, ShouldEqual, 10*time.Minute) }) }) } diff --git a/kinetic.go b/kinetic.go index f84069a..65fc899 100644 --- a/kinetic.go +++ b/kinetic.go @@ -28,14 +28,14 @@ var ( ErrNilStreamDescription = errors.New("DescribeStream returned a nil StreamDescription") ) -type kineticConfig struct { +type kineticOptions struct { logLevel aws.LogLevelType } // Kinetic represents a kinesis and firehose client and provides some utility // methods for interacting with the AWS services. type Kinetic struct { - *kineticConfig + *kineticOptions session *session.Session fclient firehoseiface.FirehoseAPI @@ -44,14 +44,16 @@ type Kinetic struct { } // New creates a new instance of Kientic. -func New(config *Config) (*Kinetic, error) { +func New(fn func(*Config)) (*Kinetic, error) { + config := NewConfig() + fn(config) session, err := config.GetSession() if err != nil { return nil, err } return &Kinetic{ - kineticConfig: config.kineticConfig, - session: session, + kineticOptions: config.kineticOptions, + session: session, }, nil } diff --git a/kinetic_test.go b/kinetic_test.go index befb61b..4f33a16 100644 --- a/kinetic_test.go +++ b/kinetic_test.go @@ -14,10 +14,11 @@ import ( func TestKinetic(t *testing.T) { Convey("given a kinetic object", t, func() { - k, err := New(NewConfig(). - WithCredentials("some-access-key", "some-secret-key", "some-security-token"). - WithRegion("some-region"). - WithEndpoint("http://127.0.0.1:4567")) + k, err := New(func(c *Config) { + c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") + c.SetRegion("some-region") + c.SetEndpoint("http://127.0.0.1:4567") + }) So(k, ShouldNotBeNil) So(err, ShouldBeNil) diff --git a/listener/config.go b/listener/config.go index a309b12..4e7b25d 100644 --- a/listener/config.go +++ b/listener/config.go @@ -1,11 +1,9 @@ package listener import ( - "net/http" "time" "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/rewardStyle/kinetic" @@ -14,17 +12,15 @@ import ( // Config is used to configure a Listener instance type Config struct { - awsConfig *aws.Config - *listenerConfig + *kinetic.AwsOptions + *listenerOptions } // NewConfig creates a new instance of Config func NewConfig(stream, shard string) *Config { return &Config{ - awsConfig: aws.NewConfig().WithHTTPClient(&http.Client{ - Timeout: 5 * time.Minute, - }), - listenerConfig: &listenerConfig{ + AwsOptions: kinetic.DefaultAwsOptions(), + listenerOptions: &listenerOptions{ stream: stream, shard: shard, batchSize: 10000, @@ -37,93 +33,48 @@ func NewConfig(stream, shard string) *Config { } } -// WithCredentials configures AWS credentials. Leave blank to use environment -// variables, IAM roles or ~/.aws/credentials. See aws-sdk-go docs for more -// details. -func (c *Config) WithCredentials(accessKey, secretKey, securityToken string) *Config { - c.awsConfig.WithCredentials( - credentials.NewStaticCredentials(accessKey, secretKey, securityToken), - ) - return c -} - -// WithRegion configures the AWS region. Leave blank to use environment -// variables. See aws-sdk-go for more details. -func (c *Config) WithRegion(region string) *Config { - c.awsConfig.WithRegion(region) - return c -} - -// WithEndpoint sets the endpoint to be used by aws-sdk-go. -func (c *Config) WithEndpoint(endpoint string) *Config { - c.awsConfig.WithEndpoint(endpoint) - return c -} - -// WithLogger configures the logger for Kinetic and the aws-sdk-go. -func (c *Config) WithLogger(logger aws.Logger) *Config { - c.awsConfig.WithLogger(logger) - return c +// SetLogLevel configures both the SDK and Kinetic log levels. +func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { + c.AwsOptions.SetLogLevel(logLevel) + c.listenerOptions.logLevel = logLevel & 0xffff0000 } -// WithLogLevel configures the log levels for Kinetic and the aws-sdk-go. Note -// that log levels for the SDK can be found in aws-sdk-go/aws package. Kinetic -// log levels are found in the kinetic/logging package. -func (c *Config) WithLogLevel(logLevel aws.LogLevelType) *Config { - c.awsConfig.WithLogLevel(logLevel & 0xffff) - c.logLevel = logLevel & 0xffff0000 - return c -} - -// WithHTTPClientTimeout configures the HTTP timeout for the SDK. -func (c *Config) WithHTTPClientTimeout(timeout time.Duration) *Config { - c.awsConfig.WithHTTPClient(&http.Client{ - Timeout: timeout, - }) - return c -} - -// GetSession creates an instance of the session.Session to be used when creating service -// clients in aws-sdk-go. -func (c *Config) GetSession() (*session.Session, error) { - return session.NewSession(c.awsConfig) -} - -// FromKinetic configures the session from Kinetic. -func (c *Config) FromKinetic(k *kinetic.Kinetic) *Config { - c.awsConfig = k.GetSession().Config - return c -} - -// WithBatchSize configures the batch size of the GetRecords call. -func (c *Config) WithBatchSize(batchSize int) *Config { +// SetBatchSize configures the batch size of the GetRecords call. +func (c *Config) SetBatchSize(batchSize int) { c.batchSize = batchSize - return c } -// WithConcurrency controls the number of goroutines the Listener will spawn to +// SetConcurrency controls the number of goroutines the Listener will spawn to // process messages. -func (c *Config) WithConcurrency(concurrency int) *Config { +func (c *Config) SetConcurrency(concurrency int) { c.concurrency = concurrency - return c } -// WithInitialShardIterator configures the settings used to retrieve initial +// SetInitialShardIterator configures the settings used to retrieve initial // shard iterator via the GetShardIterator call. -func (c *Config) WithInitialShardIterator(shardIterator *ShardIterator) *Config { +func (c *Config) SetInitialShardIterator(shardIterator *ShardIterator) { c.shardIterator = shardIterator - return c } -// WithGetRecordsReadTimeout configures the time to wait for each successive +// SetGetRecordsReadTimeout configures the time to wait for each successive // Read operation on the GetRecords response payload. -func (c *Config) WithGetRecordsReadTimeout(timouet time.Duration) *Config { +func (c *Config) SetGetRecordsReadTimeout(timouet time.Duration) { c.getRecordsReadTimeout = timouet - return c } -// WithStatsListener configures a listener to handle metrics. -func (c *Config) WithStatsListener(stats StatsListener) *Config { +// SetStatsListener configures a listener to handle metrics. +func (c *Config) SetStatsListener(stats StatsListener) { c.stats = stats +} + +// FromKinetic configures the session from Kinetic. +func (c *Config) FromKinetic(k *kinetic.Kinetic) *Config { + c.AwsConfig = k.GetSession().Config return c } + +// GetSession creates an instance of the session.Session to be used when creating service +// clients in aws-sdk-go. +func (c *Config) GetSession() (*session.Session, error) { + return session.NewSession(c.AwsConfig) +} diff --git a/listener/config_test.go b/listener/config_test.go index 76b72b0..b0ac1e7 100644 --- a/listener/config_test.go +++ b/listener/config_test.go @@ -3,7 +3,6 @@ package listener import ( . "github.com/smartystreets/goconvey/convey" - "fmt" "testing" "time" @@ -50,82 +49,26 @@ func TestNewConfig(t *testing.T) { So(config.shardIterator.shardIteratorType, ShouldEqual, "TRIM_HORIZON") So(config.getRecordsReadTimeout, ShouldEqual, 1*time.Second) So(config.stats, ShouldHaveSameTypeAs, &NilStatsListener{}) - So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 5*time.Minute) + So(config.logLevel.Value(), ShouldEqual, logging.LogOff) }) Convey("check that we can retrieve an aws.Session from it ", func() { getSession(config) }) - Convey("check that we can set credentials", func() { - config = config.WithCredentials("access-key", "secret-key", "security-token") - sess := getSession(config) - creds, err := sess.Config.Credentials.Get() - So(err, ShouldBeNil) - So(creds.AccessKeyID, ShouldEqual, "access-key") - So(creds.SecretAccessKey, ShouldEqual, "secret-key") - So(creds.SessionToken, ShouldEqual, "security-token") - }) - - Convey("check that we can set the region", func() { - config = config.WithRegion("my-region") - sess := getSession(config) - So(aws.StringValue(sess.Config.Region), ShouldEqual, "my-region") - }) - - Convey("check that we can set the endpoint", func() { - config = config.WithEndpoint("my-endpoint") - sess := getSession(config) - So(aws.StringValue(sess.Config.Endpoint), ShouldEqual, "my-endpoint") - }) - - Convey("check that we can configure a logger", func() { - var logs []string - loggerFn := func(args ...interface{}) { - logs = append(logs, fmt.Sprint(args...)) - } - config = config.WithLogger(aws.LoggerFunc(loggerFn)) - sess := getSession(config) - - Convey("check that basic logging should work", func() { - sess.Config.Logger.Log("one") - sess.Config.Logger.Log("two") - sess.Config.Logger.Log("three") - So(len(logs), ShouldEqual, 3) - So(logs, ShouldContain, "one") - So(logs, ShouldContain, "two") - So(logs, ShouldContain, "three") - Reset(func() { - logs = nil - }) - }) - }) - - Convey("check that the default log level is off for both the sdk and kinetic", func() { - sess := getSession(config) - So(sess.Config.LogLevel.Value(), ShouldEqual, aws.LogOff) - So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeFalse) - So(config.logLevel.Value(), ShouldEqual, aws.LogOff) - So(config.logLevel.AtLeast(aws.LogDebug), ShouldBeFalse) - }) - Convey("check that we can set both the sdk and kinetic log level", func() { ll := aws.LogDebug | aws.LogDebugWithSigning | logging.LogDebug - config = config.WithLogLevel(ll) + config.SetLogLevel(ll) sess := getSession(config) So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) So(config.logLevel.AtLeast(logging.LogDebug), ShouldBeTrue) }) - Convey("check that we can set the http.Client Timeout", func() { - config = config.WithHTTPClientTimeout(10 * time.Minute) - So(config.awsConfig.HTTPClient.Timeout, ShouldEqual, 10*time.Minute) - }) - Convey("check that we can import configuration from kinetic", func() { - k, err := kinetic.New(kinetic.NewConfig(). - WithEndpoint("bogus-endpoint")) + k, err := kinetic.New(func(c *kinetic.Config) { + c.SetEndpoint("bogus-endpoint") + }) So(err, ShouldBeNil) config = config.FromKinetic(k) sess := getSession(config) @@ -133,36 +76,36 @@ func TestNewConfig(t *testing.T) { }) Convey("check that we can set the batch size", func() { - config = config.WithBatchSize(1000) + config.SetBatchSize(1000) So(config.batchSize, ShouldEqual, 1000) }) Convey("check that we can set the concurrency limit", func() { - config = config.WithConcurrency(50) + config.SetConcurrency(50) So(config.concurrency, ShouldEqual, 50) }) Convey("check that the default shard iterator is TRIM_HORIZON", func() { - config = config.WithInitialShardIterator(NewShardIterator()) + config.SetInitialShardIterator(NewShardIterator()) So(config.shardIterator.shardIteratorType, ShouldEqual, "TRIM_HORIZON") So(config.shardIterator.getStartingSequenceNumber(), ShouldBeNil) So(config.shardIterator.getTimestamp(), ShouldBeNil) }) Convey("check that we can set the initial shard iterator (to LATEST)", func() { - config = config.WithInitialShardIterator(NewShardIterator().Latest()) + config.SetInitialShardIterator(NewShardIterator().Latest()) So(config.shardIterator.shardIteratorType, ShouldEqual, "LATEST") So(config.shardIterator.getStartingSequenceNumber(), ShouldBeNil) So(config.shardIterator.getTimestamp(), ShouldBeNil) }) Convey("check that we can set the read timeout for the GetRecords request", func() { - config = config.WithGetRecordsReadTimeout(10 * time.Second) + config.SetGetRecordsReadTimeout(10 * time.Second) So(config.getRecordsReadTimeout, ShouldEqual, 10*time.Second) }) Convey("check that we can configure a stats listener", func() { - config = config.WithStatsListener(&DebugStatsListener{}) + config.SetStatsListener(&DebugStatsListener{}) So(config.stats, ShouldHaveSameTypeAs, &DebugStatsListener{}) }) }) diff --git a/listener/listener.go b/listener/listener.go index 0ba0d75..ff77703 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -125,7 +125,7 @@ func (it *ShardIterator) getTimestamp() *time.Time { return aws.Time(it.timestamp) } -type listenerConfig struct { +type listenerOptions struct { stream string shard string @@ -140,7 +140,7 @@ type listenerConfig struct { // Listener polls the Kinesis stream for messages. type Listener struct { - *listenerConfig + *listenerOptions nextShardIterator string @@ -159,17 +159,19 @@ type Listener struct { // NewListener creates a new listener for listening to message on a Kinesis // stream. -func NewListener(config *Config) (*Listener, error) { +func NewListener(stream, shard string, fn func(*Config)) (*Listener, error) { + config := NewConfig(stream, shard) + fn(config) session, err := config.GetSession() if err != nil { return nil, err } return &Listener{ - listenerConfig: config.listenerConfig, - concurrencySem: make(chan Empty, config.concurrency), - throttleSem: make(chan Empty, 5), - pipeOfDeath: make(chan Empty), - session: session, + listenerOptions: config.listenerOptions, + concurrencySem: make(chan Empty, config.concurrency), + throttleSem: make(chan Empty, 5), + pipeOfDeath: make(chan Empty), + session: session, }, nil } diff --git a/listener/listener_test.go b/listener/listener_test.go index f2ece02..8c5ac51 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -82,10 +82,11 @@ func TestShardIterator(t *testing.T) { func TestListener(t *testing.T) { Convey("given a listener", t, func() { - k, err := kinetic.New(kinetic.NewConfig(). - WithCredentials("some-access-key", "some-secret-key", "some-security-token"). - WithRegion("some-region"). - WithEndpoint("http://127.0.0.1:4567")) + k, err := kinetic.New(func(c *kinetic.Config) { + c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") + c.SetRegion("some-region") + c.SetEndpoint("http://127.0.0.1:4567") + }) stream := "some-listener-stream" @@ -99,9 +100,10 @@ func TestListener(t *testing.T) { So(err, ShouldBeNil) So(len(shards), ShouldEqual, 1) - l, err := NewListener(NewConfig(stream, shards[0]). - FromKinetic(k). - WithConcurrency(10)) + l, err := NewListener(stream, shards[0], func(c *Config) { + c.FromKinetic(k) + c.SetConcurrency(10) + }) So(l, ShouldNotBeNil) So(err, ShouldBeNil) From 60f11c020126133515ccc3121b41645e7dc0dd67 Mon Sep 17 00:00:00 2001 From: Tommy Wang Date: Fri, 31 Mar 2017 09:10:56 -0500 Subject: [PATCH 007/109] Disable ListenFn test * There is probably a race condition in one of the tests. I've disabled it for now. It may not be worth adding complexity to fix the test. --- listener/listener_test.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/listener/listener_test.go b/listener/listener_test.go index 8c5ac51..84e479d 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -250,7 +250,7 @@ func TestListener(t *testing.T) { So(called, ShouldBeTrue) }) - Convey("check that listen can deliver messages to fn", func(c C) { + SkipConvey("check that listen can deliver messages to fn", func(c C) { planets := []string{"mercury", "venus", "earth", "mars", "jupiter", "saturn", "neptune", "uranus"} var count int64 var wg sync.WaitGroup @@ -266,6 +266,8 @@ func TestListener(t *testing.T) { _, err := putRecord(l, []byte(planet)) So(err, ShouldBeNil) } + // FIXME: probably a race condition here as consume may + // not have grabbed all data from the channel yet. close(l.pipeOfDeath) wg.Wait() So(count, ShouldEqual, len(planets)) From c78e1291a3a2cb765eed8b3735523658643487ce Mon Sep 17 00:00:00 2001 From: Tommy Wang Date: Fri, 31 Mar 2017 09:23:02 -0500 Subject: [PATCH 008/109] Fix race condition in ListenFn test - Use a timeout (bleh!) of 10s and a polling rate of 1s to give consume a chance to see the last planet. --- listener/listener_test.go | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/listener/listener_test.go b/listener/listener_test.go index 84e479d..b702a28 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -250,7 +250,7 @@ func TestListener(t *testing.T) { So(called, ShouldBeTrue) }) - SkipConvey("check that listen can deliver messages to fn", func(c C) { + Convey("check that listen can deliver messages to fn", func(c C) { planets := []string{"mercury", "venus", "earth", "mars", "jupiter", "saturn", "neptune", "uranus"} var count int64 var wg sync.WaitGroup @@ -266,11 +266,30 @@ func TestListener(t *testing.T) { _, err := putRecord(l, []byte(planet)) So(err, ShouldBeNil) } + timeout := time.After(10 * time.Second) + // FIXME: Not too thrilled with this implementation, but + // there is probably a race condition between when the + // last planet is put onto the Kinesis stream (and + // subsequently read by consume) with when closing the + // pipeOfDeath (which will shut down the consume loop) + // such that we may not see all the planets inside + // Listen. + stop: + for { + select { + case <-time.After(1 * time.Second): + if atomic.LoadInt64(&count) == int64(len(planets)) { + break stop + } + case <-timeout: + break stop + } + } // FIXME: probably a race condition here as consume may // not have grabbed all data from the channel yet. close(l.pipeOfDeath) wg.Wait() - So(count, ShouldEqual, len(planets)) + So(atomic.LoadInt64(&count), ShouldEqual, len(planets)) }) Convey("check that listen can be cancelled by context", func(c C) { From 2fcc21d1eac2ce7de64b280daaec751b022be273 Mon Sep 17 00:00:00 2001 From: Tommy Wang Date: Wed, 5 Apr 2017 17:08:14 -0500 Subject: [PATCH 009/109] producer -- work in progress --- config.go | 16 +- config_test.go | 4 +- kinetic.go | 25 +-- kinetic_test.go | 5 - listener/config.go | 31 ++- listener/config_test.go | 40 ++-- listener/listener.go | 89 ++++---- listener/stats.go | 56 +++-- producer/config.go | 81 +++++++ producer/firehosewriter/config.go | 1 + producer/firehosewriter/firehosewriter.go | 1 + producer/kinesiswriter/config.go | 13 ++ producer/kinesiswriter/kinesiswriter.go | 147 +++++++++++++ producer/producer.go | 247 ++++++++++++++++++++++ producer/stats.go | 79 +++++++ 15 files changed, 688 insertions(+), 147 deletions(-) create mode 100644 producer/config.go create mode 100644 producer/firehosewriter/config.go create mode 100644 producer/firehosewriter/firehosewriter.go create mode 100644 producer/kinesiswriter/config.go create mode 100644 producer/kinesiswriter/kinesiswriter.go create mode 100644 producer/producer.go create mode 100644 producer/stats.go diff --git a/config.go b/config.go index e97a2d1..1ddf48a 100644 --- a/config.go +++ b/config.go @@ -59,6 +59,12 @@ func (c *AwsOptions) SetHTTPClientTimeout(timeout time.Duration) { }) } +// GetSession creates an instance of the session.Session to be used when creating service +// clients in aws-sdk-go. +func (c *AwsOptions) GetSession() (*session.Session, error) { + return session.NewSession(c.AwsConfig) +} + // Config is used to configure a Kinetic instance type Config struct { *AwsOptions @@ -70,7 +76,7 @@ func NewConfig() *Config { return &Config{ AwsOptions: DefaultAwsOptions(), kineticOptions: &kineticOptions{ - logLevel: aws.LogOff, + LogLevel: aws.LogOff, }, } } @@ -78,11 +84,5 @@ func NewConfig() *Config { // SetLogLevel configures both the SDK and Kinetic log levels. func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { c.AwsOptions.SetLogLevel(logLevel) - c.kineticOptions.logLevel = logLevel & 0xffff0000 -} - -// GetSession creates an instance of the session.Session to be used when creating service -// clients in aws-sdk-go. -func (c *Config) GetSession() (*session.Session, error) { - return session.NewSession(c.AwsConfig) + c.LogLevel = logLevel & 0xffff0000 } diff --git a/config_test.go b/config_test.go index c95abba..e15ec3f 100644 --- a/config_test.go +++ b/config_test.go @@ -27,7 +27,7 @@ func TestNewConfig(t *testing.T) { Convey("check the default values for its non-zero config", func() { So(config.AwsConfig.HTTPClient.Timeout, ShouldEqual, 10*time.Minute) So(config.AwsConfig.LogLevel.Value(), ShouldEqual, aws.LogOff) - So(config.logLevel.Value(), ShouldEqual, logging.LogOff) + So(config.LogLevel.Value(), ShouldEqual, logging.LogOff) }) Convey("check that we can retrieve an aws.Session from it ", func() { @@ -84,7 +84,7 @@ func TestNewConfig(t *testing.T) { sess := getSession(config) So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) - So(config.logLevel.AtLeast(logging.LogDebug), ShouldBeTrue) + So(config.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) }) Convey("check that we can set the http.Client Timeout", func() { diff --git a/kinetic.go b/kinetic.go index 65fc899..814eb27 100644 --- a/kinetic.go +++ b/kinetic.go @@ -29,18 +29,18 @@ var ( ) type kineticOptions struct { - logLevel aws.LogLevelType + LogLevel aws.LogLevelType } // Kinetic represents a kinesis and firehose client and provides some utility // methods for interacting with the AWS services. type Kinetic struct { *kineticOptions - session *session.Session + clientMu sync.Mutex fclient firehoseiface.FirehoseAPI kclient kinesisiface.KinesisAPI - clientMu sync.Mutex + Session *session.Session } // New creates a new instance of Kientic. @@ -53,14 +53,14 @@ func New(fn func(*Config)) (*Kinetic, error) { } return &Kinetic{ kineticOptions: config.kineticOptions, - session: session, + Session: session, }, nil } // Log logs a message if LogDebug is set. func (k *Kinetic) Log(args ...interface{}) { - if k.logLevel.Matches(logging.LogDebug) { - k.session.Config.Logger.Log(args...) + if k.LogLevel.Matches(logging.LogDebug) { + k.Session.Config.Logger.Log(args...) } } @@ -68,7 +68,7 @@ func (k *Kinetic) ensureKinesisClient() { k.clientMu.Lock() defer k.clientMu.Unlock() if k.kclient == nil { - k.kclient = kinesis.New(k.session) + k.kclient = kinesis.New(k.Session) } } @@ -121,7 +121,7 @@ func (k *Kinetic) WaitUntilStreamDeleted(ctx context.Context, stream string, opt Expected: kinesis.ErrCodeResourceNotFoundException, }, }, - Logger: k.session.Config.Logger, + Logger: k.Session.Config.Logger, NewRequest: func(opts []request.Option) (*request.Request, error) { req, _ := k.kclient.DescribeStreamRequest(&kinesis.DescribeStreamInput{ StreamName: aws.String(stream), // Required @@ -159,12 +159,3 @@ func (k *Kinetic) GetShards(stream string) ([]string, error) { } return shards, nil } - -// GetSession returns the aws-sdk-go session.Session object. -func (k *Kinetic) GetSession() *session.Session { - return k.session -} - -// func (k *Kinetic) NewListener(config *listener.Config) (*listener.Listener, error) { -// return listener.NewListener(config, k, k.session, k.kclient) -// } diff --git a/kinetic_test.go b/kinetic_test.go index 4f33a16..a59948c 100644 --- a/kinetic_test.go +++ b/kinetic_test.go @@ -22,11 +22,6 @@ func TestKinetic(t *testing.T) { So(k, ShouldNotBeNil) So(err, ShouldBeNil) - Convey("check that we can get the session", func() { - sess := k.GetSession() - So(sess, ShouldNotBeNil) - }) - stream := "some-kinetic-stream" Convey("check that calling ensureKinesisClient twice doesn't overwrite existing client", func() { diff --git a/listener/config.go b/listener/config.go index 4e7b25d..acce032 100644 --- a/listener/config.go +++ b/listener/config.go @@ -4,7 +4,6 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/aws/session" "github.com/rewardStyle/kinetic" "github.com/rewardStyle/kinetic/logging" @@ -27,18 +26,12 @@ func NewConfig(stream, shard string) *Config { concurrency: 10000, shardIterator: NewShardIterator(), getRecordsReadTimeout: 1 * time.Second, - stats: &NilStatsListener{}, - logLevel: logging.LogOff, + LogLevel: logging.LogOff, + Stats: &NilStatsCollector{}, }, } } -// SetLogLevel configures both the SDK and Kinetic log levels. -func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { - c.AwsOptions.SetLogLevel(logLevel) - c.listenerOptions.logLevel = logLevel & 0xffff0000 -} - // SetBatchSize configures the batch size of the GetRecords call. func (c *Config) SetBatchSize(batchSize int) { c.batchSize = batchSize @@ -62,19 +55,19 @@ func (c *Config) SetGetRecordsReadTimeout(timouet time.Duration) { c.getRecordsReadTimeout = timouet } -// SetStatsListener configures a listener to handle metrics. -func (c *Config) SetStatsListener(stats StatsListener) { - c.stats = stats +// SetLogLevel configures both the SDK and Kinetic log levels. +func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { + c.AwsOptions.SetLogLevel(logLevel) + c.LogLevel = logLevel & 0xffff0000 +} + +// SetStatsCollector configures a listener to handle listener metrics. +func (c *Config) SetStatsCollector(stats StatsCollector) { + c.Stats = stats } // FromKinetic configures the session from Kinetic. func (c *Config) FromKinetic(k *kinetic.Kinetic) *Config { - c.AwsConfig = k.GetSession().Config + c.AwsConfig = k.Session.Config return c } - -// GetSession creates an instance of the session.Session to be used when creating service -// clients in aws-sdk-go. -func (c *Config) GetSession() (*session.Session, error) { - return session.NewSession(c.AwsConfig) -} diff --git a/listener/config_test.go b/listener/config_test.go index b0ac1e7..82953ba 100644 --- a/listener/config_test.go +++ b/listener/config_test.go @@ -13,20 +13,20 @@ import ( "github.com/rewardStyle/kinetic/logging" ) -type DebugStatsListener struct{} - -func (l *DebugStatsListener) AddConsumedSample(int) {} -func (l *DebugStatsListener) AddDeliveredSample(int) {} -func (l *DebugStatsListener) AddProcessedSample(int) {} -func (l *DebugStatsListener) AddBatchSizeSample(int) {} -func (l *DebugStatsListener) AddGetRecordsCalled(int) {} -func (l *DebugStatsListener) AddProvisionedThroughputExceeded(int) {} -func (l *DebugStatsListener) AddGetRecordsTimeout(int) {} -func (l *DebugStatsListener) AddGetRecordsReadTimeout(int) {} -func (l *DebugStatsListener) AddProcessedTime(time.Duration) {} -func (l *DebugStatsListener) AddGetRecordsTime(time.Duration) {} -func (l *DebugStatsListener) AddGetRecordsReadResponseTime(time.Duration) {} -func (l *DebugStatsListener) AddGetRecordsUnmarshalTime(time.Duration) {} +type DebugStatsCollector struct{} + +func (l *DebugStatsCollector) AddConsumedSample(int) {} +func (l *DebugStatsCollector) AddDeliveredSample(int) {} +func (l *DebugStatsCollector) AddProcessedSample(int) {} +func (l *DebugStatsCollector) AddBatchSizeSample(int) {} +func (l *DebugStatsCollector) AddGetRecordsCalled(int) {} +func (l *DebugStatsCollector) AddProvisionedThroughputExceeded(int) {} +func (l *DebugStatsCollector) AddGetRecordsTimeout(int) {} +func (l *DebugStatsCollector) AddGetRecordsReadTimeout(int) {} +func (l *DebugStatsCollector) AddProcessedTime(time.Duration) {} +func (l *DebugStatsCollector) AddGetRecordsTime(time.Duration) {} +func (l *DebugStatsCollector) AddGetRecordsReadResponseTime(time.Duration) {} +func (l *DebugStatsCollector) AddGetRecordsUnmarshalTime(time.Duration) {} func getSession(config *Config) *session.Session { sess, err := config.GetSession() @@ -48,8 +48,8 @@ func TestNewConfig(t *testing.T) { So(config.concurrency, ShouldEqual, 10000) So(config.shardIterator.shardIteratorType, ShouldEqual, "TRIM_HORIZON") So(config.getRecordsReadTimeout, ShouldEqual, 1*time.Second) - So(config.stats, ShouldHaveSameTypeAs, &NilStatsListener{}) - So(config.logLevel.Value(), ShouldEqual, logging.LogOff) + So(config.Stats, ShouldHaveSameTypeAs, &NilStatsCollector{}) + So(config.LogLevel.Value(), ShouldEqual, logging.LogOff) }) Convey("check that we can retrieve an aws.Session from it ", func() { @@ -62,7 +62,7 @@ func TestNewConfig(t *testing.T) { sess := getSession(config) So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) - So(config.logLevel.AtLeast(logging.LogDebug), ShouldBeTrue) + So(config.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) }) Convey("check that we can import configuration from kinetic", func() { @@ -104,9 +104,9 @@ func TestNewConfig(t *testing.T) { So(config.getRecordsReadTimeout, ShouldEqual, 10*time.Second) }) - Convey("check that we can configure a stats listener", func() { - config.SetStatsListener(&DebugStatsListener{}) - So(config.stats, ShouldHaveSameTypeAs, &DebugStatsListener{}) + Convey("check that we can configure a stats collector", func() { + config.SetStatsCollector(&DebugStatsCollector{}) + So(config.Stats, ShouldHaveSameTypeAs, &DebugStatsCollector{}) }) }) } diff --git a/listener/listener.go b/listener/listener.go index ff77703..c666cd8 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -133,9 +133,9 @@ type listenerOptions struct { concurrency int shardIterator *ShardIterator getRecordsReadTimeout time.Duration - stats StatsListener - logLevel aws.LogLevelType + LogLevel aws.LogLevelType + Stats StatsCollector } // Listener polls the Kinesis stream for messages. @@ -152,9 +152,9 @@ type Listener struct { consuming bool consumingMu sync.Mutex - session *session.Session - client kinesisiface.KinesisAPI clientMu sync.Mutex + client kinesisiface.KinesisAPI + Session *session.Session } // NewListener creates a new listener for listening to message on a Kinesis @@ -171,23 +171,23 @@ func NewListener(stream, shard string, fn func(*Config)) (*Listener, error) { concurrencySem: make(chan Empty, config.concurrency), throttleSem: make(chan Empty, 5), pipeOfDeath: make(chan Empty), - session: session, + Session: session, }, nil } // Log a debug message using the AWS SDK logger. func (l *Listener) Log(args ...interface{}) { - if l.session.Config.LogLevel.Matches(logging.LogDebug) { - l.session.Config.Logger.Log(args...) + if l.Session.Config.LogLevel.Matches(logging.LogDebug) { + l.Session.Config.Logger.Log(args...) } } // setNextShardIterator sets the nextShardIterator to use when calling // GetRecords. // -// Not thread-safe. Only called from fetchBatch (and ensureShardIterator, -// which is called from fetchBatch). Care must be taken to ensure that only -// one call to Listen and Retrieve/RetrieveFn can be running at a time. +// Not thread-safe. Only called from getRecords (and ensureShardIterator, which +// is called from getRecords). Care must be taken to ensure that only one call +// to Listen and Retrieve/RetrieveFn can be running at a time. func (l *Listener) setNextShardIterator(shardIterator string) error { if len(shardIterator) == 0 { return ErrEmptyShardIterator @@ -201,7 +201,7 @@ func (l *Listener) setNextShardIterator(shardIterator string) error { // This is only used when we need to call getShardIterator (say, to refresh the // shard iterator). // -// Not thread-safe. Only called from fetchBatch. Care must be taken to ensure +// Not thread-safe. Only called from getRecords. Care must be taken to ensure // that only one call to Listen and Retrieve/RetrieveFn can be running at a // time. func (l *Listener) setSequenceNumber(sequenceNumber string) error { @@ -238,7 +238,7 @@ func (l *Listener) ensureClient() { l.clientMu.Lock() defer l.clientMu.Unlock() if l.client == nil { - l.client = kinesis.New(l.session) + l.client = kinesis.New(l.Session) } } @@ -246,7 +246,7 @@ func (l *Listener) ensureClient() { // calling the GetShardIterator API with the configured ShardIteratorType (with // any applicable StartingSequenceNumber or Timestamp) if necessary. // -// Not thread-safe. Only called from fetchBatch Care must be taken to ensure +// Not thread-safe. Only called from getRecords Care must be taken to ensure // that only one call to Listen and Retrieve/RetrieveFn can be running at a // time. func (l *Listener) ensureShardIterator() error { @@ -284,7 +284,7 @@ func (l *Listener) throttle(sem chan Empty) { }) } -// fetchBatch calls GetRecords and delivers each record into the messages +// getRecords calls GetRecords and delivers each record into the messages // channel. // TODO: Convert timeout implementation to use context.Context // FIXME: Need to investigate that the timeout implementation doesn't result in @@ -296,7 +296,7 @@ func (l *Listener) throttle(sem chan Empty) { // down the TCP connection. Worst case scenario is that our client Timeout // eventually fires and closes the socket, but this can be susceptible to FD // exhaustion. -func (l *Listener) fetchBatch(size int) (int, error) { +func (l *Listener) getRecords(size int) (int, error) { l.ensureClient() if err := l.ensureShardIterator(); err != nil { return 0, err @@ -315,7 +315,7 @@ func (l *Listener) fetchBatch(size int) (int, error) { }) // If debug is turned on, add some handlers for GetRecords logging - if l.session.Config.LogLevel.AtLeast(aws.LogDebug) { + if l.LogLevel.AtLeast(logging.LogDebug) { req.Handlers.Send.PushBack(func(r *request.Request) { l.Log("Finished GetRecords Send, took", time.Since(start)) }) @@ -395,7 +395,7 @@ func (l *Listener) fetchBatch(size int) (int, error) { return }, OnCloseFn: func() { - l.stats.AddGetRecordsReadResponseTime(time.Since(startReadTime)) + l.Stats.AddGetRecordsReadResponseTime(time.Since(startReadTime)) l.Log("Finished GetRecords body read, took", time.Since(start)) startUnmarshalTime = time.Now() }, @@ -403,18 +403,18 @@ func (l *Listener) fetchBatch(size int) (int, error) { }) req.Handlers.Unmarshal.PushBack(func(r *request.Request) { - l.stats.AddGetRecordsUnmarshalTime(time.Since(startUnmarshalTime)) + l.Stats.AddGetRecordsUnmarshalTime(time.Since(startUnmarshalTime)) l.Log("Finished GetRecords Unmarshal, took", time.Since(start)) }) // Send the GetRecords request - l.Log("Starting GetRecords build/sign request, took", time.Since(start)) - l.stats.AddGetRecordsCalled(1) + l.Log("Starting GetRecords Build/Sign request, took", time.Since(start)) + l.Stats.AddGetRecordsCalled(1) if err := req.Send(); err != nil { l.Log("Error getting records:", err) return 0, err } - l.stats.AddGetRecordsTime(time.Since(start)) + l.Stats.AddGetRecordsTime(time.Since(start)) // Process Records l.Log(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), l.shard, time.Since(start))) @@ -422,12 +422,12 @@ func (l *Listener) fetchBatch(size int) (int, error) { return 0, ErrNilGetRecordsResponse } delivered := 0 - l.stats.AddBatchSizeSample(len(resp.Records)) + l.Stats.AddBatchSizeSample(len(resp.Records)) for _, record := range resp.Records { if record != nil { delivered++ l.messages <- &message.Message{Record: *record} - l.stats.AddConsumedSample(1) + l.Stats.AddConsumedSample(1) } if record.SequenceNumber != nil { // We can safely ignore if this call returns @@ -530,12 +530,12 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, e if !ok { return nil, err } - n, err := l.fetchBatch(1) + n, err := l.getRecords(1) if err != nil { return nil, err } if n > 0 { - l.stats.AddDeliveredSample(1) + l.Stats.AddDeliveredSample(1) return <-l.messages, nil } } @@ -559,8 +559,8 @@ func (l *Listener) RetrieveFnWithContext(ctx context.Context, fn MessageFn) erro start := time.Now() go fn(msg.Value(), &wg) wg.Wait() - l.stats.AddProcessedTime(time.Since(start)) - l.stats.AddProcessedSample(1) + l.Stats.AddProcessedTime(time.Since(start)) + l.Stats.AddProcessedSample(1) return nil } @@ -570,7 +570,7 @@ func (l *Listener) RetrieveFn(fn MessageFn) error { return l.RetrieveFnWithContext(context.TODO(), fn) } -// consume calls fetchBatch with configured batch size in a loop until the +// consume calls getRecords with configured batch size in a loop until the // listener is stopped. func (l *Listener) consume(ctx context.Context) { // We need to run blockConsumers & startConsuming to make sure that we @@ -589,17 +589,16 @@ func (l *Listener) consume(ctx context.Context) { }() stop: for { - ok, err := l.shouldConsume(ctx) + ok, _ := l.shouldConsume(ctx) if !ok { break stop } - _, err = l.fetchBatch(l.batchSize) - + _, err = l.getRecords(l.batchSize) if err != nil { switch err := err.(type) { case net.Error: if err.Timeout() { - l.stats.AddGetRecordsTimeout(1) + l.Stats.AddGetRecordsTimeout(1) l.Log("Received net error:", err.Error()) } else { l.Log("Received unknown net error:", err.Error()) @@ -607,31 +606,15 @@ func (l *Listener) consume(ctx context.Context) { case error: switch err { case ErrTimeoutReadResponseBody: - l.stats.AddGetRecordsReadTimeout(1) + l.Stats.AddGetRecordsReadTimeout(1) l.Log("Received error:", err.Error()) - case ErrEmptySequenceNumber: - fallthrough - case ErrEmptyShardIterator: - fallthrough - case ErrNilGetShardIteratorResponse: - fallthrough - case ErrNilShardIterator: - fallthrough - case ErrNilGetRecordsResponse: - fallthrough default: l.Log("Received error:", err.Error()) } case awserr.Error: switch err.Code() { case kinesis.ErrCodeProvisionedThroughputExceededException: - l.stats.AddProvisionedThroughputExceeded(1) - case kinesis.ErrCodeResourceNotFoundException: - fallthrough - case kinesis.ErrCodeInvalidArgumentException: - fallthrough - case kinesis.ErrCodeExpiredIteratorException: - fallthrough + l.Stats.AddProvisionedThroughputExceeded(1) default: l.Log("Received AWS error:", err.Error()) } @@ -652,7 +635,7 @@ func (l *Listener) ListenWithContext(ctx context.Context, fn MessageFn) { defer wg.Wait() for msg := range l.messages { - l.stats.AddDeliveredSample(1) + l.Stats.AddDeliveredSample(1) l.concurrencySem <- Empty{} wg.Add(1) go func(msg *message.Message) { @@ -664,8 +647,8 @@ func (l *Listener) ListenWithContext(ctx context.Context, fn MessageFn) { start := time.Now() fn(msg.Value(), &fnWg) fnWg.Wait() - l.stats.AddProcessedTime(time.Since(start)) - l.stats.AddProcessedSample(1) + l.Stats.AddProcessedTime(time.Since(start)) + l.Stats.AddProcessedSample(1) wg.Done() }(msg) } diff --git a/listener/stats.go b/listener/stats.go index c7eeb73..491ab27 100644 --- a/listener/stats.go +++ b/listener/stats.go @@ -4,80 +4,90 @@ import ( "time" ) -// StatsListener allows for a listener to "listen" for various metrics produced -// by the Kinetic library. This was really built with rcrowley/go-metrics in -// mind. -type StatsListener interface { - // histograms +// StatsCollector allows for a collector to collect various metrics produced by +// the Kinetic Listener library. This was really built with rcrowley/go-metrics +// in mind. +type StatsCollector interface { + // histograms give us the count, sum, min, max, mean, percentiles, + // standard deviation, and variance of the data. these metrics should + // give us the total number (sum) of messages consumed, delivered, and + // processed, as well as the average (mean) batch size. AddConsumedSample(int) AddDeliveredSample(int) AddProcessedSample(int) AddBatchSizeSample(int) - // meters + // meters give us the count and rate of the data. these metrics should + // give us the average number of times: + // - GetRecords was called per second + // - ProvisionedThroughputExceeded was received per second + // - GetRecords timed out per second + // - GetRecords read timed out per second AddGetRecordsCalled(int) AddProvisionedThroughputExceeded(int) AddGetRecordsTimeout(int) AddGetRecordsReadTimeout(int) - // timers + // timers give us the count, sum, min, max, mean, percentiles, standard + // deviation, variance, as well as the rate of the data. + // TODO: describe these metrics better AddProcessedTime(time.Duration) AddGetRecordsTime(time.Duration) AddGetRecordsReadResponseTime(time.Duration) AddGetRecordsUnmarshalTime(time.Duration) } -// NilStatsListener is a stats listener that ignores all metrics. -type NilStatsListener struct{} +// NilStatsCollector is a stats listener that ignores all metrics. +type NilStatsCollector struct{} // AddConsumedSample records a count of the number of messages received from AWS // Kinesis by the listener. -func (l *NilStatsListener) AddConsumedSample(int) {} +func (l *NilStatsCollector) AddConsumedSample(int) {} // AddDeliveredSample records a count of the number of messages delivered to the // application by the listener. -func (l *NilStatsListener) AddDeliveredSample(int) {} +func (l *NilStatsCollector) AddDeliveredSample(int) {} // AddProcessedSample records a count of the number of messages processed by the // application by the listener. This is based on a WaitGroup that is sent to // the RetrieveFn and Listen functions. Retrieve does not count processed // messages. -func (l *NilStatsListener) AddProcessedSample(int) {} +func (l *NilStatsCollector) AddProcessedSample(int) {} // AddBatchSizeSample records a count of the number of messages returned by // GetRecords in the listener. -func (l *NilStatsListener) AddBatchSizeSample(int) {} +func (l *NilStatsCollector) AddBatchSizeSample(int) {} // AddGetRecordsCalled records the number of times the GetRecords API was called // by the listener. -func (l *NilStatsListener) AddGetRecordsCalled(int) {} +func (l *NilStatsCollector) AddGetRecordsCalled(int) {} // AddProvisionedThroughputExceeded records the number of times the GetRecords // API returned a ErrCodeProvisionedThroughputExceededException by the listener. -func (l *NilStatsListener) AddProvisionedThroughputExceeded(int) {} +func (l *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} // AddGetRecordsTimeout records the number of times the GetRecords API timed out // on the HTTP level. This is influenced by the WithHTTPClientTimeout // configuration. -func (l *NilStatsListener) AddGetRecordsTimeout(int) {} +func (l *NilStatsCollector) AddGetRecordsTimeout(int) {} -// AddGetRecordsReadTimeout records the number of times the GetRecords API timed out -// while reading the response body. This is influenced by the +// AddGetRecordsReadTimeout records the number of times the GetRecords API timed +// out while reading the response body. This is influenced by the // WithGetRecordsReadTimeout configuration. -func (l *NilStatsListener) AddGetRecordsReadTimeout(int) {} +func (l *NilStatsCollector) AddGetRecordsReadTimeout(int) {} // AddProcessedTime records the duration to process a record. See notes on // AddProcessedSample. -func (l *NilStatsListener) AddProcessedTime(time.Duration) {} +func (l *NilStatsCollector) AddProcessedTime(time.Duration) {} // AddGetRecordsTime records the duration that the GetRecords API request took. // Only the times of successful calls are measured. -func (l *NilStatsListener) AddGetRecordsTime(time.Duration) {} +func (l *NilStatsCollector) AddGetRecordsTime(time.Duration) {} // AddGetRecordsReadResponseTime records the duration that it took to read the // response body of a GetRecords API request. -func (l *NilStatsListener) AddGetRecordsReadResponseTime(time.Duration) {} +func (l *NilStatsCollector) AddGetRecordsReadResponseTime(time.Duration) {} // AddGetRecordsUnmarshalTime records the duration that it took to unmarshal the // response body of a GetRecords API request. -func (l *NilStatsListener) AddGetRecordsUnmarshalTime(time.Duration) {} +func (l *NilStatsCollector) AddGetRecordsUnmarshalTime(time.Duration) {} diff --git a/producer/config.go b/producer/config.go new file mode 100644 index 0000000..979665c --- /dev/null +++ b/producer/config.go @@ -0,0 +1,81 @@ +package producer + +import ( + "time" + + "github.com/aws/aws-sdk-go/aws" + + "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/logging" +) + +// Config is used to configure a Producer instance +type Config struct { + *kinetic.AwsOptions + *producerOptions +} + +// NewConfig creates a new instance of Config +func NewConfig() *Config { + return &Config{ + AwsOptions: kinetic.DefaultAwsOptions(), + producerOptions: &producerOptions{ + batchSize: 500, + batchTimeout: 1 * time.Second, + concurrency: 1, + queueDepth: 500, + LogLevel: logging.LogOff, + Stats: &NilStatsCollector{}, + }, + } +} + +// SetBatchSize configures the batch size to flush pending records to the +// PutRecords call. +func (c *Config) SetBatchSize(batchSize int) { + c.batchSize = batchSize +} + +// SetBatchTimeout configures the timeout to flush pending records to the +// PutRecords call. +func (c *Config) SetBatchTimeout(timeout time.Duration) { + c.batchTimeout = timeout +} + +// SetConcurrency controls the number of outstanding PutRecords calls may be +// active at a time. +func (c *Config) SetConcurrency(concurrency int) { + c.concurrency = concurrency +} + +// SetQueueDepth controls the number of messages that can be in the channel +// to be processed by produce at a given time. +func (c *Config) SetQueueDepth(queueDepth int) { + c.queueDepth = queueDepth +} + +// SetWriter sets the underlying stream writer (Kinesis or Firehose) for the +// producer. There can only be a single writer associated with a producer (as +// various retry logic/state is not easily shared between multiple writers). If +// multiple streams are desired, create two different producers and write to +// both. +func (c *Config) SetWriter(writer StreamWriter) { + c.writer = writer +} + +// SetLogLevel configures both the SDK and Kinetic log levels. +func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { + c.AwsOptions.SetLogLevel(logLevel) + c.LogLevel = logLevel & 0xffff0000 +} + +// SetStatsCollector configures a listener to handle producer metrics. +func (c *Config) SetStatsCollector(stats StatsCollector) { + c.Stats = stats +} + +// FromKinetic configures the session from Kinetic. +func (c *Config) FromKinetic(k *kinetic.Kinetic) *Config { + c.AwsConfig = k.Session.Config + return c +} diff --git a/producer/firehosewriter/config.go b/producer/firehosewriter/config.go new file mode 100644 index 0000000..4f15d66 --- /dev/null +++ b/producer/firehosewriter/config.go @@ -0,0 +1 @@ +package firehosewriter diff --git a/producer/firehosewriter/firehosewriter.go b/producer/firehosewriter/firehosewriter.go new file mode 100644 index 0000000..4f15d66 --- /dev/null +++ b/producer/firehosewriter/firehosewriter.go @@ -0,0 +1 @@ +package firehosewriter diff --git a/producer/kinesiswriter/config.go b/producer/kinesiswriter/config.go new file mode 100644 index 0000000..ff972d8 --- /dev/null +++ b/producer/kinesiswriter/config.go @@ -0,0 +1,13 @@ +package kinesiswriter + +import () + +type Config struct { + *kinesisWriterOptions +} + +func NewConfig() *Config { + return &Config{ + kinesisWriterOptions: &kinesisWriterOptions{}, + } +} diff --git a/producer/kinesiswriter/kinesiswriter.go b/producer/kinesiswriter/kinesiswriter.go new file mode 100644 index 0000000..a13728e --- /dev/null +++ b/producer/kinesiswriter/kinesiswriter.go @@ -0,0 +1,147 @@ +package kinesiswriter + +import ( + "errors" + "fmt" + "sync" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/request" + "github.com/aws/aws-sdk-go/service/kinesis" + "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" + + "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic/producer" +) + +var ( + // ErrNilPutRecordsResponse is returned when the PutRecords call returns + // a nil response. + ErrNilPutRecordsResponse = errors.New("PutRecords returned a nil response") + + // ErrNilFailedRecordCount is returned when the PutRecords call returns + // a nil FailedRecordCount. + ErrNilFailedRecordCount = errors.New("GetFailedRecordCount returned a nil FailedRecordCount") +) + +type kinesisWriterOptions struct { +} + +// KinesisWriter handles the API to send records to Kinesis. +type KinesisWriter struct { + *kinesisWriterOptions + + producer *producer.Producer + client kinesisiface.KinesisAPI + clientMu sync.Mutex +} + +// NewKinesisWriter creates a new stream writer to write records to a Kinesis. +func NewKinesisWriter(fn func(*Config)) (*KinesisWriter, error) { + config := NewConfig() + fn(config) + return &KinesisWriter{ + kinesisWriterOptions: config.kinesisWriterOptions, + }, nil +} + +// Log a debug message using the Producer logger. +func (w *KinesisWriter) Log(args ...interface{}) { + w.producer.Log(args...) +} + +// ensureClient will lazily make sure we have an AWS Kinesis client. +func (w *KinesisWriter) ensureClient() error { + w.clientMu.Lock() + defer w.clientMu.Unlock() + if w.client == nil { + if w.producer == nil { + return producer.ErrNilProducer + } + w.client = kinesis.New(w.producer.Session) + } + return nil +} + +// AssociateProducer associates the Kinesis stream writer to a producer. +func (w *KinesisWriter) AssociateProducer(p *producer.Producer) error { + w.clientMu.Lock() + defer w.clientMu.Unlock() + if w.producer != nil { + return producer.ErrProducerAlreadyAssociated + } + w.producer = p + return nil +} + +// PutRecords sends a batch of records to Kinesis and returns a list of records +// that need to be retried. +func (w *KinesisWriter) PutRecords(messages []*message.Message) ([]*message.Message, error) { + if err := w.ensureClient(); err != nil { + return nil, err + } + + var startSendTime time.Time + var startBuildTime time.Time + + start := time.Now() + req, resp := w.client.PutRecordsRequest(&kinesis.PutRecordsInput{}) + + req.Handlers.Build.PushFront(func(r *request.Request) { + startBuildTime = time.Now() + w.Log("Start PutRecords Build, took", time.Since(start)) + }) + + req.Handlers.Build.PushBack(func(r *request.Request) { + w.producer.Stats.AddPutRecordsBuildTime(time.Since(startBuildTime)) + w.Log("Finished PutRecords Build, took", time.Since(start)) + }) + + req.Handlers.Send.PushFront(func(r *request.Request) { + startSendTime = time.Now() + w.Log("Start PutRecords Send took", time.Since(start)) + }) + + req.Handlers.Build.PushBack(func(r *request.Request) { + w.producer.Stats.AddPutRecordsSendTime(time.Since(startSendTime)) + w.Log("Finished PutRecords Send, took", time.Since(start)) + }) + + w.Log("Starting PutRecords Build/Sign request, took", time.Since(start)) + w.producer.Stats.AddPutRecordsCalled(1) + if err := req.Send(); err != nil { + w.Log("Error putting records:", err) + return nil, err + } + w.producer.Stats.AddPutRecordsTime(time.Since(start)) + + if resp == nil { + return nil, ErrNilPutRecordsResponse + } + if resp.FailedRecordCount == nil { + return nil, ErrNilFailedRecordCount + } + attempted := len(messages) + failed := int(aws.Int64Value(resp.FailedRecordCount)) + sent := attempted - failed + w.producer.Stats.AddBatchSizeSample(len(messages)) + w.producer.Stats.AddSentSample(sent) + w.producer.Stats.AddFailedSample(failed) + w.Log(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) + + var retries []*message.Message + var err error + for idx, record := range resp.Records { + if record.SequenceNumber != nil && record.ShardId != nil { + // TODO: per-shard metrics + } else { + // TODO metrics on failure rates + retries = append(retries, messages[idx]) + } + } + if len(retries) > 0 { + err = producer.ErrRetryRecords + } + return retries, err +} diff --git a/producer/producer.go b/producer/producer.go new file mode 100644 index 0000000..a825ba0 --- /dev/null +++ b/producer/producer.go @@ -0,0 +1,247 @@ +package producer + +import ( + "context" + "errors" + "net" + "sync" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/awserr" + "github.com/aws/aws-sdk-go/aws/session" + "github.com/aws/aws-sdk-go/service/kinesis" + + "github.com/rewardStyle/kinetic/logging" + "github.com/rewardStyle/kinetic/message" +) + +// StreamWriter is an interface that abstracts the differences in API between +// Kinesis and Firehose. +type StreamWriter interface { + AssociateProducer(producer *Producer) error + PutRecords(message []*message.Message) ([]*message.Message, error) +} + +var ( + // ErrRetryRecords is returned when the PutRecords calls requires some + // records of the batch to be retried. This failure is considered part + // of normal behavior of the Kinesis stream. + ErrRetryRecords = errors.New("PutRecords requires retry of some records in batch") + + // ErrNilProducer is returned by a StreamWriter when it has not been + // correctly associated with a Producer. + ErrNilProducer = errors.New("StreamWriter not associated with a producer") + + // ErrProducerAlreadyAssociated is returned by a StreamWriter attempting + // to associate it with a Producer when it already has an association + // with a producer. + ErrProducerAlreadyAssociated = errors.New("StreamWriter already associated with a producer") + + // ErrPipeOfDeath returns when the pipe of death is closed. + ErrPipeOfDeath = errors.New("Received pipe of death") +) + +// Empty is used a as a dummy type for counting semaphore channels. +type Empty struct{} + +type producerOptions struct { + batchSize int + batchTimeout time.Duration + concurrency int + queueDepth int + writer StreamWriter + + LogLevel aws.LogLevelType + Stats StatsCollector +} + +// Producer sends records to Kinesis or Firehose. +type Producer struct { + *producerOptions + + messages chan *message.Message + retries chan []*message.Message + concurrencySem chan Empty + pipeOfDeath chan Empty + + producing bool + producingMu sync.Mutex + + Session *session.Session +} + +// NewProducer creates a new producer for writing records to a Kinesis or +// Firehose stream. +func NewProducer(fn func(*Config)) (*Producer, error) { + config := NewConfig() + fn(config) + session, err := config.GetSession() + if err != nil { + return nil, err + } + p := &Producer{ + producerOptions: config.producerOptions, + concurrencySem: make(chan Empty, config.concurrency), + pipeOfDeath: make(chan Empty), + Session: session, + } + if err := p.writer.AssociateProducer(p); err != nil { + return nil, err + } + return p, nil +} + +// Log a debug message using the AWS SDK logger. +func (p *Producer) Log(args ...interface{}) { + if p.Session.Config.LogLevel.Matches(logging.LogDebug) { + p.Session.Config.Logger.Log(args...) + } +} + +// blockProducers will set producing to true if there is not already another +// produce loop running. +func (p *Producer) blockProducers() bool { + p.producingMu.Lock() + defer p.producingMu.Unlock() + if !p.producing { + p.producing = true + return true + } + return false +} + +// startProducing handles any initialization needed in preparation to start +// producing. +func (p *Producer) startProducing() { + p.messages = make(chan *message.Message, p.queueDepth) +} + +// shouldProduce is a convenience function that allows functions to break their +// loops if the context receives a cancellation. +func (p *Producer) shouldProduce(ctx context.Context) (bool, error) { + select { + case <-p.pipeOfDeath: + return false, ErrPipeOfDeath + case <-ctx.Done(): + return false, ctx.Err() + default: + return true, nil + } +} + +// stopProducing handles any cleanup after a producing has stopped. +func (p *Producer) stopProducing() { + close(p.messages) +} + +// allowProducers allows producing. Called after blockProducers to release the +// lock on producing. +func (p *Producer) allowProducers() { + p.producingMu.Lock() + defer p.producingMu.Unlock() + p.producing = false +} + +// IsProducing returns true while producing. +func (p *Producer) IsProducing() bool { + p.producingMu.Lock() + defer p.producingMu.Unlock() + return p.producing +} + +func (p *Producer) getBatch(ctx context.Context) ([]*message.Message, error) { + ctx, cancel := context.WithTimeout(ctx, p.batchTimeout) + defer cancel() + + var batch []*message.Message + select { + case batch = <-p.retries: + default: + } + + for len(batch) < p.batchSize { + select { + case msg := <-p.messages: + batch = append(batch, msg) + case <-ctx.Done(): + return batch, ctx.Err() + } + } + return batch, nil +} + +func (p *Producer) produce(ctx context.Context) { + if !p.blockProducers() { + return + } + p.startProducing() + go func() { + defer func() { + p.stopProducing() + p.allowProducers() + }() + stop: + for { + ok, _ := p.shouldProduce(ctx) + if !ok { + break stop + } + batch, _ := p.getBatch(ctx) + if len(batch) > 0 { + p.concurrencySem <- Empty{} + go func() { + retries, err := p.writer.PutRecords(batch) + <-p.concurrencySem + if err != nil { + switch err := err.(type) { + case net.Error: + if err.Timeout() { + p.Stats.AddPutRecordsTimeout(1) + p.Log("Received net error:", err.Error()) + } else { + p.Log("Received unknown net error:", err.Error()) + } + case error: + switch err { + case ErrRetryRecords: + p.retries <- retries + p.Log("Received error:", err.Error()) + default: + p.Log("Received error:", err.Error()) + } + case awserr.Error: + switch err.Code() { + case kinesis.ErrCodeProvisionedThroughputExceededException: + p.Stats.AddProvisionedThroughputExceeded(1) + default: + p.Log("Received AWS error:", err.Error()) + } + default: + p.Log("Received unknown error:", err.Error()) + } + } + }() + } + } + }() +} + +func (p *Producer) SendWithContext(ctx context.Context, msg *message.Message) error { + select { + case p.messages <- msg: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (p *Producer) Send(msg *message.Message) error { + return p.SendWithContext(context.TODO(), msg) +} + +func (p *Producer) TryToSend(msg *message.Message) error { + ctx, cancel := context.WithTimeout(context.TODO(), 0*time.Second) + defer cancel() + return p.SendWithContext(ctx, msg) +} diff --git a/producer/stats.go b/producer/stats.go new file mode 100644 index 0000000..7047786 --- /dev/null +++ b/producer/stats.go @@ -0,0 +1,79 @@ +package producer + +import ( + "time" +) + +// StatsCollector allows for a collector to collect various metrics produced by +// the Kinetic producer library. This was really built with rcrowley/go-metrics +// in mind. +type StatsCollector interface { + // histograms give us the count, sum, min, max, mean, percentiles, + // standard deviation, and variance of the data. these metrics should + // give us the total number (sum) of messages sent, failed, and dropped, + // as well as the average (mean) batch size. + AddSentSample(int) + AddFailedSample(int) + AddDroppedSample(int) + AddBatchSizeSample(int) + + // meters give us the count and rate of the data. these metrics should + // give us the average number of times: + // - PutRecords was called per second + // - ProvisionedThroughputExceeded was received per second + // - PutRecords timed out per second + AddPutRecordsCalled(int) + AddProvisionedThroughputExceeded(int) + AddPutRecordsTimeout(int) + + // timers give us the count, sum, min, max, mean, percentiles, standard + // deviation, variance, as well as the rate of the data. + // TODO: describe these metrics better + AddPutRecordsTime(time.Duration) + AddPutRecordsBuildTime(time.Duration) + AddPutRecordsSendTime(time.Duration) +} + +// NilStatsCollector is a stats listener that ignores all metrics. +type NilStatsCollector struct{} + +// AddSentSample records a count of the number of messages sent to AWS Kinesis +// by the producer. +func (l *NilStatsCollector) AddSentSample(int) {} + +// AddFailedSample records a count of the number of messages that failed to be +// sent to AWS Kinesis by the producer. +func (l *NilStatsCollector) AddFailedSample(int) {} + +// AddDroppedSample records a count of the number of messages dropped by the +// application after multiple failures. +func (l *NilStatsCollector) AddDroppedSample(int) {} + +// AddBatchSizeSample records a count of the number of messages attempted by +// PutRecords in the producer. +func (l *NilStatsCollector) AddBatchSizeSample(int) {} + +// AddPutRecordsCalled records the number of times the PutRecords API was called +// by the producer. +func (l *NilStatsCollector) AddPutRecordsCalled(int) {} + +// AddProvisionedThroughputExceeded records the number of times the PutRecords +// API returned a ErrCodeProvisionedThroughputExceededException by the producer. +func (l *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} + +// AddPutRecordsTimeout records the number of times the PutRecords API timed out +// on the HTTP level. This is influenced by the WithHTTPClientTimeout +// configuration. +func (l *NilStatsCollector) AddPutRecordsTimeout(int) {} + +// AddPutRecordsTime records the duration that the PutRecords API request took. +// Only the times of successful calls are measured. +func (l *NilStatsCollector) AddPutRecordsTime(time.Duration) {} + +// AddPutRecordsBuildTime records the duration that it took to build the +// PutRecords API request payload. +func (l *NilStatsCollector) AddPutRecordsBuildTime(time.Duration) {} + +// AddPutRecordsSendTime records the duration that it took to send the +// PutRecords API request payload. +func (l *NilStatsCollector) AddPutRecordsSendTime(time.Duration) {} From 0d2750134d667d601cce1ffcdca22c22121ea31a Mon Sep 17 00:00:00 2001 From: Tommy Wang Date: Thu, 6 Apr 2017 00:27:01 -0500 Subject: [PATCH 010/109] lint & build fixes --- listener/listener.go | 2 +- listener/listener_test.go | 2 +- producer/config.go | 4 ++-- producer/kinesiswriter/config.go | 2 ++ producer/producer.go | 7 +++++++ 5 files changed, 13 insertions(+), 4 deletions(-) diff --git a/listener/listener.go b/listener/listener.go index c666cd8..e9a28ca 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -593,7 +593,7 @@ func (l *Listener) consume(ctx context.Context) { if !ok { break stop } - _, err = l.getRecords(l.batchSize) + _, err := l.getRecords(l.batchSize) if err != nil { switch err := err.(type) { case net.Error: diff --git a/listener/listener_test.go b/listener/listener_test.go index b702a28..a42bb2a 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -226,7 +226,7 @@ func TestListener(t *testing.T) { secs := []float64{} for i := 1; i <= 6; i++ { start := time.Now() - l.fetchBatch(1) + l.getRecords(1) secs = append(secs, time.Since(start).Seconds()) } elapsed := time.Since(start).Seconds() diff --git a/producer/config.go b/producer/config.go index 979665c..dabb462 100644 --- a/producer/config.go +++ b/producer/config.go @@ -9,13 +9,13 @@ import ( "github.com/rewardStyle/kinetic/logging" ) -// Config is used to configure a Producer instance +// Config is used to configure a Producer instance. type Config struct { *kinetic.AwsOptions *producerOptions } -// NewConfig creates a new instance of Config +// NewConfig creates a new instance of Config. func NewConfig() *Config { return &Config{ AwsOptions: kinetic.DefaultAwsOptions(), diff --git a/producer/kinesiswriter/config.go b/producer/kinesiswriter/config.go index ff972d8..a6a58e2 100644 --- a/producer/kinesiswriter/config.go +++ b/producer/kinesiswriter/config.go @@ -2,10 +2,12 @@ package kinesiswriter import () +// Config is used to configure a Kinesis Writer instance. type Config struct { *kinesisWriterOptions } +// NewConfig creates a new instance of Config. func NewConfig() *Config { return &Config{ kinesisWriterOptions: &kinesisWriterOptions{}, diff --git a/producer/producer.go b/producer/producer.go index a825ba0..ce6c239 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -227,6 +227,8 @@ func (p *Producer) produce(ctx context.Context) { }() } +// SendWithContext sends a message to the stream. Cancellation supported +// through contexts. func (p *Producer) SendWithContext(ctx context.Context, msg *message.Message) error { select { case p.messages <- msg: @@ -236,10 +238,15 @@ func (p *Producer) SendWithContext(ctx context.Context, msg *message.Message) er } } +// Send a message to the stream, waiting on the message to be put into the +// channel. func (p *Producer) Send(msg *message.Message) error { return p.SendWithContext(context.TODO(), msg) } +// TryToSend will attempt to send a message to the stream if the channel has +// capacity for a message, or will immediately return with an error if the +// channel is full. func (p *Producer) TryToSend(msg *message.Message) error { ctx, cancel := context.WithTimeout(context.TODO(), 0*time.Second) defer cancel() From 29483f6b7fc93d43733b47c5e61fade8df72b8c9 Mon Sep 17 00:00:00 2001 From: Tommy Wang Date: Fri, 7 Apr 2017 15:58:54 -0500 Subject: [PATCH 011/109] producer refactor, wip --- listener/listener.go | 157 +++++++------ listener/listener_test.go | 1 - message/message.go | 54 +++-- producer/config.go | 31 ++- producer/kinesiswriter/kinesiswriter.go | 21 +- producer/producer.go | 290 ++++++++++++++++-------- producer/stats.go | 19 +- 7 files changed, 368 insertions(+), 205 deletions(-) diff --git a/listener/listener.go b/listener/listener.go index e9a28ca..70f0cb3 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -286,7 +286,6 @@ func (l *Listener) throttle(sem chan Empty) { // getRecords calls GetRecords and delivers each record into the messages // channel. -// TODO: Convert timeout implementation to use context.Context // FIXME: Need to investigate that the timeout implementation doesn't result in // an fd leak. Since we call Read on the HTTPResonse.Body in a select with a // timeout channel, we do prevent ourself from blocking. Once we timeout, we @@ -296,7 +295,7 @@ func (l *Listener) throttle(sem chan Empty) { // down the TCP connection. Worst case scenario is that our client Timeout // eventually fires and closes the socket, but this can be susceptible to FD // exhaustion. -func (l *Listener) getRecords(size int) (int, error) { +func (l *Listener) getRecords(batchSize int) (int, error) { l.ensureClient() if err := l.ensureShardIterator(); err != nil { return 0, err @@ -310,7 +309,7 @@ func (l *Listener) getRecords(size int) (int, error) { var startUnmarshalTime time.Time start := time.Now() req, resp := l.client.GetRecordsRequest(&kinesis.GetRecordsInput{ - Limit: aws.Int64(int64(size)), + Limit: aws.Int64(int64(batchSize)), ShardIterator: aws.String(l.nextShardIterator), }) @@ -381,7 +380,7 @@ func (l *Listener) getRecords(size int) (int, error) { // the Read() timer.Reset(l.getRecordsReadTimeout) n, err = result.n, result.err - l.Log(fmt.Sprintf("DEBUG: read %d bytes, took %v", n, time.Since(readStart))) + l.Log(fmt.Sprintf("GetRecords read %d bytes, took %v", n, time.Since(readStart))) case <-timer.C: // If we timeout, we return an error // that will unblock ioutil.ReadAll(). @@ -389,8 +388,13 @@ func (l *Listener) getRecords(size int) (int, error) { // to return an error. This error will // propogate to the original req.Send() // call (below) - l.Log(fmt.Sprintf("DEBUG: read timed out after %v", time.Since(readStart))) + l.Log(fmt.Sprintf("GetRecords read timed out after %v", time.Since(readStart))) err = ErrTimeoutReadResponseBody + case <-l.pipeOfDeath: + // The pipe of death will abort any pending + // reads on a GetRecords call. + l.Log(fmt.Sprintf("GetRecords received pipe of death after %v", time.Since(readStart))) + err = ErrPipeOfDeath } return }, @@ -425,18 +429,28 @@ func (l *Listener) getRecords(size int) (int, error) { l.Stats.AddBatchSizeSample(len(resp.Records)) for _, record := range resp.Records { if record != nil { - delivered++ - l.messages <- &message.Message{Record: *record} - l.Stats.AddConsumedSample(1) - } - if record.SequenceNumber != nil { - // We can safely ignore if this call returns - // error, as if we somehow receive an empty - // sequence number from AWS, we will simply not - // set it. At worst, this causes us to - // reprocess this record if we happen to refresh - // the iterator. - l.setSequenceNumber(*record.SequenceNumber) + // Allow (only) a pipeOfDeath to trigger an instance + // shutdown of the loop to deliver messages. Otherwise, + // a normal cancellation will not prevent getRecords + // from completing the delivery of the current batch of + // records. + select { + case l.messages <- message.FromKinesisRecord(record): + delivered++ + l.Stats.AddConsumedSample(1) + if record.SequenceNumber != nil { + // We can safely ignore if this call returns + // error, as if we somehow receive an empty + // sequence number from AWS, we will simply not + // set it. At worst, this causes us to + // reprocess this record if we happen to refresh + // the iterator. + l.setSequenceNumber(*record.SequenceNumber) + } + case <-l.pipeOfDeath: + l.Log(fmt.Sprintf("getRecords received pipe of death while delivering messages, %d delivered, ~%d dropped", delivered, len(resp.Records)-delivered)) + return delivered, ErrPipeOfDeath + } } } if resp.NextShardIterator != nil { @@ -463,24 +477,19 @@ func (l *Listener) getRecords(size int) (int, error) { return delivered, nil } -// blockConsumers will set consuming to true if there is not already another -// consume loop running. -func (l *Listener) blockConsumers() bool { +// startConsuming will initialize the consumer and set consuming to true if +// there is not already another consume loop running. +func (l *Listener) startConsuming() bool { l.consumingMu.Lock() defer l.consumingMu.Unlock() if !l.consuming { l.consuming = true + l.messages = make(chan *message.Message, l.batchSize) return true } return false } -// startConsuming handles any initialization needed in preparation to start -// consuming. -func (l *Listener) startConsuming() { - l.messages = make(chan *message.Message, l.batchSize) -} - // shouldConsume is a convenience function that allows functions to break their // loops if the context receives a cancellation. func (l *Listener) shouldConsume(ctx context.Context) (bool, error) { @@ -496,36 +505,29 @@ func (l *Listener) shouldConsume(ctx context.Context) (bool, error) { // stopConsuming handles any cleanup after a consuming has stopped. func (l *Listener) stopConsuming() { - close(l.messages) -} - -// allowConsumers allows consuming. Called after blockConsumers to release the -// lock on consuming. -func (l *Listener) allowConsumers() { l.consumingMu.Lock() defer l.consumingMu.Unlock() + if l.messages != nil { + close(l.messages) + } l.consuming = false } -// IsConsuming returns true while consuming. -func (l *Listener) IsConsuming() bool { - l.consumingMu.Lock() - defer l.consumingMu.Unlock() - return l.consuming -} - // RetrieveWithContext waits for a message from the stream and return the value. // Cancellation supported through contexts. func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, error) { - if !l.blockConsumers() { + if !l.startConsuming() { return nil, ErrAlreadyConsuming } - l.startConsuming() defer func() { l.stopConsuming() - l.allowConsumers() }() for { + // A cancellation or closing the pipe of death will cause + // Retrieve (and related functions) to abort in between + // getRecord calls. Note, that this would only occur when there + // are no new records to retrieve. Otherwise, getRecords will + // be allowed to run to completion and deliver one record. ok, err := l.shouldConsume(ctx) if !ok { return nil, err @@ -573,22 +575,29 @@ func (l *Listener) RetrieveFn(fn MessageFn) error { // consume calls getRecords with configured batch size in a loop until the // listener is stopped. func (l *Listener) consume(ctx context.Context) { - // We need to run blockConsumers & startConsuming to make sure that we - // are okay and ready to start consuming. This is mainly to avoid a - // race condition where Listen() will attempt to read the messages - // channel prior to consume() initializing it. We can then launch a - // goroutine to handle the actual consume operation. - if !l.blockConsumers() { + // We need to run startConsuming to make sure that we are okay and ready + // to start consuming. This is mainly to avoid a race condition where + // Listen() will attempt to read the messages channel prior to consume() + // initializing it. We can then launch a goroutine to handle the actual + // consume operation. + if !l.startConsuming() { return } - l.startConsuming() go func() { defer func() { l.stopConsuming() - l.allowConsumers() }() stop: for { + // The consume loop can be cancelled by a calling the + // cancellation function on the context or by closing + // the pipe of death. Note that in the case of context + // cancellation, the getRecords call below will be + // allowed to complete (as getRecords does not regard + // context cancellation). In the case of cancellation + // by pipe of death, however, the getRecords will + // immediately abort and allow the consume function to + // immediately abort as well. ok, _ := l.shouldConsume(ctx) if !ok { break stop @@ -634,23 +643,37 @@ func (l *Listener) ListenWithContext(ctx context.Context, fn MessageFn) { var wg sync.WaitGroup defer wg.Wait() - for msg := range l.messages { - l.Stats.AddDeliveredSample(1) - l.concurrencySem <- Empty{} - wg.Add(1) - go func(msg *message.Message) { - defer func() { - <-l.concurrencySem - }() - var fnWg sync.WaitGroup - fnWg.Add(1) - start := time.Now() - fn(msg.Value(), &fnWg) - fnWg.Wait() - l.Stats.AddProcessedTime(time.Since(start)) - l.Stats.AddProcessedSample(1) - wg.Done() - }(msg) +stop: + for { + select { + case msg, ok := <-l.messages: + if !ok { + break stop + } + l.Stats.AddDeliveredSample(1) + // For simplicity, did not do the pipe of death here. + // If POD is received, we may deliver a couple more + // messages (especially since select is random in which + // channel is read from). + l.concurrencySem <- Empty{} + wg.Add(1) + go func(msg *message.Message) { + defer func() { + <-l.concurrencySem + }() + var fnWg sync.WaitGroup + fnWg.Add(1) + start := time.Now() + fn(msg.Value(), &fnWg) + fnWg.Wait() + l.Stats.AddProcessedTime(time.Since(start)) + l.Stats.AddProcessedSample(1) + wg.Done() + }(msg) + case <-l.pipeOfDeath: + l.Log("ListenWithContext received pipe of death") + break stop + } } } diff --git a/listener/listener_test.go b/listener/listener_test.go index a42bb2a..85aeb04 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -215,7 +215,6 @@ func TestListener(t *testing.T) { wg.Done() }() <-time.After(10 * time.Millisecond) - So(l.IsConsuming(), ShouldBeTrue) _, err := l.Retrieve() So(err, ShouldEqual, ErrAlreadyConsuming) wg.Wait() diff --git a/message/message.go b/message/message.go index 333b7b9..38c3a76 100644 --- a/message/message.go +++ b/message/message.go @@ -1,37 +1,43 @@ package message -import kinesis "github.com/aws/aws-sdk-go/service/kinesis" +import ( + "time" + + "github.com/aws/aws-sdk-go/service/kinesis" +) // Message represents an item on the Kinesis stream type Message struct { - kinesis.Record -} + // For kinesis.Record + ApproximateArrivalTimestamp *time.Time + Data []byte + PartitionKey *string + SequenceNumber *string -// Init initializes a Message. -// Currently we are ignoring sequenceNumber. -func (k *Message) Init(msg []byte, key string) *Message { - return &Message{ - kinesis.Record{ - Data: msg, - PartitionKey: &key, - }, - } -} + // For kinesis.PutRecordRequestEntry + ExplicitHashKey *string + + // For kinesis.PutRecordResultEntry + ErrorCode *string + ErrorMessage *string + ShardId *string -// SetValue sets the underlying value of the underlying record -func (k *Message) SetValue(value []byte) { - k.Record.Data = value + FailCount int } -// Value gets the underlying value of the underlying record -func (k *Message) Value() []byte { - return k.Record.Data +func FromRecord(record *kinesis.Record) *Message { + return &Message{ + ApproximateArrivalTimestamp: record.ApproximateArrivalTimestamp, + Data: record.Data, + PartitionKey: record.PartitionKey, + SequenceNumber: record.SequenceNumber, + } } -// Key gets the partion key of the message -func (k *Message) Key() []byte { - if k.Record.PartitionKey != nil { - return []byte(*k.PartitionKey) +func (m *Message) MakeRequestEntry() *kinesis.PutRecordsRequestEntry { + return &kinesis.PutRecordsRequestEntry{ + Data: m.Data, + ExplicitHashKey: m.ExplicitHashKey, + PartitionKey: m.PartitionKey, } - return nil } diff --git a/producer/config.go b/producer/config.go index dabb462..34af3fb 100644 --- a/producer/config.go +++ b/producer/config.go @@ -20,12 +20,13 @@ func NewConfig() *Config { return &Config{ AwsOptions: kinetic.DefaultAwsOptions(), producerOptions: &producerOptions{ - batchSize: 500, - batchTimeout: 1 * time.Second, - concurrency: 1, - queueDepth: 500, - LogLevel: logging.LogOff, - Stats: &NilStatsCollector{}, + batchSize: 500, + batchTimeout: 1 * time.Second, + queueDepth: 500, + maxRetryAttempts: 10, + concurrency: 1, + LogLevel: logging.LogOff, + Stats: &NilStatsCollector{}, }, } } @@ -42,18 +43,24 @@ func (c *Config) SetBatchTimeout(timeout time.Duration) { c.batchTimeout = timeout } -// SetConcurrency controls the number of outstanding PutRecords calls may be -// active at a time. -func (c *Config) SetConcurrency(concurrency int) { - c.concurrency = concurrency -} - // SetQueueDepth controls the number of messages that can be in the channel // to be processed by produce at a given time. func (c *Config) SetQueueDepth(queueDepth int) { c.queueDepth = queueDepth } +// SetMaxRetryAttempts controls the number of times a message can be retried +// before it is discarded. +func (c *Config) SetMaxRetryAttempts(attempts int) { + c.maxRetryAttempts = attempts +} + +// SetConcurrency controls the number of outstanding PutRecords calls may be +// active at a time. +func (c *Config) SetConcurrency(concurrency int) { + c.concurrency = concurrency +} + // SetWriter sets the underlying stream writer (Kinesis or Firehose) for the // producer. There can only be a single writer associated with a producer (as // various retry logic/state is not easily shared between multiple writers). If diff --git a/producer/kinesiswriter/kinesiswriter.go b/producer/kinesiswriter/kinesiswriter.go index a13728e..f68f8bf 100644 --- a/producer/kinesiswriter/kinesiswriter.go +++ b/producer/kinesiswriter/kinesiswriter.go @@ -86,6 +86,10 @@ func (w *KinesisWriter) PutRecords(messages []*message.Message) ([]*message.Mess var startBuildTime time.Time start := time.Now() + var records []*kinesis.PutRecordsRequestEntry + for _, msg := range messages { + records = append(records, msg.MakeRequestEntry()) + } req, resp := w.client.PutRecordsRequest(&kinesis.PutRecordsInput{}) req.Handlers.Build.PushFront(func(r *request.Request) { @@ -111,7 +115,7 @@ func (w *KinesisWriter) PutRecords(messages []*message.Message) ([]*message.Mess w.Log("Starting PutRecords Build/Sign request, took", time.Since(start)) w.producer.Stats.AddPutRecordsCalled(1) if err := req.Send(); err != nil { - w.Log("Error putting records:", err) + w.Log("Error putting records:", err.Error()) return nil, err } w.producer.Stats.AddPutRecordsTime(time.Since(start)) @@ -125,9 +129,6 @@ func (w *KinesisWriter) PutRecords(messages []*message.Message) ([]*message.Mess attempted := len(messages) failed := int(aws.Int64Value(resp.FailedRecordCount)) sent := attempted - failed - w.producer.Stats.AddBatchSizeSample(len(messages)) - w.producer.Stats.AddSentSample(sent) - w.producer.Stats.AddFailedSample(failed) w.Log(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) var retries []*message.Message @@ -135,8 +136,18 @@ func (w *KinesisWriter) PutRecords(messages []*message.Message) ([]*message.Mess for idx, record := range resp.Records { if record.SequenceNumber != nil && record.ShardId != nil { // TODO: per-shard metrics + message[idx].SequenceNumber = record.SequenceNumber + message[idx].ShardId = record.ShardId } else { - // TODO metrics on failure rates + switch record.ErrorCode { + case kinesis.ErrCodeProvisionedThroughputExceededException: + w.producer.Stats.AddProvisionedThroughputExceeded(1) + default: + w.Log("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(*record.ErrorMessage)) + } + message[idx].ErrorCode = record.ErrorCode + message[idx].ErrorMessage = record.ErrorMessage + message[idx].FailCount++ retries = append(retries, messages[idx]) } } diff --git a/producer/producer.go b/producer/producer.go index ce6c239..8f3064b 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -38,6 +38,16 @@ var ( // with a producer. ErrProducerAlreadyAssociated = errors.New("StreamWriter already associated with a producer") + // ErrBatchTimeout is returned by getBatch whenever the batchTimeout + // elapses prior to receiving batchSize messages. This is *not* an + // error in the sense of a failure, but is used to distinguish the + // reason getBatch has exited. + ErrBatchTimeout = errors.New("A timeout has occurred before batch has reached optimal size") + + // ErrProducerShutdown is returend by getBatch whenever both the message + // and retry channel have been closed. + ErrProducerShutdown = errors.New("The producer has shut down") + // ErrPipeOfDeath returns when the pipe of death is closed. ErrPipeOfDeath = errors.New("Received pipe of death") ) @@ -46,11 +56,12 @@ var ( type Empty struct{} type producerOptions struct { - batchSize int - batchTimeout time.Duration - concurrency int - queueDepth int - writer StreamWriter + batchSize int + batchTimeout time.Duration + queueDepth int + maxRetryAttempts int + concurrency int + writer StreamWriter LogLevel aws.LogLevelType Stats StatsCollector @@ -61,9 +72,13 @@ type Producer struct { *producerOptions messages chan *message.Message - retries chan []*message.Message + retries chan *message.Message concurrencySem chan Empty pipeOfDeath chan Empty + producerWg sync.WaitGroup + + outstanding int64 + flushCond sync.Cond producing bool producingMu sync.Mutex @@ -99,139 +114,226 @@ func (p *Producer) Log(args ...interface{}) { } } -// blockProducers will set producing to true if there is not already another -// produce loop running. -func (p *Producer) blockProducers() bool { +// startConsuming will initialize the producer and set producing to true if +// there is not already another consume loop running. +func (p *Producer) startProducing() bool { p.producingMu.Lock() defer p.producingMu.Unlock() if !p.producing { p.producing = true + p.messages = make(chan *message.Message, p.queueDepth) + p.retries = make(chan *message.Message) // TODO: should we use a buffered channel? + p.outstanding = 0 + p.flushCond = sync.Cond{L: &sync.Mutex{}} return true } return false } -// startProducing handles any initialization needed in preparation to start -// producing. -func (p *Producer) startProducing() { - p.messages = make(chan *message.Message, p.queueDepth) -} - -// shouldProduce is a convenience function that allows functions to break their -// loops if the context receives a cancellation. -func (p *Producer) shouldProduce(ctx context.Context) (bool, error) { - select { - case <-p.pipeOfDeath: - return false, ErrPipeOfDeath - case <-ctx.Done(): - return false, ctx.Err() - default: - return true, nil - } -} - // stopProducing handles any cleanup after a producing has stopped. func (p *Producer) stopProducing() { - close(p.messages) -} - -// allowProducers allows producing. Called after blockProducers to release the -// lock on producing. -func (p *Producer) allowProducers() { p.producingMu.Lock() defer p.producingMu.Unlock() + if p.messages != nil { + close(p.messages) + } p.producing = false } -// IsProducing returns true while producing. -func (p *Producer) IsProducing() bool { - p.producingMu.Lock() - defer p.producingMu.Unlock() - return p.producing -} - -func (p *Producer) getBatch(ctx context.Context) ([]*message.Message, error) { - ctx, cancel := context.WithTimeout(ctx, p.batchTimeout) - defer cancel() - +// getBatch will retrieve a batch of messages by batchSize and batchTimeout for +// delivery. +func (p *Producer) getBatch() ([]*message.Message, error) { + var err error var batch []*message.Message - select { - case batch = <-p.retries: - default: + var timer <-chan time.Time +stop: + for len(batch) <= p.batchSize { + select { + // Using the select, retry messages will interleave with new + // messages. This is preferable to putting the messages at the + // end of the channel as it minimizes the delay in the delivery + // of retry messages. + case msg, ok := <-p.retries: + if !ok { + p.retries = nil + } else { + batch = append(batch, msg) + if timer != nil { + timer = time.After(p.batchTimeout) + } + } + case msg, ok := <-p.messages: + if !ok { + p.messages = nil + } else { + batch = append(batch, msg) + if timer != nil { + timer = time.After(p.batchTimeout) + } + } + case <-timer: + err = ErrBatchTimeout + break stop + case <-p.pipeOfDeath: + return nil, ErrPipeOfDeath + } + if p.messages == nil && p.retries == nil { + err = ErrProducerShutdown + break stop + } } + p.Stats.AddBatchSizeSample(len(batch)) + return batch, err +} - for len(batch) < p.batchSize { - select { - case msg := <-p.messages: - batch = append(batch, msg) - case <-ctx.Done(): - return batch, ctx.Err() +func (p *Producer) dispatchBatch(batch []*message.Message) { + defer p.flushCond.Signal() +stop: + for { + retries, err := p.writer.PutRecords(batch) + failed := len(retries) + sent := len(batch) - failed + p.Stats.AddSentSample(sent) + p.Stats.AddFailedSample(failed) + p.decOutstanding(int64(sent)) + // This frees up another dispatchBatch to run to allow drainage + // of the messages / retry queue. This should improve + // throughput as well as prevent a potential deadlock in which + // all batches are blocked on sending retries to the retries + // channel, and thus no batches are allowed to drain the retry + // channel. + <-p.concurrencySem + if err == nil { + break stop + } + switch err := err.(type) { + case net.Error: + if err.Timeout() { + p.Stats.AddPutRecordsTimeout(1) + p.Log("Received net error:", err.Error()) + } else { + p.Log("Received unknown net error:", err.Error()) + } + case error: + switch err { + case ErrRetryRecords: + for _, msg := range retries { + if msg.FailCount < p.maxRetryAttempts { + msg.FailCount++ + select { + case p.retries <- msg: + case <-p.pipeOfDeath: + break stop + } + } else { + p.decOutstanding(1) + p.Stats.AddDroppedSample(1) + } + } + p.Log("Received error:", err.Error()) + default: + p.Log("Received error:", err.Error()) + } + case awserr.Error: + switch err.Code() { + case kinesis.ErrCodeProvisionedThroughputExceededException: + // FIXME: It is not clear to me whether + // PutRecords would ever return a + // ProvisionedThroughputExceeded error. It + // seems that it would instead return a valid + // response in which some or all the records + // within the response will contain an error + // code and error message of + // ProvisionedThroughputExceeded. The current + // assumption is that if we receive an + // ProvisionedThroughputExceeded error, that the + // entire batch should be retried. Note we only + // increment the PutRecord stat, instead of the + // per-message stat. Furthermore, we do not + // increment the FailCount of the messages (as + // the retry mechanism is different). + p.Stats.AddPutRecordsProvisionedThroughputExceeded(1) + default: + p.Log("Received AWS error:", err.Error()) + } + default: + p.Log("Received unknown error:", err.Error()) } } - return batch, nil } -func (p *Producer) produce(ctx context.Context) { - if !p.blockProducers() { +// produce calls the underlying writer's PutRecords implementation to deliver +// batches of messages to the target stream until the producer is stopped. +func (p *Producer) produce() { + if !p.startProducing() { return } - p.startProducing() + p.producerWg.Add(1) go func() { defer func() { p.stopProducing() - p.allowProducers() + p.producerWg.Done() }() stop: for { - ok, _ := p.shouldProduce(ctx) - if !ok { + batch, err := p.getBatch() + // If getBatch aborted due to pipe of death, we will + // immediately exit the loop. + if err == ErrPipeOfDeath { break stop } - batch, _ := p.getBatch(ctx) + // Regardless if getBatch produced an error (as long as + // its not the pipe of death), we will send the messages + // via PutRecords. if len(batch) > 0 { p.concurrencySem <- Empty{} - go func() { - retries, err := p.writer.PutRecords(batch) - <-p.concurrencySem - if err != nil { - switch err := err.(type) { - case net.Error: - if err.Timeout() { - p.Stats.AddPutRecordsTimeout(1) - p.Log("Received net error:", err.Error()) - } else { - p.Log("Received unknown net error:", err.Error()) - } - case error: - switch err { - case ErrRetryRecords: - p.retries <- retries - p.Log("Received error:", err.Error()) - default: - p.Log("Received error:", err.Error()) - } - case awserr.Error: - switch err.Code() { - case kinesis.ErrCodeProvisionedThroughputExceededException: - p.Stats.AddProvisionedThroughputExceeded(1) - default: - p.Log("Received AWS error:", err.Error()) - } - default: - p.Log("Received unknown error:", err.Error()) - } - } - }() + go p.dispatchBatch(batch) + } + // If we exited getBatch earlier with a + // ErrProducerShutdown we shut down the producer. + if err == ErrProducerShutdown { + break stop } } }() } +// incOutstanding increments the number of outstanding messages that are to be +// delivered. +func (p *Producer) incOutstanding(i int64) { + p.flushCond.L.Lock() + defer p.flushCond.L.Unlock() + p.outstanding += i +} + +// decOutstanding decrements the number of outstanding messages that are to be +// delivered. +func (p *Producer) decOutstanding(i int64) { + p.flushCond.L.Lock() + defer p.flushCond.L.Unlock() + p.outstanding -= i +} + +// Close shuts down the producer, waiting for all outstanding messages and retries +// to flush. +func (p *Producer) Close() { + close(p.messages) + p.flushCond.L.Lock() + for p.outstanding != 0 { + p.flushCond.Wait() + } + close(p.retries) + p.flushCond.L.Unlock() + p.producerWg.Wait() +} + // SendWithContext sends a message to the stream. Cancellation supported // through contexts. func (p *Producer) SendWithContext(ctx context.Context, msg *message.Message) error { + p.produce() select { case p.messages <- msg: + p.incOutstanding(1) return nil case <-ctx.Done(): return ctx.Err() diff --git a/producer/stats.go b/producer/stats.go index 7047786..1cc1cbb 100644 --- a/producer/stats.go +++ b/producer/stats.go @@ -12,6 +12,8 @@ type StatsCollector interface { // standard deviation, and variance of the data. these metrics should // give us the total number (sum) of messages sent, failed, and dropped, // as well as the average (mean) batch size. + + // for producer AddSentSample(int) AddFailedSample(int) AddDroppedSample(int) @@ -19,9 +21,15 @@ type StatsCollector interface { // meters give us the count and rate of the data. these metrics should // give us the average number of times: - // - PutRecords was called per second // - ProvisionedThroughputExceeded was received per second + // - PutRecords was called per second + // - PutRecordProvisionedThroughputExceeded was received per second // - PutRecords timed out per second + + // for producer + AddPutRecordsProvisionedThroughputExceeded(int) + + // for kinesiswriter AddPutRecordsCalled(int) AddProvisionedThroughputExceeded(int) AddPutRecordsTimeout(int) @@ -29,6 +37,8 @@ type StatsCollector interface { // timers give us the count, sum, min, max, mean, percentiles, standard // deviation, variance, as well as the rate of the data. // TODO: describe these metrics better + + // for kinesis writer AddPutRecordsTime(time.Duration) AddPutRecordsBuildTime(time.Duration) AddPutRecordsSendTime(time.Duration) @@ -53,12 +63,17 @@ func (l *NilStatsCollector) AddDroppedSample(int) {} // PutRecords in the producer. func (l *NilStatsCollector) AddBatchSizeSample(int) {} +// AddProvisionedThroughputExceeded records the number of times the PutRecords +// API returned a ErrCodeProvisionedThroughputExceededException by the producer. +func (l *NilStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} + // AddPutRecordsCalled records the number of times the PutRecords API was called // by the producer. func (l *NilStatsCollector) AddPutRecordsCalled(int) {} // AddProvisionedThroughputExceeded records the number of times the PutRecords -// API returned a ErrCodeProvisionedThroughputExceededException by the producer. +// API response contained a record which contained an +// ErrCodeProvisionedThroughputExceededException error. func (l *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} // AddPutRecordsTimeout records the number of times the PutRecords API timed out From 69c785ebcf54a9df85a39a5527d4f57cbe16c491 Mon Sep 17 00:00:00 2001 From: Tommy Wang Date: Fri, 5 May 2017 09:43:17 -0500 Subject: [PATCH 012/109] producer (kinesis + firehose) --- config.go | 73 +---- config/config.go | 66 ++++ errs/errors.go | 76 +++++ kinetic.go | 38 +-- listener/config.go | 23 +- listener/config_test.go | 28 +- listener/listener.go | 263 ++++++++++------ listener/listener_test.go | 18 +- {utils => listener}/readcloser.go | 2 +- listener/stats.go | 70 ++--- logging/logging.go | 49 +++ message/message.go | 17 +- producer/config.go | 36 ++- producer/config_test.go | 122 +++++++ producer/firehosewriter.go | 139 ++++++++ producer/firehosewriter/config.go | 1 - producer/firehosewriter/firehosewriter.go | 1 - producer/{kinesiswriter => }/kinesiswriter.go | 86 ++--- producer/kinesiswriter/config.go | 15 - producer/producer.go | 297 ++++++++---------- producer/producer_test.go | 123 ++++++++ producer/stats.go | 72 ++--- 22 files changed, 1052 insertions(+), 563 deletions(-) create mode 100644 config/config.go create mode 100644 errs/errors.go rename {utils => listener}/readcloser.go (97%) create mode 100644 producer/config_test.go create mode 100644 producer/firehosewriter.go delete mode 100644 producer/firehosewriter/config.go delete mode 100644 producer/firehosewriter/firehosewriter.go rename producer/{kinesiswriter => }/kinesiswriter.go (53%) delete mode 100644 producer/kinesiswriter/config.go create mode 100644 producer/producer_test.go diff --git a/config.go b/config.go index 1ddf48a..f6fb166 100644 --- a/config.go +++ b/config.go @@ -1,83 +1,24 @@ package kinetic import ( - "net/http" - "time" - "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/aws/credentials" - "github.com/aws/aws-sdk-go/aws/session" -) - -// AwsOptions helps configure an aws.Config and session.Session -type AwsOptions struct { - AwsConfig *aws.Config -} - -// DefaultAwsOptions initializes the default aws.Config struct -func DefaultAwsOptions() *AwsOptions { - return &AwsOptions{ - AwsConfig: aws.NewConfig().WithHTTPClient( - &http.Client{ - Timeout: 10 * time.Minute, - }, - ), - } -} - -// SetCredentials configures AWS credentials. -func (c *AwsOptions) SetCredentials(accessKey, secretKey, sessionToken string) { - c.AwsConfig.WithCredentials( - credentials.NewStaticCredentials(accessKey, secretKey, sessionToken), - ) -} - -// SetRegion configures the AWS region. -func (c *AwsOptions) SetRegion(region string) { - c.AwsConfig.WithRegion(region) -} -// SetEndpoint sets the endpoint to be used by aws-sdk-go. -func (c *AwsOptions) SetEndpoint(endpoint string) { - c.AwsConfig.WithEndpoint(endpoint) -} - -// SetLogger configures the logger for Kinetic and the aws-sdk-go. -func (c *AwsOptions) SetLogger(logger aws.Logger) { - c.AwsConfig.WithLogger(logger) -} - -// SetLogLevel configures the log levels for the SDK. -func (c *AwsOptions) SetLogLevel(logLevel aws.LogLevelType) { - c.AwsConfig.WithLogLevel(logLevel & 0xffff) -} - -// SetHTTPClientTimeout configures the HTTP timeout for the SDK. -func (c *AwsOptions) SetHTTPClientTimeout(timeout time.Duration) { - c.AwsConfig.WithHTTPClient(&http.Client{ - Timeout: timeout, - }) -} - -// GetSession creates an instance of the session.Session to be used when creating service -// clients in aws-sdk-go. -func (c *AwsOptions) GetSession() (*session.Session, error) { - return session.NewSession(c.AwsConfig) -} + "github.com/rewardStyle/kinetic/config" +) // Config is used to configure a Kinetic instance type Config struct { - *AwsOptions + *config.AwsOptions *kineticOptions + LogLevel aws.LogLevelType } // NewConfig creates a new instance of Config func NewConfig() *Config { return &Config{ - AwsOptions: DefaultAwsOptions(), - kineticOptions: &kineticOptions{ - LogLevel: aws.LogOff, - }, + AwsOptions: config.DefaultAwsOptions(), + kineticOptions: &kineticOptions{}, + LogLevel: aws.LogOff, } } diff --git a/config/config.go b/config/config.go new file mode 100644 index 0000000..6bf3d11 --- /dev/null +++ b/config/config.go @@ -0,0 +1,66 @@ +package config + +import ( + "net/http" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/credentials" + "github.com/aws/aws-sdk-go/aws/session" +) + +// AwsOptions helps configure an aws.Config and session.Session +type AwsOptions struct { + AwsConfig *aws.Config +} + +// DefaultAwsOptions initializes the default aws.Config struct +func DefaultAwsOptions() *AwsOptions { + return &AwsOptions{ + AwsConfig: aws.NewConfig().WithHTTPClient( + &http.Client{ + Timeout: 10 * time.Minute, + }, + ), + } +} + +// SetCredentials configures AWS credentials. +func (c *AwsOptions) SetCredentials(accessKey, secretKey, sessionToken string) { + c.AwsConfig.WithCredentials( + credentials.NewStaticCredentials(accessKey, secretKey, sessionToken), + ) +} + +// SetRegion configures the AWS region. +func (c *AwsOptions) SetRegion(region string) { + c.AwsConfig.WithRegion(region) +} + +// SetEndpoint sets the endpoint to be used by aws-sdk-go. +func (c *AwsOptions) SetEndpoint(endpoint string) { + c.AwsConfig.WithEndpoint(endpoint) +} + +// SetLogger configures the logger for Kinetic and the aws-sdk-go. +func (c *AwsOptions) SetLogger(logger aws.Logger) { + c.AwsConfig.WithLogger(logger) +} + +// SetLogLevel configures the log levels for the SDK. +func (c *AwsOptions) SetLogLevel(logLevel aws.LogLevelType) { + c.AwsConfig.WithLogLevel(logLevel & 0xffff) +} + +// SetHTTPClientTimeout configures the HTTP timeout for the SDK. +func (c *AwsOptions) SetHTTPClientTimeout(timeout time.Duration) { + c.AwsConfig.WithHTTPClient(&http.Client{ + Timeout: timeout, + }) +} + +// GetSession creates an instance of the session.Session to be used when creating service +// clients in aws-sdk-go. +func (c *AwsOptions) GetSession() (*session.Session, error) { + return session.NewSession(c.AwsConfig) +} diff --git a/errs/errors.go b/errs/errors.go new file mode 100644 index 0000000..c9e49db --- /dev/null +++ b/errs/errors.go @@ -0,0 +1,76 @@ +package errs + +import ( + "errors" +) + +var ( + // ErrNilDescribeStreamResponse is an error returned by GetShards when + // the DescribeStream request returns a nil response + ErrNilDescribeStreamResponse = errors.New("DescribeStream returned a nil response") + + // ErrNilStreamDescription is an error returned by GetShards when the + // DescribeStream request returns a response with a nil + // StreamDescription + ErrNilStreamDescription = errors.New("DescribeStream returned a nil StreamDescription") + + // ErrPipeOfDeath returns when the pipe of death is closed. + ErrPipeOfDeath = errors.New("Received pipe of death") +) + +var ( + // ErrAlreadyConsuming is returned when attempting to consume when the + // Listener is already consuming. May be returned by + // Retrieve/RetrieveFn. + ErrAlreadyConsuming = errors.New("Listener already consuming. Only one Listen, Retrieve, or RetrieveFn may be active at a time") + + // ErrEmptySequenceNumber is returned when attempting to set an empty + // sequence number. + ErrEmptySequenceNumber = errors.New("Attempted to set sequence number with empty value") + + // ErrEmptyShardIterator is returned when attempting to set an empty + // sequence number. + ErrEmptyShardIterator = errors.New("Attempted to set shard iterator with empty value") + + // ErrNilGetShardIteratorResponse is returned when the GetShardIterator + // call returns a nil response. + ErrNilGetShardIteratorResponse = errors.New("GetShardIteratore returned a nil response") + + // ErrNilShardIterator is returned when the GetShardIterator call + // returns a nil shard iterator. + ErrNilShardIterator = errors.New("GetShardIterator returned a nil ShardIterator") + + // ErrNilGetRecordsResponse is returned when the GetRecords calls + // returns a nil response. + ErrNilGetRecordsResponse = errors.New("GetRecords returned an nil response") + + // ErrTimeoutReadResponseBody is returned when a timeout occurs while + // reading the GetRecords response body. + ErrTimeoutReadResponseBody = errors.New("Timeout while reading response body") +) + +var ( + // ErrRetryRecords is returned when the PutRecords calls requires some + // records of the batch to be retried. This failure is considered part + // of normal behavior of the Kinesis stream. + ErrRetryRecords = errors.New("PutRecords requires retry of some records in batch") + + // ErrNilProducer is returned by a StreamWriter when it has not been + // correctly associated with a Producer. + ErrNilProducer = errors.New("StreamWriter not associated with a producer") + + // ErrProducerAlreadyAssociated is returned by a StreamWriter attempting + // to associate it with a Producer when it already has an association + // with a producer. + ErrProducerAlreadyAssociated = errors.New("StreamWriter already associated with a producer") +) + +var ( + // ErrNilPutRecordsResponse is returned when the PutRecords call returns + // a nil response. + ErrNilPutRecordsResponse = errors.New("PutRecords returned a nil response") + + // ErrNilFailedRecordCount is returned when the PutRecords call returns + // a nil FailedRecordCount. + ErrNilFailedRecordCount = errors.New("GetFailedRecordCount returned a nil FailedRecordCount") +) diff --git a/kinetic.go b/kinetic.go index 814eb27..957cdce 100644 --- a/kinetic.go +++ b/kinetic.go @@ -2,7 +2,6 @@ package kinetic import ( "context" - "errors" "sync" "time" @@ -14,28 +13,18 @@ import ( "github.com/aws/aws-sdk-go/service/kinesis" "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" + "github.com/rewardStyle/kinetic/errs" "github.com/rewardStyle/kinetic/logging" ) -var ( - // ErrNilDescribeStreamResponse is an error returned by GetShards when - // the DescribeStream request returns a nil response - ErrNilDescribeStreamResponse = errors.New("DescribeStream returned a nil response") - - // ErrNilStreamDescription is an error returned by GetShards when the - // DescribeStream request returns a response with a nil - // StreamDescription - ErrNilStreamDescription = errors.New("DescribeStream returned a nil StreamDescription") -) - type kineticOptions struct { - LogLevel aws.LogLevelType } // Kinetic represents a kinesis and firehose client and provides some utility // methods for interacting with the AWS services. type Kinetic struct { *kineticOptions + *logging.LogHelper clientMu sync.Mutex fclient firehoseiface.FirehoseAPI @@ -53,17 +42,14 @@ func New(fn func(*Config)) (*Kinetic, error) { } return &Kinetic{ kineticOptions: config.kineticOptions, - Session: session, + LogHelper: &logging.LogHelper{ + LogLevel: config.LogLevel, + Logger: session.Config.Logger, + }, + Session: session, }, nil } -// Log logs a message if LogDebug is set. -func (k *Kinetic) Log(args ...interface{}) { - if k.LogLevel.Matches(logging.LogDebug) { - k.Session.Config.Logger.Log(args...) - } -} - func (k *Kinetic) ensureKinesisClient() { k.clientMu.Lock() defer k.clientMu.Unlock() @@ -80,7 +66,7 @@ func (k *Kinetic) CreateStream(stream string, shards int) error { ShardCount: aws.Int64(int64(shards)), }) if err != nil { - k.Log("Error creating kinesis stream:", err) + k.LogError("Error creating kinesis stream:", err) } return err } @@ -101,7 +87,7 @@ func (k *Kinetic) DeleteStream(stream string) error { StreamName: aws.String(stream), }) if err != nil { - k.Log("Error deleting kinesis stream:", err) + k.LogError("Error deleting kinesis stream:", err) } return err } @@ -142,14 +128,14 @@ func (k *Kinetic) GetShards(stream string) ([]string, error) { StreamName: aws.String(stream), }) if err != nil { - k.Log("Error describing kinesis stream:", err) + k.LogError("Error describing kinesis stream:", err) return nil, err } if resp == nil { - return nil, ErrNilDescribeStreamResponse + return nil, errs.ErrNilDescribeStreamResponse } if resp.StreamDescription == nil { - return nil, ErrNilStreamDescription + return nil, errs.ErrNilStreamDescription } var shards []string for _, shard := range resp.StreamDescription.Shards { diff --git a/listener/config.go b/listener/config.go index acce032..0b0394f 100644 --- a/listener/config.go +++ b/listener/config.go @@ -5,20 +5,21 @@ import ( "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/config" "github.com/rewardStyle/kinetic/logging" ) // Config is used to configure a Listener instance type Config struct { - *kinetic.AwsOptions + *config.AwsOptions *listenerOptions + LogLevel aws.LogLevelType } // NewConfig creates a new instance of Config func NewConfig(stream, shard string) *Config { return &Config{ - AwsOptions: kinetic.DefaultAwsOptions(), + AwsOptions: config.DefaultAwsOptions(), listenerOptions: &listenerOptions{ stream: stream, shard: shard, @@ -26,12 +27,18 @@ func NewConfig(stream, shard string) *Config { concurrency: 10000, shardIterator: NewShardIterator(), getRecordsReadTimeout: 1 * time.Second, - LogLevel: logging.LogOff, - Stats: &NilStatsCollector{}, + Stats: &NilStatsCollector{}, }, + LogLevel: logging.LogOff, } } +// SetAwsConfig configures the AWS Config used to create Sessions (and therefore +// kinesis clients). +func (c *Config) SetAwsConfig(config *aws.Config) { + c.AwsConfig = config +} + // SetBatchSize configures the batch size of the GetRecords call. func (c *Config) SetBatchSize(batchSize int) { c.batchSize = batchSize @@ -65,9 +72,3 @@ func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { func (c *Config) SetStatsCollector(stats StatsCollector) { c.Stats = stats } - -// FromKinetic configures the session from Kinetic. -func (c *Config) FromKinetic(k *kinetic.Kinetic) *Config { - c.AwsConfig = k.Session.Config - return c -} diff --git a/listener/config_test.go b/listener/config_test.go index 82953ba..c7168e2 100644 --- a/listener/config_test.go +++ b/listener/config_test.go @@ -15,18 +15,18 @@ import ( type DebugStatsCollector struct{} -func (l *DebugStatsCollector) AddConsumedSample(int) {} -func (l *DebugStatsCollector) AddDeliveredSample(int) {} -func (l *DebugStatsCollector) AddProcessedSample(int) {} -func (l *DebugStatsCollector) AddBatchSizeSample(int) {} -func (l *DebugStatsCollector) AddGetRecordsCalled(int) {} -func (l *DebugStatsCollector) AddProvisionedThroughputExceeded(int) {} -func (l *DebugStatsCollector) AddGetRecordsTimeout(int) {} -func (l *DebugStatsCollector) AddGetRecordsReadTimeout(int) {} -func (l *DebugStatsCollector) AddProcessedTime(time.Duration) {} -func (l *DebugStatsCollector) AddGetRecordsTime(time.Duration) {} -func (l *DebugStatsCollector) AddGetRecordsReadResponseTime(time.Duration) {} -func (l *DebugStatsCollector) AddGetRecordsUnmarshalTime(time.Duration) {} +func (l *DebugStatsCollector) AddConsumed(int) {} +func (l *DebugStatsCollector) AddDelivered(int) {} +func (l *DebugStatsCollector) AddProcessed(int) {} +func (l *DebugStatsCollector) AddBatchSize(int) {} +func (l *DebugStatsCollector) AddGetRecordsCalled(int) {} +func (l *DebugStatsCollector) AddProvisionedThroughputExceeded(int) {} +func (l *DebugStatsCollector) AddGetRecordsTimeout(int) {} +func (l *DebugStatsCollector) AddGetRecordsReadTimeout(int) {} +func (l *DebugStatsCollector) AddProcessedDuration(time.Duration) {} +func (l *DebugStatsCollector) AddGetRecordsDuration(time.Duration) {} +func (l *DebugStatsCollector) AddGetRecordsReadResponseDuration(time.Duration) {} +func (l *DebugStatsCollector) AddGetRecordsUnmarshalDuration(time.Duration) {} func getSession(config *Config) *session.Session { sess, err := config.GetSession() @@ -65,12 +65,12 @@ func TestNewConfig(t *testing.T) { So(config.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) }) - Convey("check that we can import configuration from kinetic", func() { + Convey("check that we can set the AWS configuration", func() { k, err := kinetic.New(func(c *kinetic.Config) { c.SetEndpoint("bogus-endpoint") }) So(err, ShouldBeNil) - config = config.FromKinetic(k) + config.SetAwsConfig(k.Session.Config) sess := getSession(config) So(aws.StringValue(sess.Config.Endpoint), ShouldEqual, "bogus-endpoint") }) diff --git a/listener/listener.go b/listener/listener.go index 70f0cb3..35a8bcf 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -2,7 +2,6 @@ package listener import ( "context" - "errors" "fmt" "io" "net" @@ -16,43 +15,9 @@ import ( "github.com/aws/aws-sdk-go/service/kinesis" "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" + "github.com/rewardStyle/kinetic/errs" "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" - "github.com/rewardStyle/kinetic/utils" -) - -var ( - // ErrAlreadyConsuming is returned when attempting to consume when the - // Listener is already consuming. May be returned by - // Retrieve/RetrieveFn. - ErrAlreadyConsuming = errors.New("Listener already consuming. Only one Listen, Retrieve, or RetrieveFn may be active at a time") - - // ErrEmptySequenceNumber is returned when attempting to set an empty - // sequence number. - ErrEmptySequenceNumber = errors.New("Attempted to set sequence number with empty value") - - // ErrEmptyShardIterator is returned when attempting to set an empty - // sequence number. - ErrEmptyShardIterator = errors.New("Attempted to set shard iterator with empty value") - - // ErrNilGetShardIteratorResponse is returned when the GetShardIterator - // call returns a nil response. - ErrNilGetShardIteratorResponse = errors.New("GetShardIteratore returned a nil response") - - // ErrNilShardIterator is returned when the GetShardIterator call - // returns a nil shard iterator. - ErrNilShardIterator = errors.New("GetShardIterator returned a nil ShardIterator") - - // ErrNilGetRecordsResponse is returned when the GetRecords calls - // returns a nil response. - ErrNilGetRecordsResponse = errors.New("GetRecords returned an nil response") - - // ErrTimeoutReadResponseBody is returned when a timeout occurs while - // reading the GetRecords response body. - ErrTimeoutReadResponseBody = errors.New("Timeout while reading response body") - - // ErrPipeOfDeath returns when the pipe of death is closed. - ErrPipeOfDeath = errors.New("Received pipe of death") ) // Empty is used a as a dummy type for counting semaphore channels. @@ -134,13 +99,13 @@ type listenerOptions struct { shardIterator *ShardIterator getRecordsReadTimeout time.Duration - LogLevel aws.LogLevelType - Stats StatsCollector + Stats StatsCollector } // Listener polls the Kinesis stream for messages. type Listener struct { *listenerOptions + *logging.LogHelper nextShardIterator string @@ -168,20 +133,17 @@ func NewListener(stream, shard string, fn func(*Config)) (*Listener, error) { } return &Listener{ listenerOptions: config.listenerOptions, - concurrencySem: make(chan Empty, config.concurrency), - throttleSem: make(chan Empty, 5), - pipeOfDeath: make(chan Empty), - Session: session, + LogHelper: &logging.LogHelper{ + LogLevel: config.LogLevel, + Logger: session.Config.Logger, + }, + concurrencySem: make(chan Empty, config.concurrency), + throttleSem: make(chan Empty, 5), + pipeOfDeath: make(chan Empty), + Session: session, }, nil } -// Log a debug message using the AWS SDK logger. -func (l *Listener) Log(args ...interface{}) { - if l.Session.Config.LogLevel.Matches(logging.LogDebug) { - l.Session.Config.Logger.Log(args...) - } -} - // setNextShardIterator sets the nextShardIterator to use when calling // GetRecords. // @@ -190,7 +152,7 @@ func (l *Listener) Log(args ...interface{}) { // to Listen and Retrieve/RetrieveFn can be running at a time. func (l *Listener) setNextShardIterator(shardIterator string) error { if len(shardIterator) == 0 { - return ErrEmptyShardIterator + return errs.ErrEmptyShardIterator } l.nextShardIterator = shardIterator return nil @@ -206,7 +168,7 @@ func (l *Listener) setNextShardIterator(shardIterator string) error { // time. func (l *Listener) setSequenceNumber(sequenceNumber string) error { if len(sequenceNumber) == 0 { - return ErrEmptySequenceNumber + return errs.ErrEmptySequenceNumber } l.shardIterator.AtSequenceNumber(sequenceNumber) return nil @@ -263,14 +225,14 @@ func (l *Listener) ensureShardIterator() error { Timestamp: l.shardIterator.getTimestamp(), }) if err != nil { - l.Log(err) + l.LogError(err) return err } if resp == nil { - return ErrNilGetShardIteratorResponse + return errs.ErrNilGetShardIteratorResponse } if resp.ShardIterator == nil { - return ErrNilShardIterator + return errs.ErrNilShardIterator } return l.setNextShardIterator(*resp.ShardIterator) } @@ -284,6 +246,119 @@ func (l *Listener) throttle(sem chan Empty) { }) } +/* +func (l *Listener) newGetRecords(batchSize int) (int, error) { + l.ensureClient() + if err := l.ensureShardIterator(); err != nil { + return 0, err + } + l.throttle(l.throttleSem) + + var startUnmarshalTime time.Time + start := time.Now() + + // We use the GetRecordsRequest method of creating requests to allow for + // registering custom handlers for better control over the API request. + req, resp := l.client.GetRecordsRequest(&kinesis.GetRecordsInput{ + Limit: aws.Int64(int64(batchSize)), + ShardIterator: aws.String(l.nextShardIterator), + }) + + // If debug is turned on, add some handlers for GetRecords logging + if l.LogLevel.AtLeast(logging.LogDebug) { + req.Handlers.Send.PushBack(func(r *request.Request) { + l.LogDebug("Finished GetRecords Send, took", time.Since(start)) + }) + } + + // Add some profiling timers for metrics + // TODO: Original implementation was able to measure read time and + // unmarshal time separately. Because we no longer hook into the + // ReadCloser's Close method, we can no longer measure them separately. + // Instead, the UnmarshalDuration will measure the total time to read + // and unmarshal the data. We should delete the corresponding stat for + // the read time. + req.Handlers.Unmarshal.PushFront(func(r *request.Request) { + l.LogDebug("Started GetRecords Unmarshal, took", time.Since(start)) + startUnmarshalTime = time.Now() + }) + + req.Handlers.Unmarshal.PushBack(func(r *request.Request) { + l.LogDebug("Finished GetRecords Unmarshal, took", time.Since(start)) + l.Stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) + }) + + // This should replace any WithResponseReadTimeout calls within the AWS + // SDK in customizations.go. + req.ApplyOptions(request.WithResponseReadTimeout(l.getRecordsReadTimeout)) + + // Send the GetRecords request + l.LogDebug("Starting GetRecords Build/Sign request, took", time.Since(start)) + l.Stats.AddGetRecordsCalled(1) + if err := req.Send(); err != nil { + l.LogError("Error getting records:", err) + return 0, err + } + l.Stats.AddGetRecordsDuration(time.Since(start)) + + // Process Records + l.LogDebug(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), l.shard, time.Since(start))) + if resp == nil { + return 0, errs.ErrNilGetRecordsResponse + } + delivered := 0 + l.Stats.AddBatchSize(len(resp.Records)) + for _, record := range resp.Records { + if record != nil { + // Allow (only) a pipeOfDeath to trigger an instance + // shutdown of the loop to deliver messages. Otherwise, + // a normal cancellation will not prevent getRecords + // from completing the delivery of the current batch of + // records. + select { + case l.messages <- message.FromRecord(record): + delivered++ + l.Stats.AddConsumed(1) + if record.SequenceNumber != nil { + // We can safely ignore if this call returns + // error, as if we somehow receive an empty + // sequence number from AWS, we will simply not + // set it. At worst, this causes us to + // reprocess this record if we happen to refresh + // the iterator. + l.setSequenceNumber(*record.SequenceNumber) + } + case <-l.pipeOfDeath: + l.LogInfo(fmt.Sprintf("getRecords received pipe of death while delivering messages, %d delivered, ~%d dropped", delivered, len(resp.Records)-delivered)) + return delivered, errs.ErrPipeOfDeath + } + } + } + if resp.NextShardIterator != nil { + // TODO: According to AWS docs: + // http://docs.aws.amazon.com/sdk-for-go/api/service/kinesis/#GetRecordsOutput + // + // NextShardIterator: The next position in the shard + // from which to start sequentially reading data + // records. If set to null, the shard has been closed + // and the requested iterator will not return any more + // data. + // + // When dealing with streams that will merge or split, + // we need to detect that the shard has closed and + // notify the client library. + // + // TODO: I don't know if we should be ignoring an error returned + // by setShardIterator in case of an empty shard iterator in the + // response. There isn't much we can do, and the best path for + // recovery may be simply to reprocess the batch and see if we + // get a valid NextShardIterator from AWS the next time around. + l.setNextShardIterator(*resp.NextShardIterator) + } + return delivered, nil +} +*/ + // getRecords calls GetRecords and delivers each record into the messages // channel. // FIXME: Need to investigate that the timeout implementation doesn't result in @@ -316,7 +391,7 @@ func (l *Listener) getRecords(batchSize int) (int, error) { // If debug is turned on, add some handlers for GetRecords logging if l.LogLevel.AtLeast(logging.LogDebug) { req.Handlers.Send.PushBack(func(r *request.Request) { - l.Log("Finished GetRecords Send, took", time.Since(start)) + l.LogDebug("Finished GetRecords Send, took", time.Since(start)) }) } @@ -334,14 +409,14 @@ func (l *Listener) getRecords(batchSize int) (int, error) { // that we can implement a timeout mechanism on the Read() call (which // is called by the ioutil.ReadAll() function) req.Handlers.Unmarshal.PushFront(func(r *request.Request) { - l.Log("Started GetRecords Unmarshal, took", time.Since(start)) + l.LogDebug("Started GetRecords Unmarshal, took", time.Since(start)) // Here, we set a timer that the initial Read() call on // HTTPResponse.Body must return by. Note that the normal // http.Client Timeout is still in effect. startReadTime = time.Now() timer := time.NewTimer(l.getRecordsReadTimeout) - r.HTTPResponse.Body = &utils.ReadCloserWrapper{ + r.HTTPResponse.Body = &ReadCloserWrapper{ ReadCloser: r.HTTPResponse.Body, OnReadFn: func(stream io.ReadCloser, b []byte) (n int, err error) { // The OnReadFn will be called each time @@ -380,7 +455,7 @@ func (l *Listener) getRecords(batchSize int) (int, error) { // the Read() timer.Reset(l.getRecordsReadTimeout) n, err = result.n, result.err - l.Log(fmt.Sprintf("GetRecords read %d bytes, took %v", n, time.Since(readStart))) + l.LogDebug(fmt.Sprintf("GetRecords read %d bytes, took %v", n, time.Since(readStart))) case <-timer.C: // If we timeout, we return an error // that will unblock ioutil.ReadAll(). @@ -388,45 +463,45 @@ func (l *Listener) getRecords(batchSize int) (int, error) { // to return an error. This error will // propogate to the original req.Send() // call (below) - l.Log(fmt.Sprintf("GetRecords read timed out after %v", time.Since(readStart))) - err = ErrTimeoutReadResponseBody + l.LogDebug(fmt.Sprintf("GetRecords read timed out after %v", time.Since(readStart))) + err = errs.ErrTimeoutReadResponseBody case <-l.pipeOfDeath: // The pipe of death will abort any pending // reads on a GetRecords call. - l.Log(fmt.Sprintf("GetRecords received pipe of death after %v", time.Since(readStart))) - err = ErrPipeOfDeath + l.LogDebug(fmt.Sprintf("GetRecords received pipe of death after %v", time.Since(readStart))) + err = errs.ErrPipeOfDeath } return }, OnCloseFn: func() { - l.Stats.AddGetRecordsReadResponseTime(time.Since(startReadTime)) - l.Log("Finished GetRecords body read, took", time.Since(start)) + l.Stats.AddGetRecordsReadResponseDuration(time.Since(startReadTime)) + l.LogDebug("Finished GetRecords body read, took", time.Since(start)) startUnmarshalTime = time.Now() }, } }) req.Handlers.Unmarshal.PushBack(func(r *request.Request) { - l.Stats.AddGetRecordsUnmarshalTime(time.Since(startUnmarshalTime)) - l.Log("Finished GetRecords Unmarshal, took", time.Since(start)) + l.Stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) + l.LogDebug("Finished GetRecords Unmarshal, took", time.Since(start)) }) // Send the GetRecords request - l.Log("Starting GetRecords Build/Sign request, took", time.Since(start)) + l.LogDebug("Starting GetRecords Build/Sign request, took", time.Since(start)) l.Stats.AddGetRecordsCalled(1) if err := req.Send(); err != nil { - l.Log("Error getting records:", err) + l.LogError("Error getting records:", err) return 0, err } - l.Stats.AddGetRecordsTime(time.Since(start)) + l.Stats.AddGetRecordsDuration(time.Since(start)) // Process Records - l.Log(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), l.shard, time.Since(start))) + l.LogDebug(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), l.shard, time.Since(start))) if resp == nil { - return 0, ErrNilGetRecordsResponse + return 0, errs.ErrNilGetRecordsResponse } delivered := 0 - l.Stats.AddBatchSizeSample(len(resp.Records)) + l.Stats.AddBatchSize(len(resp.Records)) for _, record := range resp.Records { if record != nil { // Allow (only) a pipeOfDeath to trigger an instance @@ -435,9 +510,9 @@ func (l *Listener) getRecords(batchSize int) (int, error) { // from completing the delivery of the current batch of // records. select { - case l.messages <- message.FromKinesisRecord(record): + case l.messages <- message.FromRecord(record): delivered++ - l.Stats.AddConsumedSample(1) + l.Stats.AddConsumed(1) if record.SequenceNumber != nil { // We can safely ignore if this call returns // error, as if we somehow receive an empty @@ -448,8 +523,8 @@ func (l *Listener) getRecords(batchSize int) (int, error) { l.setSequenceNumber(*record.SequenceNumber) } case <-l.pipeOfDeath: - l.Log(fmt.Sprintf("getRecords received pipe of death while delivering messages, %d delivered, ~%d dropped", delivered, len(resp.Records)-delivered)) - return delivered, ErrPipeOfDeath + l.LogInfo(fmt.Sprintf("getRecords received pipe of death while delivering messages, %d delivered, ~%d dropped", delivered, len(resp.Records)-delivered)) + return delivered, errs.ErrPipeOfDeath } } } @@ -495,7 +570,7 @@ func (l *Listener) startConsuming() bool { func (l *Listener) shouldConsume(ctx context.Context) (bool, error) { select { case <-l.pipeOfDeath: - return false, ErrPipeOfDeath + return false, errs.ErrPipeOfDeath case <-ctx.Done(): return false, ctx.Err() default: @@ -517,7 +592,7 @@ func (l *Listener) stopConsuming() { // Cancellation supported through contexts. func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, error) { if !l.startConsuming() { - return nil, ErrAlreadyConsuming + return nil, errs.ErrAlreadyConsuming } defer func() { l.stopConsuming() @@ -537,7 +612,7 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, e return nil, err } if n > 0 { - l.Stats.AddDeliveredSample(1) + l.Stats.AddDelivered(1) return <-l.messages, nil } } @@ -559,10 +634,10 @@ func (l *Listener) RetrieveFnWithContext(ctx context.Context, fn MessageFn) erro var wg sync.WaitGroup wg.Add(1) start := time.Now() - go fn(msg.Value(), &wg) + go fn(msg.Data, &wg) wg.Wait() - l.Stats.AddProcessedTime(time.Since(start)) - l.Stats.AddProcessedSample(1) + l.Stats.AddProcessedDuration(time.Since(start)) + l.Stats.AddProcessed(1) return nil } @@ -608,27 +683,27 @@ func (l *Listener) consume(ctx context.Context) { case net.Error: if err.Timeout() { l.Stats.AddGetRecordsTimeout(1) - l.Log("Received net error:", err.Error()) + l.LogError("Received net error:", err.Error()) } else { - l.Log("Received unknown net error:", err.Error()) + l.LogError("Received unknown net error:", err.Error()) } case error: switch err { - case ErrTimeoutReadResponseBody: + case errs.ErrTimeoutReadResponseBody: l.Stats.AddGetRecordsReadTimeout(1) - l.Log("Received error:", err.Error()) + l.LogError("Received error:", err.Error()) default: - l.Log("Received error:", err.Error()) + l.LogError("Received error:", err.Error()) } case awserr.Error: switch err.Code() { case kinesis.ErrCodeProvisionedThroughputExceededException: l.Stats.AddProvisionedThroughputExceeded(1) default: - l.Log("Received AWS error:", err.Error()) + l.LogError("Received AWS error:", err.Error()) } default: - l.Log("Received unknown error:", err.Error()) + l.LogError("Received unknown error:", err.Error()) } } } @@ -650,7 +725,7 @@ stop: if !ok { break stop } - l.Stats.AddDeliveredSample(1) + l.Stats.AddDelivered(1) // For simplicity, did not do the pipe of death here. // If POD is received, we may deliver a couple more // messages (especially since select is random in which @@ -664,14 +739,14 @@ stop: var fnWg sync.WaitGroup fnWg.Add(1) start := time.Now() - fn(msg.Value(), &fnWg) + fn(msg.Data, &fnWg) fnWg.Wait() - l.Stats.AddProcessedTime(time.Since(start)) - l.Stats.AddProcessedSample(1) + l.Stats.AddProcessedDuration(time.Since(start)) + l.Stats.AddProcessed(1) wg.Done() }(msg) case <-l.pipeOfDeath: - l.Log("ListenWithContext received pipe of death") + l.LogInfo("ListenWithContext received pipe of death") break stop } } diff --git a/listener/listener_test.go b/listener/listener_test.go index 85aeb04..5f5dbea 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -17,6 +17,7 @@ import ( "github.com/aws/aws-sdk-go/service/kinesis" "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/errs" ) func putRecord(l *Listener, b []byte) (*string, error) { @@ -101,7 +102,7 @@ func TestListener(t *testing.T) { So(len(shards), ShouldEqual, 1) l, err := NewListener(stream, shards[0], func(c *Config) { - c.FromKinetic(k) + c.SetAwsConfig(k.Session.Config) c.SetConcurrency(10) }) So(l, ShouldNotBeNil) @@ -118,12 +119,12 @@ func TestListener(t *testing.T) { Convey("check that setting an empty shard iterator returns an error", func() { err := l.setNextShardIterator("") - So(err, ShouldEqual, ErrEmptyShardIterator) + So(err, ShouldEqual, errs.ErrEmptyShardIterator) }) Convey("check that setting an empty sequence number returns an error", func() { err := l.setSequenceNumber("") - So(err, ShouldEqual, ErrEmptySequenceNumber) + So(err, ShouldEqual, errs.ErrEmptySequenceNumber) }) Convey("check that we can get the TRIM_HORIZON shard iterator", func() { @@ -140,7 +141,7 @@ func TestListener(t *testing.T) { So(seq, ShouldNotBeNil) msg, err := l.Retrieve() So(err, ShouldBeNil) - So(string(msg.Value()), ShouldEqual, datum) + So(string(msg.Data), ShouldEqual, datum) Convey(fmt.Sprintf("check that iteration %d properly advanced the shard iterator", n), func() { So(l.shardIterator.shardIteratorType, ShouldEqual, "AT_SEQUENCE_NUMBER") So(l.shardIterator.sequenceNumber, ShouldEqual, *seq) @@ -160,7 +161,7 @@ func TestListener(t *testing.T) { elapsed := time.Since(start) Printf("(it blocked %f seconds)\n", elapsed.Seconds()) So(err, ShouldBeNil) - So(string(msg.Value()), ShouldEqual, data) + So(string(msg.Data), ShouldEqual, data) So(elapsed.Seconds(), ShouldBeGreaterThan, 1) }) @@ -200,7 +201,7 @@ func TestListener(t *testing.T) { }() msg, err := l.RetrieveWithContext(ctx) So(err, ShouldBeNil) - So(string(msg.Value()), ShouldEqual, data) + So(string(msg.Data), ShouldEqual, data) }) Convey("check that retrieve properly blocks other retrieves and attempts to set the shard id", func(c C) { @@ -216,7 +217,7 @@ func TestListener(t *testing.T) { }() <-time.After(10 * time.Millisecond) _, err := l.Retrieve() - So(err, ShouldEqual, ErrAlreadyConsuming) + So(err, ShouldEqual, errs.ErrAlreadyConsuming) wg.Wait() }) @@ -319,10 +320,11 @@ func TestListener(t *testing.T) { Printf("(count was %d)", atomic.LoadInt64(&count)) }) + // TODO: test get records read timeout + Reset(func() { k.DeleteStream(stream) k.WaitUntilStreamDeleted(context.TODO(), stream, request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) }) - }) } diff --git a/utils/readcloser.go b/listener/readcloser.go similarity index 97% rename from utils/readcloser.go rename to listener/readcloser.go index 2365f3f..252bdce 100644 --- a/utils/readcloser.go +++ b/listener/readcloser.go @@ -1,4 +1,4 @@ -package utils +package listener import ( "io" diff --git a/listener/stats.go b/listener/stats.go index 491ab27..0d897af 100644 --- a/listener/stats.go +++ b/listener/stats.go @@ -8,55 +8,41 @@ import ( // the Kinetic Listener library. This was really built with rcrowley/go-metrics // in mind. type StatsCollector interface { - // histograms give us the count, sum, min, max, mean, percentiles, - // standard deviation, and variance of the data. these metrics should - // give us the total number (sum) of messages consumed, delivered, and - // processed, as well as the average (mean) batch size. - AddConsumedSample(int) - AddDeliveredSample(int) - AddProcessedSample(int) - AddBatchSizeSample(int) - - // meters give us the count and rate of the data. these metrics should - // give us the average number of times: - // - GetRecords was called per second - // - ProvisionedThroughputExceeded was received per second - // - GetRecords timed out per second - // - GetRecords read timed out per second + AddConsumed(int) + AddDelivered(int) + AddProcessed(int) + AddBatchSize(int) AddGetRecordsCalled(int) AddProvisionedThroughputExceeded(int) AddGetRecordsTimeout(int) AddGetRecordsReadTimeout(int) - // timers give us the count, sum, min, max, mean, percentiles, standard - // deviation, variance, as well as the rate of the data. - // TODO: describe these metrics better - AddProcessedTime(time.Duration) - AddGetRecordsTime(time.Duration) - AddGetRecordsReadResponseTime(time.Duration) - AddGetRecordsUnmarshalTime(time.Duration) + AddProcessedDuration(time.Duration) + AddGetRecordsDuration(time.Duration) + AddGetRecordsReadResponseDuration(time.Duration) + AddGetRecordsUnmarshalDuration(time.Duration) } // NilStatsCollector is a stats listener that ignores all metrics. type NilStatsCollector struct{} -// AddConsumedSample records a count of the number of messages received from AWS +// AddConsumed records a count of the number of messages received from AWS // Kinesis by the listener. -func (l *NilStatsCollector) AddConsumedSample(int) {} +func (l *NilStatsCollector) AddConsumed(int) {} -// AddDeliveredSample records a count of the number of messages delivered to the +// AddDelivered records a count of the number of messages delivered to the // application by the listener. -func (l *NilStatsCollector) AddDeliveredSample(int) {} +func (l *NilStatsCollector) AddDelivered(int) {} -// AddProcessedSample records a count of the number of messages processed by the +// AddProcessed records a count of the number of messages processed by the // application by the listener. This is based on a WaitGroup that is sent to // the RetrieveFn and Listen functions. Retrieve does not count processed // messages. -func (l *NilStatsCollector) AddProcessedSample(int) {} +func (l *NilStatsCollector) AddProcessed(int) {} -// AddBatchSizeSample records a count of the number of messages returned by +// AddBatchSize records a count of the number of messages returned by // GetRecords in the listener. -func (l *NilStatsCollector) AddBatchSizeSample(int) {} +func (l *NilStatsCollector) AddBatchSize(int) {} // AddGetRecordsCalled records the number of times the GetRecords API was called // by the listener. @@ -76,18 +62,18 @@ func (l *NilStatsCollector) AddGetRecordsTimeout(int) {} // WithGetRecordsReadTimeout configuration. func (l *NilStatsCollector) AddGetRecordsReadTimeout(int) {} -// AddProcessedTime records the duration to process a record. See notes on -// AddProcessedSample. -func (l *NilStatsCollector) AddProcessedTime(time.Duration) {} +// AddProcessedDuration records the duration to process a record. See notes on +// AddProcessed. +func (l *NilStatsCollector) AddProcessedDuration(time.Duration) {} -// AddGetRecordsTime records the duration that the GetRecords API request took. -// Only the times of successful calls are measured. -func (l *NilStatsCollector) AddGetRecordsTime(time.Duration) {} +// AddGetRecordsDuration records the duration that the GetRecords API request +// took. Only the times of successful calls are measured. +func (l *NilStatsCollector) AddGetRecordsDuration(time.Duration) {} -// AddGetRecordsReadResponseTime records the duration that it took to read the -// response body of a GetRecords API request. -func (l *NilStatsCollector) AddGetRecordsReadResponseTime(time.Duration) {} +// AddGetRecordsReadResponseDuration records the duration that it took to read +// the response body of a GetRecords API request. +func (l *NilStatsCollector) AddGetRecordsReadResponseDuration(time.Duration) {} -// AddGetRecordsUnmarshalTime records the duration that it took to unmarshal the -// response body of a GetRecords API request. -func (l *NilStatsCollector) AddGetRecordsUnmarshalTime(time.Duration) {} +// AddGetRecordsUnmarshalDuration records the duration that it took to unmarshal +// the response body of a GetRecords API request. +func (l *NilStatsCollector) AddGetRecordsUnmarshalDuration(time.Duration) {} diff --git a/logging/logging.go b/logging/logging.go index 16c67f9..37b502c 100644 --- a/logging/logging.go +++ b/logging/logging.go @@ -8,6 +8,15 @@ const ( // LogOff disables all logging. LogOff aws.LogLevelType = (iota * 0x1000) << 16 + // LogBug enables logging of bugs in code. + LogBug + + // LogError enables logging of errors. + LogError + + // LogInfo enables logging of informational messages.. + LogInfo + // LogDebug enables debug logging. LogDebug ) @@ -16,3 +25,43 @@ const ( // LogPlaceholder is a placeholder debug sub-level. LogPlaceholder aws.LogLevelType = LogDebug | (1 << (16 + iota)) ) + +type ILogHelper interface { + Log(aws.LogLevelType, ...interface{}) + LogBug(...interface{}) + LogError(...interface{}) + LogInfo(...interface{}) + LogDebug(...interface{}) +} + +type LogHelper struct { + LogLevel aws.LogLevelType + Logger aws.Logger +} + +// log handles levelled logging +func (l *LogHelper) Log(level aws.LogLevelType, args ...interface{}) { + if l.LogLevel.Matches(level) { + l.Logger.Log(args...) + } +} + +// LogBug logs a BUG in the code. +func (l *LogHelper) LogBug(args ...interface{}) { + l.Log(LogBug, args...) +} + +// LogBug logs an error. +func (l *LogHelper) LogError(args ...interface{}) { + l.Log(LogError, args...) +} + +// LogInfo logs an informational message. +func (l *LogHelper) LogInfo(args ...interface{}) { + l.Log(LogInfo, args...) +} + +// Log a debug message using the AWS SDK logger. +func (l *LogHelper) LogDebug(args ...interface{}) { + l.Log(LogDebug, args...) +} diff --git a/message/message.go b/message/message.go index 38c3a76..79691bb 100644 --- a/message/message.go +++ b/message/message.go @@ -3,6 +3,7 @@ package message import ( "time" + "github.com/aws/aws-sdk-go/service/firehose" "github.com/aws/aws-sdk-go/service/kinesis" ) @@ -20,11 +21,15 @@ type Message struct { // For kinesis.PutRecordResultEntry ErrorCode *string ErrorMessage *string - ShardId *string + ShardID *string + + // For firehose.PutRecordBatchResponseEntry + RecordID *string FailCount int } +// FromRecord creates a message from the kinesis.Record returned from GetRecords func FromRecord(record *kinesis.Record) *Message { return &Message{ ApproximateArrivalTimestamp: record.ApproximateArrivalTimestamp, @@ -34,6 +39,8 @@ func FromRecord(record *kinesis.Record) *Message { } } +// MakeRequestEntry creates a kinesis.PutRecordsRequestEntry to be used in the +// kinesis.PutRecords API call. func (m *Message) MakeRequestEntry() *kinesis.PutRecordsRequestEntry { return &kinesis.PutRecordsRequestEntry{ Data: m.Data, @@ -41,3 +48,11 @@ func (m *Message) MakeRequestEntry() *kinesis.PutRecordsRequestEntry { PartitionKey: m.PartitionKey, } } + +// MakeFirehoseRecord creates a firehose.Record to be used in the +// firehose.PutRecordBatch API call. +func (m *Message) MakeFirehoseRecord() *firehose.Record { + return &firehose.Record{ + Data: m.Data, + } +} diff --git a/producer/config.go b/producer/config.go index 34af3fb..974cbd1 100644 --- a/producer/config.go +++ b/producer/config.go @@ -5,32 +5,39 @@ import ( "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/config" "github.com/rewardStyle/kinetic/logging" ) // Config is used to configure a Producer instance. type Config struct { - *kinetic.AwsOptions + *config.AwsOptions *producerOptions + LogLevel aws.LogLevelType } // NewConfig creates a new instance of Config. func NewConfig() *Config { return &Config{ - AwsOptions: kinetic.DefaultAwsOptions(), + AwsOptions: config.DefaultAwsOptions(), producerOptions: &producerOptions{ batchSize: 500, batchTimeout: 1 * time.Second, queueDepth: 500, maxRetryAttempts: 10, concurrency: 1, - LogLevel: logging.LogOff, Stats: &NilStatsCollector{}, }, + LogLevel: logging.LogOff, } } +// SetAwsConfig configures the AWS Config used to create Sessions (and therefore +// kinesis clients). +func (c *Config) SetAwsConfig(config *aws.Config) { + c.AwsConfig = config +} + // SetBatchSize configures the batch size to flush pending records to the // PutRecords call. func (c *Config) SetBatchSize(batchSize int) { @@ -61,15 +68,16 @@ func (c *Config) SetConcurrency(concurrency int) { c.concurrency = concurrency } -// SetWriter sets the underlying stream writer (Kinesis or Firehose) for the -// producer. There can only be a single writer associated with a producer (as -// various retry logic/state is not easily shared between multiple writers). If -// multiple streams are desired, create two different producers and write to -// both. -func (c *Config) SetWriter(writer StreamWriter) { - c.writer = writer +// SetKinesisStream sets the producer to write to the given Kinesis stream. +func (c *Config) SetKinesisStream(stream string) { + c.writer = NewKinesisWriter(stream) } +// SetFirehoseStream sets the producer to write to the given Firehose stream. +// func (c *Config) SetFirehoseStream(stream string) { +// c.writer = NewFirehoseWriter(stream) +// } + // SetLogLevel configures both the SDK and Kinetic log levels. func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { c.AwsOptions.SetLogLevel(logLevel) @@ -80,9 +88,3 @@ func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { func (c *Config) SetStatsCollector(stats StatsCollector) { c.Stats = stats } - -// FromKinetic configures the session from Kinetic. -func (c *Config) FromKinetic(k *kinetic.Kinetic) *Config { - c.AwsConfig = k.Session.Config - return c -} diff --git a/producer/config_test.go b/producer/config_test.go new file mode 100644 index 0000000..271476f --- /dev/null +++ b/producer/config_test.go @@ -0,0 +1,122 @@ +package producer + +import ( + . "github.com/smartystreets/goconvey/convey" + + "testing" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/session" + + "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/logging" + "github.com/rewardStyle/kinetic/message" +) + +type DebugStatsCollector struct{} + +func (l *DebugStatsCollector) AddSent(int) {} +func (l *DebugStatsCollector) AddFailed(int) {} +func (l *DebugStatsCollector) AddDropped(int) {} +func (l *DebugStatsCollector) AddBatchSize(int) {} +func (l *DebugStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} +func (l *DebugStatsCollector) AddPutRecordsCalled(int) {} +func (l *DebugStatsCollector) AddProvisionedThroughputExceeded(int) {} +func (l *DebugStatsCollector) AddPutRecordsTimeout(int) {} +func (l *DebugStatsCollector) AddPutRecordsDuration(time.Duration) {} +func (l *DebugStatsCollector) AddPutRecordsBuildDuration(time.Duration) {} +func (l *DebugStatsCollector) AddPutRecordsSendDuration(time.Duration) {} + +func getSession(config *Config) *session.Session { + sess, err := config.GetSession() + So(err, ShouldBeNil) + So(sess, ShouldNotBeNil) + return sess +} + +type DebugStreamWriter struct{} + +func (w *DebugStreamWriter) PutRecords(batch []*message.Message) ([]*message.Message, error) { + return nil, nil +} + +func (w *DebugStreamWriter) AssociateProducer(producer Producer) error { + return nil +} + +func TestNewConfig(t *testing.T) { + Convey("given a new producer config", t, func() { + config := NewConfig() + + Convey("check the default values for its non-zero config", func() { + So(config.batchSize, ShouldEqual, 500) + So(config.batchTimeout, ShouldEqual, 1*time.Second) + So(config.queueDepth, ShouldEqual, 500) + So(config.maxRetryAttempts, ShouldEqual, 10) + So(config.concurrency, ShouldEqual, 1) + So(config.Stats, ShouldHaveSameTypeAs, &NilStatsCollector{}) + So(config.LogLevel.Value(), ShouldEqual, logging.LogOff) + }) + + Convey("check that we can retrieve an aws.Session from it ", func() { + getSession(config) + }) + + Convey("check that we can set both the sdk and kinetic log level", func() { + ll := aws.LogDebug | aws.LogDebugWithSigning | logging.LogDebug + config.SetLogLevel(ll) + sess := getSession(config) + So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) + So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) + So(config.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) + }) + + Convey("check that we can set the AWS configuration", func() { + k, err := kinetic.New(func(c *kinetic.Config) { + c.SetEndpoint("bogus-endpoint") + }) + So(err, ShouldBeNil) + config.SetAwsConfig(k.Session.Config) + sess := getSession(config) + So(aws.StringValue(sess.Config.Endpoint), ShouldEqual, "bogus-endpoint") + }) + + Convey("check that we can set the batch size", func() { + config.SetBatchSize(1000) + So(config.batchSize, ShouldEqual, 1000) + }) + + Convey("check that we can set the batch timeout", func() { + config.SetBatchTimeout(10 * time.Second) + So(config.batchTimeout, ShouldEqual, 10*time.Second) + }) + + Convey("check that we can set the queue depth", func() { + config.SetQueueDepth(1000) + So(config.queueDepth, ShouldEqual, 1000) + }) + + Convey("check that we can set the max retries", func() { + config.SetMaxRetryAttempts(100) + So(config.maxRetryAttempts, ShouldEqual, 100) + }) + + Convey("check that we can set the concurrency limit", func() { + config.SetConcurrency(50) + So(config.concurrency, ShouldEqual, 50) + }) + + Convey("check that we can configure a stats collector", func() { + config.SetStatsCollector(&DebugStatsCollector{}) + So(config.Stats, ShouldHaveSameTypeAs, &DebugStatsCollector{}) + }) + + Convey("check that we can configure a kinesis stream", func() { + stream := "some-stream" + config.SetKinesisStream(stream) + So(config.writer, ShouldHaveSameTypeAs, &KinesisWriter{}) + So(config.writer.(*KinesisWriter).stream, ShouldEqual, stream) + }) + }) +} diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go new file mode 100644 index 0000000..9f5b0ab --- /dev/null +++ b/producer/firehosewriter.go @@ -0,0 +1,139 @@ +package producer + +import ( + "fmt" + "sync" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/request" + "github.com/aws/aws-sdk-go/service/firehose" + "github.com/aws/aws-sdk-go/service/firehose/firehoseiface" + + "github.com/rewardStyle/kinetic/errs" + "github.com/rewardStyle/kinetic/message" +) + +// FirehoseWriter handles the API to send records to Kinesis. +type FirehoseWriter struct { + stream string + + producer *Producer + client firehoseiface.FirehoseAPI + clientMu sync.Mutex +} + +// NewFirehoseWriter creates a new stream writer to write records to a Kinesis. +func NewFirehoseWriter(stream string) *FirehoseWriter { + return &FirehoseWriter{ + stream: stream, + } +} + +// ensureClient will lazily make sure we have an AWS Kinesis client. +func (w *FirehoseWriter) ensureClient() error { + w.clientMu.Lock() + defer w.clientMu.Unlock() + if w.client == nil { + if w.producer == nil { + return errs.ErrNilProducer + } + w.client = firehose.New(w.producer.Session) + } + return nil +} + +// AssociateProducer associates the Firehose stream writer to a producer. +func (w *FirehoseWriter) AssociateProducer(p *Producer) error { + w.clientMu.Lock() + defer w.clientMu.Unlock() + if w.producer != nil { + return errs.ErrProducerAlreadyAssociated + } + w.producer = p + return nil +} + +// PutRecords sends a batch of records to Firehose and returns a list of records +// that need to be retried. +func (w *FirehoseWriter) PutRecords(messages []*message.Message) ([]*message.Message, error) { + if err := w.ensureClient(); err != nil { + return nil, err + } + + var startSendTime time.Time + var startBuildTime time.Time + + start := time.Now() + var records []*firehose.Record + for _, msg := range messages { + records = append(records, msg.MakeFirehoseRecord()) + } + req, resp := w.client.PutRecordBatchRequest(&firehose.PutRecordBatchInput{ + DeliveryStreamName: aws.String(w.stream), + Records: records, + }) + + req.Handlers.Build.PushFront(func(r *request.Request) { + startBuildTime = time.Now() + w.producer.LogDebug("Start PutRecords Build, took", time.Since(start)) + }) + + req.Handlers.Build.PushBack(func(r *request.Request) { + w.producer.Stats.AddPutRecordsBuildDuration(time.Since(startBuildTime)) + w.producer.LogDebug("Finished PutRecords Build, took", time.Since(start)) + }) + + req.Handlers.Send.PushFront(func(r *request.Request) { + startSendTime = time.Now() + w.producer.LogDebug("Start PutRecords Send took", time.Since(start)) + }) + + req.Handlers.Build.PushBack(func(r *request.Request) { + w.producer.Stats.AddPutRecordsSendDuration(time.Since(startSendTime)) + w.producer.LogDebug("Finished PutRecords Send, took", time.Since(start)) + }) + + w.producer.LogDebug("Starting PutRecords Build/Sign request, took", time.Since(start)) + w.producer.Stats.AddPutRecordsCalled(1) + if err := req.Send(); err != nil { + w.producer.LogError("Error putting records:", err.Error()) + return nil, err + } + w.producer.Stats.AddPutRecordsDuration(time.Since(start)) + + if resp == nil { + return nil, errs.ErrNilPutRecordsResponse + } + if resp.FailedPutCount == nil { + return nil, errs.ErrNilFailedRecordCount + } + attempted := len(messages) + failed := int(aws.Int64Value(resp.FailedPutCount)) + sent := attempted - failed + w.producer.LogDebug(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) + + var retries []*message.Message + var err error + for idx, record := range resp.RequestResponses { + if record.RecordId != nil { + // TODO: per-shard metrics + messages[idx].RecordID = record.RecordId + } else { + switch aws.StringValue(record.ErrorCode) { + case firehose.ErrCodeLimitExceededException: + w.producer.Stats.AddProvisionedThroughputExceeded(1) + default: + w.producer.LogDebug("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(record.ErrorMessage)) + } + messages[idx].ErrorCode = record.ErrorCode + messages[idx].ErrorMessage = record.ErrorMessage + messages[idx].FailCount++ + retries = append(retries, messages[idx]) + } + } + if len(retries) > 0 { + err = errs.ErrRetryRecords + } + return retries, err +} diff --git a/producer/firehosewriter/config.go b/producer/firehosewriter/config.go deleted file mode 100644 index 4f15d66..0000000 --- a/producer/firehosewriter/config.go +++ /dev/null @@ -1 +0,0 @@ -package firehosewriter diff --git a/producer/firehosewriter/firehosewriter.go b/producer/firehosewriter/firehosewriter.go deleted file mode 100644 index 4f15d66..0000000 --- a/producer/firehosewriter/firehosewriter.go +++ /dev/null @@ -1 +0,0 @@ -package firehosewriter diff --git a/producer/kinesiswriter/kinesiswriter.go b/producer/kinesiswriter.go similarity index 53% rename from producer/kinesiswriter/kinesiswriter.go rename to producer/kinesiswriter.go index f68f8bf..b16aae5 100644 --- a/producer/kinesiswriter/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -1,7 +1,6 @@ -package kinesiswriter +package producer import ( - "errors" "fmt" "sync" "time" @@ -11,44 +10,24 @@ import ( "github.com/aws/aws-sdk-go/service/kinesis" "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" + "github.com/rewardStyle/kinetic/errs" "github.com/rewardStyle/kinetic/message" - "github.com/rewardStyle/kinetic/producer" ) -var ( - // ErrNilPutRecordsResponse is returned when the PutRecords call returns - // a nil response. - ErrNilPutRecordsResponse = errors.New("PutRecords returned a nil response") - - // ErrNilFailedRecordCount is returned when the PutRecords call returns - // a nil FailedRecordCount. - ErrNilFailedRecordCount = errors.New("GetFailedRecordCount returned a nil FailedRecordCount") -) - -type kinesisWriterOptions struct { -} - // KinesisWriter handles the API to send records to Kinesis. type KinesisWriter struct { - *kinesisWriterOptions + stream string - producer *producer.Producer + producer *Producer client kinesisiface.KinesisAPI clientMu sync.Mutex } // NewKinesisWriter creates a new stream writer to write records to a Kinesis. -func NewKinesisWriter(fn func(*Config)) (*KinesisWriter, error) { - config := NewConfig() - fn(config) +func NewKinesisWriter(stream string) *KinesisWriter { return &KinesisWriter{ - kinesisWriterOptions: config.kinesisWriterOptions, - }, nil -} - -// Log a debug message using the Producer logger. -func (w *KinesisWriter) Log(args ...interface{}) { - w.producer.Log(args...) + stream: stream, + } } // ensureClient will lazily make sure we have an AWS Kinesis client. @@ -57,7 +36,7 @@ func (w *KinesisWriter) ensureClient() error { defer w.clientMu.Unlock() if w.client == nil { if w.producer == nil { - return producer.ErrNilProducer + return errs.ErrNilProducer } w.client = kinesis.New(w.producer.Session) } @@ -65,11 +44,11 @@ func (w *KinesisWriter) ensureClient() error { } // AssociateProducer associates the Kinesis stream writer to a producer. -func (w *KinesisWriter) AssociateProducer(p *producer.Producer) error { +func (w *KinesisWriter) AssociateProducer(p *Producer) error { w.clientMu.Lock() defer w.clientMu.Unlock() if w.producer != nil { - return producer.ErrProducerAlreadyAssociated + return errs.ErrProducerAlreadyAssociated } w.producer = p return nil @@ -90,69 +69,72 @@ func (w *KinesisWriter) PutRecords(messages []*message.Message) ([]*message.Mess for _, msg := range messages { records = append(records, msg.MakeRequestEntry()) } - req, resp := w.client.PutRecordsRequest(&kinesis.PutRecordsInput{}) + req, resp := w.client.PutRecordsRequest(&kinesis.PutRecordsInput{ + StreamName: aws.String(w.stream), + Records: records, + }) req.Handlers.Build.PushFront(func(r *request.Request) { startBuildTime = time.Now() - w.Log("Start PutRecords Build, took", time.Since(start)) + w.producer.LogDebug("Start PutRecords Build, took", time.Since(start)) }) req.Handlers.Build.PushBack(func(r *request.Request) { - w.producer.Stats.AddPutRecordsBuildTime(time.Since(startBuildTime)) - w.Log("Finished PutRecords Build, took", time.Since(start)) + w.producer.Stats.AddPutRecordsBuildDuration(time.Since(startBuildTime)) + w.producer.LogDebug("Finished PutRecords Build, took", time.Since(start)) }) req.Handlers.Send.PushFront(func(r *request.Request) { startSendTime = time.Now() - w.Log("Start PutRecords Send took", time.Since(start)) + w.producer.LogDebug("Start PutRecords Send took", time.Since(start)) }) req.Handlers.Build.PushBack(func(r *request.Request) { - w.producer.Stats.AddPutRecordsSendTime(time.Since(startSendTime)) - w.Log("Finished PutRecords Send, took", time.Since(start)) + w.producer.Stats.AddPutRecordsSendDuration(time.Since(startSendTime)) + w.producer.LogDebug("Finished PutRecords Send, took", time.Since(start)) }) - w.Log("Starting PutRecords Build/Sign request, took", time.Since(start)) + w.producer.LogDebug("Starting PutRecords Build/Sign request, took", time.Since(start)) w.producer.Stats.AddPutRecordsCalled(1) if err := req.Send(); err != nil { - w.Log("Error putting records:", err.Error()) + w.producer.LogError("Error putting records:", err.Error()) return nil, err } - w.producer.Stats.AddPutRecordsTime(time.Since(start)) + w.producer.Stats.AddPutRecordsDuration(time.Since(start)) if resp == nil { - return nil, ErrNilPutRecordsResponse + return nil, errs.ErrNilPutRecordsResponse } if resp.FailedRecordCount == nil { - return nil, ErrNilFailedRecordCount + return nil, errs.ErrNilFailedRecordCount } attempted := len(messages) failed := int(aws.Int64Value(resp.FailedRecordCount)) sent := attempted - failed - w.Log(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) + w.producer.LogDebug(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) var retries []*message.Message var err error for idx, record := range resp.Records { if record.SequenceNumber != nil && record.ShardId != nil { // TODO: per-shard metrics - message[idx].SequenceNumber = record.SequenceNumber - message[idx].ShardId = record.ShardId + messages[idx].SequenceNumber = record.SequenceNumber + messages[idx].ShardID = record.ShardId } else { - switch record.ErrorCode { + switch aws.StringValue(record.ErrorCode) { case kinesis.ErrCodeProvisionedThroughputExceededException: w.producer.Stats.AddProvisionedThroughputExceeded(1) default: - w.Log("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(*record.ErrorMessage)) + w.producer.LogDebug("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(record.ErrorMessage)) } - message[idx].ErrorCode = record.ErrorCode - message[idx].ErrorMessage = record.ErrorMessage - message[idx].FailCount++ + messages[idx].ErrorCode = record.ErrorCode + messages[idx].ErrorMessage = record.ErrorMessage + messages[idx].FailCount++ retries = append(retries, messages[idx]) } } if len(retries) > 0 { - err = producer.ErrRetryRecords + err = errs.ErrRetryRecords } return retries, err } diff --git a/producer/kinesiswriter/config.go b/producer/kinesiswriter/config.go deleted file mode 100644 index a6a58e2..0000000 --- a/producer/kinesiswriter/config.go +++ /dev/null @@ -1,15 +0,0 @@ -package kinesiswriter - -import () - -// Config is used to configure a Kinesis Writer instance. -type Config struct { - *kinesisWriterOptions -} - -// NewConfig creates a new instance of Config. -func NewConfig() *Config { - return &Config{ - kinesisWriterOptions: &kinesisWriterOptions{}, - } -} diff --git a/producer/producer.go b/producer/producer.go index 8f3064b..eaf486a 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -2,16 +2,16 @@ package producer import ( "context" - "errors" + "fmt" "net" "sync" "time" - "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/awserr" "github.com/aws/aws-sdk-go/aws/session" "github.com/aws/aws-sdk-go/service/kinesis" + "github.com/rewardStyle/kinetic/errs" "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" ) @@ -21,37 +21,9 @@ import ( type StreamWriter interface { AssociateProducer(producer *Producer) error PutRecords(message []*message.Message) ([]*message.Message, error) + ensureClient() error } -var ( - // ErrRetryRecords is returned when the PutRecords calls requires some - // records of the batch to be retried. This failure is considered part - // of normal behavior of the Kinesis stream. - ErrRetryRecords = errors.New("PutRecords requires retry of some records in batch") - - // ErrNilProducer is returned by a StreamWriter when it has not been - // correctly associated with a Producer. - ErrNilProducer = errors.New("StreamWriter not associated with a producer") - - // ErrProducerAlreadyAssociated is returned by a StreamWriter attempting - // to associate it with a Producer when it already has an association - // with a producer. - ErrProducerAlreadyAssociated = errors.New("StreamWriter already associated with a producer") - - // ErrBatchTimeout is returned by getBatch whenever the batchTimeout - // elapses prior to receiving batchSize messages. This is *not* an - // error in the sense of a failure, but is used to distinguish the - // reason getBatch has exited. - ErrBatchTimeout = errors.New("A timeout has occurred before batch has reached optimal size") - - // ErrProducerShutdown is returend by getBatch whenever both the message - // and retry channel have been closed. - ErrProducerShutdown = errors.New("The producer has shut down") - - // ErrPipeOfDeath returns when the pipe of death is closed. - ErrPipeOfDeath = errors.New("Received pipe of death") -) - // Empty is used a as a dummy type for counting semaphore channels. type Empty struct{} @@ -63,23 +35,22 @@ type producerOptions struct { concurrency int writer StreamWriter - LogLevel aws.LogLevelType - Stats StatsCollector + Stats StatsCollector } // Producer sends records to Kinesis or Firehose. type Producer struct { *producerOptions + *logging.LogHelper messages chan *message.Message retries chan *message.Message concurrencySem chan Empty pipeOfDeath chan Empty + outstanding int + shutdownCond sync.Cond producerWg sync.WaitGroup - outstanding int64 - flushCond sync.Cond - producing bool producingMu sync.Mutex @@ -97,9 +68,13 @@ func NewProducer(fn func(*Config)) (*Producer, error) { } p := &Producer{ producerOptions: config.producerOptions, - concurrencySem: make(chan Empty, config.concurrency), - pipeOfDeath: make(chan Empty), - Session: session, + LogHelper: &logging.LogHelper{ + LogLevel: config.LogLevel, + Logger: session.Config.Logger, + }, + concurrencySem: make(chan Empty, config.concurrency), + pipeOfDeath: make(chan Empty), + Session: session, } if err := p.writer.AssociateProducer(p); err != nil { return nil, err @@ -107,13 +82,6 @@ func NewProducer(fn func(*Config)) (*Producer, error) { return p, nil } -// Log a debug message using the AWS SDK logger. -func (p *Producer) Log(args ...interface{}) { - if p.Session.Config.LogLevel.Matches(logging.LogDebug) { - p.Session.Config.Logger.Log(args...) - } -} - // startConsuming will initialize the producer and set producing to true if // there is not already another consume loop running. func (p *Producer) startProducing() bool { @@ -123,8 +91,8 @@ func (p *Producer) startProducing() bool { p.producing = true p.messages = make(chan *message.Message, p.queueDepth) p.retries = make(chan *message.Message) // TODO: should we use a buffered channel? + p.shutdownCond = sync.Cond{L: &sync.Mutex{}} p.outstanding = 0 - p.flushCond = sync.Cond{L: &sync.Mutex{}} return true } return false @@ -140,69 +108,21 @@ func (p *Producer) stopProducing() { p.producing = false } -// getBatch will retrieve a batch of messages by batchSize and batchTimeout for -// delivery. -func (p *Producer) getBatch() ([]*message.Message, error) { +func (p *Producer) sendBatch(batch []*message.Message) { + defer func() { + p.shutdownCond.L.Lock() + p.outstanding-- + p.shutdownCond.L.Unlock() + }() + attempts := 0 + var retries []*message.Message var err error - var batch []*message.Message - var timer <-chan time.Time -stop: - for len(batch) <= p.batchSize { - select { - // Using the select, retry messages will interleave with new - // messages. This is preferable to putting the messages at the - // end of the channel as it minimizes the delay in the delivery - // of retry messages. - case msg, ok := <-p.retries: - if !ok { - p.retries = nil - } else { - batch = append(batch, msg) - if timer != nil { - timer = time.After(p.batchTimeout) - } - } - case msg, ok := <-p.messages: - if !ok { - p.messages = nil - } else { - batch = append(batch, msg) - if timer != nil { - timer = time.After(p.batchTimeout) - } - } - case <-timer: - err = ErrBatchTimeout - break stop - case <-p.pipeOfDeath: - return nil, ErrPipeOfDeath - } - if p.messages == nil && p.retries == nil { - err = ErrProducerShutdown - break stop - } - } - p.Stats.AddBatchSizeSample(len(batch)) - return batch, err -} - -func (p *Producer) dispatchBatch(batch []*message.Message) { - defer p.flushCond.Signal() stop: for { - retries, err := p.writer.PutRecords(batch) + retries, err = p.writer.PutRecords(batch) failed := len(retries) - sent := len(batch) - failed - p.Stats.AddSentSample(sent) - p.Stats.AddFailedSample(failed) - p.decOutstanding(int64(sent)) - // This frees up another dispatchBatch to run to allow drainage - // of the messages / retry queue. This should improve - // throughput as well as prevent a potential deadlock in which - // all batches are blocked on sending retries to the retries - // channel, and thus no batches are allowed to drain the retry - // channel. - <-p.concurrencySem + p.Stats.AddSent(len(batch) - failed) + p.Stats.AddFailed(failed) if err == nil { break stop } @@ -210,29 +130,9 @@ stop: case net.Error: if err.Timeout() { p.Stats.AddPutRecordsTimeout(1) - p.Log("Received net error:", err.Error()) + p.LogError("Received net error:", err.Error()) } else { - p.Log("Received unknown net error:", err.Error()) - } - case error: - switch err { - case ErrRetryRecords: - for _, msg := range retries { - if msg.FailCount < p.maxRetryAttempts { - msg.FailCount++ - select { - case p.retries <- msg: - case <-p.pipeOfDeath: - break stop - } - } else { - p.decOutstanding(1) - p.Stats.AddDroppedSample(1) - } - } - p.Log("Received error:", err.Error()) - default: - p.Log("Received error:", err.Error()) + p.LogError("Received unknown net error:", err.Error()) } case awserr.Error: switch err.Code() { @@ -254,10 +154,47 @@ stop: // the retry mechanism is different). p.Stats.AddPutRecordsProvisionedThroughputExceeded(1) default: - p.Log("Received AWS error:", err.Error()) + p.LogError("Received AWS error:", err.Error()) + } + case error: + switch err { + case errs.ErrRetryRecords: + break stop + default: + p.LogError("Received error:", err.Error()) } default: - p.Log("Received unknown error:", err.Error()) + p.LogError("Received unknown error:", err.Error()) + } + // NOTE: We may want to go through and increment the FailCount + // for each of the records and allow the batch to be retried + // rather than retrying the batch as-is. With this approach, we + // can kill the "stop" for loop, and set the entire batch to + // retries to allow the below code to handle retrying the + // messages. + attempts++ + if attempts > p.maxRetryAttempts { + p.LogError(fmt.Sprintf("Dropping batch after %d failed attempts to deliver to stream", attempts)) + p.Stats.AddDropped(len(batch)) + break stop + } + } + // This frees up another sendBatch to run to allow drainage of the + // messages / retry queue. This should improve throughput as well as + // prevent a potential deadlock in which all batches are blocked on + // sending retries to the retries channel, and thus no batches are + // allowed to drain the retry channel. + <-p.concurrencySem + for _, msg := range retries { + if msg.FailCount < p.maxRetryAttempts { + msg.FailCount++ + select { + case p.retries <- msg: + case <-p.pipeOfDeath: + return + } + } else { + p.Stats.AddDropped(1) } } } @@ -276,55 +213,82 @@ func (p *Producer) produce() { }() stop: for { - batch, err := p.getBatch() - // If getBatch aborted due to pipe of death, we will - // immediately exit the loop. - if err == ErrPipeOfDeath { - break stop + var batch []*message.Message + timer := time.After(p.batchTimeout) + batch: + for len(batch) <= p.batchSize { + select { + // Using the select, retry messages will + // interleave with new messages. This is + // preferable to putting the messages at the end + // of the channel as it minimizes the delay in + // the delivery of retry messages. + case msg, ok := <-p.messages: + if !ok { + p.messages = nil + } else { + batch = append(batch, msg) + } + case msg := <-p.retries: + batch = append(batch, msg) + case <-timer: + break batch + case <-p.pipeOfDeath: + break stop + } } - // Regardless if getBatch produced an error (as long as - // its not the pipe of death), we will send the messages - // via PutRecords. + p.shutdownCond.L.Lock() if len(batch) > 0 { + p.outstanding++ + p.Stats.AddBatchSize(len(batch)) p.concurrencySem <- Empty{} - go p.dispatchBatch(batch) - } - // If we exited getBatch earlier with a - // ErrProducerShutdown we shut down the producer. - if err == ErrProducerShutdown { - break stop + go p.sendBatch(batch) + } else if len(batch) == 0 { + // We did not get any records -- check if we may + // be (gracefully) shutting down the producer. + // We can exit when: + // - The messages channel is nil and no new messages + // can be enqueued + // - There are no outstanding sendBatch goroutines + // and can therefore not produce any more messages + // to retry + // - The retry channel is empty + if p.messages == nil && p.outstanding == 0 && len(p.retries) == 0 { + close(p.retries) + p.shutdownCond.Signal() + break stop + } } + p.shutdownCond.L.Unlock() } }() } -// incOutstanding increments the number of outstanding messages that are to be -// delivered. -func (p *Producer) incOutstanding(i int64) { - p.flushCond.L.Lock() - defer p.flushCond.L.Unlock() - p.outstanding += i -} - -// decOutstanding decrements the number of outstanding messages that are to be -// delivered. -func (p *Producer) decOutstanding(i int64) { - p.flushCond.L.Lock() - defer p.flushCond.L.Unlock() - p.outstanding -= i +// CloseWithContext shuts down the producer, waiting for all outstanding +// messages and retries to flush. Cancellation supported through contexts. +func (p *Producer) CloseWithContext(ctx context.Context) { + c := make(chan Empty, 1) + go func() { + close(p.messages) + p.shutdownCond.L.Lock() + for p.outstanding != 0 { + p.shutdownCond.Wait() + } + p.shutdownCond.L.Unlock() + p.producerWg.Wait() + c <- Empty{} + }() + select { + case <-c: + case <-ctx.Done(): + close(p.pipeOfDeath) + } } // Close shuts down the producer, waiting for all outstanding messages and retries // to flush. func (p *Producer) Close() { - close(p.messages) - p.flushCond.L.Lock() - for p.outstanding != 0 { - p.flushCond.Wait() - } - close(p.retries) - p.flushCond.L.Unlock() - p.producerWg.Wait() + p.CloseWithContext(context.TODO()) } // SendWithContext sends a message to the stream. Cancellation supported @@ -333,7 +297,6 @@ func (p *Producer) SendWithContext(ctx context.Context, msg *message.Message) er p.produce() select { case p.messages <- msg: - p.incOutstanding(1) return nil case <-ctx.Done(): return ctx.Err() diff --git a/producer/producer_test.go b/producer/producer_test.go new file mode 100644 index 0000000..2c127b5 --- /dev/null +++ b/producer/producer_test.go @@ -0,0 +1,123 @@ +package producer + +import ( + . "github.com/smartystreets/goconvey/convey" + + "context" + "testing" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/request" + + "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/listener" + "github.com/rewardStyle/kinetic/message" +) + +func TestProducer(t *testing.T) { + Convey("given a producer", t, func() { + k, err := kinetic.New(func(c *kinetic.Config) { + c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") + c.SetRegion("some-region") + c.SetEndpoint("http://127.0.0.1:4567") + }) + + stream := "some-producer-stream" + + err = k.CreateStream(stream, 1) + So(err, ShouldBeNil) + + err = k.WaitUntilStreamExists(context.TODO(), stream, request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + So(err, ShouldBeNil) + + shards, err := k.GetShards(stream) + So(err, ShouldBeNil) + So(len(shards), ShouldEqual, 1) + + p, err := NewProducer(func(c *Config) { + c.SetAwsConfig(k.Session.Config) + c.SetKinesisStream(stream) + c.SetBatchSize(5) + c.SetBatchTimeout(1 * time.Second) + So(err, ShouldBeNil) + }) + So(p, ShouldNotBeNil) + So(err, ShouldBeNil) + + l, err := listener.NewListener(stream, shards[0], func(c *listener.Config) { + c.SetAwsConfig(k.Session.Config) + c.SetConcurrency(10) + }) + So(l, ShouldNotBeNil) + So(err, ShouldBeNil) + + Convey("given a kinesis writer", func() { + w := p.writer.(*KinesisWriter) + + Convey("check that the writer was initialized with the correct stream name", func() { + So(w.stream, ShouldEqual, stream) + }) + + Convey("check that the writer has a valid reference to the producer", func() { + So(w.producer, ShouldEqual, p) + }) + + Convey("check that calling ensureClient twice doesn't overwrite existing client", func() { + So(w.client, ShouldBeNil) + w.ensureClient() + So(w.client, ShouldNotBeNil) + client := w.client + w.ensureClient() + So(w.client, ShouldEqual, client) + }) + }) + + Convey("check that we can send a single message after batch timeout elapses", func() { + start := time.Now() + data := "hello" + p.Send(&message.Message{ + PartitionKey: aws.String("key"), + Data: []byte(data), + }) + msg, err := l.Retrieve() + elapsed := time.Since(start) + Printf("(send took %f seconds)\n", elapsed.Seconds()) + So(err, ShouldBeNil) + So(string(msg.Data), ShouldEqual, data) + So(elapsed.Seconds(), ShouldBeGreaterThan, 1) + }) + + Convey("check that we can send a batch of messages after batch size is reached", func() { + start := time.Now() + data := []string{"hello1", "hello2", "hello3", "hello4", "hello5", "hello6"} + for _, datum := range data { + p.Send(&message.Message{ + PartitionKey: aws.String("key"), + Data: []byte(datum), + }) + } + + for i := 0; i < 5; i++ { + msg, err := l.Retrieve() + So(err, ShouldBeNil) + So(string(msg.Data), ShouldEqual, data[i]) + } + elapsed := time.Since(start) + So(elapsed.Seconds(), ShouldBeLessThan, 1) + Printf("(first 5 took %f seconds)\n", elapsed.Seconds()) + + msg, err := l.Retrieve() + So(err, ShouldBeNil) + So(string(msg.Data), ShouldEqual, data[5]) + elapsed = time.Since(start) + So(elapsed.Seconds(), ShouldBeGreaterThan, 1) + Printf("(last took %f seconds)\n", elapsed.Seconds()) + }) + + Reset(func() { + k.DeleteStream(stream) + k.WaitUntilStreamDeleted(context.TODO(), stream, request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + }) + }) +} diff --git a/producer/stats.go b/producer/stats.go index 1cc1cbb..678b43f 100644 --- a/producer/stats.go +++ b/producer/stats.go @@ -8,62 +8,40 @@ import ( // the Kinetic producer library. This was really built with rcrowley/go-metrics // in mind. type StatsCollector interface { - // histograms give us the count, sum, min, max, mean, percentiles, - // standard deviation, and variance of the data. these metrics should - // give us the total number (sum) of messages sent, failed, and dropped, - // as well as the average (mean) batch size. - - // for producer - AddSentSample(int) - AddFailedSample(int) - AddDroppedSample(int) - AddBatchSizeSample(int) - - // meters give us the count and rate of the data. these metrics should - // give us the average number of times: - // - ProvisionedThroughputExceeded was received per second - // - PutRecords was called per second - // - PutRecordProvisionedThroughputExceeded was received per second - // - PutRecords timed out per second - - // for producer + AddSent(int) + AddFailed(int) + AddDropped(int) + AddBatchSize(int) AddPutRecordsProvisionedThroughputExceeded(int) - - // for kinesiswriter AddPutRecordsCalled(int) AddProvisionedThroughputExceeded(int) AddPutRecordsTimeout(int) - // timers give us the count, sum, min, max, mean, percentiles, standard - // deviation, variance, as well as the rate of the data. - // TODO: describe these metrics better - - // for kinesis writer - AddPutRecordsTime(time.Duration) - AddPutRecordsBuildTime(time.Duration) - AddPutRecordsSendTime(time.Duration) + AddPutRecordsDuration(time.Duration) + AddPutRecordsBuildDuration(time.Duration) + AddPutRecordsSendDuration(time.Duration) } // NilStatsCollector is a stats listener that ignores all metrics. type NilStatsCollector struct{} -// AddSentSample records a count of the number of messages sent to AWS Kinesis -// by the producer. -func (l *NilStatsCollector) AddSentSample(int) {} +// AddSent records a count of the number of messages sent to AWS Kinesis by the +// producer. +func (l *NilStatsCollector) AddSent(int) {} -// AddFailedSample records a count of the number of messages that failed to be -// sent to AWS Kinesis by the producer. -func (l *NilStatsCollector) AddFailedSample(int) {} +// AddFailed records a count of the number of messages that failed to be sent to +// AWS Kinesis by the producer. +func (l *NilStatsCollector) AddFailed(int) {} -// AddDroppedSample records a count of the number of messages dropped by the +// AddDropped records a count of the number of messages dropped by the // application after multiple failures. -func (l *NilStatsCollector) AddDroppedSample(int) {} +func (l *NilStatsCollector) AddDropped(int) {} -// AddBatchSizeSample records a count of the number of messages attempted by +// AddBatchSize records a count of the number of messages attempted by // PutRecords in the producer. -func (l *NilStatsCollector) AddBatchSizeSample(int) {} +func (l *NilStatsCollector) AddBatchSize(int) {} -// AddProvisionedThroughputExceeded records the number of times the PutRecords +// AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords // API returned a ErrCodeProvisionedThroughputExceededException by the producer. func (l *NilStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} @@ -81,14 +59,14 @@ func (l *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} // configuration. func (l *NilStatsCollector) AddPutRecordsTimeout(int) {} -// AddPutRecordsTime records the duration that the PutRecords API request took. -// Only the times of successful calls are measured. -func (l *NilStatsCollector) AddPutRecordsTime(time.Duration) {} +// AddPutRecordsDuration records the duration that the PutRecords API request +// took. Only the times of successful calls are measured. +func (l *NilStatsCollector) AddPutRecordsDuration(time.Duration) {} -// AddPutRecordsBuildTime records the duration that it took to build the +// AddPutRecordsBuildDuration records the duration that it took to build the // PutRecords API request payload. -func (l *NilStatsCollector) AddPutRecordsBuildTime(time.Duration) {} +func (l *NilStatsCollector) AddPutRecordsBuildDuration(time.Duration) {} -// AddPutRecordsSendTime records the duration that it took to send the +// AddPutRecordsSendDuration records the duration that it took to send the // PutRecords API request payload. -func (l *NilStatsCollector) AddPutRecordsSendTime(time.Duration) {} +func (l *NilStatsCollector) AddPutRecordsSendDuration(time.Duration) {} From 5805f55635da7277f1b2ef1911db507e63b45bf0 Mon Sep 17 00:00:00 2001 From: Tommy Wang Date: Tue, 9 May 2017 13:43:44 -0500 Subject: [PATCH 013/109] wip --- errs/errors.go | 21 +- listener/config.go | 17 +- listener/kinesisreader.go | 357 ++++++++++++++++++++++ listener/kinesisreaderconfig.go | 27 ++ listener/listener.go | 504 +------------------------------- listener/sharditerator.go | 68 +++++ producer/config.go | 6 +- 7 files changed, 484 insertions(+), 516 deletions(-) create mode 100644 listener/kinesisreader.go create mode 100644 listener/kinesisreaderconfig.go create mode 100644 listener/sharditerator.go diff --git a/errs/errors.go b/errs/errors.go index c9e49db..3651c0b 100644 --- a/errs/errors.go +++ b/errs/errors.go @@ -19,6 +19,15 @@ var ( ) var ( + // ErrNilProducer is returned by a StreamWriter when it has not been + // correctly associated with a Producer. + ErrNilListener = errors.New("StreamReader not associated with a listener") + + // ErrProducerAlreadyAssociated is returned by a StreamWriter attempting + // to associate it with a Producer when it already has an association + // with a producer. + ErrListenerAlreadyAssociated = errors.New("StreamReader already associated with a listener") + // ErrAlreadyConsuming is returned when attempting to consume when the // Listener is already consuming. May be returned by // Retrieve/RetrieveFn. @@ -50,11 +59,6 @@ var ( ) var ( - // ErrRetryRecords is returned when the PutRecords calls requires some - // records of the batch to be retried. This failure is considered part - // of normal behavior of the Kinesis stream. - ErrRetryRecords = errors.New("PutRecords requires retry of some records in batch") - // ErrNilProducer is returned by a StreamWriter when it has not been // correctly associated with a Producer. ErrNilProducer = errors.New("StreamWriter not associated with a producer") @@ -63,9 +67,7 @@ var ( // to associate it with a Producer when it already has an association // with a producer. ErrProducerAlreadyAssociated = errors.New("StreamWriter already associated with a producer") -) -var ( // ErrNilPutRecordsResponse is returned when the PutRecords call returns // a nil response. ErrNilPutRecordsResponse = errors.New("PutRecords returned a nil response") @@ -73,4 +75,9 @@ var ( // ErrNilFailedRecordCount is returned when the PutRecords call returns // a nil FailedRecordCount. ErrNilFailedRecordCount = errors.New("GetFailedRecordCount returned a nil FailedRecordCount") + + // ErrRetryRecords is returned when the PutRecords calls requires some + // records of the batch to be retried. This failure is considered part + // of normal behavior of the Kinesis stream. + ErrRetryRecords = errors.New("PutRecords requires retry of some records in batch") ) diff --git a/listener/config.go b/listener/config.go index 0b0394f..329fbcd 100644 --- a/listener/config.go +++ b/listener/config.go @@ -17,15 +17,11 @@ type Config struct { } // NewConfig creates a new instance of Config -func NewConfig(stream, shard string) *Config { +func NewConfig() *Config { return &Config{ AwsOptions: config.DefaultAwsOptions(), listenerOptions: &listenerOptions{ - stream: stream, - shard: shard, - batchSize: 10000, concurrency: 10000, - shardIterator: NewShardIterator(), getRecordsReadTimeout: 1 * time.Second, Stats: &NilStatsCollector{}, }, @@ -39,23 +35,12 @@ func (c *Config) SetAwsConfig(config *aws.Config) { c.AwsConfig = config } -// SetBatchSize configures the batch size of the GetRecords call. -func (c *Config) SetBatchSize(batchSize int) { - c.batchSize = batchSize -} - // SetConcurrency controls the number of goroutines the Listener will spawn to // process messages. func (c *Config) SetConcurrency(concurrency int) { c.concurrency = concurrency } -// SetInitialShardIterator configures the settings used to retrieve initial -// shard iterator via the GetShardIterator call. -func (c *Config) SetInitialShardIterator(shardIterator *ShardIterator) { - c.shardIterator = shardIterator -} - // SetGetRecordsReadTimeout configures the time to wait for each successive // Read operation on the GetRecords response payload. func (c *Config) SetGetRecordsReadTimeout(timouet time.Duration) { diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go new file mode 100644 index 0000000..15a9889 --- /dev/null +++ b/listener/kinesisreader.go @@ -0,0 +1,357 @@ +package listener + +import ( + "fmt" + "io" + "sync" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/request" + "github.com/aws/aws-sdk-go/service/kinesis" + "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" + + "github.com/rewardStyle/kinetic/errs" + "github.com/rewardStyle/kinetic/logging" + "github.com/rewardStyle/kinetic/message" +) + +type kinesisReaderOptions struct { + stream string + shard string + + batchSize int + shardIterator *ShardIterator +} + +type KinesisReader struct { + *kinesisReaderOptions + + throttleSem chan Empty + nextShardIterator string + + listener *Listener + client kinesisiface.KinesisAPI + clientMu sync.Mutex +} + +func (r *KinesisReader) NewKinesisReader(stream, shard string, fn func(*KinesisReaderConfig)) (*KinesisReader, error) { + config := NewKinesisReaderConfig(stream, shard) + fn(config) + return &KinesisReader{ + kinesisReaderOptions: config.kinesisReaderOptions, + throttleSem: make(chan Empty, 5), + }, nil +} + +// AssociateListener associates the Kinesis stream writer to a producer. +func (r *KinesisReader) AssociateListener(l *Listener) error { + r.clientMu.Lock() + defer r.clientMu.Unlock() + if r.listener != nil { + return errs.ErrListenerAlreadyAssociated + } + r.listener = l + return nil +} + +// ensureClient will lazily make sure we have an AWS Kinesis client. +func (r *KinesisReader) ensureClient() error { + // From the aws-go-sdk documentation: + // http://docs.aws.amazon.com/sdk-for-go/api/aws/session/ + // + // Concurrency: + // Sessions are safe to use concurrently as long as the Session is not + // being modified. The SDK will not modify the Session once the Session + // has been created. Creating service clients concurrently from a + // shared Session is safe. + // + // We need to think through the impact of creating a new client (for + // example, after receiving an error from Kinesis) while there may be + // outstanding goroutines still processing messages. My cursory thought + // is that this is safe to do, as any outstanding messages will likely + // not interact with the Kinesis stream. At worst, we would need a lock + // around the ensureClient method to make sure that no two goroutines + // are trying to ensure the client at the same time. + // + // As we don't expose any methods (or in fact, even the Listener object + // itself) to the client through the API, I don't forsee needing to add + // this lock unless something dramatically changes about the design of + // this library. + r.clientMu.Lock() + defer r.clientMu.Unlock() + if r.client == nil { + if r.listener == nil { + return errs.ErrNilListener + } + r.client = kinesis.New(r.listener.Session) + } + return nil +} + +// ensureShardIterator will lazily make sure that we have a valid ShardIterator, +// calling the GetShardIterator API with the configured ShardIteratorType (with +// any applicable StartingSequenceNumber or Timestamp) if necessary. +// +// Not thread-safe. Only called from getRecords Care must be taken to ensure +// that only one call to Listen and Retrieve/RetrieveFn can be running at a +// time. +func (r *KinesisReader) ensureShardIterator() error { + r.ensureClient() + if r.nextShardIterator != "" { + return nil + } + + resp, err := r.client.GetShardIterator(&kinesis.GetShardIteratorInput{ + ShardId: aws.String(r.shard), // Required + ShardIteratorType: aws.String(r.shardIterator.shardIteratorType), // Required + StreamName: aws.String(r.stream), // Required + StartingSequenceNumber: r.shardIterator.getStartingSequenceNumber(), + Timestamp: r.shardIterator.getTimestamp(), + }) + if err != nil { + r.listener.LogError(err) + return err + } + if resp == nil { + return errs.ErrNilGetShardIteratorResponse + } + if resp.ShardIterator == nil { + return errs.ErrNilShardIterator + } + return r.setNextShardIterator(*resp.ShardIterator) +} + +// setNextShardIterator sets the nextShardIterator to use when calling +// GetRecords. +// +// Not thread-safe. Only called from getRecords (and ensureShardIterator, which +// is called from getRecords). Care must be taken to ensure that only one call +// to Listen and Retrieve/RetrieveFn can be running at a time. +func (r *KinesisReader) setNextShardIterator(shardIterator string) error { + if len(shardIterator) == 0 { + return errs.ErrEmptyShardIterator + } + r.nextShardIterator = shardIterator + return nil +} + +// setSequenceNumber sets the sequenceNumber of shardIterator to the last +// delivered message and updates the shardIteratorType to AT_SEQUENCE_NUMBER. +// This is only used when we need to call getShardIterator (say, to refresh the +// shard iterator). +// +// Not thread-safe. Only called from getRecords. Care must be taken to ensure +// that only one call to Listen and Retrieve/RetrieveFn can be running at a +// time. +func (r *KinesisReader) setSequenceNumber(sequenceNumber string) error { + if len(sequenceNumber) == 0 { + return errs.ErrEmptySequenceNumber + } + r.shardIterator.AtSequenceNumber(sequenceNumber) + return nil +} + +// Kinesis allows five read ops per second per shard. +// http://docs.aws.amazon.com/kinesis/latest/dev/service-sizes-and-limits.html +func (r *KinesisReader) throttle(sem chan Empty) { + sem <- Empty{} + time.AfterFunc(1*time.Second, func() { + <-sem + }) +} + +// GetRecords calls GetRecords and delivers each record into the messages +// channel. +// FIXME: Need to investigate that the timeout implementation doesn't result in +// an fd leak. Since we call Read on the HTTPResonse.Body in a select with a +// timeout channel, we do prevent ourself from blocking. Once we timeout, we +// return an error to the outer ioutil.ReadAll, which should result in a call +// to our io.ReadCloser's Close function. This will in turn call Close on the +// underlying HTTPResponse.Body. The question is whether this actually shuts +// down the TCP connection. Worst case scenario is that our client Timeout +// eventually fires and closes the socket, but this can be susceptible to FD +// exhaustion. +func (r *KinesisReader) GetRecords(batchSize int) (int, error) { + if err := r.ensureClient(); err != nil { + return 0, err + } + if err := r.ensureShardIterator(); err != nil { + return 0, err + } + + r.throttle(r.throttleSem) + + // We use the GetRecordsRequest method of creating requests to allow for + // registering custom handlers for better control over the API request. + var startReadTime time.Time + var startUnmarshalTime time.Time + start := time.Now() + req, resp := r.client.GetRecordsRequest(&kinesis.GetRecordsInput{ + Limit: aws.Int64(int64(batchSize)), + ShardIterator: aws.String(r.nextShardIterator), + }) + + // If debug is turned on, add some handlers for GetRecords logging + if r.listener.LogLevel.AtLeast(logging.LogDebug) { + req.Handlers.Send.PushBack(func(req *request.Request) { + r.listener.LogDebug("Finished GetRecords Send, took", time.Since(start)) + }) + } + + // Here, we insert a handler to be called after the Send handler and + // before the the Unmarshal handler in the aws-go-sdk library. + // + // The Send handler will call http.Client.Do() on the request, which + // blocks until the response headers have been read before returning an + // HTTPResponse. + // + // The Unmarshal handler will ultimately call ioutil.ReadAll() on the + // HTTPResponse.Body stream. + // + // Our handler wraps the HTTPResponse.Body with our own ReadCloser so + // that we can implement a timeout mechanism on the Read() call (which + // is called by the ioutil.ReadAll() function) + req.Handlers.Unmarshal.PushFront(func(req *request.Request) { + r.listener.LogDebug("Started GetRecords Unmarshal, took", time.Since(start)) + // Here, we set a timer that the initial Read() call on + // HTTPResponse.Body must return by. Note that the normal + // http.Client Timeout is still in effect. + startReadTime = time.Now() + timer := time.NewTimer(r.listener.getRecordsReadTimeout) + + req.HTTPResponse.Body = &ReadCloserWrapper{ + ReadCloser: req.HTTPResponse.Body, + OnReadFn: func(stream io.ReadCloser, b []byte) (n int, err error) { + // The OnReadFn will be called each time + // ioutil.ReadAll calls Read on the + // ReadCloserWrapper. + + // First, we set up a struct that to hold the + // results of the Read() call that can go + // through a channel + type Result struct { + n int + err error + } + + // Next, we build a channel with which to pass + // the Read() results + c := make(chan Result, 1) + + // Now, we call the Read() on the + // HTTPResponse.Body in a goroutine and feed the + // results into the channel + readStart := time.Now() + go func() { + var result Result + result.n, result.err = stream.Read(b) + c <- result + }() + + // Finally, we poll for the Read() to complete + // or the timer to elapse. + select { + case result := <-c: + // If we sucessfully Read() from the + // HTTPResponse.Body, we reset our + // timeout and return the results from + // the Read() + timer.Reset(r.listener.getRecordsReadTimeout) + n, err = result.n, result.err + r.listener.LogDebug(fmt.Sprintf("GetRecords read %d bytes, took %v", n, time.Since(readStart))) + case <-timer.C: + // If we timeout, we return an error + // that will unblock ioutil.ReadAll(). + // This will cause the Unmarshal handler + // to return an error. This error will + // propogate to the original req.Send() + // call (below) + r.listener.LogDebug(fmt.Sprintf("GetRecords read timed out after %v", time.Since(readStart))) + err = errs.ErrTimeoutReadResponseBody + case <-r.listener.pipeOfDeath: + // The pipe of death will abort any pending + // reads on a GetRecords call. + r.listener.LogDebug(fmt.Sprintf("GetRecords received pipe of death after %v", time.Since(readStart))) + err = errs.ErrPipeOfDeath + } + return + }, + OnCloseFn: func() { + r.listener.Stats.AddGetRecordsReadResponseDuration(time.Since(startReadTime)) + r.listener.LogDebug("Finished GetRecords body read, took", time.Since(start)) + startUnmarshalTime = time.Now() + }, + } + }) + + req.Handlers.Unmarshal.PushBack(func(req *request.Request) { + r.listener.Stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) + r.listener.LogDebug("Finished GetRecords Unmarshal, took", time.Since(start)) + }) + + // Send the GetRecords request + r.listener.LogDebug("Starting GetRecords Build/Sign request, took", time.Since(start)) + r.listener.Stats.AddGetRecordsCalled(1) + if err := req.Send(); err != nil { + r.listener.LogError("Error getting records:", err) + return 0, err + } + r.listener.Stats.AddGetRecordsDuration(time.Since(start)) + + // Process Records + r.listener.LogDebug(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), r.shard, time.Since(start))) + if resp == nil { + return 0, errs.ErrNilGetRecordsResponse + } + delivered := 0 + r.listener.Stats.AddBatchSize(len(resp.Records)) + for _, record := range resp.Records { + if record != nil { + // Allow (only) a pipeOfDeath to trigger an instance + // shutdown of the loop to deliver messages. Otherwise, + // a normal cancellation will not prevent getRecords + // from completing the delivery of the current batch of + // records. + select { + case r.listener.messages <- message.FromRecord(record): + delivered++ + r.listener.Stats.AddConsumed(1) + if record.SequenceNumber != nil { + // We can safely ignore if this call returns + // error, as if we somehow receive an empty + // sequence number from AWS, we will simply not + // set it. At worst, this causes us to + // reprocess this record if we happen to refresh + // the iterator. + r.setSequenceNumber(*record.SequenceNumber) + } + case <-r.listener.pipeOfDeath: + r.listener.LogInfo(fmt.Sprintf("getRecords received pipe of death while delivering messages, %d delivered, ~%d dropped", delivered, len(resp.Records)-delivered)) + return delivered, errs.ErrPipeOfDeath + } + } + } + if resp.NextShardIterator != nil { + // TODO: According to AWS docs: + // http://docs.aws.amazon.com/sdk-for-go/api/service/kinesis/#GetRecordsOutput + // + // NextShardIterator: The next position in the shard + // from which to start sequentially reading data + // records. If set to null, the shard has been closed + // and the requested iterator will not return any more + // data. + // + // When dealing with streams that will merge or split, + // we need to detect that the shard has closed and + // notify the client library. + // + // TODO: I don't know if we should be ignoring an error returned + // by setShardIterator in case of an empty shard iterator in the + // response. There isn't much we can do, and the best path for + // recovery may be simply to reprocess the batch and see if we + // get a valid NextShardIterator from AWS the next time around. + r.setNextShardIterator(*resp.NextShardIterator) + } + return delivered, nil +} diff --git a/listener/kinesisreaderconfig.go b/listener/kinesisreaderconfig.go new file mode 100644 index 0000000..0ed091a --- /dev/null +++ b/listener/kinesisreaderconfig.go @@ -0,0 +1,27 @@ +package listener + +type KinesisReaderConfig struct { + *kinesisReaderOptions +} + +func NewKinesisReaderConfig(stream, shard string) *KinesisReaderConfig { + return &KinesisReaderConfig{ + kinesisReaderOptions: &kinesisReaderOptions{ + stream: stream, + shard: shard, + batchSize: 10000, + shardIterator: NewShardIterator(), + }, + } +} + +// SetBatchSize configures the batch size of the GetRecords call. +func (c *KinesisReaderConfig) SetBatchSize(batchSize int) { + c.batchSize = batchSize +} + +// SetInitialShardIterator configures the settings used to retrieve initial +// shard iterator via the GetShardIterator call. +func (c *KinesisReaderConfig) SetInitialShardIterator(shardIterator *ShardIterator) { + c.shardIterator = shardIterator +} diff --git a/listener/listener.go b/listener/listener.go index 35a8bcf..bd8ac7a 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -2,24 +2,25 @@ package listener import ( "context" - "fmt" - "io" "net" "sync" "time" - "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/awserr" - "github.com/aws/aws-sdk-go/aws/request" "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" "github.com/rewardStyle/kinetic/errs" "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" ) +type StreamReader interface { + AssociateListener(listener *Listener) error + GetRecords(batchSize int) (int, error) + ensureClient() error +} + // Empty is used a as a dummy type for counting semaphore channels. type Empty struct{} @@ -27,77 +28,10 @@ type Empty struct{} // RetrieveFn. type MessageFn func([]byte, *sync.WaitGroup) -// ShardIterator represents the settings used to retrieve a shard iterator from -// the GetShardIterator API. -type ShardIterator struct { - shardIteratorType string - sequenceNumber string - timestamp time.Time -} - -// NewShardIterator creates a new ShardIterator. The default shard iterator -// type is TRIM_HORIZON. -func NewShardIterator() *ShardIterator { - return &ShardIterator{ - shardIteratorType: "TRIM_HORIZON", - } -} - -// TrimHorizon sets the shard iterator to TRIM_HORIZON. -func (it *ShardIterator) TrimHorizon() *ShardIterator { - it.shardIteratorType = "TRIM_HORIZON" - return it -} - -// Latest sets the shard iterator to LATEST. -func (it *ShardIterator) Latest() *ShardIterator { - it.shardIteratorType = "LATEST" - return it -} - -// AtSequenceNumber sets the shard iterator to AT_SEQUENCE_NUMBER. -func (it *ShardIterator) AtSequenceNumber(sequenceNumber string) *ShardIterator { - it.shardIteratorType = "AT_SEQUENCE_NUMBER" - it.sequenceNumber = sequenceNumber - return it -} - -// AfterSequenceNumber sets the shard iterator to AFTER_SEQUENCE_NUMBER. -func (it *ShardIterator) AfterSequenceNumber(sequenceNumber string) *ShardIterator { - it.shardIteratorType = "AFTER_SEQUENCE_NUMBER" - it.sequenceNumber = sequenceNumber - return it -} - -// AtTimestamp sets the shard iterator to AT_TIMESTAMP. -func (it *ShardIterator) AtTimestamp(timestamp time.Time) *ShardIterator { - it.shardIteratorType = "AT_TIMESTAMP" - it.timestamp = timestamp - return it -} - -func (it *ShardIterator) getStartingSequenceNumber() *string { - if it.sequenceNumber == "" { - return nil - } - return aws.String(it.sequenceNumber) -} - -func (it *ShardIterator) getTimestamp() *time.Time { - if it.timestamp.IsZero() { - return nil - } - return aws.Time(it.timestamp) -} - type listenerOptions struct { - stream string - shard string - - batchSize int concurrency int - shardIterator *ShardIterator getRecordsReadTimeout time.Duration + reader StreamReader Stats StatsCollector } @@ -107,449 +41,39 @@ type Listener struct { *listenerOptions *logging.LogHelper - nextShardIterator string - messages chan *message.Message concurrencySem chan Empty - throttleSem chan Empty pipeOfDeath chan Empty consuming bool consumingMu sync.Mutex - clientMu sync.Mutex - client kinesisiface.KinesisAPI - Session *session.Session + Session *session.Session } // NewListener creates a new listener for listening to message on a Kinesis // stream. -func NewListener(stream, shard string, fn func(*Config)) (*Listener, error) { - config := NewConfig(stream, shard) +func NewListener(fn func(*Config)) (*Listener, error) { + config := NewConfig() fn(config) session, err := config.GetSession() if err != nil { return nil, err } - return &Listener{ + l := &Listener{ listenerOptions: config.listenerOptions, LogHelper: &logging.LogHelper{ LogLevel: config.LogLevel, Logger: session.Config.Logger, }, concurrencySem: make(chan Empty, config.concurrency), - throttleSem: make(chan Empty, 5), pipeOfDeath: make(chan Empty), Session: session, - }, nil -} - -// setNextShardIterator sets the nextShardIterator to use when calling -// GetRecords. -// -// Not thread-safe. Only called from getRecords (and ensureShardIterator, which -// is called from getRecords). Care must be taken to ensure that only one call -// to Listen and Retrieve/RetrieveFn can be running at a time. -func (l *Listener) setNextShardIterator(shardIterator string) error { - if len(shardIterator) == 0 { - return errs.ErrEmptyShardIterator - } - l.nextShardIterator = shardIterator - return nil -} - -// setSequenceNumber sets the sequenceNumber of shardIterator to the last -// delivered message and updates the shardIteratorType to AT_SEQUENCE_NUMBER. -// This is only used when we need to call getShardIterator (say, to refresh the -// shard iterator). -// -// Not thread-safe. Only called from getRecords. Care must be taken to ensure -// that only one call to Listen and Retrieve/RetrieveFn can be running at a -// time. -func (l *Listener) setSequenceNumber(sequenceNumber string) error { - if len(sequenceNumber) == 0 { - return errs.ErrEmptySequenceNumber - } - l.shardIterator.AtSequenceNumber(sequenceNumber) - return nil -} - -// ensureClient will lazily make sure we have an AWS Kinesis client. -func (l *Listener) ensureClient() { - // From the aws-go-sdk documentation: - // http://docs.aws.amazon.com/sdk-for-go/api/aws/session/ - // - // Concurrency: - // Sessions are safe to use concurrently as long as the Session is not - // being modified. The SDK will not modify the Session once the Session - // has been created. Creating service clients concurrently from a - // shared Session is safe. - // - // We need to think through the impact of creating a new client (for - // example, after receiving an error from Kinesis) while there may be - // outstanding goroutines still processing messages. My cursory thought - // is that this is safe to do, as any outstanding messages will likely - // not interact with the Kinesis stream. At worst, we would need a lock - // around the ensureClient method to make sure that no two goroutines - // are trying to ensure the client at the same time. - // - // As we don't expose any methods (or in fact, even the Listener object - // itself) to the client through the API, I don't forsee needing to add - // this lock unless something dramatically changes about the design of - // this library. - l.clientMu.Lock() - defer l.clientMu.Unlock() - if l.client == nil { - l.client = kinesis.New(l.Session) - } -} - -// ensureShardIterator will lazily make sure that we have a valid ShardIterator, -// calling the GetShardIterator API with the configured ShardIteratorType (with -// any applicable StartingSequenceNumber or Timestamp) if necessary. -// -// Not thread-safe. Only called from getRecords Care must be taken to ensure -// that only one call to Listen and Retrieve/RetrieveFn can be running at a -// time. -func (l *Listener) ensureShardIterator() error { - l.ensureClient() - if l.nextShardIterator != "" { - return nil - } - - resp, err := l.client.GetShardIterator(&kinesis.GetShardIteratorInput{ - ShardId: aws.String(l.shard), // Required - ShardIteratorType: aws.String(l.shardIterator.shardIteratorType), // Required - StreamName: aws.String(l.stream), // Required - StartingSequenceNumber: l.shardIterator.getStartingSequenceNumber(), - Timestamp: l.shardIterator.getTimestamp(), - }) - if err != nil { - l.LogError(err) - return err - } - if resp == nil { - return errs.ErrNilGetShardIteratorResponse - } - if resp.ShardIterator == nil { - return errs.ErrNilShardIterator - } - return l.setNextShardIterator(*resp.ShardIterator) -} - -// 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(sem chan Empty) { - sem <- Empty{} - time.AfterFunc(1*time.Second, func() { - <-sem - }) -} - -/* -func (l *Listener) newGetRecords(batchSize int) (int, error) { - l.ensureClient() - if err := l.ensureShardIterator(); err != nil { - return 0, err - } - l.throttle(l.throttleSem) - - var startUnmarshalTime time.Time - start := time.Now() - - // We use the GetRecordsRequest method of creating requests to allow for - // registering custom handlers for better control over the API request. - req, resp := l.client.GetRecordsRequest(&kinesis.GetRecordsInput{ - Limit: aws.Int64(int64(batchSize)), - ShardIterator: aws.String(l.nextShardIterator), - }) - - // If debug is turned on, add some handlers for GetRecords logging - if l.LogLevel.AtLeast(logging.LogDebug) { - req.Handlers.Send.PushBack(func(r *request.Request) { - l.LogDebug("Finished GetRecords Send, took", time.Since(start)) - }) - } - - // Add some profiling timers for metrics - // TODO: Original implementation was able to measure read time and - // unmarshal time separately. Because we no longer hook into the - // ReadCloser's Close method, we can no longer measure them separately. - // Instead, the UnmarshalDuration will measure the total time to read - // and unmarshal the data. We should delete the corresponding stat for - // the read time. - req.Handlers.Unmarshal.PushFront(func(r *request.Request) { - l.LogDebug("Started GetRecords Unmarshal, took", time.Since(start)) - startUnmarshalTime = time.Now() - }) - - req.Handlers.Unmarshal.PushBack(func(r *request.Request) { - l.LogDebug("Finished GetRecords Unmarshal, took", time.Since(start)) - l.Stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) - }) - - // This should replace any WithResponseReadTimeout calls within the AWS - // SDK in customizations.go. - req.ApplyOptions(request.WithResponseReadTimeout(l.getRecordsReadTimeout)) - - // Send the GetRecords request - l.LogDebug("Starting GetRecords Build/Sign request, took", time.Since(start)) - l.Stats.AddGetRecordsCalled(1) - if err := req.Send(); err != nil { - l.LogError("Error getting records:", err) - return 0, err } - l.Stats.AddGetRecordsDuration(time.Since(start)) - - // Process Records - l.LogDebug(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), l.shard, time.Since(start))) - if resp == nil { - return 0, errs.ErrNilGetRecordsResponse - } - delivered := 0 - l.Stats.AddBatchSize(len(resp.Records)) - for _, record := range resp.Records { - if record != nil { - // Allow (only) a pipeOfDeath to trigger an instance - // shutdown of the loop to deliver messages. Otherwise, - // a normal cancellation will not prevent getRecords - // from completing the delivery of the current batch of - // records. - select { - case l.messages <- message.FromRecord(record): - delivered++ - l.Stats.AddConsumed(1) - if record.SequenceNumber != nil { - // We can safely ignore if this call returns - // error, as if we somehow receive an empty - // sequence number from AWS, we will simply not - // set it. At worst, this causes us to - // reprocess this record if we happen to refresh - // the iterator. - l.setSequenceNumber(*record.SequenceNumber) - } - case <-l.pipeOfDeath: - l.LogInfo(fmt.Sprintf("getRecords received pipe of death while delivering messages, %d delivered, ~%d dropped", delivered, len(resp.Records)-delivered)) - return delivered, errs.ErrPipeOfDeath - } - } - } - if resp.NextShardIterator != nil { - // TODO: According to AWS docs: - // http://docs.aws.amazon.com/sdk-for-go/api/service/kinesis/#GetRecordsOutput - // - // NextShardIterator: The next position in the shard - // from which to start sequentially reading data - // records. If set to null, the shard has been closed - // and the requested iterator will not return any more - // data. - // - // When dealing with streams that will merge or split, - // we need to detect that the shard has closed and - // notify the client library. - // - // TODO: I don't know if we should be ignoring an error returned - // by setShardIterator in case of an empty shard iterator in the - // response. There isn't much we can do, and the best path for - // recovery may be simply to reprocess the batch and see if we - // get a valid NextShardIterator from AWS the next time around. - l.setNextShardIterator(*resp.NextShardIterator) - } - return delivered, nil -} -*/ - -// getRecords calls GetRecords and delivers each record into the messages -// channel. -// FIXME: Need to investigate that the timeout implementation doesn't result in -// an fd leak. Since we call Read on the HTTPResonse.Body in a select with a -// timeout channel, we do prevent ourself from blocking. Once we timeout, we -// return an error to the outer ioutil.ReadAll, which should result in a call -// to our io.ReadCloser's Close function. This will in turn call Close on the -// underlying HTTPResponse.Body. The question is whether this actually shuts -// down the TCP connection. Worst case scenario is that our client Timeout -// eventually fires and closes the socket, but this can be susceptible to FD -// exhaustion. -func (l *Listener) getRecords(batchSize int) (int, error) { - l.ensureClient() - if err := l.ensureShardIterator(); err != nil { - return 0, err - } - - l.throttle(l.throttleSem) - - // We use the GetRecordsRequest method of creating requests to allow for - // registering custom handlers for better control over the API request. - var startReadTime time.Time - var startUnmarshalTime time.Time - start := time.Now() - req, resp := l.client.GetRecordsRequest(&kinesis.GetRecordsInput{ - Limit: aws.Int64(int64(batchSize)), - ShardIterator: aws.String(l.nextShardIterator), - }) - - // If debug is turned on, add some handlers for GetRecords logging - if l.LogLevel.AtLeast(logging.LogDebug) { - req.Handlers.Send.PushBack(func(r *request.Request) { - l.LogDebug("Finished GetRecords Send, took", time.Since(start)) - }) - } - - // Here, we insert a handler to be called after the Send handler and - // before the the Unmarshal handler in the aws-go-sdk library. - // - // The Send handler will call http.Client.Do() on the request, which - // blocks until the response headers have been read before returning an - // HTTPResponse. - // - // The Unmarshal handler will ultimately call ioutil.ReadAll() on the - // HTTPResponse.Body stream. - // - // Our handler wraps the HTTPResponse.Body with our own ReadCloser so - // that we can implement a timeout mechanism on the Read() call (which - // is called by the ioutil.ReadAll() function) - req.Handlers.Unmarshal.PushFront(func(r *request.Request) { - l.LogDebug("Started GetRecords Unmarshal, took", time.Since(start)) - // Here, we set a timer that the initial Read() call on - // HTTPResponse.Body must return by. Note that the normal - // http.Client Timeout is still in effect. - startReadTime = time.Now() - timer := time.NewTimer(l.getRecordsReadTimeout) - - r.HTTPResponse.Body = &ReadCloserWrapper{ - ReadCloser: r.HTTPResponse.Body, - OnReadFn: func(stream io.ReadCloser, b []byte) (n int, err error) { - // The OnReadFn will be called each time - // ioutil.ReadAll calls Read on the - // ReadCloserWrapper. - - // First, we set up a struct that to hold the - // results of the Read() call that can go - // through a channel - type Result struct { - n int - err error - } - - // Next, we build a channel with which to pass - // the Read() results - c := make(chan Result, 1) - - // Now, we call the Read() on the - // HTTPResponse.Body in a goroutine and feed the - // results into the channel - readStart := time.Now() - go func() { - var result Result - result.n, result.err = stream.Read(b) - c <- result - }() - - // Finally, we poll for the Read() to complete - // or the timer to elapse. - select { - case result := <-c: - // If we sucessfully Read() from the - // HTTPResponse.Body, we reset our - // timeout and return the results from - // the Read() - timer.Reset(l.getRecordsReadTimeout) - n, err = result.n, result.err - l.LogDebug(fmt.Sprintf("GetRecords read %d bytes, took %v", n, time.Since(readStart))) - case <-timer.C: - // If we timeout, we return an error - // that will unblock ioutil.ReadAll(). - // This will cause the Unmarshal handler - // to return an error. This error will - // propogate to the original req.Send() - // call (below) - l.LogDebug(fmt.Sprintf("GetRecords read timed out after %v", time.Since(readStart))) - err = errs.ErrTimeoutReadResponseBody - case <-l.pipeOfDeath: - // The pipe of death will abort any pending - // reads on a GetRecords call. - l.LogDebug(fmt.Sprintf("GetRecords received pipe of death after %v", time.Since(readStart))) - err = errs.ErrPipeOfDeath - } - return - }, - OnCloseFn: func() { - l.Stats.AddGetRecordsReadResponseDuration(time.Since(startReadTime)) - l.LogDebug("Finished GetRecords body read, took", time.Since(start)) - startUnmarshalTime = time.Now() - }, - } - }) - - req.Handlers.Unmarshal.PushBack(func(r *request.Request) { - l.Stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) - l.LogDebug("Finished GetRecords Unmarshal, took", time.Since(start)) - }) - - // Send the GetRecords request - l.LogDebug("Starting GetRecords Build/Sign request, took", time.Since(start)) - l.Stats.AddGetRecordsCalled(1) - if err := req.Send(); err != nil { - l.LogError("Error getting records:", err) - return 0, err - } - l.Stats.AddGetRecordsDuration(time.Since(start)) - - // Process Records - l.LogDebug(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), l.shard, time.Since(start))) - if resp == nil { - return 0, errs.ErrNilGetRecordsResponse - } - delivered := 0 - l.Stats.AddBatchSize(len(resp.Records)) - for _, record := range resp.Records { - if record != nil { - // Allow (only) a pipeOfDeath to trigger an instance - // shutdown of the loop to deliver messages. Otherwise, - // a normal cancellation will not prevent getRecords - // from completing the delivery of the current batch of - // records. - select { - case l.messages <- message.FromRecord(record): - delivered++ - l.Stats.AddConsumed(1) - if record.SequenceNumber != nil { - // We can safely ignore if this call returns - // error, as if we somehow receive an empty - // sequence number from AWS, we will simply not - // set it. At worst, this causes us to - // reprocess this record if we happen to refresh - // the iterator. - l.setSequenceNumber(*record.SequenceNumber) - } - case <-l.pipeOfDeath: - l.LogInfo(fmt.Sprintf("getRecords received pipe of death while delivering messages, %d delivered, ~%d dropped", delivered, len(resp.Records)-delivered)) - return delivered, errs.ErrPipeOfDeath - } - } - } - if resp.NextShardIterator != nil { - // TODO: According to AWS docs: - // http://docs.aws.amazon.com/sdk-for-go/api/service/kinesis/#GetRecordsOutput - // - // NextShardIterator: The next position in the shard - // from which to start sequentially reading data - // records. If set to null, the shard has been closed - // and the requested iterator will not return any more - // data. - // - // When dealing with streams that will merge or split, - // we need to detect that the shard has closed and - // notify the client library. - // - // TODO: I don't know if we should be ignoring an error returned - // by setShardIterator in case of an empty shard iterator in the - // response. There isn't much we can do, and the best path for - // recovery may be simply to reprocess the batch and see if we - // get a valid NextShardIterator from AWS the next time around. - l.setNextShardIterator(*resp.NextShardIterator) + if err := l.reader.AssociateListener(l); err != nil { + return nil, err } - return delivered, nil + return l, nil } // startConsuming will initialize the consumer and set consuming to true if diff --git a/listener/sharditerator.go b/listener/sharditerator.go new file mode 100644 index 0000000..df5cad8 --- /dev/null +++ b/listener/sharditerator.go @@ -0,0 +1,68 @@ +package listener + +import ( + "time" +) + +// ShardIterator represents the settings used to retrieve a shard iterator from +// the GetShardIterator API. +type ShardIterator struct { + shardIteratorType string + sequenceNumber string + timestamp time.Time +} + +// NewShardIterator creates a new ShardIterator. The default shard iterator +// type is TRIM_HORIZON. +func NewShardIterator() *ShardIterator { + return &ShardIterator{ + shardIteratorType: "TRIM_HORIZON", + } +} + +// TrimHorizon sets the shard iterator to TRIM_HORIZON. +func (it *ShardIterator) TrimHorizon() *ShardIterator { + it.shardIteratorType = "TRIM_HORIZON" + return it +} + +// Latest sets the shard iterator to LATEST. +func (it *ShardIterator) Latest() *ShardIterator { + it.shardIteratorType = "LATEST" + return it +} + +// AtSequenceNumber sets the shard iterator to AT_SEQUENCE_NUMBER. +func (it *ShardIterator) AtSequenceNumber(sequenceNumber string) *ShardIterator { + it.shardIteratorType = "AT_SEQUENCE_NUMBER" + it.sequenceNumber = sequenceNumber + return it +} + +// AfterSequenceNumber sets the shard iterator to AFTER_SEQUENCE_NUMBER. +func (it *ShardIterator) AfterSequenceNumber(sequenceNumber string) *ShardIterator { + it.shardIteratorType = "AFTER_SEQUENCE_NUMBER" + it.sequenceNumber = sequenceNumber + return it +} + +// AtTimestamp sets the shard iterator to AT_TIMESTAMP. +func (it *ShardIterator) AtTimestamp(timestamp time.Time) *ShardIterator { + it.shardIteratorType = "AT_TIMESTAMP" + it.timestamp = timestamp + return it +} + +func (it *ShardIterator) getStartingSequenceNumber() *string { + if it.sequenceNumber == "" { + return nil + } + return aws.String(it.sequenceNumber) +} + +func (it *ShardIterator) getTimestamp() *time.Time { + if it.timestamp.IsZero() { + return nil + } + return aws.Time(it.timestamp) +} diff --git a/producer/config.go b/producer/config.go index 974cbd1..107b16d 100644 --- a/producer/config.go +++ b/producer/config.go @@ -74,9 +74,9 @@ func (c *Config) SetKinesisStream(stream string) { } // SetFirehoseStream sets the producer to write to the given Firehose stream. -// func (c *Config) SetFirehoseStream(stream string) { -// c.writer = NewFirehoseWriter(stream) -// } +func (c *Config) SetFirehoseStream(stream string) { + c.writer = NewFirehoseWriter(stream) +} // SetLogLevel configures both the SDK and Kinetic log levels. func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { From d4de4d0fa1c85f1a0369370a3bf1bcdc303835a8 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 10 May 2017 16:09:25 -0500 Subject: [PATCH 014/109] Stabilized StreamReader abstraction changes --- listener/config_test.go | 27 +---------- listener/kinesisreader.go | 16 ++++++- listener/listener.go | 26 +++-------- listener/listener_test.go | 83 ++++++++-------------------------- listener/sharditerator.go | 2 + listener/sharditerator_test.go | 57 +++++++++++++++++++++++ producer/producer_test.go | 2 +- 7 files changed, 99 insertions(+), 114 deletions(-) create mode 100644 listener/sharditerator_test.go diff --git a/listener/config_test.go b/listener/config_test.go index c7168e2..91bb2b2 100644 --- a/listener/config_test.go +++ b/listener/config_test.go @@ -37,16 +37,10 @@ func getSession(config *Config) *session.Session { func TestNewConfig(t *testing.T) { Convey("given a new listener config", t, func() { - stream := "some-stream" - shard := "some-shard" - config := NewConfig(stream, shard) + config := NewConfig() Convey("check the default values for its non-zero config", func() { - So(config.stream, ShouldEqual, stream) - So(config.shard, ShouldEqual, shard) - So(config.batchSize, ShouldEqual, 10000) So(config.concurrency, ShouldEqual, 10000) - So(config.shardIterator.shardIteratorType, ShouldEqual, "TRIM_HORIZON") So(config.getRecordsReadTimeout, ShouldEqual, 1*time.Second) So(config.Stats, ShouldHaveSameTypeAs, &NilStatsCollector{}) So(config.LogLevel.Value(), ShouldEqual, logging.LogOff) @@ -75,30 +69,11 @@ func TestNewConfig(t *testing.T) { So(aws.StringValue(sess.Config.Endpoint), ShouldEqual, "bogus-endpoint") }) - Convey("check that we can set the batch size", func() { - config.SetBatchSize(1000) - So(config.batchSize, ShouldEqual, 1000) - }) - Convey("check that we can set the concurrency limit", func() { config.SetConcurrency(50) So(config.concurrency, ShouldEqual, 50) }) - Convey("check that the default shard iterator is TRIM_HORIZON", func() { - config.SetInitialShardIterator(NewShardIterator()) - So(config.shardIterator.shardIteratorType, ShouldEqual, "TRIM_HORIZON") - So(config.shardIterator.getStartingSequenceNumber(), ShouldBeNil) - So(config.shardIterator.getTimestamp(), ShouldBeNil) - }) - - Convey("check that we can set the initial shard iterator (to LATEST)", func() { - config.SetInitialShardIterator(NewShardIterator().Latest()) - So(config.shardIterator.shardIteratorType, ShouldEqual, "LATEST") - So(config.shardIterator.getStartingSequenceNumber(), ShouldBeNil) - So(config.shardIterator.getTimestamp(), ShouldBeNil) - }) - Convey("check that we can set the read timeout for the GetRecords request", func() { config.SetGetRecordsReadTimeout(10 * time.Second) So(config.getRecordsReadTimeout, ShouldEqual, 10*time.Second) diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index 15a9889..a5b0b6a 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -7,6 +7,7 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/awserr" "github.com/aws/aws-sdk-go/aws/request" "github.com/aws/aws-sdk-go/service/kinesis" "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" @@ -172,7 +173,7 @@ func (r *KinesisReader) throttle(sem chan Empty) { // down the TCP connection. Worst case scenario is that our client Timeout // eventually fires and closes the socket, but this can be susceptible to FD // exhaustion. -func (r *KinesisReader) GetRecords(batchSize int) (int, error) { +func (r *KinesisReader) GetRecords() (int, error) { if err := r.ensureClient(); err != nil { return 0, err } @@ -187,8 +188,9 @@ func (r *KinesisReader) GetRecords(batchSize int) (int, error) { var startReadTime time.Time var startUnmarshalTime time.Time start := time.Now() + req, resp := r.client.GetRecordsRequest(&kinesis.GetRecordsInput{ - Limit: aws.Int64(int64(batchSize)), + Limit: aws.Int64(int64(r.batchSize)), ShardIterator: aws.String(r.nextShardIterator), }) @@ -295,6 +297,12 @@ func (r *KinesisReader) GetRecords(batchSize int) (int, error) { r.listener.Stats.AddGetRecordsCalled(1) if err := req.Send(); err != nil { r.listener.LogError("Error getting records:", err) + switch err.(awserr.Error).Code() { + case kinesis.ErrCodeProvisionedThroughputExceededException: + r.listener.Stats.AddProvisionedThroughputExceeded(1) + default: + r.listener.LogDebug("Received AWS error:", err.Error()) + } return 0, err } r.listener.Stats.AddGetRecordsDuration(time.Since(start)) @@ -355,3 +363,7 @@ func (r *KinesisReader) GetRecords(batchSize int) (int, error) { } return delivered, nil } + +func (r *KinesisReader) getBatchSize() int { + return r.batchSize +} diff --git a/listener/listener.go b/listener/listener.go index bd8ac7a..366f871 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -6,10 +6,6 @@ import ( "sync" "time" - "github.com/aws/aws-sdk-go/aws/awserr" - "github.com/aws/aws-sdk-go/aws/session" - "github.com/aws/aws-sdk-go/service/kinesis" - "github.com/rewardStyle/kinetic/errs" "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" @@ -19,6 +15,7 @@ type StreamReader interface { AssociateListener(listener *Listener) error GetRecords(batchSize int) (int, error) ensureClient() error + getBatchSize() int } // Empty is used a as a dummy type for counting semaphore channels. @@ -36,7 +33,7 @@ type listenerOptions struct { Stats StatsCollector } -// Listener polls the Kinesis stream for messages. +// Listener polls the StreamReader for messages. type Listener struct { *listenerOptions *logging.LogHelper @@ -47,12 +44,9 @@ type Listener struct { consuming bool consumingMu sync.Mutex - - Session *session.Session } -// NewListener creates a new listener for listening to message on a Kinesis -// stream. +// NewListener creates a new listener for listening to message on a StreamReader. func NewListener(fn func(*Config)) (*Listener, error) { config := NewConfig() fn(config) @@ -68,7 +62,6 @@ func NewListener(fn func(*Config)) (*Listener, error) { }, concurrencySem: make(chan Empty, config.concurrency), pipeOfDeath: make(chan Empty), - Session: session, } if err := l.reader.AssociateListener(l); err != nil { return nil, err @@ -83,7 +76,7 @@ func (l *Listener) startConsuming() bool { defer l.consumingMu.Unlock() if !l.consuming { l.consuming = true - l.messages = make(chan *message.Message, l.batchSize) + l.messages = make(chan *message.Message, l.reader.getBatchSize()) return true } return false @@ -131,7 +124,7 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, e if !ok { return nil, err } - n, err := l.getRecords(1) + n, err := l.reader.GetRecords(1) if err != nil { return nil, err } @@ -201,7 +194,7 @@ func (l *Listener) consume(ctx context.Context) { if !ok { break stop } - _, err := l.getRecords(l.batchSize) + _, err := l.reader.GetRecords(l.reader.getBatchSize()) if err != nil { switch err := err.(type) { case net.Error: @@ -219,13 +212,6 @@ func (l *Listener) consume(ctx context.Context) { default: l.LogError("Received error:", err.Error()) } - case awserr.Error: - switch err.Code() { - case kinesis.ErrCodeProvisionedThroughputExceededException: - l.Stats.AddProvisionedThroughputExceeded(1) - default: - l.LogError("Received AWS error:", err.Error()) - } default: l.LogError("Received unknown error:", err.Error()) } diff --git a/listener/listener_test.go b/listener/listener_test.go index 5f5dbea..5a902c2 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -21,11 +21,12 @@ import ( ) func putRecord(l *Listener, b []byte) (*string, error) { - l.ensureClient() - resp, err := l.client.PutRecord(&kinesis.PutRecordInput{ + kReader := l.reader.(*KinesisReader) + l.reader.ensureClient() + resp, err := kReader.client.PutRecord(&kinesis.PutRecordInput{ Data: b, PartitionKey: aws.String("dummy"), - StreamName: aws.String(l.stream), + StreamName: aws.String(kReader.stream), }) if err != nil { return nil, err @@ -33,54 +34,6 @@ func putRecord(l *Listener, b []byte) (*string, error) { return resp.SequenceNumber, nil } -func TestShardIterator(t *testing.T) { - Convey("given a new shard iterator", t, func() { - it := NewShardIterator() - - Convey("check that the default shard iterator type is TRIM_HORIZON", func() { - So(it.shardIteratorType, ShouldEqual, "TRIM_HORIZON") - So(it.getStartingSequenceNumber(), ShouldBeNil) - So(it.getTimestamp(), ShouldBeNil) - }) - - Convey("check that we can explicitly set it to TRIM_HORIZON", func() { - it = it.TrimHorizon() - So(it.shardIteratorType, ShouldEqual, "TRIM_HORIZON") - So(it.getStartingSequenceNumber(), ShouldBeNil) - So(it.getTimestamp(), ShouldBeNil) - }) - - Convey("check that we can explicitly set it to LATEST", func() { - it = it.Latest() - So(it.shardIteratorType, ShouldEqual, "LATEST") - So(it.getStartingSequenceNumber(), ShouldBeNil) - So(it.getTimestamp(), ShouldBeNil) - }) - - Convey("check that we can explicitly set it to AT_SEQEUENCE_NUMBER", func() { - it = it.AtSequenceNumber("some-sequence") - So(it.shardIteratorType, ShouldEqual, "AT_SEQUENCE_NUMBER") - So(aws.StringValue(it.getStartingSequenceNumber()), ShouldEqual, "some-sequence") - So(it.getTimestamp(), ShouldBeNil) - }) - - Convey("check that we can explicitly set it to AFTER_SEQEUENCE_NUMBER", func() { - it = it.AfterSequenceNumber("some-sequence") - So(it.shardIteratorType, ShouldEqual, "AFTER_SEQUENCE_NUMBER") - So(aws.StringValue(it.getStartingSequenceNumber()), ShouldEqual, "some-sequence") - So(it.getTimestamp(), ShouldBeNil) - }) - - Convey("check that we can explicitly set it to AT_TIMESTAMP", func() { - n := time.Now() - it = it.AtTimestamp(n) - So(it.shardIteratorType, ShouldEqual, "AT_TIMESTAMP") - So(aws.TimeValue(it.getTimestamp()).Equal(n), ShouldBeTrue) - So(it.getStartingSequenceNumber(), ShouldBeNil) - }) - }) -} - func TestListener(t *testing.T) { Convey("given a listener", t, func() { k, err := kinetic.New(func(c *kinetic.Config) { @@ -101,7 +54,7 @@ func TestListener(t *testing.T) { So(err, ShouldBeNil) So(len(shards), ShouldEqual, 1) - l, err := NewListener(stream, shards[0], func(c *Config) { + l, err := NewListener(func(c *Config) { c.SetAwsConfig(k.Session.Config) c.SetConcurrency(10) }) @@ -109,28 +62,28 @@ func TestListener(t *testing.T) { So(err, ShouldBeNil) Convey("check that calling ensureClient twice doesn't overwrite existing client", func() { - So(l.client, ShouldBeNil) - l.ensureClient() - So(l.client, ShouldNotBeNil) - client := l.client - l.ensureClient() - So(l.client, ShouldEqual, client) + So(l.reader.(*KinesisReader).client, ShouldBeNil) + l.reader.ensureClient() + So(l.reader.(*KinesisReader).client, ShouldNotBeNil) + client := l.reader.(*KinesisReader).client + l.reader.ensureClient() + So(l.reader.(*KinesisReader).client, ShouldEqual, client) }) Convey("check that setting an empty shard iterator returns an error", func() { - err := l.setNextShardIterator("") + err := l.reader.(*KinesisReader).setNextShardIterator("") So(err, ShouldEqual, errs.ErrEmptyShardIterator) }) Convey("check that setting an empty sequence number returns an error", func() { - err := l.setSequenceNumber("") + err := l.reader.(*KinesisReader).setSequenceNumber("") So(err, ShouldEqual, errs.ErrEmptySequenceNumber) }) Convey("check that we can get the TRIM_HORIZON shard iterator", func() { - err := l.ensureShardIterator() + err := l.reader.(*KinesisReader).ensureShardIterator() So(err, ShouldBeNil) - So(l.nextShardIterator, ShouldNotBeEmpty) + So(l.reader.(*KinesisReader).nextShardIterator, ShouldNotBeEmpty) }) Convey("check that we can retrieve records one by one", func() { @@ -143,8 +96,8 @@ func TestListener(t *testing.T) { So(err, ShouldBeNil) So(string(msg.Data), ShouldEqual, datum) Convey(fmt.Sprintf("check that iteration %d properly advanced the shard iterator", n), func() { - So(l.shardIterator.shardIteratorType, ShouldEqual, "AT_SEQUENCE_NUMBER") - So(l.shardIterator.sequenceNumber, ShouldEqual, *seq) + So(l.reader.(*KinesisReader).shardIterator.shardIteratorType, ShouldEqual, "AT_SEQUENCE_NUMBER") + So(l.reader.(*KinesisReader).shardIterator.sequenceNumber, ShouldEqual, *seq) }) } }) @@ -226,7 +179,7 @@ func TestListener(t *testing.T) { secs := []float64{} for i := 1; i <= 6; i++ { start := time.Now() - l.getRecords(1) + l.reader.GetRecords(1) secs = append(secs, time.Since(start).Seconds()) } elapsed := time.Since(start).Seconds() diff --git a/listener/sharditerator.go b/listener/sharditerator.go index df5cad8..3e017ee 100644 --- a/listener/sharditerator.go +++ b/listener/sharditerator.go @@ -2,6 +2,8 @@ package listener import ( "time" + + "github.com/aws/aws-sdk-go/aws" ) // ShardIterator represents the settings used to retrieve a shard iterator from diff --git a/listener/sharditerator_test.go b/listener/sharditerator_test.go new file mode 100644 index 0000000..55c3dda --- /dev/null +++ b/listener/sharditerator_test.go @@ -0,0 +1,57 @@ +package listener + +import ( + . "github.com/smartystreets/goconvey/convey" + "github.com/aws/aws-sdk-go/aws" + + "testing" + "time" +) + +func TestShardIterator(t *testing.T) { + Convey("given a new shard iterator", t, func() { + it := NewShardIterator() + + Convey("check that the default shard iterator type is TRIM_HORIZON", func() { + So(it.shardIteratorType, ShouldEqual, "TRIM_HORIZON") + So(it.getStartingSequenceNumber(), ShouldBeNil) + So(it.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can explicitly set it to TRIM_HORIZON", func() { + it = it.TrimHorizon() + So(it.shardIteratorType, ShouldEqual, "TRIM_HORIZON") + So(it.getStartingSequenceNumber(), ShouldBeNil) + So(it.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can explicitly set it to LATEST", func() { + it = it.Latest() + So(it.shardIteratorType, ShouldEqual, "LATEST") + So(it.getStartingSequenceNumber(), ShouldBeNil) + So(it.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can explicitly set it to AT_SEQEUENCE_NUMBER", func() { + it = it.AtSequenceNumber("some-sequence") + So(it.shardIteratorType, ShouldEqual, "AT_SEQUENCE_NUMBER") + So(aws.StringValue(it.getStartingSequenceNumber()), ShouldEqual, "some-sequence") + So(it.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can explicitly set it to AFTER_SEQEUENCE_NUMBER", func() { + it = it.AfterSequenceNumber("some-sequence") + So(it.shardIteratorType, ShouldEqual, "AFTER_SEQUENCE_NUMBER") + So(aws.StringValue(it.getStartingSequenceNumber()), ShouldEqual, "some-sequence") + So(it.getTimestamp(), ShouldBeNil) + }) + + Convey("check that we can explicitly set it to AT_TIMESTAMP", func() { + n := time.Now() + it = it.AtTimestamp(n) + So(it.shardIteratorType, ShouldEqual, "AT_TIMESTAMP") + So(aws.TimeValue(it.getTimestamp()).Equal(n), ShouldBeTrue) + So(it.getStartingSequenceNumber(), ShouldBeNil) + }) + }) +} diff --git a/producer/producer_test.go b/producer/producer_test.go index 2c127b5..645802f 100644 --- a/producer/producer_test.go +++ b/producer/producer_test.go @@ -45,7 +45,7 @@ func TestProducer(t *testing.T) { So(p, ShouldNotBeNil) So(err, ShouldBeNil) - l, err := listener.NewListener(stream, shards[0], func(c *listener.Config) { + l, err := listener.NewListener(func(c *listener.Config) { c.SetAwsConfig(k.Session.Config) c.SetConcurrency(10) }) From 1f865bf07f80c9042e1012b2d930a2345d95f1d7 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 11 May 2017 14:41:11 -0500 Subject: [PATCH 015/109] Fix errors and golint warnings --- errs/errors.go | 12 ++++++------ listener/kinesisreader.go | 16 ++++++++++------ listener/kinesisreaderconfig.go | 2 ++ listener/listener.go | 24 +++++++++++++++--------- listener/listener_test.go | 5 ++--- logging/logging.go | 8 +++++--- producer/producer.go | 6 +++--- 7 files changed, 43 insertions(+), 30 deletions(-) diff --git a/errs/errors.go b/errs/errors.go index 3651c0b..e17a8e1 100644 --- a/errs/errors.go +++ b/errs/errors.go @@ -19,13 +19,13 @@ var ( ) var ( - // ErrNilProducer is returned by a StreamWriter when it has not been - // correctly associated with a Producer. + // ErrNilListener is returned when the Listener is nil when it is not + // supposed to be ErrNilListener = errors.New("StreamReader not associated with a listener") - // ErrProducerAlreadyAssociated is returned by a StreamWriter attempting - // to associate it with a Producer when it already has an association - // with a producer. + // ErrListenerAlreadyAssociated is returned by a StreamReader attempting + // to associate it with a Listener when it already has an association + // with a listener ErrListenerAlreadyAssociated = errors.New("StreamReader already associated with a listener") // ErrAlreadyConsuming is returned when attempting to consume when the @@ -43,7 +43,7 @@ var ( // ErrNilGetShardIteratorResponse is returned when the GetShardIterator // call returns a nil response. - ErrNilGetShardIteratorResponse = errors.New("GetShardIteratore returned a nil response") + ErrNilGetShardIteratorResponse = errors.New("GetShardIterator returned a nil response") // ErrNilShardIterator is returned when the GetShardIterator call // returns a nil shard iterator. diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index a5b0b6a..17e64b2 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -25,6 +25,7 @@ type kinesisReaderOptions struct { shardIterator *ShardIterator } +// KinesisReader handles the API to read records from Kinesis. type KinesisReader struct { *kinesisReaderOptions @@ -36,6 +37,7 @@ type KinesisReader struct { clientMu sync.Mutex } +// NewKinesisReader creates a new stream reader to read records from Kinesis. func (r *KinesisReader) NewKinesisReader(stream, shard string, fn func(*KinesisReaderConfig)) (*KinesisReader, error) { config := NewKinesisReaderConfig(stream, shard) fn(config) @@ -162,7 +164,13 @@ func (r *KinesisReader) throttle(sem chan Empty) { }) } -// GetRecords calls GetRecords and delivers each record into the messages +// GetRecords calls GetNRecords and delivers each record into the messages +// channel. +func (r *KinesisReader) GetRecords() (int, error) { + return r.GetNRecords(r.batchSize) +} + +// GetNRecords calls GetRecords and delivers each record into the messages // channel. // FIXME: Need to investigate that the timeout implementation doesn't result in // an fd leak. Since we call Read on the HTTPResonse.Body in a select with a @@ -173,7 +181,7 @@ func (r *KinesisReader) throttle(sem chan Empty) { // down the TCP connection. Worst case scenario is that our client Timeout // eventually fires and closes the socket, but this can be susceptible to FD // exhaustion. -func (r *KinesisReader) GetRecords() (int, error) { +func (r *KinesisReader) GetNRecords(batchSize int) (int, error) { if err := r.ensureClient(); err != nil { return 0, err } @@ -363,7 +371,3 @@ func (r *KinesisReader) GetRecords() (int, error) { } return delivered, nil } - -func (r *KinesisReader) getBatchSize() int { - return r.batchSize -} diff --git a/listener/kinesisreaderconfig.go b/listener/kinesisreaderconfig.go index 0ed091a..e71e24f 100644 --- a/listener/kinesisreaderconfig.go +++ b/listener/kinesisreaderconfig.go @@ -1,9 +1,11 @@ package listener +// KinesisReaderConfig is used to configure a KinesisReader type KinesisReaderConfig struct { *kinesisReaderOptions } +// NewKinesisReaderConfig creates a new instance of KinesisReaderConfig func NewKinesisReaderConfig(stream, shard string) *KinesisReaderConfig { return &KinesisReaderConfig{ kinesisReaderOptions: &kinesisReaderOptions{ diff --git a/listener/listener.go b/listener/listener.go index 366f871..c608172 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -6,16 +6,19 @@ import ( "sync" "time" + "github.com/aws/aws-sdk-go/aws/session" + "github.com/rewardStyle/kinetic/errs" "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" ) +// StreamReader is an interface that abstracts out a stream reader type StreamReader interface { AssociateListener(listener *Listener) error - GetRecords(batchSize int) (int, error) + GetRecords() (int, error) + GetNRecords(batchSize int) (int, error) ensureClient() error - getBatchSize() int } // Empty is used a as a dummy type for counting semaphore channels. @@ -26,11 +29,11 @@ type Empty struct{} type MessageFn func([]byte, *sync.WaitGroup) type listenerOptions struct { + queueDepth int concurrency int getRecordsReadTimeout time.Duration reader StreamReader - - Stats StatsCollector + Stats StatsCollector } // Listener polls the StreamReader for messages. @@ -42,8 +45,10 @@ type Listener struct { concurrencySem chan Empty pipeOfDeath chan Empty - consuming bool - consumingMu sync.Mutex + consuming bool + consumingMu sync.Mutex + + Session session.Session } // NewListener creates a new listener for listening to message on a StreamReader. @@ -62,6 +67,7 @@ func NewListener(fn func(*Config)) (*Listener, error) { }, concurrencySem: make(chan Empty, config.concurrency), pipeOfDeath: make(chan Empty), + Session: session, } if err := l.reader.AssociateListener(l); err != nil { return nil, err @@ -76,7 +82,7 @@ func (l *Listener) startConsuming() bool { defer l.consumingMu.Unlock() if !l.consuming { l.consuming = true - l.messages = make(chan *message.Message, l.reader.getBatchSize()) + l.messages = make(chan *message.Message, l.queueDepth) return true } return false @@ -124,7 +130,7 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, e if !ok { return nil, err } - n, err := l.reader.GetRecords(1) + n, err := l.reader.GetNRecords(1) if err != nil { return nil, err } @@ -194,7 +200,7 @@ func (l *Listener) consume(ctx context.Context) { if !ok { break stop } - _, err := l.reader.GetRecords(l.reader.getBatchSize()) + _, err := l.reader.GetRecords() if err != nil { switch err := err.(type) { case net.Error: diff --git a/listener/listener_test.go b/listener/listener_test.go index 5a902c2..e4f0f3f 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -21,12 +21,11 @@ import ( ) func putRecord(l *Listener, b []byte) (*string, error) { - kReader := l.reader.(*KinesisReader) l.reader.ensureClient() - resp, err := kReader.client.PutRecord(&kinesis.PutRecordInput{ + resp, err := l.reader.(*KinesisReader).client.PutRecord(&kinesis.PutRecordInput{ Data: b, PartitionKey: aws.String("dummy"), - StreamName: aws.String(kReader.stream), + StreamName: aws.String(l.reader.(*KinesisReader).stream), }) if err != nil { return nil, err diff --git a/logging/logging.go b/logging/logging.go index 37b502c..3b66722 100644 --- a/logging/logging.go +++ b/logging/logging.go @@ -26,6 +26,7 @@ const ( LogPlaceholder aws.LogLevelType = LogDebug | (1 << (16 + iota)) ) +// ILogHelper is an interface for LogHelper type ILogHelper interface { Log(aws.LogLevelType, ...interface{}) LogBug(...interface{}) @@ -34,12 +35,13 @@ type ILogHelper interface { LogDebug(...interface{}) } +// LogHelper is used for defining log configuration type LogHelper struct { LogLevel aws.LogLevelType Logger aws.Logger } -// log handles levelled logging +// Log handles levelled logging func (l *LogHelper) Log(level aws.LogLevelType, args ...interface{}) { if l.LogLevel.Matches(level) { l.Logger.Log(args...) @@ -51,7 +53,7 @@ func (l *LogHelper) LogBug(args ...interface{}) { l.Log(LogBug, args...) } -// LogBug logs an error. +// LogError logs an error. func (l *LogHelper) LogError(args ...interface{}) { l.Log(LogError, args...) } @@ -61,7 +63,7 @@ func (l *LogHelper) LogInfo(args ...interface{}) { l.Log(LogInfo, args...) } -// Log a debug message using the AWS SDK logger. +// LogDebug a debug message using the AWS SDK logger. func (l *LogHelper) LogDebug(args ...interface{}) { l.Log(LogDebug, args...) } diff --git a/producer/producer.go b/producer/producer.go index eaf486a..2461571 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -48,8 +48,8 @@ type Producer struct { concurrencySem chan Empty pipeOfDeath chan Empty outstanding int - shutdownCond sync.Cond - producerWg sync.WaitGroup + shutdownCond *sync.Cond + producerWg *sync.WaitGroup producing bool producingMu sync.Mutex @@ -91,7 +91,7 @@ func (p *Producer) startProducing() bool { p.producing = true p.messages = make(chan *message.Message, p.queueDepth) p.retries = make(chan *message.Message) // TODO: should we use a buffered channel? - p.shutdownCond = sync.Cond{L: &sync.Mutex{}} + p.shutdownCond = &sync.Cond{L: &sync.Mutex{}} p.outstanding = 0 return true } From d9e62afdc1586885ba4dfb8cefbbe0bec1a130c0 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 11 May 2017 15:00:31 -0500 Subject: [PATCH 016/109] Fix more errors --- listener/kinesisreader.go | 2 +- listener/listener.go | 2 +- listener/listener_test.go | 2 +- producer/config_test.go | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index 17e64b2..415e0fd 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -78,7 +78,7 @@ func (r *KinesisReader) ensureClient() error { // are trying to ensure the client at the same time. // // As we don't expose any methods (or in fact, even the Listener object - // itself) to the client through the API, I don't forsee needing to add + // itself) to the client through the API, I don't foresee needing to add // this lock unless something dramatically changes about the design of // this library. r.clientMu.Lock() diff --git a/listener/listener.go b/listener/listener.go index c608172..948ef31 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -48,7 +48,7 @@ type Listener struct { consuming bool consumingMu sync.Mutex - Session session.Session + Session *session.Session } // NewListener creates a new listener for listening to message on a StreamReader. diff --git a/listener/listener_test.go b/listener/listener_test.go index e4f0f3f..2b08d57 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -178,7 +178,7 @@ func TestListener(t *testing.T) { secs := []float64{} for i := 1; i <= 6; i++ { start := time.Now() - l.reader.GetRecords(1) + l.reader.GetNRecords(1) secs = append(secs, time.Since(start).Seconds()) } elapsed := time.Since(start).Seconds() diff --git a/producer/config_test.go b/producer/config_test.go index 271476f..7878bc8 100644 --- a/producer/config_test.go +++ b/producer/config_test.go @@ -41,7 +41,7 @@ func (w *DebugStreamWriter) PutRecords(batch []*message.Message) ([]*message.Mes return nil, nil } -func (w *DebugStreamWriter) AssociateProducer(producer Producer) error { +func (w *DebugStreamWriter) AssociateProducer(producer *Producer) error { return nil } From 4b69801fa64af01d23d68a5d1521dc06511dc23f Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 12 May 2017 11:40:46 -0500 Subject: [PATCH 017/109] Fix a Nil reference exception --- circle.yml | 2 +- listener/config.go | 5 +++++ listener/kinesisreader.go | 26 ++++++++++++++------------ listener/listener.go | 5 ++--- listener/listener_test.go | 39 +++++++++++++++++++++++++++++++-------- producer/producer.go | 3 ++- producer/producer_test.go | 3 ++- 7 files changed, 57 insertions(+), 26 deletions(-) diff --git a/circle.yml b/circle.yml index 4ffb89a..a7ecae2 100644 --- a/circle.yml +++ b/circle.yml @@ -13,7 +13,7 @@ dependencies: - go get honnef.co/go/tools/cmd/staticcheck - go get honnef.co/go/tools/cmd/gosimple - go get -u github.com/golang/lint/golint - - docker run -d -t -p 4567:4567 dlsniper/kinesalite:1.8.0 + - docker run -d -t -p 4567:4567 instructure/kinesalite:latest override: - mkdir -p "$GOPATH/src/$IMPORT_PATH" - rsync -azvrC --delete $HOME/$CIRCLE_PROJECT_REPONAME "$GOPATH/src/github.com/$CIRCLE_PROJECT_USERNAME/" diff --git a/listener/config.go b/listener/config.go index 329fbcd..35b3417 100644 --- a/listener/config.go +++ b/listener/config.go @@ -47,6 +47,11 @@ func (c *Config) SetGetRecordsReadTimeout(timouet time.Duration) { c.getRecordsReadTimeout = timouet } +// SetKinesisStream sets the listener to read to the given Kinesis stream. +func (c *Config) SetKinesisStream(stream string, shard string, fn ...func(*KinesisReaderConfig)) { + c.reader = NewKinesisReader(stream, shard, fn...) +} + // SetLogLevel configures both the SDK and Kinetic log levels. func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { c.AwsOptions.SetLogLevel(logLevel) diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index 415e0fd..4b157af 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -38,13 +38,15 @@ type KinesisReader struct { } // NewKinesisReader creates a new stream reader to read records from Kinesis. -func (r *KinesisReader) NewKinesisReader(stream, shard string, fn func(*KinesisReaderConfig)) (*KinesisReader, error) { +func NewKinesisReader(stream, shard string, fn ...func(*KinesisReaderConfig)) *KinesisReader { config := NewKinesisReaderConfig(stream, shard) - fn(config) + for _, f := range fn { + f(config) + } return &KinesisReader{ kinesisReaderOptions: config.kinesisReaderOptions, throttleSem: make(chan Empty, 5), - }, nil + } } // AssociateListener associates the Kinesis stream writer to a producer. @@ -164,13 +166,7 @@ func (r *KinesisReader) throttle(sem chan Empty) { }) } -// GetRecords calls GetNRecords and delivers each record into the messages -// channel. -func (r *KinesisReader) GetRecords() (int, error) { - return r.GetNRecords(r.batchSize) -} - -// GetNRecords calls GetRecords and delivers each record into the messages +// GetRecords calls GetRecords and delivers each record into the messages // channel. // FIXME: Need to investigate that the timeout implementation doesn't result in // an fd leak. Since we call Read on the HTTPResonse.Body in a select with a @@ -181,7 +177,7 @@ func (r *KinesisReader) GetRecords() (int, error) { // down the TCP connection. Worst case scenario is that our client Timeout // eventually fires and closes the socket, but this can be susceptible to FD // exhaustion. -func (r *KinesisReader) GetNRecords(batchSize int) (int, error) { +func (r *KinesisReader) GetRecords(batchSize ...int) (int, error) { if err := r.ensureClient(); err != nil { return 0, err } @@ -197,8 +193,14 @@ func (r *KinesisReader) GetNRecords(batchSize int) (int, error) { var startUnmarshalTime time.Time start := time.Now() + getRecordsBatchSize := r.batchSize + for _, size := range batchSize { + getRecordsBatchSize = size + break + } + req, resp := r.client.GetRecordsRequest(&kinesis.GetRecordsInput{ - Limit: aws.Int64(int64(r.batchSize)), + Limit: aws.Int64(int64(getRecordsBatchSize)), ShardIterator: aws.String(r.nextShardIterator), }) diff --git a/listener/listener.go b/listener/listener.go index 948ef31..c1dd201 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -16,8 +16,7 @@ import ( // StreamReader is an interface that abstracts out a stream reader type StreamReader interface { AssociateListener(listener *Listener) error - GetRecords() (int, error) - GetNRecords(batchSize int) (int, error) + GetRecords(batchSize ...int) (int, error) ensureClient() error } @@ -130,7 +129,7 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, e if !ok { return nil, err } - n, err := l.reader.GetNRecords(1) + n, err := l.reader.GetRecords(1) if err != nil { return nil, err } diff --git a/listener/listener_test.go b/listener/listener_test.go index 2b08d57..e97eb14 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -18,6 +18,7 @@ import ( "github.com/rewardStyle/kinetic" "github.com/rewardStyle/kinetic/errs" + "github.com/rewardStyle/kinetic/producer" ) func putRecord(l *Listener, b []byte) (*string, error) { @@ -53,20 +54,42 @@ func TestListener(t *testing.T) { So(err, ShouldBeNil) So(len(shards), ShouldEqual, 1) + p, err := producer.NewProducer(func(c *producer.Config) { + c.SetAwsConfig(k.Session.Config) + c.SetKinesisStream(stream) + c.SetBatchSize(5) + c.SetBatchTimeout(1 * time.Second) + }) + So(p, ShouldNotBeNil) + So(err, ShouldBeNil) + l, err := NewListener(func(c *Config) { c.SetAwsConfig(k.Session.Config) + c.SetKinesisStream(stream, shards[0]) c.SetConcurrency(10) }) So(l, ShouldNotBeNil) So(err, ShouldBeNil) - Convey("check that calling ensureClient twice doesn't overwrite existing client", func() { - So(l.reader.(*KinesisReader).client, ShouldBeNil) - l.reader.ensureClient() - So(l.reader.(*KinesisReader).client, ShouldNotBeNil) - client := l.reader.(*KinesisReader).client - l.reader.ensureClient() - So(l.reader.(*KinesisReader).client, ShouldEqual, client) + Convey("given a kinesis reader", func() { + r := l.reader.(*KinesisReader) + + Convey("check that the reader was initialized with the correct stream name", func() { + So(r.stream, ShouldEqual, stream) + }) + + Convey("check that the reader has a valid reference to the listener", func() { + So(r.listener, ShouldEqual, l) + }) + + Convey("check that calling ensureClient twice doesn't overwrite existing client", func() { + So(r.client, ShouldBeNil) + r.ensureClient() + So(r.client, ShouldNotBeNil) + client := r.client + r.ensureClient() + So(r.client, ShouldEqual, client) + }) }) Convey("check that setting an empty shard iterator returns an error", func() { @@ -178,7 +201,7 @@ func TestListener(t *testing.T) { secs := []float64{} for i := 1; i <= 6; i++ { start := time.Now() - l.reader.GetNRecords(1) + l.reader.GetRecords(1) secs = append(secs, time.Since(start).Seconds()) } elapsed := time.Since(start).Seconds() diff --git a/producer/producer.go b/producer/producer.go index 2461571..9566202 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -92,6 +92,7 @@ func (p *Producer) startProducing() bool { p.messages = make(chan *message.Message, p.queueDepth) p.retries = make(chan *message.Message) // TODO: should we use a buffered channel? p.shutdownCond = &sync.Cond{L: &sync.Mutex{}} + p.producerWg = new(sync.WaitGroup) p.outstanding = 0 return true } @@ -123,7 +124,7 @@ stop: failed := len(retries) p.Stats.AddSent(len(batch) - failed) p.Stats.AddFailed(failed) - if err == nil { + if err == nil && failed == 0 { break stop } switch err := err.(type) { diff --git a/producer/producer_test.go b/producer/producer_test.go index 645802f..554a7e4 100644 --- a/producer/producer_test.go +++ b/producer/producer_test.go @@ -40,14 +40,15 @@ func TestProducer(t *testing.T) { c.SetKinesisStream(stream) c.SetBatchSize(5) c.SetBatchTimeout(1 * time.Second) - So(err, ShouldBeNil) }) So(p, ShouldNotBeNil) So(err, ShouldBeNil) l, err := listener.NewListener(func(c *listener.Config) { c.SetAwsConfig(k.Session.Config) + c.SetKinesisStream(stream, shards[0]) c.SetConcurrency(10) + c.SetGetRecordsReadTimeout(1 * time.Second) }) So(l, ShouldNotBeNil) So(err, ShouldBeNil) From 1acb629ca5d8c1755ae3f5b2377ecd8325fd490a Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Sat, 13 May 2017 14:46:42 -0500 Subject: [PATCH 018/109] Added a default queueDepth which was causing the listener to hang --- listener/config.go | 6 ++++++ producer/producer_test.go | 16 ++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/listener/config.go b/listener/config.go index 35b3417..f85886c 100644 --- a/listener/config.go +++ b/listener/config.go @@ -21,6 +21,7 @@ func NewConfig() *Config { return &Config{ AwsOptions: config.DefaultAwsOptions(), listenerOptions: &listenerOptions{ + queueDepth: 10000, concurrency: 10000, getRecordsReadTimeout: 1 * time.Second, Stats: &NilStatsCollector{}, @@ -35,6 +36,11 @@ func (c *Config) SetAwsConfig(config *aws.Config) { c.AwsConfig = config } +// SetQueueDepth controls the depth of the listener queue +func (c *Config) SetQueueDepth(queueDepth int) { + c.queueDepth = queueDepth +} + // SetConcurrency controls the number of goroutines the Listener will spawn to // process messages. func (c *Config) SetConcurrency(concurrency int) { diff --git a/producer/producer_test.go b/producer/producer_test.go index 554a7e4..fbca2f5 100644 --- a/producer/producer_test.go +++ b/producer/producer_test.go @@ -47,6 +47,7 @@ func TestProducer(t *testing.T) { l, err := listener.NewListener(func(c *listener.Config) { c.SetAwsConfig(k.Session.Config) c.SetKinesisStream(stream, shards[0]) + c.SetQueueDepth(10) c.SetConcurrency(10) c.SetGetRecordsReadTimeout(1 * time.Second) }) @@ -74,6 +75,21 @@ func TestProducer(t *testing.T) { }) }) + Convey("check that we can send and receive a single message", func(){ + start := time.Now() + data := "hello" + p.Send(&message.Message{ + PartitionKey: aws.String("key"), + Data: []byte(data), + }) + msg, err := l.RetrieveWithContext(context.TODO()) + elapsed := time.Since(start) + So(err, ShouldBeNil) + So(string(msg.Data), ShouldEqual, data) + So(elapsed.Seconds(), ShouldBeGreaterThan, 1) + + }) + Convey("check that we can send a single message after batch timeout elapses", func() { start := time.Now() data := "hello" From 7f564b25620cd680ac14e07abb969a7a385ad6ba Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 16 May 2017 10:25:42 -0500 Subject: [PATCH 019/109] Refactored StreamReader interface based on new design considerations. Introduced first draft of KclReader. Code cleanup based on code review. --- circle.yml | 2 +- listener/kclreader.go | 183 ++++++++++++++++++++++++++++++++++++ listener/kclreaderconfig.go | 23 +++++ listener/kinesisreader.go | 22 +++-- listener/listener.go | 7 +- listener/listener_test.go | 12 +-- multilang/action.go | 37 ++++++++ 7 files changed, 262 insertions(+), 24 deletions(-) create mode 100644 listener/kclreader.go create mode 100644 listener/kclreaderconfig.go create mode 100644 multilang/action.go diff --git a/circle.yml b/circle.yml index a7ecae2..4ffb89a 100644 --- a/circle.yml +++ b/circle.yml @@ -13,7 +13,7 @@ dependencies: - go get honnef.co/go/tools/cmd/staticcheck - go get honnef.co/go/tools/cmd/gosimple - go get -u github.com/golang/lint/golint - - docker run -d -t -p 4567:4567 instructure/kinesalite:latest + - docker run -d -t -p 4567:4567 dlsniper/kinesalite:1.8.0 override: - mkdir -p "$GOPATH/src/$IMPORT_PATH" - rsync -azvrC --delete $HOME/$CIRCLE_PROJECT_REPONAME "$GOPATH/src/github.com/$CIRCLE_PROJECT_USERNAME/" diff --git a/listener/kclreader.go b/listener/kclreader.go new file mode 100644 index 0000000..977c7b5 --- /dev/null +++ b/listener/kclreader.go @@ -0,0 +1,183 @@ +package listener + +import ( + "bufio" + "encoding/json" + "math" + "os" + "sync" + + "github.com/rewardStyle/kinetic/errs" + "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic/multilang" +) + +type kclReaderOptions struct { + onInitCallbackFn func() error + onCheckpointCallbackFn func() error + onShutdownCallbackFn func() error +} + +// KclReader +type KclReader struct { + *kclReaderOptions + throttleSem chan Empty + listener *Listener + scanner *bufio.Scanner + mutex *sync.Mutex + msgBuffer []message.Message + ackPending bool +} + +func NewKclReader(fn ...func(*KclReaderConfig)) *KclReader { + config := NewKclReaderConfig() + for _, f := range fn { + f(config) + } + return &KclReader{ + kclReaderOptions: config.kclReaderOptions, + throttleSem: make(chan Empty, 5), + msgBuffer: []message.Message{}, + mutex: &sync.Mutex{}, + } +} + +// AssociateListener +func (r *KclReader) AssociateListener(l *Listener) error { + r.mutex.Lock() + defer r.mutex.Unlock() + if r.listener != nil { + return errs.ErrListenerAlreadyAssociated + } + r.listener = l + return nil +} + +func (r *KclReader) ensureClient() error { + r.mutex.Lock() + defer r.mutex.Unlock() + if r.scanner == nil { + if r.listener == nil { + return errs.ErrNilListener + } + r.scanner = bufio.NewScanner(os.Stdin) + bufio.NewReader(os.Stdin) + } + return nil +} + +// GetRecord +func (r *KclReader) GetRecord() (int, error) { + return r.processRecords(1) +} + +// GetRecords +func (r *KclReader) GetRecords() (int, error) { + return r.processRecords(-1) +} + +func (r *KclReader) processRecords(numRecords int) (int, error) { + r.mutex.Lock() + defer r.mutex.Unlock() + + // Define the batchSize + batchSize := 0; + if len(r.msgBuffer) > 0 { + if numRecords < 0 { + batchSize = len(r.msgBuffer) + } else { + batchSize = int(math.Min(float64(len(r.msgBuffer)), float64(numRecords))) + } + } + + // Loop through the message buffer and put the correct number of messages on the listener's message channel + for i := 0; i < batchSize; i++ { + r.listener.messages <- &r.msgBuffer[0] + r.msgBuffer = r.msgBuffer[1:] + } + + // Send an acknowledgement that the 'ProcessRecords' message was received/processed + if len(r.msgBuffer) == 0 && r.ackPending { + r.sendMessage(multilang.NewStatusMessage(multilang.ProcessRecords)) + } + + return batchSize, nil +} + +func (r *KclReader) processAction() error { + if err := r.ensureClient(); err != nil { + return err + } + + actionMessage := &multilang.ActionMessage{} + for r.scanner.Scan() { + err := json.Unmarshal(r.scanner.Bytes(), actionMessage) + if err != nil { + return err + } + + switch actionMessage.Action { + case multilang.Initialize: + r.onInit() + r.sendMessage(multilang.NewStatusMessage(multilang.Initialize)) + case multilang.Checkpoint: + r.onCheckpoint() + r.sendMessage(multilang.NewStatusMessage(multilang.Checkpoint)) + case multilang.Shutdown: + r.onShutdown() + r.sendMessage(multilang.NewStatusMessage(multilang.Shutdown)) + case multilang.ProcessRecords: + go func(){ + r.mutex.Lock() + defer r.mutex.Unlock() + + if r.ackPending { + // TODO: error out + // This is an error according to the Multilang protocol + } + + r.msgBuffer = append(r.msgBuffer, actionMessage.Records...) + r.ackPending = true; + }() + default: + } + } + + return nil +} + +func (r *KclReader) sendMessage(msg *multilang.ActionMessage) error { + json.Marshal(msg) + + return nil +} + +func (r *KclReader) onInit() error { + if r.onInitCallbackFn != nil { + err := r.onInitCallbackFn() + if err != nil { + // TODO: + } + } + return nil +} + +func (r *KclReader) onCheckpoint() error { + if r.onCheckpointCallbackFn != nil { + err := r.onCheckpointCallbackFn() + if err != nil { + // TODO: + } + } + return nil +} + +func (r *KclReader) onShutdown() error { + if r.onShutdownCallbackFn != nil { + err := r.onShutdownCallbackFn() + if err != nil { + // TODO: + } + } + return nil +} diff --git a/listener/kclreaderconfig.go b/listener/kclreaderconfig.go new file mode 100644 index 0000000..89147f5 --- /dev/null +++ b/listener/kclreaderconfig.go @@ -0,0 +1,23 @@ +package listener + +type KclReaderConfig struct { + *kclReaderOptions +} + +func NewKclReaderConfig() *KclReaderConfig { + return &KclReaderConfig{ + kclReaderOptions: &kclReaderOptions{}, + } +} + +func (c *KclReaderConfig) SetOnInitCallbackFn(fn func() error) { + c.onInitCallbackFn = fn +} + +func (c *KclReaderConfig) SetOnCheckpointCallbackFn(fn func() error) { + c.onCheckpointCallbackFn = fn +} + +func (c *KclReaderConfig) SetOnShutdownCallbackFn(fn func() error) { + c.onShutdownCallbackFn = fn +} \ No newline at end of file diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index 4b157af..129c87f 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -166,8 +166,18 @@ func (r *KinesisReader) throttle(sem chan Empty) { }) } -// GetRecords calls GetRecords and delivers each record into the messages +// GetRecord calls getRecords and delivers one record into the messages // channel. +func (r *KinesisReader) GetRecord() (int, error) { + return r.getRecords(1) +} + +// GetRecords calls getRecords and delivers each record into the messages +// channel. +func (r *KinesisReader) GetRecords() (int, error) { + return r.getRecords(r.batchSize) +} + // FIXME: Need to investigate that the timeout implementation doesn't result in // an fd leak. Since we call Read on the HTTPResonse.Body in a select with a // timeout channel, we do prevent ourself from blocking. Once we timeout, we @@ -177,7 +187,7 @@ func (r *KinesisReader) throttle(sem chan Empty) { // down the TCP connection. Worst case scenario is that our client Timeout // eventually fires and closes the socket, but this can be susceptible to FD // exhaustion. -func (r *KinesisReader) GetRecords(batchSize ...int) (int, error) { +func (r *KinesisReader) getRecords(batchSize int) (int, error) { if err := r.ensureClient(); err != nil { return 0, err } @@ -193,14 +203,8 @@ func (r *KinesisReader) GetRecords(batchSize ...int) (int, error) { var startUnmarshalTime time.Time start := time.Now() - getRecordsBatchSize := r.batchSize - for _, size := range batchSize { - getRecordsBatchSize = size - break - } - req, resp := r.client.GetRecordsRequest(&kinesis.GetRecordsInput{ - Limit: aws.Int64(int64(getRecordsBatchSize)), + Limit: aws.Int64(int64(batchSize)), ShardIterator: aws.String(r.nextShardIterator), }) diff --git a/listener/listener.go b/listener/listener.go index c1dd201..925a267 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -16,7 +16,8 @@ import ( // StreamReader is an interface that abstracts out a stream reader type StreamReader interface { AssociateListener(listener *Listener) error - GetRecords(batchSize ...int) (int, error) + GetRecord() (int, error) + GetRecords() (int, error) ensureClient() error } @@ -129,12 +130,12 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, e if !ok { return nil, err } - n, err := l.reader.GetRecords(1) + n, err := l.reader.GetRecord() if err != nil { return nil, err } if n > 0 { - l.Stats.AddDelivered(1) + l.Stats.AddDelivered(n) return <-l.messages, nil } } diff --git a/listener/listener_test.go b/listener/listener_test.go index e97eb14..1343680 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -18,7 +18,6 @@ import ( "github.com/rewardStyle/kinetic" "github.com/rewardStyle/kinetic/errs" - "github.com/rewardStyle/kinetic/producer" ) func putRecord(l *Listener, b []byte) (*string, error) { @@ -54,15 +53,6 @@ func TestListener(t *testing.T) { So(err, ShouldBeNil) So(len(shards), ShouldEqual, 1) - p, err := producer.NewProducer(func(c *producer.Config) { - c.SetAwsConfig(k.Session.Config) - c.SetKinesisStream(stream) - c.SetBatchSize(5) - c.SetBatchTimeout(1 * time.Second) - }) - So(p, ShouldNotBeNil) - So(err, ShouldBeNil) - l, err := NewListener(func(c *Config) { c.SetAwsConfig(k.Session.Config) c.SetKinesisStream(stream, shards[0]) @@ -201,7 +191,7 @@ func TestListener(t *testing.T) { secs := []float64{} for i := 1; i <= 6; i++ { start := time.Now() - l.reader.GetRecords(1) + l.reader.GetRecord() secs = append(secs, time.Since(start).Seconds()) } elapsed := time.Since(start).Seconds() diff --git a/multilang/action.go b/multilang/action.go new file mode 100644 index 0000000..8c43f78 --- /dev/null +++ b/multilang/action.go @@ -0,0 +1,37 @@ +package multilang + +import "github.com/rewardStyle/kinetic/message" + +type ActionType string + +const ( + Initialize ActionType = "initialize" + ProcessRecords ActionType = "processRecords" + Checkpoint ActionType = "checkpoint" + Shutdown ActionType = "shutdown" + Status ActionType = "status" +) + +type ActionMessage struct { + Action ActionType `json:"action"` + ShardId string `json:"shardId,omitempty"` + Records []message.Message `json:"records,omitempty"` + Checkpoint int `json:"checkpoint,omitempty"` + Error string `json:"error,omitempty"` + Reason string `json:"reason,omitempty"` + ResponseFor ActionType `json:"responseFor"` +} + +func NewCheckpointMessage(seqNum int) *ActionMessage { + return &ActionMessage{ + Action: Checkpoint, + Checkpoint: seqNum, + } +} + +func NewStatusMessage(actionType ActionType) *ActionMessage { + return &ActionMessage{ + Action: Status, + ResponseFor: actionType, + } +} \ No newline at end of file From 6edab8483684bb4eb11d2c4c64ce82f3ebca5de4 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 16 May 2017 11:00:26 -0500 Subject: [PATCH 020/109] Fixed golint errors. Implemented sendMessage function to write to stdout. --- listener/kclreader.go | 29 +++++++++++++++++++++++------ listener/kclreaderconfig.go | 8 ++++++++ listener/kinesisreader.go | 2 +- multilang/action.go | 7 ++++++- 4 files changed, 38 insertions(+), 8 deletions(-) diff --git a/listener/kclreader.go b/listener/kclreader.go index 977c7b5..43dc4b1 100644 --- a/listener/kclreader.go +++ b/listener/kclreader.go @@ -18,7 +18,7 @@ type kclReaderOptions struct { onShutdownCallbackFn func() error } -// KclReader +// KclReader handles the KCL Multilang Protocol to read records from KCL type KclReader struct { *kclReaderOptions throttleSem chan Empty @@ -29,6 +29,7 @@ type KclReader struct { ackPending bool } +// NewKclReader creates a new stream reader to read records from KCL func NewKclReader(fn ...func(*KclReaderConfig)) *KclReader { config := NewKclReaderConfig() for _, f := range fn { @@ -42,7 +43,7 @@ func NewKclReader(fn ...func(*KclReaderConfig)) *KclReader { } } -// AssociateListener +// AssociateListener associates the KCL stream reader to a listener func (r *KclReader) AssociateListener(l *Listener) error { r.mutex.Lock() defer r.mutex.Unlock() @@ -66,16 +67,22 @@ func (r *KclReader) ensureClient() error { return nil } -// GetRecord +// GetRecord calls processRecords to attempt to put one message from message buffer to the listener's message +// channel func (r *KclReader) GetRecord() (int, error) { return r.processRecords(1) } -// GetRecords +// GetRecords calls processRecords to attempt to put all messages on the message buffer on the listener's +// message channel func (r *KclReader) GetRecords() (int, error) { return r.processRecords(-1) } +// processRecords is a helper method which loops through the message buffer and puts messages on the listener's +// message channel. After all the messages on the message buffer have been moved to the listener's message +// channel, a message is sent (following the Multilang protocol) to acknowledge that the processRecords message +// has been received / processed func (r *KclReader) processRecords(numRecords int) (int, error) { r.mutex.Lock() defer r.mutex.Unlock() @@ -98,7 +105,10 @@ func (r *KclReader) processRecords(numRecords int) (int, error) { // Send an acknowledgement that the 'ProcessRecords' message was received/processed if len(r.msgBuffer) == 0 && r.ackPending { - r.sendMessage(multilang.NewStatusMessage(multilang.ProcessRecords)) + err := r.sendMessage(multilang.NewStatusMessage(multilang.ProcessRecords)) + if err != nil { + // TODO: What to do if the ack status message fails? + } } return batchSize, nil @@ -147,7 +157,14 @@ func (r *KclReader) processAction() error { } func (r *KclReader) sendMessage(msg *multilang.ActionMessage) error { - json.Marshal(msg) + b, err := json.Marshal(msg) + if err != nil { + // TODO: + } + _, err = os.Stdout.Write(b) + if err != nil { + // TODO: + } return nil } diff --git a/listener/kclreaderconfig.go b/listener/kclreaderconfig.go index 89147f5..13f32ca 100644 --- a/listener/kclreaderconfig.go +++ b/listener/kclreaderconfig.go @@ -1,23 +1,31 @@ package listener +// KclReaderConfig is used to configure KclReader type KclReaderConfig struct { *kclReaderOptions } +// NewKclReaderConfig creates a new instance of KclReaderConfig func NewKclReaderConfig() *KclReaderConfig { return &KclReaderConfig{ kclReaderOptions: &kclReaderOptions{}, } } +// SetOnInitCallbackFn configures a callback function which is run prior to sending a status message +// acknowledging an 'initialize' message was received / processed func (c *KclReaderConfig) SetOnInitCallbackFn(fn func() error) { c.onInitCallbackFn = fn } +// SetOnCheckpointCallbackFn configures a callback function which is run prior to sending a status message +// acknowledging an 'checkpoint' message was received / processed func (c *KclReaderConfig) SetOnCheckpointCallbackFn(fn func() error) { c.onCheckpointCallbackFn = fn } +// SetOnShutdownCallbackFn configures a callback function which is run prior to sending a status message +// acknowledging a 'shutdown' message was received / processed func (c *KclReaderConfig) SetOnShutdownCallbackFn(fn func() error) { c.onShutdownCallbackFn = fn } \ No newline at end of file diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index 129c87f..4971d8e 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -49,7 +49,7 @@ func NewKinesisReader(stream, shard string, fn ...func(*KinesisReaderConfig)) *K } } -// AssociateListener associates the Kinesis stream writer to a producer. +// AssociateListener associates the Kinesis stream reader to a listener. func (r *KinesisReader) AssociateListener(l *Listener) error { r.clientMu.Lock() defer r.clientMu.Unlock() diff --git a/multilang/action.go b/multilang/action.go index 8c43f78..345b389 100644 --- a/multilang/action.go +++ b/multilang/action.go @@ -2,8 +2,10 @@ package multilang import "github.com/rewardStyle/kinetic/message" +// ActionType is used as an enum for KCL Multilang protocol action message types type ActionType string +// These are the enumerated KCL Multilang protocol action message types const ( Initialize ActionType = "initialize" ProcessRecords ActionType = "processRecords" @@ -12,9 +14,10 @@ const ( Status ActionType = "status" ) +// ActionMessage is a struct used to marshal / unmarshall KCL Multilang protocol action messages type ActionMessage struct { Action ActionType `json:"action"` - ShardId string `json:"shardId,omitempty"` + ShardID string `json:"shardId,omitempty"` Records []message.Message `json:"records,omitempty"` Checkpoint int `json:"checkpoint,omitempty"` Error string `json:"error,omitempty"` @@ -22,6 +25,7 @@ type ActionMessage struct { ResponseFor ActionType `json:"responseFor"` } +// NewCheckpointMessage is used to create a new checkpoint message func NewCheckpointMessage(seqNum int) *ActionMessage { return &ActionMessage{ Action: Checkpoint, @@ -29,6 +33,7 @@ func NewCheckpointMessage(seqNum int) *ActionMessage { } } +// NewStatusMessage is used to create a new status message func NewStatusMessage(actionType ActionType) *ActionMessage { return &ActionMessage{ Action: Status, From 98dfd11cd58063d61b1882631e4a9b32d0abae26 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 16 May 2017 13:43:15 -0500 Subject: [PATCH 021/109] Call processAction from ensureClient. Handle decoding and unmarshalling correctly. --- listener/kclreader.go | 65 ++++++++++++++++++++++++++++++++----------- multilang/action.go | 4 +-- 2 files changed, 50 insertions(+), 19 deletions(-) diff --git a/listener/kclreader.go b/listener/kclreader.go index 43dc4b1..f10ca48 100644 --- a/listener/kclreader.go +++ b/listener/kclreader.go @@ -2,7 +2,9 @@ package listener import ( "bufio" + "encoding/base64" "encoding/json" + "errors" "math" "os" "sync" @@ -54,6 +56,7 @@ func (r *KclReader) AssociateListener(l *Listener) error { return nil } +// ensureClient will lazily ensure that we are reading from STDIN. func (r *KclReader) ensureClient() error { r.mutex.Lock() defer r.mutex.Unlock() @@ -62,7 +65,13 @@ func (r *KclReader) ensureClient() error { return errs.ErrNilListener } r.scanner = bufio.NewScanner(os.Stdin) - bufio.NewReader(os.Stdin) + go func() error { + err := r.processAction() + if err != nil { + return err + } + return nil + }() } return nil } @@ -84,6 +93,10 @@ func (r *KclReader) GetRecords() (int, error) { // channel, a message is sent (following the Multilang protocol) to acknowledge that the processRecords message // has been received / processed func (r *KclReader) processRecords(numRecords int) (int, error) { + if err := r.ensureClient(); err != nil { + return 0, err + } + r.mutex.Lock() defer r.mutex.Unlock() @@ -107,20 +120,18 @@ func (r *KclReader) processRecords(numRecords int) (int, error) { if len(r.msgBuffer) == 0 && r.ackPending { err := r.sendMessage(multilang.NewStatusMessage(multilang.ProcessRecords)) if err != nil { - // TODO: What to do if the ack status message fails? + r.listener.LogError(err) + return batchSize, err } } return batchSize, nil } +// processAction listens to STDIN and processes action messages based on the Multilanguage protocol from KCL func (r *KclReader) processAction() error { - if err := r.ensureClient(); err != nil { - return err - } - - actionMessage := &multilang.ActionMessage{} for r.scanner.Scan() { + actionMessage := &multilang.ActionMessage{} err := json.Unmarshal(r.scanner.Bytes(), actionMessage) if err != nil { return err @@ -137,17 +148,34 @@ func (r *KclReader) processAction() error { r.onShutdown() r.sendMessage(multilang.NewStatusMessage(multilang.Shutdown)) case multilang.ProcessRecords: - go func(){ + go func() error { r.mutex.Lock() defer r.mutex.Unlock() if r.ackPending { - // TODO: error out - // This is an error according to the Multilang protocol + return errors.New("Receieved a processRecords action message from KCL " + + "unexpectedly") + } + + // Decode the base64 encoded message + decodedMsg, err := base64.StdEncoding.DecodeString(actionMessage.Records) + if err != nil { + r.listener.LogError(err) + return err } - r.msgBuffer = append(r.msgBuffer, actionMessage.Records...) + // Unmarshall the decoded message + msgs := []message.Message{} + err = json.Unmarshal(decodedMsg, msgs) + if err != nil { + r.listener.LogError(err) + return err + } + + r.msgBuffer = append(r.msgBuffer, msgs...) r.ackPending = true; + + return nil }() default: } @@ -159,11 +187,13 @@ func (r *KclReader) processAction() error { func (r *KclReader) sendMessage(msg *multilang.ActionMessage) error { b, err := json.Marshal(msg) if err != nil { - // TODO: + r.listener.LogError(err) + return err } _, err = os.Stdout.Write(b) if err != nil { - // TODO: + r.listener.LogError(err) + return err } return nil @@ -173,7 +203,8 @@ func (r *KclReader) onInit() error { if r.onInitCallbackFn != nil { err := r.onInitCallbackFn() if err != nil { - // TODO: + r.listener.LogError(err) + return err } } return nil @@ -183,7 +214,8 @@ func (r *KclReader) onCheckpoint() error { if r.onCheckpointCallbackFn != nil { err := r.onCheckpointCallbackFn() if err != nil { - // TODO: + r.listener.LogError(err) + return err } } return nil @@ -193,7 +225,8 @@ func (r *KclReader) onShutdown() error { if r.onShutdownCallbackFn != nil { err := r.onShutdownCallbackFn() if err != nil { - // TODO: + r.listener.LogError(err) + return err } } return nil diff --git a/multilang/action.go b/multilang/action.go index 345b389..752b9b8 100644 --- a/multilang/action.go +++ b/multilang/action.go @@ -1,7 +1,5 @@ package multilang -import "github.com/rewardStyle/kinetic/message" - // ActionType is used as an enum for KCL Multilang protocol action message types type ActionType string @@ -18,7 +16,7 @@ const ( type ActionMessage struct { Action ActionType `json:"action"` ShardID string `json:"shardId,omitempty"` - Records []message.Message `json:"records,omitempty"` + Records string `json:"records,omitempty"` Checkpoint int `json:"checkpoint,omitempty"` Error string `json:"error,omitempty"` Reason string `json:"reason,omitempty"` From d41dd3623366cb51c83083e23c1440c710bee61f Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 16 May 2017 13:48:48 -0500 Subject: [PATCH 022/109] Fix gosimple error by simplifying error handling logic --- listener/kclreader.go | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/listener/kclreader.go b/listener/kclreader.go index f10ca48..4bcbc90 100644 --- a/listener/kclreader.go +++ b/listener/kclreader.go @@ -66,11 +66,7 @@ func (r *KclReader) ensureClient() error { } r.scanner = bufio.NewScanner(os.Stdin) go func() error { - err := r.processAction() - if err != nil { - return err - } - return nil + return r.processAction() }() } return nil From af64e1aaaa0fe497465f5e459d26f82c18affbef Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 18 May 2017 14:01:14 -0500 Subject: [PATCH 023/109] Fix unmarshalling / base64 decoding step of the KCL Multilang Daemon's processRecord message --- .gitignore | 2 ++ listener/kclreader.go | 20 +++--------------- multilang/action.go | 49 ++++++++++++++++++++++++++++++++++++++++--- 3 files changed, 51 insertions(+), 20 deletions(-) diff --git a/.gitignore b/.gitignore index 838c953..4b7ea69 100644 --- a/.gitignore +++ b/.gitignore @@ -35,3 +35,5 @@ tags # custom # external packages folder vendor/ +.idea/* +kinetic.iml diff --git a/listener/kclreader.go b/listener/kclreader.go index 4bcbc90..2737430 100644 --- a/listener/kclreader.go +++ b/listener/kclreader.go @@ -2,7 +2,6 @@ package listener import ( "bufio" - "encoding/base64" "encoding/json" "errors" "math" @@ -149,26 +148,13 @@ func (r *KclReader) processAction() error { defer r.mutex.Unlock() if r.ackPending { - return errors.New("Receieved a processRecords action message from KCL " + + return errors.New("Received a processRecords action message from KCL " + "unexpectedly") } - // Decode the base64 encoded message - decodedMsg, err := base64.StdEncoding.DecodeString(actionMessage.Records) - if err != nil { - r.listener.LogError(err) - return err + for _, msg := range actionMessage.Records { + r.msgBuffer = append(r.msgBuffer, *msg.ToMessage()) } - - // Unmarshall the decoded message - msgs := []message.Message{} - err = json.Unmarshal(decodedMsg, msgs) - if err != nil { - r.listener.LogError(err) - return err - } - - r.msgBuffer = append(r.msgBuffer, msgs...) r.ackPending = true; return nil diff --git a/multilang/action.go b/multilang/action.go index 752b9b8..b4aacb1 100644 --- a/multilang/action.go +++ b/multilang/action.go @@ -1,5 +1,13 @@ package multilang +import ( + "encoding/base64" + "encoding/json" + "time" + + "github.com/rewardStyle/kinetic/message" +) + // ActionType is used as an enum for KCL Multilang protocol action message types type ActionType string @@ -12,15 +20,50 @@ const ( Status ActionType = "status" ) -// ActionMessage is a struct used to marshal / unmarshall KCL Multilang protocol action messages +// ActionMessage is a struct used to marshal / unmarshal KCL Multilang protocol action messages type ActionMessage struct { Action ActionType `json:"action"` ShardID string `json:"shardId,omitempty"` - Records string `json:"records,omitempty"` + Records []Record `json:"records,omitempty"` Checkpoint int `json:"checkpoint,omitempty"` Error string `json:"error,omitempty"` Reason string `json:"reason,omitempty"` - ResponseFor ActionType `json:"responseFor"` + ResponseFor ActionType `json:"responseFor,omitempty"` +} + +// Record is a struct used to marshal / unmarshal kinesis records from KCL Multilang protocol +type Record struct { + ApproximateArrivalTimestamp time.Time `json:"approximateArrivalTimestamp,omitempty"` + Data []byte `json:"data,omitempty"` + PartitionKey string `json:"partitionKey,omitempty"` + SequenceNumber string `json:"sequenceNumber,omitempty"` +} + +// UnmarshalJSON is used as a custom unmarshaller to base64 decode the data field of the KCL Multilang +// processRecord message +func (r *Record) UnmarshalJSON(data []byte) error { + record := &Record{} + if err := json.Unmarshal(data, record); err != nil { + return err + } + + decodedMsg, err := base64.StdEncoding.DecodeString(record.Data) + if err != nil { + return err + } + r.Data = decodedMsg + + return nil +} + +// ToMessage is used to transform a multilang.Record struct into a message.Message struct +func (r *Record) ToMessage() *message.Message { + return &message.Message{ + ApproximateArrivalTimestamp: &r.ApproximateArrivalTimestamp, + Data: r.Data, + PartitionKey: &r.PartitionKey, + SequenceNumber: &r.SequenceNumber, + } } // NewCheckpointMessage is used to create a new checkpoint message From c408a5f9839fe8951268cfd67f30e1b9d7f9f51a Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 18 May 2017 16:09:25 -0500 Subject: [PATCH 024/109] Added a custom unmarshaller to handle html angle brackets --- multilang/action.go | 48 ++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 43 insertions(+), 5 deletions(-) diff --git a/multilang/action.go b/multilang/action.go index b4aacb1..b5ff8ee 100644 --- a/multilang/action.go +++ b/multilang/action.go @@ -25,7 +25,7 @@ type ActionMessage struct { Action ActionType `json:"action"` ShardID string `json:"shardId,omitempty"` Records []Record `json:"records,omitempty"` - Checkpoint int `json:"checkpoint,omitempty"` + Checkpoint string `json:"checkpoint,omitempty"` Error string `json:"error,omitempty"` Reason string `json:"reason,omitempty"` ResponseFor ActionType `json:"responseFor,omitempty"` @@ -39,6 +39,31 @@ type Record struct { SequenceNumber string `json:"sequenceNumber,omitempty"` } +// UnmarshalJSON is used as a custom unmarshaller to unmarshal the KCL Multilang ActionMessage +func (a *ActionMessage) UnmarshalJSON(data []byte) error { + am := &ActionMessage{} + if err := json.Unmarshal(data, am); err != nil { + return err + } + + a.Action = am.Action + a.ShardID = am.ShardID + a.Records = am.Records + if len(am.Checkpoint) > 2 { + a.Checkpoint = am.Checkpoint[1:len(am.Checkpoint)-1] + } + + if len(am.Error) > 2 { + a.Error = am.Error[1:len(am.Error)-1] + } + + if len(am.Reason) > 2 { + a.Reason = am.Reason[1:len(am.Reason)-1] + } + + return nil +} + // UnmarshalJSON is used as a custom unmarshaller to base64 decode the data field of the KCL Multilang // processRecord message func (r *Record) UnmarshalJSON(data []byte) error { @@ -47,11 +72,24 @@ func (r *Record) UnmarshalJSON(data []byte) error { return err } - decodedMsg, err := base64.StdEncoding.DecodeString(record.Data) - if err != nil { - return err + r.ApproximateArrivalTimestamp = record.ApproximateArrivalTimestamp + + if len(record.Data) > 2 { + encodedString := string(record.Data[1:len(record.Data)-1]) + decodedMsg, err := base64.StdEncoding.DecodeString(encodedString) + if err != nil { + return err + } + r.Data = decodedMsg + } + + if len(record.PartitionKey) > 2 { + r.PartitionKey = record.PartitionKey[1:len(record.PartitionKey)-1] + } + + if len(record.SequenceNumber) > 2 { + r.SequenceNumber = record.SequenceNumber[1:len(record.SequenceNumber)-1] } - r.Data = decodedMsg return nil } From cc096472b09cd06aa9fb5c501a5d83ba0b879de2 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 18 May 2017 16:20:48 -0500 Subject: [PATCH 025/109] Fix gosimple error --- multilang/action.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/multilang/action.go b/multilang/action.go index b5ff8ee..b8aea03 100644 --- a/multilang/action.go +++ b/multilang/action.go @@ -105,7 +105,7 @@ func (r *Record) ToMessage() *message.Message { } // NewCheckpointMessage is used to create a new checkpoint message -func NewCheckpointMessage(seqNum int) *ActionMessage { +func NewCheckpointMessage(seqNum string) *ActionMessage { return &ActionMessage{ Action: Checkpoint, Checkpoint: seqNum, From f5a6907deedc88024a4c7c45c753d27239e2a4e4 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 22 May 2017 09:32:05 -0500 Subject: [PATCH 026/109] Added a function to set the stream reader instead of setting the kinetic stream --- listener/config.go | 10 +++++----- listener/listener_test.go | 2 +- producer/producer_test.go | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/listener/config.go b/listener/config.go index f85886c..f8a2cf6 100644 --- a/listener/config.go +++ b/listener/config.go @@ -49,13 +49,13 @@ func (c *Config) SetConcurrency(concurrency int) { // SetGetRecordsReadTimeout configures the time to wait for each successive // Read operation on the GetRecords response payload. -func (c *Config) SetGetRecordsReadTimeout(timouet time.Duration) { - c.getRecordsReadTimeout = timouet +func (c *Config) SetGetRecordsReadTimeout(timeout time.Duration) { + c.getRecordsReadTimeout = timeout } -// SetKinesisStream sets the listener to read to the given Kinesis stream. -func (c *Config) SetKinesisStream(stream string, shard string, fn ...func(*KinesisReaderConfig)) { - c.reader = NewKinesisReader(stream, shard, fn...) +// SetReader sets the listener's stream reader +func (c *Config) SetReader(reader StreamReader) { + c.reader = reader } // SetLogLevel configures both the SDK and Kinetic log levels. diff --git a/listener/listener_test.go b/listener/listener_test.go index 1343680..213c6c2 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -55,7 +55,7 @@ func TestListener(t *testing.T) { l, err := NewListener(func(c *Config) { c.SetAwsConfig(k.Session.Config) - c.SetKinesisStream(stream, shards[0]) + c.SetReader(NewKinesisReader(stream, shards[0])) c.SetConcurrency(10) }) So(l, ShouldNotBeNil) diff --git a/producer/producer_test.go b/producer/producer_test.go index fbca2f5..66ff6cf 100644 --- a/producer/producer_test.go +++ b/producer/producer_test.go @@ -46,7 +46,7 @@ func TestProducer(t *testing.T) { l, err := listener.NewListener(func(c *listener.Config) { c.SetAwsConfig(k.Session.Config) - c.SetKinesisStream(stream, shards[0]) + c.SetReader(listener.NewKinesisReader(stream, shards[0])) c.SetQueueDepth(10) c.SetConcurrency(10) c.SetGetRecordsReadTimeout(1 * time.Second) From 5cdc2a29a736f0e915cdb26dba73c886c108f56e Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 22 May 2017 11:53:42 -0500 Subject: [PATCH 027/109] Git branching clean up --- listener/kclreader.go | 215 ------------------------------------ listener/kclreaderconfig.go | 31 ------ 2 files changed, 246 deletions(-) delete mode 100644 listener/kclreader.go delete mode 100644 listener/kclreaderconfig.go diff --git a/listener/kclreader.go b/listener/kclreader.go deleted file mode 100644 index 2737430..0000000 --- a/listener/kclreader.go +++ /dev/null @@ -1,215 +0,0 @@ -package listener - -import ( - "bufio" - "encoding/json" - "errors" - "math" - "os" - "sync" - - "github.com/rewardStyle/kinetic/errs" - "github.com/rewardStyle/kinetic/message" - "github.com/rewardStyle/kinetic/multilang" -) - -type kclReaderOptions struct { - onInitCallbackFn func() error - onCheckpointCallbackFn func() error - onShutdownCallbackFn func() error -} - -// KclReader handles the KCL Multilang Protocol to read records from KCL -type KclReader struct { - *kclReaderOptions - throttleSem chan Empty - listener *Listener - scanner *bufio.Scanner - mutex *sync.Mutex - msgBuffer []message.Message - ackPending bool -} - -// NewKclReader creates a new stream reader to read records from KCL -func NewKclReader(fn ...func(*KclReaderConfig)) *KclReader { - config := NewKclReaderConfig() - for _, f := range fn { - f(config) - } - return &KclReader{ - kclReaderOptions: config.kclReaderOptions, - throttleSem: make(chan Empty, 5), - msgBuffer: []message.Message{}, - mutex: &sync.Mutex{}, - } -} - -// AssociateListener associates the KCL stream reader to a listener -func (r *KclReader) AssociateListener(l *Listener) error { - r.mutex.Lock() - defer r.mutex.Unlock() - if r.listener != nil { - return errs.ErrListenerAlreadyAssociated - } - r.listener = l - return nil -} - -// ensureClient will lazily ensure that we are reading from STDIN. -func (r *KclReader) ensureClient() error { - r.mutex.Lock() - defer r.mutex.Unlock() - if r.scanner == nil { - if r.listener == nil { - return errs.ErrNilListener - } - r.scanner = bufio.NewScanner(os.Stdin) - go func() error { - return r.processAction() - }() - } - return nil -} - -// GetRecord calls processRecords to attempt to put one message from message buffer to the listener's message -// channel -func (r *KclReader) GetRecord() (int, error) { - return r.processRecords(1) -} - -// GetRecords calls processRecords to attempt to put all messages on the message buffer on the listener's -// message channel -func (r *KclReader) GetRecords() (int, error) { - return r.processRecords(-1) -} - -// processRecords is a helper method which loops through the message buffer and puts messages on the listener's -// message channel. After all the messages on the message buffer have been moved to the listener's message -// channel, a message is sent (following the Multilang protocol) to acknowledge that the processRecords message -// has been received / processed -func (r *KclReader) processRecords(numRecords int) (int, error) { - if err := r.ensureClient(); err != nil { - return 0, err - } - - r.mutex.Lock() - defer r.mutex.Unlock() - - // Define the batchSize - batchSize := 0; - if len(r.msgBuffer) > 0 { - if numRecords < 0 { - batchSize = len(r.msgBuffer) - } else { - batchSize = int(math.Min(float64(len(r.msgBuffer)), float64(numRecords))) - } - } - - // Loop through the message buffer and put the correct number of messages on the listener's message channel - for i := 0; i < batchSize; i++ { - r.listener.messages <- &r.msgBuffer[0] - r.msgBuffer = r.msgBuffer[1:] - } - - // Send an acknowledgement that the 'ProcessRecords' message was received/processed - if len(r.msgBuffer) == 0 && r.ackPending { - err := r.sendMessage(multilang.NewStatusMessage(multilang.ProcessRecords)) - if err != nil { - r.listener.LogError(err) - return batchSize, err - } - } - - return batchSize, nil -} - -// processAction listens to STDIN and processes action messages based on the Multilanguage protocol from KCL -func (r *KclReader) processAction() error { - for r.scanner.Scan() { - actionMessage := &multilang.ActionMessage{} - err := json.Unmarshal(r.scanner.Bytes(), actionMessage) - if err != nil { - return err - } - - switch actionMessage.Action { - case multilang.Initialize: - r.onInit() - r.sendMessage(multilang.NewStatusMessage(multilang.Initialize)) - case multilang.Checkpoint: - r.onCheckpoint() - r.sendMessage(multilang.NewStatusMessage(multilang.Checkpoint)) - case multilang.Shutdown: - r.onShutdown() - r.sendMessage(multilang.NewStatusMessage(multilang.Shutdown)) - case multilang.ProcessRecords: - go func() error { - r.mutex.Lock() - defer r.mutex.Unlock() - - if r.ackPending { - return errors.New("Received a processRecords action message from KCL " + - "unexpectedly") - } - - for _, msg := range actionMessage.Records { - r.msgBuffer = append(r.msgBuffer, *msg.ToMessage()) - } - r.ackPending = true; - - return nil - }() - default: - } - } - - return nil -} - -func (r *KclReader) sendMessage(msg *multilang.ActionMessage) error { - b, err := json.Marshal(msg) - if err != nil { - r.listener.LogError(err) - return err - } - _, err = os.Stdout.Write(b) - if err != nil { - r.listener.LogError(err) - return err - } - - return nil -} - -func (r *KclReader) onInit() error { - if r.onInitCallbackFn != nil { - err := r.onInitCallbackFn() - if err != nil { - r.listener.LogError(err) - return err - } - } - return nil -} - -func (r *KclReader) onCheckpoint() error { - if r.onCheckpointCallbackFn != nil { - err := r.onCheckpointCallbackFn() - if err != nil { - r.listener.LogError(err) - return err - } - } - return nil -} - -func (r *KclReader) onShutdown() error { - if r.onShutdownCallbackFn != nil { - err := r.onShutdownCallbackFn() - if err != nil { - r.listener.LogError(err) - return err - } - } - return nil -} diff --git a/listener/kclreaderconfig.go b/listener/kclreaderconfig.go deleted file mode 100644 index 13f32ca..0000000 --- a/listener/kclreaderconfig.go +++ /dev/null @@ -1,31 +0,0 @@ -package listener - -// KclReaderConfig is used to configure KclReader -type KclReaderConfig struct { - *kclReaderOptions -} - -// NewKclReaderConfig creates a new instance of KclReaderConfig -func NewKclReaderConfig() *KclReaderConfig { - return &KclReaderConfig{ - kclReaderOptions: &kclReaderOptions{}, - } -} - -// SetOnInitCallbackFn configures a callback function which is run prior to sending a status message -// acknowledging an 'initialize' message was received / processed -func (c *KclReaderConfig) SetOnInitCallbackFn(fn func() error) { - c.onInitCallbackFn = fn -} - -// SetOnCheckpointCallbackFn configures a callback function which is run prior to sending a status message -// acknowledging an 'checkpoint' message was received / processed -func (c *KclReaderConfig) SetOnCheckpointCallbackFn(fn func() error) { - c.onCheckpointCallbackFn = fn -} - -// SetOnShutdownCallbackFn configures a callback function which is run prior to sending a status message -// acknowledging a 'shutdown' message was received / processed -func (c *KclReaderConfig) SetOnShutdownCallbackFn(fn func() error) { - c.onShutdownCallbackFn = fn -} \ No newline at end of file From 77f05e9c2b447e2d61c2e7956c3732ec4bd9a6b8 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 25 May 2017 11:23:56 -0500 Subject: [PATCH 028/109] Added integration testing to test sending multiple messages through a kinesis stream using the kinetic object with a producer and a listener --- kinetic_integration_test.go | 199 ++++++++++++++++++++++++++++++++++++ 1 file changed, 199 insertions(+) create mode 100644 kinetic_integration_test.go diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go new file mode 100644 index 0000000..ad0ec4f --- /dev/null +++ b/kinetic_integration_test.go @@ -0,0 +1,199 @@ +package kinetic + +import ( + "encoding/json" + "context" + "log" + "math/rand" + "strconv" + "sync" + "testing" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/request" + "github.com/rewardStyle/kinetic/listener" + "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic/producer" + "github.com/stretchr/testify/assert" +) + +type Message struct { + Id int `json:"id"` + Message string `json:"message"` +} + +type StreamData struct { + mutex sync.Mutex + Frequencies map[int]int + Messages map[int][]string +} + +func NewStreamData() *StreamData { + return &StreamData{ + mutex: sync.Mutex{}, + Frequencies: make(map[int]int), + Messages: make(map[int][]string), + } +} + +func (m *StreamData) put(key int, value string) { + m.mutex.Lock() + defer m.mutex.Unlock() + + m.Frequencies[key] = m.Frequencies[key] + 1 + m.Messages[key] = append(m.Messages[key], value) +} + +func (m *StreamData) exists(key int) bool { + m.mutex.Lock() + defer m.mutex.Unlock() + + return m.Frequencies[key] > 0 +} + +func (m *StreamData) size() int { + m.mutex.Lock() + defer m.mutex.Unlock() + + return len(m.Messages) +} + +func (m *StreamData) hasDuplicates() bool { + m.mutex.Lock() + defer m.mutex.Unlock() + + for _, freq := range m.Frequencies { + if freq > 1 { + return true + } + } + return false +} + +func (m *StreamData) printSummary() { + m.mutex.Lock() + defer m.mutex.Unlock() + + log.Println("***** Stream Data Summary *****") + log.Printf("Total messages sent/received: [%d]\n", len(m.Messages)) + for index, freq := range m.Frequencies { + if freq > 1 { + log.Printf("Message [%d] occurred [%d] times\n", index, freq) + } + } + log.Println("***** Stream Data Summary *****") +} + +func TestKineticIntegration(t *testing.T) { + + // Set the RNG Seed based on current time (in order to randomize the RNG) + rand.Seed(time.Now().UTC().UnixNano()) + + // Instantiate a new kinentic object + k, err := New(func(c *Config) { + c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") + c.SetRegion("some-region") + c.SetEndpoint("http://127.0.0.1:4567") + }) + assert.NotNil(t, k) + assert.Nil(t, err) + + // Create a kinetic stream + stream := "some-stream-" + strconv.Itoa(rand.Int()) + err = k.CreateStream(stream, 1) + assert.Nil(t, err) + + // Delete the kinetic stream if no dups were found (this is for debugging the kinetic stream) + duplicate := false + defer func() { + if !duplicate { + k.DeleteStream(stream) + k.WaitUntilStreamDeleted(context.TODO(), stream, + request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + } + }() + + // Wait until the stream is ready to go + err = k.WaitUntilStreamExists(context.TODO(), stream, + request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + assert.Nil(t, err) + + // Determine the shard name + shards, err := k.GetShards(stream) + assert.Equal(t, len(shards), 1) + assert.Nil(t, err) + + log.Printf("Stream Name: %s\n", stream) + log.Printf("Shard Name: %s\n", shards[0]) + + // Create a new kinetic producer + p, err := producer.NewProducer(func(c *producer.Config) { + c.SetAwsConfig(k.Session.Config) + c.SetKinesisStream(stream) + c.SetBatchSize(5) + c.SetBatchTimeout(1000 * time.Millisecond) + }) + assert.NotNil(t, p) + assert.Nil(t, err) + + // Create a new kinetic listener + l, err := listener.NewListener(func(c *listener.Config) { + c.SetAwsConfig(k.Session.Config) + c.SetReader(listener.NewKinesisReader(stream, shards[0])) + c.SetQueueDepth(20) + c.SetConcurrency(10) + c.SetGetRecordsReadTimeout(1000 * time.Millisecond) + //c.SetLogLevel(aws.LogDebug) + }) + assert.NotNil(t, l) + assert.Nil(t, err) + + numMsg := 1000 + numSent := 0 + streamData := NewStreamData() + + // Use the producer to write messages to the kinetic stream + wg := sync.WaitGroup{} + wg.Add(numMsg + 1) + go func(sent *int) { + for i := 0; i < numMsg; i++ { + msg := &Message{ + Id: i, + Message: "hello_" + strconv.Itoa(i), + } + jsonStr, _ := json.Marshal(msg) + if err := p.Send(&message.Message { + PartitionKey: aws.String("key"), + Data: []byte(jsonStr), + }); err == nil { + *sent++ + } + } + wg.Done() + }(&numSent) + + // Use the listener to read messages from the kinetic stream + go func() { + l.Listen(func(b []byte, fnwg *sync.WaitGroup) { + msg := &Message{} + json.Unmarshal(b, msg) + + if !streamData.exists(msg.Id) { + wg.Done() + } else { + log.Printf("WARNING: Duplicate message: %v\n", msg) + } + + streamData.put(msg.Id, msg.Message) + fnwg.Done() + }) + }() + + wg.Wait() + assert.Equal(t, int(numSent), numMsg, "Number of message sent should equal the number of messages") + assert.Equal(t, streamData.size(), numMsg, "Number of messages") + + streamData.printSummary() + duplicate = streamData.hasDuplicates() +} From 28d14bd6d5a541b32c32a9e22f4fa746a225784b Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 25 May 2017 11:36:42 -0500 Subject: [PATCH 029/109] Update glide files to pull in assertion package. Fix golint issues. --- glide.lock | 18 +++++++++++++++--- glide.yaml | 2 ++ kinetic_integration_test.go | 8 ++++---- 3 files changed, 21 insertions(+), 7 deletions(-) diff --git a/glide.lock b/glide.lock index c884078..01a6021 100644 --- a/glide.lock +++ b/glide.lock @@ -1,5 +1,5 @@ -hash: 9f381497df21cd2b2e97a96a4526bb2f48443a7d63baaf81e306f4b6c3cdad98 -updated: 2017-03-23T00:24:46.690555736-05:00 +hash: c70e5dd0bac15fcbacf3d07e58cad9df3b23ba556bc583f3a850f79f60ea2bae +updated: 2017-05-25T11:30:54.735250056-05:00 imports: - name: github.com/aws/aws-sdk-go version: 72df63b404d3f9820db08c73176c1b277d9f614f @@ -54,6 +54,10 @@ imports: - convey - convey/gotest - convey/reporting +- name: github.com/stretchr/testify + version: 69483b4bd14f5845b5a1e55bca19e954e827f1d0 + subpackages: + - assert - name: gopkg.in/gcfg.v1 version: 27e4946190b4a327b539185f2b5b1f7c84730728 subpackages: @@ -62,4 +66,12 @@ imports: - types - name: gopkg.in/warnings.v0 version: 8a331561fe74dadba6edfc59f3be66c22c3b065d -testImports: [] +testImports: +- name: github.com/davecgh/go-spew + version: 6d212800a42e8ab5c146b8ace3490ee17e5225f9 + subpackages: + - spew +- name: github.com/pmezard/go-difflib + version: d8ed2627bdf02c080bf22230dbb337003b7aba2d + subpackages: + - difflib diff --git a/glide.yaml b/glide.yaml index afcdda7..fcdcd82 100644 --- a/glide.yaml +++ b/glide.yaml @@ -61,3 +61,5 @@ import: - types - package: gopkg.in/warnings.v0 version: 8a331561fe74dadba6edfc59f3be66c22c3b065d +- package: github.com/stretchr/testify/assert + version: v1.1.4 diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index ad0ec4f..4619760 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -19,7 +19,7 @@ import ( ) type Message struct { - Id int `json:"id"` + ID int `json:"id"` Message string `json:"message"` } @@ -159,7 +159,7 @@ func TestKineticIntegration(t *testing.T) { go func(sent *int) { for i := 0; i < numMsg; i++ { msg := &Message{ - Id: i, + ID: i, Message: "hello_" + strconv.Itoa(i), } jsonStr, _ := json.Marshal(msg) @@ -179,13 +179,13 @@ func TestKineticIntegration(t *testing.T) { msg := &Message{} json.Unmarshal(b, msg) - if !streamData.exists(msg.Id) { + if !streamData.exists(msg.ID) { wg.Done() } else { log.Printf("WARNING: Duplicate message: %v\n", msg) } - streamData.put(msg.Id, msg.Message) + streamData.put(msg.ID, msg.Message) fnwg.Done() }) }() From 3282271ccbe3408974e8a3f93b129f3c522c1312 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 25 May 2017 13:15:39 -0500 Subject: [PATCH 030/109] Adjust the kinetic producer / listener configs to fix race condition causing sporadic unit test failures --- producer/producer_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/producer/producer_test.go b/producer/producer_test.go index 66ff6cf..e80e598 100644 --- a/producer/producer_test.go +++ b/producer/producer_test.go @@ -39,7 +39,7 @@ func TestProducer(t *testing.T) { c.SetAwsConfig(k.Session.Config) c.SetKinesisStream(stream) c.SetBatchSize(5) - c.SetBatchTimeout(1 * time.Second) + c.SetBatchTimeout(1000 * time.Millisecond) }) So(p, ShouldNotBeNil) So(err, ShouldBeNil) @@ -49,7 +49,7 @@ func TestProducer(t *testing.T) { c.SetReader(listener.NewKinesisReader(stream, shards[0])) c.SetQueueDepth(10) c.SetConcurrency(10) - c.SetGetRecordsReadTimeout(1 * time.Second) + c.SetGetRecordsReadTimeout(100 * time.Millisecond) }) So(l, ShouldNotBeNil) So(err, ShouldBeNil) From 636b47d740a7182cec9f8c2e087eb8bfb42d4a72 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 26 May 2017 11:11:55 -0500 Subject: [PATCH 031/109] Added sample application to create a stand-alone executable for running/testing the kinetic producer and listener. Also made enhancements to unit tests. --- example/main.go | 136 ++++++++++++++++++++++++++++++++++++ example/messages.go | 22 ++++++ example/stats.go | 68 ++++++++++++++++++ kinetic_integration_test.go | 18 ++--- listener/listener_test.go | 1 + 5 files changed, 237 insertions(+), 8 deletions(-) create mode 100644 example/main.go create mode 100644 example/messages.go create mode 100644 example/stats.go diff --git a/example/main.go b/example/main.go new file mode 100644 index 0000000..d78199c --- /dev/null +++ b/example/main.go @@ -0,0 +1,136 @@ +package main + +import ( + "encoding/json" + "context" + "log" + "math/rand" + "strconv" + "sync" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/request" + "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/listener" + "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic/producer" +) + +func main() { + + // Set the RNG Seed based on current time (in order to randomize the RNG) + rand.Seed(time.Now().UTC().UnixNano()) + + // Instantiate StreamData Object to keep stats + streamData := NewStreamData() + + // Instantiate a new kinentic object + k, err := kinetic.New(func(c *kinetic.Config) { + c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") + c.SetRegion("some-region") + c.SetEndpoint("http://127.0.0.1:4567") + }) + if err != nil { + // TODO: process error + } + + // Create a kinetic stream + stream := "some-stream-" + strconv.Itoa(rand.Int()) + err = k.CreateStream(stream, 1) + if err != nil { + // TODO: process error + } + + // Delete the kinetic stream if no dups were found (this is for debugging the kinetic stream) + defer func(s *StreamData) { + if !s.hasDuplicates() { + k.DeleteStream(stream) + k.WaitUntilStreamDeleted(context.TODO(), stream, + request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + } + }(streamData) + + // Wait until the stream is ready to go + err = k.WaitUntilStreamExists(context.TODO(), stream, + request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + if err != nil { + // TODO: process error + } + + // Determine the shard name + shards, err := k.GetShards(stream) + if err != nil { + // TODO: process error + } + + log.Printf("Stream Name: %s\n", stream) + log.Printf("Shard Name: %s\n", shards[0]) + + // Create a new kinetic producer + p, err := producer.NewProducer(func(c *producer.Config) { + c.SetAwsConfig(k.Session.Config) + c.SetKinesisStream(stream) + c.SetBatchSize(5) + c.SetBatchTimeout(1000 * time.Millisecond) + }) + if err != nil { + // TODO: process error + } + + // Create a new kinetic listener + l, err := listener.NewListener(func(c *listener.Config) { + c.SetAwsConfig(k.Session.Config) + c.SetReader(listener.NewKinesisReader(stream, shards[0])) + c.SetQueueDepth(20) + c.SetConcurrency(10) + c.SetGetRecordsReadTimeout(1000 * time.Millisecond) + //c.SetLogLevel(aws.LogDebug) + }) + if err != nil { + // TODO: process error + } + + numMsg := 100 + numSent := 0 + + // Use the producer to write messages to the kinetic stream + wg := sync.WaitGroup{} + wg.Add(numMsg + 1) + go func(sent *int) { + for i := 0; i < numMsg; i++ { + jsonStr, _ := json.Marshal(NewMessage()) + if err := p.Send(&message.Message { + PartitionKey: aws.String("key"), + Data: []byte(jsonStr), + }); err == nil { + *sent++ + } + } + wg.Done() + }(&numSent) + + // Use the listener to read messages from the kinetic stream + go func() { + l.Listen(func(b []byte, fnwg *sync.WaitGroup) { + // Unmarshal data + msg := &Message{} + json.Unmarshal(b, msg) + + // Only mark "done" if the message isn't a duplicate + if !streamData.exists(msg.Id) { + wg.Done() + } else { + log.Printf("WARNING: Duplicate message: %v\n", msg) + } + + // Record message regardless if it is a duplicate + streamData.put(msg.Id, msg.Message) + fnwg.Done() + }) + }() + + wg.Wait() + + streamData.printSummary() +} diff --git a/example/messages.go b/example/messages.go new file mode 100644 index 0000000..f726ab9 --- /dev/null +++ b/example/messages.go @@ -0,0 +1,22 @@ +package main + +import ( + "sync/atomic" + "strconv" +) + +var msgCount uint64 = 0 + +type Message struct { + Id int `json:"id"` + Message string `json:"message"` +} + +func NewMessage() *Message { + atomic.AddUint64(&msgCount, 1) + id := int(atomic.LoadUint64(&msgCount)) + return &Message{ + Id: id, + Message: "hello_" + strconv.Itoa(id), + } +} diff --git a/example/stats.go b/example/stats.go new file mode 100644 index 0000000..e8c5953 --- /dev/null +++ b/example/stats.go @@ -0,0 +1,68 @@ +package main + +import ( + "log" + "sync" +) + +type StreamData struct { + mutex sync.Mutex + Frequencies map[int]int + Messages map[int][]string +} + +func NewStreamData() *StreamData { + return &StreamData{ + mutex: sync.Mutex{}, + Frequencies: make(map[int]int), + Messages: make(map[int][]string), + } +} + +func (m *StreamData) put(key int, value string) { + m.mutex.Lock() + defer m.mutex.Unlock() + + m.Frequencies[key] = m.Frequencies[key] + 1 + m.Messages[key] = append(m.Messages[key], value) +} + +func (m *StreamData) exists(key int) bool { + m.mutex.Lock() + defer m.mutex.Unlock() + + return m.Frequencies[key] > 0 +} + +func (m *StreamData) size() int { + m.mutex.Lock() + defer m.mutex.Unlock() + + return len(m.Messages) +} + +func (m *StreamData) hasDuplicates() bool { + m.mutex.Lock() + defer m.mutex.Unlock() + + for _, freq := range m.Frequencies { + if freq > 1 { + return true + } + } + return false +} + +func (m *StreamData) printSummary() { + m.mutex.Lock() + defer m.mutex.Unlock() + + log.Println("***** Stream Data Summary *****") + log.Printf("Total messages sent/received: [%d]\n", len(m.Messages)) + for index, freq := range m.Frequencies { + if freq > 1 { + log.Printf("Message [%d] was received [%d] times\n", index, freq) + } + } + log.Println("***** Stream Data Summary *****") +} diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index 4619760..f266ef6 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -90,6 +90,9 @@ func TestKineticIntegration(t *testing.T) { // Set the RNG Seed based on current time (in order to randomize the RNG) rand.Seed(time.Now().UTC().UnixNano()) + // Instantiate StreamData Object to keep stats + streamData := NewStreamData() + // Instantiate a new kinentic object k, err := New(func(c *Config) { c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") @@ -105,14 +108,13 @@ func TestKineticIntegration(t *testing.T) { assert.Nil(t, err) // Delete the kinetic stream if no dups were found (this is for debugging the kinetic stream) - duplicate := false - defer func() { - if !duplicate { + defer func(s *StreamData) { + if !s.hasDuplicates() { k.DeleteStream(stream) k.WaitUntilStreamDeleted(context.TODO(), stream, request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) } - }() + }(streamData) // Wait until the stream is ready to go err = k.WaitUntilStreamExists(context.TODO(), stream, @@ -151,12 +153,13 @@ func TestKineticIntegration(t *testing.T) { numMsg := 1000 numSent := 0 - streamData := NewStreamData() + // Use the producer to write messages to the kinetic stream wg := sync.WaitGroup{} wg.Add(numMsg + 1) go func(sent *int) { + defer wg.Done() for i := 0; i < numMsg; i++ { msg := &Message{ ID: i, @@ -170,12 +173,13 @@ func TestKineticIntegration(t *testing.T) { *sent++ } } - wg.Done() }(&numSent) // Use the listener to read messages from the kinetic stream go func() { l.Listen(func(b []byte, fnwg *sync.WaitGroup) { + defer fnwg.Done() + msg := &Message{} json.Unmarshal(b, msg) @@ -186,7 +190,6 @@ func TestKineticIntegration(t *testing.T) { } streamData.put(msg.ID, msg.Message) - fnwg.Done() }) }() @@ -195,5 +198,4 @@ func TestKineticIntegration(t *testing.T) { assert.Equal(t, streamData.size(), numMsg, "Number of messages") streamData.printSummary() - duplicate = streamData.hasDuplicates() } diff --git a/listener/listener_test.go b/listener/listener_test.go index 213c6c2..bd81d2b 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -56,6 +56,7 @@ func TestListener(t *testing.T) { l, err := NewListener(func(c *Config) { c.SetAwsConfig(k.Session.Config) c.SetReader(NewKinesisReader(stream, shards[0])) + c.SetQueueDepth(10) c.SetConcurrency(10) }) So(l, ShouldNotBeNil) From 90b6288c47fb84e6fd489738afc055c3b1c54003 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 26 May 2017 12:28:51 -0500 Subject: [PATCH 032/109] Fix golint errors --- example/main.go | 6 +++--- example/messages.go | 8 +++++--- example/stats.go | 2 ++ 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/example/main.go b/example/main.go index d78199c..f27ab6f 100644 --- a/example/main.go +++ b/example/main.go @@ -91,7 +91,7 @@ func main() { // TODO: process error } - numMsg := 100 + numMsg := 1000000 numSent := 0 // Use the producer to write messages to the kinetic stream @@ -118,14 +118,14 @@ func main() { json.Unmarshal(b, msg) // Only mark "done" if the message isn't a duplicate - if !streamData.exists(msg.Id) { + if !streamData.exists(msg.ID) { wg.Done() } else { log.Printf("WARNING: Duplicate message: %v\n", msg) } // Record message regardless if it is a duplicate - streamData.put(msg.Id, msg.Message) + streamData.put(msg.ID, msg.Message) fnwg.Done() }) }() diff --git a/example/messages.go b/example/messages.go index f726ab9..88b8e15 100644 --- a/example/messages.go +++ b/example/messages.go @@ -5,18 +5,20 @@ import ( "strconv" ) -var msgCount uint64 = 0 +var msgCount uint64 +// Message is a data structure for sending / receiving messages over a kinetic stream type Message struct { - Id int `json:"id"` + ID int `json:"id"` Message string `json:"message"` } +// NewMessage creates a new Message struct with a unique identifier func NewMessage() *Message { atomic.AddUint64(&msgCount, 1) id := int(atomic.LoadUint64(&msgCount)) return &Message{ - Id: id, + ID: id, Message: "hello_" + strconv.Itoa(id), } } diff --git a/example/stats.go b/example/stats.go index e8c5953..c0e852b 100644 --- a/example/stats.go +++ b/example/stats.go @@ -5,12 +5,14 @@ import ( "sync" ) +// StreamData is used to collect stream stats type StreamData struct { mutex sync.Mutex Frequencies map[int]int Messages map[int][]string } +// NewStreamData instantiates a new StreamData struct func NewStreamData() *StreamData { return &StreamData{ mutex: sync.Mutex{}, From 6a24e8e7b449277a1adb3178fb0378749130baec Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 26 May 2017 14:54:59 -0500 Subject: [PATCH 033/109] Added exception logging and more code clean up --- example/main.go | 27 ++++++++++++++------------- example/stats.go | 13 ++++++++++++- kinetic_integration_test.go | 10 +++++----- 3 files changed, 31 insertions(+), 19 deletions(-) diff --git a/example/main.go b/example/main.go index f27ab6f..9cbb350 100644 --- a/example/main.go +++ b/example/main.go @@ -32,14 +32,21 @@ func main() { c.SetEndpoint("http://127.0.0.1:4567") }) if err != nil { - // TODO: process error + log.Fatalf("Unable to create new kinetic object due to: %v\n", err) } // Create a kinetic stream stream := "some-stream-" + strconv.Itoa(rand.Int()) err = k.CreateStream(stream, 1) if err != nil { - // TODO: process error + log.Fatalf("Unable to create new stream %s due to: %v\n", stream, err) + } + + // Wait until the stream is ready to go + err = k.WaitUntilStreamExists(context.TODO(), stream, + request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + if err != nil { + log.Fatalf("Unable to wait until stream %s exists due to: %v\n", stream, err) } // Delete the kinetic stream if no dups were found (this is for debugging the kinetic stream) @@ -51,17 +58,10 @@ func main() { } }(streamData) - // Wait until the stream is ready to go - err = k.WaitUntilStreamExists(context.TODO(), stream, - request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) - if err != nil { - // TODO: process error - } - // Determine the shard name shards, err := k.GetShards(stream) if err != nil { - // TODO: process error + log.Fatalf("Unable to get shards for stream %s due to: %v\n", stream, err) } log.Printf("Stream Name: %s\n", stream) @@ -75,7 +75,7 @@ func main() { c.SetBatchTimeout(1000 * time.Millisecond) }) if err != nil { - // TODO: process error + log.Fatalf("Unable to create a new producer due to: %v\n", err) } // Create a new kinetic listener @@ -88,15 +88,15 @@ func main() { //c.SetLogLevel(aws.LogDebug) }) if err != nil { - // TODO: process error + log.Fatalf("Unable to create a new listener due to: %v\n", err) } numMsg := 1000000 - numSent := 0 // Use the producer to write messages to the kinetic stream wg := sync.WaitGroup{} wg.Add(numMsg + 1) + numSent := 0 go func(sent *int) { for i := 0; i < numMsg; i++ { jsonStr, _ := json.Marshal(NewMessage()) @@ -132,5 +132,6 @@ func main() { wg.Wait() + streamData.setMsgCount(numSent) streamData.printSummary() } diff --git a/example/stats.go b/example/stats.go index c0e852b..5fa7f7e 100644 --- a/example/stats.go +++ b/example/stats.go @@ -8,6 +8,7 @@ import ( // StreamData is used to collect stream stats type StreamData struct { mutex sync.Mutex + MsgCount int Frequencies map[int]int Messages map[int][]string } @@ -21,6 +22,13 @@ func NewStreamData() *StreamData { } } +func (m *StreamData) setMsgCount(i int) { + m.mutex.Lock() + defer m.mutex.Unlock() + + m.MsgCount = i +} + func (m *StreamData) put(key int, value string) { m.mutex.Lock() defer m.mutex.Unlock() @@ -59,12 +67,15 @@ func (m *StreamData) printSummary() { m.mutex.Lock() defer m.mutex.Unlock() + log.Println() log.Println("***** Stream Data Summary *****") - log.Printf("Total messages sent/received: [%d]\n", len(m.Messages)) + log.Printf("Total messages sent: [%d]\n", m.MsgCount) + log.Printf("Total messages received: [%d]\n", len(m.Messages)) for index, freq := range m.Frequencies { if freq > 1 { log.Printf("Message [%d] was received [%d] times\n", index, freq) } } log.Println("***** Stream Data Summary *****") + log.Println() } diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index f266ef6..58aa095 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -107,6 +107,11 @@ func TestKineticIntegration(t *testing.T) { err = k.CreateStream(stream, 1) assert.Nil(t, err) + // Wait until the stream is ready to go + err = k.WaitUntilStreamExists(context.TODO(), stream, + request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + assert.Nil(t, err) + // Delete the kinetic stream if no dups were found (this is for debugging the kinetic stream) defer func(s *StreamData) { if !s.hasDuplicates() { @@ -116,11 +121,6 @@ func TestKineticIntegration(t *testing.T) { } }(streamData) - // Wait until the stream is ready to go - err = k.WaitUntilStreamExists(context.TODO(), stream, - request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) - assert.Nil(t, err) - // Determine the shard name shards, err := k.GetShards(stream) assert.Equal(t, len(shards), 1) From 538f5b8b1137f7e5beaa9a10328ce9ca49756752 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 1 Jun 2017 10:44:15 -0500 Subject: [PATCH 034/109] Refactored executable test program to add the ability to 1) stream to an AWS Kinesis stream and 2) run for either a fixed or indefinite duration of time. Fixed potential deadlock conditions and early termination conditions. --- example/main.go | 137 ------------ testexec/README.md | 64 ++++++ testexec/cli.go | 75 +++++++ testexec/main.go | 361 ++++++++++++++++++++++++++++++ {example => testexec}/messages.go | 6 +- {example => testexec}/stats.go | 49 ++-- 6 files changed, 533 insertions(+), 159 deletions(-) delete mode 100644 example/main.go create mode 100644 testexec/README.md create mode 100644 testexec/cli.go create mode 100644 testexec/main.go rename {example => testexec}/messages.go (77%) rename {example => testexec}/stats.go (54%) diff --git a/example/main.go b/example/main.go deleted file mode 100644 index 9cbb350..0000000 --- a/example/main.go +++ /dev/null @@ -1,137 +0,0 @@ -package main - -import ( - "encoding/json" - "context" - "log" - "math/rand" - "strconv" - "sync" - "time" - - "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/aws/request" - "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/listener" - "github.com/rewardStyle/kinetic/message" - "github.com/rewardStyle/kinetic/producer" -) - -func main() { - - // Set the RNG Seed based on current time (in order to randomize the RNG) - rand.Seed(time.Now().UTC().UnixNano()) - - // Instantiate StreamData Object to keep stats - streamData := NewStreamData() - - // Instantiate a new kinentic object - k, err := kinetic.New(func(c *kinetic.Config) { - c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") - c.SetRegion("some-region") - c.SetEndpoint("http://127.0.0.1:4567") - }) - if err != nil { - log.Fatalf("Unable to create new kinetic object due to: %v\n", err) - } - - // Create a kinetic stream - stream := "some-stream-" + strconv.Itoa(rand.Int()) - err = k.CreateStream(stream, 1) - if err != nil { - log.Fatalf("Unable to create new stream %s due to: %v\n", stream, err) - } - - // Wait until the stream is ready to go - err = k.WaitUntilStreamExists(context.TODO(), stream, - request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) - if err != nil { - log.Fatalf("Unable to wait until stream %s exists due to: %v\n", stream, err) - } - - // Delete the kinetic stream if no dups were found (this is for debugging the kinetic stream) - defer func(s *StreamData) { - if !s.hasDuplicates() { - k.DeleteStream(stream) - k.WaitUntilStreamDeleted(context.TODO(), stream, - request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) - } - }(streamData) - - // Determine the shard name - shards, err := k.GetShards(stream) - if err != nil { - log.Fatalf("Unable to get shards for stream %s due to: %v\n", stream, err) - } - - log.Printf("Stream Name: %s\n", stream) - log.Printf("Shard Name: %s\n", shards[0]) - - // Create a new kinetic producer - p, err := producer.NewProducer(func(c *producer.Config) { - c.SetAwsConfig(k.Session.Config) - c.SetKinesisStream(stream) - c.SetBatchSize(5) - c.SetBatchTimeout(1000 * time.Millisecond) - }) - if err != nil { - log.Fatalf("Unable to create a new producer due to: %v\n", err) - } - - // Create a new kinetic listener - l, err := listener.NewListener(func(c *listener.Config) { - c.SetAwsConfig(k.Session.Config) - c.SetReader(listener.NewKinesisReader(stream, shards[0])) - c.SetQueueDepth(20) - c.SetConcurrency(10) - c.SetGetRecordsReadTimeout(1000 * time.Millisecond) - //c.SetLogLevel(aws.LogDebug) - }) - if err != nil { - log.Fatalf("Unable to create a new listener due to: %v\n", err) - } - - numMsg := 1000000 - - // Use the producer to write messages to the kinetic stream - wg := sync.WaitGroup{} - wg.Add(numMsg + 1) - numSent := 0 - go func(sent *int) { - for i := 0; i < numMsg; i++ { - jsonStr, _ := json.Marshal(NewMessage()) - if err := p.Send(&message.Message { - PartitionKey: aws.String("key"), - Data: []byte(jsonStr), - }); err == nil { - *sent++ - } - } - wg.Done() - }(&numSent) - - // Use the listener to read messages from the kinetic stream - go func() { - l.Listen(func(b []byte, fnwg *sync.WaitGroup) { - // Unmarshal data - msg := &Message{} - json.Unmarshal(b, msg) - - // Only mark "done" if the message isn't a duplicate - if !streamData.exists(msg.ID) { - wg.Done() - } else { - log.Printf("WARNING: Duplicate message: %v\n", msg) - } - - // Record message regardless if it is a duplicate - streamData.put(msg.ID, msg.Message) - fnwg.Done() - }) - }() - - wg.Wait() - - streamData.setMsgCount(numSent) - streamData.printSummary() -} diff --git a/testexec/README.md b/testexec/README.md new file mode 100644 index 0000000..bc3cbb3 --- /dev/null +++ b/testexec/README.md @@ -0,0 +1,64 @@ +# testexec +testexec is an executable CLI program to test the kinetic producer / lister. + +The executable program creates a kinetic object connecting to either a local kinesalte instance or to AWS Kinesis using local AWS credentials/config. + +A new stream will be created (with a random stream name) unless a stream name is provided (See useage). + +A kinetic producer will stream dummy data to the kinesis stream and a kinetic listener will stream from the kinesis stream. + +The program logs stream data stats to the console log periodically. + +## Installation +```sh +cd ${GOHOME}/src/github.com/rewardStyle/kinetic/testexec` +go install +``` + +## Requirements + +- local + - requires kinesalite running on http://127.0.0.1:4567 + +- aws + - ~/.aws/credentials to exist and contain valid aws credentials + - ~/.aws/config to exist and contain AWS configuration settings + - environment variable `AWS_SDK_LOAD_CONFIG` should exist and be set to `true` + - environment variable `AWS_PROFILE` should exist and be set to the preferred AWS profile + + +## Usage + +```text +Usage of ./testexec: + -cleanup + used to specify whether or not to delete a newly created kinesis stream (default true) + -duration int + used to specify the duration (in seconds) the program should run. Use a value of -1 to run indefinitely. Either -num-msgs or -duration must be set. + -location string + used to specify the location of the kinesis stream. Accepted values are (local|aws). For local, run kinesalite on http://127.0.0.1:4567. For aws, your aws credentials and configuration needs to be defined at ~/.aws (default "local") + -num-msgs int + used to specify the number of messages to (attempt to) send / receive. Either -num-msgs or -duration must be set. + -stream-name string + used to specify a pre-existing stream to be used for testing. A new stream will be created if not defined. + -verbose + used to specify whether or not to log in verbose mode +``` + +## Examples + +To run kinetic testexec on a local kinesalite instance to stream a fixed number of messages to a new kinesis stream: +```sh +./testexec -num-msgs 1000 -verbose +``` + +To run kinetic testexec on a local kinesalite instance to stream for a fixed duration of time to an existing kinesis stream: +```sh +./testexec -location local -stream-name some-stream -duration 1000 +``` + +To run kinetic testexec on an AWS Kinesis Stream to stream indefinitely: +```sh +./testexec -location aws -duration -1 -cleanup=false -verbose +``` + diff --git a/testexec/cli.go b/testexec/cli.go new file mode 100644 index 0000000..ef40815 --- /dev/null +++ b/testexec/cli.go @@ -0,0 +1,75 @@ +package main + +import ( + "flag" + "log" +) + +// Config is a data structure used to hold this program's configuration info +type Config struct { + Location *string + StreamName *string + NumMsgs *int + Duration *int + Cleanup *bool + Verbose *bool +} + +func parseCommandLineArgs() *Config { + + // Define command line flags + locationPtr := flag.String("location", "local", "used to specify the location of the kinesis stream. " + + "Accepted values are (local|aws). For local, run kinesalite on http://127.0.0.1:4567. For aws, your " + + "aws credentials and configuration needs to be defined at ~/.aws") + streamNamePtr := flag.String("stream-name", "", "used to specify a pre-existing stream to be used for " + + "testing. A new stream will be created if not defined.") + numMsgsPtr := flag.Int("num-msgs", 0, "used to specify the number of messages to (attempt to) send / " + + "receive. Either -num-msgs or -duration must be set.") + durationPtr := flag.Int("duration", 0, "used to specify the duration (in seconds) the program should run. " + + "Use a value of -1 to run indefinitely. Either -num-msgs or -duration must be set.") + cleanupPtr := flag.Bool("cleanup", true, "used to specify whether or not to delete a newly created kinesis " + + "stream") + verbosePtr := flag.Bool("verbose", false, "used to specify whether or not to log in verbose mode") + + // Parse command line arguments + flag.Parse() + + // Process command line arguments + if *numMsgsPtr == 0 && *durationPtr == 0 { + log.Fatal("Either -num-msgs or -duration must be set.") + } else if *numMsgsPtr != 0 && *durationPtr != 0 { + log.Fatal("Both -num-msgs and -duration were set. Only one may be set.") + } else if *durationPtr != 0 { + numMsgsPtr = nil + } else if *numMsgsPtr < 0 { + log.Fatal("Number of messages value must be greater than 0") + } else { + durationPtr = nil + } + + return &Config{ + StreamName: streamNamePtr, + Duration: durationPtr, + NumMsgs: numMsgsPtr, + Location: locationPtr, + Cleanup: cleanupPtr, + Verbose: verbosePtr, + } +} + +func (c *Config) printConfigs() { + if *c.Verbose { + log.Println("Command Line Arguments:") + log.Println("-location: ", *c.Location) + log.Println("-stream-name: ", *c.StreamName) + if c.NumMsgs != nil { + log.Println("-num-msgs: ", *c.NumMsgs) + } + if c.Duration != nil { + log.Println("-duration: ", *c.Duration) + } + log.Println("-cleanup: ", *c.Cleanup) + log.Println("-verbose: ", *c.Verbose) + log.Println() + } +} diff --git a/testexec/main.go b/testexec/main.go new file mode 100644 index 0000000..4381fe2 --- /dev/null +++ b/testexec/main.go @@ -0,0 +1,361 @@ +package main + +import ( + "context" + "encoding/json" + "fmt" + "log" + "math/rand" + "os" + "os/signal" + "os/user" + "strings" + "sync" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/request" + "github.com/aws/aws-sdk-go/aws/session" + "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/listener" + "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic/producer" +) + +// Declare global variables +var startTime time.Time +var streamStart time.Time +var streamStop time.Time +var pipeOfDeath chan os.Signal +var stopDisplay chan struct{} +var stopProducing chan struct{} +var config *Config + +func init() { + // Start the stopwatch + startTime = time.Now() + + // Set the RNG Seed based on current time (in order to randomize the RNG) + rand.Seed(startTime.UTC().UnixNano()) + + // Instantiate channels for communicating between threads + pipeOfDeath = make(chan os.Signal, 1) + stopDisplay = make(chan struct{}, 1) + stopProducing = make(chan struct{}, 1) + + // Set up pipeOfDeath channel to receive os signals + signal.Notify(pipeOfDeath, os.Interrupt) +} + +func cleanup(k *kinetic.Kinetic, streamName string) { + if *config.Cleanup { + if *config.Verbose { + log.Println() + log.Printf("Cleaning up by deleting stream [%s] ...\n", streamName) + } + k.DeleteStream(streamName) + if *config.Verbose { + log.Printf("Waiting for stream [%s] to be deleted ...\n", streamName) + } + k.WaitUntilStreamDeleted(context.TODO(), streamName, + request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + if *config.Verbose { + log.Println("Clean up complete") + log.Println() + } + } +} + +func main() { + // Process command line arguments + config = parseCommandLineArgs() + config.printConfigs() + + // Instantiate a new kinetic object based on the location of the kinesis stream + var k *kinetic.Kinetic + switch strings.ToLower(*config.Location) { + case "local": + k = newDefaultKinetic() + case "aws": + k = newAwsKinetic() + default : + log.Fatalf("Unknown location for kinesis stream: %s\n", *config.Location) + } + + // Set streamName from config or generate a random one + streamName := *config.StreamName + if streamName == "" { + streamName = generateRandomStreamName() + } + + // Create and wait for a new kinetic stream (if the stream name doesn't already exist) + prepareKinesisStream(k, streamName) + + // Create a new kinetic producer + p := newKineticProducer(k, streamName) + + // Create a new kinetic listener + l := newKineticListener(k, streamName) + + // Instantiate StreamData Object to keep stats + streamData := NewStreamData() + + // Display stream data statistics to the console + displayWg := sync.WaitGroup{} + displayWg.Add(1) + go func(sd *StreamData) { + displayLoop: + for { + select { + case <-pipeOfDeath: + stopProducing <- struct{}{} + if *config.Verbose { + log.Println() + log.Println("display: Received pipeOfDeath ...") + } + case <-stopDisplay: + if *config.Verbose { + log.Println() + log.Print("display: Received stopDisplay ...") + } + break displayLoop + default: + time.Sleep(1000 * time.Millisecond) + log.Println() + log.Printf("Stream name: %s\n", streamName) + log.Printf("Elapsed Time: %v\n", time.Since(startTime)) + if streamStop.IsZero() { + log.Printf("Streaming Time: %v\n", time.Since(streamStart)) + } else { + log.Printf("Streaming Time: %v\n", streamStop.Sub(streamStart)) + } + sd.printStats() + } + } + streamData.printSummary() + displayWg.Done() + }(streamData) + + // Use the producer to write messages to the kinetic stream + streamWg := sync.WaitGroup{} + streamWg.Add(1) + go func(sd *StreamData) { + streamStart = time.Now() + if config.Duration != nil { + // Set a timeout based on the config + var timeout <-chan time.Time + if *config.Duration > 0 { + timeout = time.After(time.Duration(*config.Duration) * time.Second) + } else { + timeout = make(chan time.Time, 1) + } + + produceLoop: + for { + select { + case <-stopProducing: + if *config.Verbose { + log.Println() + log.Print("producer: Received stopProducing ...") + } + streamStop = time.Now() + break produceLoop + case <-timeout: + if *config.Verbose { + log.Println() + log.Print("producer: Timed out ...") + } + streamStop = time.Now() + break produceLoop + default: + jsonStr, _ := json.Marshal(NewMessage()) + if err := p.Send(&message.Message { + PartitionKey: aws.String("key"), + Data: []byte(jsonStr), + }); err == nil { + streamWg.Add(1) + sd.incrementMsgCount() + } + } + } + } else if config.NumMsgs != nil { + for i := 0; i < *config.NumMsgs; i++ { + jsonStr, _ := json.Marshal(NewMessage()) + if err := p.Send(&message.Message { + PartitionKey: aws.String("key"), + Data: []byte(jsonStr), + }); err == nil { + streamWg.Add(1) + sd.incrementMsgCount() + } + } + streamStop = time.Now() + } + streamWg.Done() + }(streamData) + + // Use the listener to read messages from the kinetic stream + go func(sd *StreamData) { + l.Listen(func(b []byte, fnwg *sync.WaitGroup) { + // Unmarshal data + msg := &Message{} + json.Unmarshal(b, msg) + + // Only mark "done" if the message isn't a duplicate + if !sd.exists(msg.ID) { + streamWg.Done() + } else { + if *config.Verbose { + log.Printf("listner: Duplicate message: %v\n", msg) + } + } + + // Record message regardless if it is a duplicate + sd.mark(msg.ID) + fnwg.Done() + }) + }(streamData) + + // Wait until streaming is complete + streamWg.Wait() + stopDisplay <- struct{}{} + + // Wait until output display is complete + displayWg.Wait() + cleanup(k, streamName) +} + +func newDefaultKinetic() *kinetic.Kinetic { + k, err := kinetic.New(func(c *kinetic.Config) { + c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") + c.SetRegion("some-region") + c.SetEndpoint("http://127.0.0.1:4567") + }) + if err != nil { + log.Fatalf("Unable to create new default kinetic object due to: %v\n", err) + } + + return k +} + +func newAwsKinetic() *kinetic.Kinetic { + // Verify that ~/.aws/credentials file exists + // TODO: Fix this bug + //if _, err := os.Stat("~/.aws/credentials"); os.IsNotExist(err) { + // log.Fatal("~/.aws/credentials does not exist") + //} + + // Verify that ~/.aws/config file exists + // TODO: Fix this bug + //if _, err := os.Stat("~/.aws/config"); os.IsNotExist(err) { + // log.Fatal("~/.aws/config does not exist") + //} + + // Verify that AWS_SDK_LOAD_CONFIG is set as an environment variable + if val, found := os.LookupEnv("AWS_SDK_LOAD_CONFIG"); val != "true" || !found { + log.Fatal("Environemnt variable AWS_SDK_LOAD_CONFIG must be defined and true") + } + + // Verify that AWS_PROFILE is set as an environment variable + if _, found := os.LookupEnv("AWS_PROFILE"); !found { + log.Fatal("Environemnt variable AWS_PROFILE must be defined") + } + + // Establish an AWS session + sess := session.Must(session.NewSession()) + creds, _ := sess.Config.Credentials.Get() + + // Instantiate a new kinetic object configured with appropriate configs + k, err := kinetic.New(func(c *kinetic.Config) { + c.SetCredentials(creds.AccessKeyID, creds.SecretAccessKey, creds.SessionToken) + c.SetRegion(*sess.Config.Region) + }) + if err != nil { + log.Fatalf("Unable to create new aws kinetic object due to: %v\n", err) + } + + return k +} + +func generateRandomStreamName() string { + baseName := "test-" + if cu, err := user.Current(); err == nil { + baseName += cu.Username + } + + return fmt.Sprintf("%s-%09d", baseName, rand.Intn(999999999)) +} + +func prepareKinesisStream(k *kinetic.Kinetic, streamName string) { + if *config.Verbose { + log.Printf("Preparing kinesis stream: [%s] ...\n", streamName) + } + + // Determine if the kinesis stream exists + err := k.WaitUntilStreamExists(context.TODO(), streamName, + request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + if err != nil { + // If not, create a kinetic stream + if *config.Verbose { + log.Printf("Creating a new kinesis stream: [%s] ...\n", streamName) + } + err := k.CreateStream(streamName, 1) + if err != nil { + log.Fatalf("Unable to create new stream %s due to: %v\n", streamName, err) + } + + // And wait until the stream is ready to go + if *config.Verbose { + log.Printf("Waiting for new kinesis stream: [%s]\n", streamName) + } + err = k.WaitUntilStreamExists(context.TODO(), streamName, + request.WithWaiterDelay(request.ConstantWaiterDelay(2*time.Second))) + if err != nil { + log.Fatalf("Unable to wait until stream %s exists due to: %v\n", streamName, err) + } + } +} + +func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Producer { + if *config.Verbose { + log.Println("Creating a kinetic producer ...") + } + + p, err := producer.NewProducer(func(c *producer.Config) { + c.SetAwsConfig(k.Session.Config) + c.SetKinesisStream(streamName) + c.SetBatchSize(5) + c.SetBatchTimeout(1000 * time.Millisecond) + }) + if err != nil { + log.Fatalf("Unable to create a new producer due to: %v\n", err) + } + + return p +} + +func newKineticListener(k *kinetic.Kinetic, streamName string) *listener.Listener { + if *config.Verbose { + log.Println("Creating a kinetic listener ...") + } + + // Determine the shard name + shards, err := k.GetShards(streamName) + if err != nil { + log.Fatalf("Unable to get shards for stream %s due to: %v\n", streamName, err) + } + + l, err := listener.NewListener(func(c *listener.Config) { + c.SetAwsConfig(k.Session.Config) + c.SetReader(listener.NewKinesisReader(streamName, shards[0])) + c.SetQueueDepth(20) + c.SetConcurrency(10) + c.SetGetRecordsReadTimeout(1000 * time.Millisecond) + //c.SetLogLevel(aws.LogDebug) + }) + if err != nil { + log.Fatalf("Unable to create a new listener due to: %v\n", err) + } + + return l +} diff --git a/example/messages.go b/testexec/messages.go similarity index 77% rename from example/messages.go rename to testexec/messages.go index 88b8e15..730e50d 100644 --- a/example/messages.go +++ b/testexec/messages.go @@ -9,16 +9,16 @@ var msgCount uint64 // Message is a data structure for sending / receiving messages over a kinetic stream type Message struct { - ID int `json:"id"` + ID uint64 `json:"id"` Message string `json:"message"` } // NewMessage creates a new Message struct with a unique identifier func NewMessage() *Message { atomic.AddUint64(&msgCount, 1) - id := int(atomic.LoadUint64(&msgCount)) + id := atomic.LoadUint64(&msgCount) return &Message{ ID: id, - Message: "hello_" + strconv.Itoa(id), + Message: "hello_" + strconv.Itoa(int(id)), } } diff --git a/example/stats.go b/testexec/stats.go similarity index 54% rename from example/stats.go rename to testexec/stats.go index 5fa7f7e..610229f 100644 --- a/example/stats.go +++ b/testexec/stats.go @@ -3,52 +3,53 @@ package main import ( "log" "sync" + "sync/atomic" ) // StreamData is used to collect stream stats type StreamData struct { mutex sync.Mutex - MsgCount int - Frequencies map[int]int - Messages map[int][]string + MsgCount uint64 + Frequencies map[uint64]uint64 + Duplicates uint64 } // NewStreamData instantiates a new StreamData struct func NewStreamData() *StreamData { return &StreamData{ mutex: sync.Mutex{}, - Frequencies: make(map[int]int), - Messages: make(map[int][]string), + Frequencies: make(map[uint64]uint64), } } -func (m *StreamData) setMsgCount(i int) { - m.mutex.Lock() - defer m.mutex.Unlock() - - m.MsgCount = i +func (m *StreamData) incrementMsgCount() { + atomic.AddUint64(&m.MsgCount, 1) } -func (m *StreamData) put(key int, value string) { +func (m *StreamData) mark(key uint64) { m.mutex.Lock() defer m.mutex.Unlock() m.Frequencies[key] = m.Frequencies[key] + 1 - m.Messages[key] = append(m.Messages[key], value) + if m.Frequencies[key] > 1 { + m.Duplicates++ + } } -func (m *StreamData) exists(key int) bool { +func (m *StreamData) exists(key uint64) bool { m.mutex.Lock() defer m.mutex.Unlock() - return m.Frequencies[key] > 0 + _, ok := m.Frequencies[key] + + return ok } -func (m *StreamData) size() int { +func (m *StreamData) size() uint64 { m.mutex.Lock() defer m.mutex.Unlock() - return len(m.Messages) + return uint64(len(m.Frequencies)) } func (m *StreamData) hasDuplicates() bool { @@ -63,6 +64,17 @@ func (m *StreamData) hasDuplicates() bool { return false } +func (m *StreamData) printStats() { + m.mutex.Lock() + defer m.mutex.Unlock() + + log.Println() + log.Println("***** Stream Data Stats *****") + log.Printf("Messages sent: [%d]\n", m.MsgCount) + log.Printf("Messages received: [%d]\n", len(m.Frequencies)) + log.Printf("Number of duplicated messages: [%d]\n", m.Duplicates) +} + func (m *StreamData) printSummary() { m.mutex.Lock() defer m.mutex.Unlock() @@ -70,12 +82,11 @@ func (m *StreamData) printSummary() { log.Println() log.Println("***** Stream Data Summary *****") log.Printf("Total messages sent: [%d]\n", m.MsgCount) - log.Printf("Total messages received: [%d]\n", len(m.Messages)) + log.Printf("Total messages received: [%d]\n", len(m.Frequencies)) + log.Printf("Total Number of duplicated messages: [%d]\n", m.Duplicates) for index, freq := range m.Frequencies { if freq > 1 { log.Printf("Message [%d] was received [%d] times\n", index, freq) } } - log.Println("***** Stream Data Summary *****") - log.Println() } From 9702c17645558adf5e4ccc154f5366aa90861595 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 2 Jun 2017 10:45:08 -0500 Subject: [PATCH 035/109] Modified the README.md file to update usage examples for the producer / listener --- README.md | 62 ++++++++++++++++++++++++---------------------- testexec/README.md | 2 +- 2 files changed, 34 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index ffb16c6..2ebcbdd 100644 --- a/README.md +++ b/README.md @@ -31,45 +31,49 @@ Kinetic can be used to interface with kinesis like so: ```go -import "github.com/rewardStyle/kinetic" - -// Use configuration in /etc/kinetic.conf -listener, _ := new(kinetic.Listener).Init() +import ( + "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/listener" + "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic/producer" + "sync" +) + +// Create a kinetic object associated with a local kinesalite stream +k, _ := kinetic.New(func(c *kinetic.Config) { + c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") + c.SetRegion("some-region") + c.SetEndpoint("http://127.0.0.1:4567") +}) -// Use custom configuration -producer, _ := new(kinetic.Producer).InitC("your-stream", "0", "shard-type", "accesskey", "secretkey", "region", 10) +// Create a kinetic producer +p, _ := producer.NewProducer(func(c *producer.Config) { + c.SetAwsConfig(k.Session.Config) + c.SetKinesisStream("stream-name") +}) -producer.Send(new(kinetic.Message).Init([]byte(`{"foo":"bar"}`), "test")) +// Create a kinetic listener +l, _ := listener.NewListener(func(c *listener.Config) { + c.SetAwsConfig(k.Session.Config) + c.SetReader(listener.NewKinesisReader("stream-name", "shard-name")) +}) -// Using Retrieve -msg, err := listener.Retrieve() +msg, err := l.Retrieve() if err != nil { println(err) } -println(string(msg)) - // Using Listen - will block unless sent in goroutine -go listener.Listen(func(msg []byte, wg *sync.WaitGroup) { - println(string(msg)) - wg.Done() +go l.Listen(func(b []byte, fnwg *sync.WaitGroup){ + println(string(b)) + fnwg.Done() }) -producer.Send(new(KinesisMessage).Init([]byte(`{"foo":"bar"}`), "test")) - -listener.Close() -producer.Close() - -// Or with Kinesis Firehose -firehose, err := new(kinetic.Producer).Firehose() -if err != nil { - println(err) -} - -firehose.Send(new(KinesisMessage).Init([]byte(`{"foo":"bar"}`), "test")) - -firehose.Close() +// Send a message using the producer +p.Send(&message.Message{ + Data: []byte(`{"foo":"bar"}`), +}) ``` -For more examples take a look at the tests. API documentation can be found [here](https://godoc.org/github.com/rewardStyle/kinetic). +For more examples take a look at the tests or the test program in the `testexec` directory. API documentation can be found [here](https://godoc.org/github.com/rewardStyle/kinetic). diff --git a/testexec/README.md b/testexec/README.md index bc3cbb3..149627a 100644 --- a/testexec/README.md +++ b/testexec/README.md @@ -3,7 +3,7 @@ testexec is an executable CLI program to test the kinetic producer / lister. The executable program creates a kinetic object connecting to either a local kinesalte instance or to AWS Kinesis using local AWS credentials/config. -A new stream will be created (with a random stream name) unless a stream name is provided (See useage). +A new stream will be created (with a random stream name) unless a stream name is provided (See usage). A kinetic producer will stream dummy data to the kinesis stream and a kinetic listener will stream from the kinesis stream. From a957c86368b754b2935951015eaafc129d933871 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 7 Jun 2017 16:39:49 -0500 Subject: [PATCH 036/109] Added back KCL reader implementation and fixed bugs encountered when interacting with the MultilangDaemon --- listener/config.go | 1 - listener/kclreader.go | 241 ++++++++++++++++++++++++++++++++++++ listener/kclreaderconfig.go | 31 +++++ multilang/action.go | 109 ++++++---------- testexec/cli.go | 2 +- 5 files changed, 314 insertions(+), 70 deletions(-) create mode 100644 listener/kclreader.go create mode 100644 listener/kclreaderconfig.go diff --git a/listener/config.go b/listener/config.go index f8a2cf6..d5c5d55 100644 --- a/listener/config.go +++ b/listener/config.go @@ -4,7 +4,6 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic/config" "github.com/rewardStyle/kinetic/logging" ) diff --git a/listener/kclreader.go b/listener/kclreader.go new file mode 100644 index 0000000..22c7c62 --- /dev/null +++ b/listener/kclreader.go @@ -0,0 +1,241 @@ +package listener + +import ( + "bufio" + "bytes" + "encoding/json" + "errors" + "fmt" + "math" + "os" + "sync" + + "github.com/rewardStyle/kinetic/errs" + "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic/multilang" +) + +type kclReaderOptions struct { + onInitCallbackFn func() error + onCheckpointCallbackFn func() error + onShutdownCallbackFn func() error +} + +// KclReader handles the KCL Multilang Protocol to read records from KCL +type KclReader struct { + *kclReaderOptions + throttleSem chan Empty + listener *Listener + scanner *bufio.Scanner + reader *bufio.Reader + mutex *sync.Mutex + msgBuffer []message.Message + ackPending bool +} + +// NewKclReader creates a new stream reader to read records from KCL +func NewKclReader(fn ...func(*KclReaderConfig)) *KclReader { + config := NewKclReaderConfig() + for _, f := range fn { + f(config) + } + return &KclReader{ + kclReaderOptions: config.kclReaderOptions, + throttleSem: make(chan Empty, 5), + msgBuffer: []message.Message{}, + mutex: &sync.Mutex{}, + } +} + +// AssociateListener associates the KCL stream reader to a listener +func (r *KclReader) AssociateListener(l *Listener) error { + r.mutex.Lock() + defer r.mutex.Unlock() + if r.listener != nil { + return errs.ErrListenerAlreadyAssociated + } + r.listener = l + return nil +} + +// ensureClient will lazily ensure that we are reading from STDIN. +func (r *KclReader) ensureClient() error { + r.mutex.Lock() + defer r.mutex.Unlock() + if r.scanner == nil { + if r.listener == nil { + return errs.ErrNilListener + } + r.scanner = bufio.NewScanner(os.Stdin) + r.reader = bufio.NewReader(os.Stdin) + go func() error { + return r.processAction() + }() + } + return nil +} + +// GetRecord calls processRecords to attempt to put one message from message buffer to the listener's message +// channel +func (r *KclReader) GetRecord() (int, error) { + return r.processRecords(1) +} + +// GetRecords calls processRecords to attempt to put all messages on the message buffer on the listener's +// message channel +func (r *KclReader) GetRecords() (int, error) { + return r.processRecords(-1) +} + +// processRecords is a helper method which loops through the message buffer and puts messages on the listener's +// message channel. After all the messages on the message buffer have been moved to the listener's message +// channel, a message is sent (following the Multilang protocol) to acknowledge that the processRecords message +// has been received / processed +func (r *KclReader) processRecords(numRecords int) (int, error) { + if err := r.ensureClient(); err != nil { + return 0, err + } + + r.mutex.Lock() + defer r.mutex.Unlock() + + // Define the batchSize + batchSize := 0; + if len(r.msgBuffer) > 0 { + if numRecords < 0 { + batchSize = len(r.msgBuffer) + } else { + batchSize = int(math.Min(float64(len(r.msgBuffer)), float64(numRecords))) + } + } + + // Loop through the message buffer and put the correct number of messages on the listener's message channel + for i := 0; i < batchSize; i++ { + r.listener.messages <- &r.msgBuffer[0] + r.msgBuffer = r.msgBuffer[1:] + } + + // Send an acknowledgement that the 'ProcessRecords' message was received/processed + if len(r.msgBuffer) == 0 && r.ackPending { + err := r.sendMessage(multilang.NewStatusMessage(multilang.PROCESSRECORDS)) + if err != nil { + r.listener.LogError(err) + return batchSize, err + } + } + + return batchSize, nil +} + +func (r *KclReader) getAction() (*multilang.ActionMessage, error) { + buffer := &bytes.Buffer{} + for { + line, isPrefix, err := r.reader.ReadLine() + if err != nil { + panic("Unable to read line from stdin " + err.Error()) + } + buffer.Write(line) + if !isPrefix { + break + } + } + + actionMsg := &multilang.ActionMessage{} + err := json.Unmarshal(buffer.Bytes(), actionMsg) + if err != nil { + fmt.Fprintf(os.Stderr, "Could not understand line read from input: %s\n", buffer.String()) + } + return actionMsg, nil +} + +// processAction listens to STDIN and processes action messages based on the Multilang protocol from KCL +func (r *KclReader) processAction() error { + for { + + // Retrieve action message + actionMessage, err := r.getAction() + if err != nil { + return err + } + if actionMessage == nil { + break + } + + switch actionMessage.Action { + case multilang.INITIALIZE: + r.onInit() + r.sendMessage(multilang.NewStatusMessage(multilang.INITIALIZE)) + case multilang.CHECKPOINT: + r.onCheckpoint() + r.sendMessage(multilang.NewStatusMessage(multilang.CHECKPOINT)) + case multilang.SHUTDOWN: + r.onShutdown() + r.sendMessage(multilang.NewStatusMessage(multilang.SHUTDOWN)) + case multilang.PROCESSRECORDS: + go func() error { + r.mutex.Lock() + defer r.mutex.Unlock() + + if r.ackPending { + return errors.New("Received a processRecords action message from KCL " + + "unexpectedly") + } + + for _, msg := range actionMessage.Records { + r.msgBuffer = append(r.msgBuffer, *msg.ToMessage()) + } + r.ackPending = true; + + return nil + }() + default: + } + } + + return nil +} + +func (r *KclReader) sendMessage(msg interface{}) error { + b, err := json.Marshal(msg) + if err != nil { + r.listener.LogError(err) + return err + } + + fmt.Fprintln(os.Stdout, string(b)) + + return nil +} + +func (r *KclReader) onInit() error { + if r.onInitCallbackFn != nil { + err := r.onInitCallbackFn() + if err != nil { + r.listener.LogError(err) + return err + } + } + return nil +} + +func (r *KclReader) onCheckpoint() error { + if r.onCheckpointCallbackFn != nil { + err := r.onCheckpointCallbackFn() + if err != nil { + r.listener.LogError(err) + return err + } + } + return nil +} + +func (r *KclReader) onShutdown() error { + if r.onShutdownCallbackFn != nil { + err := r.onShutdownCallbackFn() + if err != nil { + r.listener.LogError(err) + return err + } + } + return nil +} diff --git a/listener/kclreaderconfig.go b/listener/kclreaderconfig.go new file mode 100644 index 0000000..ef2bdb4 --- /dev/null +++ b/listener/kclreaderconfig.go @@ -0,0 +1,31 @@ +package listener + +// KclReaderConfig is used to configure KclReader +type KclReaderConfig struct { + *kclReaderOptions +} + +// NewKclReaderConfig creates a new instance of KclReaderConfig +func NewKclReaderConfig() *KclReaderConfig { + return &KclReaderConfig{ + kclReaderOptions: &kclReaderOptions{}, + } +} + +// SetOnInitCallbackFn configures a callback function which is run prior to sending a status message +// acknowledging an 'initialize' message was received / processed +func (c *KclReaderConfig) SetOnInitCallbackFn(fn func() error) { + c.onInitCallbackFn = fn +} + +// SetOnCheckpointCallbackFn configures a callback function which is run prior to sending a status message +// acknowledging an 'checkpoint' message was received / processed +func (c *KclReaderConfig) SetOnCheckpointCallbackFn(fn func() error) { + c.onCheckpointCallbackFn = fn +} + +// SetOnShutdownCallbackFn configures a callback function which is run prior to sending a status message +// acknowledging a 'shutdown' message was received / processed +func (c *KclReaderConfig) SetOnShutdownCallbackFn(fn func() error) { + c.onShutdownCallbackFn = fn +} diff --git a/multilang/action.go b/multilang/action.go index b8aea03..386895d 100644 --- a/multilang/action.go +++ b/multilang/action.go @@ -2,8 +2,8 @@ package multilang import ( "encoding/base64" - "encoding/json" "time" + "strconv" "github.com/rewardStyle/kinetic/message" ) @@ -13,92 +13,65 @@ type ActionType string // These are the enumerated KCL Multilang protocol action message types const ( - Initialize ActionType = "initialize" - ProcessRecords ActionType = "processRecords" - Checkpoint ActionType = "checkpoint" - Shutdown ActionType = "shutdown" - Status ActionType = "status" + INITIALIZE ActionType = "initialize" + PROCESSRECORDS ActionType = "processRecords" + RECORD ActionType = "record" + CHECKPOINT ActionType = "checkpoint" + SHUTDOWN ActionType = "shutdown" + STATUS ActionType = "status" ) // ActionMessage is a struct used to marshal / unmarshal KCL Multilang protocol action messages type ActionMessage struct { - Action ActionType `json:"action"` - ShardID string `json:"shardId,omitempty"` - Records []Record `json:"records,omitempty"` - Checkpoint string `json:"checkpoint,omitempty"` - Error string `json:"error,omitempty"` - Reason string `json:"reason,omitempty"` - ResponseFor ActionType `json:"responseFor,omitempty"` + Action ActionType `json:"action"` + ShardID string `json:"shardId,omitempty"` + SequenceNumber string `json:"sequenceNumber,omitempty"` + Records []Record `json:"records,omitempty"` + Checkpoint string `json:"checkpoint,omitempty"` + Error string `json:"error,omitempty"` + Reason string `json:"reason,omitempty"` + ResponseFor ActionType `json:"responseFor,omitempty"` } // Record is a struct used to marshal / unmarshal kinesis records from KCL Multilang protocol type Record struct { - ApproximateArrivalTimestamp time.Time `json:"approximateArrivalTimestamp,omitempty"` - Data []byte `json:"data,omitempty"` - PartitionKey string `json:"partitionKey,omitempty"` - SequenceNumber string `json:"sequenceNumber,omitempty"` + Action ActionType `json:"action"` + ApproximateArrival Timestamp `json:"approximateArrivalTimestamp"` + Data string `json:"data,omitempty"` + PartitionKey string `json:"partitionKey,omitempty"` + SequenceNumber string `json:"sequenceNumber,omitempty"` + SubSequenceNumber int `json:"subSequenceNumber,omitempty"` } -// UnmarshalJSON is used as a custom unmarshaller to unmarshal the KCL Multilang ActionMessage -func (a *ActionMessage) UnmarshalJSON(data []byte) error { - am := &ActionMessage{} - if err := json.Unmarshal(data, am); err != nil { - return err - } - - a.Action = am.Action - a.ShardID = am.ShardID - a.Records = am.Records - if len(am.Checkpoint) > 2 { - a.Checkpoint = am.Checkpoint[1:len(am.Checkpoint)-1] - } - - if len(am.Error) > 2 { - a.Error = am.Error[1:len(am.Error)-1] - } - - if len(am.Reason) > 2 { - a.Reason = am.Reason[1:len(am.Reason)-1] - } - - return nil +// Timestamp is a time.Time type +type Timestamp struct { + time.Time } -// UnmarshalJSON is used as a custom unmarshaller to base64 decode the data field of the KCL Multilang -// processRecord message -func (r *Record) UnmarshalJSON(data []byte) error { - record := &Record{} - if err := json.Unmarshal(data, record); err != nil { +// UnmarshalJSON is used as a custom unmarshaller unmarshal unix time stamps +func (t *Timestamp) UnmarshalJSON(b []byte) error { + ts, err := strconv.Atoi(string(b)) + if err != nil { return err } - r.ApproximateArrivalTimestamp = record.ApproximateArrivalTimestamp - - if len(record.Data) > 2 { - encodedString := string(record.Data[1:len(record.Data)-1]) - decodedMsg, err := base64.StdEncoding.DecodeString(encodedString) - if err != nil { - return err - } - r.Data = decodedMsg - } - - if len(record.PartitionKey) > 2 { - r.PartitionKey = record.PartitionKey[1:len(record.PartitionKey)-1] - } - - if len(record.SequenceNumber) > 2 { - r.SequenceNumber = record.SequenceNumber[1:len(record.SequenceNumber)-1] - } + milliseconds := ts % 1000 + seconds := (ts - milliseconds) / 1000 + t.Time = time.Unix(int64(seconds), int64(milliseconds * 1000)) return nil } // ToMessage is used to transform a multilang.Record struct into a message.Message struct func (r *Record) ToMessage() *message.Message { + b, err := base64.StdEncoding.DecodeString(r.Data) + if err != nil { + panic("There was a problem decoding kcl data") + } + return &message.Message{ - ApproximateArrivalTimestamp: &r.ApproximateArrivalTimestamp, - Data: r.Data, + ApproximateArrivalTimestamp: &r.ApproximateArrival.Time, + Data: b, PartitionKey: &r.PartitionKey, SequenceNumber: &r.SequenceNumber, } @@ -107,7 +80,7 @@ func (r *Record) ToMessage() *message.Message { // NewCheckpointMessage is used to create a new checkpoint message func NewCheckpointMessage(seqNum string) *ActionMessage { return &ActionMessage{ - Action: Checkpoint, + Action: CHECKPOINT, Checkpoint: seqNum, } } @@ -115,7 +88,7 @@ func NewCheckpointMessage(seqNum string) *ActionMessage { // NewStatusMessage is used to create a new status message func NewStatusMessage(actionType ActionType) *ActionMessage { return &ActionMessage{ - Action: Status, + Action: STATUS, ResponseFor: actionType, } -} \ No newline at end of file +} diff --git a/testexec/cli.go b/testexec/cli.go index ef40815..1131507 100644 --- a/testexec/cli.go +++ b/testexec/cli.go @@ -20,7 +20,7 @@ func parseCommandLineArgs() *Config { // Define command line flags locationPtr := flag.String("location", "local", "used to specify the location of the kinesis stream. " + "Accepted values are (local|aws). For local, run kinesalite on http://127.0.0.1:4567. For aws, your " + - "aws credentials and configuration needs to be defined at ~/.aws") + "aws credentials and configuration need to be defined at ~/.aws") streamNamePtr := flag.String("stream-name", "", "used to specify a pre-existing stream to be used for " + "testing. A new stream will be created if not defined.") numMsgsPtr := flag.Int("num-msgs", 0, "used to specify the number of messages to (attempt to) send / " + From 0bd3a3c05a3af9cf873b116b4144afd0211f8c63 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 13 Jun 2017 16:24:34 -0500 Subject: [PATCH 037/109] Producer refactor changes after team code review - make TryToSend() drop message and return error if message channel is full - add exponential back off to sendBatch() retries - move close channel call out of goroutine in CloseWithContext() - fix potential deadlock in produce() --- errs/errors.go | 3 +++ producer/producer.go | 41 +++++++++++++++++++++++++---------------- 2 files changed, 28 insertions(+), 16 deletions(-) diff --git a/errs/errors.go b/errs/errors.go index e17a8e1..c81cd17 100644 --- a/errs/errors.go +++ b/errs/errors.go @@ -80,4 +80,7 @@ var ( // records of the batch to be retried. This failure is considered part // of normal behavior of the Kinesis stream. ErrRetryRecords = errors.New("PutRecords requires retry of some records in batch") + + // ErrDroppedMessage is returned when the message channel is full and messages are being dropped + ErrDroppedMessage = errors.New("Channel is full, dropped message") ) diff --git a/producer/producer.go b/producer/producer.go index 9566202..3906163 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -35,7 +35,7 @@ type producerOptions struct { concurrency int writer StreamWriter - Stats StatsCollector + Stats StatsCollector } // Producer sends records to Kinesis or Firehose. @@ -51,10 +51,10 @@ type Producer struct { shutdownCond *sync.Cond producerWg *sync.WaitGroup - producing bool - producingMu sync.Mutex + producing bool + producingMu sync.Mutex - Session *session.Session + Session *session.Session } // NewProducer creates a new producer for writing records to a Kinesis or @@ -62,7 +62,7 @@ type Producer struct { func NewProducer(fn func(*Config)) (*Producer, error) { config := NewConfig() fn(config) - session, err := config.GetSession() + sess, err := config.GetSession() if err != nil { return nil, err } @@ -70,11 +70,11 @@ func NewProducer(fn func(*Config)) (*Producer, error) { producerOptions: config.producerOptions, LogHelper: &logging.LogHelper{ LogLevel: config.LogLevel, - Logger: session.Config.Logger, + Logger: sess.Config.Logger, }, concurrencySem: make(chan Empty, config.concurrency), pipeOfDeath: make(chan Empty), - Session: session, + Session: sess, } if err := p.writer.AssociateProducer(p); err != nil { return nil, err @@ -90,7 +90,7 @@ func (p *Producer) startProducing() bool { if !p.producing { p.producing = true p.messages = make(chan *message.Message, p.queueDepth) - p.retries = make(chan *message.Message) // TODO: should we use a buffered channel? + p.retries = make(chan *message.Message, p.queueDepth) p.shutdownCond = &sync.Cond{L: &sync.Mutex{}} p.producerWg = new(sync.WaitGroup) p.outstanding = 0 @@ -115,6 +115,7 @@ func (p *Producer) sendBatch(batch []*message.Message) { p.outstanding-- p.shutdownCond.L.Unlock() }() + attempts := 0 var retries []*message.Message var err error @@ -179,6 +180,9 @@ stop: p.Stats.AddDropped(len(batch)) break stop } + + // Apply an exponential back-off before retrying + time.Sleep(time.Duration(attempts * attempts) * time.Second) } // This frees up another sendBatch to run to allow drainage of the // messages / retry queue. This should improve throughput as well as @@ -212,7 +216,7 @@ func (p *Producer) produce() { p.stopProducing() p.producerWg.Done() }() - stop: + for { var batch []*message.Message timer := time.After(p.batchTimeout) @@ -235,7 +239,7 @@ func (p *Producer) produce() { case <-timer: break batch case <-p.pipeOfDeath: - break stop + return } } p.shutdownCond.L.Lock() @@ -256,8 +260,9 @@ func (p *Producer) produce() { // - The retry channel is empty if p.messages == nil && p.outstanding == 0 && len(p.retries) == 0 { close(p.retries) - p.shutdownCond.Signal() - break stop + p.shutdownCond.Broadcast() + p.shutdownCond.L.Unlock() + return } } p.shutdownCond.L.Unlock() @@ -269,8 +274,8 @@ func (p *Producer) produce() { // messages and retries to flush. Cancellation supported through contexts. func (p *Producer) CloseWithContext(ctx context.Context) { c := make(chan Empty, 1) + close(p.messages) go func() { - close(p.messages) p.shutdownCond.L.Lock() for p.outstanding != 0 { p.shutdownCond.Wait() @@ -314,7 +319,11 @@ func (p *Producer) Send(msg *message.Message) error { // capacity for a message, or will immediately return with an error if the // channel is full. func (p *Producer) TryToSend(msg *message.Message) error { - ctx, cancel := context.WithTimeout(context.TODO(), 0*time.Second) - defer cancel() - return p.SendWithContext(ctx, msg) + select { + case p.messages <- msg: + return nil + default: + p.Stats.AddDropped(1) + return errs.ErrDroppedMessage + } } From baebb8a8e0cd4e3b6712961f54b68e982ab49526 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 22 Jun 2017 11:29:10 -0500 Subject: [PATCH 038/109] Refactored kinetic to remove the association between the producer/writer and the listener/reader by requiring the readers (kinesisreader/kclreader) and writers (kinesiswriter/firehosewriter) to have their own StatsCollector and derive a Logger and LogLevel through a provided AwsConfig. Reorganized the fields in all kinetic objects (producer/listener/reader/writer) to pass in all required fields to its respective New* method. --- config/config.go | 7 + kinetic_integration_test.go | 34 +- listener/kclreader.go | 119 +++---- listener/kclreaderconfig.go | 25 +- listener/kinesisreader.go | 316 +++++++----------- listener/kinesisreaderconfig.go | 35 +- listener/listener.go | 199 +++++------ listener/listener_test.go | 56 ++-- listener/{config.go => listenerconfig.go} | 37 +- ...{config_test.go => listenerconfig_test.go} | 29 +- producer/firehosewriter.go | 91 +++-- producer/firehosewriterconfig.go | 34 ++ producer/kinesiswriter.go | 91 +++-- producer/kinesiswriterconfig.go | 34 ++ producer/producer.go | 136 +++----- producer/producer_test.go | 40 ++- producer/{config.go => producerconfig.go} | 48 +-- ...{config_test.go => producerconfig_test.go} | 73 ++-- testexec/main.go | 28 +- 19 files changed, 679 insertions(+), 753 deletions(-) rename listener/{config.go => listenerconfig.go} (59%) rename listener/{config_test.go => listenerconfig_test.go} (80%) create mode 100644 producer/firehosewriterconfig.go create mode 100644 producer/kinesiswriterconfig.go rename producer/{config.go => producerconfig.go} (64%) rename producer/{config_test.go => producerconfig_test.go} (58%) diff --git a/config/config.go b/config/config.go index 6bf3d11..f1c2764 100644 --- a/config/config.go +++ b/config/config.go @@ -25,6 +25,13 @@ func DefaultAwsOptions() *AwsOptions { } } +// NewAwsOptionsFromConfig initializes a new AwsOption from an existing aws Config +func NewAwsOptionsFromConfig(c *aws.Config) *AwsOptions { + return &AwsOptions{ + AwsConfig: c, + } +} + // SetCredentials configures AWS credentials. func (c *AwsOptions) SetCredentials(accessKey, secretKey, sessionToken string) { c.AwsConfig.WithCredentials( diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index 58aa095..927e31b 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -129,24 +129,36 @@ func TestKineticIntegration(t *testing.T) { log.Printf("Stream Name: %s\n", stream) log.Printf("Shard Name: %s\n", shards[0]) + // Create a new kinesis stream writer + w, err := producer.NewKinesisWriter(k.Session.Config, stream, func(kwc *producer.KinesisWriterConfig) { + kwc.SetLogLevel(aws.LogDebug) + }) + if err != nil { + log.Fatalf("Unable to create a new kinesis stream writer due to: %v\n", err) + } + // Create a new kinetic producer - p, err := producer.NewProducer(func(c *producer.Config) { - c.SetAwsConfig(k.Session.Config) - c.SetKinesisStream(stream) + p, err := producer.NewProducer(k.Session.Config, w, func(c *producer.Config) { c.SetBatchSize(5) c.SetBatchTimeout(1000 * time.Millisecond) }) assert.NotNil(t, p) assert.Nil(t, err) + assert.NotNil(t, k.Session) + assert.NotNil(t, k.Session.Config) + r, err := listener.NewKinesisReader(k.Session.Config, stream, shards[0], + func(krc *listener.KinesisReaderConfig) { + krc.SetReadTimeout(1000 * time.Millisecond) + }) + assert.NotNil(t, r) + assert.NoError(t, err) + //assert.NotNil(t, r.Session) + // Create a new kinetic listener - l, err := listener.NewListener(func(c *listener.Config) { - c.SetAwsConfig(k.Session.Config) - c.SetReader(listener.NewKinesisReader(stream, shards[0])) + l, err := listener.NewListener(k.Session.Config, r, func(c *listener.Config) { c.SetQueueDepth(20) c.SetConcurrency(10) - c.SetGetRecordsReadTimeout(1000 * time.Millisecond) - //c.SetLogLevel(aws.LogDebug) }) assert.NotNil(t, l) assert.Nil(t, err) @@ -177,11 +189,11 @@ func TestKineticIntegration(t *testing.T) { // Use the listener to read messages from the kinetic stream go func() { - l.Listen(func(b []byte, fnwg *sync.WaitGroup) { + l.Listen(func(m *message.Message, fnwg *sync.WaitGroup) error { defer fnwg.Done() msg := &Message{} - json.Unmarshal(b, msg) + json.Unmarshal(m.Data, msg) if !streamData.exists(msg.ID) { wg.Done() @@ -190,6 +202,8 @@ func TestKineticIntegration(t *testing.T) { } streamData.put(msg.ID, msg.Message) + + return nil }) }() diff --git a/listener/kclreader.go b/listener/kclreader.go index 22c7c62..bdf6951 100644 --- a/listener/kclreader.go +++ b/listener/kclreader.go @@ -3,102 +3,61 @@ package listener import ( "bufio" "bytes" + "context" "encoding/json" - "errors" "fmt" "math" "os" "sync" - "github.com/rewardStyle/kinetic/errs" "github.com/rewardStyle/kinetic/message" "github.com/rewardStyle/kinetic/multilang" + "github.com/rewardStyle/kinetic/logging" + "github.com/aws/aws-sdk-go/aws" ) type kclReaderOptions struct { onInitCallbackFn func() error onCheckpointCallbackFn func() error onShutdownCallbackFn func() error + Stats StatsCollector + } // KclReader handles the KCL Multilang Protocol to read records from KCL type KclReader struct { *kclReaderOptions - throttleSem chan Empty - listener *Listener + *logging.LogHelper + + throttleSem chan empty + pipeOfDeath chan empty scanner *bufio.Scanner reader *bufio.Reader - mutex *sync.Mutex msgBuffer []message.Message - ackPending bool } // NewKclReader creates a new stream reader to read records from KCL -func NewKclReader(fn ...func(*KclReaderConfig)) *KclReader { - config := NewKclReaderConfig() +func NewKclReader(c *aws.Config, fn ...func(*KclReaderConfig)) (*KclReader, error) { + cfg := NewKclReaderConfig(c) for _, f := range fn { - f(config) + f(cfg) } return &KclReader{ - kclReaderOptions: config.kclReaderOptions, - throttleSem: make(chan Empty, 5), + kclReaderOptions: cfg.kclReaderOptions, + LogHelper: &logging.LogHelper{ + LogLevel: cfg.LogLevel, + Logger: cfg.AwsConfig.Logger, + }, + throttleSem: make(chan empty, 5), msgBuffer: []message.Message{}, - mutex: &sync.Mutex{}, - } -} - -// AssociateListener associates the KCL stream reader to a listener -func (r *KclReader) AssociateListener(l *Listener) error { - r.mutex.Lock() - defer r.mutex.Unlock() - if r.listener != nil { - return errs.ErrListenerAlreadyAssociated - } - r.listener = l - return nil -} - -// ensureClient will lazily ensure that we are reading from STDIN. -func (r *KclReader) ensureClient() error { - r.mutex.Lock() - defer r.mutex.Unlock() - if r.scanner == nil { - if r.listener == nil { - return errs.ErrNilListener - } - r.scanner = bufio.NewScanner(os.Stdin) - r.reader = bufio.NewReader(os.Stdin) - go func() error { - return r.processAction() - }() - } - return nil -} - -// GetRecord calls processRecords to attempt to put one message from message buffer to the listener's message -// channel -func (r *KclReader) GetRecord() (int, error) { - return r.processRecords(1) -} - -// GetRecords calls processRecords to attempt to put all messages on the message buffer on the listener's -// message channel -func (r *KclReader) GetRecords() (int, error) { - return r.processRecords(-1) + }, nil } // processRecords is a helper method which loops through the message buffer and puts messages on the listener's // message channel. After all the messages on the message buffer have been moved to the listener's message // channel, a message is sent (following the Multilang protocol) to acknowledge that the processRecords message // has been received / processed -func (r *KclReader) processRecords(numRecords int) (int, error) { - if err := r.ensureClient(); err != nil { - return 0, err - } - - r.mutex.Lock() - defer r.mutex.Unlock() - +func (r *KclReader) processRecords(fn MessageFn, numRecords int) (int, error) { // Define the batchSize batchSize := 0; if len(r.msgBuffer) > 0 { @@ -110,16 +69,19 @@ func (r *KclReader) processRecords(numRecords int) (int, error) { } // Loop through the message buffer and put the correct number of messages on the listener's message channel + var wg sync.WaitGroup for i := 0; i < batchSize; i++ { - r.listener.messages <- &r.msgBuffer[0] + wg.Add(1) + go fn(&r.msgBuffer[0], &wg) r.msgBuffer = r.msgBuffer[1:] } + wg.Wait() // Send an acknowledgement that the 'ProcessRecords' message was received/processed - if len(r.msgBuffer) == 0 && r.ackPending { + if len(r.msgBuffer) == 0 { err := r.sendMessage(multilang.NewStatusMessage(multilang.PROCESSRECORDS)) if err != nil { - r.listener.LogError(err) + r.LogError(err) return batchSize, err } } @@ -173,18 +135,9 @@ func (r *KclReader) processAction() error { r.sendMessage(multilang.NewStatusMessage(multilang.SHUTDOWN)) case multilang.PROCESSRECORDS: go func() error { - r.mutex.Lock() - defer r.mutex.Unlock() - - if r.ackPending { - return errors.New("Received a processRecords action message from KCL " + - "unexpectedly") - } - for _, msg := range actionMessage.Records { r.msgBuffer = append(r.msgBuffer, *msg.ToMessage()) } - r.ackPending = true; return nil }() @@ -198,7 +151,7 @@ func (r *KclReader) processAction() error { func (r *KclReader) sendMessage(msg interface{}) error { b, err := json.Marshal(msg) if err != nil { - r.listener.LogError(err) + r.LogError(err) return err } @@ -211,7 +164,7 @@ func (r *KclReader) onInit() error { if r.onInitCallbackFn != nil { err := r.onInitCallbackFn() if err != nil { - r.listener.LogError(err) + r.LogError(err) return err } } @@ -222,7 +175,7 @@ func (r *KclReader) onCheckpoint() error { if r.onCheckpointCallbackFn != nil { err := r.onCheckpointCallbackFn() if err != nil { - r.listener.LogError(err) + r.LogError(err) return err } } @@ -233,9 +186,21 @@ func (r *KclReader) onShutdown() error { if r.onShutdownCallbackFn != nil { err := r.onShutdownCallbackFn() if err != nil { - r.listener.LogError(err) + r.LogError(err) return err } } return nil } + +// GetRecord calls processRecords to attempt to put one message from message buffer to the listener's message +// channel +func (r *KclReader) GetRecord(ctx context.Context,fn MessageFn) (int, error) { + return r.processRecords(fn, 1) +} + +// GetRecords calls processRecords to attempt to put all messages on the message buffer on the listener's +// message channel +func (r *KclReader) GetRecords(ctx context.Context,fn MessageFn) (int, error) { + return r.processRecords(fn, -1) +} diff --git a/listener/kclreaderconfig.go b/listener/kclreaderconfig.go index ef2bdb4..66f8fa4 100644 --- a/listener/kclreaderconfig.go +++ b/listener/kclreaderconfig.go @@ -1,14 +1,25 @@ package listener +import ( + "github.com/aws/aws-sdk-go/aws" + "github.com/rewardStyle/kinetic/config" +) + // KclReaderConfig is used to configure KclReader type KclReaderConfig struct { + *config.AwsOptions *kclReaderOptions + LogLevel aws.LogLevelType } // NewKclReaderConfig creates a new instance of KclReaderConfig -func NewKclReaderConfig() *KclReaderConfig { +func NewKclReaderConfig(cfg *aws.Config) *KclReaderConfig { return &KclReaderConfig{ - kclReaderOptions: &kclReaderOptions{}, + AwsOptions: config.NewAwsOptionsFromConfig(cfg), + kclReaderOptions: &kclReaderOptions{ + Stats: &NilStatsCollector{}, + }, + LogLevel: *cfg.LogLevel, } } @@ -29,3 +40,13 @@ func (c *KclReaderConfig) SetOnCheckpointCallbackFn(fn func() error) { func (c *KclReaderConfig) SetOnShutdownCallbackFn(fn func() error) { c.onShutdownCallbackFn = fn } + +// SetStatsCollector configures a listener to handle listener metrics. +func (c *KclReaderConfig) SetStatsCollector(stats StatsCollector) { + c.Stats = stats +} + +// SetLogLevel configures both the SDK and Kinetic log levels. +func (c *KclReaderConfig) SetLogLevel(logLevel aws.LogLevelType) { + c.LogLevel = logLevel & 0xffff0000 +} diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index 4971d8e..108c5b2 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -1,6 +1,7 @@ package listener import ( + "context" "fmt" "io" "sync" @@ -17,92 +18,56 @@ import ( "github.com/rewardStyle/kinetic/message" ) +// kinesisReaderOptions is used to hold all of the configurable settings of a KinesisReader. type kinesisReaderOptions struct { - stream string - shard string - batchSize int shardIterator *ShardIterator + readTimeout time.Duration + Stats StatsCollector } // KinesisReader handles the API to read records from Kinesis. type KinesisReader struct { *kinesisReaderOptions + *logging.LogHelper - throttleSem chan Empty + stream string + shard string + throttleSem chan empty nextShardIterator string - - listener *Listener - client kinesisiface.KinesisAPI - clientMu sync.Mutex + client kinesisiface.KinesisAPI } -// NewKinesisReader creates a new stream reader to read records from Kinesis. -func NewKinesisReader(stream, shard string, fn ...func(*KinesisReaderConfig)) *KinesisReader { - config := NewKinesisReaderConfig(stream, shard) +// NewKinesisReader creates a new KinesisReader object which implements the StreamReader interface to read records from +// Kinesis. +func NewKinesisReader(c *aws.Config, stream string, shard string, fn ...func(*KinesisReaderConfig)) (*KinesisReader, error) { + cfg := NewKinesisReaderConfig(c) for _, f := range fn { - f(config) - } - return &KinesisReader{ - kinesisReaderOptions: config.kinesisReaderOptions, - throttleSem: make(chan Empty, 5), + f(cfg) } -} - -// AssociateListener associates the Kinesis stream reader to a listener. -func (r *KinesisReader) AssociateListener(l *Listener) error { - r.clientMu.Lock() - defer r.clientMu.Unlock() - if r.listener != nil { - return errs.ErrListenerAlreadyAssociated - } - r.listener = l - return nil -} - -// ensureClient will lazily make sure we have an AWS Kinesis client. -func (r *KinesisReader) ensureClient() error { - // From the aws-go-sdk documentation: - // http://docs.aws.amazon.com/sdk-for-go/api/aws/session/ - // - // Concurrency: - // Sessions are safe to use concurrently as long as the Session is not - // being modified. The SDK will not modify the Session once the Session - // has been created. Creating service clients concurrently from a - // shared Session is safe. - // - // We need to think through the impact of creating a new client (for - // example, after receiving an error from Kinesis) while there may be - // outstanding goroutines still processing messages. My cursory thought - // is that this is safe to do, as any outstanding messages will likely - // not interact with the Kinesis stream. At worst, we would need a lock - // around the ensureClient method to make sure that no two goroutines - // are trying to ensure the client at the same time. - // - // As we don't expose any methods (or in fact, even the Listener object - // itself) to the client through the API, I don't foresee needing to add - // this lock unless something dramatically changes about the design of - // this library. - r.clientMu.Lock() - defer r.clientMu.Unlock() - if r.client == nil { - if r.listener == nil { - return errs.ErrNilListener - } - r.client = kinesis.New(r.listener.Session) + sess, err := cfg.GetSession() + if err != nil { + return nil, err } - return nil + return &KinesisReader{ + kinesisReaderOptions: cfg.kinesisReaderOptions, + LogHelper: &logging.LogHelper{ + LogLevel: cfg.LogLevel, + Logger: cfg.AwsConfig.Logger, + }, + stream: stream, + shard: shard, + throttleSem: make(chan empty, 5), + client: kinesis.New(sess), + }, nil } -// ensureShardIterator will lazily make sure that we have a valid ShardIterator, -// calling the GetShardIterator API with the configured ShardIteratorType (with -// any applicable StartingSequenceNumber or Timestamp) if necessary. +// ensureShardIterator will lazily make sure that we have a valid ShardIterator, calling the GetShardIterator API with +// the configured ShardIteratorType (with any applicable StartingSequenceNumber or Timestamp) if necessary. // -// Not thread-safe. Only called from getRecords Care must be taken to ensure -// that only one call to Listen and Retrieve/RetrieveFn can be running at a -// time. +// Not thread-safe. Only called from getRecords Care must be taken to ensure that only one call to Listen and +// Retrieve/RetrieveFn can be running at a time. func (r *KinesisReader) ensureShardIterator() error { - r.ensureClient() if r.nextShardIterator != "" { return nil } @@ -115,7 +80,7 @@ func (r *KinesisReader) ensureShardIterator() error { Timestamp: r.shardIterator.getTimestamp(), }) if err != nil { - r.listener.LogError(err) + r.LogError(err) return err } if resp == nil { @@ -127,12 +92,10 @@ func (r *KinesisReader) ensureShardIterator() error { return r.setNextShardIterator(*resp.ShardIterator) } -// setNextShardIterator sets the nextShardIterator to use when calling -// GetRecords. +// setNextShardIterator sets the nextShardIterator to use when calling GetRecords. // -// Not thread-safe. Only called from getRecords (and ensureShardIterator, which -// is called from getRecords). Care must be taken to ensure that only one call -// to Listen and Retrieve/RetrieveFn can be running at a time. +// Not thread-safe. Only called from getRecords (and ensureShardIterator, which is called from getRecords). Care must +// be taken to ensure that only one call to Listen and Retrieve/RetrieveFn can be running at a time. func (r *KinesisReader) setNextShardIterator(shardIterator string) error { if len(shardIterator) == 0 { return errs.ErrEmptyShardIterator @@ -141,9 +104,8 @@ func (r *KinesisReader) setNextShardIterator(shardIterator string) error { return nil } -// setSequenceNumber sets the sequenceNumber of shardIterator to the last -// delivered message and updates the shardIteratorType to AT_SEQUENCE_NUMBER. -// This is only used when we need to call getShardIterator (say, to refresh the +// setSequenceNumber sets the sequenceNumber of shardIterator to the last delivered message and updates the +// shardIteratorType to AT_SEQUENCE_NUMBER. This is only used when we need to call getShardIterator (say, to refresh the // shard iterator). // // Not thread-safe. Only called from getRecords. Care must be taken to ensure @@ -159,46 +121,28 @@ func (r *KinesisReader) setSequenceNumber(sequenceNumber string) error { // Kinesis allows five read ops per second per shard. // http://docs.aws.amazon.com/kinesis/latest/dev/service-sizes-and-limits.html -func (r *KinesisReader) throttle(sem chan Empty) { - sem <- Empty{} +func (r *KinesisReader) throttle(sem chan empty) { + sem <- empty{} time.AfterFunc(1*time.Second, func() { <-sem }) } -// GetRecord calls getRecords and delivers one record into the messages -// channel. -func (r *KinesisReader) GetRecord() (int, error) { - return r.getRecords(1) -} - -// GetRecords calls getRecords and delivers each record into the messages -// channel. -func (r *KinesisReader) GetRecords() (int, error) { - return r.getRecords(r.batchSize) -} - -// FIXME: Need to investigate that the timeout implementation doesn't result in -// an fd leak. Since we call Read on the HTTPResonse.Body in a select with a -// timeout channel, we do prevent ourself from blocking. Once we timeout, we -// return an error to the outer ioutil.ReadAll, which should result in a call -// to our io.ReadCloser's Close function. This will in turn call Close on the -// underlying HTTPResponse.Body. The question is whether this actually shuts -// down the TCP connection. Worst case scenario is that our client Timeout -// eventually fires and closes the socket, but this can be susceptible to FD -// exhaustion. -func (r *KinesisReader) getRecords(batchSize int) (int, error) { - if err := r.ensureClient(); err != nil { - return 0, err - } +// FIXME: Need to investigate that the timeout implementation doesn't result in an fd leak. Since we call Read on the +// HTTPResonse.Body in a select with a timeout channel, we do prevent ourself from blocking. Once we timeout, we return +// an error to the outer ioutil.ReadAll, which should result in a call to our io.ReadCloser's Close function. This will +// in turn call Close on the underlying HTTPResponse.Body. The question is whether this actually shuts down the TCP +// connection. Worst case scenario is that our client Timeout eventually fires and closes the socket, but this can be +// susceptible to FD exhaustion. +func (r *KinesisReader) getRecords(ctx context.Context, fn MessageFn, batchSize int) (int, error) { if err := r.ensureShardIterator(); err != nil { return 0, err } r.throttle(r.throttleSem) - // We use the GetRecordsRequest method of creating requests to allow for - // registering custom handlers for better control over the API request. + // We use the GetRecordsRequest method of creating requests to allow for registering custom handlers for better + // control over the API request. var startReadTime time.Time var startUnmarshalTime time.Time start := time.Now() @@ -209,55 +153,47 @@ func (r *KinesisReader) getRecords(batchSize int) (int, error) { }) // If debug is turned on, add some handlers for GetRecords logging - if r.listener.LogLevel.AtLeast(logging.LogDebug) { + if r.LogLevel.AtLeast(logging.LogDebug) { req.Handlers.Send.PushBack(func(req *request.Request) { - r.listener.LogDebug("Finished GetRecords Send, took", time.Since(start)) + r.LogDebug("Finished getRecords Send, took", time.Since(start)) }) } - // Here, we insert a handler to be called after the Send handler and - // before the the Unmarshal handler in the aws-go-sdk library. + // Here, we insert a handler to be called after the Send handler and before the the Unmarshal handler in the + // aws-go-sdk library. // - // The Send handler will call http.Client.Do() on the request, which - // blocks until the response headers have been read before returning an - // HTTPResponse. + // The Send handler will call http.Client.Do() on the request, which blocks until the response headers have been + // read before returning an HTTPResponse. // - // The Unmarshal handler will ultimately call ioutil.ReadAll() on the - // HTTPResponse.Body stream. + // The Unmarshal handler will ultimately call ioutil.ReadAll() on the HTTPResponse.Body stream. // - // Our handler wraps the HTTPResponse.Body with our own ReadCloser so - // that we can implement a timeout mechanism on the Read() call (which - // is called by the ioutil.ReadAll() function) + // Our handler wraps the HTTPResponse.Body with our own ReadCloser so that we can implement a timeout mechanism + // on the Read() call (which is called by the ioutil.ReadAll() function) req.Handlers.Unmarshal.PushFront(func(req *request.Request) { - r.listener.LogDebug("Started GetRecords Unmarshal, took", time.Since(start)) - // Here, we set a timer that the initial Read() call on - // HTTPResponse.Body must return by. Note that the normal - // http.Client Timeout is still in effect. + r.LogDebug("Started getRecords Unmarshal, took", time.Since(start)) + // Here, we set a timer that the initial Read() call on HTTPResponse.Body must return by. Note that the + // normal http.Client Timeout is still in effect. startReadTime = time.Now() - timer := time.NewTimer(r.listener.getRecordsReadTimeout) + timer := time.NewTimer(r.readTimeout) req.HTTPResponse.Body = &ReadCloserWrapper{ ReadCloser: req.HTTPResponse.Body, OnReadFn: func(stream io.ReadCloser, b []byte) (n int, err error) { - // The OnReadFn will be called each time - // ioutil.ReadAll calls Read on the + // The OnReadFn will be called each time ioutil.ReadAll calls Read on the // ReadCloserWrapper. - // First, we set up a struct that to hold the - // results of the Read() call that can go + // First, we set up a struct that to hold the results of the Read() call that can go // through a channel type Result struct { n int err error } - // Next, we build a channel with which to pass - // the Read() results + // Next, we build a channel with which to pass the Read() results c := make(chan Result, 1) - // Now, we call the Read() on the - // HTTPResponse.Body in a goroutine and feed the - // results into the channel + // Now, we call the Read() on the HTTPResponse.Body in a goroutine and feed the results + // into the channel readStart := time.Now() go func() { var result Result @@ -265,92 +201,84 @@ func (r *KinesisReader) getRecords(batchSize int) (int, error) { c <- result }() - // Finally, we poll for the Read() to complete - // or the timer to elapse. + // Finally, we poll for the Read() to complete or the timer to elapse. select { case result := <-c: - // If we sucessfully Read() from the - // HTTPResponse.Body, we reset our - // timeout and return the results from - // the Read() - timer.Reset(r.listener.getRecordsReadTimeout) + // If we sucessfully Read() from the HTTPResponse.Body, we reset our timeout and + // return the results from the Read() + timer.Reset(r.readTimeout) n, err = result.n, result.err - r.listener.LogDebug(fmt.Sprintf("GetRecords read %d bytes, took %v", n, time.Since(readStart))) + r.LogDebug(fmt.Sprintf("getRecords read %d bytes, took %v", n, time.Since(readStart))) case <-timer.C: - // If we timeout, we return an error - // that will unblock ioutil.ReadAll(). - // This will cause the Unmarshal handler - // to return an error. This error will - // propogate to the original req.Send() - // call (below) - r.listener.LogDebug(fmt.Sprintf("GetRecords read timed out after %v", time.Since(readStart))) + // If we timeout, we return an error that will unblock ioutil.ReadAll(). This + // will cause the Unmarshal handler to return an error. This error will + // propogate to the original req.Send() call (below) + r.LogDebug(fmt.Sprintf("getRecords read timed out after %v", time.Since(readStart))) err = errs.ErrTimeoutReadResponseBody - case <-r.listener.pipeOfDeath: - // The pipe of death will abort any pending - // reads on a GetRecords call. - r.listener.LogDebug(fmt.Sprintf("GetRecords received pipe of death after %v", time.Since(readStart))) - err = errs.ErrPipeOfDeath + case <-ctx.Done(): + // The pipe of death will abort any pending reads on a GetRecords call. + r.LogDebug(fmt.Sprintf("getRecords received ctx.Done() after %v", time.Since(readStart))) + err = ctx.Err() } return }, OnCloseFn: func() { - r.listener.Stats.AddGetRecordsReadResponseDuration(time.Since(startReadTime)) - r.listener.LogDebug("Finished GetRecords body read, took", time.Since(start)) + r.Stats.AddGetRecordsReadResponseDuration(time.Since(startReadTime)) + r.LogDebug("Finished GetRecords body read, took", time.Since(start)) startUnmarshalTime = time.Now() }, } }) req.Handlers.Unmarshal.PushBack(func(req *request.Request) { - r.listener.Stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) - r.listener.LogDebug("Finished GetRecords Unmarshal, took", time.Since(start)) + r.Stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) + r.LogDebug("Finished GetRecords Unmarshal, took", time.Since(start)) }) // Send the GetRecords request - r.listener.LogDebug("Starting GetRecords Build/Sign request, took", time.Since(start)) - r.listener.Stats.AddGetRecordsCalled(1) + r.LogDebug("Starting GetRecords Build/Sign request, took", time.Since(start)) + r.Stats.AddGetRecordsCalled(1) if err := req.Send(); err != nil { - r.listener.LogError("Error getting records:", err) + r.LogError("Error getting records:", err) switch err.(awserr.Error).Code() { case kinesis.ErrCodeProvisionedThroughputExceededException: - r.listener.Stats.AddProvisionedThroughputExceeded(1) + r.Stats.AddProvisionedThroughputExceeded(1) default: - r.listener.LogDebug("Received AWS error:", err.Error()) + r.LogDebug("Received AWS error:", err.Error()) } return 0, err } - r.listener.Stats.AddGetRecordsDuration(time.Since(start)) + r.Stats.AddGetRecordsDuration(time.Since(start)) // Process Records - r.listener.LogDebug(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), r.shard, time.Since(start))) + r.LogDebug(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), r.shard, time.Since(start))) if resp == nil { return 0, errs.ErrNilGetRecordsResponse } delivered := 0 - r.listener.Stats.AddBatchSize(len(resp.Records)) + r.Stats.AddBatchSize(len(resp.Records)) for _, record := range resp.Records { if record != nil { - // Allow (only) a pipeOfDeath to trigger an instance - // shutdown of the loop to deliver messages. Otherwise, - // a normal cancellation will not prevent getRecords - // from completing the delivery of the current batch of - // records. + // Allow (only) a pipeOfDeath to trigger an instance shutdown of the loop to deliver messages. + // Otherwise, a normal cancellation will not prevent getRecords from completing the delivery of + // the current batch of records. select { - case r.listener.messages <- message.FromRecord(record): + case <-ctx.Done(): + r.LogInfo(fmt.Sprintf("getRecords received ctx.Done() while delivering messages, %d delivered, ~%d dropped", delivered, len(resp.Records)-delivered)) + return delivered, ctx.Err() + default: + var wg sync.WaitGroup + wg.Add(1) + go fn(message.FromRecord(record), &wg) + wg.Wait() delivered++ - r.listener.Stats.AddConsumed(1) + r.Stats.AddConsumed(1) if record.SequenceNumber != nil { - // We can safely ignore if this call returns - // error, as if we somehow receive an empty - // sequence number from AWS, we will simply not - // set it. At worst, this causes us to - // reprocess this record if we happen to refresh - // the iterator. + // We can safely ignore if this call returns error, as if we somehow receive an + // empty sequence number from AWS, we will simply not set it. At worst, this + // causes us to reprocess this record if we happen to refresh the iterator. r.setSequenceNumber(*record.SequenceNumber) } - case <-r.listener.pipeOfDeath: - r.listener.LogInfo(fmt.Sprintf("getRecords received pipe of death while delivering messages, %d delivered, ~%d dropped", delivered, len(resp.Records)-delivered)) - return delivered, errs.ErrPipeOfDeath } } } @@ -358,22 +286,28 @@ func (r *KinesisReader) getRecords(batchSize int) (int, error) { // TODO: According to AWS docs: // http://docs.aws.amazon.com/sdk-for-go/api/service/kinesis/#GetRecordsOutput // - // NextShardIterator: The next position in the shard - // from which to start sequentially reading data - // records. If set to null, the shard has been closed - // and the requested iterator will not return any more - // data. + // NextShardIterator: The next position in the shard from which to start sequentially reading data + // records. If set to null, the shard has been closed and the requested iterator will not return any + // more data. // - // When dealing with streams that will merge or split, - // we need to detect that the shard has closed and + // When dealing with streams that will merge or split, we need to detect that the shard has closed and // notify the client library. // - // TODO: I don't know if we should be ignoring an error returned - // by setShardIterator in case of an empty shard iterator in the - // response. There isn't much we can do, and the best path for - // recovery may be simply to reprocess the batch and see if we - // get a valid NextShardIterator from AWS the next time around. + // TODO: I don't know if we should be ignoring an error returned by setShardIterator in case of an empty + // shard iterator in the response. There isn't much we can do, and the best path for recovery may be + // simply to reprocess the batch and see if we get a valid NextShardIterator from AWS the next time + // around. r.setNextShardIterator(*resp.NextShardIterator) } return delivered, nil } + +// GetRecord calls getRecords and delivers one record into the messages channel. +func (r *KinesisReader) GetRecord(ctx context.Context, fn MessageFn) (int, error) { + return r.getRecords(ctx, fn, 1) +} + +// GetRecords calls getRecords and delivers each record into the messages channel. +func (r *KinesisReader) GetRecords(ctx context.Context, fn MessageFn) (int, error) { + return r.getRecords(ctx, fn, r.batchSize) +} diff --git a/listener/kinesisreaderconfig.go b/listener/kinesisreaderconfig.go index e71e24f..900ae13 100644 --- a/listener/kinesisreaderconfig.go +++ b/listener/kinesisreaderconfig.go @@ -1,19 +1,29 @@ package listener +import ( + "time" + "github.com/aws/aws-sdk-go/aws" + "github.com/rewardStyle/kinetic/config" +) + // KinesisReaderConfig is used to configure a KinesisReader type KinesisReaderConfig struct { + *config.AwsOptions *kinesisReaderOptions + LogLevel aws.LogLevelType } // NewKinesisReaderConfig creates a new instance of KinesisReaderConfig -func NewKinesisReaderConfig(stream, shard string) *KinesisReaderConfig { +func NewKinesisReaderConfig(cfg *aws.Config) *KinesisReaderConfig { return &KinesisReaderConfig{ + AwsOptions: config.NewAwsOptionsFromConfig(cfg), kinesisReaderOptions: &kinesisReaderOptions{ - stream: stream, - shard: shard, batchSize: 10000, shardIterator: NewShardIterator(), + readTimeout: time.Second, + Stats: &NilStatsCollector{}, }, + LogLevel: *cfg.LogLevel, } } @@ -22,8 +32,23 @@ func (c *KinesisReaderConfig) SetBatchSize(batchSize int) { c.batchSize = batchSize } -// SetInitialShardIterator configures the settings used to retrieve initial -// shard iterator via the GetShardIterator call. +// SetInitialShardIterator configures the settings used to retrieve initial shard iterator via the GetShardIterator +// call. func (c *KinesisReaderConfig) SetInitialShardIterator(shardIterator *ShardIterator) { c.shardIterator = shardIterator } + +// SetReadTimeout configures the time to wait for each successive Read operation on the GetRecords response payload. +func (c *KinesisReaderConfig) SetReadTimeout(timeout time.Duration) { + c.readTimeout = timeout +} + +// SetStatsCollector configures a listener to handle listener metrics. +func (c *KinesisReaderConfig) SetStatsCollector(stats StatsCollector) { + c.Stats = stats +} + +// SetLogLevel configures the log levels for the SDK. +func (c *KinesisReaderConfig) SetLogLevel(logLevel aws.LogLevelType) { + c.LogLevel = logLevel & 0xffff0000 +} diff --git a/listener/listener.go b/listener/listener.go index 925a267..abf0c33 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -6,34 +6,30 @@ import ( "sync" "time" - "github.com/aws/aws-sdk-go/aws/session" - + "github.com/aws/aws-sdk-go/aws" "github.com/rewardStyle/kinetic/errs" "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" ) -// StreamReader is an interface that abstracts out a stream reader +// StreamReader is an interface that abstracts out a stream reader. type StreamReader interface { - AssociateListener(listener *Listener) error - GetRecord() (int, error) - GetRecords() (int, error) - ensureClient() error + GetRecord(context.Context, MessageFn) (int, error) + GetRecords(context.Context, MessageFn) (int, error) } -// Empty is used a as a dummy type for counting semaphore channels. -type Empty struct{} +// empty is used a as a dummy type for semaphore channels and the pipe of death channel. +type empty struct{} -// MessageFn defines the signature of a message handler used by Listen and -// RetrieveFn. -type MessageFn func([]byte, *sync.WaitGroup) +// MessageFn defines the signature of a message handler used by Listen, RetrieveFn and their associated WithContext +// functions. +type MessageFn func(*message.Message, *sync.WaitGroup) error +// listenerOptions is used to hold all of the configurable settings of a Listener object. type listenerOptions struct { - queueDepth int - concurrency int - getRecordsReadTimeout time.Duration - reader StreamReader - Stats StatsCollector + queueDepth int + concurrency int + Stats StatsCollector } // Listener polls the StreamReader for messages. @@ -41,42 +37,34 @@ type Listener struct { *listenerOptions *logging.LogHelper + reader StreamReader messages chan *message.Message - concurrencySem chan Empty - pipeOfDeath chan Empty - + concurrencySem chan empty + pipeOfDeath chan empty consuming bool consumingMu sync.Mutex - - Session *session.Session } -// NewListener creates a new listener for listening to message on a StreamReader. -func NewListener(fn func(*Config)) (*Listener, error) { - config := NewConfig() - fn(config) - session, err := config.GetSession() - if err != nil { - return nil, err +// NewListener creates a new Listener object for retrieving and listening to message(s) on a StreamReader. +func NewListener(c *aws.Config, r StreamReader, fn ...func(*Config)) (*Listener, error) { + cfg := NewConfig(c) + for _, f := range fn { + f(cfg) } - l := &Listener{ - listenerOptions: config.listenerOptions, + return &Listener{ + listenerOptions: cfg.listenerOptions, LogHelper: &logging.LogHelper{ - LogLevel: config.LogLevel, - Logger: session.Config.Logger, + LogLevel: cfg.LogLevel, + Logger: cfg.AwsConfig.Logger, }, - concurrencySem: make(chan Empty, config.concurrency), - pipeOfDeath: make(chan Empty), - Session: session, - } - if err := l.reader.AssociateListener(l); err != nil { - return nil, err - } - return l, nil + reader: r, + concurrencySem: make(chan empty, cfg.concurrency), + pipeOfDeath: make(chan empty), + }, nil } -// startConsuming will initialize the consumer and set consuming to true if -// there is not already another consume loop running. +// startConsuming will initialize the message channel and set consuming to true if there is not already another consume +// loop running. func (l *Listener) startConsuming() bool { l.consumingMu.Lock() defer l.consumingMu.Unlock() @@ -88,8 +76,8 @@ func (l *Listener) startConsuming() bool { return false } -// shouldConsume is a convenience function that allows functions to break their -// loops if the context receives a cancellation. +// shouldConsume is a convenience function that allows functions to break their loops if the context receives a +// cancellation or a pipe of death. func (l *Listener) shouldConsume(ctx context.Context) (bool, error) { select { case <-l.pipeOfDeath: @@ -101,36 +89,37 @@ func (l *Listener) shouldConsume(ctx context.Context) (bool, error) { } } -// stopConsuming handles any cleanup after a consuming has stopped. +// stopConsuming handles any cleanup after consuming has stopped. func (l *Listener) stopConsuming() { l.consumingMu.Lock() defer l.consumingMu.Unlock() - if l.messages != nil { + if l.consuming && l.messages != nil { close(l.messages) } l.consuming = false } -// RetrieveWithContext waits for a message from the stream and return the value. -// Cancellation supported through contexts. +// RetrieveWithContext waits for a message from the stream and returns the message. Cancellation is supported through +// contexts. func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, error) { if !l.startConsuming() { return nil, errs.ErrAlreadyConsuming } - defer func() { - l.stopConsuming() - }() + defer l.stopConsuming() for { - // A cancellation or closing the pipe of death will cause - // Retrieve (and related functions) to abort in between - // getRecord calls. Note, that this would only occur when there - // are no new records to retrieve. Otherwise, getRecords will - // be allowed to run to completion and deliver one record. + // A cancellation or closing the pipe of death will cause Retrieve (and related functions) to abort in + // between getRecord calls. Note, that this would only occur when there are no new records to retrieve. + // Otherwise, getRecords will be allowed to run to completion and deliver one record. ok, err := l.shouldConsume(ctx) if !ok { return nil, err } - n, err := l.reader.GetRecord() + n, err := l.reader.GetRecord(context.TODO(), func(msg *message.Message, wg *sync.WaitGroup) error { + defer wg.Done() + l.messages <- msg + + return nil + }) if err != nil { return nil, err } @@ -141,66 +130,64 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, e } } -// Retrieve waits for a message from the stream and return the value. +// Retrieve waits for a message from the stream and returns the value func (l *Listener) Retrieve() (*message.Message, error) { return l.RetrieveWithContext(context.TODO()) } -// RetrieveFnWithContext retrieves a message from the stream and dispatches it -// to the supplied function. RetrieveFn will wait until the function completes. -// Cancellation supported through context. +// RetrieveFnWithContext retrieves a message from the stream and dispatches it to the supplied function. RetrieveFn +// will wait until the function completes. Cancellation is supported through context. func (l *Listener) RetrieveFnWithContext(ctx context.Context, fn MessageFn) error { msg, err := l.RetrieveWithContext(ctx) if err != nil { return err } + var wg sync.WaitGroup + defer wg.Wait() wg.Add(1) - start := time.Now() - go fn(msg.Data, &wg) - wg.Wait() - l.Stats.AddProcessedDuration(time.Since(start)) - l.Stats.AddProcessed(1) + go func() { + start := time.Now() + fn(msg, &wg) + l.Stats.AddProcessedDuration(time.Since(start)) + l.Stats.AddProcessed(1) + }() + return nil } -// RetrieveFn retrieves a message from the stream and dispatches it to the -// supplied function. RetrieveFn will wait until the function completes. +// RetrieveFn retrieves a message from the stream and dispatches it to the supplied function. RetrieveFn will wait +// until the function completes. func (l *Listener) RetrieveFn(fn MessageFn) error { return l.RetrieveFnWithContext(context.TODO(), fn) } -// consume calls getRecords with configured batch size in a loop until the -// listener is stopped. +// consume calls getRecords with configured batch size in a loop until the listener is stopped. func (l *Listener) consume(ctx context.Context) { - // We need to run startConsuming to make sure that we are okay and ready - // to start consuming. This is mainly to avoid a race condition where - // Listen() will attempt to read the messages channel prior to consume() - // initializing it. We can then launch a goroutine to handle the actual - // consume operation. - if !l.startConsuming() { - return - } + // We need to run startConsuming to make sure that we are okay and ready to start consuming. This is mainly to + // avoid a race condition where Listen() will attempt to read the messages channel prior to consume() + // initializing it. We can then launch a goroutine to handle the actual consume operation. + if !l.startConsuming() { return } go func() { - defer func() { - l.stopConsuming() - }() - stop: + defer l.stopConsuming() + for { - // The consume loop can be cancelled by a calling the - // cancellation function on the context or by closing - // the pipe of death. Note that in the case of context - // cancellation, the getRecords call below will be - // allowed to complete (as getRecords does not regard - // context cancellation). In the case of cancellation - // by pipe of death, however, the getRecords will - // immediately abort and allow the consume function to - // immediately abort as well. - ok, _ := l.shouldConsume(ctx) - if !ok { - break stop + // The consume loop can be cancelled by a calling the cancellation function on the context or by + // closing the pipe of death. Note that in the case of context cancellation, the getRecords + // call below will be allowed to complete (as getRecords does not regard context cancellation). + // In the case of cancellation by pipe of death, however, the getRecords will immediately abort + // and allow the consume function to immediately abort as well. + if ok, _ := l.shouldConsume(ctx); !ok { + return } - _, err := l.reader.GetRecords() + + _, err := l.reader.GetRecords(context.TODO(), + func(msg *message.Message, wg *sync.WaitGroup) error { + defer wg.Done() + l.messages <- msg + + return nil + }) if err != nil { switch err := err.(type) { case net.Error: @@ -226,27 +213,23 @@ func (l *Listener) consume(ctx context.Context) { }() } -// ListenWithContext listens and delivers message to the supplied function. -// Upon cancellation, Listen will stop the consumer loop and wait until the -// messages channel is closed and all messages are delivered. +// ListenWithContext listens and delivers message to the supplied function. Upon cancellation, Listen will stop the +// consumer loop and wait until the messages channel is closed and all messages are delivered. func (l *Listener) ListenWithContext(ctx context.Context, fn MessageFn) { l.consume(ctx) var wg sync.WaitGroup defer wg.Wait() -stop: for { select { case msg, ok := <-l.messages: if !ok { - break stop + return } l.Stats.AddDelivered(1) - // For simplicity, did not do the pipe of death here. - // If POD is received, we may deliver a couple more - // messages (especially since select is random in which - // channel is read from). - l.concurrencySem <- Empty{} + // For simplicity, did not do the pipe of death here. If POD is received, we may deliver a + // couple more messages (especially since select is random in which channel is read from). + l.concurrencySem <- empty{} wg.Add(1) go func(msg *message.Message) { defer func() { @@ -255,7 +238,7 @@ stop: var fnWg sync.WaitGroup fnWg.Add(1) start := time.Now() - fn(msg.Data, &fnWg) + fn(msg, &fnWg) fnWg.Wait() l.Stats.AddProcessedDuration(time.Since(start)) l.Stats.AddProcessed(1) @@ -263,7 +246,7 @@ stop: }(msg) case <-l.pipeOfDeath: l.LogInfo("ListenWithContext received pipe of death") - break stop + return } } } diff --git a/listener/listener_test.go b/listener/listener_test.go index bd81d2b..c4301e6 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -18,10 +18,10 @@ import ( "github.com/rewardStyle/kinetic" "github.com/rewardStyle/kinetic/errs" + "github.com/rewardStyle/kinetic/message" ) func putRecord(l *Listener, b []byte) (*string, error) { - l.reader.ensureClient() resp, err := l.reader.(*KinesisReader).client.PutRecord(&kinesis.PutRecordInput{ Data: b, PartitionKey: aws.String("dummy"), @@ -46,16 +46,24 @@ func TestListener(t *testing.T) { err = k.CreateStream(stream, 1) So(err, ShouldBeNil) - err = k.WaitUntilStreamExists(context.TODO(), stream, request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + err = k.WaitUntilStreamExists(context.TODO(), stream, + request.WithWaiterDelay(request.ConstantWaiterDelay(time.Second))) So(err, ShouldBeNil) shards, err := k.GetShards(stream) So(err, ShouldBeNil) So(len(shards), ShouldEqual, 1) - l, err := NewListener(func(c *Config) { - c.SetAwsConfig(k.Session.Config) - c.SetReader(NewKinesisReader(stream, shards[0])) + So(k.Session, ShouldNotBeNil) + So(k.Session.Config, ShouldNotBeNil) + r, err := NewKinesisReader(k.Session.Config, stream, shards[0], func(krc *KinesisReaderConfig) { + krc.SetReadTimeout(time.Second) + }) + So(r, ShouldNotBeNil) + So(err, ShouldBeNil) + //So(r.Session, ShouldNotBeNil) + + l, err := NewListener(k.Session.Config, r, func(c *Config) { c.SetQueueDepth(10) c.SetConcurrency(10) }) @@ -69,17 +77,12 @@ func TestListener(t *testing.T) { So(r.stream, ShouldEqual, stream) }) - Convey("check that the reader has a valid reference to the listener", func() { - So(r.listener, ShouldEqual, l) - }) + //Convey("check that the kinesis reader was initialized with an AWS session", func() { + // So(r.Session, ShouldNotBeNil) + //}) - Convey("check that calling ensureClient twice doesn't overwrite existing client", func() { - So(r.client, ShouldBeNil) - r.ensureClient() + Convey("check that the kinesis client was initialized correctly", func() { So(r.client, ShouldNotBeNil) - client := r.client - r.ensureClient() - So(r.client, ShouldEqual, client) }) }) @@ -174,7 +177,7 @@ func TestListener(t *testing.T) { var wg sync.WaitGroup wg.Add(1) go func() { - ctx, cancel := context.WithTimeout(context.TODO(), 1*time.Second) + ctx, cancel := context.WithTimeout(context.TODO(), 1000*time.Millisecond) defer cancel() _, err := l.RetrieveWithContext(ctx) c.So(err, ShouldNotBeNil) @@ -187,12 +190,17 @@ func TestListener(t *testing.T) { wg.Wait() }) + // TODO: Move this test to kinesisreader_test.go Convey("check that throttle mechanism prevents more than 5 calls to get records", func() { start := time.Now() secs := []float64{} for i := 1; i <= 6; i++ { start := time.Now() - l.reader.GetRecord() + l.reader.GetRecord(context.TODO(), func(msg *message.Message, wg *sync.WaitGroup) error { + defer wg.Done() + + return nil + }) secs = append(secs, time.Since(start).Seconds()) } elapsed := time.Since(start).Seconds() @@ -205,12 +213,14 @@ func TestListener(t *testing.T) { data := "retrieved" _, err := putRecord(l, []byte(data)) So(err, ShouldBeNil) - err = l.RetrieveFn(func(b []byte, wg *sync.WaitGroup) { + err = l.RetrieveFn(func(msg *message.Message, wg *sync.WaitGroup) error { called = true // Note that because this is called in a goroutine, we have to use // the goconvey context - c.So(string(b), ShouldEqual, data) + c.So(string(msg.Data), ShouldEqual, data) wg.Done() + + return nil }) So(err, ShouldBeNil) So(called, ShouldBeTrue) @@ -223,9 +233,11 @@ func TestListener(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - l.Listen(func(b []byte, wg *sync.WaitGroup) { + l.Listen(func(msg *message.Message, wg *sync.WaitGroup) error { defer wg.Done() atomic.AddInt64(&count, 1) + + return nil }) }() for _, planet := range planets { @@ -269,16 +281,18 @@ func TestListener(t *testing.T) { go func() { ctx, cancel := context.WithCancel(context.TODO()) defer wg.Done() - l.ListenWithContext(ctx, func(b []byte, wg *sync.WaitGroup) { + l.ListenWithContext(ctx, func(m *message.Message, wg *sync.WaitGroup) error { defer wg.Done() time.AfterFunc(time.Duration(rand.Intn(10))*time.Second, func() { - n, err := strconv.Atoi(string(b)) + n, err := strconv.Atoi(string(m.Data)) c.So(err, ShouldBeNil) atomic.AddInt64(&count, 1) if n == 15 { cancel() } }) + + return nil }) }() wg.Wait() diff --git a/listener/config.go b/listener/listenerconfig.go similarity index 59% rename from listener/config.go rename to listener/listenerconfig.go index d5c5d55..52e207a 100644 --- a/listener/config.go +++ b/listener/listenerconfig.go @@ -1,8 +1,6 @@ package listener import ( - "time" - "github.com/aws/aws-sdk-go/aws" "github.com/rewardStyle/kinetic/config" "github.com/rewardStyle/kinetic/logging" @@ -16,25 +14,18 @@ type Config struct { } // NewConfig creates a new instance of Config -func NewConfig() *Config { +func NewConfig(cfg *aws.Config) *Config { return &Config{ - AwsOptions: config.DefaultAwsOptions(), + AwsOptions: config.NewAwsOptionsFromConfig(cfg), listenerOptions: &listenerOptions{ - queueDepth: 10000, - concurrency: 10000, - getRecordsReadTimeout: 1 * time.Second, - Stats: &NilStatsCollector{}, + queueDepth: 10000, + concurrency: 10000, + Stats: &NilStatsCollector{}, }, LogLevel: logging.LogOff, } } -// SetAwsConfig configures the AWS Config used to create Sessions (and therefore -// kinesis clients). -func (c *Config) SetAwsConfig(config *aws.Config) { - c.AwsConfig = config -} - // SetQueueDepth controls the depth of the listener queue func (c *Config) SetQueueDepth(queueDepth int) { c.queueDepth = queueDepth @@ -46,24 +37,14 @@ func (c *Config) SetConcurrency(concurrency int) { c.concurrency = concurrency } -// SetGetRecordsReadTimeout configures the time to wait for each successive -// Read operation on the GetRecords response payload. -func (c *Config) SetGetRecordsReadTimeout(timeout time.Duration) { - c.getRecordsReadTimeout = timeout -} - -// SetReader sets the listener's stream reader -func (c *Config) SetReader(reader StreamReader) { - c.reader = reader +// SetStatsCollector configures a listener to handle listener metrics. +func (c *Config) SetStatsCollector(stats StatsCollector) { + c.Stats = stats } // SetLogLevel configures both the SDK and Kinetic log levels. func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { + // TODO: Do we want to change the AWS log level? c.AwsOptions.SetLogLevel(logLevel) c.LogLevel = logLevel & 0xffff0000 } - -// SetStatsCollector configures a listener to handle listener metrics. -func (c *Config) SetStatsCollector(stats StatsCollector) { - c.Stats = stats -} diff --git a/listener/config_test.go b/listener/listenerconfig_test.go similarity index 80% rename from listener/config_test.go rename to listener/listenerconfig_test.go index 91bb2b2..0d55cd0 100644 --- a/listener/config_test.go +++ b/listener/listenerconfig_test.go @@ -8,7 +8,6 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/session" - "github.com/rewardStyle/kinetic" "github.com/rewardStyle/kinetic/logging" ) @@ -37,11 +36,16 @@ func getSession(config *Config) *session.Session { func TestNewConfig(t *testing.T) { Convey("given a new listener config", t, func() { - config := NewConfig() + k, err := kinetic.New(func(c *kinetic.Config) { + c.SetEndpoint("bogus-endpoint") + }) + So(k, ShouldNotBeNil) + So(err, ShouldBeNil) + config := NewConfig(k.Session.Config) Convey("check the default values for its non-zero config", func() { + So(config.queueDepth, ShouldEqual, 10000) So(config.concurrency, ShouldEqual, 10000) - So(config.getRecordsReadTimeout, ShouldEqual, 1*time.Second) So(config.Stats, ShouldHaveSameTypeAs, &NilStatsCollector{}) So(config.LogLevel.Value(), ShouldEqual, logging.LogOff) }) @@ -59,25 +63,16 @@ func TestNewConfig(t *testing.T) { So(config.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) }) - Convey("check that we can set the AWS configuration", func() { - k, err := kinetic.New(func(c *kinetic.Config) { - c.SetEndpoint("bogus-endpoint") - }) - So(err, ShouldBeNil) - config.SetAwsConfig(k.Session.Config) - sess := getSession(config) - So(aws.StringValue(sess.Config.Endpoint), ShouldEqual, "bogus-endpoint") - }) - Convey("check that we can set the concurrency limit", func() { config.SetConcurrency(50) So(config.concurrency, ShouldEqual, 50) }) - Convey("check that we can set the read timeout for the GetRecords request", func() { - config.SetGetRecordsReadTimeout(10 * time.Second) - So(config.getRecordsReadTimeout, ShouldEqual, 10*time.Second) - }) + //TODO: Move this test to kinesisreaderconfig_test.go + //Convey("check that we can set the read timeout for the GetRecords request", func() { + // config.SetReadTimeout(10 * time.Second) + // So(config.readTimeout, ShouldEqual, 10*time.Second) + //}) Convey("check that we can configure a stats collector", func() { config.SetStatsCollector(&DebugStatsCollector{}) diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go index 9f5b0ab..66e2332 100644 --- a/producer/firehosewriter.go +++ b/producer/firehosewriter.go @@ -2,7 +2,6 @@ package producer import ( "fmt" - "sync" "time" "github.com/aws/aws-sdk-go/aws" @@ -11,56 +10,46 @@ import ( "github.com/aws/aws-sdk-go/service/firehose/firehoseiface" "github.com/rewardStyle/kinetic/errs" + "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" ) +type firehoseWriterOptions struct { + Stats StatsCollector +} + // FirehoseWriter handles the API to send records to Kinesis. type FirehoseWriter struct { - stream string + *firehoseWriterOptions + *logging.LogHelper - producer *Producer - client firehoseiface.FirehoseAPI - clientMu sync.Mutex + stream string + client firehoseiface.FirehoseAPI } // NewFirehoseWriter creates a new stream writer to write records to a Kinesis. -func NewFirehoseWriter(stream string) *FirehoseWriter { - return &FirehoseWriter{ - stream: stream, - } -} - -// ensureClient will lazily make sure we have an AWS Kinesis client. -func (w *FirehoseWriter) ensureClient() error { - w.clientMu.Lock() - defer w.clientMu.Unlock() - if w.client == nil { - if w.producer == nil { - return errs.ErrNilProducer - } - w.client = firehose.New(w.producer.Session) +func NewFirehoseWriter(c *aws.Config, stream string, fn ...func(*FirehoseWriterConfig)) (*FirehoseWriter, error) { + cfg := NewFirehoseWriterConfig(c) + for _, f := range fn { + f(cfg) } - return nil -} - -// AssociateProducer associates the Firehose stream writer to a producer. -func (w *FirehoseWriter) AssociateProducer(p *Producer) error { - w.clientMu.Lock() - defer w.clientMu.Unlock() - if w.producer != nil { - return errs.ErrProducerAlreadyAssociated + sess, err := cfg.GetSession() + if err != nil { + return nil, err } - w.producer = p - return nil + return &FirehoseWriter{ + firehoseWriterOptions: cfg.firehoseWriterOptions, + LogHelper: &logging.LogHelper{ + LogLevel: cfg.LogLevel, + Logger: cfg.AwsConfig.Logger, + }, + stream: stream, + client: firehose.New(sess), + }, nil } -// PutRecords sends a batch of records to Firehose and returns a list of records -// that need to be retried. +// PutRecords sends a batch of records to Firehose and returns a list of records that need to be retried. func (w *FirehoseWriter) PutRecords(messages []*message.Message) ([]*message.Message, error) { - if err := w.ensureClient(); err != nil { - return nil, err - } - var startSendTime time.Time var startBuildTime time.Time @@ -76,31 +65,31 @@ func (w *FirehoseWriter) PutRecords(messages []*message.Message) ([]*message.Mes req.Handlers.Build.PushFront(func(r *request.Request) { startBuildTime = time.Now() - w.producer.LogDebug("Start PutRecords Build, took", time.Since(start)) + w.LogDebug("Start PutRecords Build, took", time.Since(start)) }) req.Handlers.Build.PushBack(func(r *request.Request) { - w.producer.Stats.AddPutRecordsBuildDuration(time.Since(startBuildTime)) - w.producer.LogDebug("Finished PutRecords Build, took", time.Since(start)) + w.Stats.AddPutRecordsBuildDuration(time.Since(startBuildTime)) + w.LogDebug("Finished PutRecords Build, took", time.Since(start)) }) req.Handlers.Send.PushFront(func(r *request.Request) { startSendTime = time.Now() - w.producer.LogDebug("Start PutRecords Send took", time.Since(start)) + w.LogDebug("Start PutRecords Send took", time.Since(start)) }) - req.Handlers.Build.PushBack(func(r *request.Request) { - w.producer.Stats.AddPutRecordsSendDuration(time.Since(startSendTime)) - w.producer.LogDebug("Finished PutRecords Send, took", time.Since(start)) + req.Handlers.Send.PushBack(func(r *request.Request) { + w.Stats.AddPutRecordsSendDuration(time.Since(startSendTime)) + w.LogDebug("Finished PutRecords Send, took", time.Since(start)) }) - w.producer.LogDebug("Starting PutRecords Build/Sign request, took", time.Since(start)) - w.producer.Stats.AddPutRecordsCalled(1) + w.LogDebug("Starting PutRecords Build/Sign request, took", time.Since(start)) + w.Stats.AddPutRecordsCalled(1) if err := req.Send(); err != nil { - w.producer.LogError("Error putting records:", err.Error()) + w.LogError("Error putting records:", err.Error()) return nil, err } - w.producer.Stats.AddPutRecordsDuration(time.Since(start)) + w.Stats.AddPutRecordsDuration(time.Since(start)) if resp == nil { return nil, errs.ErrNilPutRecordsResponse @@ -111,7 +100,7 @@ func (w *FirehoseWriter) PutRecords(messages []*message.Message) ([]*message.Mes attempted := len(messages) failed := int(aws.Int64Value(resp.FailedPutCount)) sent := attempted - failed - w.producer.LogDebug(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) + w.LogDebug(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) var retries []*message.Message var err error @@ -122,9 +111,9 @@ func (w *FirehoseWriter) PutRecords(messages []*message.Message) ([]*message.Mes } else { switch aws.StringValue(record.ErrorCode) { case firehose.ErrCodeLimitExceededException: - w.producer.Stats.AddProvisionedThroughputExceeded(1) + w.Stats.AddProvisionedThroughputExceeded(1) default: - w.producer.LogDebug("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(record.ErrorMessage)) + w.LogDebug("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(record.ErrorMessage)) } messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage diff --git a/producer/firehosewriterconfig.go b/producer/firehosewriterconfig.go new file mode 100644 index 0000000..9bb4870 --- /dev/null +++ b/producer/firehosewriterconfig.go @@ -0,0 +1,34 @@ +package producer + +import ( + "github.com/aws/aws-sdk-go/aws" + "github.com/rewardStyle/kinetic/config" +) + +// FirehoseWriterConfig is used to configure FirehoseWriter +type FirehoseWriterConfig struct { + *config.AwsOptions + *firehoseWriterOptions + LogLevel aws.LogLevelType +} + +// NewFirehoseWriterConfig creates a new instance of FirehoseWriterConfig +func NewFirehoseWriterConfig(cfg *aws.Config) *FirehoseWriterConfig { + return &FirehoseWriterConfig{ + AwsOptions: config.NewAwsOptionsFromConfig(cfg), + firehoseWriterOptions: &firehoseWriterOptions{ + Stats: &NilStatsCollector{}, + }, + LogLevel: *cfg.LogLevel, + } +} + +// SetStatsCollector configures a listener to handle listener metrics. +func (c *FirehoseWriterConfig) SetStatsCollector(stats StatsCollector) { + c.Stats = stats +} + +// SetLogLevel configures the log levels for the SDK. +func (c *FirehoseWriterConfig) SetLogLevel(logLevel aws.LogLevelType) { + c.LogLevel = logLevel & 0xffff0000 +} diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index b16aae5..8ef5eb5 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -2,7 +2,6 @@ package producer import ( "fmt" - "sync" "time" "github.com/aws/aws-sdk-go/aws" @@ -11,56 +10,46 @@ import ( "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" "github.com/rewardStyle/kinetic/errs" + "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" ) +type kinesisWriterOptions struct { + Stats StatsCollector +} + // KinesisWriter handles the API to send records to Kinesis. type KinesisWriter struct { - stream string + *kinesisWriterOptions + *logging.LogHelper - producer *Producer - client kinesisiface.KinesisAPI - clientMu sync.Mutex + stream string + client kinesisiface.KinesisAPI } // NewKinesisWriter creates a new stream writer to write records to a Kinesis. -func NewKinesisWriter(stream string) *KinesisWriter { - return &KinesisWriter{ - stream: stream, - } -} - -// ensureClient will lazily make sure we have an AWS Kinesis client. -func (w *KinesisWriter) ensureClient() error { - w.clientMu.Lock() - defer w.clientMu.Unlock() - if w.client == nil { - if w.producer == nil { - return errs.ErrNilProducer - } - w.client = kinesis.New(w.producer.Session) +func NewKinesisWriter(c *aws.Config, stream string, fn ...func(*KinesisWriterConfig)) (*KinesisWriter, error) { + cfg := NewKinesisWriterConfig(c) + for _, f := range fn { + f(cfg) } - return nil -} - -// AssociateProducer associates the Kinesis stream writer to a producer. -func (w *KinesisWriter) AssociateProducer(p *Producer) error { - w.clientMu.Lock() - defer w.clientMu.Unlock() - if w.producer != nil { - return errs.ErrProducerAlreadyAssociated + sess, err := cfg.GetSession() + if err != nil { + return nil, err } - w.producer = p - return nil + return &KinesisWriter{ + kinesisWriterOptions: cfg.kinesisWriterOptions, + LogHelper: &logging.LogHelper{ + LogLevel: cfg.LogLevel, + Logger: cfg.AwsConfig.Logger, + }, + stream: stream, + client: kinesis.New(sess), + }, nil } -// PutRecords sends a batch of records to Kinesis and returns a list of records -// that need to be retried. +// PutRecords sends a batch of records to Kinesis and returns a list of records that need to be retried. func (w *KinesisWriter) PutRecords(messages []*message.Message) ([]*message.Message, error) { - if err := w.ensureClient(); err != nil { - return nil, err - } - var startSendTime time.Time var startBuildTime time.Time @@ -76,31 +65,31 @@ func (w *KinesisWriter) PutRecords(messages []*message.Message) ([]*message.Mess req.Handlers.Build.PushFront(func(r *request.Request) { startBuildTime = time.Now() - w.producer.LogDebug("Start PutRecords Build, took", time.Since(start)) + w.LogDebug("Start PutRecords Build, took", time.Since(start)) }) req.Handlers.Build.PushBack(func(r *request.Request) { - w.producer.Stats.AddPutRecordsBuildDuration(time.Since(startBuildTime)) - w.producer.LogDebug("Finished PutRecords Build, took", time.Since(start)) + w.Stats.AddPutRecordsBuildDuration(time.Since(startBuildTime)) + w.LogDebug("Finished PutRecords Build, took", time.Since(start)) }) req.Handlers.Send.PushFront(func(r *request.Request) { startSendTime = time.Now() - w.producer.LogDebug("Start PutRecords Send took", time.Since(start)) + w.LogDebug("Start PutRecords Send took", time.Since(start)) }) - req.Handlers.Build.PushBack(func(r *request.Request) { - w.producer.Stats.AddPutRecordsSendDuration(time.Since(startSendTime)) - w.producer.LogDebug("Finished PutRecords Send, took", time.Since(start)) + req.Handlers.Send.PushBack(func(r *request.Request) { + w.Stats.AddPutRecordsSendDuration(time.Since(startSendTime)) + w.LogDebug("Finished PutRecords Send, took", time.Since(start)) }) - w.producer.LogDebug("Starting PutRecords Build/Sign request, took", time.Since(start)) - w.producer.Stats.AddPutRecordsCalled(1) + w.LogDebug("Starting PutRecords Build/Sign request, took", time.Since(start)) + w.Stats.AddPutRecordsCalled(1) if err := req.Send(); err != nil { - w.producer.LogError("Error putting records:", err.Error()) + w.LogError("Error putting records:", err.Error()) return nil, err } - w.producer.Stats.AddPutRecordsDuration(time.Since(start)) + w.Stats.AddPutRecordsDuration(time.Since(start)) if resp == nil { return nil, errs.ErrNilPutRecordsResponse @@ -111,7 +100,7 @@ func (w *KinesisWriter) PutRecords(messages []*message.Message) ([]*message.Mess attempted := len(messages) failed := int(aws.Int64Value(resp.FailedRecordCount)) sent := attempted - failed - w.producer.LogDebug(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) + w.LogDebug(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) var retries []*message.Message var err error @@ -123,9 +112,9 @@ func (w *KinesisWriter) PutRecords(messages []*message.Message) ([]*message.Mess } else { switch aws.StringValue(record.ErrorCode) { case kinesis.ErrCodeProvisionedThroughputExceededException: - w.producer.Stats.AddProvisionedThroughputExceeded(1) + w.Stats.AddProvisionedThroughputExceeded(1) default: - w.producer.LogDebug("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(record.ErrorMessage)) + w.LogDebug("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(record.ErrorMessage)) } messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage diff --git a/producer/kinesiswriterconfig.go b/producer/kinesiswriterconfig.go new file mode 100644 index 0000000..f619351 --- /dev/null +++ b/producer/kinesiswriterconfig.go @@ -0,0 +1,34 @@ +package producer + +import ( + "github.com/aws/aws-sdk-go/aws" + "github.com/rewardStyle/kinetic/config" +) + +// KinesisWriterConfig is used to configure KinesisWriter +type KinesisWriterConfig struct { + *config.AwsOptions + *kinesisWriterOptions + LogLevel aws.LogLevelType +} + +// NewKinesisWriterConfig creates a new instance of KinesisWriterConfig +func NewKinesisWriterConfig(cfg *aws.Config) *KinesisWriterConfig { + return &KinesisWriterConfig{ + AwsOptions: config.NewAwsOptionsFromConfig(cfg), + kinesisWriterOptions: &kinesisWriterOptions{ + Stats: &NilStatsCollector{}, + }, + LogLevel: *cfg.LogLevel, + } +} + +// SetStatsCollector configures a listener to handle listener metrics. +func (c *KinesisWriterConfig) SetStatsCollector(stats StatsCollector) { + c.Stats = stats +} + +// SetLogLevel configures the log levels for the SDK. +func (c *KinesisWriterConfig) SetLogLevel(logLevel aws.LogLevelType) { + c.LogLevel = logLevel & 0xffff0000 +} diff --git a/producer/producer.go b/producer/producer.go index 3906163..7d37144 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -7,8 +7,8 @@ import ( "sync" "time" + "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/awserr" - "github.com/aws/aws-sdk-go/aws/session" "github.com/aws/aws-sdk-go/service/kinesis" "github.com/rewardStyle/kinetic/errs" @@ -16,16 +16,13 @@ import ( "github.com/rewardStyle/kinetic/message" ) -// StreamWriter is an interface that abstracts the differences in API between -// Kinesis and Firehose. +// StreamWriter is an interface that abstracts the differences in API between Kinesis and Firehose. type StreamWriter interface { - AssociateProducer(producer *Producer) error PutRecords(message []*message.Message) ([]*message.Message, error) - ensureClient() error } -// Empty is used a as a dummy type for counting semaphore channels. -type Empty struct{} +// empty is used a as a dummy type for counting semaphore channels. +type empty struct{} type producerOptions struct { batchSize int @@ -33,8 +30,6 @@ type producerOptions struct { queueDepth int maxRetryAttempts int concurrency int - writer StreamWriter - Stats StatsCollector } @@ -43,47 +38,38 @@ type Producer struct { *producerOptions *logging.LogHelper + writer StreamWriter messages chan *message.Message retries chan *message.Message - concurrencySem chan Empty - pipeOfDeath chan Empty + concurrencySem chan empty + pipeOfDeath chan empty outstanding int shutdownCond *sync.Cond producerWg *sync.WaitGroup - producing bool producingMu sync.Mutex - - Session *session.Session } -// NewProducer creates a new producer for writing records to a Kinesis or -// Firehose stream. -func NewProducer(fn func(*Config)) (*Producer, error) { - config := NewConfig() - fn(config) - sess, err := config.GetSession() - if err != nil { - return nil, err +// NewProducer creates a new producer for writing records to a Kinesis or Firehose stream. +func NewProducer(c *aws.Config, w StreamWriter, fn ...func(*Config)) (*Producer, error) { + cfg := NewConfig(c) + for _, f := range fn { + f(cfg) } - p := &Producer{ - producerOptions: config.producerOptions, + return &Producer{ + producerOptions: cfg.producerOptions, LogHelper: &logging.LogHelper{ - LogLevel: config.LogLevel, - Logger: sess.Config.Logger, + LogLevel: cfg.LogLevel, + Logger: cfg.AwsConfig.Logger, }, - concurrencySem: make(chan Empty, config.concurrency), - pipeOfDeath: make(chan Empty), - Session: sess, - } - if err := p.writer.AssociateProducer(p); err != nil { - return nil, err - } - return p, nil + writer: w, + concurrencySem: make(chan empty, cfg.concurrency), + pipeOfDeath: make(chan empty), + }, nil } -// startConsuming will initialize the producer and set producing to true if -// there is not already another consume loop running. +// startConsuming will initialize the producer and set producing to true if there is not already another consume loop +// running. func (p *Producer) startProducing() bool { p.producingMu.Lock() defer p.producingMu.Unlock() @@ -139,20 +125,13 @@ stop: case awserr.Error: switch err.Code() { case kinesis.ErrCodeProvisionedThroughputExceededException: - // FIXME: It is not clear to me whether - // PutRecords would ever return a - // ProvisionedThroughputExceeded error. It - // seems that it would instead return a valid - // response in which some or all the records - // within the response will contain an error - // code and error message of - // ProvisionedThroughputExceeded. The current - // assumption is that if we receive an - // ProvisionedThroughputExceeded error, that the - // entire batch should be retried. Note we only - // increment the PutRecord stat, instead of the - // per-message stat. Furthermore, we do not - // increment the FailCount of the messages (as + // FIXME: It is not clear to me whether PutRecords would ever return a + // ProvisionedThroughputExceeded error. It seems that it would instead return a valid + // response in which some or all the records within the response will contain an error + // code and error message of ProvisionedThroughputExceeded. The current assumption is + // that if we receive an ProvisionedThroughputExceeded error, that the entire batch + // should be retried. Note we only increment the PutRecord stat, instead of the per- + // message stat. Furthermore, we do not increment the FailCount of the messages (as // the retry mechanism is different). p.Stats.AddPutRecordsProvisionedThroughputExceeded(1) default: @@ -168,11 +147,9 @@ stop: default: p.LogError("Received unknown error:", err.Error()) } - // NOTE: We may want to go through and increment the FailCount - // for each of the records and allow the batch to be retried - // rather than retrying the batch as-is. With this approach, we - // can kill the "stop" for loop, and set the entire batch to - // retries to allow the below code to handle retrying the + // NOTE: We may want to go through and increment the FailCount for each of the records and allow the + // batch to be retried rather than retrying the batch as-is. With this approach, we can kill the "stop" + // for loop, and set the entire batch to retries to allow the below code to handle retrying the // messages. attempts++ if attempts > p.maxRetryAttempts { @@ -184,11 +161,9 @@ stop: // Apply an exponential back-off before retrying time.Sleep(time.Duration(attempts * attempts) * time.Second) } - // This frees up another sendBatch to run to allow drainage of the - // messages / retry queue. This should improve throughput as well as - // prevent a potential deadlock in which all batches are blocked on - // sending retries to the retries channel, and thus no batches are - // allowed to drain the retry channel. + // This frees up another sendBatch to run to allow drainage of the messages / retry queue. This should improve + // throughput as well as prevent a potential deadlock in which all batches are blocked on sending retries to the + // retries channel, and thus no batches are allowed to drain the retry channel. <-p.concurrencySem for _, msg := range retries { if msg.FailCount < p.maxRetryAttempts { @@ -204,8 +179,8 @@ stop: } } -// produce calls the underlying writer's PutRecords implementation to deliver -// batches of messages to the target stream until the producer is stopped. +// produce calls the underlying writer's PutRecords implementation to deliver batches of messages to the target stream +// until the producer is stopped. func (p *Producer) produce() { if !p.startProducing() { return @@ -246,17 +221,14 @@ func (p *Producer) produce() { if len(batch) > 0 { p.outstanding++ p.Stats.AddBatchSize(len(batch)) - p.concurrencySem <- Empty{} + p.concurrencySem <- empty{} go p.sendBatch(batch) } else if len(batch) == 0 { - // We did not get any records -- check if we may - // be (gracefully) shutting down the producer. - // We can exit when: - // - The messages channel is nil and no new messages - // can be enqueued - // - There are no outstanding sendBatch goroutines - // and can therefore not produce any more messages - // to retry + // We did not get any records -- check if we may be (gracefully) shutting down the + // producer. We can exit when: + // - The messages channel is nil and no new messages can be enqueued + // - There are no outstanding sendBatch goroutines and can therefore not produce any + // more messages to retry // - The retry channel is empty if p.messages == nil && p.outstanding == 0 && len(p.retries) == 0 { close(p.retries) @@ -270,10 +242,10 @@ func (p *Producer) produce() { }() } -// CloseWithContext shuts down the producer, waiting for all outstanding -// messages and retries to flush. Cancellation supported through contexts. +// CloseWithContext shuts down the producer, waiting for all outstanding messages and retries to flush. Cancellation +// is supported through contexts. func (p *Producer) CloseWithContext(ctx context.Context) { - c := make(chan Empty, 1) + c := make(chan empty, 1) close(p.messages) go func() { p.shutdownCond.L.Lock() @@ -282,7 +254,7 @@ func (p *Producer) CloseWithContext(ctx context.Context) { } p.shutdownCond.L.Unlock() p.producerWg.Wait() - c <- Empty{} + c <- empty{} }() select { case <-c: @@ -291,14 +263,12 @@ func (p *Producer) CloseWithContext(ctx context.Context) { } } -// Close shuts down the producer, waiting for all outstanding messages and retries -// to flush. +// Close shuts down the producer, waiting for all outstanding messages and retries to flush. func (p *Producer) Close() { p.CloseWithContext(context.TODO()) } -// SendWithContext sends a message to the stream. Cancellation supported -// through contexts. +// SendWithContext sends a message to the stream. Cancellation supported through contexts. func (p *Producer) SendWithContext(ctx context.Context, msg *message.Message) error { p.produce() select { @@ -309,15 +279,13 @@ func (p *Producer) SendWithContext(ctx context.Context, msg *message.Message) er } } -// Send a message to the stream, waiting on the message to be put into the -// channel. +// Send a message to the stream, waiting on the message to be put into the channel. func (p *Producer) Send(msg *message.Message) error { return p.SendWithContext(context.TODO(), msg) } -// TryToSend will attempt to send a message to the stream if the channel has -// capacity for a message, or will immediately return with an error if the -// channel is full. +// TryToSend will attempt to send a message to the stream if the channel has capacity for a message, or will immediately +// return with an error if the channel is full. func (p *Producer) TryToSend(msg *message.Message) error { select { case p.messages <- msg: diff --git a/producer/producer_test.go b/producer/producer_test.go index e80e598..7b2a1f2 100644 --- a/producer/producer_test.go +++ b/producer/producer_test.go @@ -28,28 +28,41 @@ func TestProducer(t *testing.T) { err = k.CreateStream(stream, 1) So(err, ShouldBeNil) - err = k.WaitUntilStreamExists(context.TODO(), stream, request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + err = k.WaitUntilStreamExists(context.TODO(), stream, + request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) So(err, ShouldBeNil) shards, err := k.GetShards(stream) So(err, ShouldBeNil) So(len(shards), ShouldEqual, 1) - p, err := NewProducer(func(c *Config) { - c.SetAwsConfig(k.Session.Config) - c.SetKinesisStream(stream) + So(k.Session, ShouldNotBeNil) + So(k.Session.Config, ShouldNotBeNil) + w, err := NewKinesisWriter(k.Session.Config, stream) + So(w, ShouldNotBeNil) + So(err, ShouldBeNil) + + p, err := NewProducer(k.Session.Config, w, func(c *Config) { c.SetBatchSize(5) c.SetBatchTimeout(1000 * time.Millisecond) + c.SetConcurrency(10) + c.SetQueueDepth(10) }) So(p, ShouldNotBeNil) So(err, ShouldBeNil) - l, err := listener.NewListener(func(c *listener.Config) { - c.SetAwsConfig(k.Session.Config) - c.SetReader(listener.NewKinesisReader(stream, shards[0])) + So(k.Session, ShouldNotBeNil) + So(k.Session.Config, ShouldNotBeNil) + r, err := listener.NewKinesisReader(k.Session.Config, stream, shards[0], + func(krc *listener.KinesisReaderConfig) { + krc.SetReadTimeout(1000 * time.Millisecond) + }) + So(r, ShouldNotBeNil) + So(err, ShouldBeNil) + + l, err := listener.NewListener(k.Session.Config, r, func(c *listener.Config) { c.SetQueueDepth(10) c.SetConcurrency(10) - c.SetGetRecordsReadTimeout(100 * time.Millisecond) }) So(l, ShouldNotBeNil) So(err, ShouldBeNil) @@ -61,17 +74,8 @@ func TestProducer(t *testing.T) { So(w.stream, ShouldEqual, stream) }) - Convey("check that the writer has a valid reference to the producer", func() { - So(w.producer, ShouldEqual, p) - }) - - Convey("check that calling ensureClient twice doesn't overwrite existing client", func() { - So(w.client, ShouldBeNil) - w.ensureClient() + Convey("check that the writer was initialized correctly", func() { So(w.client, ShouldNotBeNil) - client := w.client - w.ensureClient() - So(w.client, ShouldEqual, client) }) }) diff --git a/producer/config.go b/producer/producerconfig.go similarity index 64% rename from producer/config.go rename to producer/producerconfig.go index 107b16d..71c9ef0 100644 --- a/producer/config.go +++ b/producer/producerconfig.go @@ -4,9 +4,7 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic/config" - "github.com/rewardStyle/kinetic/logging" ) // Config is used to configure a Producer instance. @@ -17,74 +15,54 @@ type Config struct { } // NewConfig creates a new instance of Config. -func NewConfig() *Config { +func NewConfig(cfg *aws.Config) *Config { return &Config{ - AwsOptions: config.DefaultAwsOptions(), + AwsOptions: config.NewAwsOptionsFromConfig(cfg), producerOptions: &producerOptions{ batchSize: 500, - batchTimeout: 1 * time.Second, + batchTimeout: time.Second, queueDepth: 500, maxRetryAttempts: 10, concurrency: 1, Stats: &NilStatsCollector{}, }, - LogLevel: logging.LogOff, + LogLevel: *cfg.LogLevel, } } -// SetAwsConfig configures the AWS Config used to create Sessions (and therefore -// kinesis clients). -func (c *Config) SetAwsConfig(config *aws.Config) { - c.AwsConfig = config -} - -// SetBatchSize configures the batch size to flush pending records to the -// PutRecords call. +// SetBatchSize configures the batch size to flush pending records to the PutRecords call. func (c *Config) SetBatchSize(batchSize int) { c.batchSize = batchSize } -// SetBatchTimeout configures the timeout to flush pending records to the -// PutRecords call. +// SetBatchTimeout configures the timeout to flush pending records to the PutRecords call. func (c *Config) SetBatchTimeout(timeout time.Duration) { c.batchTimeout = timeout } -// SetQueueDepth controls the number of messages that can be in the channel -// to be processed by produce at a given time. +// SetQueueDepth controls the number of messages that can be in the channel to be processed by produce at a given time. func (c *Config) SetQueueDepth(queueDepth int) { c.queueDepth = queueDepth } -// SetMaxRetryAttempts controls the number of times a message can be retried -// before it is discarded. +// SetMaxRetryAttempts controls the number of times a message can be retried before it is discarded. func (c *Config) SetMaxRetryAttempts(attempts int) { c.maxRetryAttempts = attempts } -// SetConcurrency controls the number of outstanding PutRecords calls may be -// active at a time. +// SetConcurrency controls the number of outstanding PutRecords calls may be active at a time. func (c *Config) SetConcurrency(concurrency int) { c.concurrency = concurrency } -// SetKinesisStream sets the producer to write to the given Kinesis stream. -func (c *Config) SetKinesisStream(stream string) { - c.writer = NewKinesisWriter(stream) -} - -// SetFirehoseStream sets the producer to write to the given Firehose stream. -func (c *Config) SetFirehoseStream(stream string) { - c.writer = NewFirehoseWriter(stream) +// SetStatsCollector configures a listener to handle producer metrics. +func (c *Config) SetStatsCollector(stats StatsCollector) { + c.Stats = stats } // SetLogLevel configures both the SDK and Kinetic log levels. func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { + // TODO: Do we want to change the AWS log level? c.AwsOptions.SetLogLevel(logLevel) c.LogLevel = logLevel & 0xffff0000 } - -// SetStatsCollector configures a listener to handle producer metrics. -func (c *Config) SetStatsCollector(stats StatsCollector) { - c.Stats = stats -} diff --git a/producer/config_test.go b/producer/producerconfig_test.go similarity index 58% rename from producer/config_test.go rename to producer/producerconfig_test.go index 7878bc8..5174ded 100644 --- a/producer/config_test.go +++ b/producer/producerconfig_test.go @@ -41,82 +41,65 @@ func (w *DebugStreamWriter) PutRecords(batch []*message.Message) ([]*message.Mes return nil, nil } -func (w *DebugStreamWriter) AssociateProducer(producer *Producer) error { - return nil -} - func TestNewConfig(t *testing.T) { Convey("given a new producer config", t, func() { - config := NewConfig() + k, err := kinetic.New(func(c *kinetic.Config) { + c.SetEndpoint("bogus-endpoint") + }) + So(err, ShouldBeNil) + cfg := NewConfig(k.Session.Config) Convey("check the default values for its non-zero config", func() { - So(config.batchSize, ShouldEqual, 500) - So(config.batchTimeout, ShouldEqual, 1*time.Second) - So(config.queueDepth, ShouldEqual, 500) - So(config.maxRetryAttempts, ShouldEqual, 10) - So(config.concurrency, ShouldEqual, 1) - So(config.Stats, ShouldHaveSameTypeAs, &NilStatsCollector{}) - So(config.LogLevel.Value(), ShouldEqual, logging.LogOff) + So(cfg.batchSize, ShouldEqual, 500) + So(cfg.batchTimeout, ShouldEqual, 1*time.Second) + So(cfg.queueDepth, ShouldEqual, 500) + So(cfg.maxRetryAttempts, ShouldEqual, 10) + So(cfg.concurrency, ShouldEqual, 1) + So(cfg.Stats, ShouldHaveSameTypeAs, &NilStatsCollector{}) + So(cfg.LogLevel.Value(), ShouldEqual, logging.LogOff) }) Convey("check that we can retrieve an aws.Session from it ", func() { - getSession(config) + getSession(cfg) }) Convey("check that we can set both the sdk and kinetic log level", func() { ll := aws.LogDebug | aws.LogDebugWithSigning | logging.LogDebug - config.SetLogLevel(ll) - sess := getSession(config) + cfg.SetLogLevel(ll) + sess := getSession(cfg) So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) - So(config.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) - }) - - Convey("check that we can set the AWS configuration", func() { - k, err := kinetic.New(func(c *kinetic.Config) { - c.SetEndpoint("bogus-endpoint") - }) - So(err, ShouldBeNil) - config.SetAwsConfig(k.Session.Config) - sess := getSession(config) - So(aws.StringValue(sess.Config.Endpoint), ShouldEqual, "bogus-endpoint") + So(cfg.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) }) Convey("check that we can set the batch size", func() { - config.SetBatchSize(1000) - So(config.batchSize, ShouldEqual, 1000) + cfg.SetBatchSize(1000) + So(cfg.batchSize, ShouldEqual, 1000) }) Convey("check that we can set the batch timeout", func() { - config.SetBatchTimeout(10 * time.Second) - So(config.batchTimeout, ShouldEqual, 10*time.Second) + cfg.SetBatchTimeout(10 * time.Second) + So(cfg.batchTimeout, ShouldEqual, 10*time.Second) }) Convey("check that we can set the queue depth", func() { - config.SetQueueDepth(1000) - So(config.queueDepth, ShouldEqual, 1000) + cfg.SetQueueDepth(1000) + So(cfg.queueDepth, ShouldEqual, 1000) }) Convey("check that we can set the max retries", func() { - config.SetMaxRetryAttempts(100) - So(config.maxRetryAttempts, ShouldEqual, 100) + cfg.SetMaxRetryAttempts(100) + So(cfg.maxRetryAttempts, ShouldEqual, 100) }) Convey("check that we can set the concurrency limit", func() { - config.SetConcurrency(50) - So(config.concurrency, ShouldEqual, 50) + cfg.SetConcurrency(50) + So(cfg.concurrency, ShouldEqual, 50) }) Convey("check that we can configure a stats collector", func() { - config.SetStatsCollector(&DebugStatsCollector{}) - So(config.Stats, ShouldHaveSameTypeAs, &DebugStatsCollector{}) - }) - - Convey("check that we can configure a kinesis stream", func() { - stream := "some-stream" - config.SetKinesisStream(stream) - So(config.writer, ShouldHaveSameTypeAs, &KinesisWriter{}) - So(config.writer.(*KinesisWriter).stream, ShouldEqual, stream) + cfg.SetStatsCollector(&DebugStatsCollector{}) + So(cfg.Stats, ShouldHaveSameTypeAs, &DebugStatsCollector{}) }) }) } diff --git a/testexec/main.go b/testexec/main.go index 4381fe2..e1158d3 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -196,10 +196,10 @@ func main() { // Use the listener to read messages from the kinetic stream go func(sd *StreamData) { - l.Listen(func(b []byte, fnwg *sync.WaitGroup) { + l.Listen(func(m *message.Message, fnwg *sync.WaitGroup) error { // Unmarshal data msg := &Message{} - json.Unmarshal(b, msg) + json.Unmarshal(m.Data, msg) // Only mark "done" if the message isn't a duplicate if !sd.exists(msg.ID) { @@ -213,6 +213,8 @@ func main() { // Record message regardless if it is a duplicate sd.mark(msg.ID) fnwg.Done() + + return nil }) }(streamData) @@ -321,9 +323,14 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce log.Println("Creating a kinetic producer ...") } - p, err := producer.NewProducer(func(c *producer.Config) { - c.SetAwsConfig(k.Session.Config) - c.SetKinesisStream(streamName) + w, err := producer.NewKinesisWriter(k.Session.Config, streamName, func(kwc *producer.KinesisWriterConfig) { + kwc.SetLogLevel(aws.LogDebug) + }) + if err != nil { + log.Fatalf("Unable to create a new kinesis stream writer due to: %v\n", err) + } + + p, err := producer.NewProducer(k.Session.Config, w, func(c *producer.Config) { c.SetBatchSize(5) c.SetBatchTimeout(1000 * time.Millisecond) }) @@ -345,13 +352,14 @@ func newKineticListener(k *kinetic.Kinetic, streamName string) *listener.Listene log.Fatalf("Unable to get shards for stream %s due to: %v\n", streamName, err) } - l, err := listener.NewListener(func(c *listener.Config) { - c.SetAwsConfig(k.Session.Config) - c.SetReader(listener.NewKinesisReader(streamName, shards[0])) + r, err := listener.NewKinesisReader(k.Session.Config, streamName, shards[0], + func(krc *listener.KinesisReaderConfig) { + krc.SetReadTimeout(1000 * time.Millisecond) + }) + + l, err := listener.NewListener(k.Session.Config, r, func(c *listener.Config) { c.SetQueueDepth(20) c.SetConcurrency(10) - c.SetGetRecordsReadTimeout(1000 * time.Millisecond) - //c.SetLogLevel(aws.LogDebug) }) if err != nil { log.Fatalf("Unable to create a new listener due to: %v\n", err) From 251fc65135dcd98175c91aebc5ced2d18f3320ad Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 26 Jun 2017 18:03:26 -0500 Subject: [PATCH 039/109] Modified the signature of PutRecords to accept a context variable and a lambda function for processing failed messages. Refactored the firehosewriter and streamwriter to implement the new stream writer interface. Added additional StatsCollector methods for AddDroppedCapacity and AddDroppedRetry to differentiate between dropped messages due to the channel being full versus because the number of retry attempts was exceeded. Renamed Make*() functions to To*() in message.go. --- listener/listener.go | 13 +++++--- listener/listener_test.go | 26 ++++++++++++---- message/message.go | 10 +++--- producer/firehosewriter.go | 30 ++++++++++-------- producer/kinesiswriter.go | 30 ++++++++++-------- producer/producer.go | 54 +++++++++++++++++++++------------ producer/producerconfig_test.go | 4 ++- producer/stats.go | 53 ++++++++++++++++---------------- 8 files changed, 133 insertions(+), 87 deletions(-) diff --git a/listener/listener.go b/listener/listener.go index abf0c33..8f7f024 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -21,10 +21,15 @@ type StreamReader interface { // empty is used a as a dummy type for semaphore channels and the pipe of death channel. type empty struct{} -// MessageFn defines the signature of a message handler used by Listen, RetrieveFn and their associated WithContext -// functions. +// MessageFn defines the signature of a message handler used by Listen, RetrieveFn and their associated *WithContext +// functions. MessageFn accepts a WaitGroup so the function can be run as a blocking operation as opposed to +// MessageFnAsync. type MessageFn func(*message.Message, *sync.WaitGroup) error +// MessageFnAsync defines the signature of a message handler used by Listen, RetrieveFn and their associated +// *WithContext functions. MessageFnAsync is meant to be run asynchronously. +type MessageFnAsync func(*message.Message) error + // listenerOptions is used to hold all of the configurable settings of a Listener object. type listenerOptions struct { queueDepth int @@ -223,9 +228,7 @@ func (l *Listener) ListenWithContext(ctx context.Context, fn MessageFn) { for { select { case msg, ok := <-l.messages: - if !ok { - return - } + if !ok { return } l.Stats.AddDelivered(1) // For simplicity, did not do the pipe of death here. If POD is received, we may deliver a // couple more messages (especially since select is random in which channel is read from). diff --git a/listener/listener_test.go b/listener/listener_test.go index c4301e6..eb12c70 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -61,7 +61,6 @@ func TestListener(t *testing.T) { }) So(r, ShouldNotBeNil) So(err, ShouldBeNil) - //So(r.Session, ShouldNotBeNil) l, err := NewListener(k.Session.Config, r, func(c *Config) { c.SetQueueDepth(10) @@ -77,10 +76,6 @@ func TestListener(t *testing.T) { So(r.stream, ShouldEqual, stream) }) - //Convey("check that the kinesis reader was initialized with an AWS session", func() { - // So(r.Session, ShouldNotBeNil) - //}) - Convey("check that the kinesis client was initialized correctly", func() { So(r.client, ShouldNotBeNil) }) @@ -190,6 +185,24 @@ func TestListener(t *testing.T) { wg.Wait() }) + Convey("check that listen and retrieve can not be called concurrently", func(c C) { + var wg sync.WaitGroup + wg.Add(1) + go func() { + ctx, cancel := context.WithTimeout(context.TODO(), 1000*time.Millisecond) + defer cancel() + l.ListenWithContext(ctx, func(msg *message.Message, wg *sync.WaitGroup) error { + defer wg.Done() + return nil + }) + wg.Done() + }() + <-time.After(10 * time.Millisecond) + _, err := l.Retrieve() + So(err, ShouldEqual, errs.ErrAlreadyConsuming) + wg.Wait() + }) + // TODO: Move this test to kinesisreader_test.go Convey("check that throttle mechanism prevents more than 5 calls to get records", func() { start := time.Now() @@ -214,11 +227,12 @@ func TestListener(t *testing.T) { _, err := putRecord(l, []byte(data)) So(err, ShouldBeNil) err = l.RetrieveFn(func(msg *message.Message, wg *sync.WaitGroup) error { + defer wg.Done() + called = true // Note that because this is called in a goroutine, we have to use // the goconvey context c.So(string(msg.Data), ShouldEqual, data) - wg.Done() return nil }) diff --git a/message/message.go b/message/message.go index 79691bb..2c51ba6 100644 --- a/message/message.go +++ b/message/message.go @@ -39,9 +39,8 @@ func FromRecord(record *kinesis.Record) *Message { } } -// MakeRequestEntry creates a kinesis.PutRecordsRequestEntry to be used in the -// kinesis.PutRecords API call. -func (m *Message) MakeRequestEntry() *kinesis.PutRecordsRequestEntry { +// ToRequestEntry creates a kinesis.PutRecordsRequestEntry to be used in the kinesis.PutRecords API call. +func (m *Message) ToRequestEntry() *kinesis.PutRecordsRequestEntry { return &kinesis.PutRecordsRequestEntry{ Data: m.Data, ExplicitHashKey: m.ExplicitHashKey, @@ -49,9 +48,8 @@ func (m *Message) MakeRequestEntry() *kinesis.PutRecordsRequestEntry { } } -// MakeFirehoseRecord creates a firehose.Record to be used in the -// firehose.PutRecordBatch API call. -func (m *Message) MakeFirehoseRecord() *firehose.Record { +// ToFirehoseRecord creates a firehose.Record to be used in the firehose.PutRecordBatch API call. +func (m *Message) ToFirehoseRecord() *firehose.Record { return &firehose.Record{ Data: m.Data, } diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go index 66e2332..9a9618f 100644 --- a/producer/firehosewriter.go +++ b/producer/firehosewriter.go @@ -1,8 +1,10 @@ package producer import ( + "context" "fmt" "time" + "sync" "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" @@ -49,14 +51,14 @@ func NewFirehoseWriter(c *aws.Config, stream string, fn ...func(*FirehoseWriterC } // PutRecords sends a batch of records to Firehose and returns a list of records that need to be retried. -func (w *FirehoseWriter) PutRecords(messages []*message.Message) ([]*message.Message, error) { +func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageFn) error { var startSendTime time.Time var startBuildTime time.Time start := time.Now() var records []*firehose.Record for _, msg := range messages { - records = append(records, msg.MakeFirehoseRecord()) + records = append(records, msg.ToFirehoseRecord()) } req, resp := w.client.PutRecordBatchRequest(&firehose.PutRecordBatchInput{ DeliveryStreamName: aws.String(w.stream), @@ -87,28 +89,30 @@ func (w *FirehoseWriter) PutRecords(messages []*message.Message) ([]*message.Mes w.Stats.AddPutRecordsCalled(1) if err := req.Send(); err != nil { w.LogError("Error putting records:", err.Error()) - return nil, err + return err } w.Stats.AddPutRecordsDuration(time.Since(start)) if resp == nil { - return nil, errs.ErrNilPutRecordsResponse + return errs.ErrNilPutRecordsResponse } if resp.FailedPutCount == nil { - return nil, errs.ErrNilFailedRecordCount + return errs.ErrNilFailedRecordCount } attempted := len(messages) failed := int(aws.Int64Value(resp.FailedPutCount)) sent := attempted - failed w.LogDebug(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) - var retries []*message.Message - var err error + var retries int + var wg sync.WaitGroup for idx, record := range resp.RequestResponses { if record.RecordId != nil { // TODO: per-shard metrics messages[idx].RecordID = record.RecordId } else { + retries++ + switch aws.StringValue(record.ErrorCode) { case firehose.ErrCodeLimitExceededException: w.Stats.AddProvisionedThroughputExceeded(1) @@ -117,12 +121,14 @@ func (w *FirehoseWriter) PutRecords(messages []*message.Message) ([]*message.Mes } messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage - messages[idx].FailCount++ - retries = append(retries, messages[idx]) + + wg.Add(1) + go fn(messages[idx], &wg) } } - if len(retries) > 0 { - err = errs.ErrRetryRecords + wg.Wait() + if retries > 0 { + return errs.ErrRetryRecords } - return retries, err + return nil } diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index 8ef5eb5..aca9a23 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -1,8 +1,10 @@ package producer import ( + "context" "fmt" "time" + "sync" "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" @@ -49,14 +51,14 @@ func NewKinesisWriter(c *aws.Config, stream string, fn ...func(*KinesisWriterCon } // PutRecords sends a batch of records to Kinesis and returns a list of records that need to be retried. -func (w *KinesisWriter) PutRecords(messages []*message.Message) ([]*message.Message, error) { +func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageFn) error { var startSendTime time.Time var startBuildTime time.Time start := time.Now() var records []*kinesis.PutRecordsRequestEntry for _, msg := range messages { - records = append(records, msg.MakeRequestEntry()) + records = append(records, msg.ToRequestEntry()) } req, resp := w.client.PutRecordsRequest(&kinesis.PutRecordsInput{ StreamName: aws.String(w.stream), @@ -87,29 +89,31 @@ func (w *KinesisWriter) PutRecords(messages []*message.Message) ([]*message.Mess w.Stats.AddPutRecordsCalled(1) if err := req.Send(); err != nil { w.LogError("Error putting records:", err.Error()) - return nil, err + return err } w.Stats.AddPutRecordsDuration(time.Since(start)) if resp == nil { - return nil, errs.ErrNilPutRecordsResponse + return errs.ErrNilPutRecordsResponse } if resp.FailedRecordCount == nil { - return nil, errs.ErrNilFailedRecordCount + return errs.ErrNilFailedRecordCount } attempted := len(messages) failed := int(aws.Int64Value(resp.FailedRecordCount)) sent := attempted - failed w.LogDebug(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) - var retries []*message.Message - var err error + var retries int + var wg sync.WaitGroup for idx, record := range resp.Records { if record.SequenceNumber != nil && record.ShardId != nil { // TODO: per-shard metrics messages[idx].SequenceNumber = record.SequenceNumber messages[idx].ShardID = record.ShardId } else { + retries++ + switch aws.StringValue(record.ErrorCode) { case kinesis.ErrCodeProvisionedThroughputExceededException: w.Stats.AddProvisionedThroughputExceeded(1) @@ -118,12 +122,14 @@ func (w *KinesisWriter) PutRecords(messages []*message.Message) ([]*message.Mess } messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage - messages[idx].FailCount++ - retries = append(retries, messages[idx]) + + wg.Add(1) + go fn(messages[idx], &wg) } } - if len(retries) > 0 { - err = errs.ErrRetryRecords + wg.Wait() + if retries > 0 { + return errs.ErrRetryRecords } - return retries, err + return nil } diff --git a/producer/producer.go b/producer/producer.go index 7d37144..a5b0ebd 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -18,12 +18,21 @@ import ( // StreamWriter is an interface that abstracts the differences in API between Kinesis and Firehose. type StreamWriter interface { - PutRecords(message []*message.Message) ([]*message.Message, error) + PutRecords(context.Context, []*message.Message, MessageFn) error } // empty is used a as a dummy type for counting semaphore channels. type empty struct{} +// MessageFn defines the signature of a message handler used by Listen, RetrieveFn and their associated *WithContext +// functions. MessageFn accepts a WaitGroup so the function can be run as a blocking operation as opposed to +// MessageFnAsync. +type MessageFn func(*message.Message, *sync.WaitGroup) error + +// MessageFnAsync defines the signature of a message handler used by Listen, RetrieveFn and their associated +// *WithContext functions. MessageFnAsync is meant to be run asynchronously. +type MessageFnAsync func(*message.Message) error + type producerOptions struct { batchSize int batchTimeout time.Duration @@ -104,10 +113,27 @@ func (p *Producer) sendBatch(batch []*message.Message) { attempts := 0 var retries []*message.Message - var err error + stop: for { - retries, err = p.writer.PutRecords(batch) + err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message, wg *sync.WaitGroup) error { + defer wg.Done() + var err error + if msg.FailCount < p.maxRetryAttempts { + msg.FailCount++ + select { + case p.retries <- msg: + case <-p.pipeOfDeath: + err = errs.ErrPipeOfDeath + } + } else { + p.Stats.AddDroppedTotal(1) + p.Stats.AddDroppedRetries(1) + } + retries = append(retries, msg) + + return err + }) failed := len(retries) p.Stats.AddSent(len(batch) - failed) p.Stats.AddFailed(failed) @@ -151,32 +177,21 @@ stop: // batch to be retried rather than retrying the batch as-is. With this approach, we can kill the "stop" // for loop, and set the entire batch to retries to allow the below code to handle retrying the // messages. - attempts++ if attempts > p.maxRetryAttempts { p.LogError(fmt.Sprintf("Dropping batch after %d failed attempts to deliver to stream", attempts)) - p.Stats.AddDropped(len(batch)) + p.Stats.AddDroppedTotal(len(batch)) + p.Stats.AddDroppedRetries(len(batch)) break stop } + attempts++ // Apply an exponential back-off before retrying - time.Sleep(time.Duration(attempts * attempts) * time.Second) + time.Sleep(time.Duration(attempts * 10) * time.Millisecond) } // This frees up another sendBatch to run to allow drainage of the messages / retry queue. This should improve // throughput as well as prevent a potential deadlock in which all batches are blocked on sending retries to the // retries channel, and thus no batches are allowed to drain the retry channel. <-p.concurrencySem - for _, msg := range retries { - if msg.FailCount < p.maxRetryAttempts { - msg.FailCount++ - select { - case p.retries <- msg: - case <-p.pipeOfDeath: - return - } - } else { - p.Stats.AddDropped(1) - } - } } // produce calls the underlying writer's PutRecords implementation to deliver batches of messages to the target stream @@ -291,7 +306,8 @@ func (p *Producer) TryToSend(msg *message.Message) error { case p.messages <- msg: return nil default: - p.Stats.AddDropped(1) + p.Stats.AddDroppedTotal(1) + p.Stats.AddDroppedCapacity(1) return errs.ErrDroppedMessage } } diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go index 5174ded..c3fdac8 100644 --- a/producer/producerconfig_test.go +++ b/producer/producerconfig_test.go @@ -18,7 +18,9 @@ type DebugStatsCollector struct{} func (l *DebugStatsCollector) AddSent(int) {} func (l *DebugStatsCollector) AddFailed(int) {} -func (l *DebugStatsCollector) AddDropped(int) {} +func (l *DebugStatsCollector) AddDroppedTotal(int) {} +func (l *DebugStatsCollector) AddDroppedCapacity(int) {} +func (l *DebugStatsCollector) AddDroppedRetries(int) {} func (l *DebugStatsCollector) AddBatchSize(int) {} func (l *DebugStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} func (l *DebugStatsCollector) AddPutRecordsCalled(int) {} diff --git a/producer/stats.go b/producer/stats.go index 678b43f..6c2a6bc 100644 --- a/producer/stats.go +++ b/producer/stats.go @@ -10,7 +10,9 @@ import ( type StatsCollector interface { AddSent(int) AddFailed(int) - AddDropped(int) + AddDroppedTotal(int) + AddDroppedCapacity(int) + AddDroppedRetries(int) AddBatchSize(int) AddPutRecordsProvisionedThroughputExceeded(int) AddPutRecordsCalled(int) @@ -25,48 +27,47 @@ type StatsCollector interface { // NilStatsCollector is a stats listener that ignores all metrics. type NilStatsCollector struct{} -// AddSent records a count of the number of messages sent to AWS Kinesis by the -// producer. +// AddSent records a count of the number of messages sent to AWS Kinesis by the producer. func (l *NilStatsCollector) AddSent(int) {} -// AddFailed records a count of the number of messages that failed to be sent to -// AWS Kinesis by the producer. +// AddFailed records a count of the number of messages that failed to be sent to AWS Kinesis by the producer. func (l *NilStatsCollector) AddFailed(int) {} -// AddDropped records a count of the number of messages dropped by the -// application after multiple failures. -func (l *NilStatsCollector) AddDropped(int) {} +// AddDroppedTotal records a count of the total number of messages dropped by the application after multiple failures. +func (l *NilStatsCollector) AddDroppedTotal(int) {} -// AddBatchSize records a count of the number of messages attempted by -// PutRecords in the producer. +// AddDroppedCapacity records a count of the number of messages that were dropped by the application due to the stream +// writer being at capacity. +func (l *NilStatsCollector) AddDroppedCapacity(int) {} + +// AddDroppedRetries records a count of the number of retry messages dropped by the application after the max number of +// retries was exceeded. +func (l *NilStatsCollector) AddDroppedRetries(int) {} + +// AddBatchSize records a count of the number of messages attempted by PutRecords in the producer. func (l *NilStatsCollector) AddBatchSize(int) {} -// AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords -// API returned a ErrCodeProvisionedThroughputExceededException by the producer. +// AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords API returned a +// ErrCodeProvisionedThroughputExceededException by the producer. func (l *NilStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} -// AddPutRecordsCalled records the number of times the PutRecords API was called -// by the producer. +// AddPutRecordsCalled records the number of times the PutRecords API was called by the producer. func (l *NilStatsCollector) AddPutRecordsCalled(int) {} -// AddProvisionedThroughputExceeded records the number of times the PutRecords -// API response contained a record which contained an -// ErrCodeProvisionedThroughputExceededException error. +// AddProvisionedThroughputExceeded records the number of times the PutRecords API response contained a record which +// contained an ErrCodeProvisionedThroughputExceededException error. func (l *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} -// AddPutRecordsTimeout records the number of times the PutRecords API timed out -// on the HTTP level. This is influenced by the WithHTTPClientTimeout -// configuration. +// AddPutRecordsTimeout records the number of times the PutRecords API timed out on the HTTP level. This is influenced +// by the WithHTTPClientTimeout configuration. func (l *NilStatsCollector) AddPutRecordsTimeout(int) {} -// AddPutRecordsDuration records the duration that the PutRecords API request -// took. Only the times of successful calls are measured. +// AddPutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls +// are measured. func (l *NilStatsCollector) AddPutRecordsDuration(time.Duration) {} -// AddPutRecordsBuildDuration records the duration that it took to build the -// PutRecords API request payload. +// AddPutRecordsBuildDuration records the duration that it took to build the PutRecords API request payload. func (l *NilStatsCollector) AddPutRecordsBuildDuration(time.Duration) {} -// AddPutRecordsSendDuration records the duration that it took to send the -// PutRecords API request payload. +// AddPutRecordsSendDuration records the duration that it took to send the PutRecords API request payload. func (l *NilStatsCollector) AddPutRecordsSendDuration(time.Duration) {} From 04d6790c31b051c68dd20e71dcbb5ecdcf87319c Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 26 Jun 2017 18:35:36 -0500 Subject: [PATCH 040/109] Refactor readers and writers to have an AwsConfig field instead of using AwsOptions as an embedded type in order to simplify code. Note: The producer and listener still have AwsOptions as an embedded type but only to allow configuration changes provided by the AwsOptions struct. --- listener/kclreaderconfig.go | 5 ++--- listener/kinesisreader.go | 3 ++- listener/kinesisreaderconfig.go | 5 ++--- producer/firehosewriter.go | 3 ++- producer/firehosewriterconfig.go | 5 ++--- producer/kinesiswriter.go | 3 ++- producer/kinesiswriterconfig.go | 5 ++--- 7 files changed, 14 insertions(+), 15 deletions(-) diff --git a/listener/kclreaderconfig.go b/listener/kclreaderconfig.go index 66f8fa4..11a54fc 100644 --- a/listener/kclreaderconfig.go +++ b/listener/kclreaderconfig.go @@ -2,20 +2,19 @@ package listener import ( "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic/config" ) // KclReaderConfig is used to configure KclReader type KclReaderConfig struct { - *config.AwsOptions *kclReaderOptions + AwsConfig *aws.Config LogLevel aws.LogLevelType } // NewKclReaderConfig creates a new instance of KclReaderConfig func NewKclReaderConfig(cfg *aws.Config) *KclReaderConfig { return &KclReaderConfig{ - AwsOptions: config.NewAwsOptionsFromConfig(cfg), + AwsConfig: cfg, kclReaderOptions: &kclReaderOptions{ Stats: &NilStatsCollector{}, }, diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index 108c5b2..0c7d37c 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -10,6 +10,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/awserr" "github.com/aws/aws-sdk-go/aws/request" + "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" @@ -45,7 +46,7 @@ func NewKinesisReader(c *aws.Config, stream string, shard string, fn ...func(*Ki for _, f := range fn { f(cfg) } - sess, err := cfg.GetSession() + sess, err := session.NewSession(cfg.AwsConfig) if err != nil { return nil, err } diff --git a/listener/kinesisreaderconfig.go b/listener/kinesisreaderconfig.go index 900ae13..db3cd16 100644 --- a/listener/kinesisreaderconfig.go +++ b/listener/kinesisreaderconfig.go @@ -3,20 +3,19 @@ package listener import ( "time" "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic/config" ) // KinesisReaderConfig is used to configure a KinesisReader type KinesisReaderConfig struct { - *config.AwsOptions *kinesisReaderOptions + AwsConfig *aws.Config LogLevel aws.LogLevelType } // NewKinesisReaderConfig creates a new instance of KinesisReaderConfig func NewKinesisReaderConfig(cfg *aws.Config) *KinesisReaderConfig { return &KinesisReaderConfig{ - AwsOptions: config.NewAwsOptionsFromConfig(cfg), + AwsConfig: cfg, kinesisReaderOptions: &kinesisReaderOptions{ batchSize: 10000, shardIterator: NewShardIterator(), diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go index 9a9618f..5d537eb 100644 --- a/producer/firehosewriter.go +++ b/producer/firehosewriter.go @@ -8,6 +8,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" + "github.com/aws/aws-sdk-go/aws/session" "github.com/aws/aws-sdk-go/service/firehose" "github.com/aws/aws-sdk-go/service/firehose/firehoseiface" @@ -35,7 +36,7 @@ func NewFirehoseWriter(c *aws.Config, stream string, fn ...func(*FirehoseWriterC for _, f := range fn { f(cfg) } - sess, err := cfg.GetSession() + sess, err := session.NewSession(cfg.AwsConfig) if err != nil { return nil, err } diff --git a/producer/firehosewriterconfig.go b/producer/firehosewriterconfig.go index 9bb4870..8092f71 100644 --- a/producer/firehosewriterconfig.go +++ b/producer/firehosewriterconfig.go @@ -2,20 +2,19 @@ package producer import ( "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic/config" ) // FirehoseWriterConfig is used to configure FirehoseWriter type FirehoseWriterConfig struct { - *config.AwsOptions *firehoseWriterOptions + AwsConfig *aws.Config LogLevel aws.LogLevelType } // NewFirehoseWriterConfig creates a new instance of FirehoseWriterConfig func NewFirehoseWriterConfig(cfg *aws.Config) *FirehoseWriterConfig { return &FirehoseWriterConfig{ - AwsOptions: config.NewAwsOptionsFromConfig(cfg), + AwsConfig: cfg, firehoseWriterOptions: &firehoseWriterOptions{ Stats: &NilStatsCollector{}, }, diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index aca9a23..ce6f4dd 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -8,6 +8,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" + "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" @@ -35,7 +36,7 @@ func NewKinesisWriter(c *aws.Config, stream string, fn ...func(*KinesisWriterCon for _, f := range fn { f(cfg) } - sess, err := cfg.GetSession() + sess, err := session.NewSession(cfg.AwsConfig) if err != nil { return nil, err } diff --git a/producer/kinesiswriterconfig.go b/producer/kinesiswriterconfig.go index f619351..704d329 100644 --- a/producer/kinesiswriterconfig.go +++ b/producer/kinesiswriterconfig.go @@ -2,20 +2,19 @@ package producer import ( "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic/config" ) // KinesisWriterConfig is used to configure KinesisWriter type KinesisWriterConfig struct { - *config.AwsOptions *kinesisWriterOptions + AwsConfig *aws.Config LogLevel aws.LogLevelType } // NewKinesisWriterConfig creates a new instance of KinesisWriterConfig func NewKinesisWriterConfig(cfg *aws.Config) *KinesisWriterConfig { return &KinesisWriterConfig{ - AwsOptions: config.NewAwsOptionsFromConfig(cfg), + AwsConfig: cfg, kinesisWriterOptions: &kinesisWriterOptions{ Stats: &NilStatsCollector{}, }, From 2381724d2dea89958e1fa4cc9546eadd8b2a2e63 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 27 Jun 2017 15:32:51 -0500 Subject: [PATCH 041/109] Renamed MessageFn to MessageHandler and created a new type called MessageProcessor for handling the user-supplied functions for Listen, Retrieve, etc.. Changed PutRecords interface to process the message callback function asynchronously. Fixed potential race condition with the failed counter in the PutRecords lambda function. --- config/config.go | 2 +- config_test.go | 2 +- listener/kclreader.go | 6 ++--- listener/kinesisreader.go | 7 ++--- listener/listener.go | 30 ++++++++++++---------- producer/firehosewriter.go | 14 +++------- producer/kinesiswriter.go | 17 +++---------- producer/producer.go | 52 ++++++++++++++++++-------------------- 8 files changed, 58 insertions(+), 72 deletions(-) diff --git a/config/config.go b/config/config.go index f1c2764..0195f6d 100644 --- a/config/config.go +++ b/config/config.go @@ -19,7 +19,7 @@ func DefaultAwsOptions() *AwsOptions { return &AwsOptions{ AwsConfig: aws.NewConfig().WithHTTPClient( &http.Client{ - Timeout: 10 * time.Minute, + Timeout: 2 * time.Minute, }, ), } diff --git a/config_test.go b/config_test.go index e15ec3f..37b0027 100644 --- a/config_test.go +++ b/config_test.go @@ -25,7 +25,7 @@ func TestNewConfig(t *testing.T) { config := NewConfig() Convey("check the default values for its non-zero config", func() { - So(config.AwsConfig.HTTPClient.Timeout, ShouldEqual, 10*time.Minute) + So(config.AwsConfig.HTTPClient.Timeout, ShouldEqual, 2*time.Minute) So(config.AwsConfig.LogLevel.Value(), ShouldEqual, aws.LogOff) So(config.LogLevel.Value(), ShouldEqual, logging.LogOff) }) diff --git a/listener/kclreader.go b/listener/kclreader.go index bdf6951..6024c63 100644 --- a/listener/kclreader.go +++ b/listener/kclreader.go @@ -57,7 +57,7 @@ func NewKclReader(c *aws.Config, fn ...func(*KclReaderConfig)) (*KclReader, erro // message channel. After all the messages on the message buffer have been moved to the listener's message // channel, a message is sent (following the Multilang protocol) to acknowledge that the processRecords message // has been received / processed -func (r *KclReader) processRecords(fn MessageFn, numRecords int) (int, error) { +func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, error) { // Define the batchSize batchSize := 0; if len(r.msgBuffer) > 0 { @@ -195,12 +195,12 @@ func (r *KclReader) onShutdown() error { // GetRecord calls processRecords to attempt to put one message from message buffer to the listener's message // channel -func (r *KclReader) GetRecord(ctx context.Context,fn MessageFn) (int, error) { +func (r *KclReader) GetRecord(ctx context.Context,fn MessageHandler) (int, error) { return r.processRecords(fn, 1) } // GetRecords calls processRecords to attempt to put all messages on the message buffer on the listener's // message channel -func (r *KclReader) GetRecords(ctx context.Context,fn MessageFn) (int, error) { +func (r *KclReader) GetRecords(ctx context.Context,fn MessageHandler) (int, error) { return r.processRecords(fn, -1) } diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index 0c7d37c..f89a1fa 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -135,7 +135,7 @@ func (r *KinesisReader) throttle(sem chan empty) { // in turn call Close on the underlying HTTPResponse.Body. The question is whether this actually shuts down the TCP // connection. Worst case scenario is that our client Timeout eventually fires and closes the socket, but this can be // susceptible to FD exhaustion. -func (r *KinesisReader) getRecords(ctx context.Context, fn MessageFn, batchSize int) (int, error) { +func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batchSize int) (int, error) { if err := r.ensureShardIterator(); err != nil { return 0, err } @@ -239,6 +239,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageFn, batchSize // Send the GetRecords request r.LogDebug("Starting GetRecords Build/Sign request, took", time.Since(start)) r.Stats.AddGetRecordsCalled(1) + //req.SetContext(ctx) if err := req.Send(); err != nil { r.LogError("Error getting records:", err) switch err.(awserr.Error).Code() { @@ -304,11 +305,11 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageFn, batchSize } // GetRecord calls getRecords and delivers one record into the messages channel. -func (r *KinesisReader) GetRecord(ctx context.Context, fn MessageFn) (int, error) { +func (r *KinesisReader) GetRecord(ctx context.Context, fn MessageHandler) (int, error) { return r.getRecords(ctx, fn, 1) } // GetRecords calls getRecords and delivers each record into the messages channel. -func (r *KinesisReader) GetRecords(ctx context.Context, fn MessageFn) (int, error) { +func (r *KinesisReader) GetRecords(ctx context.Context, fn MessageHandler) (int, error) { return r.getRecords(ctx, fn, r.batchSize) } diff --git a/listener/listener.go b/listener/listener.go index 8f7f024..34c0f32 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -14,21 +14,25 @@ import ( // StreamReader is an interface that abstracts out a stream reader. type StreamReader interface { - GetRecord(context.Context, MessageFn) (int, error) - GetRecords(context.Context, MessageFn) (int, error) + GetRecord(context.Context, MessageHandler) (int, error) + GetRecords(context.Context, MessageHandler) (int, error) } // empty is used a as a dummy type for semaphore channels and the pipe of death channel. type empty struct{} -// MessageFn defines the signature of a message handler used by Listen, RetrieveFn and their associated *WithContext -// functions. MessageFn accepts a WaitGroup so the function can be run as a blocking operation as opposed to -// MessageFnAsync. -type MessageFn func(*message.Message, *sync.WaitGroup) error +// MessageProcessor defines the signature of a message handler used by Listen, RetrieveFn and their associated +// *WithContext functions. MessageHandler accepts a WaitGroup so the function can be run as a blocking operation as +// opposed to MessageHandlerAsync. +type MessageProcessor func(*message.Message, *sync.WaitGroup) error -// MessageFnAsync defines the signature of a message handler used by Listen, RetrieveFn and their associated -// *WithContext functions. MessageFnAsync is meant to be run asynchronously. -type MessageFnAsync func(*message.Message) error +// MessageHandler defines the signature of a message handler used by GetRecord() and GetRecords(). MessageHandler +// accepts a WaitGroup so the function can be run as a blocking operation as opposed to MessageHandlerAsync. +type MessageHandler func(*message.Message, *sync.WaitGroup) error + +// MessageHandlerAsync defines the signature of a message handler used by GetRecord() and GetRecords(). +// MessageHandlerAsync is meant to be run asynchronously. +type MessageHandlerAsync func(*message.Message) error // listenerOptions is used to hold all of the configurable settings of a Listener object. type listenerOptions struct { @@ -142,7 +146,7 @@ func (l *Listener) Retrieve() (*message.Message, error) { // RetrieveFnWithContext retrieves a message from the stream and dispatches it to the supplied function. RetrieveFn // will wait until the function completes. Cancellation is supported through context. -func (l *Listener) RetrieveFnWithContext(ctx context.Context, fn MessageFn) error { +func (l *Listener) RetrieveFnWithContext(ctx context.Context, fn MessageProcessor) error { msg, err := l.RetrieveWithContext(ctx) if err != nil { return err @@ -163,7 +167,7 @@ func (l *Listener) RetrieveFnWithContext(ctx context.Context, fn MessageFn) erro // RetrieveFn retrieves a message from the stream and dispatches it to the supplied function. RetrieveFn will wait // until the function completes. -func (l *Listener) RetrieveFn(fn MessageFn) error { +func (l *Listener) RetrieveFn(fn MessageProcessor) error { return l.RetrieveFnWithContext(context.TODO(), fn) } @@ -220,7 +224,7 @@ func (l *Listener) consume(ctx context.Context) { // ListenWithContext listens and delivers message to the supplied function. Upon cancellation, Listen will stop the // consumer loop and wait until the messages channel is closed and all messages are delivered. -func (l *Listener) ListenWithContext(ctx context.Context, fn MessageFn) { +func (l *Listener) ListenWithContext(ctx context.Context, fn MessageProcessor) { l.consume(ctx) var wg sync.WaitGroup defer wg.Wait() @@ -255,6 +259,6 @@ func (l *Listener) ListenWithContext(ctx context.Context, fn MessageFn) { } // Listen listens and delivers message to the supplied function. -func (l *Listener) Listen(fn MessageFn) { +func (l *Listener) Listen(fn MessageProcessor) { l.ListenWithContext(context.TODO(), fn) } diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go index 5d537eb..81239c0 100644 --- a/producer/firehosewriter.go +++ b/producer/firehosewriter.go @@ -4,7 +4,6 @@ import ( "context" "fmt" "time" - "sync" "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" @@ -27,7 +26,7 @@ type FirehoseWriter struct { *logging.LogHelper stream string - client firehoseiface.FirehoseAPI + client firehoseiface.FirehoseAPI } // NewFirehoseWriter creates a new stream writer to write records to a Kinesis. @@ -52,7 +51,7 @@ func NewFirehoseWriter(c *aws.Config, stream string, fn ...func(*FirehoseWriterC } // PutRecords sends a batch of records to Firehose and returns a list of records that need to be retried. -func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageFn) error { +func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageHandlerAsync) error { var startSendTime time.Time var startBuildTime time.Time @@ -106,7 +105,6 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Mes w.LogDebug(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) var retries int - var wg sync.WaitGroup for idx, record := range resp.RequestResponses { if record.RecordId != nil { // TODO: per-shard metrics @@ -123,13 +121,9 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Mes messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage - wg.Add(1) - go fn(messages[idx], &wg) + go fn(messages[idx]) } } - wg.Wait() - if retries > 0 { - return errs.ErrRetryRecords - } + return nil } diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index ce6f4dd..539e047 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -4,7 +4,6 @@ import ( "context" "fmt" "time" - "sync" "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" @@ -27,7 +26,7 @@ type KinesisWriter struct { *logging.LogHelper stream string - client kinesisiface.KinesisAPI + client kinesisiface.KinesisAPI } // NewKinesisWriter creates a new stream writer to write records to a Kinesis. @@ -52,7 +51,7 @@ func NewKinesisWriter(c *aws.Config, stream string, fn ...func(*KinesisWriterCon } // PutRecords sends a batch of records to Kinesis and returns a list of records that need to be retried. -func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageFn) error { +func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageHandlerAsync) error { var startSendTime time.Time var startBuildTime time.Time @@ -105,16 +104,12 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess sent := attempted - failed w.LogDebug(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) - var retries int - var wg sync.WaitGroup for idx, record := range resp.Records { if record.SequenceNumber != nil && record.ShardId != nil { // TODO: per-shard metrics messages[idx].SequenceNumber = record.SequenceNumber messages[idx].ShardID = record.ShardId } else { - retries++ - switch aws.StringValue(record.ErrorCode) { case kinesis.ErrCodeProvisionedThroughputExceededException: w.Stats.AddProvisionedThroughputExceeded(1) @@ -124,13 +119,9 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage - wg.Add(1) - go fn(messages[idx], &wg) + go fn(messages[idx]) } } - wg.Wait() - if retries > 0 { - return errs.ErrRetryRecords - } + return nil } diff --git a/producer/producer.go b/producer/producer.go index a5b0ebd..647b66f 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -5,6 +5,7 @@ import ( "fmt" "net" "sync" + "sync/atomic" "time" "github.com/aws/aws-sdk-go/aws" @@ -18,20 +19,19 @@ import ( // StreamWriter is an interface that abstracts the differences in API between Kinesis and Firehose. type StreamWriter interface { - PutRecords(context.Context, []*message.Message, MessageFn) error + PutRecords(context.Context, []*message.Message, MessageHandlerAsync) error } // empty is used a as a dummy type for counting semaphore channels. type empty struct{} -// MessageFn defines the signature of a message handler used by Listen, RetrieveFn and their associated *WithContext -// functions. MessageFn accepts a WaitGroup so the function can be run as a blocking operation as opposed to -// MessageFnAsync. -type MessageFn func(*message.Message, *sync.WaitGroup) error +// MessageHandler defines the signature of a message handler used by PutRecords(). MessageHandler accepts a WaitGroup +// so the function can be run as a blocking operation as opposed to MessageHandlerAsync. +type MessageHandler func(*message.Message, *sync.WaitGroup) error -// MessageFnAsync defines the signature of a message handler used by Listen, RetrieveFn and their associated -// *WithContext functions. MessageFnAsync is meant to be run asynchronously. -type MessageFnAsync func(*message.Message) error +// MessageHandlerAsync defines the signature of a message handler used by PutRecords(). MessageHandlerAsync is meant to +// be run asynchronously. +type MessageHandlerAsync func(*message.Message) error type producerOptions struct { batchSize int @@ -86,9 +86,8 @@ func (p *Producer) startProducing() bool { p.producing = true p.messages = make(chan *message.Message, p.queueDepth) p.retries = make(chan *message.Message, p.queueDepth) - p.shutdownCond = &sync.Cond{L: &sync.Mutex{}} + p.shutdownCond = sync.NewCond(new(sync.Mutex)) p.producerWg = new(sync.WaitGroup) - p.outstanding = 0 return true } return false @@ -111,35 +110,34 @@ func (p *Producer) sendBatch(batch []*message.Message) { p.shutdownCond.L.Unlock() }() - attempts := 0 - var retries []*message.Message + var attempts int + var failed uint64 stop: for { - err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message, wg *sync.WaitGroup) error { - defer wg.Done() - var err error + err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message) error { if msg.FailCount < p.maxRetryAttempts { msg.FailCount++ select { case p.retries <- msg: + atomic.AddUint64(&failed, 1) case <-p.pipeOfDeath: - err = errs.ErrPipeOfDeath + return errs.ErrPipeOfDeath } } else { p.Stats.AddDroppedTotal(1) p.Stats.AddDroppedRetries(1) } - retries = append(retries, msg) - return err + return nil }) - failed := len(retries) - p.Stats.AddSent(len(batch) - failed) - p.Stats.AddFailed(failed) - if err == nil && failed == 0 { + p.Stats.AddSent(len(batch) - int(failed)) + p.Stats.AddFailed(int(failed)) + if err == nil { break stop } + + // The call failed so we need to retry the batch switch err := err.(type) { case net.Error: if err.Timeout() { @@ -185,7 +183,7 @@ stop: } attempts++ - // Apply an exponential back-off before retrying + // Apply a delay before retrying time.Sleep(time.Duration(attempts * 10) * time.Millisecond) } // This frees up another sendBatch to run to allow drainage of the messages / retry queue. This should improve @@ -213,11 +211,9 @@ func (p *Producer) produce() { batch: for len(batch) <= p.batchSize { select { - // Using the select, retry messages will - // interleave with new messages. This is - // preferable to putting the messages at the end - // of the channel as it minimizes the delay in - // the delivery of retry messages. + // Using the select, retry messages will interleave with new messages. This is + // preferable to putting the messages at the end of the channel as it minimizes the + // delay in the delivery of retry messages. case msg, ok := <-p.messages: if !ok { p.messages = nil From 76666ac98d022875c8ec6b7099bf8a11c394d269 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 27 Jun 2017 22:30:03 -0500 Subject: [PATCH 042/109] Upgraded to the latest version of the aws-sdk-go library and simplified the ReadCloserWrapper in getRecords(). Added per-request response timeouts to both GetRecordsRequest and PutRecordsRequest. Passed in a child context to GetRecord(s)() calls. --- glide.lock | 7 ++-- glide.yaml | 2 +- kinetic_integration_test.go | 3 +- listener/kinesisreader.go | 61 +++++---------------------------- listener/kinesisreaderconfig.go | 13 +++---- listener/listener.go | 9 +++-- listener/listener_test.go | 2 +- producer/kinesiswriter.go | 4 ++- producer/kinesiswriterconfig.go | 8 +++++ producer/producer_test.go | 2 +- testexec/main.go | 3 +- 11 files changed, 44 insertions(+), 70 deletions(-) diff --git a/glide.lock b/glide.lock index 01a6021..e481c5a 100644 --- a/glide.lock +++ b/glide.lock @@ -1,8 +1,8 @@ -hash: c70e5dd0bac15fcbacf3d07e58cad9df3b23ba556bc583f3a850f79f60ea2bae -updated: 2017-05-25T11:30:54.735250056-05:00 +hash: 16ed7c907af98983bf7cbe350e9ff03fb1208a4ad54664a83e488935889e6207 +updated: 2017-06-27T15:49:17.807220643-05:00 imports: - name: github.com/aws/aws-sdk-go - version: 72df63b404d3f9820db08c73176c1b277d9f614f + version: 72e42b13da62269f68308fb6068b7ea691a416a4 subpackages: - aws - aws/awserr @@ -20,6 +20,7 @@ imports: - aws/request - aws/session - aws/signer/v4 + - internal/shareddefaults - private/protocol - private/protocol/json/jsonutil - private/protocol/jsonrpc diff --git a/glide.yaml b/glide.yaml index fcdcd82..ebe4cdf 100644 --- a/glide.yaml +++ b/glide.yaml @@ -1,7 +1,7 @@ package: github.com/rewardStyle/kinetic import: - package: github.com/aws/aws-sdk-go - version: v1.8.1 + version: v1.10.3 subpackages: - aws - aws/awserr diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index 927e31b..6a7cc38 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -132,6 +132,7 @@ func TestKineticIntegration(t *testing.T) { // Create a new kinesis stream writer w, err := producer.NewKinesisWriter(k.Session.Config, stream, func(kwc *producer.KinesisWriterConfig) { kwc.SetLogLevel(aws.LogDebug) + kwc.SetResponseReadTimeout(time.Second) }) if err != nil { log.Fatalf("Unable to create a new kinesis stream writer due to: %v\n", err) @@ -149,7 +150,7 @@ func TestKineticIntegration(t *testing.T) { assert.NotNil(t, k.Session.Config) r, err := listener.NewKinesisReader(k.Session.Config, stream, shards[0], func(krc *listener.KinesisReaderConfig) { - krc.SetReadTimeout(1000 * time.Millisecond) + krc.SetResponseReadTimeout(time.Second) }) assert.NotNil(t, r) assert.NoError(t, err) diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index f89a1fa..ad4a900 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -21,10 +21,10 @@ import ( // kinesisReaderOptions is used to hold all of the configurable settings of a KinesisReader. type kinesisReaderOptions struct { - batchSize int - shardIterator *ShardIterator - readTimeout time.Duration - Stats StatsCollector + batchSize int + shardIterator *ShardIterator + responseReadTimeout time.Duration + Stats StatsCollector } // KinesisReader handles the API to read records from Kinesis. @@ -152,6 +152,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch Limit: aws.Int64(int64(batchSize)), ShardIterator: aws.String(r.nextShardIterator), }) + req.ApplyOptions(request.WithResponseReadTimeout(r.responseReadTimeout)) // If debug is turned on, add some handlers for GetRecords logging if r.LogLevel.AtLeast(logging.LogDebug) { @@ -168,60 +169,15 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch // // The Unmarshal handler will ultimately call ioutil.ReadAll() on the HTTPResponse.Body stream. // - // Our handler wraps the HTTPResponse.Body with our own ReadCloser so that we can implement a timeout mechanism - // on the Read() call (which is called by the ioutil.ReadAll() function) + // Our handler wraps the HTTPResponse.Body with our own ReadCloser so that we can implement stats collection req.Handlers.Unmarshal.PushFront(func(req *request.Request) { r.LogDebug("Started getRecords Unmarshal, took", time.Since(start)) - // Here, we set a timer that the initial Read() call on HTTPResponse.Body must return by. Note that the - // normal http.Client Timeout is still in effect. startReadTime = time.Now() - timer := time.NewTimer(r.readTimeout) req.HTTPResponse.Body = &ReadCloserWrapper{ ReadCloser: req.HTTPResponse.Body, - OnReadFn: func(stream io.ReadCloser, b []byte) (n int, err error) { - // The OnReadFn will be called each time ioutil.ReadAll calls Read on the - // ReadCloserWrapper. - - // First, we set up a struct that to hold the results of the Read() call that can go - // through a channel - type Result struct { - n int - err error - } - - // Next, we build a channel with which to pass the Read() results - c := make(chan Result, 1) - - // Now, we call the Read() on the HTTPResponse.Body in a goroutine and feed the results - // into the channel - readStart := time.Now() - go func() { - var result Result - result.n, result.err = stream.Read(b) - c <- result - }() - - // Finally, we poll for the Read() to complete or the timer to elapse. - select { - case result := <-c: - // If we sucessfully Read() from the HTTPResponse.Body, we reset our timeout and - // return the results from the Read() - timer.Reset(r.readTimeout) - n, err = result.n, result.err - r.LogDebug(fmt.Sprintf("getRecords read %d bytes, took %v", n, time.Since(readStart))) - case <-timer.C: - // If we timeout, we return an error that will unblock ioutil.ReadAll(). This - // will cause the Unmarshal handler to return an error. This error will - // propogate to the original req.Send() call (below) - r.LogDebug(fmt.Sprintf("getRecords read timed out after %v", time.Since(readStart))) - err = errs.ErrTimeoutReadResponseBody - case <-ctx.Done(): - // The pipe of death will abort any pending reads on a GetRecords call. - r.LogDebug(fmt.Sprintf("getRecords received ctx.Done() after %v", time.Since(readStart))) - err = ctx.Err() - } - return + OnReadFn: func(stream io.ReadCloser, b []byte) (int, error) { + return stream.Read(b) }, OnCloseFn: func() { r.Stats.AddGetRecordsReadResponseDuration(time.Since(startReadTime)) @@ -239,7 +195,6 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch // Send the GetRecords request r.LogDebug("Starting GetRecords Build/Sign request, took", time.Since(start)) r.Stats.AddGetRecordsCalled(1) - //req.SetContext(ctx) if err := req.Send(); err != nil { r.LogError("Error getting records:", err) switch err.(awserr.Error).Code() { diff --git a/listener/kinesisreaderconfig.go b/listener/kinesisreaderconfig.go index db3cd16..e8bbd35 100644 --- a/listener/kinesisreaderconfig.go +++ b/listener/kinesisreaderconfig.go @@ -2,6 +2,7 @@ package listener import ( "time" + "github.com/aws/aws-sdk-go/aws" ) @@ -17,10 +18,10 @@ func NewKinesisReaderConfig(cfg *aws.Config) *KinesisReaderConfig { return &KinesisReaderConfig{ AwsConfig: cfg, kinesisReaderOptions: &kinesisReaderOptions{ - batchSize: 10000, + batchSize: 10000, shardIterator: NewShardIterator(), - readTimeout: time.Second, - Stats: &NilStatsCollector{}, + responseReadTimeout: time.Second, + Stats: &NilStatsCollector{}, }, LogLevel: *cfg.LogLevel, } @@ -37,9 +38,9 @@ func (c *KinesisReaderConfig) SetInitialShardIterator(shardIterator *ShardIterat c.shardIterator = shardIterator } -// SetReadTimeout configures the time to wait for each successive Read operation on the GetRecords response payload. -func (c *KinesisReaderConfig) SetReadTimeout(timeout time.Duration) { - c.readTimeout = timeout +// SetResponseReadTimeout configures the time to wait for each successive Read operation on the GetRecords response payload. +func (c *KinesisReaderConfig) SetResponseReadTimeout(timeout time.Duration) { + c.responseReadTimeout = timeout } // SetStatsCollector configures a listener to handle listener metrics. diff --git a/listener/listener.go b/listener/listener.go index 34c0f32..d359a6c 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -115,6 +115,9 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, e return nil, errs.ErrAlreadyConsuming } defer l.stopConsuming() + + childCtx, cancel := context.WithCancel(ctx) + defer cancel() for { // A cancellation or closing the pipe of death will cause Retrieve (and related functions) to abort in // between getRecord calls. Note, that this would only occur when there are no new records to retrieve. @@ -123,7 +126,7 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, e if !ok { return nil, err } - n, err := l.reader.GetRecord(context.TODO(), func(msg *message.Message, wg *sync.WaitGroup) error { + n, err := l.reader.GetRecord(childCtx, func(msg *message.Message, wg *sync.WaitGroup) error { defer wg.Done() l.messages <- msg @@ -180,6 +183,8 @@ func (l *Listener) consume(ctx context.Context) { go func() { defer l.stopConsuming() + childCtx, cancel := context.WithCancel(ctx) + defer cancel() for { // The consume loop can be cancelled by a calling the cancellation function on the context or by // closing the pipe of death. Note that in the case of context cancellation, the getRecords @@ -190,7 +195,7 @@ func (l *Listener) consume(ctx context.Context) { return } - _, err := l.reader.GetRecords(context.TODO(), + _, err := l.reader.GetRecords(childCtx, func(msg *message.Message, wg *sync.WaitGroup) error { defer wg.Done() l.messages <- msg diff --git a/listener/listener_test.go b/listener/listener_test.go index eb12c70..d79cc7d 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -57,7 +57,7 @@ func TestListener(t *testing.T) { So(k.Session, ShouldNotBeNil) So(k.Session.Config, ShouldNotBeNil) r, err := NewKinesisReader(k.Session.Config, stream, shards[0], func(krc *KinesisReaderConfig) { - krc.SetReadTimeout(time.Second) + krc.SetResponseReadTimeout(time.Second) }) So(r, ShouldNotBeNil) So(err, ShouldBeNil) diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index 539e047..688b50a 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -17,7 +17,8 @@ import ( ) type kinesisWriterOptions struct { - Stats StatsCollector + responseReadTimeout time.Duration + Stats StatsCollector } // KinesisWriter handles the API to send records to Kinesis. @@ -64,6 +65,7 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess StreamName: aws.String(w.stream), Records: records, }) + req.ApplyOptions(request.WithResponseReadTimeout(w.responseReadTimeout)) req.Handlers.Build.PushFront(func(r *request.Request) { startBuildTime = time.Now() diff --git a/producer/kinesiswriterconfig.go b/producer/kinesiswriterconfig.go index 704d329..878c34e 100644 --- a/producer/kinesiswriterconfig.go +++ b/producer/kinesiswriterconfig.go @@ -1,6 +1,8 @@ package producer import ( + "time" + "github.com/aws/aws-sdk-go/aws" ) @@ -16,12 +18,18 @@ func NewKinesisWriterConfig(cfg *aws.Config) *KinesisWriterConfig { return &KinesisWriterConfig{ AwsConfig: cfg, kinesisWriterOptions: &kinesisWriterOptions{ + responseReadTimeout: time.Second, Stats: &NilStatsCollector{}, }, LogLevel: *cfg.LogLevel, } } +// SetResponseReadTimeout configures the time to wait for each successive Read operation on the GetRecords response payload. +func (c *KinesisWriterConfig) SetResponseReadTimeout(timeout time.Duration) { + c.responseReadTimeout = timeout +} + // SetStatsCollector configures a listener to handle listener metrics. func (c *KinesisWriterConfig) SetStatsCollector(stats StatsCollector) { c.Stats = stats diff --git a/producer/producer_test.go b/producer/producer_test.go index 7b2a1f2..e5b927a 100644 --- a/producer/producer_test.go +++ b/producer/producer_test.go @@ -55,7 +55,7 @@ func TestProducer(t *testing.T) { So(k.Session.Config, ShouldNotBeNil) r, err := listener.NewKinesisReader(k.Session.Config, stream, shards[0], func(krc *listener.KinesisReaderConfig) { - krc.SetReadTimeout(1000 * time.Millisecond) + krc.SetResponseReadTimeout(1000 * time.Millisecond) }) So(r, ShouldNotBeNil) So(err, ShouldBeNil) diff --git a/testexec/main.go b/testexec/main.go index e1158d3..1aa484c 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -325,6 +325,7 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce w, err := producer.NewKinesisWriter(k.Session.Config, streamName, func(kwc *producer.KinesisWriterConfig) { kwc.SetLogLevel(aws.LogDebug) + kwc.SetResponseReadTimeout(time.Second) }) if err != nil { log.Fatalf("Unable to create a new kinesis stream writer due to: %v\n", err) @@ -354,7 +355,7 @@ func newKineticListener(k *kinetic.Kinetic, streamName string) *listener.Listene r, err := listener.NewKinesisReader(k.Session.Config, streamName, shards[0], func(krc *listener.KinesisReaderConfig) { - krc.SetReadTimeout(1000 * time.Millisecond) + krc.SetResponseReadTimeout(1000 * time.Millisecond) }) l, err := listener.NewListener(k.Session.Config, r, func(c *listener.Config) { From 4217118810d71908a893ef4bbf5906971e5c6871 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 29 Jun 2017 17:13:45 -0500 Subject: [PATCH 043/109] Fixed a bug with setting the batchSize to 500. Added validation logic to prevent the user from setting the batchSize to an invalid value. Add stats to kclreader.go --- errs/errors.go | 3 ++ kinetic_integration_test.go | 1 - listener/kclreader.go | 8 +++-- listener/kinesisreader.go | 6 ---- listener/kinesisreader_test.go | 64 +++++++++++++++++++++++++++++++++ listener/listener_test.go | 2 -- listener/listenerconfig_test.go | 6 ---- producer/producer.go | 2 +- producer/producerconfig.go | 10 +++++- producer/producerconfig_test.go | 4 +-- testexec/cli.go | 15 ++++++++ testexec/main.go | 2 +- 12 files changed, 100 insertions(+), 23 deletions(-) create mode 100644 listener/kinesisreader_test.go diff --git a/errs/errors.go b/errs/errors.go index c81cd17..bf9954d 100644 --- a/errs/errors.go +++ b/errs/errors.go @@ -83,4 +83,7 @@ var ( // ErrDroppedMessage is returned when the message channel is full and messages are being dropped ErrDroppedMessage = errors.New("Channel is full, dropped message") + + // ErrInvalidBatchSize is returned when the batchSize is invalid + ErrInvalidBatchSize = errors.New("PutRecordsBatch supports batch sizes less than or equal to 500") ) diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index 6a7cc38..3e0f30e 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -154,7 +154,6 @@ func TestKineticIntegration(t *testing.T) { }) assert.NotNil(t, r) assert.NoError(t, err) - //assert.NotNil(t, r.Session) // Create a new kinetic listener l, err := listener.NewListener(k.Session.Config, r, func(c *listener.Config) { diff --git a/listener/kclreader.go b/listener/kclreader.go index 6024c63..2df9e0a 100644 --- a/listener/kclreader.go +++ b/listener/kclreader.go @@ -67,13 +67,15 @@ func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, erro batchSize = int(math.Min(float64(len(r.msgBuffer)), float64(numRecords))) } } + r.Stats.AddBatchSize(batchSize) - // Loop through the message buffer and put the correct number of messages on the listener's message channel + // Loop through the message buffer and call the message handler function on each message var wg sync.WaitGroup for i := 0; i < batchSize; i++ { wg.Add(1) go fn(&r.msgBuffer[0], &wg) r.msgBuffer = r.msgBuffer[1:] + r.Stats.AddConsumed(1) } wg.Wait() @@ -195,12 +197,12 @@ func (r *KclReader) onShutdown() error { // GetRecord calls processRecords to attempt to put one message from message buffer to the listener's message // channel -func (r *KclReader) GetRecord(ctx context.Context,fn MessageHandler) (int, error) { +func (r *KclReader) GetRecord(ctx context.Context, fn MessageHandler) (int, error) { return r.processRecords(fn, 1) } // GetRecords calls processRecords to attempt to put all messages on the message buffer on the listener's // message channel -func (r *KclReader) GetRecords(ctx context.Context,fn MessageHandler) (int, error) { +func (r *KclReader) GetRecords(ctx context.Context, fn MessageHandler) (int, error) { return r.processRecords(fn, -1) } diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index ad4a900..59c516a 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -129,12 +129,6 @@ func (r *KinesisReader) throttle(sem chan empty) { }) } -// FIXME: Need to investigate that the timeout implementation doesn't result in an fd leak. Since we call Read on the -// HTTPResonse.Body in a select with a timeout channel, we do prevent ourself from blocking. Once we timeout, we return -// an error to the outer ioutil.ReadAll, which should result in a call to our io.ReadCloser's Close function. This will -// in turn call Close on the underlying HTTPResponse.Body. The question is whether this actually shuts down the TCP -// connection. Worst case scenario is that our client Timeout eventually fires and closes the socket, but this can be -// susceptible to FD exhaustion. func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batchSize int) (int, error) { if err := r.ensureShardIterator(); err != nil { return 0, err diff --git a/listener/kinesisreader_test.go b/listener/kinesisreader_test.go new file mode 100644 index 0000000..d661973 --- /dev/null +++ b/listener/kinesisreader_test.go @@ -0,0 +1,64 @@ +package listener + +import ( + . "github.com/smartystreets/goconvey/convey" + + "math/rand" + "testing" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/logging" +) + +func TestNewKinesisReader(t *testing.T) { + // Set the RNG Seed based on current time (in order to randomize the RNG) + rand.Seed(time.Now().UTC().UnixNano()) + + Convey("given a kinetic object and kinesis stream/shard names", t, func() { + // Instantiate a new kinentic object + k, err := kinetic.New(func(c *kinetic.Config) { + c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") + c.SetRegion("some-region") + c.SetEndpoint("http://127.0.0.1:4567") + }) + So(k, ShouldNotBeNil) + So(err, ShouldBeNil) + + stream := "some-stream" + shard := "some-shard" + + Convey("check that we can create a new KinesisReader with default values", func() { + r, err := NewKinesisReader(k.Session.Config, stream, shard) + So(r, ShouldNotBeNil) + So(err, ShouldBeNil) + So(r.batchSize, ShouldEqual, 10000) + So(r.shardIterator, ShouldNotBeNil) + So(r.responseReadTimeout, ShouldEqual, time.Second) + So(r.Stats, ShouldNotBeNil) + }) + + Convey("check that we can create a new KinesisReader with configured values", func() { + batchSize := rand.Int() + respReadTimeout := time.Duration(rand.Int()) * time.Second + logLevel := aws.LogDebug | aws.LogDebugWithSigning | logging.LogDebug + shardIterator := NewShardIterator() + myStatsCollector := &NilStatsCollector{} + r, err := NewKinesisReader(k.Session.Config, stream, shard, func(krc *KinesisReaderConfig) { + krc.SetBatchSize(batchSize) + krc.SetResponseReadTimeout(respReadTimeout) + krc.SetLogLevel(logLevel) + krc.SetInitialShardIterator(shardIterator) + krc.SetStatsCollector(myStatsCollector) + }) + So(r, ShouldNotBeNil) + So(err, ShouldBeNil) + So(r.batchSize, ShouldEqual, batchSize) + So(r.responseReadTimeout, ShouldEqual, respReadTimeout) + So(r.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) + So(r.shardIterator, ShouldEqual, shardIterator) + So(r.Stats, ShouldEqual, myStatsCollector) + }) + }) +} diff --git a/listener/listener_test.go b/listener/listener_test.go index d79cc7d..c142de1 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -314,8 +314,6 @@ func TestListener(t *testing.T) { Printf("(count was %d)", atomic.LoadInt64(&count)) }) - // TODO: test get records read timeout - Reset(func() { k.DeleteStream(stream) k.WaitUntilStreamDeleted(context.TODO(), stream, request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) diff --git a/listener/listenerconfig_test.go b/listener/listenerconfig_test.go index 0d55cd0..79b57f5 100644 --- a/listener/listenerconfig_test.go +++ b/listener/listenerconfig_test.go @@ -68,12 +68,6 @@ func TestNewConfig(t *testing.T) { So(config.concurrency, ShouldEqual, 50) }) - //TODO: Move this test to kinesisreaderconfig_test.go - //Convey("check that we can set the read timeout for the GetRecords request", func() { - // config.SetReadTimeout(10 * time.Second) - // So(config.readTimeout, ShouldEqual, 10*time.Second) - //}) - Convey("check that we can configure a stats collector", func() { config.SetStatsCollector(&DebugStatsCollector{}) So(config.Stats, ShouldHaveSameTypeAs, &DebugStatsCollector{}) diff --git a/producer/producer.go b/producer/producer.go index 647b66f..ee3f902 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -209,7 +209,7 @@ func (p *Producer) produce() { var batch []*message.Message timer := time.After(p.batchTimeout) batch: - for len(batch) <= p.batchSize { + for len(batch) < p.batchSize { select { // Using the select, retry messages will interleave with new messages. This is // preferable to putting the messages at the end of the channel as it minimizes the diff --git a/producer/producerconfig.go b/producer/producerconfig.go index 71c9ef0..02c5fd3 100644 --- a/producer/producerconfig.go +++ b/producer/producerconfig.go @@ -1,10 +1,12 @@ package producer import ( + "log" "time" "github.com/aws/aws-sdk-go/aws" "github.com/rewardStyle/kinetic/config" + "github.com/rewardStyle/kinetic/errs" ) // Config is used to configure a Producer instance. @@ -32,7 +34,13 @@ func NewConfig(cfg *aws.Config) *Config { // SetBatchSize configures the batch size to flush pending records to the PutRecords call. func (c *Config) SetBatchSize(batchSize int) { - c.batchSize = batchSize + if batchSize > 0 && batchSize <= 500 { + c.batchSize = batchSize + } else { + // http://docs.aws.amazon.com/firehose/latest/APIReference/API_PutRecordBatch.html + log.Fatal("BatchSize must be less than or equal to 500 ") + panic(errs.ErrInvalidBatchSize) + } } // SetBatchTimeout configures the timeout to flush pending records to the PutRecords call. diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go index c3fdac8..25bfbc7 100644 --- a/producer/producerconfig_test.go +++ b/producer/producerconfig_test.go @@ -75,8 +75,8 @@ func TestNewConfig(t *testing.T) { }) Convey("check that we can set the batch size", func() { - cfg.SetBatchSize(1000) - So(cfg.batchSize, ShouldEqual, 1000) + cfg.SetBatchSize(100) + So(cfg.batchSize, ShouldEqual, 100) }) Convey("check that we can set the batch timeout", func() { diff --git a/testexec/cli.go b/testexec/cli.go index 1131507..2b30430 100644 --- a/testexec/cli.go +++ b/testexec/cli.go @@ -3,10 +3,12 @@ package main import ( "flag" "log" + "strings" ) // Config is a data structure used to hold this program's configuration info type Config struct { + Mode *string Location *string StreamName *string NumMsgs *int @@ -18,6 +20,7 @@ type Config struct { func parseCommandLineArgs() *Config { // Define command line flags + modePtr := flag.String("mode", "write", "used to specify the mode in which to run; either read or write") locationPtr := flag.String("location", "local", "used to specify the location of the kinesis stream. " + "Accepted values are (local|aws). For local, run kinesalite on http://127.0.0.1:4567. For aws, your " + "aws credentials and configuration need to be defined at ~/.aws") @@ -47,7 +50,18 @@ func parseCommandLineArgs() *Config { durationPtr = nil } + var mode string + switch strings.ToLower(*modePtr) { + case "read": + fallthrough + case "write": + mode = strings.ToLower(*modePtr) + default: + log.Fatal("Mode must be defined as either 'read' or 'write'") + } + return &Config{ + Mode: &mode, StreamName: streamNamePtr, Duration: durationPtr, NumMsgs: numMsgsPtr, @@ -60,6 +74,7 @@ func parseCommandLineArgs() *Config { func (c *Config) printConfigs() { if *c.Verbose { log.Println("Command Line Arguments:") + log.Println("-mode: ", *c.Mode) log.Println("-location: ", *c.Location) log.Println("-stream-name: ", *c.StreamName) if c.NumMsgs != nil { diff --git a/testexec/main.go b/testexec/main.go index 1aa484c..9e55a19 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -332,7 +332,7 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce } p, err := producer.NewProducer(k.Session.Config, w, func(c *producer.Config) { - c.SetBatchSize(5) + c.SetBatchSize(500) c.SetBatchTimeout(1000 * time.Millisecond) }) if err != nil { From 75d928961751f47197688301e581aaddea5f775d Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 5 Jul 2017 17:50:10 -0500 Subject: [PATCH 044/109] Simplified ReadCloserWrapper by removing the OnReadFn field which is no longer needed because of the aws-sdk-go version upgrade. Listenerconfig and producerconfig now have aws.Config as a field instead of AwsOption. Major refactor of test executable to be able to run in 'read', 'write' and 'readwrite' mode in addition to 'throttle' and 'blast' mode for the producer. --- config/config.go | 7 - listener/kinesisreader.go | 4 - listener/listenerconfig.go | 7 +- listener/listenerconfig_test.go | 15 - listener/readcloser.go | 7 - producer/kinesiswriter.go | 1 + producer/producer.go | 15 +- producer/producerconfig.go | 7 +- producer/producerconfig_test.go | 16 - testexec/cli.go | 19 +- testexec/listenerstats.go | 117 +++++++ testexec/main.go | 471 ++++++++++++++++++--------- testexec/producerstats.go | 117 +++++++ testexec/{stats.go => streamdata.go} | 14 +- 14 files changed, 587 insertions(+), 230 deletions(-) create mode 100644 testexec/listenerstats.go create mode 100644 testexec/producerstats.go rename testexec/{stats.go => streamdata.go} (69%) diff --git a/config/config.go b/config/config.go index 0195f6d..a86d1d9 100644 --- a/config/config.go +++ b/config/config.go @@ -25,13 +25,6 @@ func DefaultAwsOptions() *AwsOptions { } } -// NewAwsOptionsFromConfig initializes a new AwsOption from an existing aws Config -func NewAwsOptionsFromConfig(c *aws.Config) *AwsOptions { - return &AwsOptions{ - AwsConfig: c, - } -} - // SetCredentials configures AWS credentials. func (c *AwsOptions) SetCredentials(accessKey, secretKey, sessionToken string) { c.AwsConfig.WithCredentials( diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index 59c516a..e3cdb5b 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -3,7 +3,6 @@ package listener import ( "context" "fmt" - "io" "sync" "time" @@ -170,9 +169,6 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch req.HTTPResponse.Body = &ReadCloserWrapper{ ReadCloser: req.HTTPResponse.Body, - OnReadFn: func(stream io.ReadCloser, b []byte) (int, error) { - return stream.Read(b) - }, OnCloseFn: func() { r.Stats.AddGetRecordsReadResponseDuration(time.Since(startReadTime)) r.LogDebug("Finished GetRecords body read, took", time.Since(start)) diff --git a/listener/listenerconfig.go b/listener/listenerconfig.go index 52e207a..6f5e893 100644 --- a/listener/listenerconfig.go +++ b/listener/listenerconfig.go @@ -2,21 +2,20 @@ package listener import ( "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic/config" "github.com/rewardStyle/kinetic/logging" ) // Config is used to configure a Listener instance type Config struct { - *config.AwsOptions *listenerOptions + AwsConfig *aws.Config LogLevel aws.LogLevelType } // NewConfig creates a new instance of Config func NewConfig(cfg *aws.Config) *Config { return &Config{ - AwsOptions: config.NewAwsOptionsFromConfig(cfg), + AwsConfig: cfg, listenerOptions: &listenerOptions{ queueDepth: 10000, concurrency: 10000, @@ -44,7 +43,5 @@ func (c *Config) SetStatsCollector(stats StatsCollector) { // SetLogLevel configures both the SDK and Kinetic log levels. func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { - // TODO: Do we want to change the AWS log level? - c.AwsOptions.SetLogLevel(logLevel) c.LogLevel = logLevel & 0xffff0000 } diff --git a/listener/listenerconfig_test.go b/listener/listenerconfig_test.go index 79b57f5..531056f 100644 --- a/listener/listenerconfig_test.go +++ b/listener/listenerconfig_test.go @@ -7,7 +7,6 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/aws/session" "github.com/rewardStyle/kinetic" "github.com/rewardStyle/kinetic/logging" ) @@ -27,13 +26,6 @@ func (l *DebugStatsCollector) AddGetRecordsDuration(time.Duration) { func (l *DebugStatsCollector) AddGetRecordsReadResponseDuration(time.Duration) {} func (l *DebugStatsCollector) AddGetRecordsUnmarshalDuration(time.Duration) {} -func getSession(config *Config) *session.Session { - sess, err := config.GetSession() - So(err, ShouldBeNil) - So(sess, ShouldNotBeNil) - return sess -} - func TestNewConfig(t *testing.T) { Convey("given a new listener config", t, func() { k, err := kinetic.New(func(c *kinetic.Config) { @@ -50,16 +42,9 @@ func TestNewConfig(t *testing.T) { So(config.LogLevel.Value(), ShouldEqual, logging.LogOff) }) - Convey("check that we can retrieve an aws.Session from it ", func() { - getSession(config) - }) - Convey("check that we can set both the sdk and kinetic log level", func() { ll := aws.LogDebug | aws.LogDebugWithSigning | logging.LogDebug config.SetLogLevel(ll) - sess := getSession(config) - So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) - So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) So(config.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) }) diff --git a/listener/readcloser.go b/listener/readcloser.go index 252bdce..57b011b 100644 --- a/listener/readcloser.go +++ b/listener/readcloser.go @@ -9,16 +9,9 @@ import ( // the aws-sdk-go. type ReadCloserWrapper struct { io.ReadCloser - OnReadFn func(io.ReadCloser, []byte) (int, error) OnCloseFn func() } -// Read is called to read the wrapped stream. The supplied OnReadFn is -// responsible for making the read to the wrapped stream. -func (r *ReadCloserWrapper) Read(b []byte) (int, error) { - return r.OnReadFn(r.ReadCloser, b) -} - // Close is called to close the wrapped stream. The supplied OnCloseFn is not // responsible for closing the stream. func (r *ReadCloserWrapper) Close() error { diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index 688b50a..cb7ece2 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -120,6 +120,7 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess } messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage + messages[idx].FailCount++ go fn(messages[idx]) } diff --git a/producer/producer.go b/producer/producer.go index ee3f902..6135b38 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -116,8 +116,10 @@ func (p *Producer) sendBatch(batch []*message.Message) { stop: for { err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message) error { - if msg.FailCount < p.maxRetryAttempts { - msg.FailCount++ + if msg.FailCount <= p.maxRetryAttempts { + // Apply a delay before retrying + time.Sleep(time.Duration(msg.FailCount * msg.FailCount) * time.Second) + select { case p.retries <- msg: atomic.AddUint64(&failed, 1) @@ -184,11 +186,12 @@ stop: attempts++ // Apply a delay before retrying - time.Sleep(time.Duration(attempts * 10) * time.Millisecond) + time.Sleep(time.Duration(attempts * attempts) * time.Second) } - // This frees up another sendBatch to run to allow drainage of the messages / retry queue. This should improve - // throughput as well as prevent a potential deadlock in which all batches are blocked on sending retries to the - // retries channel, and thus no batches are allowed to drain the retry channel. + + // This frees up another sendBatch to run to allow drainage of the messages / retry queue. This should + // improve throughput as well as prevent a potential deadlock in which all batches are blocked on + // sending retries to the retries channel, and thus no batches are allowed to drain the retry channel. <-p.concurrencySem } diff --git a/producer/producerconfig.go b/producer/producerconfig.go index 02c5fd3..e7b7526 100644 --- a/producer/producerconfig.go +++ b/producer/producerconfig.go @@ -5,21 +5,20 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic/config" "github.com/rewardStyle/kinetic/errs" ) // Config is used to configure a Producer instance. type Config struct { - *config.AwsOptions *producerOptions + AwsConfig *aws.Config LogLevel aws.LogLevelType } // NewConfig creates a new instance of Config. func NewConfig(cfg *aws.Config) *Config { return &Config{ - AwsOptions: config.NewAwsOptionsFromConfig(cfg), + AwsConfig: cfg, producerOptions: &producerOptions{ batchSize: 500, batchTimeout: time.Second, @@ -70,7 +69,5 @@ func (c *Config) SetStatsCollector(stats StatsCollector) { // SetLogLevel configures both the SDK and Kinetic log levels. func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { - // TODO: Do we want to change the AWS log level? - c.AwsOptions.SetLogLevel(logLevel) c.LogLevel = logLevel & 0xffff0000 } diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go index 25bfbc7..8704ea8 100644 --- a/producer/producerconfig_test.go +++ b/producer/producerconfig_test.go @@ -7,8 +7,6 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/aws/session" - "github.com/rewardStyle/kinetic" "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" @@ -30,13 +28,6 @@ func (l *DebugStatsCollector) AddPutRecordsDuration(time.Duration) {} func (l *DebugStatsCollector) AddPutRecordsBuildDuration(time.Duration) {} func (l *DebugStatsCollector) AddPutRecordsSendDuration(time.Duration) {} -func getSession(config *Config) *session.Session { - sess, err := config.GetSession() - So(err, ShouldBeNil) - So(sess, ShouldNotBeNil) - return sess -} - type DebugStreamWriter struct{} func (w *DebugStreamWriter) PutRecords(batch []*message.Message) ([]*message.Message, error) { @@ -61,16 +52,9 @@ func TestNewConfig(t *testing.T) { So(cfg.LogLevel.Value(), ShouldEqual, logging.LogOff) }) - Convey("check that we can retrieve an aws.Session from it ", func() { - getSession(cfg) - }) - Convey("check that we can set both the sdk and kinetic log level", func() { ll := aws.LogDebug | aws.LogDebugWithSigning | logging.LogDebug cfg.SetLogLevel(ll) - sess := getSession(cfg) - So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) - So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) So(cfg.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) }) diff --git a/testexec/cli.go b/testexec/cli.go index 2b30430..968809d 100644 --- a/testexec/cli.go +++ b/testexec/cli.go @@ -13,6 +13,7 @@ type Config struct { StreamName *string NumMsgs *int Duration *int + Throttle *bool Cleanup *bool Verbose *bool } @@ -20,7 +21,8 @@ type Config struct { func parseCommandLineArgs() *Config { // Define command line flags - modePtr := flag.String("mode", "write", "used to specify the mode in which to run; either read or write") + modePtr := flag.String("mode", "readwrite", "used to specify the mode in which to run; either 'r', 'read', 'w', " + + "'write', 'rw' or 'readwrite'") locationPtr := flag.String("location", "local", "used to specify the location of the kinesis stream. " + "Accepted values are (local|aws). For local, run kinesalite on http://127.0.0.1:4567. For aws, your " + "aws credentials and configuration need to be defined at ~/.aws") @@ -30,6 +32,7 @@ func parseCommandLineArgs() *Config { "receive. Either -num-msgs or -duration must be set.") durationPtr := flag.Int("duration", 0, "used to specify the duration (in seconds) the program should run. " + "Use a value of -1 to run indefinitely. Either -num-msgs or -duration must be set.") + throttlePtr := flag.Bool("throttle", true, "used to specify whether to throttle PutRecord requests by 1 ms. ") cleanupPtr := flag.Bool("cleanup", true, "used to specify whether or not to delete a newly created kinesis " + "stream") verbosePtr := flag.Bool("verbose", false, "used to specify whether or not to log in verbose mode") @@ -52,12 +55,20 @@ func parseCommandLineArgs() *Config { var mode string switch strings.ToLower(*modePtr) { + case "r": + fallthrough case "read": + mode = ModeRead + case "w": fallthrough case "write": - mode = strings.ToLower(*modePtr) + mode = ModeWrite + case "rw": + fallthrough + case "readwrite": + mode = ModeReadWrite default: - log.Fatal("Mode must be defined as either 'read' or 'write'") + log.Fatal("Mode must be defined as either 'r', 'read', 'w', 'write', 'rw' or 'readwrite'") } return &Config{ @@ -66,6 +77,7 @@ func parseCommandLineArgs() *Config { Duration: durationPtr, NumMsgs: numMsgsPtr, Location: locationPtr, + Throttle: throttlePtr, Cleanup: cleanupPtr, Verbose: verbosePtr, } @@ -83,6 +95,7 @@ func (c *Config) printConfigs() { if c.Duration != nil { log.Println("-duration: ", *c.Duration) } + log.Println("-throttle: ", *c.Throttle) log.Println("-cleanup: ", *c.Cleanup) log.Println("-verbose: ", *c.Verbose) log.Println() diff --git a/testexec/listenerstats.go b/testexec/listenerstats.go new file mode 100644 index 0000000..8e70d1a --- /dev/null +++ b/testexec/listenerstats.go @@ -0,0 +1,117 @@ +package main + +import ( + "sync/atomic" + "time" + "log" +) + +// ListenerStatsCollector implements the listener's StatsCollector +type ListenerStatsCollector struct { + Consumed uint64 + Delivered uint64 + Processed uint64 + BatchSize uint64 + GetRecordsCalled uint64 + ProvisionedThroughputExceeded uint64 + GetRecordsTimeout uint64 + GetRecordsReadTimeout uint64 + + ProcessedDuration time.Duration + GetRecordsDuration time.Duration + GetRecordsReadResonseDuration time.Duration + GetRecordsUnmarshalDuration time.Duration +} + +// AddConsumed records a count of the number of messages received from AWS +// Kinesis by the listener. +func (lsc *ListenerStatsCollector) AddConsumed(count int) { + atomic.AddUint64(&lsc.Consumed, uint64(count)) +} + +// AddDelivered records a count of the number of messages delivered to the +// application by the listener. +func (lsc *ListenerStatsCollector) AddDelivered(count int) { + atomic.AddUint64(&lsc.Delivered, uint64(count)) +} + +// AddProcessed records a count of the number of messages processed by the +// application by the listener. This is based on a WaitGroup that is sent to +// the RetrieveFn and Listen functions. Retrieve does not count processed +// messages. +func (lsc *ListenerStatsCollector) AddProcessed(count int) { + atomic.AddUint64(&lsc.Processed, uint64(count)) +} + +// AddBatchSize records a count of the number of messages returned by +// GetRecords in the listener. +func (lsc *ListenerStatsCollector) AddBatchSize(count int) { + atomic.AddUint64(&lsc.BatchSize, uint64(count)) +} + +// AddGetRecordsCalled records the number of times the GetRecords API was called +// by the listener. +func (lsc *ListenerStatsCollector) AddGetRecordsCalled(count int) { + atomic.AddUint64(&lsc.GetRecordsCalled, uint64(count)) +} + +// AddProvisionedThroughputExceeded records the number of times the GetRecords +// API returned a ErrCodeProvisionedThroughputExceededException by the listener. +func (lsc *ListenerStatsCollector) AddProvisionedThroughputExceeded(count int) { + atomic.AddUint64(&lsc.ProvisionedThroughputExceeded, uint64(count)) +} + +// AddGetRecordsTimeout records the number of times the GetRecords API timed out +// on the HTTP level. This is influenced by the WithHTTPClientTimeout +// configuration. +func (lsc *ListenerStatsCollector) AddGetRecordsTimeout(count int) { + atomic.AddUint64(&lsc.GetRecordsTimeout, uint64(count)) +} + +// AddGetRecordsReadTimeout records the number of times the GetRecords API timed +// out while reading the response body. This is influenced by the +// WithGetRecordsReadTimeout configuration. +func (lsc *ListenerStatsCollector) AddGetRecordsReadTimeout(count int) { + atomic.AddUint64(&lsc.GetRecordsReadTimeout, uint64(count)) +} + +// AddProcessedDuration records the duration to process a record. See notes on +// AddProcessed. +func (lsc *ListenerStatsCollector) AddProcessedDuration(duration time.Duration) { + // TODO: Not threadsafe + lsc.ProcessedDuration = duration +} + +// AddGetRecordsDuration records the duration that the GetRecords API request +// took. Only the times of successful calls are measured. +func (lsc *ListenerStatsCollector) AddGetRecordsDuration(duration time.Duration) { + // TODO: Not threadsafe + lsc.GetRecordsDuration = duration + +} + +// AddGetRecordsReadResponseDuration records the duration that it took to read +// the response body of a GetRecords API request. +func (lsc *ListenerStatsCollector) AddGetRecordsReadResponseDuration(duration time.Duration) { + // TODO: Not threadsafe + lsc.GetRecordsReadResonseDuration = duration +} + +// AddGetRecordsUnmarshalDuration records the duration that it took to unmarshal +// the response body of a GetRecords API request. +func (lsc *ListenerStatsCollector) AddGetRecordsUnmarshalDuration(duration time.Duration) { + // TODO: Not threadsafe + lsc.GetRecordsUnmarshalDuration = duration +} + +// PrintStats logs stats atomically +func (lsc *ListenerStatsCollector) PrintStats() { + log.Printf("Listener stats: Consumed: [%d]\n", atomic.LoadUint64(&lsc.Consumed)) + log.Printf("Listener stats: Delivered: [%d]\n", atomic.LoadUint64(&lsc.Delivered)) + log.Printf("Listener stats: Processed: [%d]\n", atomic.LoadUint64(&lsc.Processed)) + log.Printf("Listener stats: Batch Size: [%d]\n", atomic.LoadUint64(&lsc.BatchSize)) + log.Printf("Listener stats: GetRecords Called: [%d]\n", atomic.LoadUint64(&lsc.GetRecordsCalled)) + log.Printf("Listener stats: Provisioned Throughput Exceeded: [%d]\n", atomic.LoadUint64(&lsc.ProvisionedThroughputExceeded)) + log.Printf("Listener stats: GetRecords Timeout: [%d]\n", atomic.LoadUint64(&lsc.GetRecordsTimeout)) + log.Printf("Listener stats: GetRecords Read Timeout: [%d]\n", atomic.LoadUint64(&lsc.GetRecordsReadTimeout)) +} \ No newline at end of file diff --git a/testexec/main.go b/testexec/main.go index 9e55a19..19e41b5 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -11,6 +11,7 @@ import ( "os/user" "strings" "sync" + "sync/atomic" "time" "github.com/aws/aws-sdk-go/aws" @@ -20,18 +21,41 @@ import ( "github.com/rewardStyle/kinetic/listener" "github.com/rewardStyle/kinetic/message" "github.com/rewardStyle/kinetic/producer" + + _ "net/http/pprof" + "net/http" +) + +// Define constants for Kinesis stream location +const ( + LocationLocal = "local" + LocationAws = "aws" +) + +// Define operation modes +const ( + ModeRead = "read" + ModeWrite = "write" + ModeReadWrite = "readwrite" ) // Declare global variables +var streamName string var startTime time.Time var streamStart time.Time var streamStop time.Time var pipeOfDeath chan os.Signal var stopDisplay chan struct{} -var stopProducing chan struct{} -var config *Config +var stopProduce chan struct{} +var stopListen chan struct{} +var cfg *Config func init() { + // Set up Http server for pprof + go func() { + log.Println(http.ListenAndServe("localhost:6060", nil)) + }() + // Start the stopwatch startTime = time.Now() @@ -41,25 +65,26 @@ func init() { // Instantiate channels for communicating between threads pipeOfDeath = make(chan os.Signal, 1) stopDisplay = make(chan struct{}, 1) - stopProducing = make(chan struct{}, 1) + stopProduce = make(chan struct{}, 1) + stopListen = make(chan struct{}, 1) // Set up pipeOfDeath channel to receive os signals signal.Notify(pipeOfDeath, os.Interrupt) } -func cleanup(k *kinetic.Kinetic, streamName string) { - if *config.Cleanup { - if *config.Verbose { +func cleanup(k *kinetic.Kinetic, stream string) { + if *cfg.Cleanup { + if *cfg.Verbose { log.Println() - log.Printf("Cleaning up by deleting stream [%s] ...\n", streamName) + log.Printf("Cleaning up by deleting stream [%s] ...\n", stream) } - k.DeleteStream(streamName) - if *config.Verbose { - log.Printf("Waiting for stream [%s] to be deleted ...\n", streamName) + k.DeleteStream(stream) + if *cfg.Verbose { + log.Printf("Waiting for stream [%s] to be deleted ...\n", stream) } - k.WaitUntilStreamDeleted(context.TODO(), streamName, + k.WaitUntilStreamDeleted(context.TODO(), stream, request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) - if *config.Verbose { + if *cfg.Verbose { log.Println("Clean up complete") log.Println() } @@ -68,28 +93,29 @@ func cleanup(k *kinetic.Kinetic, streamName string) { func main() { // Process command line arguments - config = parseCommandLineArgs() - config.printConfigs() + cfg = parseCommandLineArgs() + cfg.printConfigs() // Instantiate a new kinetic object based on the location of the kinesis stream var k *kinetic.Kinetic - switch strings.ToLower(*config.Location) { - case "local": + switch strings.ToLower(*cfg.Location) { + case LocationLocal: k = newDefaultKinetic() - case "aws": + case LocationAws: k = newAwsKinetic() default : - log.Fatalf("Unknown location for kinesis stream: %s\n", *config.Location) + log.Fatalf("Unknown location for kinesis stream: %s\n", *cfg.Location) } // Set streamName from config or generate a random one - streamName := *config.StreamName + streamName = *cfg.StreamName if streamName == "" { streamName = generateRandomStreamName() } // Create and wait for a new kinetic stream (if the stream name doesn't already exist) prepareKinesisStream(k, streamName) + defer cleanup(k, streamName) // Create a new kinetic producer p := newKineticProducer(k, streamName) @@ -100,131 +126,14 @@ func main() { // Instantiate StreamData Object to keep stats streamData := NewStreamData() - // Display stream data statistics to the console - displayWg := sync.WaitGroup{} - displayWg.Add(1) - go func(sd *StreamData) { - displayLoop: - for { - select { - case <-pipeOfDeath: - stopProducing <- struct{}{} - if *config.Verbose { - log.Println() - log.Println("display: Received pipeOfDeath ...") - } - case <-stopDisplay: - if *config.Verbose { - log.Println() - log.Print("display: Received stopDisplay ...") - } - break displayLoop - default: - time.Sleep(1000 * time.Millisecond) - log.Println() - log.Printf("Stream name: %s\n", streamName) - log.Printf("Elapsed Time: %v\n", time.Since(startTime)) - if streamStop.IsZero() { - log.Printf("Streaming Time: %v\n", time.Since(streamStart)) - } else { - log.Printf("Streaming Time: %v\n", streamStop.Sub(streamStart)) - } - sd.printStats() - } - } - streamData.printSummary() - displayWg.Done() - }(streamData) - - // Use the producer to write messages to the kinetic stream - streamWg := sync.WaitGroup{} - streamWg.Add(1) - go func(sd *StreamData) { - streamStart = time.Now() - if config.Duration != nil { - // Set a timeout based on the config - var timeout <-chan time.Time - if *config.Duration > 0 { - timeout = time.After(time.Duration(*config.Duration) * time.Second) - } else { - timeout = make(chan time.Time, 1) - } - - produceLoop: - for { - select { - case <-stopProducing: - if *config.Verbose { - log.Println() - log.Print("producer: Received stopProducing ...") - } - streamStop = time.Now() - break produceLoop - case <-timeout: - if *config.Verbose { - log.Println() - log.Print("producer: Timed out ...") - } - streamStop = time.Now() - break produceLoop - default: - jsonStr, _ := json.Marshal(NewMessage()) - if err := p.Send(&message.Message { - PartitionKey: aws.String("key"), - Data: []byte(jsonStr), - }); err == nil { - streamWg.Add(1) - sd.incrementMsgCount() - } - } - } - } else if config.NumMsgs != nil { - for i := 0; i < *config.NumMsgs; i++ { - jsonStr, _ := json.Marshal(NewMessage()) - if err := p.Send(&message.Message { - PartitionKey: aws.String("key"), - Data: []byte(jsonStr), - }); err == nil { - streamWg.Add(1) - sd.incrementMsgCount() - } - } - streamStop = time.Now() - } - streamWg.Done() - }(streamData) - - // Use the listener to read messages from the kinetic stream - go func(sd *StreamData) { - l.Listen(func(m *message.Message, fnwg *sync.WaitGroup) error { - // Unmarshal data - msg := &Message{} - json.Unmarshal(m.Data, msg) - - // Only mark "done" if the message isn't a duplicate - if !sd.exists(msg.ID) { - streamWg.Done() - } else { - if *config.Verbose { - log.Printf("listner: Duplicate message: %v\n", msg) - } - } - - // Record message regardless if it is a duplicate - sd.mark(msg.ID) - fnwg.Done() - - return nil - }) - }(streamData) - - // Wait until streaming is complete - streamWg.Wait() - stopDisplay <- struct{}{} - - // Wait until output display is complete - displayWg.Wait() - cleanup(k, streamName) + // Run all the things concurrently + mainWg := sync.WaitGroup{} + mainWg.Add(3) // Wait for display, produce and listen + go handlePoD() + go display(streamData, p, l, &mainWg) + go produce(streamData, p, &mainWg) + go listen(streamData, l, &mainWg) + mainWg.Wait() } func newDefaultKinetic() *kinetic.Kinetic { @@ -289,7 +198,7 @@ func generateRandomStreamName() string { } func prepareKinesisStream(k *kinetic.Kinetic, streamName string) { - if *config.Verbose { + if *cfg.Verbose { log.Printf("Preparing kinesis stream: [%s] ...\n", streamName) } @@ -298,7 +207,7 @@ func prepareKinesisStream(k *kinetic.Kinetic, streamName string) { request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) if err != nil { // If not, create a kinetic stream - if *config.Verbose { + if *cfg.Verbose { log.Printf("Creating a new kinesis stream: [%s] ...\n", streamName) } err := k.CreateStream(streamName, 1) @@ -307,8 +216,8 @@ func prepareKinesisStream(k *kinetic.Kinetic, streamName string) { } // And wait until the stream is ready to go - if *config.Verbose { - log.Printf("Waiting for new kinesis stream: [%s]\n", streamName) + if *cfg.Verbose { + log.Printf("Waiting for new kinesis stream: [%s] ...\n", streamName) } err = k.WaitUntilStreamExists(context.TODO(), streamName, request.WithWaiterDelay(request.ConstantWaiterDelay(2*time.Second))) @@ -319,13 +228,15 @@ func prepareKinesisStream(k *kinetic.Kinetic, streamName string) { } func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Producer { - if *config.Verbose { + if *cfg.Verbose { log.Println("Creating a kinetic producer ...") } + psc := new(ProducerStatsCollector) w, err := producer.NewKinesisWriter(k.Session.Config, streamName, func(kwc *producer.KinesisWriterConfig) { kwc.SetLogLevel(aws.LogDebug) kwc.SetResponseReadTimeout(time.Second) + kwc.SetStatsCollector(psc) }) if err != nil { log.Fatalf("Unable to create a new kinesis stream writer due to: %v\n", err) @@ -334,6 +245,7 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce p, err := producer.NewProducer(k.Session.Config, w, func(c *producer.Config) { c.SetBatchSize(500) c.SetBatchTimeout(1000 * time.Millisecond) + c.SetStatsCollector(psc) }) if err != nil { log.Fatalf("Unable to create a new producer due to: %v\n", err) @@ -343,7 +255,7 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce } func newKineticListener(k *kinetic.Kinetic, streamName string) *listener.Listener { - if *config.Verbose { + if *cfg.Verbose { log.Println("Creating a kinetic listener ...") } @@ -353,14 +265,20 @@ func newKineticListener(k *kinetic.Kinetic, streamName string) *listener.Listene log.Fatalf("Unable to get shards for stream %s due to: %v\n", streamName, err) } + lsc := new(ListenerStatsCollector) r, err := listener.NewKinesisReader(k.Session.Config, streamName, shards[0], func(krc *listener.KinesisReaderConfig) { krc.SetResponseReadTimeout(1000 * time.Millisecond) + krc.SetStatsCollector(lsc) }) + if err != nil { + log.Fatalf("Unable to create a new kinesis reader due to: %v\n", err) + } l, err := listener.NewListener(k.Session.Config, r, func(c *listener.Config) { - c.SetQueueDepth(20) + c.SetQueueDepth(500) c.SetConcurrency(10) + c.SetStatsCollector(lsc) }) if err != nil { log.Fatalf("Unable to create a new listener due to: %v\n", err) @@ -368,3 +286,254 @@ func newKineticListener(k *kinetic.Kinetic, streamName string) *listener.Listene return l } + +func handlePoD() { + <-pipeOfDeath + if *cfg.Verbose { + log.Println() + log.Println("Received pipeOfDeath ...") + } + if *cfg.Mode == ModeRead { + stopListen <- struct{}{} + } else { + stopProduce <- struct{}{} + } +} + +func display(sd *StreamData, p *producer.Producer, l *listener.Listener, wg *sync.WaitGroup) { + defer wg.Done() + + for { + select { + case <-stopDisplay: + if *cfg.Verbose { + log.Println() + log.Print("display: Received stopDisplay ...") + } + log.Println() + log.Println("***** Stream Data Summary *****") + if *cfg.Mode != ModeRead { + p.Stats.(*ProducerStatsCollector).PrintStats() + } + if *cfg.Mode != ModeWrite { + l.Stats.(*ListenerStatsCollector).PrintStats() + sd.printSummary() + } + return + case <-time.After(time.Second): + log.Println() + log.Printf("Stream name: %s\n", streamName) + log.Printf("Elapsed Time: %v\n", time.Since(startTime)) + if !streamStart.IsZero() && !streamStop.IsZero() { + log.Printf("Streaming Time: %v\n", streamStop.Sub(streamStart)) + } + log.Println() + log.Println("***** Stream Data Stats *****") + if *cfg.Mode != ModeRead { + p.Stats.(*ProducerStatsCollector).PrintStats() + } + if *cfg.Mode != ModeWrite { + l.Stats.(*ListenerStatsCollector).PrintStats() + sd.printStats() + } + } + } +} + +func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { + defer wg.Done() + + // Return early if we don't need to produce + if *cfg.Mode == ModeRead { + return + } + + defer func() { + if *cfg.Verbose { + log.Println() + log.Println("producer: Exiting produce ...") + } + + // In write mode the producer controls when to stop displaying + if *cfg.Mode == ModeWrite { + stopDisplay <- struct{}{} + } + }() + + // Define a timeout channel if the duration is set + streamStart = time.Now() + var timeout <-chan time.Time + if cfg.Duration != nil && *cfg.Duration > 0 { + timeout = time.After(time.Duration(*cfg.Duration) * time.Second) + } else { + timeout = make(chan time.Time, 1) + } + + var sendSignal = make(chan struct{}, 1) + go func() { + for { + <-sendSignal + jsonStr, _ := json.Marshal(NewMessage()) + if err := p.Send(&message.Message{ + PartitionKey: aws.String("key"), + Data: []byte(jsonStr), + }); err == nil { + sd.incrementMsgCount() + } else { + log.Println("producer: Uh oh, something bad happened!!!!") + } + } + }() + + // Control when to exit produce + produceWg := sync.WaitGroup{} + produceWg.Add(1) + go func(){ + defer produceWg.Done() + + var sent uint64 + var sendTicker *time.Ticker + if *cfg.Throttle { + sendTicker = time.NewTicker(time.Millisecond) + } else { + sendTicker = time.NewTicker(time.Nanosecond) + } + produce: + for { + select { + case <-stopProduce: + if *cfg.Verbose { + log.Println() + log.Println("produce: Received stop produce ...") + } + break produce + case <-timeout: + if *cfg.Verbose { + log.Println() + log.Print("producer: Duration time out ...") + } + break produce + case <-sendTicker.C: + // Break from the loop if we have sent the correct number of messages + if cfg.NumMsgs != nil { + if atomic.LoadUint64(&sent) >= uint64(*cfg.NumMsgs) { + break produce + } + } + sendSignal <- struct{}{} + atomic.AddUint64(&sent, 1) + } + } + + // We may need to wait for Send to finish so we add a delay before exiting produce + var staleTimeout time.Duration + switch strings.ToLower(*cfg.Location) { + case LocationLocal: + staleTimeout = time.Duration(2 * time.Second) + case LocationAws: + staleTimeout = time.Duration(60 * time.Second) + } + staleTime := time.NewTimer(staleTimeout) + + for { + select { + case <-staleTime.C: + if *cfg.Verbose { + log.Println() + log.Println("producer: No more outgoing messages from producer ...") + } + return + case <-time.After(time.Second): + newSent := atomic.LoadUint64(&p.Stats.(*ProducerStatsCollector).Sent); + if sent != newSent { + staleTime.Reset(staleTimeout) + sent = newSent + } + } + } + }() + produceWg.Wait() +} + +func listen(sd *StreamData, l *listener.Listener, wg *sync.WaitGroup) { + defer wg.Done() + + // Return early if we don't need to produce + if *cfg.Mode == ModeWrite { + return + } + + defer func() { + if *cfg.Verbose { + log.Println() + log.Println("listener: Exiting listen ...") + } + + // In read and readwrite mode the listener controls when to stop displaying + stopDisplay <- struct{}{} + }() + + // Call Listen within a go routine + go func(){ + l.Listen(func(m *message.Message, wg *sync.WaitGroup) error { + defer wg.Done() + + // Unmarshal data + msg := &Message{} + json.Unmarshal(m.Data, msg) + + // Only mark "done" if the message isn't a duplicate + if sd.exists(msg.ID) { + if *cfg.Verbose { + log.Printf("listener: Duplicate message: %v\n", msg) + } + } + + // Record message regardless if it is a duplicate + sd.mark(msg.ID) + + return nil + }) + }() + + // Control when to exit listen + listenWg := sync.WaitGroup{} + listenWg.Add(1) + go func(){ + defer listenWg.Done() + + var staleTimeout time.Duration + switch strings.ToLower(*cfg.Location) { + case LocationLocal: + staleTimeout = time.Duration(10 * time.Second) + case LocationAws: + staleTimeout = time.Duration(60 * time.Second) + } + staleTime := time.NewTimer(staleTimeout) + + var consumed uint64 + for { + select { + case <-stopListen: + if *cfg.Verbose { + log.Println() + log.Println("listener: Received stop listen ...") + } + return + case <-staleTime.C: + if *cfg.Verbose { + log.Println() + log.Println("listener: No more incoming messages from listener ...") + } + return + case <-time.After(time.Second): + newConsumed := atomic.LoadUint64(&l.Stats.(*ListenerStatsCollector).Consumed); + if consumed != newConsumed { + staleTime.Reset(staleTimeout) + consumed = newConsumed + } + } + } + }() + listenWg.Wait() +} diff --git a/testexec/producerstats.go b/testexec/producerstats.go new file mode 100644 index 0000000..83fec67 --- /dev/null +++ b/testexec/producerstats.go @@ -0,0 +1,117 @@ +package main + +import ( + "sync/atomic" + "time" + "log" +) + +// ProducerStatsCollector implements the producer's StatsCollector +type ProducerStatsCollector struct { + Sent uint64 + Failed uint64 + DroppedTotal uint64 + DroppedCapacity uint64 + DroppedRetries uint64 + BatchSize uint64 + PutRecordsProvisionedThroughputExceeded uint64 + PutRecordsCalled uint64 + ProvisionedThroughputExceeded uint64 + PutRecordsTimeout uint64 + + PutRecordsDuration time.Duration + PutRecordsBuildDuration time.Duration + PutRecordsSendDuration time.Duration +} + +// GetSent returns the number of messages sent to AWS Kinesis by the producer. +func (psc *ProducerStatsCollector) GetSent() uint64 { + return atomic.LoadUint64(&psc.Sent) +} + +// AddSent records a count of the number of messages sent to AWS Kinesis by the producer. +func (psc *ProducerStatsCollector) AddSent(count int) { + atomic.AddUint64(&psc.Sent, uint64(count)) +} + +// AddFailed records a count of the number of messages that failed to be sent to AWS Kinesis by the producer. +func (psc *ProducerStatsCollector) AddFailed(count int) { + atomic.AddUint64(&psc.Failed, uint64(count)) +} + +// AddDroppedTotal records a count of the total number of messages dropped by the application after multiple failures. +func (psc *ProducerStatsCollector) AddDroppedTotal(count int) { + atomic.AddUint64(&psc.DroppedTotal, uint64(count)) +} + +// AddDroppedCapacity records a count of the number of messages that were dropped by the application due to the stream +// writer being at capacity. +func (psc *ProducerStatsCollector) AddDroppedCapacity(count int) { + atomic.AddUint64(&psc.DroppedCapacity, uint64(count)) +} + +// AddDroppedRetries records a count of the number of retry messages dropped by the application after the max number of +// retries was exceeded. +func (psc *ProducerStatsCollector) AddDroppedRetries(count int) { + atomic.AddUint64(&psc.DroppedRetries, uint64(count)) +} + +// AddBatchSize records a count of the number of messages attempted by PutRecords in the producer. +func (psc *ProducerStatsCollector) AddBatchSize(count int) { + atomic.AddUint64(&psc.BatchSize, uint64(count)) +} + +// AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords API returned a +// ErrCodeProvisionedThroughputExceededException by the producer. +func (psc *ProducerStatsCollector) AddPutRecordsProvisionedThroughputExceeded(count int) { + atomic.AddUint64(&psc.PutRecordsProvisionedThroughputExceeded, uint64(count)) +} + +// AddPutRecordsCalled records the number of times the PutRecords API was called by the producer. +func (psc *ProducerStatsCollector) AddPutRecordsCalled(count int) { + atomic.AddUint64(&psc.PutRecordsCalled, uint64(count)) +} + +// AddProvisionedThroughputExceeded records the number of times the PutRecords API response contained a record which +// contained an ErrCodeProvisionedThroughputExceededException error. +func (psc *ProducerStatsCollector) AddProvisionedThroughputExceeded(count int) { + atomic.AddUint64(&psc.ProvisionedThroughputExceeded, uint64(count)) +} + +// AddPutRecordsTimeout records the number of times the PutRecords API timed out on the HTTP level. This is influenced +// by the WithHTTPClientTimeout configuration. +func (psc *ProducerStatsCollector) AddPutRecordsTimeout(count int) { + atomic.AddUint64(&psc.PutRecordsTimeout, uint64(count)) +} + +// AddPutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls +// are measured. +func (psc *ProducerStatsCollector) AddPutRecordsDuration(duration time.Duration) { + // TODO: Not threadsafe + psc.PutRecordsDuration = duration +} + +// AddPutRecordsBuildDuration records the duration that it took to build the PutRecords API request payload. +func (psc *ProducerStatsCollector) AddPutRecordsBuildDuration(duration time.Duration) { + // TODO: Not threadsafe + psc.PutRecordsBuildDuration = duration +} + +// AddPutRecordsSendDuration records the duration that it took to send the PutRecords API request payload. +func (psc *ProducerStatsCollector) AddPutRecordsSendDuration(duration time.Duration) { + // TODO: Not threadsafe + psc.PutRecordsSendDuration = duration +} + +// PrintStats logs stats atomically +func (psc *ProducerStatsCollector) PrintStats() { + log.Printf("Producer Stats: Sent: [%d]\n", atomic.LoadUint64(&psc.Sent)) + log.Printf("Producer Stats: Failed: [%d]\n", atomic.LoadUint64(&psc.Failed)) + log.Printf("Producer Stats: Dropped Total: [%d]\n", atomic.LoadUint64(&psc.DroppedTotal)) + log.Printf("Producer Stats: Dropped Retries: [%d]\n", atomic.LoadUint64(&psc.DroppedRetries)) + log.Printf("Producer Stats: Dropped Capacity: [%d]\n", atomic.LoadUint64(&psc.DroppedCapacity)) + log.Printf("Producer Stats: Batch Size: [%d]\n", atomic.LoadUint64(&psc.BatchSize)) + log.Printf("Producer Stats: PutRecords Called: [%d]\n", atomic.LoadUint64(&psc.PutRecordsCalled)) + log.Printf("Producer Stats: Provisioned Throughput Exceeded: [%d]\n", atomic.LoadUint64(&psc.ProvisionedThroughputExceeded)) + log.Printf("Producer Stats: PutRecords Timeout: [%d]\n", atomic.LoadUint64(&psc.PutRecordsTimeout)) +} diff --git a/testexec/stats.go b/testexec/streamdata.go similarity index 69% rename from testexec/stats.go rename to testexec/streamdata.go index 610229f..cdf4e28 100644 --- a/testexec/stats.go +++ b/testexec/streamdata.go @@ -68,25 +68,17 @@ func (m *StreamData) printStats() { m.mutex.Lock() defer m.mutex.Unlock() - log.Println() - log.Println("***** Stream Data Stats *****") - log.Printf("Messages sent: [%d]\n", m.MsgCount) - log.Printf("Messages received: [%d]\n", len(m.Frequencies)) - log.Printf("Number of duplicated messages: [%d]\n", m.Duplicates) + log.Printf("Stream Data: Number Duplicated: [%d]\n", m.Duplicates) } func (m *StreamData) printSummary() { m.mutex.Lock() defer m.mutex.Unlock() - log.Println() - log.Println("***** Stream Data Summary *****") - log.Printf("Total messages sent: [%d]\n", m.MsgCount) - log.Printf("Total messages received: [%d]\n", len(m.Frequencies)) - log.Printf("Total Number of duplicated messages: [%d]\n", m.Duplicates) + log.Printf("Stream Data: Total Number of duplicated messages: [%d]\n", m.Duplicates) for index, freq := range m.Frequencies { if freq > 1 { - log.Printf("Message [%d] was received [%d] times\n", index, freq) + log.Printf("Stream Data: Message [%d] was received [%d] times\n", index, freq) } } } From 8ad9b48b9b23e7f5712467327fd04ab4053e0a8d Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 6 Jul 2017 15:04:20 -0500 Subject: [PATCH 045/109] Make attempts counter in producer "thread safe". Changes to testexec for output logging, command line interface and documentation. --- producer/producer.go | 6 +++--- testexec/README.md | 33 ++++++++++++++++++++++----------- testexec/cli.go | 36 +++++++++++++++++++++--------------- testexec/main.go | 22 +++++++++++++++------- 4 files changed, 61 insertions(+), 36 deletions(-) diff --git a/producer/producer.go b/producer/producer.go index 6135b38..5bb87c2 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -110,7 +110,7 @@ func (p *Producer) sendBatch(batch []*message.Message) { p.shutdownCond.L.Unlock() }() - var attempts int + var attempts uint64 var failed uint64 stop: @@ -177,13 +177,13 @@ stop: // batch to be retried rather than retrying the batch as-is. With this approach, we can kill the "stop" // for loop, and set the entire batch to retries to allow the below code to handle retrying the // messages. - if attempts > p.maxRetryAttempts { + if atomic.LoadUint64(&attempts) > uint64(p.maxRetryAttempts) { p.LogError(fmt.Sprintf("Dropping batch after %d failed attempts to deliver to stream", attempts)) p.Stats.AddDroppedTotal(len(batch)) p.Stats.AddDroppedRetries(len(batch)) break stop } - attempts++ + atomic.AddUint64(&attempts, 1) // Apply a delay before retrying time.Sleep(time.Duration(attempts * attempts) * time.Second) diff --git a/testexec/README.md b/testexec/README.md index 149627a..22c5e87 100644 --- a/testexec/README.md +++ b/testexec/README.md @@ -1,11 +1,13 @@ # testexec -testexec is an executable CLI program to test the kinetic producer / lister. - -The executable program creates a kinetic object connecting to either a local kinesalte instance or to AWS Kinesis using local AWS credentials/config. +testexec is an executable CLI program to test the kinetic producer / listener. +The executable program creates a kinetic object connecting to either a local kinesalte instance or to AWS Kinesis using local AWS credentials/config. + A new stream will be created (with a random stream name) unless a stream name is provided (See usage). -A kinetic producer will stream dummy data to the kinesis stream and a kinetic listener will stream from the kinesis stream. +There are three modes of operations permitted: 'read', 'write' and 'readwrite'. Write mode is produce only where as read mode is consume only. Readwrite mode is both produce and consume concurrently. + +So depending on the mode of operation, a kinetic producer will stream dummy data to the kinesis stream and a kinetic listener will stream from the kinesis stream. The program logs stream data stats to the console log periodically. @@ -32,17 +34,21 @@ go install ```text Usage of ./testexec: -cleanup - used to specify whether or not to delete a newly created kinesis stream (default true) + used to specify whether or not to delete the kinesis stream after processing is complete. (default true) -duration int - used to specify the duration (in seconds) the program should run. Use a value of -1 to run indefinitely. Either -num-msgs or -duration must be set. + used to specify the duration (in seconds) the program should run. This flag is only applicable to 'write' and 'readwrite' modes. Use zero or negative number to run indefinitely. (default 0) -location string - used to specify the location of the kinesis stream. Accepted values are (local|aws). For local, run kinesalite on http://127.0.0.1:4567. For aws, your aws credentials and configuration needs to be defined at ~/.aws (default "local") + used to specify the location of the kinesis stream. Accepted values are (local|aws). For local, run kinesalite on http://127.0.0.1:4567. For aws, your aws credentials and configuration need to be defined at ~/.aws (default "local") + -mode string + used to specify the mode in which to run; either 'r', 'read', 'w', 'write', 'rw' or 'readwrite' (default "readwrite") -num-msgs int - used to specify the number of messages to (attempt to) send / receive. Either -num-msgs or -duration must be set. + used to specify the number of messages to (attempt to) send. This flag is only applicable to 'write' and 'readwrite' modes. Use zero or a negative number to produce indefinitely (default 0) -stream-name string used to specify a pre-existing stream to be used for testing. A new stream will be created if not defined. + -throttle + used to specify whether to throttle PutRecord requests by 1 ms. (default true) -verbose - used to specify whether or not to log in verbose mode + used to specify whether or not to log in verbose mode. (default false) ``` ## Examples @@ -57,8 +63,13 @@ To run kinetic testexec on a local kinesalite instance to stream for a fixed dur ./testexec -location local -stream-name some-stream -duration 1000 ``` -To run kinetic testexec on an AWS Kinesis Stream to stream indefinitely: +To run kinetic testexec on an AWS Kinesis Stream to stream indefinitely (Ctrl-C to stop producing): ```sh -./testexec -location aws -duration -1 -cleanup=false -verbose +./testexec -location aws -duration -1 -mode write -stream-name test-stream -cleanup=false -verbose ``` +To run kinetic testexec in read mode an AWS KinesisStream: +To run kinetic testexec on an AWS Kinesis Stream to stream indefinitely (Ctrl-C to stop producing): +```sh +./testexec -location aws -mode read -stream-name test-stream -cleanup=false -verbose +``` diff --git a/testexec/cli.go b/testexec/cli.go index 968809d..283c96d 100644 --- a/testexec/cli.go +++ b/testexec/cli.go @@ -28,31 +28,29 @@ func parseCommandLineArgs() *Config { "aws credentials and configuration need to be defined at ~/.aws") streamNamePtr := flag.String("stream-name", "", "used to specify a pre-existing stream to be used for " + "testing. A new stream will be created if not defined.") - numMsgsPtr := flag.Int("num-msgs", 0, "used to specify the number of messages to (attempt to) send / " + - "receive. Either -num-msgs or -duration must be set.") + numMsgsPtr := flag.Int("num-msgs", 0, "used to specify the number of messages to (attempt to) send. This " + + "flag is only applicable to 'write' and 'readwrite' modes. Use zero or a negative number to produce " + + "indefinitely") durationPtr := flag.Int("duration", 0, "used to specify the duration (in seconds) the program should run. " + - "Use a value of -1 to run indefinitely. Either -num-msgs or -duration must be set.") + "This flag is only applicable to 'write' and 'readwrite' modes. Use zero or negative number to run " + + "indefinitely.") throttlePtr := flag.Bool("throttle", true, "used to specify whether to throttle PutRecord requests by 1 ms. ") - cleanupPtr := flag.Bool("cleanup", true, "used to specify whether or not to delete a newly created kinesis " + - "stream") + cleanupPtr := flag.Bool("cleanup", true, "used to specify whether or not to delete the kinesis stream after " + + "processing is complete.") verbosePtr := flag.Bool("verbose", false, "used to specify whether or not to log in verbose mode") // Parse command line arguments flag.Parse() // Process command line arguments - if *numMsgsPtr == 0 && *durationPtr == 0 { - log.Fatal("Either -num-msgs or -duration must be set.") - } else if *numMsgsPtr != 0 && *durationPtr != 0 { - log.Fatal("Both -num-msgs and -duration were set. Only one may be set.") - } else if *durationPtr != 0 { - numMsgsPtr = nil - } else if *numMsgsPtr < 0 { - log.Fatal("Number of messages value must be greater than 0") - } else { + if *durationPtr <= 0 { durationPtr = nil } + if *numMsgsPtr <= 0 { + numMsgsPtr = nil + } + var mode string switch strings.ToLower(*modePtr) { case "r": @@ -88,12 +86,20 @@ func (c *Config) printConfigs() { log.Println("Command Line Arguments:") log.Println("-mode: ", *c.Mode) log.Println("-location: ", *c.Location) - log.Println("-stream-name: ", *c.StreamName) + if len(*c.StreamName) == 0 { + log.Println("-stream-name: (randomly generated)") + } else { + log.Println("-stream-name: ", *c.StreamName) + } if c.NumMsgs != nil { log.Println("-num-msgs: ", *c.NumMsgs) + } else { + log.Println("-num-msgs: (unbounded)") } if c.Duration != nil { log.Println("-duration: ", *c.Duration) + } else { + log.Println("-duration: (indefinite)") } log.Println("-throttle: ", *c.Throttle) log.Println("-cleanup: ", *c.Cleanup) diff --git a/testexec/main.go b/testexec/main.go index 19e41b5..7505555 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -245,6 +245,9 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce p, err := producer.NewProducer(k.Session.Config, w, func(c *producer.Config) { c.SetBatchSize(500) c.SetBatchTimeout(1000 * time.Millisecond) + c.SetConcurrency(10) + c.SetMaxRetryAttempts(2) + c.SetQueueDepth(100) c.SetStatsCollector(psc) }) if err != nil { @@ -291,7 +294,7 @@ func handlePoD() { <-pipeOfDeath if *cfg.Verbose { log.Println() - log.Println("Received pipeOfDeath ...") + log.Println("main: Received pipeOfDeath ...") } if *cfg.Mode == ModeRead { stopListen <- struct{}{} @@ -311,6 +314,12 @@ func display(sd *StreamData, p *producer.Producer, l *listener.Listener, wg *syn log.Print("display: Received stopDisplay ...") } log.Println() + log.Printf("Stream name: %s\n", streamName) + log.Printf("Elapsed Time: %v\n", time.Since(startTime)) + if !streamStart.IsZero() && !streamStop.IsZero() { + log.Printf("Streaming Time: %v\n", streamStop.Sub(streamStart)) + } + log.Println() log.Println("***** Stream Data Summary *****") if *cfg.Mode != ModeRead { p.Stats.(*ProducerStatsCollector).PrintStats() @@ -324,9 +333,6 @@ func display(sd *StreamData, p *producer.Producer, l *listener.Listener, wg *syn log.Println() log.Printf("Stream name: %s\n", streamName) log.Printf("Elapsed Time: %v\n", time.Since(startTime)) - if !streamStart.IsZero() && !streamStop.IsZero() { - log.Printf("Streaming Time: %v\n", streamStop.Sub(streamStart)) - } log.Println() log.Println("***** Stream Data Stats *****") if *cfg.Mode != ModeRead { @@ -369,6 +375,7 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { timeout = make(chan time.Time, 1) } + // Run Send in a separate go routine listening for the sendSignal var sendSignal = make(chan struct{}, 1) go func() { for { @@ -404,7 +411,7 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { case <-stopProduce: if *cfg.Verbose { log.Println() - log.Println("produce: Received stop produce ...") + log.Println("producer: Received stop produce ...") } break produce case <-timeout: @@ -424,6 +431,7 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { atomic.AddUint64(&sent, 1) } } + streamStop = time.Now() // We may need to wait for Send to finish so we add a delay before exiting produce var staleTimeout time.Duration @@ -431,7 +439,7 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { case LocationLocal: staleTimeout = time.Duration(2 * time.Second) case LocationAws: - staleTimeout = time.Duration(60 * time.Second) + staleTimeout = time.Duration(10 * time.Second) } staleTime := time.NewTimer(staleTimeout) @@ -505,7 +513,7 @@ func listen(sd *StreamData, l *listener.Listener, wg *sync.WaitGroup) { var staleTimeout time.Duration switch strings.ToLower(*cfg.Location) { case LocationLocal: - staleTimeout = time.Duration(10 * time.Second) + staleTimeout = time.Duration(3 * time.Second) case LocationAws: staleTimeout = time.Duration(60 * time.Second) } From eda4de4f6be374de9dc5980d77eeecda6ac3945e Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 6 Jul 2017 15:36:08 -0500 Subject: [PATCH 046/109] Run gofmt. --- kinetic_integration_test.go | 15 ++++++------ listener/kclreader.go | 11 ++++----- listener/kclreaderconfig.go | 2 +- listener/kinesisreader.go | 8 +++---- listener/kinesisreaderconfig.go | 8 +++---- listener/listener.go | 14 +++++++---- listener/listenerconfig.go | 2 +- listener/sharditerator_test.go | 2 +- multilang/action.go | 40 ++++++++++++++++---------------- producer/firehosewriter.go | 2 +- producer/firehosewriterconfig.go | 2 +- producer/kinesiswriter.go | 2 +- producer/kinesiswriterconfig.go | 4 ++-- producer/producer.go | 10 ++++---- producer/producer_test.go | 6 ++--- producer/producerconfig.go | 2 +- producer/producerconfig_test.go | 6 ++--- testexec/cli.go | 32 ++++++++++++------------- testexec/listenerstats.go | 6 ++--- testexec/main.go | 24 +++++++++---------- testexec/messages.go | 4 ++-- testexec/producerstats.go | 8 +++---- testexec/streamdata.go | 2 +- 23 files changed, 107 insertions(+), 105 deletions(-) diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index 3e0f30e..548e1f0 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -1,8 +1,8 @@ package kinetic import ( - "encoding/json" "context" + "encoding/json" "log" "math/rand" "strconv" @@ -31,9 +31,9 @@ type StreamData struct { func NewStreamData() *StreamData { return &StreamData{ - mutex: sync.Mutex{}, + mutex: sync.Mutex{}, Frequencies: make(map[int]int), - Messages: make(map[int][]string), + Messages: make(map[int][]string), } } @@ -151,7 +151,7 @@ func TestKineticIntegration(t *testing.T) { r, err := listener.NewKinesisReader(k.Session.Config, stream, shards[0], func(krc *listener.KinesisReaderConfig) { krc.SetResponseReadTimeout(time.Second) - }) + }) assert.NotNil(t, r) assert.NoError(t, err) @@ -166,7 +166,6 @@ func TestKineticIntegration(t *testing.T) { numMsg := 1000 numSent := 0 - // Use the producer to write messages to the kinetic stream wg := sync.WaitGroup{} wg.Add(numMsg + 1) @@ -174,13 +173,13 @@ func TestKineticIntegration(t *testing.T) { defer wg.Done() for i := 0; i < numMsg; i++ { msg := &Message{ - ID: i, + ID: i, Message: "hello_" + strconv.Itoa(i), } jsonStr, _ := json.Marshal(msg) - if err := p.Send(&message.Message { + if err := p.Send(&message.Message{ PartitionKey: aws.String("key"), - Data: []byte(jsonStr), + Data: []byte(jsonStr), }); err == nil { *sent++ } diff --git a/listener/kclreader.go b/listener/kclreader.go index 2df9e0a..5e4f9d6 100644 --- a/listener/kclreader.go +++ b/listener/kclreader.go @@ -10,10 +10,10 @@ import ( "os" "sync" + "github.com/aws/aws-sdk-go/aws" + "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" "github.com/rewardStyle/kinetic/multilang" - "github.com/rewardStyle/kinetic/logging" - "github.com/aws/aws-sdk-go/aws" ) type kclReaderOptions struct { @@ -21,7 +21,6 @@ type kclReaderOptions struct { onCheckpointCallbackFn func() error onShutdownCallbackFn func() error Stats StatsCollector - } // KclReader handles the KCL Multilang Protocol to read records from KCL @@ -46,10 +45,10 @@ func NewKclReader(c *aws.Config, fn ...func(*KclReaderConfig)) (*KclReader, erro kclReaderOptions: cfg.kclReaderOptions, LogHelper: &logging.LogHelper{ LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, + Logger: cfg.AwsConfig.Logger, }, throttleSem: make(chan empty, 5), - msgBuffer: []message.Message{}, + msgBuffer: []message.Message{}, }, nil } @@ -59,7 +58,7 @@ func NewKclReader(c *aws.Config, fn ...func(*KclReaderConfig)) (*KclReader, erro // has been received / processed func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, error) { // Define the batchSize - batchSize := 0; + batchSize := 0 if len(r.msgBuffer) > 0 { if numRecords < 0 { batchSize = len(r.msgBuffer) diff --git a/listener/kclreaderconfig.go b/listener/kclreaderconfig.go index 11a54fc..a4b0f43 100644 --- a/listener/kclreaderconfig.go +++ b/listener/kclreaderconfig.go @@ -8,7 +8,7 @@ import ( type KclReaderConfig struct { *kclReaderOptions AwsConfig *aws.Config - LogLevel aws.LogLevelType + LogLevel aws.LogLevelType } // NewKclReaderConfig creates a new instance of KclReaderConfig diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index e3cdb5b..dc0f5bd 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -53,12 +53,12 @@ func NewKinesisReader(c *aws.Config, stream string, shard string, fn ...func(*Ki kinesisReaderOptions: cfg.kinesisReaderOptions, LogHelper: &logging.LogHelper{ LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, + Logger: cfg.AwsConfig.Logger, }, - stream: stream, - shard: shard, + stream: stream, + shard: shard, throttleSem: make(chan empty, 5), - client: kinesis.New(sess), + client: kinesis.New(sess), }, nil } diff --git a/listener/kinesisreaderconfig.go b/listener/kinesisreaderconfig.go index e8bbd35..7531313 100644 --- a/listener/kinesisreaderconfig.go +++ b/listener/kinesisreaderconfig.go @@ -10,7 +10,7 @@ import ( type KinesisReaderConfig struct { *kinesisReaderOptions AwsConfig *aws.Config - LogLevel aws.LogLevelType + LogLevel aws.LogLevelType } // NewKinesisReaderConfig creates a new instance of KinesisReaderConfig @@ -18,10 +18,10 @@ func NewKinesisReaderConfig(cfg *aws.Config) *KinesisReaderConfig { return &KinesisReaderConfig{ AwsConfig: cfg, kinesisReaderOptions: &kinesisReaderOptions{ - batchSize: 10000, - shardIterator: NewShardIterator(), + batchSize: 10000, + shardIterator: NewShardIterator(), responseReadTimeout: time.Second, - Stats: &NilStatsCollector{}, + Stats: &NilStatsCollector{}, }, LogLevel: *cfg.LogLevel, } diff --git a/listener/listener.go b/listener/listener.go index d359a6c..0140474 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -64,9 +64,9 @@ func NewListener(c *aws.Config, r StreamReader, fn ...func(*Config)) (*Listener, listenerOptions: cfg.listenerOptions, LogHelper: &logging.LogHelper{ LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, + Logger: cfg.AwsConfig.Logger, }, - reader: r, + reader: r, concurrencySem: make(chan empty, cfg.concurrency), pipeOfDeath: make(chan empty), }, nil @@ -179,7 +179,9 @@ func (l *Listener) consume(ctx context.Context) { // We need to run startConsuming to make sure that we are okay and ready to start consuming. This is mainly to // avoid a race condition where Listen() will attempt to read the messages channel prior to consume() // initializing it. We can then launch a goroutine to handle the actual consume operation. - if !l.startConsuming() { return } + if !l.startConsuming() { + return + } go func() { defer l.stopConsuming() @@ -201,7 +203,7 @@ func (l *Listener) consume(ctx context.Context) { l.messages <- msg return nil - }) + }) if err != nil { switch err := err.(type) { case net.Error: @@ -237,7 +239,9 @@ func (l *Listener) ListenWithContext(ctx context.Context, fn MessageProcessor) { for { select { case msg, ok := <-l.messages: - if !ok { return } + if !ok { + return + } l.Stats.AddDelivered(1) // For simplicity, did not do the pipe of death here. If POD is received, we may deliver a // couple more messages (especially since select is random in which channel is read from). diff --git a/listener/listenerconfig.go b/listener/listenerconfig.go index 6f5e893..5a2ba0f 100644 --- a/listener/listenerconfig.go +++ b/listener/listenerconfig.go @@ -9,7 +9,7 @@ import ( type Config struct { *listenerOptions AwsConfig *aws.Config - LogLevel aws.LogLevelType + LogLevel aws.LogLevelType } // NewConfig creates a new instance of Config diff --git a/listener/sharditerator_test.go b/listener/sharditerator_test.go index 55c3dda..3efa843 100644 --- a/listener/sharditerator_test.go +++ b/listener/sharditerator_test.go @@ -1,8 +1,8 @@ package listener import ( - . "github.com/smartystreets/goconvey/convey" "github.com/aws/aws-sdk-go/aws" + . "github.com/smartystreets/goconvey/convey" "testing" "time" diff --git a/multilang/action.go b/multilang/action.go index 386895d..e2547a0 100644 --- a/multilang/action.go +++ b/multilang/action.go @@ -2,8 +2,8 @@ package multilang import ( "encoding/base64" - "time" "strconv" + "time" "github.com/rewardStyle/kinetic/message" ) @@ -13,29 +13,29 @@ type ActionType string // These are the enumerated KCL Multilang protocol action message types const ( - INITIALIZE ActionType = "initialize" + INITIALIZE ActionType = "initialize" PROCESSRECORDS ActionType = "processRecords" - RECORD ActionType = "record" - CHECKPOINT ActionType = "checkpoint" - SHUTDOWN ActionType = "shutdown" - STATUS ActionType = "status" + RECORD ActionType = "record" + CHECKPOINT ActionType = "checkpoint" + SHUTDOWN ActionType = "shutdown" + STATUS ActionType = "status" ) // ActionMessage is a struct used to marshal / unmarshal KCL Multilang protocol action messages type ActionMessage struct { - Action ActionType `json:"action"` - ShardID string `json:"shardId,omitempty"` - SequenceNumber string `json:"sequenceNumber,omitempty"` - Records []Record `json:"records,omitempty"` - Checkpoint string `json:"checkpoint,omitempty"` - Error string `json:"error,omitempty"` - Reason string `json:"reason,omitempty"` - ResponseFor ActionType `json:"responseFor,omitempty"` + Action ActionType `json:"action"` + ShardID string `json:"shardId,omitempty"` + SequenceNumber string `json:"sequenceNumber,omitempty"` + Records []Record `json:"records,omitempty"` + Checkpoint string `json:"checkpoint,omitempty"` + Error string `json:"error,omitempty"` + Reason string `json:"reason,omitempty"` + ResponseFor ActionType `json:"responseFor,omitempty"` } // Record is a struct used to marshal / unmarshal kinesis records from KCL Multilang protocol type Record struct { - Action ActionType `json:"action"` + Action ActionType `json:"action"` ApproximateArrival Timestamp `json:"approximateArrivalTimestamp"` Data string `json:"data,omitempty"` PartitionKey string `json:"partitionKey,omitempty"` @@ -57,7 +57,7 @@ func (t *Timestamp) UnmarshalJSON(b []byte) error { milliseconds := ts % 1000 seconds := (ts - milliseconds) / 1000 - t.Time = time.Unix(int64(seconds), int64(milliseconds * 1000)) + t.Time = time.Unix(int64(seconds), int64(milliseconds*1000)) return nil } @@ -71,8 +71,8 @@ func (r *Record) ToMessage() *message.Message { return &message.Message{ ApproximateArrivalTimestamp: &r.ApproximateArrival.Time, - Data: b, - PartitionKey: &r.PartitionKey, + Data: b, + PartitionKey: &r.PartitionKey, SequenceNumber: &r.SequenceNumber, } } @@ -80,7 +80,7 @@ func (r *Record) ToMessage() *message.Message { // NewCheckpointMessage is used to create a new checkpoint message func NewCheckpointMessage(seqNum string) *ActionMessage { return &ActionMessage{ - Action: CHECKPOINT, + Action: CHECKPOINT, Checkpoint: seqNum, } } @@ -88,7 +88,7 @@ func NewCheckpointMessage(seqNum string) *ActionMessage { // NewStatusMessage is used to create a new status message func NewStatusMessage(actionType ActionType) *ActionMessage { return &ActionMessage{ - Action: STATUS, + Action: STATUS, ResponseFor: actionType, } } diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go index 81239c0..07abc34 100644 --- a/producer/firehosewriter.go +++ b/producer/firehosewriter.go @@ -43,7 +43,7 @@ func NewFirehoseWriter(c *aws.Config, stream string, fn ...func(*FirehoseWriterC firehoseWriterOptions: cfg.firehoseWriterOptions, LogHelper: &logging.LogHelper{ LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, + Logger: cfg.AwsConfig.Logger, }, stream: stream, client: firehose.New(sess), diff --git a/producer/firehosewriterconfig.go b/producer/firehosewriterconfig.go index 8092f71..2e2c36b 100644 --- a/producer/firehosewriterconfig.go +++ b/producer/firehosewriterconfig.go @@ -8,7 +8,7 @@ import ( type FirehoseWriterConfig struct { *firehoseWriterOptions AwsConfig *aws.Config - LogLevel aws.LogLevelType + LogLevel aws.LogLevelType } // NewFirehoseWriterConfig creates a new instance of FirehoseWriterConfig diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index cb7ece2..236a66a 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -44,7 +44,7 @@ func NewKinesisWriter(c *aws.Config, stream string, fn ...func(*KinesisWriterCon kinesisWriterOptions: cfg.kinesisWriterOptions, LogHelper: &logging.LogHelper{ LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, + Logger: cfg.AwsConfig.Logger, }, stream: stream, client: kinesis.New(sess), diff --git a/producer/kinesiswriterconfig.go b/producer/kinesiswriterconfig.go index 878c34e..320e3fe 100644 --- a/producer/kinesiswriterconfig.go +++ b/producer/kinesiswriterconfig.go @@ -10,7 +10,7 @@ import ( type KinesisWriterConfig struct { *kinesisWriterOptions AwsConfig *aws.Config - LogLevel aws.LogLevelType + LogLevel aws.LogLevelType } // NewKinesisWriterConfig creates a new instance of KinesisWriterConfig @@ -19,7 +19,7 @@ func NewKinesisWriterConfig(cfg *aws.Config) *KinesisWriterConfig { AwsConfig: cfg, kinesisWriterOptions: &kinesisWriterOptions{ responseReadTimeout: time.Second, - Stats: &NilStatsCollector{}, + Stats: &NilStatsCollector{}, }, LogLevel: *cfg.LogLevel, } diff --git a/producer/producer.go b/producer/producer.go index 5bb87c2..b3f2fe0 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -39,7 +39,7 @@ type producerOptions struct { queueDepth int maxRetryAttempts int concurrency int - Stats StatsCollector + Stats StatsCollector } // Producer sends records to Kinesis or Firehose. @@ -69,9 +69,9 @@ func NewProducer(c *aws.Config, w StreamWriter, fn ...func(*Config)) (*Producer, producerOptions: cfg.producerOptions, LogHelper: &logging.LogHelper{ LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, + Logger: cfg.AwsConfig.Logger, }, - writer: w, + writer: w, concurrencySem: make(chan empty, cfg.concurrency), pipeOfDeath: make(chan empty), }, nil @@ -118,7 +118,7 @@ stop: err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message) error { if msg.FailCount <= p.maxRetryAttempts { // Apply a delay before retrying - time.Sleep(time.Duration(msg.FailCount * msg.FailCount) * time.Second) + time.Sleep(time.Duration(msg.FailCount*msg.FailCount) * time.Second) select { case p.retries <- msg: @@ -186,7 +186,7 @@ stop: atomic.AddUint64(&attempts, 1) // Apply a delay before retrying - time.Sleep(time.Duration(attempts * attempts) * time.Second) + time.Sleep(time.Duration(attempts*attempts) * time.Second) } // This frees up another sendBatch to run to allow drainage of the messages / retry queue. This should diff --git a/producer/producer_test.go b/producer/producer_test.go index e5b927a..181d638 100644 --- a/producer/producer_test.go +++ b/producer/producer_test.go @@ -56,7 +56,7 @@ func TestProducer(t *testing.T) { r, err := listener.NewKinesisReader(k.Session.Config, stream, shards[0], func(krc *listener.KinesisReaderConfig) { krc.SetResponseReadTimeout(1000 * time.Millisecond) - }) + }) So(r, ShouldNotBeNil) So(err, ShouldBeNil) @@ -79,12 +79,12 @@ func TestProducer(t *testing.T) { }) }) - Convey("check that we can send and receive a single message", func(){ + Convey("check that we can send and receive a single message", func() { start := time.Now() data := "hello" p.Send(&message.Message{ PartitionKey: aws.String("key"), - Data: []byte(data), + Data: []byte(data), }) msg, err := l.RetrieveWithContext(context.TODO()) elapsed := time.Since(start) diff --git a/producer/producerconfig.go b/producer/producerconfig.go index e7b7526..d2ec8ce 100644 --- a/producer/producerconfig.go +++ b/producer/producerconfig.go @@ -12,7 +12,7 @@ import ( type Config struct { *producerOptions AwsConfig *aws.Config - LogLevel aws.LogLevelType + LogLevel aws.LogLevelType } // NewConfig creates a new instance of Config. diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go index 8704ea8..ab40742 100644 --- a/producer/producerconfig_test.go +++ b/producer/producerconfig_test.go @@ -16,9 +16,9 @@ type DebugStatsCollector struct{} func (l *DebugStatsCollector) AddSent(int) {} func (l *DebugStatsCollector) AddFailed(int) {} -func (l *DebugStatsCollector) AddDroppedTotal(int) {} -func (l *DebugStatsCollector) AddDroppedCapacity(int) {} -func (l *DebugStatsCollector) AddDroppedRetries(int) {} +func (l *DebugStatsCollector) AddDroppedTotal(int) {} +func (l *DebugStatsCollector) AddDroppedCapacity(int) {} +func (l *DebugStatsCollector) AddDroppedRetries(int) {} func (l *DebugStatsCollector) AddBatchSize(int) {} func (l *DebugStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} func (l *DebugStatsCollector) AddPutRecordsCalled(int) {} diff --git a/testexec/cli.go b/testexec/cli.go index 283c96d..32bdc92 100644 --- a/testexec/cli.go +++ b/testexec/cli.go @@ -21,21 +21,21 @@ type Config struct { func parseCommandLineArgs() *Config { // Define command line flags - modePtr := flag.String("mode", "readwrite", "used to specify the mode in which to run; either 'r', 'read', 'w', " + + modePtr := flag.String("mode", "readwrite", "used to specify the mode in which to run; either 'r', 'read', 'w', "+ "'write', 'rw' or 'readwrite'") - locationPtr := flag.String("location", "local", "used to specify the location of the kinesis stream. " + - "Accepted values are (local|aws). For local, run kinesalite on http://127.0.0.1:4567. For aws, your " + + locationPtr := flag.String("location", "local", "used to specify the location of the kinesis stream. "+ + "Accepted values are (local|aws). For local, run kinesalite on http://127.0.0.1:4567. For aws, your "+ "aws credentials and configuration need to be defined at ~/.aws") - streamNamePtr := flag.String("stream-name", "", "used to specify a pre-existing stream to be used for " + + streamNamePtr := flag.String("stream-name", "", "used to specify a pre-existing stream to be used for "+ "testing. A new stream will be created if not defined.") - numMsgsPtr := flag.Int("num-msgs", 0, "used to specify the number of messages to (attempt to) send. This " + - "flag is only applicable to 'write' and 'readwrite' modes. Use zero or a negative number to produce " + + numMsgsPtr := flag.Int("num-msgs", 0, "used to specify the number of messages to (attempt to) send. This "+ + "flag is only applicable to 'write' and 'readwrite' modes. Use zero or a negative number to produce "+ "indefinitely") - durationPtr := flag.Int("duration", 0, "used to specify the duration (in seconds) the program should run. " + - "This flag is only applicable to 'write' and 'readwrite' modes. Use zero or negative number to run " + + durationPtr := flag.Int("duration", 0, "used to specify the duration (in seconds) the program should run. "+ + "This flag is only applicable to 'write' and 'readwrite' modes. Use zero or negative number to run "+ "indefinitely.") throttlePtr := flag.Bool("throttle", true, "used to specify whether to throttle PutRecord requests by 1 ms. ") - cleanupPtr := flag.Bool("cleanup", true, "used to specify whether or not to delete the kinesis stream after " + + cleanupPtr := flag.Bool("cleanup", true, "used to specify whether or not to delete the kinesis stream after "+ "processing is complete.") verbosePtr := flag.Bool("verbose", false, "used to specify whether or not to log in verbose mode") @@ -70,14 +70,14 @@ func parseCommandLineArgs() *Config { } return &Config{ - Mode: &mode, + Mode: &mode, StreamName: streamNamePtr, - Duration: durationPtr, - NumMsgs: numMsgsPtr, - Location: locationPtr, - Throttle: throttlePtr, - Cleanup: cleanupPtr, - Verbose: verbosePtr, + Duration: durationPtr, + NumMsgs: numMsgsPtr, + Location: locationPtr, + Throttle: throttlePtr, + Cleanup: cleanupPtr, + Verbose: verbosePtr, } } diff --git a/testexec/listenerstats.go b/testexec/listenerstats.go index 8e70d1a..9d05a2d 100644 --- a/testexec/listenerstats.go +++ b/testexec/listenerstats.go @@ -1,9 +1,9 @@ package main import ( + "log" "sync/atomic" "time" - "log" ) // ListenerStatsCollector implements the listener's StatsCollector @@ -17,7 +17,7 @@ type ListenerStatsCollector struct { GetRecordsTimeout uint64 GetRecordsReadTimeout uint64 - ProcessedDuration time.Duration + ProcessedDuration time.Duration GetRecordsDuration time.Duration GetRecordsReadResonseDuration time.Duration GetRecordsUnmarshalDuration time.Duration @@ -114,4 +114,4 @@ func (lsc *ListenerStatsCollector) PrintStats() { log.Printf("Listener stats: Provisioned Throughput Exceeded: [%d]\n", atomic.LoadUint64(&lsc.ProvisionedThroughputExceeded)) log.Printf("Listener stats: GetRecords Timeout: [%d]\n", atomic.LoadUint64(&lsc.GetRecordsTimeout)) log.Printf("Listener stats: GetRecords Read Timeout: [%d]\n", atomic.LoadUint64(&lsc.GetRecordsReadTimeout)) -} \ No newline at end of file +} diff --git a/testexec/main.go b/testexec/main.go index 7505555..bc2e2aa 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -22,20 +22,20 @@ import ( "github.com/rewardStyle/kinetic/message" "github.com/rewardStyle/kinetic/producer" - _ "net/http/pprof" "net/http" + _ "net/http/pprof" ) // Define constants for Kinesis stream location const ( LocationLocal = "local" - LocationAws = "aws" + LocationAws = "aws" ) // Define operation modes const ( - ModeRead = "read" - ModeWrite = "write" + ModeRead = "read" + ModeWrite = "write" ModeReadWrite = "readwrite" ) @@ -103,7 +103,7 @@ func main() { k = newDefaultKinetic() case LocationAws: k = newAwsKinetic() - default : + default: log.Fatalf("Unknown location for kinesis stream: %s\n", *cfg.Location) } @@ -273,7 +273,7 @@ func newKineticListener(k *kinetic.Kinetic, streamName string) *listener.Listene func(krc *listener.KinesisReaderConfig) { krc.SetResponseReadTimeout(1000 * time.Millisecond) krc.SetStatsCollector(lsc) - }) + }) if err != nil { log.Fatalf("Unable to create a new kinesis reader due to: %v\n", err) } @@ -383,7 +383,7 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { jsonStr, _ := json.Marshal(NewMessage()) if err := p.Send(&message.Message{ PartitionKey: aws.String("key"), - Data: []byte(jsonStr), + Data: []byte(jsonStr), }); err == nil { sd.incrementMsgCount() } else { @@ -395,7 +395,7 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { // Control when to exit produce produceWg := sync.WaitGroup{} produceWg.Add(1) - go func(){ + go func() { defer produceWg.Done() var sent uint64 @@ -452,7 +452,7 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { } return case <-time.After(time.Second): - newSent := atomic.LoadUint64(&p.Stats.(*ProducerStatsCollector).Sent); + newSent := atomic.LoadUint64(&p.Stats.(*ProducerStatsCollector).Sent) if sent != newSent { staleTime.Reset(staleTimeout) sent = newSent @@ -482,7 +482,7 @@ func listen(sd *StreamData, l *listener.Listener, wg *sync.WaitGroup) { }() // Call Listen within a go routine - go func(){ + go func() { l.Listen(func(m *message.Message, wg *sync.WaitGroup) error { defer wg.Done() @@ -507,7 +507,7 @@ func listen(sd *StreamData, l *listener.Listener, wg *sync.WaitGroup) { // Control when to exit listen listenWg := sync.WaitGroup{} listenWg.Add(1) - go func(){ + go func() { defer listenWg.Done() var staleTimeout time.Duration @@ -535,7 +535,7 @@ func listen(sd *StreamData, l *listener.Listener, wg *sync.WaitGroup) { } return case <-time.After(time.Second): - newConsumed := atomic.LoadUint64(&l.Stats.(*ListenerStatsCollector).Consumed); + newConsumed := atomic.LoadUint64(&l.Stats.(*ListenerStatsCollector).Consumed) if consumed != newConsumed { staleTime.Reset(staleTimeout) consumed = newConsumed diff --git a/testexec/messages.go b/testexec/messages.go index 730e50d..b963dad 100644 --- a/testexec/messages.go +++ b/testexec/messages.go @@ -1,8 +1,8 @@ package main import ( - "sync/atomic" "strconv" + "sync/atomic" ) var msgCount uint64 @@ -18,7 +18,7 @@ func NewMessage() *Message { atomic.AddUint64(&msgCount, 1) id := atomic.LoadUint64(&msgCount) return &Message{ - ID: id, + ID: id, Message: "hello_" + strconv.Itoa(int(id)), } } diff --git a/testexec/producerstats.go b/testexec/producerstats.go index 83fec67..2f420ea 100644 --- a/testexec/producerstats.go +++ b/testexec/producerstats.go @@ -1,9 +1,9 @@ package main import ( + "log" "sync/atomic" "time" - "log" ) // ProducerStatsCollector implements the producer's StatsCollector @@ -19,9 +19,9 @@ type ProducerStatsCollector struct { ProvisionedThroughputExceeded uint64 PutRecordsTimeout uint64 - PutRecordsDuration time.Duration - PutRecordsBuildDuration time.Duration - PutRecordsSendDuration time.Duration + PutRecordsDuration time.Duration + PutRecordsBuildDuration time.Duration + PutRecordsSendDuration time.Duration } // GetSent returns the number of messages sent to AWS Kinesis by the producer. diff --git a/testexec/streamdata.go b/testexec/streamdata.go index cdf4e28..253bcf8 100644 --- a/testexec/streamdata.go +++ b/testexec/streamdata.go @@ -17,7 +17,7 @@ type StreamData struct { // NewStreamData instantiates a new StreamData struct func NewStreamData() *StreamData { return &StreamData{ - mutex: sync.Mutex{}, + mutex: sync.Mutex{}, Frequencies: make(map[uint64]uint64), } } From 2a6685c30cfd75cbe66898e847b708ff1c9124c3 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 10 Jul 2017 12:00:29 -0500 Subject: [PATCH 047/109] Added a DefaultStatsCollector to both the listener package and the producer package which both implement the StatsCollector interface using the rcrowley/go-metrics library. Added unit tests and modified the testexec to use the new DefaultStatusCollectors. --- glide.lock | 6 +- glide.yaml | 1 + listener/stats.go | 152 +++++++++++++++++++++++++++++++++---- listener/stats_test.go | 143 +++++++++++++++++++++++++++++++++++ producer/stats.go | 152 +++++++++++++++++++++++++++++++++---- producer/stats_test.go | 154 ++++++++++++++++++++++++++++++++++++++ testexec/cli.go | 2 +- testexec/listenerstats.go | 117 ----------------------------- testexec/main.go | 30 +++++--- testexec/producerstats.go | 117 ----------------------------- 10 files changed, 598 insertions(+), 276 deletions(-) create mode 100644 listener/stats_test.go create mode 100644 producer/stats_test.go delete mode 100644 testexec/listenerstats.go delete mode 100644 testexec/producerstats.go diff --git a/glide.lock b/glide.lock index e481c5a..fa3df35 100644 --- a/glide.lock +++ b/glide.lock @@ -1,5 +1,5 @@ -hash: 16ed7c907af98983bf7cbe350e9ff03fb1208a4ad54664a83e488935889e6207 -updated: 2017-06-27T15:49:17.807220643-05:00 +hash: 70d3e36401b9615bc762c364243c4177b3b4b51d36e3fdeaa3ec923996430e24 +updated: 2017-07-06T15:44:39.094878239-05:00 imports: - name: github.com/aws/aws-sdk-go version: 72e42b13da62269f68308fb6068b7ea691a416a4 @@ -44,6 +44,8 @@ imports: version: bd40a432e4c76585ef6b72d3fd96fb9b6dc7b68d - name: github.com/jtolds/gls version: f1ac7f4f24f50328e6bc838ca4437d1612a0243c +- name: github.com/rcrowley/go-metrics + version: ab2277b1c5d15c3cba104e9cbddbdfc622df5ad8 - name: github.com/smartystreets/assertions version: c9ee7d9e9a2aeec0bee7c4a516f3e0ad7cb7e558 subpackages: diff --git a/glide.yaml b/glide.yaml index ebe4cdf..7c867c6 100644 --- a/glide.yaml +++ b/glide.yaml @@ -63,3 +63,4 @@ import: version: 8a331561fe74dadba6edfc59f3be66c22c3b065d - package: github.com/stretchr/testify/assert version: v1.1.4 +- package: github.com/rcrowley/go-metrics diff --git a/listener/stats.go b/listener/stats.go index 0d897af..14c913f 100644 --- a/listener/stats.go +++ b/listener/stats.go @@ -1,7 +1,10 @@ package listener import ( + "log" "time" + + "github.com/rcrowley/go-metrics" ) // StatsCollector allows for a collector to collect various metrics produced by @@ -16,7 +19,6 @@ type StatsCollector interface { AddProvisionedThroughputExceeded(int) AddGetRecordsTimeout(int) AddGetRecordsReadTimeout(int) - AddProcessedDuration(time.Duration) AddGetRecordsDuration(time.Duration) AddGetRecordsReadResponseDuration(time.Duration) @@ -28,52 +30,176 @@ type NilStatsCollector struct{} // AddConsumed records a count of the number of messages received from AWS // Kinesis by the listener. -func (l *NilStatsCollector) AddConsumed(int) {} +func (nsc *NilStatsCollector) AddConsumed(int) {} + +// AddDelivered records a count of the number of messages delivered to the +// application by the listener. +func (nsc *NilStatsCollector) AddDelivered(int) {} + +// AddProcessed records a count of the number of messages processed by the +// application by the listener. This is based on a WaitGroup that is sent to +// the RetrieveFn and Listen functions. Retrieve does not count processed +// messages. +func (nsc *NilStatsCollector) AddProcessed(int) {} + +// AddBatchSize records a count of the number of messages returned by +// GetRecords in the listener. +func (nsc *NilStatsCollector) AddBatchSize(int) {} + +// AddGetRecordsCalled records the number of times the GetRecords API was called +// by the listener. +func (nsc *NilStatsCollector) AddGetRecordsCalled(int) {} + +// AddProvisionedThroughputExceeded records the number of times the GetRecords +// API returned a ErrCodeProvisionedThroughputExceededException by the listener. +func (nsc *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} + +// AddGetRecordsTimeout records the number of times the GetRecords API timed out +// on the HTTP level. This is influenced by the WithHTTPClientTimeout +// configuration. +func (nsc *NilStatsCollector) AddGetRecordsTimeout(int) {} + +// AddGetRecordsReadTimeout records the number of times the GetRecords API timed +// out while reading the response body. This is influenced by the +// WithGetRecordsReadTimeout configuration. +func (nsc *NilStatsCollector) AddGetRecordsReadTimeout(int) {} + +// AddProcessedDuration records the duration to process a record. See notes on +// AddProcessed. +func (nsc *NilStatsCollector) AddProcessedDuration(time.Duration) {} + +// AddGetRecordsDuration records the duration that the GetRecords API request +// took. Only the times of successful calls are measured. +func (nsc *NilStatsCollector) AddGetRecordsDuration(time.Duration) {} + +// AddGetRecordsReadResponseDuration records the duration that it took to read +// the response body of a GetRecords API request. +func (nsc *NilStatsCollector) AddGetRecordsReadResponseDuration(time.Duration) {} + +// AddGetRecordsUnmarshalDuration records the duration that it took to unmarshal +// the response body of a GetRecords API request. +func (nsc *NilStatsCollector) AddGetRecordsUnmarshalDuration(time.Duration) {} + +// DefaultStatsCollector is a type that implements the listener's StatsCollector interface using the +// rcrowley/go-metrics library +type DefaultStatsCollector struct { + Consumed metrics.Counter + Delivered metrics.Counter + Processed metrics.Counter + BatchSize metrics.Counter + GetRecordsCalled metrics.Counter + ProvisionedThroughputExceeded metrics.Counter + GetRecordsTimeout metrics.Counter + GetRecordsReadTimeout metrics.Counter + ProcessedDuration metrics.Gauge + GetRecordsDuration metrics.Gauge + GetRecordsReadResponseDuration metrics.Gauge + GetRecordsUnmarshalDuration metrics.Gauge +} + +// NewDefaultStatsCollector instantiates a new DefaultStatsCollector object +func NewDefaultStatsCollector(r metrics.Registry) *DefaultStatsCollector { + return &DefaultStatsCollector{ + Consumed: metrics.GetOrRegisterCounter("listener.consumed", r), + Delivered: metrics.GetOrRegisterCounter("listener.delivered", r), + Processed: metrics.GetOrRegisterCounter("listener.processed", r), + BatchSize: metrics.GetOrRegisterCounter("listener.batchsize", r), + GetRecordsCalled: metrics.GetOrRegisterCounter("listener.sent", r), + ProvisionedThroughputExceeded: metrics.GetOrRegisterCounter("listener.provisionedthroughputexceeded", r), + GetRecordsTimeout: metrics.GetOrRegisterCounter("listener.getrecords.timeout", r), + GetRecordsReadTimeout: metrics.GetOrRegisterCounter("listener.getrecords.readtimeout", r), + ProcessedDuration: metrics.GetOrRegisterGauge("listener.processed.duration", r), + GetRecordsDuration: metrics.GetOrRegisterGauge("listener.getrecords.duration", r), + GetRecordsReadResponseDuration: metrics.GetOrRegisterGauge("listener.getrecords.readresponse.duration", r), + GetRecordsUnmarshalDuration: metrics.GetOrRegisterGauge("listener.getrecords.unmarshal.duration", r), + } +} + +// AddConsumed records a count of the number of messages received from AWS +// Kinesis by the listener. +func (dsc *DefaultStatsCollector) AddConsumed(count int) { + dsc.Consumed.Inc(int64(count)) +} // AddDelivered records a count of the number of messages delivered to the // application by the listener. -func (l *NilStatsCollector) AddDelivered(int) {} +func (dsc *DefaultStatsCollector) AddDelivered(count int) { + dsc.Delivered.Inc(int64(count)) +} // AddProcessed records a count of the number of messages processed by the // application by the listener. This is based on a WaitGroup that is sent to // the RetrieveFn and Listen functions. Retrieve does not count processed // messages. -func (l *NilStatsCollector) AddProcessed(int) {} +func (dsc *DefaultStatsCollector) AddProcessed(count int) { + dsc.Processed.Inc(int64(count)) +} // AddBatchSize records a count of the number of messages returned by // GetRecords in the listener. -func (l *NilStatsCollector) AddBatchSize(int) {} +func (dsc *DefaultStatsCollector) AddBatchSize(count int) { + dsc.BatchSize.Inc(int64(count)) +} // AddGetRecordsCalled records the number of times the GetRecords API was called // by the listener. -func (l *NilStatsCollector) AddGetRecordsCalled(int) {} +func (dsc *DefaultStatsCollector) AddGetRecordsCalled(count int) { + dsc.GetRecordsCalled.Inc(int64(count)) +} // AddProvisionedThroughputExceeded records the number of times the GetRecords // API returned a ErrCodeProvisionedThroughputExceededException by the listener. -func (l *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} +func (dsc *DefaultStatsCollector) AddProvisionedThroughputExceeded(count int) { + dsc.ProvisionedThroughputExceeded.Inc(int64(count)) +} // AddGetRecordsTimeout records the number of times the GetRecords API timed out // on the HTTP level. This is influenced by the WithHTTPClientTimeout // configuration. -func (l *NilStatsCollector) AddGetRecordsTimeout(int) {} +func (dsc *DefaultStatsCollector) AddGetRecordsTimeout(count int) { + dsc.GetRecordsTimeout.Inc(int64(count)) +} // AddGetRecordsReadTimeout records the number of times the GetRecords API timed // out while reading the response body. This is influenced by the // WithGetRecordsReadTimeout configuration. -func (l *NilStatsCollector) AddGetRecordsReadTimeout(int) {} +func (dsc *DefaultStatsCollector) AddGetRecordsReadTimeout(count int) { + dsc.GetRecordsReadTimeout.Inc(int64(count)) +} // AddProcessedDuration records the duration to process a record. See notes on // AddProcessed. -func (l *NilStatsCollector) AddProcessedDuration(time.Duration) {} +func (dsc *DefaultStatsCollector) AddProcessedDuration(duration time.Duration) { + dsc.ProcessedDuration.Update(duration.Nanoseconds()) +} // AddGetRecordsDuration records the duration that the GetRecords API request // took. Only the times of successful calls are measured. -func (l *NilStatsCollector) AddGetRecordsDuration(time.Duration) {} +func (dsc *DefaultStatsCollector) AddGetRecordsDuration(duration time.Duration) { + dsc.GetRecordsDuration.Update(duration.Nanoseconds()) +} // AddGetRecordsReadResponseDuration records the duration that it took to read // the response body of a GetRecords API request. -func (l *NilStatsCollector) AddGetRecordsReadResponseDuration(time.Duration) {} +func (dsc *DefaultStatsCollector) AddGetRecordsReadResponseDuration(duration time.Duration) { + dsc.GetRecordsReadResponseDuration.Update(duration.Nanoseconds()) +} // AddGetRecordsUnmarshalDuration records the duration that it took to unmarshal // the response body of a GetRecords API request. -func (l *NilStatsCollector) AddGetRecordsUnmarshalDuration(time.Duration) {} +func (dsc *DefaultStatsCollector) AddGetRecordsUnmarshalDuration(duration time.Duration) { + dsc.GetRecordsUnmarshalDuration.Update(duration.Nanoseconds()) +} + +// PrintStats logs the stats +func (dsc *DefaultStatsCollector) PrintStats() { + log.Printf("Listener stats: Consumed: [%d]\n", dsc.Consumed.Count()) + log.Printf("Listener stats: Delivered: [%d]\n", dsc.Delivered.Count()) + log.Printf("Listener stats: Processed: [%d]\n", dsc.Processed.Count()) + log.Printf("Listener stats: Batch Size: [%d]\n", dsc.BatchSize.Count()) + log.Printf("Listener stats: GetRecords Called: [%d]\n", dsc.GetRecordsCalled.Count()) + log.Printf("Listener stats: Provisioned Throughput Exceeded: [%d]\n", dsc.ProvisionedThroughputExceeded.Count()) + log.Printf("Listener stats: GetRecords Timeout: [%d]\n", dsc.GetRecordsTimeout.Count()) + log.Printf("Listener stats: GetRecords Read Timeout: [%d]\n", dsc.GetRecordsReadTimeout.Count()) +} + diff --git a/listener/stats_test.go b/listener/stats_test.go new file mode 100644 index 0000000..c5b4650 --- /dev/null +++ b/listener/stats_test.go @@ -0,0 +1,143 @@ +package listener + +import ( + . "github.com/smartystreets/goconvey/convey" + "math/rand" + "testing" + "time" + + "github.com/rcrowley/go-metrics" +) + +func init() { + rand.Seed(time.Now().UnixNano()) +} + +func TestStatsCollector(t *testing.T) { + Convey("given a NilStatsCollector", t, func() { + var sc StatsCollector = &NilStatsCollector{} + So(sc, ShouldNotBeNil) + + Convey("check that AddConsumed does not error", func() { + sc.AddConsumed(1) + }) + + Convey("check that AddDelivered does not error", func() { + sc.AddDelivered(1) + }) + + Convey("check that AddProcessed does not error", func() { + sc.AddProcessed(1) + }) + + Convey("check that AddBatchSize does not error", func() { + sc.AddBatchSize(1) + }) + + Convey("check that AddGetRecordsCalled does not error", func() { + sc.AddGetRecordsCalled(1) + }) + + Convey("check that AddProvisionedThroughputExceeded does not error", func() { + sc.AddProvisionedThroughputExceeded(1) + }) + + Convey("check that AddGetRecordsTimeout does not error", func() { + sc.AddGetRecordsTimeout(1) + }) + + Convey("check that AddGetRecordsReadTimeout does not error", func() { + sc.AddGetRecordsReadTimeout(1) + }) + + Convey("check that AddProcessedDuration does not error", func() { + sc.AddProcessedDuration(1) + }) + + Convey("check that AddGetRecordsDuration does not error", func() { + sc.AddGetRecordsDuration(1) + }) + + Convey("check that AddGetRecordsReadResponseDuration does not error", func() { + sc.AddGetRecordsReadResponseDuration(1) + }) + + Convey("check that AddGetRecordsUnmarshalDuration does not error", func() { + sc.AddGetRecordsUnmarshalDuration(1) + }) + }) + + Convey("given a DefaultStatsCollector", t, func() { + r := metrics.NewRegistry() + var sc StatsCollector = NewDefaultStatsCollector(r) + So(sc, ShouldNotBeNil) + + Convey("check that AddConsumed does not error", func() { + count := rand.Int() + sc.AddConsumed(count) + So(sc.(*DefaultStatsCollector).Consumed.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddDelivered does not error", func() { + count := rand.Int() + sc.AddDelivered(count) + So(sc.(*DefaultStatsCollector).Delivered.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddProcessed does not error", func() { + count := rand.Int() + sc.AddProcessed(count) + So(sc.(*DefaultStatsCollector).Processed.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddBatchSize does not error", func() { + count := rand.Int() + sc.AddBatchSize(count) + So(sc.(*DefaultStatsCollector).BatchSize.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddGetRecordsCalled does not error", func() { + count := rand.Int() + sc.AddGetRecordsCalled(count) + So(sc.(*DefaultStatsCollector).GetRecordsCalled.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddProvisionedThroughputExceeded does not error", func() { + count := rand.Int() + sc.AddProvisionedThroughputExceeded(count) + So(sc.(*DefaultStatsCollector).ProvisionedThroughputExceeded.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddGetRecordsTimeout does not error", func() { + count := rand.Int() + sc.AddGetRecordsTimeout(count) + So(sc.(*DefaultStatsCollector).GetRecordsTimeout.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddGetRecordsReadTimeout does not error", func() { + count := rand.Int() + sc.AddGetRecordsReadTimeout(count) + So(sc.(*DefaultStatsCollector).GetRecordsReadTimeout.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddProcessedDuration does not error", func() { + sc.AddProcessedDuration(time.Second) + So(sc.(*DefaultStatsCollector).ProcessedDuration.Value(), ShouldEqual, 1000000000) + }) + + Convey("check that AddGetRecordsDuration does not error", func() { + sc.AddGetRecordsDuration(time.Second) + So(sc.(*DefaultStatsCollector).GetRecordsDuration.Value(), ShouldEqual, 1000000000) + }) + + Convey("check that AddGetRecordsReadResponseDuration does not error", func() { + sc.AddGetRecordsReadResponseDuration(time.Second) + So(sc.(*DefaultStatsCollector).GetRecordsReadResponseDuration.Value(), ShouldEqual, 1000000000) + }) + + Convey("check that AddGetRecordsUnmarshalDuration does not error", func() { + sc.AddGetRecordsUnmarshalDuration(time.Second) + So(sc.(*DefaultStatsCollector).GetRecordsUnmarshalDuration.Value(), ShouldEqual, 1000000000) + }) + }) +} diff --git a/producer/stats.go b/producer/stats.go index 6c2a6bc..657149a 100644 --- a/producer/stats.go +++ b/producer/stats.go @@ -1,7 +1,10 @@ package producer import ( + "log" "time" + + "github.com/rcrowley/go-metrics" ) // StatsCollector allows for a collector to collect various metrics produced by @@ -18,7 +21,6 @@ type StatsCollector interface { AddPutRecordsCalled(int) AddProvisionedThroughputExceeded(int) AddPutRecordsTimeout(int) - AddPutRecordsDuration(time.Duration) AddPutRecordsBuildDuration(time.Duration) AddPutRecordsSendDuration(time.Duration) @@ -28,46 +30,168 @@ type StatsCollector interface { type NilStatsCollector struct{} // AddSent records a count of the number of messages sent to AWS Kinesis by the producer. -func (l *NilStatsCollector) AddSent(int) {} +func (nsc *NilStatsCollector) AddSent(int) {} // AddFailed records a count of the number of messages that failed to be sent to AWS Kinesis by the producer. -func (l *NilStatsCollector) AddFailed(int) {} +func (nsc *NilStatsCollector) AddFailed(int) {} // AddDroppedTotal records a count of the total number of messages dropped by the application after multiple failures. -func (l *NilStatsCollector) AddDroppedTotal(int) {} +func (nsc *NilStatsCollector) AddDroppedTotal(int) {} // AddDroppedCapacity records a count of the number of messages that were dropped by the application due to the stream // writer being at capacity. -func (l *NilStatsCollector) AddDroppedCapacity(int) {} +func (nsc *NilStatsCollector) AddDroppedCapacity(int) {} // AddDroppedRetries records a count of the number of retry messages dropped by the application after the max number of // retries was exceeded. -func (l *NilStatsCollector) AddDroppedRetries(int) {} +func (nsc *NilStatsCollector) AddDroppedRetries(int) {} // AddBatchSize records a count of the number of messages attempted by PutRecords in the producer. -func (l *NilStatsCollector) AddBatchSize(int) {} +func (nsc *NilStatsCollector) AddBatchSize(int) {} // AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords API returned a // ErrCodeProvisionedThroughputExceededException by the producer. -func (l *NilStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} +func (nsc *NilStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} // AddPutRecordsCalled records the number of times the PutRecords API was called by the producer. -func (l *NilStatsCollector) AddPutRecordsCalled(int) {} +func (nsc *NilStatsCollector) AddPutRecordsCalled(int) {} // AddProvisionedThroughputExceeded records the number of times the PutRecords API response contained a record which // contained an ErrCodeProvisionedThroughputExceededException error. -func (l *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} +func (nsc *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} // AddPutRecordsTimeout records the number of times the PutRecords API timed out on the HTTP level. This is influenced // by the WithHTTPClientTimeout configuration. -func (l *NilStatsCollector) AddPutRecordsTimeout(int) {} +func (nsc *NilStatsCollector) AddPutRecordsTimeout(int) {} // AddPutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls // are measured. -func (l *NilStatsCollector) AddPutRecordsDuration(time.Duration) {} +func (nsc *NilStatsCollector) AddPutRecordsDuration(time.Duration) {} // AddPutRecordsBuildDuration records the duration that it took to build the PutRecords API request payload. -func (l *NilStatsCollector) AddPutRecordsBuildDuration(time.Duration) {} +func (nsc *NilStatsCollector) AddPutRecordsBuildDuration(time.Duration) {} // AddPutRecordsSendDuration records the duration that it took to send the PutRecords API request payload. -func (l *NilStatsCollector) AddPutRecordsSendDuration(time.Duration) {} +func (nsc *NilStatsCollector) AddPutRecordsSendDuration(time.Duration) {} + +// DefaultStatsCollector is a type that implements the producers's StatsCollector interface using the +// rcrowley/go-metrics library +type DefaultStatsCollector struct { + Sent metrics.Counter + Failed metrics.Counter + DroppedTotal metrics.Counter + DroppedCapacity metrics.Counter + DroppedRetries metrics.Counter + BatchSize metrics.Counter + PutRecordsProvisionedThroughputExceeded metrics.Counter + PutRecordsCalled metrics.Counter + ProvisionedThroughputExceeded metrics.Counter + PutRecordsTimeout metrics.Counter + PutRecordsDuration metrics.Gauge + PutRecordsBuildDuration metrics.Gauge + PutRecordsSendDuration metrics.Gauge +} + +// NewDefaultStatsCollector instantiates a new DefaultStatsCollector object +func NewDefaultStatsCollector(r metrics.Registry) *DefaultStatsCollector { + return &DefaultStatsCollector{ + Sent: metrics.GetOrRegisterCounter("producer.sent", r), + Failed: metrics.GetOrRegisterCounter("producer.failed", r), + DroppedTotal: metrics.GetOrRegisterCounter("producer.dropped.total", r), + DroppedCapacity: metrics.GetOrRegisterCounter("producer.dropped.capacity", r), + DroppedRetries: metrics.GetOrRegisterCounter("producer.dropped.retries", r), + BatchSize: metrics.GetOrRegisterCounter("producer.batchsize", r), + PutRecordsProvisionedThroughputExceeded: metrics.GetOrRegisterCounter("producer.putrecords.provisionedthroughputexceeded", r), + PutRecordsCalled: metrics.GetOrRegisterCounter("producer.putrecords.called", r), + ProvisionedThroughputExceeded: metrics.GetOrRegisterCounter("producer.provisionedthroughputexceeded", r), + PutRecordsTimeout: metrics.GetOrRegisterCounter("producer.putrecords.timeout", r), + PutRecordsDuration: metrics.GetOrRegisterGauge("producer.putrecords.duration", r), + PutRecordsBuildDuration: metrics.GetOrRegisterGauge("producer.putrecords.build.duration", r), + PutRecordsSendDuration: metrics.GetOrRegisterGauge("producer.putrecords.send.duration", r), + } +} + +// AddSent records a count of the number of messages sent to AWS Kinesis by the producer. +func (dsc *DefaultStatsCollector) AddSent(count int) { + dsc.Sent.Inc(int64(count)) +} + +// AddFailed records a count of the number of messages that failed to be sent to AWS Kinesis by the producer. +func (dsc *DefaultStatsCollector) AddFailed(count int) { + dsc.Failed.Inc(int64(count)) +} + +// AddDroppedTotal records a count of the total number of messages dropped by the application after multiple failures. +func (dsc *DefaultStatsCollector) AddDroppedTotal(count int) { + dsc.DroppedTotal.Inc(int64(count)) +} + +// AddDroppedCapacity records a count of the number of messages that were dropped by the application due to the stream +// writer being at capacity. +func (dsc *DefaultStatsCollector) AddDroppedCapacity(count int) { + dsc.DroppedCapacity.Inc(int64(count)) +} + +// AddDroppedRetries records a count of the number of retry messages dropped by the application after the max number of +// retries was exceeded. +func (dsc *DefaultStatsCollector) AddDroppedRetries(count int) { + dsc.DroppedRetries.Inc(int64(count)) +} + +// AddBatchSize records a count of the number of messages attempted by PutRecords in the producer. +func (dsc *DefaultStatsCollector) AddBatchSize(count int) { + dsc.BatchSize.Inc(int64(count)) +} + +// AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords API returned a +// ErrCodeProvisionedThroughputExceededException by the producer. +func (dsc *DefaultStatsCollector) AddPutRecordsProvisionedThroughputExceeded(count int) { + dsc.PutRecordsProvisionedThroughputExceeded.Inc(int64(count)) +} + +// AddPutRecordsCalled records the number of times the PutRecords API was called by the producer. +func (dsc *DefaultStatsCollector) AddPutRecordsCalled(count int) { + dsc.PutRecordsCalled.Inc(int64(count)) +} + +// AddProvisionedThroughputExceeded records the number of times the PutRecords API response contained a record which +// contained an ErrCodeProvisionedThroughputExceededException error. +func (dsc *DefaultStatsCollector) AddProvisionedThroughputExceeded(count int) { + dsc.ProvisionedThroughputExceeded.Inc(int64(count)) +} + +// AddPutRecordsTimeout records the number of times the PutRecords API timed out on the HTTP level. This is influenced +// by the WithHTTPClientTimeout configuration. +func (dsc *DefaultStatsCollector) AddPutRecordsTimeout(count int) { + dsc.PutRecordsTimeout.Inc(int64(count)) +} + +// AddPutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls +// are measured. +func (dsc *DefaultStatsCollector) AddPutRecordsDuration(duration time.Duration) { + dsc.PutRecordsDuration.Update(duration.Nanoseconds()) +} + +// AddPutRecordsBuildDuration records the duration that it took to build the PutRecords API request payload. +func (dsc *DefaultStatsCollector) AddPutRecordsBuildDuration(duration time.Duration) { + dsc.PutRecordsBuildDuration.Update(duration.Nanoseconds()) +} + +// AddPutRecordsSendDuration records the duration that it took to send the PutRecords API request payload. +func (dsc *DefaultStatsCollector) AddPutRecordsSendDuration(duration time.Duration) { + dsc.PutRecordsSendDuration.Update(duration.Nanoseconds()) +} + +// PrintStats logs the stats +func (dsc *DefaultStatsCollector) PrintStats() { + log.Printf("Producer Stats: Sent: [%d]\n", dsc.Sent.Count()) + log.Printf("Producer Stats: Failed: [%d]\n", dsc.Failed.Count()) + log.Printf("Producer Stats: Dropped Total: [%d]\n", dsc.DroppedTotal.Count()) + log.Printf("Producer Stats: Dropped Retries: [%d]\n", dsc.DroppedRetries.Count()) + log.Printf("Producer Stats: Dropped Capacity: [%d]\n", dsc.DroppedCapacity.Count()) + log.Printf("Producer Stats: Batch Size: [%d]\n", dsc.BatchSize.Count()) + log.Printf("Producer Stats: PutRecords Called: [%d]\n", dsc.PutRecordsCalled.Count()) + log.Printf("Producer Stats: Provisioned Throughput Exceeded: [%d]\n", dsc.ProvisionedThroughputExceeded.Count()) + log.Printf("Producer Stats: PutRecords Timeout: [%d]\n", dsc.PutRecordsTimeout.Count()) +} + diff --git a/producer/stats_test.go b/producer/stats_test.go new file mode 100644 index 0000000..953afe0 --- /dev/null +++ b/producer/stats_test.go @@ -0,0 +1,154 @@ +package producer + +import ( + . "github.com/smartystreets/goconvey/convey" + "math/rand" + "testing" + "time" + + "github.com/rcrowley/go-metrics" +) + +func init() { + rand.Seed(time.Now().UnixNano()) +} + +func TestStatsCollector(t *testing.T) { + Convey("given a NilStatsCollector", t, func() { + var sc StatsCollector = &NilStatsCollector{} + So(sc, ShouldNotBeNil) + + Convey("check that AddSent does not error", func() { + sc.AddSent(1) + }) + + Convey("check that AddFailed does not error", func() { + sc.AddFailed(1) + }) + + Convey("check that AddDroppedTotal does not error", func() { + sc.AddDroppedTotal(1) + }) + + Convey("check that AddDroppedCapacity does not error", func() { + sc.AddDroppedCapacity(1) + }) + + Convey("check that AddDroppedRetries does not error", func() { + sc.AddDroppedRetries(1) + }) + + Convey("check that AddBatchSize does not error", func() { + sc.AddBatchSize(1) + }) + + Convey("check that AddPutRecordsProvisionedThroughputExceeded does not error", func() { + sc.AddPutRecordsProvisionedThroughputExceeded(1) + }) + + Convey("check that AddPutRecordsCalled does not eroror", func() { + sc.AddPutRecordsCalled(1) + }) + + Convey("check that AddProvisionedThroughputExceeded does not erro", func() { + sc.AddProvisionedThroughputExceeded(1) + }) + + Convey("check that AddPutRecordsTimeout does not error", func() { + sc.AddPutRecordsTimeout(1) + }) + + Convey("check that AddPutRecordsDuration does not error", func() { + sc.AddPutRecordsDuration(time.Second) + }) + + Convey("check that AddPutRecordsBuildDuration does not error", func() { + sc.AddPutRecordsBuildDuration(time.Second) + }) + + Convey("check that AddPutRecordsSendDuration does not error", func() { + sc.AddPutRecordsSendDuration(time.Second) + }) + }) + + Convey("given a DefaulStatsCollector", t, func() { + r := metrics.NewRegistry() + var sc StatsCollector = NewDefaultStatsCollector(r) + So(sc, ShouldNotBeNil) + + Convey("check that AddSent does not error", func() { + count := rand.Int() + sc.AddSent(count) + So(sc.(*DefaultStatsCollector).Sent.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddFailed does not error", func() { + count := rand.Int() + sc.AddFailed(count) + So(sc.(*DefaultStatsCollector).Failed.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddDroppedTotal does not error", func() { + count := rand.Int() + sc.AddDroppedTotal(count) + So(sc.(*DefaultStatsCollector).DroppedTotal.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddDroppedCapacity does not error", func() { + count := rand.Int() + sc.AddDroppedCapacity(count) + So(sc.(*DefaultStatsCollector).DroppedCapacity.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddDroppedRetries does not error", func() { + count := rand.Int() + sc.AddDroppedRetries(count) + So(sc.(*DefaultStatsCollector).DroppedRetries.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddBatchSize does not error", func() { + count := rand.Int() + sc.AddBatchSize(count) + So(sc.(*DefaultStatsCollector).BatchSize.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddPutRecordsProvisionedThroughputExceeded does not error", func() { + count := rand.Int() + sc.AddPutRecordsProvisionedThroughputExceeded(count) + So(sc.(*DefaultStatsCollector).PutRecordsProvisionedThroughputExceeded.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddPutRecordsCalled does not eroror", func() { + count := rand.Int() + sc.AddPutRecordsCalled(count) + So(sc.(*DefaultStatsCollector).PutRecordsCalled.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddProvisionedThroughputExceeded does not erro", func() { + count := rand.Int() + sc.AddProvisionedThroughputExceeded(count) + So(sc.(*DefaultStatsCollector).ProvisionedThroughputExceeded.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddPutRecordsTimeout does not error", func() { + count := rand.Int() + sc.AddPutRecordsTimeout(count) + So(sc.(*DefaultStatsCollector).PutRecordsTimeout.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddPutRecordsDuration does not error", func() { + sc.AddPutRecordsDuration(time.Second) + So(sc.(*DefaultStatsCollector).PutRecordsDuration.Value(), ShouldEqual, 1000000000) + }) + + Convey("check that AddPutRecordsBuildDuration does not error", func() { + sc.AddPutRecordsBuildDuration(time.Second) + So(sc.(*DefaultStatsCollector).PutRecordsBuildDuration.Value(), ShouldEqual, 1000000000) + }) + + Convey("check that AddPutRecordsSendDuration does not error", func() { + sc.AddPutRecordsSendDuration(time.Second) + So(sc.(*DefaultStatsCollector).PutRecordsSendDuration.Value(), ShouldEqual, 1000000000) + }) + }) +} diff --git a/testexec/cli.go b/testexec/cli.go index 32bdc92..b16b832 100644 --- a/testexec/cli.go +++ b/testexec/cli.go @@ -37,7 +37,7 @@ func parseCommandLineArgs() *Config { throttlePtr := flag.Bool("throttle", true, "used to specify whether to throttle PutRecord requests by 1 ms. ") cleanupPtr := flag.Bool("cleanup", true, "used to specify whether or not to delete the kinesis stream after "+ "processing is complete.") - verbosePtr := flag.Bool("verbose", false, "used to specify whether or not to log in verbose mode") + verbosePtr := flag.Bool("verbose", true, "used to specify whether or not to log in verbose mode") // Parse command line arguments flag.Parse() diff --git a/testexec/listenerstats.go b/testexec/listenerstats.go deleted file mode 100644 index 9d05a2d..0000000 --- a/testexec/listenerstats.go +++ /dev/null @@ -1,117 +0,0 @@ -package main - -import ( - "log" - "sync/atomic" - "time" -) - -// ListenerStatsCollector implements the listener's StatsCollector -type ListenerStatsCollector struct { - Consumed uint64 - Delivered uint64 - Processed uint64 - BatchSize uint64 - GetRecordsCalled uint64 - ProvisionedThroughputExceeded uint64 - GetRecordsTimeout uint64 - GetRecordsReadTimeout uint64 - - ProcessedDuration time.Duration - GetRecordsDuration time.Duration - GetRecordsReadResonseDuration time.Duration - GetRecordsUnmarshalDuration time.Duration -} - -// AddConsumed records a count of the number of messages received from AWS -// Kinesis by the listener. -func (lsc *ListenerStatsCollector) AddConsumed(count int) { - atomic.AddUint64(&lsc.Consumed, uint64(count)) -} - -// AddDelivered records a count of the number of messages delivered to the -// application by the listener. -func (lsc *ListenerStatsCollector) AddDelivered(count int) { - atomic.AddUint64(&lsc.Delivered, uint64(count)) -} - -// AddProcessed records a count of the number of messages processed by the -// application by the listener. This is based on a WaitGroup that is sent to -// the RetrieveFn and Listen functions. Retrieve does not count processed -// messages. -func (lsc *ListenerStatsCollector) AddProcessed(count int) { - atomic.AddUint64(&lsc.Processed, uint64(count)) -} - -// AddBatchSize records a count of the number of messages returned by -// GetRecords in the listener. -func (lsc *ListenerStatsCollector) AddBatchSize(count int) { - atomic.AddUint64(&lsc.BatchSize, uint64(count)) -} - -// AddGetRecordsCalled records the number of times the GetRecords API was called -// by the listener. -func (lsc *ListenerStatsCollector) AddGetRecordsCalled(count int) { - atomic.AddUint64(&lsc.GetRecordsCalled, uint64(count)) -} - -// AddProvisionedThroughputExceeded records the number of times the GetRecords -// API returned a ErrCodeProvisionedThroughputExceededException by the listener. -func (lsc *ListenerStatsCollector) AddProvisionedThroughputExceeded(count int) { - atomic.AddUint64(&lsc.ProvisionedThroughputExceeded, uint64(count)) -} - -// AddGetRecordsTimeout records the number of times the GetRecords API timed out -// on the HTTP level. This is influenced by the WithHTTPClientTimeout -// configuration. -func (lsc *ListenerStatsCollector) AddGetRecordsTimeout(count int) { - atomic.AddUint64(&lsc.GetRecordsTimeout, uint64(count)) -} - -// AddGetRecordsReadTimeout records the number of times the GetRecords API timed -// out while reading the response body. This is influenced by the -// WithGetRecordsReadTimeout configuration. -func (lsc *ListenerStatsCollector) AddGetRecordsReadTimeout(count int) { - atomic.AddUint64(&lsc.GetRecordsReadTimeout, uint64(count)) -} - -// AddProcessedDuration records the duration to process a record. See notes on -// AddProcessed. -func (lsc *ListenerStatsCollector) AddProcessedDuration(duration time.Duration) { - // TODO: Not threadsafe - lsc.ProcessedDuration = duration -} - -// AddGetRecordsDuration records the duration that the GetRecords API request -// took. Only the times of successful calls are measured. -func (lsc *ListenerStatsCollector) AddGetRecordsDuration(duration time.Duration) { - // TODO: Not threadsafe - lsc.GetRecordsDuration = duration - -} - -// AddGetRecordsReadResponseDuration records the duration that it took to read -// the response body of a GetRecords API request. -func (lsc *ListenerStatsCollector) AddGetRecordsReadResponseDuration(duration time.Duration) { - // TODO: Not threadsafe - lsc.GetRecordsReadResonseDuration = duration -} - -// AddGetRecordsUnmarshalDuration records the duration that it took to unmarshal -// the response body of a GetRecords API request. -func (lsc *ListenerStatsCollector) AddGetRecordsUnmarshalDuration(duration time.Duration) { - // TODO: Not threadsafe - lsc.GetRecordsUnmarshalDuration = duration -} - -// PrintStats logs stats atomically -func (lsc *ListenerStatsCollector) PrintStats() { - log.Printf("Listener stats: Consumed: [%d]\n", atomic.LoadUint64(&lsc.Consumed)) - log.Printf("Listener stats: Delivered: [%d]\n", atomic.LoadUint64(&lsc.Delivered)) - log.Printf("Listener stats: Processed: [%d]\n", atomic.LoadUint64(&lsc.Processed)) - log.Printf("Listener stats: Batch Size: [%d]\n", atomic.LoadUint64(&lsc.BatchSize)) - log.Printf("Listener stats: GetRecords Called: [%d]\n", atomic.LoadUint64(&lsc.GetRecordsCalled)) - log.Printf("Listener stats: Provisioned Throughput Exceeded: [%d]\n", atomic.LoadUint64(&lsc.ProvisionedThroughputExceeded)) - log.Printf("Listener stats: GetRecords Timeout: [%d]\n", atomic.LoadUint64(&lsc.GetRecordsTimeout)) - log.Printf("Listener stats: GetRecords Read Timeout: [%d]\n", atomic.LoadUint64(&lsc.GetRecordsReadTimeout)) -} diff --git a/testexec/main.go b/testexec/main.go index bc2e2aa..c0be555 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -17,6 +17,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" "github.com/aws/aws-sdk-go/aws/session" + "github.com/rcrowley/go-metrics" "github.com/rewardStyle/kinetic" "github.com/rewardStyle/kinetic/listener" "github.com/rewardStyle/kinetic/message" @@ -49,6 +50,7 @@ var stopDisplay chan struct{} var stopProduce chan struct{} var stopListen chan struct{} var cfg *Config +var registry metrics.Registry func init() { // Set up Http server for pprof @@ -70,6 +72,9 @@ func init() { // Set up pipeOfDeath channel to receive os signals signal.Notify(pipeOfDeath, os.Interrupt) + + // Set up rcrowley metrics registry + registry = metrics.NewRegistry() } func cleanup(k *kinetic.Kinetic, stream string) { @@ -232,7 +237,7 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce log.Println("Creating a kinetic producer ...") } - psc := new(ProducerStatsCollector) + psc := producer.NewDefaultStatsCollector(registry) w, err := producer.NewKinesisWriter(k.Session.Config, streamName, func(kwc *producer.KinesisWriterConfig) { kwc.SetLogLevel(aws.LogDebug) kwc.SetResponseReadTimeout(time.Second) @@ -268,7 +273,7 @@ func newKineticListener(k *kinetic.Kinetic, streamName string) *listener.Listene log.Fatalf("Unable to get shards for stream %s due to: %v\n", streamName, err) } - lsc := new(ListenerStatsCollector) + lsc := listener.NewDefaultStatsCollector(registry) r, err := listener.NewKinesisReader(k.Session.Config, streamName, shards[0], func(krc *listener.KinesisReaderConfig) { krc.SetResponseReadTimeout(1000 * time.Millisecond) @@ -322,10 +327,10 @@ func display(sd *StreamData, p *producer.Producer, l *listener.Listener, wg *syn log.Println() log.Println("***** Stream Data Summary *****") if *cfg.Mode != ModeRead { - p.Stats.(*ProducerStatsCollector).PrintStats() + p.Stats.(*producer.DefaultStatsCollector).PrintStats() } if *cfg.Mode != ModeWrite { - l.Stats.(*ListenerStatsCollector).PrintStats() + l.Stats.(*listener.DefaultStatsCollector).PrintStats() sd.printSummary() } return @@ -336,10 +341,10 @@ func display(sd *StreamData, p *producer.Producer, l *listener.Listener, wg *syn log.Println() log.Println("***** Stream Data Stats *****") if *cfg.Mode != ModeRead { - p.Stats.(*ProducerStatsCollector).PrintStats() + p.Stats.(*producer.DefaultStatsCollector).PrintStats() } if *cfg.Mode != ModeWrite { - l.Stats.(*ListenerStatsCollector).PrintStats() + l.Stats.(*listener.DefaultStatsCollector).PrintStats() sd.printStats() } } @@ -452,10 +457,11 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { } return case <-time.After(time.Second): - newSent := atomic.LoadUint64(&p.Stats.(*ProducerStatsCollector).Sent) - if sent != newSent { + //newSent := atomic.LoadUint64(&p.Stats.(*ProducerStatsCollector).Sent) + newSent := p.Stats.(*producer.DefaultStatsCollector).Sent.Count() + if sent != uint64(newSent) { staleTime.Reset(staleTimeout) - sent = newSent + sent = uint64(newSent) } } } @@ -535,10 +541,10 @@ func listen(sd *StreamData, l *listener.Listener, wg *sync.WaitGroup) { } return case <-time.After(time.Second): - newConsumed := atomic.LoadUint64(&l.Stats.(*ListenerStatsCollector).Consumed) - if consumed != newConsumed { + newConsumed := l.Stats.(*listener.DefaultStatsCollector).Consumed.Count() + if consumed != uint64(newConsumed) { staleTime.Reset(staleTimeout) - consumed = newConsumed + consumed = uint64(newConsumed) } } } diff --git a/testexec/producerstats.go b/testexec/producerstats.go deleted file mode 100644 index 2f420ea..0000000 --- a/testexec/producerstats.go +++ /dev/null @@ -1,117 +0,0 @@ -package main - -import ( - "log" - "sync/atomic" - "time" -) - -// ProducerStatsCollector implements the producer's StatsCollector -type ProducerStatsCollector struct { - Sent uint64 - Failed uint64 - DroppedTotal uint64 - DroppedCapacity uint64 - DroppedRetries uint64 - BatchSize uint64 - PutRecordsProvisionedThroughputExceeded uint64 - PutRecordsCalled uint64 - ProvisionedThroughputExceeded uint64 - PutRecordsTimeout uint64 - - PutRecordsDuration time.Duration - PutRecordsBuildDuration time.Duration - PutRecordsSendDuration time.Duration -} - -// GetSent returns the number of messages sent to AWS Kinesis by the producer. -func (psc *ProducerStatsCollector) GetSent() uint64 { - return atomic.LoadUint64(&psc.Sent) -} - -// AddSent records a count of the number of messages sent to AWS Kinesis by the producer. -func (psc *ProducerStatsCollector) AddSent(count int) { - atomic.AddUint64(&psc.Sent, uint64(count)) -} - -// AddFailed records a count of the number of messages that failed to be sent to AWS Kinesis by the producer. -func (psc *ProducerStatsCollector) AddFailed(count int) { - atomic.AddUint64(&psc.Failed, uint64(count)) -} - -// AddDroppedTotal records a count of the total number of messages dropped by the application after multiple failures. -func (psc *ProducerStatsCollector) AddDroppedTotal(count int) { - atomic.AddUint64(&psc.DroppedTotal, uint64(count)) -} - -// AddDroppedCapacity records a count of the number of messages that were dropped by the application due to the stream -// writer being at capacity. -func (psc *ProducerStatsCollector) AddDroppedCapacity(count int) { - atomic.AddUint64(&psc.DroppedCapacity, uint64(count)) -} - -// AddDroppedRetries records a count of the number of retry messages dropped by the application after the max number of -// retries was exceeded. -func (psc *ProducerStatsCollector) AddDroppedRetries(count int) { - atomic.AddUint64(&psc.DroppedRetries, uint64(count)) -} - -// AddBatchSize records a count of the number of messages attempted by PutRecords in the producer. -func (psc *ProducerStatsCollector) AddBatchSize(count int) { - atomic.AddUint64(&psc.BatchSize, uint64(count)) -} - -// AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords API returned a -// ErrCodeProvisionedThroughputExceededException by the producer. -func (psc *ProducerStatsCollector) AddPutRecordsProvisionedThroughputExceeded(count int) { - atomic.AddUint64(&psc.PutRecordsProvisionedThroughputExceeded, uint64(count)) -} - -// AddPutRecordsCalled records the number of times the PutRecords API was called by the producer. -func (psc *ProducerStatsCollector) AddPutRecordsCalled(count int) { - atomic.AddUint64(&psc.PutRecordsCalled, uint64(count)) -} - -// AddProvisionedThroughputExceeded records the number of times the PutRecords API response contained a record which -// contained an ErrCodeProvisionedThroughputExceededException error. -func (psc *ProducerStatsCollector) AddProvisionedThroughputExceeded(count int) { - atomic.AddUint64(&psc.ProvisionedThroughputExceeded, uint64(count)) -} - -// AddPutRecordsTimeout records the number of times the PutRecords API timed out on the HTTP level. This is influenced -// by the WithHTTPClientTimeout configuration. -func (psc *ProducerStatsCollector) AddPutRecordsTimeout(count int) { - atomic.AddUint64(&psc.PutRecordsTimeout, uint64(count)) -} - -// AddPutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls -// are measured. -func (psc *ProducerStatsCollector) AddPutRecordsDuration(duration time.Duration) { - // TODO: Not threadsafe - psc.PutRecordsDuration = duration -} - -// AddPutRecordsBuildDuration records the duration that it took to build the PutRecords API request payload. -func (psc *ProducerStatsCollector) AddPutRecordsBuildDuration(duration time.Duration) { - // TODO: Not threadsafe - psc.PutRecordsBuildDuration = duration -} - -// AddPutRecordsSendDuration records the duration that it took to send the PutRecords API request payload. -func (psc *ProducerStatsCollector) AddPutRecordsSendDuration(duration time.Duration) { - // TODO: Not threadsafe - psc.PutRecordsSendDuration = duration -} - -// PrintStats logs stats atomically -func (psc *ProducerStatsCollector) PrintStats() { - log.Printf("Producer Stats: Sent: [%d]\n", atomic.LoadUint64(&psc.Sent)) - log.Printf("Producer Stats: Failed: [%d]\n", atomic.LoadUint64(&psc.Failed)) - log.Printf("Producer Stats: Dropped Total: [%d]\n", atomic.LoadUint64(&psc.DroppedTotal)) - log.Printf("Producer Stats: Dropped Retries: [%d]\n", atomic.LoadUint64(&psc.DroppedRetries)) - log.Printf("Producer Stats: Dropped Capacity: [%d]\n", atomic.LoadUint64(&psc.DroppedCapacity)) - log.Printf("Producer Stats: Batch Size: [%d]\n", atomic.LoadUint64(&psc.BatchSize)) - log.Printf("Producer Stats: PutRecords Called: [%d]\n", atomic.LoadUint64(&psc.PutRecordsCalled)) - log.Printf("Producer Stats: Provisioned Throughput Exceeded: [%d]\n", atomic.LoadUint64(&psc.ProvisionedThroughputExceeded)) - log.Printf("Producer Stats: PutRecords Timeout: [%d]\n", atomic.LoadUint64(&psc.PutRecordsTimeout)) -} From 1d4a236f3ee6c01b3a69d77d3fb0e64f35361b3e Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 10 Jul 2017 15:27:00 -0500 Subject: [PATCH 048/109] Fixed a bug with the producer's stats collection for the Sent and Failed counters by moving the stats collection into the writers. --- producer/firehosewriter.go | 6 +++--- producer/kinesiswriter.go | 2 ++ producer/producer.go | 6 +----- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go index 07abc34..0868982 100644 --- a/producer/firehosewriter.go +++ b/producer/firehosewriter.go @@ -104,14 +104,12 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Mes sent := attempted - failed w.LogDebug(fmt.Sprintf("Finished PutRecords request, %d records attempted, %d records successful, %d records failed, took %v\n", attempted, sent, failed, time.Since(start))) - var retries int for idx, record := range resp.RequestResponses { if record.RecordId != nil { // TODO: per-shard metrics messages[idx].RecordID = record.RecordId + w.Stats.AddSent(1) } else { - retries++ - switch aws.StringValue(record.ErrorCode) { case firehose.ErrCodeLimitExceededException: w.Stats.AddProvisionedThroughputExceeded(1) @@ -120,6 +118,8 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Mes } messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage + messages[idx].FailCount++ + w.Stats.AddFailed(1) go fn(messages[idx]) } diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index 236a66a..f1b621b 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -111,6 +111,7 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess // TODO: per-shard metrics messages[idx].SequenceNumber = record.SequenceNumber messages[idx].ShardID = record.ShardId + w.Stats.AddSent(1) } else { switch aws.StringValue(record.ErrorCode) { case kinesis.ErrCodeProvisionedThroughputExceededException: @@ -121,6 +122,7 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage messages[idx].FailCount++ + w.Stats.AddFailed(1) go fn(messages[idx]) } diff --git a/producer/producer.go b/producer/producer.go index b3f2fe0..2e1e5a6 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -111,18 +111,16 @@ func (p *Producer) sendBatch(batch []*message.Message) { }() var attempts uint64 - var failed uint64 stop: for { err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message) error { if msg.FailCount <= p.maxRetryAttempts { // Apply a delay before retrying - time.Sleep(time.Duration(msg.FailCount*msg.FailCount) * time.Second) + time.Sleep(time.Duration(msg.FailCount * msg.FailCount) * time.Second) select { case p.retries <- msg: - atomic.AddUint64(&failed, 1) case <-p.pipeOfDeath: return errs.ErrPipeOfDeath } @@ -133,8 +131,6 @@ stop: return nil }) - p.Stats.AddSent(len(batch) - int(failed)) - p.Stats.AddFailed(int(failed)) if err == nil { break stop } From 99cb2e879c6045f13cd45011378a5ed73ef9fab6 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 10 Jul 2017 17:46:38 -0500 Subject: [PATCH 049/109] Converted magic strings for stats into constants and added missing stats in PrintStats. --- listener/stats.go | 47 ++++++++++++++++-------- producer/stats.go | 92 ++++++++++++++++++++++++++++------------------- 2 files changed, 89 insertions(+), 50 deletions(-) diff --git a/listener/stats.go b/listener/stats.go index 14c913f..29a43bf 100644 --- a/listener/stats.go +++ b/listener/stats.go @@ -80,6 +80,22 @@ func (nsc *NilStatsCollector) AddGetRecordsReadResponseDuration(time.Duration) { // the response body of a GetRecords API request. func (nsc *NilStatsCollector) AddGetRecordsUnmarshalDuration(time.Duration) {} +// Metric names to be exported +const ( + MetricsConsumed = "kinetic.listener.consumed" + MetricsDelivered = "kinetic.listener.delivered" + MetricsProcessed = "kinetic.listener.processed" + MetricsBatchSize = "kinetic.listener.batchsize" + MetricsSent = "kinetic.listener.sent" + MetricsProvisionedThroughputExceeded = "kinetic.listener.provisionedthroughputexceeded" + MetricsGetRecordsTimeout = "kinetic.listener.getrecords.timeout" + MetricsGetRecordsReadTimeout = "kinetic.listener.getrecords.readtimeout" + MetricsProcessedDuration = "kinetic.listener.processed.duration" + MetricsGetRecordsDuration = "kinetic.listener.getrecords.duration" + MetricsGetRecordsReadResponseDuration = "kinetic.listener.getrecords.readresponse.duration" + MetricsGetRecordsUnmarshalDuration = "kinetic.listener.getrecords.unmarshal.duration" +) + // DefaultStatsCollector is a type that implements the listener's StatsCollector interface using the // rcrowley/go-metrics library type DefaultStatsCollector struct { @@ -100,18 +116,18 @@ type DefaultStatsCollector struct { // NewDefaultStatsCollector instantiates a new DefaultStatsCollector object func NewDefaultStatsCollector(r metrics.Registry) *DefaultStatsCollector { return &DefaultStatsCollector{ - Consumed: metrics.GetOrRegisterCounter("listener.consumed", r), - Delivered: metrics.GetOrRegisterCounter("listener.delivered", r), - Processed: metrics.GetOrRegisterCounter("listener.processed", r), - BatchSize: metrics.GetOrRegisterCounter("listener.batchsize", r), - GetRecordsCalled: metrics.GetOrRegisterCounter("listener.sent", r), - ProvisionedThroughputExceeded: metrics.GetOrRegisterCounter("listener.provisionedthroughputexceeded", r), - GetRecordsTimeout: metrics.GetOrRegisterCounter("listener.getrecords.timeout", r), - GetRecordsReadTimeout: metrics.GetOrRegisterCounter("listener.getrecords.readtimeout", r), - ProcessedDuration: metrics.GetOrRegisterGauge("listener.processed.duration", r), - GetRecordsDuration: metrics.GetOrRegisterGauge("listener.getrecords.duration", r), - GetRecordsReadResponseDuration: metrics.GetOrRegisterGauge("listener.getrecords.readresponse.duration", r), - GetRecordsUnmarshalDuration: metrics.GetOrRegisterGauge("listener.getrecords.unmarshal.duration", r), + Consumed: metrics.GetOrRegisterCounter(MetricsConsumed, r), + Delivered: metrics.GetOrRegisterCounter(MetricsDelivered, r), + Processed: metrics.GetOrRegisterCounter(MetricsProcessed, r), + BatchSize: metrics.GetOrRegisterCounter(MetricsBatchSize, r), + GetRecordsCalled: metrics.GetOrRegisterCounter(MetricsSent, r), + ProvisionedThroughputExceeded: metrics.GetOrRegisterCounter(MetricsProvisionedThroughputExceeded, r), + GetRecordsTimeout: metrics.GetOrRegisterCounter(MetricsGetRecordsTimeout, r), + GetRecordsReadTimeout: metrics.GetOrRegisterCounter(MetricsGetRecordsReadTimeout, r), + ProcessedDuration: metrics.GetOrRegisterGauge(MetricsProcessedDuration, r), + GetRecordsDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsDuration, r), + GetRecordsReadResponseDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsReadResponseDuration, r), + GetRecordsUnmarshalDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsUnmarshalDuration, r), } } @@ -198,8 +214,11 @@ func (dsc *DefaultStatsCollector) PrintStats() { log.Printf("Listener stats: Processed: [%d]\n", dsc.Processed.Count()) log.Printf("Listener stats: Batch Size: [%d]\n", dsc.BatchSize.Count()) log.Printf("Listener stats: GetRecords Called: [%d]\n", dsc.GetRecordsCalled.Count()) - log.Printf("Listener stats: Provisioned Throughput Exceeded: [%d]\n", dsc.ProvisionedThroughputExceeded.Count()) log.Printf("Listener stats: GetRecords Timeout: [%d]\n", dsc.GetRecordsTimeout.Count()) log.Printf("Listener stats: GetRecords Read Timeout: [%d]\n", dsc.GetRecordsReadTimeout.Count()) + log.Printf("Listener stats: Provisioned Throughput Exceeded: [%d]\n", dsc.ProvisionedThroughputExceeded.Count()) + log.Printf("Listener stats: ProcessedDuration (ns): [%d]\n", dsc.ProcessedDuration.Value()) + log.Printf("Listener stats: GetRecordsDuration (ns): [%d]\n", dsc.GetRecordsDuration.Value()) + log.Printf("Listener stats: GetRecordsReadResponseDuration (ns): [%d]\n", dsc.GetRecordsReadResponseDuration.Value()) + log.Printf("Listener stats: GetRecordsUnmarshalDuration (ns): [%d]\n", dsc.GetRecordsUnmarshalDuration.Value()) } - diff --git a/producer/stats.go b/producer/stats.go index 657149a..115105d 100644 --- a/producer/stats.go +++ b/producer/stats.go @@ -17,10 +17,10 @@ type StatsCollector interface { AddDroppedCapacity(int) AddDroppedRetries(int) AddBatchSize(int) - AddPutRecordsProvisionedThroughputExceeded(int) AddPutRecordsCalled(int) - AddProvisionedThroughputExceeded(int) AddPutRecordsTimeout(int) + AddProvisionedThroughputExceeded(int) + AddPutRecordsProvisionedThroughputExceeded(int) AddPutRecordsDuration(time.Duration) AddPutRecordsBuildDuration(time.Duration) AddPutRecordsSendDuration(time.Duration) @@ -49,20 +49,20 @@ func (nsc *NilStatsCollector) AddDroppedRetries(int) {} // AddBatchSize records a count of the number of messages attempted by PutRecords in the producer. func (nsc *NilStatsCollector) AddBatchSize(int) {} -// AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords API returned a -// ErrCodeProvisionedThroughputExceededException by the producer. -func (nsc *NilStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} - // AddPutRecordsCalled records the number of times the PutRecords API was called by the producer. func (nsc *NilStatsCollector) AddPutRecordsCalled(int) {} +// AddPutRecordsTimeout records the number of times the PutRecords API timed out on the HTTP level. This is influenced +// by the WithHTTPClientTimeout configuration. +func (nsc *NilStatsCollector) AddPutRecordsTimeout(int) {} + // AddProvisionedThroughputExceeded records the number of times the PutRecords API response contained a record which // contained an ErrCodeProvisionedThroughputExceededException error. func (nsc *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} -// AddPutRecordsTimeout records the number of times the PutRecords API timed out on the HTTP level. This is influenced -// by the WithHTTPClientTimeout configuration. -func (nsc *NilStatsCollector) AddPutRecordsTimeout(int) {} +// AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords API returned a +// ErrCodeProvisionedThroughputExceededException by the producer. +func (nsc *NilStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} // AddPutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls // are measured. @@ -74,6 +74,23 @@ func (nsc *NilStatsCollector) AddPutRecordsBuildDuration(time.Duration) {} // AddPutRecordsSendDuration records the duration that it took to send the PutRecords API request payload. func (nsc *NilStatsCollector) AddPutRecordsSendDuration(time.Duration) {} +// Metric names to be exported +const ( + MetricsSent = "kinetic.producer.sent" + MetricsFailed = "kinetic.producer.failed" + MetricsDroppedTotal = "kinetic.producer.dropped.total" + MetricsDroppedCapacity = "kinetic.producer.dropped.capacity" + MetricsDroppedRetries = "kinetic.producer.dropped.retries" + MetricsBatchSize = "kinetic.producer.batchsize" + MetricsPutRecordsCalled = "kinetic.producer.putrecords.called" + MetricsPutRecordsTimeout = "kinetic.producer.putrecords.timeout" + MetricsProvisionedThroughputExceeded = "kinetic.producer.provisionedthroughputexceeded" + MetricsPutRecordsProvisionedThroughputExceeded = "kinetic.producer.putrecords.provisionedthroughputexceeded" + MetricsPutRecordsDuration = "kinetic.producer.putrecords.duration" + MetricsPutRecordsBuildDuration = "kinetic.producer.putrecords.build.duration" + MetricsPutRecordsSendDuration = "kinetic.producer.putrecords.send.duration" +) + // DefaultStatsCollector is a type that implements the producers's StatsCollector interface using the // rcrowley/go-metrics library type DefaultStatsCollector struct { @@ -83,10 +100,10 @@ type DefaultStatsCollector struct { DroppedCapacity metrics.Counter DroppedRetries metrics.Counter BatchSize metrics.Counter - PutRecordsProvisionedThroughputExceeded metrics.Counter PutRecordsCalled metrics.Counter - ProvisionedThroughputExceeded metrics.Counter PutRecordsTimeout metrics.Counter + ProvisionedThroughputExceeded metrics.Counter + PutRecordsProvisionedThroughputExceeded metrics.Counter PutRecordsDuration metrics.Gauge PutRecordsBuildDuration metrics.Gauge PutRecordsSendDuration metrics.Gauge @@ -95,19 +112,19 @@ type DefaultStatsCollector struct { // NewDefaultStatsCollector instantiates a new DefaultStatsCollector object func NewDefaultStatsCollector(r metrics.Registry) *DefaultStatsCollector { return &DefaultStatsCollector{ - Sent: metrics.GetOrRegisterCounter("producer.sent", r), - Failed: metrics.GetOrRegisterCounter("producer.failed", r), - DroppedTotal: metrics.GetOrRegisterCounter("producer.dropped.total", r), - DroppedCapacity: metrics.GetOrRegisterCounter("producer.dropped.capacity", r), - DroppedRetries: metrics.GetOrRegisterCounter("producer.dropped.retries", r), - BatchSize: metrics.GetOrRegisterCounter("producer.batchsize", r), - PutRecordsProvisionedThroughputExceeded: metrics.GetOrRegisterCounter("producer.putrecords.provisionedthroughputexceeded", r), - PutRecordsCalled: metrics.GetOrRegisterCounter("producer.putrecords.called", r), - ProvisionedThroughputExceeded: metrics.GetOrRegisterCounter("producer.provisionedthroughputexceeded", r), - PutRecordsTimeout: metrics.GetOrRegisterCounter("producer.putrecords.timeout", r), - PutRecordsDuration: metrics.GetOrRegisterGauge("producer.putrecords.duration", r), - PutRecordsBuildDuration: metrics.GetOrRegisterGauge("producer.putrecords.build.duration", r), - PutRecordsSendDuration: metrics.GetOrRegisterGauge("producer.putrecords.send.duration", r), + Sent: metrics.GetOrRegisterCounter(MetricsSent, r), + Failed: metrics.GetOrRegisterCounter(MetricsFailed, r), + DroppedTotal: metrics.GetOrRegisterCounter(MetricsDroppedTotal, r), + DroppedCapacity: metrics.GetOrRegisterCounter(MetricsDroppedCapacity, r), + DroppedRetries: metrics.GetOrRegisterCounter(MetricsDroppedRetries, r), + BatchSize: metrics.GetOrRegisterCounter(MetricsBatchSize, r), + PutRecordsCalled: metrics.GetOrRegisterCounter(MetricsPutRecordsCalled, r), + PutRecordsTimeout: metrics.GetOrRegisterCounter(MetricsPutRecordsTimeout, r), + ProvisionedThroughputExceeded: metrics.GetOrRegisterCounter(MetricsProvisionedThroughputExceeded, r), + PutRecordsProvisionedThroughputExceeded: metrics.GetOrRegisterCounter(MetricsPutRecordsProvisionedThroughputExceeded, r), + PutRecordsDuration: metrics.GetOrRegisterGauge(MetricsPutRecordsDuration, r), + PutRecordsBuildDuration: metrics.GetOrRegisterGauge(MetricsPutRecordsBuildDuration, r), + PutRecordsSendDuration: metrics.GetOrRegisterGauge(MetricsPutRecordsSendDuration, r), } } @@ -143,27 +160,27 @@ func (dsc *DefaultStatsCollector) AddBatchSize(count int) { dsc.BatchSize.Inc(int64(count)) } -// AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords API returned a -// ErrCodeProvisionedThroughputExceededException by the producer. -func (dsc *DefaultStatsCollector) AddPutRecordsProvisionedThroughputExceeded(count int) { - dsc.PutRecordsProvisionedThroughputExceeded.Inc(int64(count)) -} - // AddPutRecordsCalled records the number of times the PutRecords API was called by the producer. func (dsc *DefaultStatsCollector) AddPutRecordsCalled(count int) { dsc.PutRecordsCalled.Inc(int64(count)) } +// AddPutRecordsTimeout records the number of times the PutRecords API timed out on the HTTP level. This is influenced +// by the WithHTTPClientTimeout configuration. +func (dsc *DefaultStatsCollector) AddPutRecordsTimeout(count int) { + dsc.PutRecordsTimeout.Inc(int64(count)) +} + // AddProvisionedThroughputExceeded records the number of times the PutRecords API response contained a record which // contained an ErrCodeProvisionedThroughputExceededException error. func (dsc *DefaultStatsCollector) AddProvisionedThroughputExceeded(count int) { dsc.ProvisionedThroughputExceeded.Inc(int64(count)) } -// AddPutRecordsTimeout records the number of times the PutRecords API timed out on the HTTP level. This is influenced -// by the WithHTTPClientTimeout configuration. -func (dsc *DefaultStatsCollector) AddPutRecordsTimeout(count int) { - dsc.PutRecordsTimeout.Inc(int64(count)) +// AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords API returned a +// ErrCodeProvisionedThroughputExceededException by the producer. +func (dsc *DefaultStatsCollector) AddPutRecordsProvisionedThroughputExceeded(count int) { + dsc.PutRecordsProvisionedThroughputExceeded.Inc(int64(count)) } // AddPutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls @@ -191,7 +208,10 @@ func (dsc *DefaultStatsCollector) PrintStats() { log.Printf("Producer Stats: Dropped Capacity: [%d]\n", dsc.DroppedCapacity.Count()) log.Printf("Producer Stats: Batch Size: [%d]\n", dsc.BatchSize.Count()) log.Printf("Producer Stats: PutRecords Called: [%d]\n", dsc.PutRecordsCalled.Count()) - log.Printf("Producer Stats: Provisioned Throughput Exceeded: [%d]\n", dsc.ProvisionedThroughputExceeded.Count()) log.Printf("Producer Stats: PutRecords Timeout: [%d]\n", dsc.PutRecordsTimeout.Count()) + log.Printf("Producer Stats: Provisioned Throughput Exceeded: [%d]\n", dsc.ProvisionedThroughputExceeded.Count()) + log.Printf("Producer Stats: PutRecords Provisioned Throughput Exceeded: [%d]\n", dsc.PutRecordsProvisionedThroughputExceeded.Count()) + log.Printf("Producer Stats: PutRecords Duration (ns): [%d]\n", dsc.PutRecordsDuration.Value()) + log.Printf("Producer Stats: PutRecords Build Duration (ns): [%d]\n", dsc.PutRecordsBuildDuration.Value()) + log.Printf("Producer Stats: PutRecords Send Duration (ns): [%d]\n", dsc.PutRecordsSendDuration.Value()) } - From 2e20f0c0500684a57479625b77826c6322aee651 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 11 Jul 2017 17:50:06 -0500 Subject: [PATCH 050/109] Renamed the producer stats to make more sense. --- producer/firehosewriter.go | 4 +- producer/kinesiswriter.go | 4 +- producer/producer.go | 4 +- producer/producerconfig_test.go | 6 +-- producer/stats.go | 66 ++++++++++++++++----------------- producer/stats_test.go | 40 ++++++++++---------- testexec/main.go | 3 +- 7 files changed, 63 insertions(+), 64 deletions(-) diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go index 0868982..3f3c945 100644 --- a/producer/firehosewriter.go +++ b/producer/firehosewriter.go @@ -108,7 +108,7 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Mes if record.RecordId != nil { // TODO: per-shard metrics messages[idx].RecordID = record.RecordId - w.Stats.AddSent(1) + w.Stats.AddSentSuccess(1) } else { switch aws.StringValue(record.ErrorCode) { case firehose.ErrCodeLimitExceededException: @@ -119,7 +119,7 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Mes messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage messages[idx].FailCount++ - w.Stats.AddFailed(1) + w.Stats.AddSentFailed(1) go fn(messages[idx]) } diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index f1b621b..4c7c7a8 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -111,7 +111,7 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess // TODO: per-shard metrics messages[idx].SequenceNumber = record.SequenceNumber messages[idx].ShardID = record.ShardId - w.Stats.AddSent(1) + w.Stats.AddSentSuccess(1) } else { switch aws.StringValue(record.ErrorCode) { case kinesis.ErrCodeProvisionedThroughputExceededException: @@ -122,7 +122,7 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage messages[idx].FailCount++ - w.Stats.AddFailed(1) + w.Stats.AddSentFailed(1) go fn(messages[idx]) } diff --git a/producer/producer.go b/producer/producer.go index 2e1e5a6..69b3df7 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -117,7 +117,7 @@ stop: err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message) error { if msg.FailCount <= p.maxRetryAttempts { // Apply a delay before retrying - time.Sleep(time.Duration(msg.FailCount * msg.FailCount) * time.Second) + time.Sleep(time.Duration(msg.FailCount*msg.FailCount) * time.Second) select { case p.retries <- msg: @@ -230,7 +230,7 @@ func (p *Producer) produce() { p.shutdownCond.L.Lock() if len(batch) > 0 { p.outstanding++ - p.Stats.AddBatchSize(len(batch)) + p.Stats.AddSentTotal(len(batch)) p.concurrencySem <- empty{} go p.sendBatch(batch) } else if len(batch) == 0 { diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go index ab40742..22889dd 100644 --- a/producer/producerconfig_test.go +++ b/producer/producerconfig_test.go @@ -14,12 +14,12 @@ import ( type DebugStatsCollector struct{} -func (l *DebugStatsCollector) AddSent(int) {} -func (l *DebugStatsCollector) AddFailed(int) {} +func (l *DebugStatsCollector) AddSentTotal(int) {} +func (l *DebugStatsCollector) AddSentSuccess(int) {} +func (l *DebugStatsCollector) AddSentFailed(int) {} func (l *DebugStatsCollector) AddDroppedTotal(int) {} func (l *DebugStatsCollector) AddDroppedCapacity(int) {} func (l *DebugStatsCollector) AddDroppedRetries(int) {} -func (l *DebugStatsCollector) AddBatchSize(int) {} func (l *DebugStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} func (l *DebugStatsCollector) AddPutRecordsCalled(int) {} func (l *DebugStatsCollector) AddProvisionedThroughputExceeded(int) {} diff --git a/producer/stats.go b/producer/stats.go index 115105d..664fb58 100644 --- a/producer/stats.go +++ b/producer/stats.go @@ -11,12 +11,12 @@ import ( // the Kinetic producer library. This was really built with rcrowley/go-metrics // in mind. type StatsCollector interface { - AddSent(int) - AddFailed(int) + AddSentTotal(int) + AddSentSuccess(int) + AddSentFailed(int) AddDroppedTotal(int) AddDroppedCapacity(int) AddDroppedRetries(int) - AddBatchSize(int) AddPutRecordsCalled(int) AddPutRecordsTimeout(int) AddProvisionedThroughputExceeded(int) @@ -29,11 +29,14 @@ type StatsCollector interface { // NilStatsCollector is a stats listener that ignores all metrics. type NilStatsCollector struct{} -// AddSent records a count of the number of messages sent to AWS Kinesis by the producer. -func (nsc *NilStatsCollector) AddSent(int) {} +// AddSentTotal records a count of the total number of messages attempted by PutRecords in the producer. +func (nsc *NilStatsCollector) AddSentTotal(int) {} -// AddFailed records a count of the number of messages that failed to be sent to AWS Kinesis by the producer. -func (nsc *NilStatsCollector) AddFailed(int) {} +// AddSentSuccess records a count of the number of messages sent successfully to AWS Kinesis by the producer. +func (nsc *NilStatsCollector) AddSentSuccess(int) {} + +// AddSentFailed records a count of the number of messages that failed to be sent to AWS Kinesis by the producer. +func (nsc *NilStatsCollector) AddSentFailed(int) {} // AddDroppedTotal records a count of the total number of messages dropped by the application after multiple failures. func (nsc *NilStatsCollector) AddDroppedTotal(int) {} @@ -46,9 +49,6 @@ func (nsc *NilStatsCollector) AddDroppedCapacity(int) {} // retries was exceeded. func (nsc *NilStatsCollector) AddDroppedRetries(int) {} -// AddBatchSize records a count of the number of messages attempted by PutRecords in the producer. -func (nsc *NilStatsCollector) AddBatchSize(int) {} - // AddPutRecordsCalled records the number of times the PutRecords API was called by the producer. func (nsc *NilStatsCollector) AddPutRecordsCalled(int) {} @@ -76,12 +76,12 @@ func (nsc *NilStatsCollector) AddPutRecordsSendDuration(time.Duration) {} // Metric names to be exported const ( - MetricsSent = "kinetic.producer.sent" - MetricsFailed = "kinetic.producer.failed" + MetricsSentTotal = "kinetic.producer.sent.total" + MetricsSentSuccess = "kinetic.producer.sent.success" + MetricsSentFailed = "kinetic.producer.sent.failed" MetricsDroppedTotal = "kinetic.producer.dropped.total" MetricsDroppedCapacity = "kinetic.producer.dropped.capacity" MetricsDroppedRetries = "kinetic.producer.dropped.retries" - MetricsBatchSize = "kinetic.producer.batchsize" MetricsPutRecordsCalled = "kinetic.producer.putrecords.called" MetricsPutRecordsTimeout = "kinetic.producer.putrecords.timeout" MetricsProvisionedThroughputExceeded = "kinetic.producer.provisionedthroughputexceeded" @@ -94,12 +94,12 @@ const ( // DefaultStatsCollector is a type that implements the producers's StatsCollector interface using the // rcrowley/go-metrics library type DefaultStatsCollector struct { - Sent metrics.Counter - Failed metrics.Counter + SentTotal metrics.Counter + SentSuccess metrics.Counter + SentFailed metrics.Counter DroppedTotal metrics.Counter DroppedCapacity metrics.Counter DroppedRetries metrics.Counter - BatchSize metrics.Counter PutRecordsCalled metrics.Counter PutRecordsTimeout metrics.Counter ProvisionedThroughputExceeded metrics.Counter @@ -112,12 +112,12 @@ type DefaultStatsCollector struct { // NewDefaultStatsCollector instantiates a new DefaultStatsCollector object func NewDefaultStatsCollector(r metrics.Registry) *DefaultStatsCollector { return &DefaultStatsCollector{ - Sent: metrics.GetOrRegisterCounter(MetricsSent, r), - Failed: metrics.GetOrRegisterCounter(MetricsFailed, r), + SentTotal: metrics.GetOrRegisterCounter(MetricsSentTotal, r), + SentSuccess: metrics.GetOrRegisterCounter(MetricsSentSuccess, r), + SentFailed: metrics.GetOrRegisterCounter(MetricsSentFailed, r), DroppedTotal: metrics.GetOrRegisterCounter(MetricsDroppedTotal, r), DroppedCapacity: metrics.GetOrRegisterCounter(MetricsDroppedCapacity, r), DroppedRetries: metrics.GetOrRegisterCounter(MetricsDroppedRetries, r), - BatchSize: metrics.GetOrRegisterCounter(MetricsBatchSize, r), PutRecordsCalled: metrics.GetOrRegisterCounter(MetricsPutRecordsCalled, r), PutRecordsTimeout: metrics.GetOrRegisterCounter(MetricsPutRecordsTimeout, r), ProvisionedThroughputExceeded: metrics.GetOrRegisterCounter(MetricsProvisionedThroughputExceeded, r), @@ -128,14 +128,19 @@ func NewDefaultStatsCollector(r metrics.Registry) *DefaultStatsCollector { } } -// AddSent records a count of the number of messages sent to AWS Kinesis by the producer. -func (dsc *DefaultStatsCollector) AddSent(count int) { - dsc.Sent.Inc(int64(count)) +// AddSentTotal records a count of the total number of messages attempted by PutRecords in the producer. +func (dsc *DefaultStatsCollector) AddSentTotal(count int) { + dsc.SentTotal.Inc(int64(count)) +} + +// AddSentSuccess records a count of the number of messages sent successfully to AWS Kinesis by the producer. +func (dsc *DefaultStatsCollector) AddSentSuccess(count int) { + dsc.SentSuccess.Inc(int64(count)) } -// AddFailed records a count of the number of messages that failed to be sent to AWS Kinesis by the producer. -func (dsc *DefaultStatsCollector) AddFailed(count int) { - dsc.Failed.Inc(int64(count)) +// AddSentFailed records a count of the number of messages that failed to be sent to AWS Kinesis by the producer. +func (dsc *DefaultStatsCollector) AddSentFailed(count int) { + dsc.SentFailed.Inc(int64(count)) } // AddDroppedTotal records a count of the total number of messages dropped by the application after multiple failures. @@ -155,11 +160,6 @@ func (dsc *DefaultStatsCollector) AddDroppedRetries(count int) { dsc.DroppedRetries.Inc(int64(count)) } -// AddBatchSize records a count of the number of messages attempted by PutRecords in the producer. -func (dsc *DefaultStatsCollector) AddBatchSize(count int) { - dsc.BatchSize.Inc(int64(count)) -} - // AddPutRecordsCalled records the number of times the PutRecords API was called by the producer. func (dsc *DefaultStatsCollector) AddPutRecordsCalled(count int) { dsc.PutRecordsCalled.Inc(int64(count)) @@ -201,12 +201,12 @@ func (dsc *DefaultStatsCollector) AddPutRecordsSendDuration(duration time.Durati // PrintStats logs the stats func (dsc *DefaultStatsCollector) PrintStats() { - log.Printf("Producer Stats: Sent: [%d]\n", dsc.Sent.Count()) - log.Printf("Producer Stats: Failed: [%d]\n", dsc.Failed.Count()) + log.Printf("Producer Stats: Sent Total: [%d]\n", dsc.SentTotal.Count()) + log.Printf("Producer Stats: Sent Success: [%d]\n", dsc.SentSuccess.Count()) + log.Printf("Producer Stats: Sent Failed: [%d]\n", dsc.SentFailed.Count()) log.Printf("Producer Stats: Dropped Total: [%d]\n", dsc.DroppedTotal.Count()) log.Printf("Producer Stats: Dropped Retries: [%d]\n", dsc.DroppedRetries.Count()) log.Printf("Producer Stats: Dropped Capacity: [%d]\n", dsc.DroppedCapacity.Count()) - log.Printf("Producer Stats: Batch Size: [%d]\n", dsc.BatchSize.Count()) log.Printf("Producer Stats: PutRecords Called: [%d]\n", dsc.PutRecordsCalled.Count()) log.Printf("Producer Stats: PutRecords Timeout: [%d]\n", dsc.PutRecordsTimeout.Count()) log.Printf("Producer Stats: Provisioned Throughput Exceeded: [%d]\n", dsc.ProvisionedThroughputExceeded.Count()) diff --git a/producer/stats_test.go b/producer/stats_test.go index 953afe0..dbff301 100644 --- a/producer/stats_test.go +++ b/producer/stats_test.go @@ -18,12 +18,16 @@ func TestStatsCollector(t *testing.T) { var sc StatsCollector = &NilStatsCollector{} So(sc, ShouldNotBeNil) - Convey("check that AddSent does not error", func() { - sc.AddSent(1) + Convey("check that AddSentTotal does not error", func() { + sc.AddSentTotal(1) }) - Convey("check that AddFailed does not error", func() { - sc.AddFailed(1) + Convey("check that AddSentSuccess does not error", func() { + sc.AddSentSuccess(1) + }) + + Convey("check that AddSentFailed does not error", func() { + sc.AddSentFailed(1) }) Convey("check that AddDroppedTotal does not error", func() { @@ -38,10 +42,6 @@ func TestStatsCollector(t *testing.T) { sc.AddDroppedRetries(1) }) - Convey("check that AddBatchSize does not error", func() { - sc.AddBatchSize(1) - }) - Convey("check that AddPutRecordsProvisionedThroughputExceeded does not error", func() { sc.AddPutRecordsProvisionedThroughputExceeded(1) }) @@ -76,16 +76,22 @@ func TestStatsCollector(t *testing.T) { var sc StatsCollector = NewDefaultStatsCollector(r) So(sc, ShouldNotBeNil) - Convey("check that AddSent does not error", func() { + Convey("check that AddSentTotal does not error", func() { + count := rand.Int() + sc.AddSentTotal(count) + So(sc.(*DefaultStatsCollector).SentTotal.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddSentSuccess does not error", func() { count := rand.Int() - sc.AddSent(count) - So(sc.(*DefaultStatsCollector).Sent.Count(), ShouldEqual, int64(count)) + sc.AddSentSuccess(count) + So(sc.(*DefaultStatsCollector).SentSuccess.Count(), ShouldEqual, int64(count)) }) - Convey("check that AddFailed does not error", func() { + Convey("check that AddSentFailed does not error", func() { count := rand.Int() - sc.AddFailed(count) - So(sc.(*DefaultStatsCollector).Failed.Count(), ShouldEqual, int64(count)) + sc.AddSentFailed(count) + So(sc.(*DefaultStatsCollector).SentFailed.Count(), ShouldEqual, int64(count)) }) Convey("check that AddDroppedTotal does not error", func() { @@ -106,12 +112,6 @@ func TestStatsCollector(t *testing.T) { So(sc.(*DefaultStatsCollector).DroppedRetries.Count(), ShouldEqual, int64(count)) }) - Convey("check that AddBatchSize does not error", func() { - count := rand.Int() - sc.AddBatchSize(count) - So(sc.(*DefaultStatsCollector).BatchSize.Count(), ShouldEqual, int64(count)) - }) - Convey("check that AddPutRecordsProvisionedThroughputExceeded does not error", func() { count := rand.Int() sc.AddPutRecordsProvisionedThroughputExceeded(count) diff --git a/testexec/main.go b/testexec/main.go index c0be555..eef8449 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -457,8 +457,7 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { } return case <-time.After(time.Second): - //newSent := atomic.LoadUint64(&p.Stats.(*ProducerStatsCollector).Sent) - newSent := p.Stats.(*producer.DefaultStatsCollector).Sent.Count() + newSent := p.Stats.(*producer.DefaultStatsCollector).SentSuccess.Count() if sent != uint64(newSent) { staleTime.Reset(staleTimeout) sent = uint64(newSent) From b12dc4480e8f55e955f793d4b831b6f220ca5107 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 12 Jul 2017 15:43:25 -0500 Subject: [PATCH 051/109] Added new producer stats to the StatsCollector interface to capture retried message statistics. Added a new command line input to the testexec program to allow testing for deadlocks with the producer. --- listener/stats.go | 8 ++++---- producer/producer.go | 25 +++++++------------------ producer/producerconfig_test.go | 1 + producer/stats.go | 15 +++++++++++++++ producer/stats_test.go | 10 ++++++++++ testexec/cli.go | 8 ++++++-- testexec/main.go | 13 +++++++++---- 7 files changed, 52 insertions(+), 28 deletions(-) diff --git a/listener/stats.go b/listener/stats.go index 29a43bf..a3947a8 100644 --- a/listener/stats.go +++ b/listener/stats.go @@ -217,8 +217,8 @@ func (dsc *DefaultStatsCollector) PrintStats() { log.Printf("Listener stats: GetRecords Timeout: [%d]\n", dsc.GetRecordsTimeout.Count()) log.Printf("Listener stats: GetRecords Read Timeout: [%d]\n", dsc.GetRecordsReadTimeout.Count()) log.Printf("Listener stats: Provisioned Throughput Exceeded: [%d]\n", dsc.ProvisionedThroughputExceeded.Count()) - log.Printf("Listener stats: ProcessedDuration (ns): [%d]\n", dsc.ProcessedDuration.Value()) - log.Printf("Listener stats: GetRecordsDuration (ns): [%d]\n", dsc.GetRecordsDuration.Value()) - log.Printf("Listener stats: GetRecordsReadResponseDuration (ns): [%d]\n", dsc.GetRecordsReadResponseDuration.Value()) - log.Printf("Listener stats: GetRecordsUnmarshalDuration (ns): [%d]\n", dsc.GetRecordsUnmarshalDuration.Value()) + log.Printf("Listener stats: Processed Duration (ns): [%d]\n", dsc.ProcessedDuration.Value()) + log.Printf("Listener stats: GetRecords Duration (ns): [%d]\n", dsc.GetRecordsDuration.Value()) + log.Printf("Listener stats: GetRecords Read Response Duration (ns): [%d]\n", dsc.GetRecordsReadResponseDuration.Value()) + log.Printf("Listener stats: GetRecords Unmarshal Duration (ns): [%d]\n", dsc.GetRecordsUnmarshalDuration.Value()) } diff --git a/producer/producer.go b/producer/producer.go index 69b3df7..13ca24d 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -5,12 +5,10 @@ import ( "fmt" "net" "sync" - "sync/atomic" "time" "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/awserr" - "github.com/aws/aws-sdk-go/service/kinesis" "github.com/rewardStyle/kinetic/errs" "github.com/rewardStyle/kinetic/logging" @@ -110,7 +108,7 @@ func (p *Producer) sendBatch(batch []*message.Message) { p.shutdownCond.L.Unlock() }() - var attempts uint64 + var retryAttempts int stop: for { @@ -121,6 +119,7 @@ stop: select { case p.retries <- msg: + p.Stats.AddSentRetried(1) case <-p.pipeOfDeath: return errs.ErrPipeOfDeath } @@ -132,6 +131,7 @@ stop: return nil }) if err == nil { + p.Stats.AddSentTotal(len(batch)) break stop } @@ -146,16 +146,6 @@ stop: } case awserr.Error: switch err.Code() { - case kinesis.ErrCodeProvisionedThroughputExceededException: - // FIXME: It is not clear to me whether PutRecords would ever return a - // ProvisionedThroughputExceeded error. It seems that it would instead return a valid - // response in which some or all the records within the response will contain an error - // code and error message of ProvisionedThroughputExceeded. The current assumption is - // that if we receive an ProvisionedThroughputExceeded error, that the entire batch - // should be retried. Note we only increment the PutRecord stat, instead of the per- - // message stat. Furthermore, we do not increment the FailCount of the messages (as - // the retry mechanism is different). - p.Stats.AddPutRecordsProvisionedThroughputExceeded(1) default: p.LogError("Received AWS error:", err.Error()) } @@ -173,16 +163,16 @@ stop: // batch to be retried rather than retrying the batch as-is. With this approach, we can kill the "stop" // for loop, and set the entire batch to retries to allow the below code to handle retrying the // messages. - if atomic.LoadUint64(&attempts) > uint64(p.maxRetryAttempts) { - p.LogError(fmt.Sprintf("Dropping batch after %d failed attempts to deliver to stream", attempts)) + if retryAttempts > p.maxRetryAttempts { + p.LogError(fmt.Sprintf("Dropping batch after %d failed attempts to deliver to stream", retryAttempts)) p.Stats.AddDroppedTotal(len(batch)) p.Stats.AddDroppedRetries(len(batch)) break stop } - atomic.AddUint64(&attempts, 1) + retryAttempts++ // Apply a delay before retrying - time.Sleep(time.Duration(attempts*attempts) * time.Second) + time.Sleep(time.Duration(retryAttempts * retryAttempts) * time.Second) } // This frees up another sendBatch to run to allow drainage of the messages / retry queue. This should @@ -230,7 +220,6 @@ func (p *Producer) produce() { p.shutdownCond.L.Lock() if len(batch) > 0 { p.outstanding++ - p.Stats.AddSentTotal(len(batch)) p.concurrencySem <- empty{} go p.sendBatch(batch) } else if len(batch) == 0 { diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go index 22889dd..7a5c935 100644 --- a/producer/producerconfig_test.go +++ b/producer/producerconfig_test.go @@ -17,6 +17,7 @@ type DebugStatsCollector struct{} func (l *DebugStatsCollector) AddSentTotal(int) {} func (l *DebugStatsCollector) AddSentSuccess(int) {} func (l *DebugStatsCollector) AddSentFailed(int) {} +func (l *DebugStatsCollector) AddSentRetried(int) {} func (l *DebugStatsCollector) AddDroppedTotal(int) {} func (l *DebugStatsCollector) AddDroppedCapacity(int) {} func (l *DebugStatsCollector) AddDroppedRetries(int) {} diff --git a/producer/stats.go b/producer/stats.go index 664fb58..5a42453 100644 --- a/producer/stats.go +++ b/producer/stats.go @@ -14,6 +14,7 @@ type StatsCollector interface { AddSentTotal(int) AddSentSuccess(int) AddSentFailed(int) + AddSentRetried(int) AddDroppedTotal(int) AddDroppedCapacity(int) AddDroppedRetries(int) @@ -38,6 +39,10 @@ func (nsc *NilStatsCollector) AddSentSuccess(int) {} // AddSentFailed records a count of the number of messages that failed to be sent to AWS Kinesis by the producer. func (nsc *NilStatsCollector) AddSentFailed(int) {} +// AddSentRetried records a count of the number of messages that were retried after some error occurred when sending +// to AWS Kinesis by the producer. +func (nsc *NilStatsCollector) AddSentRetried(int) {} + // AddDroppedTotal records a count of the total number of messages dropped by the application after multiple failures. func (nsc *NilStatsCollector) AddDroppedTotal(int) {} @@ -79,6 +84,7 @@ const ( MetricsSentTotal = "kinetic.producer.sent.total" MetricsSentSuccess = "kinetic.producer.sent.success" MetricsSentFailed = "kinetic.producer.sent.failed" + MetricsSentRetried = "kinetic.producer.sent.retried" MetricsDroppedTotal = "kinetic.producer.dropped.total" MetricsDroppedCapacity = "kinetic.producer.dropped.capacity" MetricsDroppedRetries = "kinetic.producer.dropped.retries" @@ -97,6 +103,7 @@ type DefaultStatsCollector struct { SentTotal metrics.Counter SentSuccess metrics.Counter SentFailed metrics.Counter + SentRetried metrics.Counter DroppedTotal metrics.Counter DroppedCapacity metrics.Counter DroppedRetries metrics.Counter @@ -115,6 +122,7 @@ func NewDefaultStatsCollector(r metrics.Registry) *DefaultStatsCollector { SentTotal: metrics.GetOrRegisterCounter(MetricsSentTotal, r), SentSuccess: metrics.GetOrRegisterCounter(MetricsSentSuccess, r), SentFailed: metrics.GetOrRegisterCounter(MetricsSentFailed, r), + SentRetried: metrics.GetOrRegisterCounter(MetricsSentRetried, r), DroppedTotal: metrics.GetOrRegisterCounter(MetricsDroppedTotal, r), DroppedCapacity: metrics.GetOrRegisterCounter(MetricsDroppedCapacity, r), DroppedRetries: metrics.GetOrRegisterCounter(MetricsDroppedRetries, r), @@ -143,6 +151,12 @@ func (dsc *DefaultStatsCollector) AddSentFailed(count int) { dsc.SentFailed.Inc(int64(count)) } +// AddSentRetried records a count of the number of messages that were retried after some error occurred when sending +// to AWS Kinesis by the producer. +func (dsc *DefaultStatsCollector) AddSentRetried(count int) { + dsc.SentRetried.Inc(int64(count)) +} + // AddDroppedTotal records a count of the total number of messages dropped by the application after multiple failures. func (dsc *DefaultStatsCollector) AddDroppedTotal(count int) { dsc.DroppedTotal.Inc(int64(count)) @@ -204,6 +218,7 @@ func (dsc *DefaultStatsCollector) PrintStats() { log.Printf("Producer Stats: Sent Total: [%d]\n", dsc.SentTotal.Count()) log.Printf("Producer Stats: Sent Success: [%d]\n", dsc.SentSuccess.Count()) log.Printf("Producer Stats: Sent Failed: [%d]\n", dsc.SentFailed.Count()) + log.Printf("Producer Stats: Sent Retried: [%d]\n", dsc.SentRetried.Count()) log.Printf("Producer Stats: Dropped Total: [%d]\n", dsc.DroppedTotal.Count()) log.Printf("Producer Stats: Dropped Retries: [%d]\n", dsc.DroppedRetries.Count()) log.Printf("Producer Stats: Dropped Capacity: [%d]\n", dsc.DroppedCapacity.Count()) diff --git a/producer/stats_test.go b/producer/stats_test.go index dbff301..3d6f320 100644 --- a/producer/stats_test.go +++ b/producer/stats_test.go @@ -30,6 +30,10 @@ func TestStatsCollector(t *testing.T) { sc.AddSentFailed(1) }) + Convey("check that AddSentRetried does not error", func() { + sc.AddSentRetried(1) + }) + Convey("check that AddDroppedTotal does not error", func() { sc.AddDroppedTotal(1) }) @@ -94,6 +98,12 @@ func TestStatsCollector(t *testing.T) { So(sc.(*DefaultStatsCollector).SentFailed.Count(), ShouldEqual, int64(count)) }) + Convey("check that AddSentRetried does not error", func() { + count := rand.Int() + sc.AddSentRetried(count) + So(sc.(*DefaultStatsCollector).SentRetried.Count(), ShouldEqual, int64(count)) + }) + Convey("check that AddDroppedTotal does not error", func() { count := rand.Int() sc.AddDroppedTotal(count) diff --git a/testexec/cli.go b/testexec/cli.go index b16b832..5ec044e 100644 --- a/testexec/cli.go +++ b/testexec/cli.go @@ -14,6 +14,7 @@ type Config struct { NumMsgs *int Duration *int Throttle *bool + Deadlock *bool Cleanup *bool Verbose *bool } @@ -34,7 +35,8 @@ func parseCommandLineArgs() *Config { durationPtr := flag.Int("duration", 0, "used to specify the duration (in seconds) the program should run. "+ "This flag is only applicable to 'write' and 'readwrite' modes. Use zero or negative number to run "+ "indefinitely.") - throttlePtr := flag.Bool("throttle", true, "used to specify whether to throttle PutRecord requests by 1 ms. ") + throttlePtr := flag.Bool("throttle", true, "used to specify whether to throttle PutRecord requests by 1 ms.") + deadlockPtr := flag.Bool("deadlock", false, "used to test potential deadlock condition for the producer.") cleanupPtr := flag.Bool("cleanup", true, "used to specify whether or not to delete the kinesis stream after "+ "processing is complete.") verbosePtr := flag.Bool("verbose", true, "used to specify whether or not to log in verbose mode") @@ -76,6 +78,7 @@ func parseCommandLineArgs() *Config { NumMsgs: numMsgsPtr, Location: locationPtr, Throttle: throttlePtr, + Deadlock: deadlockPtr, Cleanup: cleanupPtr, Verbose: verbosePtr, } @@ -97,11 +100,12 @@ func (c *Config) printConfigs() { log.Println("-num-msgs: (unbounded)") } if c.Duration != nil { - log.Println("-duration: ", *c.Duration) + log.Printf("-duration: [%d] (s)", *c.Duration) } else { log.Println("-duration: (indefinite)") } log.Println("-throttle: ", *c.Throttle) + log.Println("-deadlock: ", *c.Deadlock) log.Println("-cleanup: ", *c.Cleanup) log.Println("-verbose: ", *c.Verbose) log.Println() diff --git a/testexec/main.go b/testexec/main.go index eef8449..340e099 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -248,12 +248,17 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce } p, err := producer.NewProducer(k.Session.Config, w, func(c *producer.Config) { - c.SetBatchSize(500) c.SetBatchTimeout(1000 * time.Millisecond) - c.SetConcurrency(10) - c.SetMaxRetryAttempts(2) - c.SetQueueDepth(100) + c.SetBatchSize(500) + c.SetMaxRetryAttempts(3) c.SetStatsCollector(psc) + if *cfg.Deadlock { + c.SetConcurrency(10) + c.SetQueueDepth(1) + } else { + c.SetConcurrency(10) + c.SetQueueDepth(10000) + } }) if err != nil { log.Fatalf("Unable to create a new producer due to: %v\n", err) From 8f919fcd59d1b7dcec0a473bb2f61c742e4a6f68 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 17 Jul 2017 21:48:42 -0500 Subject: [PATCH 052/109] Initial commit of producer refactor based on producer redsign to 1) remove the retry channel completely 2) and have go routines hold on to failed messages to retry them instead and 3) add a rate limiter to throttle the number of messages that can be sent --- producer/producer.go | 447 +++++++++++++++++--------------- producer/producer_test.go | 2 +- producer/producerconfig.go | 26 +- producer/producerconfig_test.go | 10 +- producer/ratelimiter.go | 80 ++++++ producer/ratelimiter_test.go | 50 ++++ producer/sharedtypes.go | 37 +++ producer/worker.go | 153 +++++++++++ testexec/main.go | 4 +- 9 files changed, 594 insertions(+), 215 deletions(-) create mode 100644 producer/ratelimiter.go create mode 100644 producer/ratelimiter_test.go create mode 100644 producer/sharedtypes.go create mode 100644 producer/worker.go diff --git a/producer/producer.go b/producer/producer.go index 13ca24d..c4d85c0 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -2,59 +2,46 @@ package producer import ( "context" - "fmt" "net" "sync" "time" "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/awserr" - - "github.com/rewardStyle/kinetic/errs" "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic/errs" ) -// StreamWriter is an interface that abstracts the differences in API between Kinesis and Firehose. -type StreamWriter interface { - PutRecords(context.Context, []*message.Message, MessageHandlerAsync) error -} - -// empty is used a as a dummy type for counting semaphore channels. -type empty struct{} - -// MessageHandler defines the signature of a message handler used by PutRecords(). MessageHandler accepts a WaitGroup -// so the function can be run as a blocking operation as opposed to MessageHandlerAsync. -type MessageHandler func(*message.Message, *sync.WaitGroup) error - -// MessageHandlerAsync defines the signature of a message handler used by PutRecords(). MessageHandlerAsync is meant to -// be run asynchronously. -type MessageHandlerAsync func(*message.Message) error - +// producerOptions holds all of the configurable settings for a Producer type producerOptions struct { - batchSize int - batchTimeout time.Duration - queueDepth int - maxRetryAttempts int - concurrency int - Stats StatsCollector + batchSize int // maximum message capacity per request + batchTimeout time.Duration // maximum time duration to wait for incoming messages + queueDepth int // maximum number of messages to enqueue in the message queue + maxRetryAttempts int // maximum number of retry attempts for failed messages + workersPerShard int // number of concurrent workers per shard + shardCount int // initial shard size + rateLimit int // maximum records to be sent per cycle for the rate limiting model + resetFrequency time.Duration // duration of a cycle for the rate limiting model + Stats StatsCollector // stats collection mechanism } -// Producer sends records to Kinesis or Firehose. +// Producer sends records to AWS Kinesis or Firehose. type Producer struct { - *producerOptions - *logging.LogHelper - - writer StreamWriter - messages chan *message.Message - retries chan *message.Message - concurrencySem chan empty - pipeOfDeath chan empty - outstanding int - shutdownCond *sync.Cond - producerWg *sync.WaitGroup - producing bool - producingMu sync.Mutex + *producerOptions // contains all of the configuration settings for the Producer + *logging.LogHelper // object for help with logging + writer StreamWriter // interface for abstracting the PutRecords call + rateLimiter *rateLimiter // throttles the number of messages sent based on total count and size + workerRegistry map[string]*worker // roster of workers in the worker pool + messages chan *message.Message // channel for enqueuing messages to be put on the stream + statusChannel chan *statusReport // channel for workers to communicate their current status + decommChannel chan empty // channel for handling the decommissioning of a surplus of workers + stopChannel chan empty // channel for handling shutdown + pipeOfDeath chan empty // channel for handling pipe of death + startupOnce sync.Once // used to ensure that the startup function is called once + shutdownOnce sync.Once // used to ensure that the shutdown function is called once + resizeMu sync.Mutex // used to prevent resizeWorkerPool from being called synchronously with itself + noCopy noCopy // prevents the Producer from being copied } // NewProducer creates a new producer for writing records to a Kinesis or Firehose stream. @@ -67,202 +54,255 @@ func NewProducer(c *aws.Config, w StreamWriter, fn ...func(*Config)) (*Producer, producerOptions: cfg.producerOptions, LogHelper: &logging.LogHelper{ LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, + Logger: cfg.AwsConfig.Logger, }, - writer: w, - concurrencySem: make(chan empty, cfg.concurrency), - pipeOfDeath: make(chan empty), + writer: w, + workerRegistry: make(map[string]*worker), + rateLimiter: newRateLimiter(cfg.rateLimit, cfg.resetFrequency), + messages: make(chan *message.Message, cfg.queueDepth), + statusChannel: make(chan *statusReport), + decommChannel: make(chan empty), + stopChannel: make(chan empty), + pipeOfDeath: make(chan empty), }, nil } -// startConsuming will initialize the producer and set producing to true if there is not already another consume loop -// running. -func (p *Producer) startProducing() bool { - p.producingMu.Lock() - defer p.producingMu.Unlock() - if !p.producing { - p.producing = true - p.messages = make(chan *message.Message, p.queueDepth) - p.retries = make(chan *message.Message, p.queueDepth) - p.shutdownCond = sync.NewCond(new(sync.Mutex)) - p.producerWg = new(sync.WaitGroup) - return true - } - return false -} - -// stopProducing handles any cleanup after a producing has stopped. -func (p *Producer) stopProducing() { - p.producingMu.Lock() - defer p.producingMu.Unlock() - if p.messages != nil { - close(p.messages) - } - p.producing = false -} - -func (p *Producer) sendBatch(batch []*message.Message) { - defer func() { - p.shutdownCond.L.Lock() - p.outstanding-- - p.shutdownCond.L.Unlock() - }() +// produce is called once to initialize a pool of workers which send batches of messages concurrently +func (p *Producer) produce() { + p.startupOnce.Do(func() { + // Reset shutdownOnce to allow the shut down sequence to happen again + p.shutdownOnce = sync.Once{} - var retryAttempts int + // Instantiate and register new workers + p.resizeWorkerPool(p.shardCount * p.workersPerShard) -stop: - for { - err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message) error { - if msg.FailCount <= p.maxRetryAttempts { - // Apply a delay before retrying - time.Sleep(time.Duration(msg.FailCount*msg.FailCount) * time.Second) + // Instantiate and start a new rate limiter + p.rateLimiter.start() + go func(){ + // Dispatch messages to each worker depending on the worker's capacity and the rate limit + for { select { - case p.retries <- msg: - p.Stats.AddSentRetried(1) case <-p.pipeOfDeath: - return errs.ErrPipeOfDeath + return + case <-p.stopChannel: + return + case req := <-p.statusChannel: + // Use the worker registry to find the worker by workerID + worker, ok := p.workerRegistry[req.workerID]; + if !ok { + // move on to the next status if a bogus workerID is provided + break + } + + // Verify that worker is in the right state for handling a new batch of messages + workerState := worker.getWorkerState() + if !(workerState == workerStateIdle || + workerState == workerStateIdleWithRetries) { + // otherwise on to the next one + break + } + + // If we need to decommission some workers distribute the decommission command + // to idle workers, otherwise try to fill the workers capacity + var batch []*message.Message + var decommissioned bool + if len(p.decommChannel) > 0 && workerState == workerStateIdle { + <-p.decommChannel + decommissioned = true + } else { + tokenCount := p.rateLimiter.getTokenCount() + + timeout := time.After(req.timeout) + + fillBatch: + for len(batch) < req.capacity { + select { + case <-timeout: + break fillBatch + case msg := <-p.messages: + batch = append(batch, msg) + } + if len(batch) + req.failed >= tokenCount { + break fillBatch + } + } + p.rateLimiter.claimTokens(len(batch) + req.failed) + } + + // Send the command to the worker's command channel + worker.commands <-&workerCommand{ + batchMsgs: batch, + decommissioned: decommissioned, + } } - } else { - p.Stats.AddDroppedTotal(1) - p.Stats.AddDroppedRetries(1) } + }() + }) +} - return nil - }) - if err == nil { - p.Stats.AddSentTotal(len(batch)) - break stop +// shutdown is called once to handle the graceful shutdown of the produce function +func (p *Producer) shutdown() { + p.shutdownOnce.Do(func() { + // Close the messages channel to prevent any more incoming messages + if p.messages != nil { + close(p.messages) } - // The call failed so we need to retry the batch - switch err := err.(type) { - case net.Error: - if err.Timeout() { - p.Stats.AddPutRecordsTimeout(1) - p.LogError("Received net error:", err.Error()) - } else { - p.LogError("Received unknown net error:", err.Error()) - } - case awserr.Error: - switch err.Code() { - default: - p.LogError("Received AWS error:", err.Error()) - } - case error: - switch err { - case errs.ErrRetryRecords: - break stop - default: - p.LogError("Received error:", err.Error()) + // Allow the workers to drain the message channel first + wg := sync.WaitGroup{} + wg.Add(1) + go func(){ + defer wg.Done() + + staleTimeout := time.Duration(3 * time.Second) + timer := time.NewTimer(staleTimeout) + remaining := len(p.messages) + for remaining > 0 { + select { + case <-time.After(time.Second): + newRemaining := len(p.messages) + if newRemaining != remaining { + timer.Reset(staleTimeout) + remaining = newRemaining + } + case <-timer.C: + return + } } - default: - p.LogError("Received unknown error:", err.Error()) + }() + wg.Wait() + + // Decommission all the workers + p.resizeWorkerPool(0) + + // Close the decommission channel + if p.decommChannel != nil { + close(p.decommChannel) } - // NOTE: We may want to go through and increment the FailCount for each of the records and allow the - // batch to be retried rather than retrying the batch as-is. With this approach, we can kill the "stop" - // for loop, and set the entire batch to retries to allow the below code to handle retrying the - // messages. - if retryAttempts > p.maxRetryAttempts { - p.LogError(fmt.Sprintf("Dropping batch after %d failed attempts to deliver to stream", retryAttempts)) - p.Stats.AddDroppedTotal(len(batch)) - p.Stats.AddDroppedRetries(len(batch)) - break stop + + // Close the status channel + if p.statusChannel != nil { + close(p.statusChannel) } - retryAttempts++ - // Apply a delay before retrying - time.Sleep(time.Duration(retryAttempts * retryAttempts) * time.Second) - } + // Stop the rate limiter + p.rateLimiter.stop() - // This frees up another sendBatch to run to allow drainage of the messages / retry queue. This should - // improve throughput as well as prevent a potential deadlock in which all batches are blocked on - // sending retries to the retries channel, and thus no batches are allowed to drain the retry channel. - <-p.concurrencySem + // Stop the running go routine in produce + p.stopChannel <- empty{} + + // Reset startupOnce to allow the start up sequence to happen again + p.startupOnce = sync.Once{} + }) } -// produce calls the underlying writer's PutRecords implementation to deliver batches of messages to the target stream -// until the producer is stopped. -func (p *Producer) produce() { - if !p.startProducing() { - return +// resizeWorkerPool is called to instantiate new workers, decommission workers or recommission workers that have been +// deactivated +func (p *Producer) resizeWorkerPool(desiredWorkerCount int) { + p.resizeMu.Lock() + defer p.resizeMu.Unlock() + + // Get the number of total workers in the registry + totalWorkerCount := len(p.workerRegistry) + + + // Create a map of available workers + availableWorkers := make(map[string]*worker, totalWorkerCount) + for id, worker := range p.workerRegistry { + workerState := worker.getWorkerState() + if workerState == workerStateInactive || workerState == workerStateDecommissioned { + availableWorkers[id] = worker + } } - p.producerWg.Add(1) - go func() { - defer func() { - p.stopProducing() - p.producerWg.Done() - }() + activeWorkerCount := totalWorkerCount - len(availableWorkers) - for { - var batch []*message.Message - timer := time.After(p.batchTimeout) - batch: - for len(batch) < p.batchSize { - select { - // Using the select, retry messages will interleave with new messages. This is - // preferable to putting the messages at the end of the channel as it minimizes the - // delay in the delivery of retry messages. - case msg, ok := <-p.messages: - if !ok { - p.messages = nil - } else { - batch = append(batch, msg) - } - case msg := <-p.retries: - batch = append(batch, msg) - case <-timer: - break batch - case <-p.pipeOfDeath: - return - } + // We have too many workers at present than we actually need + if desiredWorkerCount < activeWorkerCount { + // Decommission the workers that we don't need + for i := activeWorkerCount; i <= desiredWorkerCount; i-- { + p.decommChannel <-empty{} + } + // We need more workers than are presently active or commissioned + } else if activeWorkerCount < desiredWorkerCount { + // Recommission those workers that are inactive + var activated int + for _, worker := range availableWorkers { + if activated > desiredWorkerCount - activeWorkerCount { + break } - p.shutdownCond.L.Lock() - if len(batch) > 0 { - p.outstanding++ - p.concurrencySem <- empty{} - go p.sendBatch(batch) - } else if len(batch) == 0 { - // We did not get any records -- check if we may be (gracefully) shutting down the - // producer. We can exit when: - // - The messages channel is nil and no new messages can be enqueued - // - There are no outstanding sendBatch goroutines and can therefore not produce any - // more messages to retry - // - The retry channel is empty - if p.messages == nil && p.outstanding == 0 && len(p.retries) == 0 { - close(p.retries) - p.shutdownCond.Broadcast() - p.shutdownCond.L.Unlock() - return - } + worker.start() + activated++ + } + + // Spawn new workers if still not enough + if desiredWorkerCount > totalWorkerCount { + for i := totalWorkerCount; i < desiredWorkerCount; i++ { + worker := newWorker(p.producerOptions, p.sendBatch, p.reportStatus) + p.workerRegistry[worker.workerID] = worker } - p.shutdownCond.L.Unlock() } - }() + } } -// CloseWithContext shuts down the producer, waiting for all outstanding messages and retries to flush. Cancellation -// is supported through contexts. -func (p *Producer) CloseWithContext(ctx context.Context) { - c := make(chan empty, 1) - close(p.messages) - go func() { - p.shutdownCond.L.Lock() - for p.outstanding != 0 { - p.shutdownCond.Wait() +// sendBatch is the function that is called by each worker to put records on the stream. sendBatch accepts a slice of +// messages to send and returns a slice of messages that failed to send +func (p *Producer) sendBatch(batch []*message.Message) []*message.Message { + var failed []*message.Message + err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message) error { + if msg.FailCount <= p.maxRetryAttempts { + failed = append(failed, msg) + p.Stats.AddSentRetried(1) + } else { + p.Stats.AddDroppedTotal(1) + p.Stats.AddDroppedRetries(1) } - p.shutdownCond.L.Unlock() - p.producerWg.Wait() - c <- empty{} - }() - select { - case <-c: - case <-ctx.Done(): - close(p.pipeOfDeath) + return nil + }) + if err == nil { + p.Stats.AddSentTotal(len(batch)) + return failed + } + + // Beyond this point the PutRecords API call failed for some reason + switch err := err.(type) { + case net.Error: + if err.Timeout() { + p.Stats.AddPutRecordsTimeout(1) + p.LogError("Received net error:", err.Error()) + } else { + p.LogError("Received unknown net error:", err.Error()) + } + case awserr.Error: + p.LogError("Received AWS error:", err.Error()) + case error: + switch err { + case errs.ErrRetryRecords: + break + default: + p.LogError("Received error:", err.Error()) + } + default: + p.LogError("Received unknown error:", err.Error()) } + + return batch +} + +// reportStatus is used as a closure for the workers to report status to +func (p *Producer) reportStatus(report *statusReport) { + p.statusChannel <- report } -// Close shuts down the producer, waiting for all outstanding messages and retries to flush. +// CloseWithContext initiates the graceful shutdown of the produce function, waiting for all outstanding messages and to +// flush. Cancellation is supported through contexts. +func (p *Producer) CloseWithContext(ctx context.Context) { + p.shutdown() + <-ctx.Done() + close(p.pipeOfDeath) +} + +// Close initiates the graceful shutdown of the produce function, waiting for all outstanding messages and to flush. func (p *Producer) Close() { p.CloseWithContext(context.TODO()) } @@ -286,6 +326,7 @@ func (p *Producer) Send(msg *message.Message) error { // TryToSend will attempt to send a message to the stream if the channel has capacity for a message, or will immediately // return with an error if the channel is full. func (p *Producer) TryToSend(msg *message.Message) error { + p.produce() select { case p.messages <- msg: return nil diff --git a/producer/producer_test.go b/producer/producer_test.go index 181d638..f2c2e14 100644 --- a/producer/producer_test.go +++ b/producer/producer_test.go @@ -45,7 +45,7 @@ func TestProducer(t *testing.T) { p, err := NewProducer(k.Session.Config, w, func(c *Config) { c.SetBatchSize(5) c.SetBatchTimeout(1000 * time.Millisecond) - c.SetConcurrency(10) + //c.SetConcurrency(10) c.SetQueueDepth(10) }) So(p, ShouldNotBeNil) diff --git a/producer/producerconfig.go b/producer/producerconfig.go index d2ec8ce..24ff443 100644 --- a/producer/producerconfig.go +++ b/producer/producerconfig.go @@ -24,7 +24,10 @@ func NewConfig(cfg *aws.Config) *Config { batchTimeout: time.Second, queueDepth: 500, maxRetryAttempts: 10, - concurrency: 1, + workersPerShard: 5, + shardCount: 1, + rateLimit: 1000, + resetFrequency: time.Second, Stats: &NilStatsCollector{}, }, LogLevel: *cfg.LogLevel, @@ -57,9 +60,24 @@ func (c *Config) SetMaxRetryAttempts(attempts int) { c.maxRetryAttempts = attempts } -// SetConcurrency controls the number of outstanding PutRecords calls may be active at a time. -func (c *Config) SetConcurrency(concurrency int) { - c.concurrency = concurrency +// SetWorkersPerShard defines the number of concurrent workers to run per active shard +func (c *Config) SetWorkersPerShard(count int) { + c.workersPerShard = count +} + +// SetShardCount defines the initial shard size +func (c *Config) SetShardCount(count int) { + c.shardCount = count +} + +// SetRateLimit defines the maximum number of message to send per cycle +func (c *Config) SetRateLimit(limit int) { + c.rateLimit = limit +} + +// SetResetFrequency defines the frequency at which the rateLimiter resets +func (c *Config) SetResetFrequency(freq time.Duration) { + c.resetFrequency = freq } // SetStatsCollector configures a listener to handle producer metrics. diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go index 7a5c935..73fa58f 100644 --- a/producer/producerconfig_test.go +++ b/producer/producerconfig_test.go @@ -48,7 +48,10 @@ func TestNewConfig(t *testing.T) { So(cfg.batchTimeout, ShouldEqual, 1*time.Second) So(cfg.queueDepth, ShouldEqual, 500) So(cfg.maxRetryAttempts, ShouldEqual, 10) - So(cfg.concurrency, ShouldEqual, 1) + So(cfg.workersPerShard, ShouldEqual, 5) + So(cfg.shardCount, ShouldEqual, 1) + So(cfg.rateLimit, ShouldEqual, 1000) + So(cfg.resetFrequency, ShouldEqual, time.Second) So(cfg.Stats, ShouldHaveSameTypeAs, &NilStatsCollector{}) So(cfg.LogLevel.Value(), ShouldEqual, logging.LogOff) }) @@ -79,10 +82,7 @@ func TestNewConfig(t *testing.T) { So(cfg.maxRetryAttempts, ShouldEqual, 100) }) - Convey("check that we can set the concurrency limit", func() { - cfg.SetConcurrency(50) - So(cfg.concurrency, ShouldEqual, 50) - }) + // TODO: Add tests for setting new producer options Convey("check that we can configure a stats collector", func() { cfg.SetStatsCollector(&DebugStatsCollector{}) diff --git a/producer/ratelimiter.go b/producer/ratelimiter.go new file mode 100644 index 0000000..55ea3b0 --- /dev/null +++ b/producer/ratelimiter.go @@ -0,0 +1,80 @@ +package producer + +import ( + "sync" + "time" +) + +// TODO: add size rate limiting in addition to message count + +// rateLimiter is used by the producer to rate limit the total number and size of records sent per cycle. +type rateLimiter struct { + limit int // upper limit of throughput per cycle + duration time.Duration // frequency with which to reset the remaining tokens count + tokenCount int // remaining tokens available for the cycle + tokenMu sync.Mutex // mutex to protect remaining token count + stopChannel chan empty // channel for communicating when to stop rate limiting + startOnce sync.Once // startOnce is used to ensure that start is called once and only once + stopOnce sync.Once // stopOnce is used to ensure that stop is called once and only once +} + +// newRateLimiter creates a new rateLimiter. +func newRateLimiter(limit int, duration time.Duration) *rateLimiter { + return &rateLimiter{ + limit: limit, + duration: duration, + tokenCount: limit, + } +} + +// start runs a timer in a go routine background which resets the the number and size counters every cycle. +func (r *rateLimiter) start() { + r.startOnce.Do(func() { + // Reset stopOnce to allow the rateLimiter to be shut down again + r.stopOnce = sync.Once{} + + r.stopChannel = make(chan empty) + ticker := time.NewTicker(r.duration) + go func(){ + for { + select { + case <-r.stopChannel: + return + case <-ticker.C: + r.reset() + } + } + }() + }) +} + +// stop sends a signal to the rateLimiter's stopChannel +func (r *rateLimiter) stop() { + r.stopOnce.Do(func(){ + r.stopChannel <- empty{} + + // Reset startOnce to allow the rateLimiter to be started again + r.startOnce = sync.Once{} + }) +} + +// reset is called to reset the rateLimiter's tokens to the initial values +func (r *rateLimiter) reset() { + r.tokenMu.Lock() + defer r.tokenMu.Unlock() + r.tokenCount = r.limit +} + +// getTokenCount is used to retrieve the current token count +func (r *rateLimiter) getTokenCount() int { + r.tokenMu.Lock() + defer r.tokenMu.Unlock() + return r.tokenCount +} + +// claimTokens is used to claim tokens prior to sending messages +func (r *rateLimiter) claimTokens(count int) { + r.tokenMu.Lock() + defer r.tokenMu.Unlock() + r.tokenCount -= count +} diff --git a/producer/ratelimiter_test.go b/producer/ratelimiter_test.go new file mode 100644 index 0000000..7513172 --- /dev/null +++ b/producer/ratelimiter_test.go @@ -0,0 +1,50 @@ +package producer + +import ( + . "github.com/smartystreets/goconvey/convey" + + "testing" + "time" +) + +func TestRateLimiter(t *testing.T) { + limit := 1000 + duration := time.Second + Convey("calling newRateLimiter should return the pointer to a rateLimiter with properly initalized values", t, func() { + rl := newRateLimiter(limit, duration) + So(rl, ShouldNotBeNil) + So(rl.limit, ShouldEqual, limit) + So(rl.duration, ShouldEqual, duration) + So(rl.tokenCount, ShouldEqual, limit) + + Convey("calling getTokenCount returns the correct token count", func() { + So(rl.getTokenCount(), ShouldEqual, limit) + }) + + Convey("calling claimTokens reduces the tokens properly", func() { + So(rl.getTokenCount(), ShouldEqual, limit) + claim := 100 + rl.claimTokens(claim) + So(rl.getTokenCount(), ShouldEqual, limit - claim) + }) + + Convey("calling reset will reset the token count back to the original limit", func() { + So(rl.getTokenCount(), ShouldEqual, limit) + claim := 100 + rl.claimTokens(claim) + So(rl.getTokenCount(), ShouldNotEqual, limit) + rl.reset() + So(rl.getTokenCount(), ShouldEqual, limit) + }) + + Convey("calling start on the rateLimiter", func() { + // TODO: Finish this unit test + rl.start() + time.Sleep(10 * time.Second) + }) + + Convey("calling stop on the rateLimiter", func() { + // TODO: Finish this unit test + }) + }) +} diff --git a/producer/sharedtypes.go b/producer/sharedtypes.go new file mode 100644 index 0000000..738c604 --- /dev/null +++ b/producer/sharedtypes.go @@ -0,0 +1,37 @@ +package producer + +import ( + "context" + "sync" + + "github.com/rewardStyle/kinetic/message" +) + +// StreamWriter is an interface that abstracts the differences in API between Kinesis and Firehose. +type StreamWriter interface { + PutRecords(context.Context, []*message.Message, MessageHandlerAsync) error +} + +// sendBatchFn defines the signature for the sendBatch function defined by the producer and passed to the worker as a +// closure for execution +type sendBatchFn func([]*message.Message) []*message.Message + +// reportStatusFn defines the signature for the worker to communicate its status to the producer +type reportStatusFn func(*statusReport) + +// MessageHandler defines the signature of a message handler used by PutRecords(). MessageHandler accepts a WaitGroup +// so the function can be run as a blocking operation as opposed to MessageHandlerAsync. +type MessageHandler func(*message.Message, *sync.WaitGroup) error + +// MessageHandlerAsync defines the signature of a message handler used by PutRecords(). MessageHandlerAsync is meant to +// be run asynchronously. +type MessageHandlerAsync func(*message.Message) error + +// empty is used a as a dummy type for counting semaphore channels. +type empty struct{} + +// noCopy is used to prevent structs from being copied +type noCopy struct{} + +// Lock is a no-op used by -copylocks checker from `go vet`. +func (*noCopy) Lock() {} diff --git a/producer/worker.go b/producer/worker.go new file mode 100644 index 0000000..95ae04b --- /dev/null +++ b/producer/worker.go @@ -0,0 +1,153 @@ +package producer + +import ( + "crypto/rand" + "fmt" + "sync" + "time" + + "github.com/rewardStyle/kinetic/message" +) + +const ( + workerStateActive = "active" + workerStateIdle = "idle" + workerStateIdleWithRetries = "idleWithErrors" + workerStateBusy = "busy" + workerStateDecommissioned = "decommissioned" + workerStateInactive = "inactive" +) + +// statusReport is used to communicate a worker's status to the producer. +type statusReport struct { + workerID string // uuid of the worker making the work request + capacity int // maximum message capacity the worker can handle + timeout time.Duration // maximum time duration to wait for incoming messages + failed int // number of previous messages that failed to send +} + +// workerCommand is used by the producer to send commands to the workers. +type workerCommand struct { + batchMsgs []*message.Message // batch of messages for the worker to send + decommissioned bool // whether or not to decommission the worker +} + +// worker is an object that, when started, runs in a separate go routine handling commands from the producer through +// the worker's command channel. The worker communicates its status to the producer and the producer responds with +// a batch of messages for the worker to send. The worker calls a closure (defined by the producer) to send a batch +// of messages containing a combination of new messages and previously failed messages. The worker may also receive +// a command to decommission itself if the worker has no previously failed messages. +type worker struct { + workerID string // unique ID to identify itself from the other workers + state string // current state of the worker + stateMu sync.Mutex // mutex for protecting against concurrent changes of state + batchSize int // maximum number of message to be sent per batch + batchTimeout time.Duration // maximum duration with which to wait for a batch of new messages + sendBatchFn sendBatchFn // a closure to call when a batch is ready to be sent + reportStatusFn reportStatusFn // a clousre to call when the worker has a status update + retries []*message.Message // a slice of messages that previously failed + commands chan *workerCommand // a channel to which the producer sends commands +} + +// newWorker creates a new worker object that automatically starts working. +func newWorker(cfg *producerOptions, sendBatchFn sendBatchFn, reportStatusFn reportStatusFn) *worker { + var w *worker + defer func() { + go w.start() + }() + + w = &worker{ + workerID: generateWorkerID(), + state: workerStateActive, + batchSize: cfg.batchSize, + batchTimeout: cfg.batchTimeout, + sendBatchFn: sendBatchFn, + reportStatusFn: reportStatusFn, + commands: make(chan *workerCommand), + } + + return w +} + +// start is called when a new worker is instantiated via newWorker. start initiates an infinite loop that communicates +// the worker's status to the producer and listens to its command channel for commands sent by the producer +func (w *worker) start() { + defer w.shutdown() + + var failed int + for { + // Update worker state + if failed > 0 { + w.setWorkerState(workerStateIdleWithRetries) + } else { + w.setWorkerState(workerStateIdle) + } + + // Send the dispatcher a status report with the worker's capacity + go func() { + req := &statusReport{ + workerID: w.workerID, + capacity: w.batchSize - failed, + timeout: w.batchTimeout, + failed: failed, + } + w.reportStatusFn(req) + }() + + // Listen to incoming commands + cmd := <-w.commands + if cmd.decommissioned { + w.setWorkerState(workerStateDecommissioned) + return + } + + // Only send if there are any messages to send + if len(cmd.batchMsgs) + len(w.retries) > 0 { + w.setWorkerState(workerStateBusy) + + // Combine new batch messages with previously failed messages + var batch []*message.Message + batch = append(batch, w.retries...) + batch = append(batch, cmd.batchMsgs...) + + // Send batch and reset retries + w.retries = w.sendBatchFn(batch) + } + } +} + +// shutdown initiates the shutdown procedure for a worker +func (w *worker) shutdown() { + // Close the command channel to prevent any more incoming commands + if w.commands != nil { + close(w.commands) + } + + // Update the worker state to inactive + w.setWorkerState(workerStateInactive) +} + +// getWorkerState returns the worker's current state +func (w *worker) getWorkerState() string { + w.stateMu.Lock() + defer w.stateMu.Unlock() + return w.state +} + +// setWorkerStatus updates the worker's current state +func (w *worker) setWorkerState(state string) { + w.stateMu.Lock() + defer w.stateMu.Unlock() + w.state = state +} + +// generateWorkerID is a helper function that generates UUIDs that are used as unique worker IDs +func generateWorkerID() string { + b := make([]byte, 16) + _, err := rand.Read(b) + if err != nil { + fmt.Println("Error: ", err) + return "" + } + return fmt.Sprintf("%X-%X-%X-%X-%X", b[0:4], b[4:6], b[6:8], b[8:10], b[10:]) +} diff --git a/testexec/main.go b/testexec/main.go index 340e099..a9f02fd 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -253,10 +253,10 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce c.SetMaxRetryAttempts(3) c.SetStatsCollector(psc) if *cfg.Deadlock { - c.SetConcurrency(10) + //c.SetConcurrency(10) c.SetQueueDepth(1) } else { - c.SetConcurrency(10) + //c.SetConcurrency(10) c.SetQueueDepth(10000) } }) From 1a8ec243186cace90a375f4514f569f96e55448b Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 18 Jul 2017 10:56:27 -0500 Subject: [PATCH 053/109] Fixed a bug that introduced nil messages into the workers retry messages by converting the closure function that handles what to do with failed messages to a synchronous call that waits on a wait group until processing is complete. Added nil message checks to the producer and the writers to prevent potential segmentation fault errors. --- producer/firehosewriter.go | 12 +++++++++--- producer/kinesiswriter.go | 12 +++++++++--- producer/producer.go | 9 ++++++--- producer/sharedtypes.go | 2 +- producer/worker.go | 2 +- 5 files changed, 26 insertions(+), 11 deletions(-) diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go index 3f3c945..dfc41c8 100644 --- a/producer/firehosewriter.go +++ b/producer/firehosewriter.go @@ -3,6 +3,7 @@ package producer import ( "context" "fmt" + "sync" "time" "github.com/aws/aws-sdk-go/aws" @@ -51,14 +52,16 @@ func NewFirehoseWriter(c *aws.Config, stream string, fn ...func(*FirehoseWriterC } // PutRecords sends a batch of records to Firehose and returns a list of records that need to be retried. -func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageHandlerAsync) error { +func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageHandler) error { var startSendTime time.Time var startBuildTime time.Time start := time.Now() var records []*firehose.Record for _, msg := range messages { - records = append(records, msg.ToFirehoseRecord()) + if msg != nil { + records = append(records, msg.ToFirehoseRecord()) + } } req, resp := w.client.PutRecordBatchRequest(&firehose.PutRecordBatchInput{ DeliveryStreamName: aws.String(w.stream), @@ -121,7 +124,10 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Mes messages[idx].FailCount++ w.Stats.AddSentFailed(1) - go fn(messages[idx]) + wg := sync.WaitGroup{} + wg.Add(1) + fn(messages[idx], &wg) + wg.Wait() } } diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index 4c7c7a8..45722d1 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -3,6 +3,7 @@ package producer import ( "context" "fmt" + "sync" "time" "github.com/aws/aws-sdk-go/aws" @@ -52,14 +53,16 @@ func NewKinesisWriter(c *aws.Config, stream string, fn ...func(*KinesisWriterCon } // PutRecords sends a batch of records to Kinesis and returns a list of records that need to be retried. -func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageHandlerAsync) error { +func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageHandler) error { var startSendTime time.Time var startBuildTime time.Time start := time.Now() var records []*kinesis.PutRecordsRequestEntry for _, msg := range messages { - records = append(records, msg.ToRequestEntry()) + if msg != nil { + records = append(records, msg.ToRequestEntry()) + } } req, resp := w.client.PutRecordsRequest(&kinesis.PutRecordsInput{ StreamName: aws.String(w.stream), @@ -124,7 +127,10 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess messages[idx].FailCount++ w.Stats.AddSentFailed(1) - go fn(messages[idx]) + wg := sync.WaitGroup{} + wg.Add(1) + fn(messages[idx], &wg) + wg.Wait() } } diff --git a/producer/producer.go b/producer/producer.go index c4d85c0..05708fd 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -112,7 +112,6 @@ func (p *Producer) produce() { decommissioned = true } else { tokenCount := p.rateLimiter.getTokenCount() - timeout := time.After(req.timeout) fillBatch: @@ -121,7 +120,9 @@ func (p *Producer) produce() { case <-timeout: break fillBatch case msg := <-p.messages: - batch = append(batch, msg) + if msg != nil { + batch = append(batch, msg) + } } if len(batch) + req.failed >= tokenCount { break fillBatch @@ -249,7 +250,9 @@ func (p *Producer) resizeWorkerPool(desiredWorkerCount int) { // messages to send and returns a slice of messages that failed to send func (p *Producer) sendBatch(batch []*message.Message) []*message.Message { var failed []*message.Message - err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message) error { + err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message, wg *sync.WaitGroup) error { + defer wg.Done() + if msg.FailCount <= p.maxRetryAttempts { failed = append(failed, msg) p.Stats.AddSentRetried(1) diff --git a/producer/sharedtypes.go b/producer/sharedtypes.go index 738c604..7e8b8f6 100644 --- a/producer/sharedtypes.go +++ b/producer/sharedtypes.go @@ -9,7 +9,7 @@ import ( // StreamWriter is an interface that abstracts the differences in API between Kinesis and Firehose. type StreamWriter interface { - PutRecords(context.Context, []*message.Message, MessageHandlerAsync) error + PutRecords(context.Context, []*message.Message, MessageHandler) error } // sendBatchFn defines the signature for the sendBatch function defined by the producer and passed to the worker as a diff --git a/producer/worker.go b/producer/worker.go index 95ae04b..61e129c 100644 --- a/producer/worker.go +++ b/producer/worker.go @@ -12,7 +12,7 @@ import ( const ( workerStateActive = "active" workerStateIdle = "idle" - workerStateIdleWithRetries = "idleWithErrors" + workerStateIdleWithRetries = "idleWithRetries" workerStateBusy = "busy" workerStateDecommissioned = "decommissioned" workerStateInactive = "inactive" From 0c6ec56bd4eebb192efbbcb73157aee213314643 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 18 Jul 2017 22:46:34 -0500 Subject: [PATCH 054/109] Removed the worker object entirely replaced it with a helper function within the producer for running concurrent worker routines to call sendBatch. Added the worker's 'inbox' channel in the the statusReport to eliminate the need for workerIDs. Adopted the recommendation for the workers to listen on a global dismiss channel to eliminate the need for worker states. Reverted back to removing the wait group from the message handler in PutRecords. Added new functionality in rateLimiter to be able to resize the limit and duration and to tryToClaimTokens. --- producer/firehosewriter.go | 8 +- producer/kinesiswriter.go | 8 +- producer/producer.go | 240 +++++++++++++++++-------------------- producer/ratelimiter.go | 40 ++++++- producer/sharedtypes.go | 16 +-- producer/worker.go | 153 ----------------------- 6 files changed, 159 insertions(+), 306 deletions(-) delete mode 100644 producer/worker.go diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go index dfc41c8..bfc6a1b 100644 --- a/producer/firehosewriter.go +++ b/producer/firehosewriter.go @@ -3,7 +3,6 @@ package producer import ( "context" "fmt" - "sync" "time" "github.com/aws/aws-sdk-go/aws" @@ -52,7 +51,7 @@ func NewFirehoseWriter(c *aws.Config, stream string, fn ...func(*FirehoseWriterC } // PutRecords sends a batch of records to Firehose and returns a list of records that need to be retried. -func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageHandler) error { +func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageHandlerAsync) error { var startSendTime time.Time var startBuildTime time.Time @@ -124,10 +123,7 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Mes messages[idx].FailCount++ w.Stats.AddSentFailed(1) - wg := sync.WaitGroup{} - wg.Add(1) - fn(messages[idx], &wg) - wg.Wait() + fn(messages[idx]) } } diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index 45722d1..6f07913 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -3,7 +3,6 @@ package producer import ( "context" "fmt" - "sync" "time" "github.com/aws/aws-sdk-go/aws" @@ -53,7 +52,7 @@ func NewKinesisWriter(c *aws.Config, stream string, fn ...func(*KinesisWriterCon } // PutRecords sends a batch of records to Kinesis and returns a list of records that need to be retried. -func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageHandler) error { +func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageHandlerAsync) error { var startSendTime time.Time var startBuildTime time.Time @@ -127,10 +126,7 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess messages[idx].FailCount++ w.Stats.AddSentFailed(1) - wg := sync.WaitGroup{} - wg.Add(1) - fn(messages[idx], &wg) - wg.Wait() + fn(messages[idx]) } } diff --git a/producer/producer.go b/producer/producer.go index 05708fd..02d695c 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -28,20 +28,20 @@ type producerOptions struct { // Producer sends records to AWS Kinesis or Firehose. type Producer struct { - *producerOptions // contains all of the configuration settings for the Producer - *logging.LogHelper // object for help with logging - writer StreamWriter // interface for abstracting the PutRecords call - rateLimiter *rateLimiter // throttles the number of messages sent based on total count and size - workerRegistry map[string]*worker // roster of workers in the worker pool - messages chan *message.Message // channel for enqueuing messages to be put on the stream - statusChannel chan *statusReport // channel for workers to communicate their current status - decommChannel chan empty // channel for handling the decommissioning of a surplus of workers - stopChannel chan empty // channel for handling shutdown - pipeOfDeath chan empty // channel for handling pipe of death - startupOnce sync.Once // used to ensure that the startup function is called once - shutdownOnce sync.Once // used to ensure that the shutdown function is called once - resizeMu sync.Mutex // used to prevent resizeWorkerPool from being called synchronously with itself - noCopy noCopy // prevents the Producer from being copied + *producerOptions // contains all of the configuration settings for the Producer + *logging.LogHelper // object for help with logging + writer StreamWriter // interface for abstracting the PutRecords call + rateLimiter *rateLimiter // throttles the number of messages sent based on total count and size + workerCount int // number of concurrent workers sending batch messages for the producer + messages chan *message.Message // channel for enqueuing messages to be put on the stream + status chan *statusReport // channel for workers to communicate their current status + dismiss chan empty // channel for handling the decommissioning of a surplus of workers + stop chan empty // channel for handling shutdown + pipeOfDeath chan empty // channel for handling pipe of death + startupOnce sync.Once // used to ensure that the startup function is called once + shutdownOnce sync.Once // used to ensure that the shutdown function is called once + resizeMu sync.Mutex // used to prevent resizeWorkerPool from being called synchronously with itself + noCopy noCopy // prevents the Producer from being copied } // NewProducer creates a new producer for writing records to a Kinesis or Firehose stream. @@ -57,13 +57,6 @@ func NewProducer(c *aws.Config, w StreamWriter, fn ...func(*Config)) (*Producer, Logger: cfg.AwsConfig.Logger, }, writer: w, - workerRegistry: make(map[string]*worker), - rateLimiter: newRateLimiter(cfg.rateLimit, cfg.resetFrequency), - messages: make(chan *message.Message, cfg.queueDepth), - statusChannel: make(chan *statusReport), - decommChannel: make(chan empty), - stopChannel: make(chan empty), - pipeOfDeath: make(chan empty), }, nil } @@ -73,6 +66,14 @@ func (p *Producer) produce() { // Reset shutdownOnce to allow the shut down sequence to happen again p.shutdownOnce = sync.Once{} + // Create communication channels + p.rateLimiter = newRateLimiter(p.rateLimit, p.resetFrequency) + p.messages = make(chan *message.Message, p.queueDepth) + p.status = make(chan *statusReport) + p.dismiss = make(chan empty) + p.stop = make(chan empty) + p.pipeOfDeath = make(chan empty) + // Instantiate and register new workers p.resizeWorkerPool(p.shardCount * p.workersPerShard) @@ -85,57 +86,27 @@ func (p *Producer) produce() { select { case <-p.pipeOfDeath: return - case <-p.stopChannel: + case <-p.stop: return - case req := <-p.statusChannel: - // Use the worker registry to find the worker by workerID - worker, ok := p.workerRegistry[req.workerID]; - if !ok { - // move on to the next status if a bogus workerID is provided - break - } - - // Verify that worker is in the right state for handling a new batch of messages - workerState := worker.getWorkerState() - if !(workerState == workerStateIdle || - workerState == workerStateIdleWithRetries) { - // otherwise on to the next one - break - } - - // If we need to decommission some workers distribute the decommission command - // to idle workers, otherwise try to fill the workers capacity + case status := <-p.status: var batch []*message.Message - var decommissioned bool - if len(p.decommChannel) > 0 && workerState == workerStateIdle { - <-p.decommChannel - decommissioned = true - } else { - tokenCount := p.rateLimiter.getTokenCount() - timeout := time.After(req.timeout) - - fillBatch: - for len(batch) < req.capacity { - select { - case <-timeout: - break fillBatch - case msg := <-p.messages: - if msg != nil { - batch = append(batch, msg) - } - } - if len(batch) + req.failed >= tokenCount { - break fillBatch + tokenCount := p.rateLimiter.getTokenCount() + timeout := time.After(p.batchTimeout) + + fillBatch: + for len(batch) < status.capacity && len(batch) + status.failed < tokenCount { + select { + case <-timeout: + break fillBatch + case msg := <-p.messages: + if msg != nil { + batch = append(batch, msg) } } - p.rateLimiter.claimTokens(len(batch) + req.failed) } + p.rateLimiter.claimTokens(len(batch) + status.failed) - // Send the command to the worker's command channel - worker.commands <-&workerCommand{ - batchMsgs: batch, - decommissioned: decommissioned, - } + status.channel <-batch } } }() @@ -151,47 +122,46 @@ func (p *Producer) shutdown() { } // Allow the workers to drain the message channel first - wg := sync.WaitGroup{} - wg.Add(1) - go func(){ - defer wg.Done() - - staleTimeout := time.Duration(3 * time.Second) - timer := time.NewTimer(staleTimeout) - remaining := len(p.messages) - for remaining > 0 { - select { - case <-time.After(time.Second): - newRemaining := len(p.messages) - if newRemaining != remaining { - timer.Reset(staleTimeout) - remaining = newRemaining - } - case <-timer.C: - return + staleTimeout := time.Duration(3 * time.Second) + timer := time.NewTimer(staleTimeout) + remaining := len(p.messages) + + drain: + for remaining > 0 { + select { + case <-time.After(time.Second): + newRemaining := len(p.messages) + if newRemaining != remaining { + timer.Reset(staleTimeout) + remaining = newRemaining + } + case <-timer.C: + timer.Stop() + if remaining > 0 { + // TODO: Send remaining messages to the data spill } + break drain } - }() - wg.Wait() + } // Decommission all the workers p.resizeWorkerPool(0) // Close the decommission channel - if p.decommChannel != nil { - close(p.decommChannel) + if p.dismiss != nil { + close(p.dismiss) } // Close the status channel - if p.statusChannel != nil { - close(p.statusChannel) + if p.status != nil { + close(p.status) } // Stop the rate limiter p.rateLimiter.stop() // Stop the running go routine in produce - p.stopChannel <- empty{} + p.stop <- empty{} // Reset startupOnce to allow the start up sequence to happen again p.startupOnce = sync.Once{} @@ -204,45 +174,60 @@ func (p *Producer) resizeWorkerPool(desiredWorkerCount int) { p.resizeMu.Lock() defer p.resizeMu.Unlock() - // Get the number of total workers in the registry - totalWorkerCount := len(p.workerRegistry) - - - // Create a map of available workers - availableWorkers := make(map[string]*worker, totalWorkerCount) - for id, worker := range p.workerRegistry { - workerState := worker.getWorkerState() - if workerState == workerStateInactive || workerState == workerStateDecommissioned { - availableWorkers[id] = worker + if p.workerCount < desiredWorkerCount { + for p.workerCount < desiredWorkerCount { + go p.newWorker() + p.workerCount++ + } + } else { + for p.workerCount > desiredWorkerCount { + p.dismiss <- empty{} + p.workerCount-- } } - activeWorkerCount := totalWorkerCount - len(availableWorkers) +} - // We have too many workers at present than we actually need - if desiredWorkerCount < activeWorkerCount { - // Decommission the workers that we don't need - for i := activeWorkerCount; i <= desiredWorkerCount; i-- { - p.decommChannel <-empty{} - } - // We need more workers than are presently active or commissioned - } else if activeWorkerCount < desiredWorkerCount { - // Recommission those workers that are inactive - var activated int - for _, worker := range availableWorkers { - if activated > desiredWorkerCount - activeWorkerCount { - break +// newWorker is a (blocking) helper function that increases the number of concurrent go routines calling the +// sendBatch function. Communications between the produce function and the newWorker function occurs on the status +// channel where the "worker" provides information regarding its previously failed message count, its capacity for +// new messages and the "worker's" channel to which the produce function should send the batches. The "workers" also +// listen on the dismiss channel which upon receiving a signal will continue sending previously failed messages only +// until all failed messages have been sent successfully or aged out. +func (p *Producer) newWorker() { + batches := make(chan []*message.Message) + defer close(batches) + + var retries []*message.Message + var dismissed bool + for ok := true; ok; ok = !dismissed || len(retries) != 0 { + // Check to see if there were any signals to dismiss workers (if eligible) + if !dismissed { + select { + case <-p.dismiss: + dismissed = true + default: } - worker.start() - activated++ } - // Spawn new workers if still not enough - if desiredWorkerCount > totalWorkerCount { - for i := totalWorkerCount; i < desiredWorkerCount; i++ { - worker := newWorker(p.producerOptions, p.sendBatch, p.reportStatus) - p.workerRegistry[worker.workerID] = worker - } + // Send a status report to the status channel based on the number of previously failed messages and + // whether the worker has been dismissed + var capacity int + if !dismissed { + capacity = p.batchSize - len(retries) } + status := &statusReport { + capacity: capacity, + failed: len(retries), + channel: batches, + } + p.status <-status + + // Receive a batch of messages and call the producer's sendBatch function + batch := <-batches + if len(batch) + len(retries) > 0 { + retries = p.sendBatch(append(retries, batch...)) + } + } } @@ -250,9 +235,7 @@ func (p *Producer) resizeWorkerPool(desiredWorkerCount int) { // messages to send and returns a slice of messages that failed to send func (p *Producer) sendBatch(batch []*message.Message) []*message.Message { var failed []*message.Message - err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message, wg *sync.WaitGroup) error { - defer wg.Done() - + err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message) error { if msg.FailCount <= p.maxRetryAttempts { failed = append(failed, msg) p.Stats.AddSentRetried(1) @@ -292,11 +275,6 @@ func (p *Producer) sendBatch(batch []*message.Message) []*message.Message { return batch } -// reportStatus is used as a closure for the workers to report status to -func (p *Producer) reportStatus(report *statusReport) { - p.statusChannel <- report -} - // CloseWithContext initiates the graceful shutdown of the produce function, waiting for all outstanding messages and to // flush. Cancellation is supported through contexts. func (p *Producer) CloseWithContext(ctx context.Context) { diff --git a/producer/ratelimiter.go b/producer/ratelimiter.go index 55ea3b0..9f387c8 100644 --- a/producer/ratelimiter.go +++ b/producer/ratelimiter.go @@ -39,6 +39,7 @@ func (r *rateLimiter) start() { for { select { case <-r.stopChannel: + ticker.Stop() return case <-ticker.C: r.reset() @@ -58,6 +59,22 @@ func (r *rateLimiter) stop() { }) } +// restart calls stop then start to restart the rate limiter +func (r *rateLimiter) restart() { + r.stop() + r.start() +} + +// resize updates the maximum message size and reset frequency +func (r *rateLimiter) resize(limit int, duration time.Duration) { + r.tokenMu.Lock() + defer r.tokenMu.Unlock() + r.limit = limit + r.duration = duration + + r.restart() +} + // reset is called to reset the rateLimiter's tokens to the initial values func (r *rateLimiter) reset() { r.tokenMu.Lock() @@ -65,16 +82,35 @@ func (r *rateLimiter) reset() { r.tokenCount = r.limit } -// getTokenCount is used to retrieve the current token count +// getTokenCount is used to retrieve the current token count. Be aware of thread safety when trying to use +// getTokenCount and claimToken back to back as they are not tread safe with each other (use tryToClaimToken instead) func (r *rateLimiter) getTokenCount() int { r.tokenMu.Lock() defer r.tokenMu.Unlock() return r.tokenCount } -// claimTokens is used to claim tokens prior to sending messages +// claimTokens is used to claim tokens prior to sending messages. Be aware of thread safety when trying to use +// getTokenCount and claimToken back to back as they are not tread safe with each other (use tryToClaimToken instead) func (r *rateLimiter) claimTokens(count int) { r.tokenMu.Lock() defer r.tokenMu.Unlock() r.tokenCount -= count } + + +// tryToClaim attempts to claim the wanted number of tokens and returns the actual number claimed based on the +// number of tokens actually availab +func (r *rateLimiter) tryToClaimTokens(want int) (got int) { + r.tokenMu.Lock() + defer r.tokenMu.Unlock() + + if want <= r.tokenCount { + r.tokenCount -= want + got = want + } else { + r.tokenCount = 0 + got = r.tokenCount + } + return got +} diff --git a/producer/sharedtypes.go b/producer/sharedtypes.go index 7e8b8f6..bc38358 100644 --- a/producer/sharedtypes.go +++ b/producer/sharedtypes.go @@ -9,16 +9,9 @@ import ( // StreamWriter is an interface that abstracts the differences in API between Kinesis and Firehose. type StreamWriter interface { - PutRecords(context.Context, []*message.Message, MessageHandler) error + PutRecords(context.Context, []*message.Message, MessageHandlerAsync) error } -// sendBatchFn defines the signature for the sendBatch function defined by the producer and passed to the worker as a -// closure for execution -type sendBatchFn func([]*message.Message) []*message.Message - -// reportStatusFn defines the signature for the worker to communicate its status to the producer -type reportStatusFn func(*statusReport) - // MessageHandler defines the signature of a message handler used by PutRecords(). MessageHandler accepts a WaitGroup // so the function can be run as a blocking operation as opposed to MessageHandlerAsync. type MessageHandler func(*message.Message, *sync.WaitGroup) error @@ -27,6 +20,13 @@ type MessageHandler func(*message.Message, *sync.WaitGroup) error // be run asynchronously. type MessageHandlerAsync func(*message.Message) error +// statusReport is used to communicate a worker's capacity for new messages and to which channel they should be sent. +type statusReport struct { + capacity int // maximum message capacity the worker can handle + failed int // number of previous messages that failed to send + channel chan []*message.Message // channel of the worker to which the batch messages should be sent +} + // empty is used a as a dummy type for counting semaphore channels. type empty struct{} diff --git a/producer/worker.go b/producer/worker.go deleted file mode 100644 index 61e129c..0000000 --- a/producer/worker.go +++ /dev/null @@ -1,153 +0,0 @@ -package producer - -import ( - "crypto/rand" - "fmt" - "sync" - "time" - - "github.com/rewardStyle/kinetic/message" -) - -const ( - workerStateActive = "active" - workerStateIdle = "idle" - workerStateIdleWithRetries = "idleWithRetries" - workerStateBusy = "busy" - workerStateDecommissioned = "decommissioned" - workerStateInactive = "inactive" -) - -// statusReport is used to communicate a worker's status to the producer. -type statusReport struct { - workerID string // uuid of the worker making the work request - capacity int // maximum message capacity the worker can handle - timeout time.Duration // maximum time duration to wait for incoming messages - failed int // number of previous messages that failed to send -} - -// workerCommand is used by the producer to send commands to the workers. -type workerCommand struct { - batchMsgs []*message.Message // batch of messages for the worker to send - decommissioned bool // whether or not to decommission the worker -} - -// worker is an object that, when started, runs in a separate go routine handling commands from the producer through -// the worker's command channel. The worker communicates its status to the producer and the producer responds with -// a batch of messages for the worker to send. The worker calls a closure (defined by the producer) to send a batch -// of messages containing a combination of new messages and previously failed messages. The worker may also receive -// a command to decommission itself if the worker has no previously failed messages. -type worker struct { - workerID string // unique ID to identify itself from the other workers - state string // current state of the worker - stateMu sync.Mutex // mutex for protecting against concurrent changes of state - batchSize int // maximum number of message to be sent per batch - batchTimeout time.Duration // maximum duration with which to wait for a batch of new messages - sendBatchFn sendBatchFn // a closure to call when a batch is ready to be sent - reportStatusFn reportStatusFn // a clousre to call when the worker has a status update - retries []*message.Message // a slice of messages that previously failed - commands chan *workerCommand // a channel to which the producer sends commands -} - -// newWorker creates a new worker object that automatically starts working. -func newWorker(cfg *producerOptions, sendBatchFn sendBatchFn, reportStatusFn reportStatusFn) *worker { - var w *worker - defer func() { - go w.start() - }() - - w = &worker{ - workerID: generateWorkerID(), - state: workerStateActive, - batchSize: cfg.batchSize, - batchTimeout: cfg.batchTimeout, - sendBatchFn: sendBatchFn, - reportStatusFn: reportStatusFn, - commands: make(chan *workerCommand), - } - - return w -} - -// start is called when a new worker is instantiated via newWorker. start initiates an infinite loop that communicates -// the worker's status to the producer and listens to its command channel for commands sent by the producer -func (w *worker) start() { - defer w.shutdown() - - var failed int - for { - // Update worker state - if failed > 0 { - w.setWorkerState(workerStateIdleWithRetries) - } else { - w.setWorkerState(workerStateIdle) - } - - // Send the dispatcher a status report with the worker's capacity - go func() { - req := &statusReport{ - workerID: w.workerID, - capacity: w.batchSize - failed, - timeout: w.batchTimeout, - failed: failed, - } - w.reportStatusFn(req) - }() - - // Listen to incoming commands - cmd := <-w.commands - if cmd.decommissioned { - w.setWorkerState(workerStateDecommissioned) - return - } - - // Only send if there are any messages to send - if len(cmd.batchMsgs) + len(w.retries) > 0 { - w.setWorkerState(workerStateBusy) - - // Combine new batch messages with previously failed messages - var batch []*message.Message - batch = append(batch, w.retries...) - batch = append(batch, cmd.batchMsgs...) - - // Send batch and reset retries - w.retries = w.sendBatchFn(batch) - } - } -} - -// shutdown initiates the shutdown procedure for a worker -func (w *worker) shutdown() { - // Close the command channel to prevent any more incoming commands - if w.commands != nil { - close(w.commands) - } - - // Update the worker state to inactive - w.setWorkerState(workerStateInactive) -} - -// getWorkerState returns the worker's current state -func (w *worker) getWorkerState() string { - w.stateMu.Lock() - defer w.stateMu.Unlock() - return w.state -} - -// setWorkerStatus updates the worker's current state -func (w *worker) setWorkerState(state string) { - w.stateMu.Lock() - defer w.stateMu.Unlock() - w.state = state -} - -// generateWorkerID is a helper function that generates UUIDs that are used as unique worker IDs -func generateWorkerID() string { - b := make([]byte, 16) - _, err := rand.Read(b) - if err != nil { - fmt.Println("Error: ", err) - return "" - } - return fmt.Sprintf("%X-%X-%X-%X-%X", b[0:4], b[4:6], b[6:8], b[8:10], b[10:]) -} From a4f7a283941be76d2caf2fbe86385045c1968c0f Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 19 Jul 2017 10:24:49 -0500 Subject: [PATCH 055/109] Clean up the command line arguments for the test executable and updated README.md --- testexec/README.md | 35 +++++++++++++++++------------------ testexec/cli.go | 45 +++++++++++++++++++++------------------------ testexec/main.go | 20 +++++++------------- 3 files changed, 45 insertions(+), 55 deletions(-) diff --git a/testexec/README.md b/testexec/README.md index 22c5e87..3daf1a8 100644 --- a/testexec/README.md +++ b/testexec/README.md @@ -33,43 +33,42 @@ go install ```text Usage of ./testexec: - -cleanup + -blast + used to specify whether to call the producer's send function at full blast. (default false) + -clean used to specify whether or not to delete the kinesis stream after processing is complete. (default true) + -count int + used to specify the number of messages to (attempt to) send. This flag is only applicable to 'write' and 'readwrite' modes. Use zero or a negative number to produce indefinitely -duration int - used to specify the duration (in seconds) the program should run. This flag is only applicable to 'write' and 'readwrite' modes. Use zero or negative number to run indefinitely. (default 0) + used to specify the duration (in seconds) the program should run. This flag is only applicable to 'write' and 'readwrite' modes. Use zero or negative number to run indefinitely. -location string used to specify the location of the kinesis stream. Accepted values are (local|aws). For local, run kinesalite on http://127.0.0.1:4567. For aws, your aws credentials and configuration need to be defined at ~/.aws (default "local") -mode string used to specify the mode in which to run; either 'r', 'read', 'w', 'write', 'rw' or 'readwrite' (default "readwrite") - -num-msgs int - used to specify the number of messages to (attempt to) send. This flag is only applicable to 'write' and 'readwrite' modes. Use zero or a negative number to produce indefinitely (default 0) - -stream-name string - used to specify a pre-existing stream to be used for testing. A new stream will be created if not defined. - -throttle - used to specify whether to throttle PutRecord requests by 1 ms. (default true) + -stream string + used to specify a specific stream to write to / read from for testing. The stream will be created if a stream name is given but does not exist. A random stream will be created if the stream is undefined. -verbose - used to specify whether or not to log in verbose mode. (default false) + used to specify whether or not to log in verbose mode (default true) ``` ## Examples -To run kinetic testexec on a local kinesalite instance to stream a fixed number of messages to a new kinesis stream: +To run kinetic testexec to stream to / from a local kinesalite instance indefinitely to a new Kinesis stream: ```sh -./testexec -num-msgs 1000 -verbose +./testexec ``` -To run kinetic testexec on a local kinesalite instance to stream for a fixed duration of time to an existing kinesis stream: +To run kinetic testexec to stream a fixed number of messages to / from a new Kinesis stream: ```sh -./testexec -location local -stream-name some-stream -duration 1000 +./testexec -location aws -count 1000 ``` -To run kinetic testexec on an AWS Kinesis Stream to stream indefinitely (Ctrl-C to stop producing): +To run a stress test on the kinetic testexec to stream to / from a specific Kinesis stream and save the stream for confirmation later (Ctrl-C to stop producing): ```sh -./testexec -location aws -duration -1 -mode write -stream-name test-stream -cleanup=false -verbose +./testexec -location aws -stream my-stress-test-stream -blast=true -clean=false ``` -To run kinetic testexec in read mode an AWS KinesisStream: -To run kinetic testexec on an AWS Kinesis Stream to stream indefinitely (Ctrl-C to stop producing): +To run kinetic testexec in listen-only mode on an existing Kinesis stream: ```sh -./testexec -location aws -mode read -stream-name test-stream -cleanup=false -verbose +./testexec -location aws -mode read -stream my-stress-test-stream -clean=false ``` diff --git a/testexec/cli.go b/testexec/cli.go index 5ec044e..daaa2f5 100644 --- a/testexec/cli.go +++ b/testexec/cli.go @@ -11,11 +11,10 @@ type Config struct { Mode *string Location *string StreamName *string - NumMsgs *int + MsgCount *int Duration *int - Throttle *bool - Deadlock *bool - Cleanup *bool + Blast *bool + Clean *bool Verbose *bool } @@ -27,17 +26,17 @@ func parseCommandLineArgs() *Config { locationPtr := flag.String("location", "local", "used to specify the location of the kinesis stream. "+ "Accepted values are (local|aws). For local, run kinesalite on http://127.0.0.1:4567. For aws, your "+ "aws credentials and configuration need to be defined at ~/.aws") - streamNamePtr := flag.String("stream-name", "", "used to specify a pre-existing stream to be used for "+ - "testing. A new stream will be created if not defined.") - numMsgsPtr := flag.Int("num-msgs", 0, "used to specify the number of messages to (attempt to) send. This "+ + streamNamePtr := flag.String("stream", "", "used to specify a specific stream to write to / read from for " + + "testing. The stream will be created if a stream name is given but does not exist. A random stream will " + + "be created if the stream is undefined.") + msgCountPtr := flag.Int("count", 0, "used to specify the number of messages to (attempt to) send. This "+ "flag is only applicable to 'write' and 'readwrite' modes. Use zero or a negative number to produce "+ "indefinitely") durationPtr := flag.Int("duration", 0, "used to specify the duration (in seconds) the program should run. "+ "This flag is only applicable to 'write' and 'readwrite' modes. Use zero or negative number to run "+ "indefinitely.") - throttlePtr := flag.Bool("throttle", true, "used to specify whether to throttle PutRecord requests by 1 ms.") - deadlockPtr := flag.Bool("deadlock", false, "used to test potential deadlock condition for the producer.") - cleanupPtr := flag.Bool("cleanup", true, "used to specify whether or not to delete the kinesis stream after "+ + blastPtr := flag.Bool("blast", false, "used to specify whether to call the producer's send function at full blast.") + cleanPtr := flag.Bool("clean", true, "used to specify whether or not to delete the kinesis stream after "+ "processing is complete.") verbosePtr := flag.Bool("verbose", true, "used to specify whether or not to log in verbose mode") @@ -49,8 +48,8 @@ func parseCommandLineArgs() *Config { durationPtr = nil } - if *numMsgsPtr <= 0 { - numMsgsPtr = nil + if *msgCountPtr <= 0 { + msgCountPtr = nil } var mode string @@ -75,11 +74,10 @@ func parseCommandLineArgs() *Config { Mode: &mode, StreamName: streamNamePtr, Duration: durationPtr, - NumMsgs: numMsgsPtr, + MsgCount: msgCountPtr, Location: locationPtr, - Throttle: throttlePtr, - Deadlock: deadlockPtr, - Cleanup: cleanupPtr, + Blast: blastPtr, + Clean: cleanPtr, Verbose: verbosePtr, } } @@ -90,23 +88,22 @@ func (c *Config) printConfigs() { log.Println("-mode: ", *c.Mode) log.Println("-location: ", *c.Location) if len(*c.StreamName) == 0 { - log.Println("-stream-name: (randomly generated)") + log.Println("-stream: (randomly generated)") } else { - log.Println("-stream-name: ", *c.StreamName) + log.Println("-stream: ", *c.StreamName) } - if c.NumMsgs != nil { - log.Println("-num-msgs: ", *c.NumMsgs) + if c.MsgCount != nil { + log.Println("-count: ", *c.MsgCount) } else { - log.Println("-num-msgs: (unbounded)") + log.Println("-count: (unbounded)") } if c.Duration != nil { log.Printf("-duration: [%d] (s)", *c.Duration) } else { log.Println("-duration: (indefinite)") } - log.Println("-throttle: ", *c.Throttle) - log.Println("-deadlock: ", *c.Deadlock) - log.Println("-cleanup: ", *c.Cleanup) + log.Println("-blast: ", *c.Blast) + log.Println("-cleanup: ", *c.Clean) log.Println("-verbose: ", *c.Verbose) log.Println() } diff --git a/testexec/main.go b/testexec/main.go index a9f02fd..ddda6a1 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -78,7 +78,7 @@ func init() { } func cleanup(k *kinetic.Kinetic, stream string) { - if *cfg.Cleanup { + if *cfg.Clean { if *cfg.Verbose { log.Println() log.Printf("Cleaning up by deleting stream [%s] ...\n", stream) @@ -252,13 +252,7 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce c.SetBatchSize(500) c.SetMaxRetryAttempts(3) c.SetStatsCollector(psc) - if *cfg.Deadlock { - //c.SetConcurrency(10) - c.SetQueueDepth(1) - } else { - //c.SetConcurrency(10) - c.SetQueueDepth(10000) - } + c.SetQueueDepth(10000) }) if err != nil { log.Fatalf("Unable to create a new producer due to: %v\n", err) @@ -410,10 +404,10 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { var sent uint64 var sendTicker *time.Ticker - if *cfg.Throttle { - sendTicker = time.NewTicker(time.Millisecond) - } else { + if *cfg.Blast { sendTicker = time.NewTicker(time.Nanosecond) + } else { + sendTicker = time.NewTicker(time.Millisecond) } produce: for { @@ -432,8 +426,8 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { break produce case <-sendTicker.C: // Break from the loop if we have sent the correct number of messages - if cfg.NumMsgs != nil { - if atomic.LoadUint64(&sent) >= uint64(*cfg.NumMsgs) { + if cfg.MsgCount != nil { + if atomic.LoadUint64(&sent) >= uint64(*cfg.MsgCount) { break produce } } From 529ef1935dfb3526236a9828f84590101b44de91 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 19 Jul 2017 16:19:29 -0500 Subject: [PATCH 056/109] Fixed a bug in the produce function where we weren't waiting for the rateLimiter to reset by added a reset notification channel in the rateLimiter that notifies when the token count has been reset. --- producer/producer.go | 15 +++++++++++++-- producer/producerconfig.go | 2 +- producer/ratelimiter.go | 29 ++++++++++++++++++++++------- testexec/main.go | 7 +++---- 4 files changed, 39 insertions(+), 14 deletions(-) diff --git a/producer/producer.go b/producer/producer.go index 02d695c..bea9331 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -89,8 +89,14 @@ func (p *Producer) produce() { case <-p.stop: return case status := <-p.status: + var tokenCount int + for tokenCount = p.rateLimiter.getTokenCount(); tokenCount == 0; { + // Wait for a reset notification from the rateLimiter if needed + <-p.rateLimiter.resetChannel + tokenCount = p.rateLimiter.getTokenCount() + } + var batch []*message.Message - tokenCount := p.rateLimiter.getTokenCount() timeout := time.After(p.batchTimeout) fillBatch: @@ -104,8 +110,13 @@ func (p *Producer) produce() { } } } - p.rateLimiter.claimTokens(len(batch) + status.failed) + // Claim tokens only if needed + if len(batch) + status.failed > 0 { + p.rateLimiter.claimTokens(len(batch) + status.failed) + } + + // Send batch regardless if it is empty or not status.channel <-batch } } diff --git a/producer/producerconfig.go b/producer/producerconfig.go index 24ff443..c495e65 100644 --- a/producer/producerconfig.go +++ b/producer/producerconfig.go @@ -22,7 +22,7 @@ func NewConfig(cfg *aws.Config) *Config { producerOptions: &producerOptions{ batchSize: 500, batchTimeout: time.Second, - queueDepth: 500, + queueDepth: 10000, maxRetryAttempts: 10, workersPerShard: 5, shardCount: 1, diff --git a/producer/ratelimiter.go b/producer/ratelimiter.go index 9f387c8..ec6fb7d 100644 --- a/producer/ratelimiter.go +++ b/producer/ratelimiter.go @@ -9,13 +9,14 @@ import ( // rateLimiter is used by the producer to rate limit the total number and size of records sent per cycle. type rateLimiter struct { - limit int // upper limit of throughput per cycle - duration time.Duration // frequency with which to reset the remaining tokens count - tokenCount int // remaining tokens available for the cycle - tokenMu sync.Mutex // mutex to protect remaining token count - stopChannel chan empty // channel for communicating when to stop rate limiting - startOnce sync.Once // startOnce is used to ensure that start is called once and only once - stopOnce sync.Once // stopOnce is used to ensure that stop is called once and only once + limit int // upper limit of throughput per cycle + duration time.Duration // frequency with which to reset the remaining tokens count + tokenCount int // remaining tokens available for the cycle + tokenMu sync.Mutex // mutex to protect remaining token count + stopChannel chan empty // channel for communicating when to stop rate limiting + resetChannel chan empty // channel for communicating when the rate limiter has been reset + startOnce sync.Once // startOnce is used to ensure that start is called once and only once + stopOnce sync.Once // stopOnce is used to ensure that stop is called once and only once } // newRateLimiter creates a new rateLimiter. @@ -34,6 +35,7 @@ func (r *rateLimiter) start() { r.stopOnce = sync.Once{} r.stopChannel = make(chan empty) + r.resetChannel = make(chan empty) ticker := time.NewTicker(r.duration) go func(){ for { @@ -54,6 +56,9 @@ func (r *rateLimiter) stop() { r.stopOnce.Do(func(){ r.stopChannel <- empty{} + close(r.stopChannel) + close(r.resetChannel) + // Reset startOnce to allow the rateLimiter to be started again r.startOnce = sync.Once{} }) @@ -80,6 +85,16 @@ func (r *rateLimiter) reset() { r.tokenMu.Lock() defer r.tokenMu.Unlock() r.tokenCount = r.limit + + // Send a signal to the resetChanel and remove it immediately if no one picked it up + select { + case r.resetChannel <-empty{}: + default: + select { + case <-r.resetChannel: + default: + } + } } // getTokenCount is used to retrieve the current token count. Be aware of thread safety when trying to use diff --git a/testexec/main.go b/testexec/main.go index ddda6a1..144e9fa 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -239,7 +239,7 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce psc := producer.NewDefaultStatsCollector(registry) w, err := producer.NewKinesisWriter(k.Session.Config, streamName, func(kwc *producer.KinesisWriterConfig) { - kwc.SetLogLevel(aws.LogDebug) + //kwc.SetLogLevel(logging.LogDebug) kwc.SetResponseReadTimeout(time.Second) kwc.SetStatsCollector(psc) }) @@ -253,6 +253,7 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce c.SetMaxRetryAttempts(3) c.SetStatsCollector(psc) c.SetQueueDepth(10000) + c.SetWorkersPerShard(2) }) if err != nil { log.Fatalf("Unable to create a new producer due to: %v\n", err) @@ -385,13 +386,11 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { for { <-sendSignal jsonStr, _ := json.Marshal(NewMessage()) - if err := p.Send(&message.Message{ + if err := p.TryToSend(&message.Message{ PartitionKey: aws.String("key"), Data: []byte(jsonStr), }); err == nil { sd.incrementMsgCount() - } else { - log.Println("producer: Uh oh, something bad happened!!!!") } } }() From 19ebef50a5ebadb2b80a4588e3668b71452845c4 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 19 Jul 2017 16:30:29 -0500 Subject: [PATCH 057/109] Fixed broken unit test --- producer/producerconfig_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go index 73fa58f..7f025ae 100644 --- a/producer/producerconfig_test.go +++ b/producer/producerconfig_test.go @@ -46,7 +46,7 @@ func TestNewConfig(t *testing.T) { Convey("check the default values for its non-zero config", func() { So(cfg.batchSize, ShouldEqual, 500) So(cfg.batchTimeout, ShouldEqual, 1*time.Second) - So(cfg.queueDepth, ShouldEqual, 500) + So(cfg.queueDepth, ShouldEqual, 10000) So(cfg.maxRetryAttempts, ShouldEqual, 10) So(cfg.workersPerShard, ShouldEqual, 5) So(cfg.shardCount, ShouldEqual, 1) From f2159adaf26df7a76a0d365cd34cd97501ff1f5e Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 21 Jul 2017 07:58:54 -0500 Subject: [PATCH 058/109] Replaced the rateLimiter implementation with a golang library (github.com/golang/time/rate) to resolve an unusually high throughput exceeded exception caused by a race condition. Added a message size rate limiter to throttle the dispatching of messages according to the message (transmission) size as well. --- glide.lock | 6 +- glide.yaml | 1 + message/message.go | 11 +++ producer/producer.go | 129 ++++++++++++++++++++----------- producer/producerconfig.go | 16 ++-- producer/producerconfig_test.go | 19 ++++- producer/ratelimiter.go | 131 -------------------------------- producer/ratelimiter_test.go | 50 ------------ producer/sharedtypes.go | 7 +- testexec/main.go | 4 +- 10 files changed, 130 insertions(+), 244 deletions(-) delete mode 100644 producer/ratelimiter.go delete mode 100644 producer/ratelimiter_test.go diff --git a/glide.lock b/glide.lock index fa3df35..00b8b30 100644 --- a/glide.lock +++ b/glide.lock @@ -1,5 +1,5 @@ -hash: 70d3e36401b9615bc762c364243c4177b3b4b51d36e3fdeaa3ec923996430e24 -updated: 2017-07-06T15:44:39.094878239-05:00 +hash: 478410f2430633b92c3939e14a69e21c2f95684bae1e74250640ccb998d23f48 +updated: 2017-07-20T11:57:12.855440218-05:00 imports: - name: github.com/aws/aws-sdk-go version: 72e42b13da62269f68308fb6068b7ea691a416a4 @@ -36,6 +36,8 @@ imports: - service/sts - name: github.com/go-ini/ini version: 1730955e3146956d6a087861380f9b4667ed5071 +- name: github.com/golang/time + version: 8be79e1e0910c292df4e79c241bb7e8f7e725959 - name: github.com/gopherjs/gopherjs version: 0cc212b4ef1943b61e894a7d56490575b3216439 subpackages: diff --git a/glide.yaml b/glide.yaml index 7c867c6..eb4b8ba 100644 --- a/glide.yaml +++ b/glide.yaml @@ -64,3 +64,4 @@ import: - package: github.com/stretchr/testify/assert version: v1.1.4 - package: github.com/rcrowley/go-metrics +- package: github.com/golang/time diff --git a/message/message.go b/message/message.go index 2c51ba6..578455a 100644 --- a/message/message.go +++ b/message/message.go @@ -1,6 +1,7 @@ package message import ( + "encoding/json" "time" "github.com/aws/aws-sdk-go/service/firehose" @@ -39,6 +40,16 @@ func FromRecord(record *kinesis.Record) *Message { } } +// RequestEntrySize calculates what the size (in bytes) of the message will be after calling ToRequestEntry on it and +// marshalling it to json +func (m *Message) RequestEntrySize() int { + buf, err := json.Marshal(m.ToRequestEntry()) + if err != nil { + + } + return len(buf) +} + // ToRequestEntry creates a kinesis.PutRecordsRequestEntry to be used in the kinesis.PutRecords API call. func (m *Message) ToRequestEntry() *kinesis.PutRecordsRequestEntry { return &kinesis.PutRecordsRequestEntry{ diff --git a/producer/producer.go b/producer/producer.go index bea9331..e584a40 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -8,6 +8,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/awserr" + "github.com/golang/time/rate" "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" "github.com/rewardStyle/kinetic/errs" @@ -21,27 +22,28 @@ type producerOptions struct { maxRetryAttempts int // maximum number of retry attempts for failed messages workersPerShard int // number of concurrent workers per shard shardCount int // initial shard size - rateLimit int // maximum records to be sent per cycle for the rate limiting model - resetFrequency time.Duration // duration of a cycle for the rate limiting model + msgCountLimit int // maximum records to be sent per cycle for the rate limiting model + msgSizeLimit int // maximum records to be sent per cycle for the rate limiting model Stats StatsCollector // stats collection mechanism } // Producer sends records to AWS Kinesis or Firehose. type Producer struct { - *producerOptions // contains all of the configuration settings for the Producer - *logging.LogHelper // object for help with logging - writer StreamWriter // interface for abstracting the PutRecords call - rateLimiter *rateLimiter // throttles the number of messages sent based on total count and size - workerCount int // number of concurrent workers sending batch messages for the producer - messages chan *message.Message // channel for enqueuing messages to be put on the stream - status chan *statusReport // channel for workers to communicate their current status - dismiss chan empty // channel for handling the decommissioning of a surplus of workers - stop chan empty // channel for handling shutdown - pipeOfDeath chan empty // channel for handling pipe of death - startupOnce sync.Once // used to ensure that the startup function is called once - shutdownOnce sync.Once // used to ensure that the shutdown function is called once - resizeMu sync.Mutex // used to prevent resizeWorkerPool from being called synchronously with itself - noCopy noCopy // prevents the Producer from being copied + *producerOptions // contains all of the configuration settings for the Producer + *logging.LogHelper // object for help with logging + writer StreamWriter // interface for abstracting the PutRecords call + msgCountLimiter *rate.Limiter // + msgSizeLimiter *rate.Limiter // + workerCount int // number of concurrent workers sending batch messages for the producer + messages chan *message.Message // channel for enqueuing messages to be put on the stream + status chan *statusReport // channel for workers to communicate their current status + dismiss chan empty // channel for handling the decommissioning of a surplus of workers + stop chan empty // channel for handling shutdown + pipeOfDeath chan empty // channel for handling pipe of death + startupOnce sync.Once // used to ensure that the startup function is called once + shutdownOnce sync.Once // used to ensure that the shutdown function is called once + resizeMu sync.Mutex // used to prevent resizeWorkerPool from being called synchronously with itself + noCopy noCopy // prevents the Producer from being copied } // NewProducer creates a new producer for writing records to a Kinesis or Firehose stream. @@ -66,8 +68,11 @@ func (p *Producer) produce() { // Reset shutdownOnce to allow the shut down sequence to happen again p.shutdownOnce = sync.Once{} + // Instantiate rate limiters + p.msgCountLimiter = rate.NewLimiter(rate.Limit(float64(p.msgCountLimit)), p.batchSize) + p.msgSizeLimiter = rate.NewLimiter(rate.Limit(float64(p.msgSizeLimit)), p.msgSizeLimit) + // Create communication channels - p.rateLimiter = newRateLimiter(p.rateLimit, p.resetFrequency) p.messages = make(chan *message.Message, p.queueDepth) p.status = make(chan *statusReport) p.dismiss = make(chan empty) @@ -77,9 +82,6 @@ func (p *Producer) produce() { // Instantiate and register new workers p.resizeWorkerPool(p.shardCount * p.workersPerShard) - // Instantiate and start a new rate limiter - p.rateLimiter.start() - go func(){ // Dispatch messages to each worker depending on the worker's capacity and the rate limit for { @@ -89,32 +91,64 @@ func (p *Producer) produce() { case <-p.stop: return case status := <-p.status: - var tokenCount int - for tokenCount = p.rateLimiter.getTokenCount(); tokenCount == 0; { - // Wait for a reset notification from the rateLimiter if needed - <-p.rateLimiter.resetChannel - tokenCount = p.rateLimiter.getTokenCount() - } - var batch []*message.Message - timeout := time.After(p.batchTimeout) - - fillBatch: - for len(batch) < status.capacity && len(batch) + status.failed < tokenCount { - select { - case <-timeout: - break fillBatch - case msg := <-p.messages: - if msg != nil { + + // Start pulling from the message channel and waiting for the bucket token to + // fill at the same time + wg := sync.WaitGroup{} + wg.Add(1) + go func() { + defer wg.Done() + + var batchMsgSize int + timeout := time.After(p.batchTimeout) + + fillBatch: + for len(batch) < status.capacity { + select { + case <-timeout: + break fillBatch + case msg := <-p.messages: + if msg == nil { + // Drop nil message + break + } + + var msgSize int + if msgSize = msg.RequestEntrySize(); msgSize > p.msgSizeLimit { + // TODO: Send this huge message to data spill + break + } + batchMsgSize += msgSize batch = append(batch, msg) } } - } - // Claim tokens only if needed - if len(batch) + status.failed > 0 { - p.rateLimiter.claimTokens(len(batch) + status.failed) - } + // Request and wait for the message size (transmission) rate limiter to + // allow this payload + ctx, cancel := context.WithTimeout(context.TODO(), p.batchTimeout) + if err := p.msgSizeLimiter.WaitN(ctx, batchMsgSize + status.failedSize); err != nil { + + } + cancel() + }() + + wg.Add(1) + go func() { + defer wg.Done() + + // Request and wait for the message counter rate limiter to allow this + // payload + ctx, cancel := context.WithTimeout(context.TODO(), p.batchTimeout) + err := p.msgCountLimiter.WaitN(ctx, status.capacity + status.failedCount); + if err != nil { + // TODO: handle error properly + } + cancel() + }() + + // Wait until the batch is ready and the rate limiter gives the go ahead + wg.Wait() // Send batch regardless if it is empty or not status.channel <-batch @@ -168,9 +202,6 @@ func (p *Producer) shutdown() { close(p.status) } - // Stop the rate limiter - p.rateLimiter.stop() - // Stop the running go routine in produce p.stop <- empty{} @@ -223,12 +254,18 @@ func (p *Producer) newWorker() { // Send a status report to the status channel based on the number of previously failed messages and // whether the worker has been dismissed var capacity int + failedCount := len(retries) if !dismissed { - capacity = p.batchSize - len(retries) + capacity = p.batchSize - failedCount + } + var failedSize int + for _, msg := range retries { + failedSize += msg.RequestEntrySize() } status := &statusReport { capacity: capacity, - failed: len(retries), + failedCount: failedCount, + failedSize: failedSize, channel: batches, } p.status <-status diff --git a/producer/producerconfig.go b/producer/producerconfig.go index c495e65..99aa3e9 100644 --- a/producer/producerconfig.go +++ b/producer/producerconfig.go @@ -26,8 +26,8 @@ func NewConfig(cfg *aws.Config) *Config { maxRetryAttempts: 10, workersPerShard: 5, shardCount: 1, - rateLimit: 1000, - resetFrequency: time.Second, + msgCountLimit: 1000, + msgSizeLimit: 1000000, Stats: &NilStatsCollector{}, }, LogLevel: *cfg.LogLevel, @@ -70,14 +70,14 @@ func (c *Config) SetShardCount(count int) { c.shardCount = count } -// SetRateLimit defines the maximum number of message to send per cycle -func (c *Config) SetRateLimit(limit int) { - c.rateLimit = limit +// SetMsgCountLimit defines the maximum number of message that can be sent per second +func (c *Config) SetMsgCountLimit(limit int) { + c.msgCountLimit = limit } -// SetResetFrequency defines the frequency at which the rateLimiter resets -func (c *Config) SetResetFrequency(freq time.Duration) { - c.resetFrequency = freq +// SetMsgSizeLimit defines the maximum size (in bytes) that can be sent per second +func (c *Config) SetMsgSizeLimit(limit int) { + c.msgSizeLimit = limit } // SetStatsCollector configures a listener to handle producer metrics. diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go index 7f025ae..15dc251 100644 --- a/producer/producerconfig_test.go +++ b/producer/producerconfig_test.go @@ -50,8 +50,8 @@ func TestNewConfig(t *testing.T) { So(cfg.maxRetryAttempts, ShouldEqual, 10) So(cfg.workersPerShard, ShouldEqual, 5) So(cfg.shardCount, ShouldEqual, 1) - So(cfg.rateLimit, ShouldEqual, 1000) - So(cfg.resetFrequency, ShouldEqual, time.Second) + So(cfg.msgCountLimit, ShouldEqual, 1000) + So(cfg.msgSizeLimit, ShouldEqual, 1000000) So(cfg.Stats, ShouldHaveSameTypeAs, &NilStatsCollector{}) So(cfg.LogLevel.Value(), ShouldEqual, logging.LogOff) }) @@ -82,7 +82,20 @@ func TestNewConfig(t *testing.T) { So(cfg.maxRetryAttempts, ShouldEqual, 100) }) - // TODO: Add tests for setting new producer options + Convey("check that we can set the workers per shard", func() { + cfg.SetWorkersPerShard(10) + So(cfg.workersPerShard, ShouldEqual, 10) + }) + + Convey("check that we can set the message count rate limit", func() { + cfg.SetMsgCountLimit(100) + So(cfg.msgCountLimit, ShouldEqual, 100) + }) + + Convey("check that we can set the message size rate limit", func() { + cfg.SetMsgSizeLimit(10) + So(cfg.msgSizeLimit, ShouldEqual, 10) + }) Convey("check that we can configure a stats collector", func() { cfg.SetStatsCollector(&DebugStatsCollector{}) diff --git a/producer/ratelimiter.go b/producer/ratelimiter.go deleted file mode 100644 index ec6fb7d..0000000 --- a/producer/ratelimiter.go +++ /dev/null @@ -1,131 +0,0 @@ -package producer - -import ( - "sync" - "time" -) - -// TODO: add size rate limiting in addition to message count - -// rateLimiter is used by the producer to rate limit the total number and size of records sent per cycle. -type rateLimiter struct { - limit int // upper limit of throughput per cycle - duration time.Duration // frequency with which to reset the remaining tokens count - tokenCount int // remaining tokens available for the cycle - tokenMu sync.Mutex // mutex to protect remaining token count - stopChannel chan empty // channel for communicating when to stop rate limiting - resetChannel chan empty // channel for communicating when the rate limiter has been reset - startOnce sync.Once // startOnce is used to ensure that start is called once and only once - stopOnce sync.Once // stopOnce is used to ensure that stop is called once and only once -} - -// newRateLimiter creates a new rateLimiter. -func newRateLimiter(limit int, duration time.Duration) *rateLimiter { - return &rateLimiter{ - limit: limit, - duration: duration, - tokenCount: limit, - } -} - -// start runs a timer in a go routine background which resets the the number and size counters every cycle. -func (r *rateLimiter) start() { - r.startOnce.Do(func() { - // Reset stopOnce to allow the rateLimiter to be shut down again - r.stopOnce = sync.Once{} - - r.stopChannel = make(chan empty) - r.resetChannel = make(chan empty) - ticker := time.NewTicker(r.duration) - go func(){ - for { - select { - case <-r.stopChannel: - ticker.Stop() - return - case <-ticker.C: - r.reset() - } - } - }() - }) -} - -// stop sends a signal to the rateLimiter's stopChannel -func (r *rateLimiter) stop() { - r.stopOnce.Do(func(){ - r.stopChannel <- empty{} - - close(r.stopChannel) - close(r.resetChannel) - - // Reset startOnce to allow the rateLimiter to be started again - r.startOnce = sync.Once{} - }) -} - -// restart calls stop then start to restart the rate limiter -func (r *rateLimiter) restart() { - r.stop() - r.start() -} - -// resize updates the maximum message size and reset frequency -func (r *rateLimiter) resize(limit int, duration time.Duration) { - r.tokenMu.Lock() - defer r.tokenMu.Unlock() - r.limit = limit - r.duration = duration - - r.restart() -} - -// reset is called to reset the rateLimiter's tokens to the initial values -func (r *rateLimiter) reset() { - r.tokenMu.Lock() - defer r.tokenMu.Unlock() - r.tokenCount = r.limit - - // Send a signal to the resetChanel and remove it immediately if no one picked it up - select { - case r.resetChannel <-empty{}: - default: - select { - case <-r.resetChannel: - default: - } - } -} - -// getTokenCount is used to retrieve the current token count. Be aware of thread safety when trying to use -// getTokenCount and claimToken back to back as they are not tread safe with each other (use tryToClaimToken instead) -func (r *rateLimiter) getTokenCount() int { - r.tokenMu.Lock() - defer r.tokenMu.Unlock() - return r.tokenCount -} - -// claimTokens is used to claim tokens prior to sending messages. Be aware of thread safety when trying to use -// getTokenCount and claimToken back to back as they are not tread safe with each other (use tryToClaimToken instead) -func (r *rateLimiter) claimTokens(count int) { - r.tokenMu.Lock() - defer r.tokenMu.Unlock() - r.tokenCount -= count -} - - -// tryToClaim attempts to claim the wanted number of tokens and returns the actual number claimed based on the -// number of tokens actually availab -func (r *rateLimiter) tryToClaimTokens(want int) (got int) { - r.tokenMu.Lock() - defer r.tokenMu.Unlock() - - if want <= r.tokenCount { - r.tokenCount -= want - got = want - } else { - r.tokenCount = 0 - got = r.tokenCount - } - return got -} diff --git a/producer/ratelimiter_test.go b/producer/ratelimiter_test.go deleted file mode 100644 index 7513172..0000000 --- a/producer/ratelimiter_test.go +++ /dev/null @@ -1,50 +0,0 @@ -package producer - -import ( - . "github.com/smartystreets/goconvey/convey" - - "testing" - "time" -) - -func TestRateLimiter(t *testing.T) { - limit := 1000 - duration := time.Second - Convey("calling newRateLimiter should return the pointer to a rateLimiter with properly initalized values", t, func() { - rl := newRateLimiter(limit, duration) - So(rl, ShouldNotBeNil) - So(rl.limit, ShouldEqual, limit) - So(rl.duration, ShouldEqual, duration) - So(rl.tokenCount, ShouldEqual, limit) - - Convey("calling getTokenCount returns the correct token count", func() { - So(rl.getTokenCount(), ShouldEqual, limit) - }) - - Convey("calling claimTokens reduces the tokens properly", func() { - So(rl.getTokenCount(), ShouldEqual, limit) - claim := 100 - rl.claimTokens(claim) - So(rl.getTokenCount(), ShouldEqual, limit - claim) - }) - - Convey("calling reset will reset the token count back to the original limit", func() { - So(rl.getTokenCount(), ShouldEqual, limit) - claim := 100 - rl.claimTokens(claim) - So(rl.getTokenCount(), ShouldNotEqual, limit) - rl.reset() - So(rl.getTokenCount(), ShouldEqual, limit) - }) - - Convey("calling start on the rateLimiter", func() { - // TODO: Finish this unit test - rl.start() - time.Sleep(10 * time.Second) - }) - - Convey("calling stop on the rateLimiter", func() { - // TODO: Finish this unit test - }) - }) -} diff --git a/producer/sharedtypes.go b/producer/sharedtypes.go index bc38358..ad99690 100644 --- a/producer/sharedtypes.go +++ b/producer/sharedtypes.go @@ -22,9 +22,10 @@ type MessageHandlerAsync func(*message.Message) error // statusReport is used to communicate a worker's capacity for new messages and to which channel they should be sent. type statusReport struct { - capacity int // maximum message capacity the worker can handle - failed int // number of previous messages that failed to send - channel chan []*message.Message // channel of the worker to which the batch messages should be sent + capacity int // maximum message capacity the worker can handle + failedCount int // number of previous messages that failed to send + failedSize int // size in bytes of the previous messages that failed to send + channel chan []*message.Message // channel of the worker to which the batch messages should be sent } // empty is used a as a dummy type for counting semaphore channels. diff --git a/testexec/main.go b/testexec/main.go index 144e9fa..299e603 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -253,7 +253,9 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce c.SetMaxRetryAttempts(3) c.SetStatsCollector(psc) c.SetQueueDepth(10000) - c.SetWorkersPerShard(2) + c.SetWorkersPerShard(3) + c.SetMsgCountLimit(1000) + c.SetMsgSizeLimit(1000000) }) if err != nil { log.Fatalf("Unable to create a new producer due to: %v\n", err) From a685a900e27ff64d8cb7010067de01ac1809ca82 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 24 Jul 2017 12:39:00 -0500 Subject: [PATCH 059/109] Implemented shard monitoring for the kinesiswriter and the produce function of the producer. Re-arranged the rate limiter properties to be part of the kineswriter and firehosewriter configuration instead of the producer's. Ran gofmt. --- producer/firehosewriter.go | 19 +++++- producer/firehosewriterconfig.go | 25 ++++++- producer/kinesiswriter.go | 41 +++++++++++- producer/kinesiswriterconfig.go | 23 +++++++ producer/producer.go | 109 ++++++++++++++++++++----------- producer/producer_test.go | 49 +++++++++----- producer/producerconfig.go | 28 +++----- producer/producerconfig_test.go | 21 ++---- producer/sharedtypes.go | 5 +- producer/stats.go | 2 +- testexec/cli.go | 6 +- testexec/main.go | 7 +- 12 files changed, 232 insertions(+), 103 deletions(-) diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go index bfc6a1b..aab0fa6 100644 --- a/producer/firehosewriter.go +++ b/producer/firehosewriter.go @@ -17,7 +17,9 @@ import ( ) type firehoseWriterOptions struct { - Stats StatsCollector + msgCountRateLimit int // maximum number of records to be sent per second + msgSizeRateLimit int // maximum (transmission) size of records to be sent per second + Stats StatsCollector // stats collection mechanism } // FirehoseWriter handles the API to send records to Kinesis. @@ -129,3 +131,18 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Mes return nil } + +// getMsgCountRateLimit returns the writer's message count rate limit +func (w *FirehoseWriter) getMsgCountRateLimit() int { + return w.msgCountRateLimit +} + +// getMsgSizeRateLimit returns the writer's message size rate limit +func (w *FirehoseWriter) getMsgSizeRateLimit() int { + return w.msgSizeRateLimit +} + +// getConcurrencyMultiplier returns the writer's concurrency multiplier. For the firehosewriter the multiplier is 1. +func (w *FirehoseWriter) getConcurrencyMultiplier() (int, error) { + return 1, nil +} diff --git a/producer/firehosewriterconfig.go b/producer/firehosewriterconfig.go index 2e2c36b..ef9c023 100644 --- a/producer/firehosewriterconfig.go +++ b/producer/firehosewriterconfig.go @@ -4,6 +4,11 @@ import ( "github.com/aws/aws-sdk-go/aws" ) +const ( + firehoseMsgCountRateLimit = 5000 // AWS Firehose limit of 5000 records/sec + firehoseMsgSizeRateLimit = 5000000 // AWS Firehose limit of 5 MB/sec +) + // FirehoseWriterConfig is used to configure FirehoseWriter type FirehoseWriterConfig struct { *firehoseWriterOptions @@ -16,12 +21,30 @@ func NewFirehoseWriterConfig(cfg *aws.Config) *FirehoseWriterConfig { return &FirehoseWriterConfig{ AwsConfig: cfg, firehoseWriterOptions: &firehoseWriterOptions{ - Stats: &NilStatsCollector{}, + msgCountRateLimit: firehoseMsgCountRateLimit, + msgSizeRateLimit: firehoseMsgSizeRateLimit, + Stats: &NilStatsCollector{}, }, LogLevel: *cfg.LogLevel, } } +// SetMsgCountRateLimit configures the maximum number of messages that can be sent per second +func (c *FirehoseWriterConfig) SetMsgCountRateLimit(limit int) { + if limit > firehoseMsgCountRateLimit { + + } + c.msgCountRateLimit = limit +} + +// SetMsgSizeRateLimit configures the maximum transmission size of the messages that can be sent per second +func (c *FirehoseWriterConfig) SetMsgSizeRateLimit(limit int) { + if limit > firehoseMsgSizeRateLimit { + + } + c.msgSizeRateLimit = limit +} + // SetStatsCollector configures a listener to handle listener metrics. func (c *FirehoseWriterConfig) SetStatsCollector(stats StatsCollector) { c.Stats = stats diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index 6f07913..1ff2ee7 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -17,8 +17,10 @@ import ( ) type kinesisWriterOptions struct { - responseReadTimeout time.Duration - Stats StatsCollector + responseReadTimeout time.Duration // maximum time to wait for PutRecords API call before timing out + msgCountRateLimit int // maximum number of records to be sent per second + msgSizeRateLimit int // maximum (transmission) size of records to be sent per second + Stats StatsCollector // stats collection mechanism } // KinesisWriter handles the API to send records to Kinesis. @@ -132,3 +134,38 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess return nil } + +// getMsgCountRateLimit returns the writer's message count rate limit +func (w *KinesisWriter) getMsgCountRateLimit() int { + return w.msgCountRateLimit +} + +// getMsgSizeRateLimit returns the writer's message size rate limit +func (w *KinesisWriter) getMsgSizeRateLimit() int { + return w.msgSizeRateLimit +} + +// getConcurrencyMultiplier returns the writer's concurrency multiplier. For the kinesiswriter the multiplier is the +// number of active shards for the Kinesis stream +func (w *KinesisWriter) getConcurrencyMultiplier() (int, error) { + resp, err := w.client.DescribeStream(&kinesis.DescribeStreamInput{ + StreamName: aws.String(w.stream), + }) + if err != nil { + w.LogError("Error describing kinesis stream: ", w.stream, err) + return 0, err + } + if resp == nil { + return 0, errs.ErrNilDescribeStreamResponse + } + if resp.StreamDescription == nil { + return 0, errs.ErrNilStreamDescription + } + var shards []string + for _, shard := range resp.StreamDescription.Shards { + if shard.ShardId != nil { + shards = append(shards, aws.StringValue(shard.ShardId)) + } + } + return len(shards), nil +} diff --git a/producer/kinesiswriterconfig.go b/producer/kinesiswriterconfig.go index 320e3fe..8b639a6 100644 --- a/producer/kinesiswriterconfig.go +++ b/producer/kinesiswriterconfig.go @@ -6,6 +6,11 @@ import ( "github.com/aws/aws-sdk-go/aws" ) +const ( + kinesisMsgCountRateLimit = 1000 // AWS Kinesis limit of 1000 records/sec + kinesisMsgSizeRateLimit = 1000000 // AWS Kinesis limit of 1 MB/sec +) + // KinesisWriterConfig is used to configure KinesisWriter type KinesisWriterConfig struct { *kinesisWriterOptions @@ -19,6 +24,8 @@ func NewKinesisWriterConfig(cfg *aws.Config) *KinesisWriterConfig { AwsConfig: cfg, kinesisWriterOptions: &kinesisWriterOptions{ responseReadTimeout: time.Second, + msgCountRateLimit: kinesisMsgCountRateLimit, + msgSizeRateLimit: kinesisMsgSizeRateLimit, Stats: &NilStatsCollector{}, }, LogLevel: *cfg.LogLevel, @@ -30,6 +37,22 @@ func (c *KinesisWriterConfig) SetResponseReadTimeout(timeout time.Duration) { c.responseReadTimeout = timeout } +// SetMsgCountRateLimit configures the maximum number of messages that can be sent per second +func (c *KinesisWriterConfig) SetMsgCountRateLimit(limit int) { + if limit > kinesisMsgCountRateLimit { + + } + c.msgCountRateLimit = limit +} + +// SetMsgSizeRateLimit configures the maximum transmission size of the messages that can be sent per second +func (c *KinesisWriterConfig) SetMsgSizeRateLimit(limit int) { + if limit > kinesisMsgSizeRateLimit { + + } + c.msgSizeRateLimit = limit +} + // SetStatsCollector configures a listener to handle listener metrics. func (c *KinesisWriterConfig) SetStatsCollector(stats StatsCollector) { c.Stats = stats diff --git a/producer/producer.go b/producer/producer.go index e584a40..3aa4d0b 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -9,9 +9,9 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/awserr" "github.com/golang/time/rate" + "github.com/rewardStyle/kinetic/errs" "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" - "github.com/rewardStyle/kinetic/errs" ) // producerOptions holds all of the configurable settings for a Producer @@ -20,30 +20,29 @@ type producerOptions struct { batchTimeout time.Duration // maximum time duration to wait for incoming messages queueDepth int // maximum number of messages to enqueue in the message queue maxRetryAttempts int // maximum number of retry attempts for failed messages - workersPerShard int // number of concurrent workers per shard - shardCount int // initial shard size - msgCountLimit int // maximum records to be sent per cycle for the rate limiting model - msgSizeLimit int // maximum records to be sent per cycle for the rate limiting model + concurrency int // number of concurrent workers per shard + shardCheckFreq time.Duration // frequency (specified as a duration) with which to check the the shard size Stats StatsCollector // stats collection mechanism } // Producer sends records to AWS Kinesis or Firehose. type Producer struct { - *producerOptions // contains all of the configuration settings for the Producer - *logging.LogHelper // object for help with logging - writer StreamWriter // interface for abstracting the PutRecords call - msgCountLimiter *rate.Limiter // - msgSizeLimiter *rate.Limiter // - workerCount int // number of concurrent workers sending batch messages for the producer - messages chan *message.Message // channel for enqueuing messages to be put on the stream - status chan *statusReport // channel for workers to communicate their current status - dismiss chan empty // channel for handling the decommissioning of a surplus of workers - stop chan empty // channel for handling shutdown - pipeOfDeath chan empty // channel for handling pipe of death - startupOnce sync.Once // used to ensure that the startup function is called once - shutdownOnce sync.Once // used to ensure that the shutdown function is called once - resizeMu sync.Mutex // used to prevent resizeWorkerPool from being called synchronously with itself - noCopy noCopy // prevents the Producer from being copied + *producerOptions // contains all of the configuration settings for the Producer + *logging.LogHelper // object for help with logging + writer StreamWriter // interface for abstracting the PutRecords call + msgCountLimiter *rate.Limiter // rate limiter to limit the number of messages dispatched per second + msgSizeLimiter *rate.Limiter // rate limiter to limit the total size (in bytes) of messages dispatched per second + workerCount int // number of concurrent workers sending batch messages for the producer + messages chan *message.Message // channel for enqueuing messages to be put on the stream + status chan *statusReport // channel for workers to communicate their current status + dismiss chan empty // channel for handling the decommissioning of a surplus of workers + stop chan empty // channel for handling shutdown + pipeOfDeath chan empty // channel for handling pipe of death + startupOnce sync.Once // used to ensure that the startup function is called once + shutdownOnce sync.Once // used to ensure that the shutdown function is called once + resizeMu sync.Mutex // used to prevent resizeWorkerPool from being called synchronously with itself + shardCountMu sync.Mutex // used to add thread-safe access to the shardCount variable + noCopy noCopy // prevents the Producer from being copied } // NewProducer creates a new producer for writing records to a Kinesis or Firehose stream. @@ -56,7 +55,7 @@ func NewProducer(c *aws.Config, w StreamWriter, fn ...func(*Config)) (*Producer, producerOptions: cfg.producerOptions, LogHelper: &logging.LogHelper{ LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, + Logger: cfg.AwsConfig.Logger, }, writer: w, }, nil @@ -69,8 +68,8 @@ func (p *Producer) produce() { p.shutdownOnce = sync.Once{} // Instantiate rate limiters - p.msgCountLimiter = rate.NewLimiter(rate.Limit(float64(p.msgCountLimit)), p.batchSize) - p.msgSizeLimiter = rate.NewLimiter(rate.Limit(float64(p.msgSizeLimit)), p.msgSizeLimit) + p.msgCountLimiter = rate.NewLimiter(rate.Limit(float64(p.writer.getMsgCountRateLimit())), p.batchSize) + p.msgSizeLimiter = rate.NewLimiter(rate.Limit(float64(p.writer.getMsgSizeRateLimit())), p.writer.getMsgSizeRateLimit()) // Create communication channels p.messages = make(chan *message.Message, p.queueDepth) @@ -79,11 +78,41 @@ func (p *Producer) produce() { p.stop = make(chan empty) p.pipeOfDeath = make(chan empty) - // Instantiate and register new workers - p.resizeWorkerPool(p.shardCount * p.workersPerShard) + // Run a separate go routine to check the shard size (throughput multiplier) and resize the worker pool + // periodically if needed + stopShardCheck := make(chan empty) + go func() { + var mult int + + timer := time.NewTicker(p.shardCheckFreq) + for { + newMult, err := p.writer.getConcurrencyMultiplier() + if err != nil { + p.LogError("Failed to call getConcurrencyMultiplier due to: ", err) + } + if newMult != mult && newMult > 0 { + p.resizeWorkerPool(newMult * p.concurrency) + mult = newMult + } + + select { + case <-stopShardCheck: + timer.Stop() + return + case <-timer.C: + break + } + } + }() + + // Dispatch messages to each worker depending on the worker's capacity and the rate limit + go func() { + defer func() { + // Stop the periodic shard size check + stopShardCheck <- empty{} + close(stopShardCheck) + }() - go func(){ - // Dispatch messages to each worker depending on the worker's capacity and the rate limit for { select { case <-p.pipeOfDeath: @@ -103,7 +132,7 @@ func (p *Producer) produce() { var batchMsgSize int timeout := time.After(p.batchTimeout) - fillBatch: + fillBatch: for len(batch) < status.capacity { select { case <-timeout: @@ -115,7 +144,7 @@ func (p *Producer) produce() { } var msgSize int - if msgSize = msg.RequestEntrySize(); msgSize > p.msgSizeLimit { + if msgSize = msg.RequestEntrySize(); msgSize > p.writer.getMsgSizeRateLimit() { // TODO: Send this huge message to data spill break } @@ -127,7 +156,8 @@ func (p *Producer) produce() { // Request and wait for the message size (transmission) rate limiter to // allow this payload ctx, cancel := context.WithTimeout(context.TODO(), p.batchTimeout) - if err := p.msgSizeLimiter.WaitN(ctx, batchMsgSize + status.failedSize); err != nil { + defer cancel() + if err := p.msgSizeLimiter.WaitN(ctx, batchMsgSize+status.failedSize); err != nil { } cancel() @@ -140,7 +170,7 @@ func (p *Producer) produce() { // Request and wait for the message counter rate limiter to allow this // payload ctx, cancel := context.WithTimeout(context.TODO(), p.batchTimeout) - err := p.msgCountLimiter.WaitN(ctx, status.capacity + status.failedCount); + err := p.msgCountLimiter.WaitN(ctx, status.capacity+status.failedCount) if err != nil { // TODO: handle error properly } @@ -151,7 +181,7 @@ func (p *Producer) produce() { wg.Wait() // Send batch regardless if it is empty or not - status.channel <-batch + status.channel <- batch } } }() @@ -171,7 +201,7 @@ func (p *Producer) shutdown() { timer := time.NewTimer(staleTimeout) remaining := len(p.messages) - drain: + drain: for remaining > 0 { select { case <-time.After(time.Second): @@ -262,17 +292,17 @@ func (p *Producer) newWorker() { for _, msg := range retries { failedSize += msg.RequestEntrySize() } - status := &statusReport { - capacity: capacity, + status := &statusReport{ + capacity: capacity, failedCount: failedCount, - failedSize: failedSize, - channel: batches, + failedSize: failedSize, + channel: batches, } - p.status <-status + p.status <- status // Receive a batch of messages and call the producer's sendBatch function batch := <-batches - if len(batch) + len(retries) > 0 { + if len(batch)+len(retries) > 0 { retries = p.sendBatch(append(retries, batch...)) } @@ -290,6 +320,7 @@ func (p *Producer) sendBatch(batch []*message.Message) []*message.Message { } else { p.Stats.AddDroppedTotal(1) p.Stats.AddDroppedRetries(1) + // TODO: Add to data spill } return nil }) diff --git a/producer/producer_test.go b/producer/producer_test.go index f2c2e14..1230680 100644 --- a/producer/producer_test.go +++ b/producer/producer_test.go @@ -4,6 +4,7 @@ import ( . "github.com/smartystreets/goconvey/convey" "context" + "sync" "testing" "time" @@ -44,9 +45,9 @@ func TestProducer(t *testing.T) { p, err := NewProducer(k.Session.Config, w, func(c *Config) { c.SetBatchSize(5) - c.SetBatchTimeout(1000 * time.Millisecond) - //c.SetConcurrency(10) + c.SetBatchTimeout(time.Second) c.SetQueueDepth(10) + c.SetConcurrency(2) }) So(p, ShouldNotBeNil) So(err, ShouldBeNil) @@ -91,7 +92,6 @@ func TestProducer(t *testing.T) { So(err, ShouldBeNil) So(string(msg.Data), ShouldEqual, data) So(elapsed.Seconds(), ShouldBeGreaterThan, 1) - }) Convey("check that we can send a single message after batch timeout elapses", func() { @@ -109,23 +109,36 @@ func TestProducer(t *testing.T) { So(elapsed.Seconds(), ShouldBeGreaterThan, 1) }) - Convey("check that we can send a batch of messages after batch size is reached", func() { + Convey("check that we can send a batch of messages after batch size is reached", func(c C) { start := time.Now() + var elapsed time.Duration data := []string{"hello1", "hello2", "hello3", "hello4", "hello5", "hello6"} - for _, datum := range data { - p.Send(&message.Message{ - PartitionKey: aws.String("key"), - Data: []byte(datum), - }) - } - - for i := 0; i < 5; i++ { - msg, err := l.Retrieve() - So(err, ShouldBeNil) - So(string(msg.Data), ShouldEqual, data[i]) - } - elapsed := time.Since(start) - So(elapsed.Seconds(), ShouldBeLessThan, 1) + + wg := sync.WaitGroup{} + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 5; i++ { + msg, err := l.Retrieve() + c.So(err, ShouldBeNil) + c.So(string(msg.Data), ShouldEqual, data[i]) + } + elapsed = time.Since(start) + }() + + wg.Add(1) + go func() { + defer wg.Done() + for _, datum := range data { + p.Send(&message.Message{ + PartitionKey: aws.String("key"), + Data: []byte(datum), + }) + } + }() + wg.Wait() + + //So(elapsed.Seconds(), ShouldBeLessThan, 1) Printf("(first 5 took %f seconds)\n", elapsed.Seconds()) msg, err := l.Retrieve() diff --git a/producer/producerconfig.go b/producer/producerconfig.go index 99aa3e9..ac54e44 100644 --- a/producer/producerconfig.go +++ b/producer/producerconfig.go @@ -24,10 +24,8 @@ func NewConfig(cfg *aws.Config) *Config { batchTimeout: time.Second, queueDepth: 10000, maxRetryAttempts: 10, - workersPerShard: 5, - shardCount: 1, - msgCountLimit: 1000, - msgSizeLimit: 1000000, + concurrency: 3, + shardCheckFreq: time.Minute, Stats: &NilStatsCollector{}, }, LogLevel: *cfg.LogLevel, @@ -60,24 +58,14 @@ func (c *Config) SetMaxRetryAttempts(attempts int) { c.maxRetryAttempts = attempts } -// SetWorkersPerShard defines the number of concurrent workers to run per active shard -func (c *Config) SetWorkersPerShard(count int) { - c.workersPerShard = count +// SetConcurrency defines the number of concurrent workers to run per active shard (multiplier) +func (c *Config) SetConcurrency(count int) { + c.concurrency = count } -// SetShardCount defines the initial shard size -func (c *Config) SetShardCount(count int) { - c.shardCount = count -} - -// SetMsgCountLimit defines the maximum number of message that can be sent per second -func (c *Config) SetMsgCountLimit(limit int) { - c.msgCountLimit = limit -} - -// SetMsgSizeLimit defines the maximum size (in bytes) that can be sent per second -func (c *Config) SetMsgSizeLimit(limit int) { - c.msgSizeLimit = limit +// SetShardCheckFreq defines a frequency (specified as a duration) with which to check the shard size +func (c *Config) SetShardCheckFreq(duration time.Duration) { + c.shardCheckFreq = duration } // SetStatsCollector configures a listener to handle producer metrics. diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go index 15dc251..54dc5fa 100644 --- a/producer/producerconfig_test.go +++ b/producer/producerconfig_test.go @@ -48,10 +48,8 @@ func TestNewConfig(t *testing.T) { So(cfg.batchTimeout, ShouldEqual, 1*time.Second) So(cfg.queueDepth, ShouldEqual, 10000) So(cfg.maxRetryAttempts, ShouldEqual, 10) - So(cfg.workersPerShard, ShouldEqual, 5) - So(cfg.shardCount, ShouldEqual, 1) - So(cfg.msgCountLimit, ShouldEqual, 1000) - So(cfg.msgSizeLimit, ShouldEqual, 1000000) + So(cfg.concurrency, ShouldEqual, 3) + So(cfg.shardCheckFreq, ShouldEqual, time.Minute) So(cfg.Stats, ShouldHaveSameTypeAs, &NilStatsCollector{}) So(cfg.LogLevel.Value(), ShouldEqual, logging.LogOff) }) @@ -83,18 +81,13 @@ func TestNewConfig(t *testing.T) { }) Convey("check that we can set the workers per shard", func() { - cfg.SetWorkersPerShard(10) - So(cfg.workersPerShard, ShouldEqual, 10) + cfg.SetConcurrency(10) + So(cfg.concurrency, ShouldEqual, 10) }) - Convey("check that we can set the message count rate limit", func() { - cfg.SetMsgCountLimit(100) - So(cfg.msgCountLimit, ShouldEqual, 100) - }) - - Convey("check that we can set the message size rate limit", func() { - cfg.SetMsgSizeLimit(10) - So(cfg.msgSizeLimit, ShouldEqual, 10) + Convey("check that we can set the check shard frequency", func() { + cfg.SetShardCheckFreq(time.Second) + So(cfg.shardCheckFreq, ShouldEqual, time.Second) }) Convey("check that we can configure a stats collector", func() { diff --git a/producer/sharedtypes.go b/producer/sharedtypes.go index ad99690..7704a34 100644 --- a/producer/sharedtypes.go +++ b/producer/sharedtypes.go @@ -10,6 +10,9 @@ import ( // StreamWriter is an interface that abstracts the differences in API between Kinesis and Firehose. type StreamWriter interface { PutRecords(context.Context, []*message.Message, MessageHandlerAsync) error + getMsgCountRateLimit() int + getMsgSizeRateLimit() int + getConcurrencyMultiplier() (int, error) } // MessageHandler defines the signature of a message handler used by PutRecords(). MessageHandler accepts a WaitGroup @@ -24,7 +27,7 @@ type MessageHandlerAsync func(*message.Message) error type statusReport struct { capacity int // maximum message capacity the worker can handle failedCount int // number of previous messages that failed to send - failedSize int // size in bytes of the previous messages that failed to send + failedSize int // size in bytes of the previous messages that failed to send channel chan []*message.Message // channel of the worker to which the batch messages should be sent } diff --git a/producer/stats.go b/producer/stats.go index 5a42453..6c7c7dd 100644 --- a/producer/stats.go +++ b/producer/stats.go @@ -84,7 +84,7 @@ const ( MetricsSentTotal = "kinetic.producer.sent.total" MetricsSentSuccess = "kinetic.producer.sent.success" MetricsSentFailed = "kinetic.producer.sent.failed" - MetricsSentRetried = "kinetic.producer.sent.retried" + MetricsSentRetried = "kinetic.producer.sent.retried" MetricsDroppedTotal = "kinetic.producer.dropped.total" MetricsDroppedCapacity = "kinetic.producer.dropped.capacity" MetricsDroppedRetries = "kinetic.producer.dropped.retries" diff --git a/testexec/cli.go b/testexec/cli.go index daaa2f5..d651565 100644 --- a/testexec/cli.go +++ b/testexec/cli.go @@ -26,8 +26,8 @@ func parseCommandLineArgs() *Config { locationPtr := flag.String("location", "local", "used to specify the location of the kinesis stream. "+ "Accepted values are (local|aws). For local, run kinesalite on http://127.0.0.1:4567. For aws, your "+ "aws credentials and configuration need to be defined at ~/.aws") - streamNamePtr := flag.String("stream", "", "used to specify a specific stream to write to / read from for " + - "testing. The stream will be created if a stream name is given but does not exist. A random stream will " + + streamNamePtr := flag.String("stream", "", "used to specify a specific stream to write to / read from for "+ + "testing. The stream will be created if a stream name is given but does not exist. A random stream will "+ "be created if the stream is undefined.") msgCountPtr := flag.Int("count", 0, "used to specify the number of messages to (attempt to) send. This "+ "flag is only applicable to 'write' and 'readwrite' modes. Use zero or a negative number to produce "+ @@ -74,7 +74,7 @@ func parseCommandLineArgs() *Config { Mode: &mode, StreamName: streamNamePtr, Duration: durationPtr, - MsgCount: msgCountPtr, + MsgCount: msgCountPtr, Location: locationPtr, Blast: blastPtr, Clean: cleanPtr, diff --git a/testexec/main.go b/testexec/main.go index 299e603..5b38064 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -242,6 +242,8 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce //kwc.SetLogLevel(logging.LogDebug) kwc.SetResponseReadTimeout(time.Second) kwc.SetStatsCollector(psc) + kwc.SetMsgCountRateLimit(1000) + kwc.SetMsgSizeRateLimit(1000000) }) if err != nil { log.Fatalf("Unable to create a new kinesis stream writer due to: %v\n", err) @@ -253,9 +255,8 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce c.SetMaxRetryAttempts(3) c.SetStatsCollector(psc) c.SetQueueDepth(10000) - c.SetWorkersPerShard(3) - c.SetMsgCountLimit(1000) - c.SetMsgSizeLimit(1000000) + c.SetConcurrency(3) + c.SetShardCheckFreq(time.Minute) }) if err != nil { log.Fatalf("Unable to create a new producer due to: %v\n", err) From e26a55c03b55e90158901bb8aa519dad0e871fa0 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 24 Jul 2017 16:38:28 -0500 Subject: [PATCH 060/109] Added new stats to track the number of concurrent workers for the producer. Fixed a bug with the determination of the number of open shards for a stream. --- producer/firehosewriter.go | 6 ++--- producer/kinesiswriter.go | 42 +++++++++++++++++++++++++++----- producer/producer.go | 1 + producer/producerconfig_test.go | 7 +++--- producer/stats.go | 43 +++++++++++++++++++++------------ producer/stats_test.go | 33 ++++++++++++++++--------- 6 files changed, 93 insertions(+), 39 deletions(-) diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go index aab0fa6..3885900 100644 --- a/producer/firehosewriter.go +++ b/producer/firehosewriter.go @@ -75,7 +75,7 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Mes }) req.Handlers.Build.PushBack(func(r *request.Request) { - w.Stats.AddPutRecordsBuildDuration(time.Since(startBuildTime)) + w.Stats.UpdatePutRecordsBuildDuration(time.Since(startBuildTime)) w.LogDebug("Finished PutRecords Build, took", time.Since(start)) }) @@ -85,7 +85,7 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Mes }) req.Handlers.Send.PushBack(func(r *request.Request) { - w.Stats.AddPutRecordsSendDuration(time.Since(startSendTime)) + w.Stats.UpdatePutRecordsSendDuration(time.Since(startSendTime)) w.LogDebug("Finished PutRecords Send, took", time.Since(start)) }) @@ -95,7 +95,7 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Mes w.LogError("Error putting records:", err.Error()) return err } - w.Stats.AddPutRecordsDuration(time.Since(start)) + w.Stats.UpdatePutRecordsDuration(time.Since(start)) if resp == nil { return errs.ErrNilPutRecordsResponse diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index 1ff2ee7..11fa37a 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -77,7 +77,7 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess }) req.Handlers.Build.PushBack(func(r *request.Request) { - w.Stats.AddPutRecordsBuildDuration(time.Since(startBuildTime)) + w.Stats.UpdatePutRecordsBuildDuration(time.Since(startBuildTime)) w.LogDebug("Finished PutRecords Build, took", time.Since(start)) }) @@ -87,7 +87,7 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess }) req.Handlers.Send.PushBack(func(r *request.Request) { - w.Stats.AddPutRecordsSendDuration(time.Since(startSendTime)) + w.Stats.UpdatePutRecordsSendDuration(time.Since(startSendTime)) w.LogDebug("Finished PutRecords Send, took", time.Since(start)) }) @@ -97,7 +97,7 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess w.LogError("Error putting records:", err.Error()) return err } - w.Stats.AddPutRecordsDuration(time.Since(start)) + w.Stats.UpdatePutRecordsDuration(time.Since(start)) if resp == nil { return errs.ErrNilPutRecordsResponse @@ -161,11 +161,41 @@ func (w *KinesisWriter) getConcurrencyMultiplier() (int, error) { if resp.StreamDescription == nil { return 0, errs.ErrNilStreamDescription } - var shards []string + + // maps shardID to a boolean that indicates whether or not the shard is a parent shard or an adjacent parent shard + shardMap := make(map[string]bool) for _, shard := range resp.StreamDescription.Shards { if shard.ShardId != nil { - shards = append(shards, aws.StringValue(shard.ShardId)) + shardID := aws.StringValue(shard.ShardId) + if _, ok := shardMap[shardID]; !ok { + shardMap[shardID] = false + } } } - return len(shards), nil + + // Loop through all the shards and mark which ones are parents + for _, shard := range resp.StreamDescription.Shards { + if shard.ParentShardId != nil { + shardID := aws.StringValue(shard.ParentShardId) + if _, ok := shardMap[shardID]; ok { + shardMap[shardID] = true + } + } + if shard.AdjacentParentShardId != nil { + shardID := aws.StringValue(shard.AdjacentParentShardId) + if _, ok := shardMap[shardID]; ok { + shardMap[shardID] = true + } + } + } + + // Determine the number of open shards by removing those shards that are reported as parents + openShardCount := len(shardMap) + for _, isParent := range shardMap { + if isParent { + openShardCount-- + } + } + + return openShardCount, nil } diff --git a/producer/producer.go b/producer/producer.go index 3aa4d0b..594107d 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -92,6 +92,7 @@ func (p *Producer) produce() { } if newMult != mult && newMult > 0 { p.resizeWorkerPool(newMult * p.concurrency) + p.Stats.UpdateProducerConcurrency(newMult * p.concurrency) mult = newMult } diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go index 54dc5fa..702bd7d 100644 --- a/producer/producerconfig_test.go +++ b/producer/producerconfig_test.go @@ -25,9 +25,10 @@ func (l *DebugStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} func (l *DebugStatsCollector) AddPutRecordsCalled(int) {} func (l *DebugStatsCollector) AddProvisionedThroughputExceeded(int) {} func (l *DebugStatsCollector) AddPutRecordsTimeout(int) {} -func (l *DebugStatsCollector) AddPutRecordsDuration(time.Duration) {} -func (l *DebugStatsCollector) AddPutRecordsBuildDuration(time.Duration) {} -func (l *DebugStatsCollector) AddPutRecordsSendDuration(time.Duration) {} +func (l *DebugStatsCollector) UpdatePutRecordsDuration(time.Duration) {} +func (l *DebugStatsCollector) UpdatePutRecordsBuildDuration(time.Duration) {} +func (l *DebugStatsCollector) UpdatePutRecordsSendDuration(time.Duration) {} +func (l *DebugStatsCollector) UpdateProducerConcurrency(int) {} type DebugStreamWriter struct{} diff --git a/producer/stats.go b/producer/stats.go index 6c7c7dd..88d3988 100644 --- a/producer/stats.go +++ b/producer/stats.go @@ -22,9 +22,10 @@ type StatsCollector interface { AddPutRecordsTimeout(int) AddProvisionedThroughputExceeded(int) AddPutRecordsProvisionedThroughputExceeded(int) - AddPutRecordsDuration(time.Duration) - AddPutRecordsBuildDuration(time.Duration) - AddPutRecordsSendDuration(time.Duration) + UpdatePutRecordsDuration(time.Duration) + UpdatePutRecordsBuildDuration(time.Duration) + UpdatePutRecordsSendDuration(time.Duration) + UpdateProducerConcurrency(int) } // NilStatsCollector is a stats listener that ignores all metrics. @@ -69,15 +70,18 @@ func (nsc *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} // ErrCodeProvisionedThroughputExceededException by the producer. func (nsc *NilStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} -// AddPutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls +// UpdatePutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls // are measured. -func (nsc *NilStatsCollector) AddPutRecordsDuration(time.Duration) {} +func (nsc *NilStatsCollector) UpdatePutRecordsDuration(time.Duration) {} -// AddPutRecordsBuildDuration records the duration that it took to build the PutRecords API request payload. -func (nsc *NilStatsCollector) AddPutRecordsBuildDuration(time.Duration) {} +// UpdatePutRecordsBuildDuration records the duration that it took to build the PutRecords API request payload. +func (nsc *NilStatsCollector) UpdatePutRecordsBuildDuration(time.Duration) {} -// AddPutRecordsSendDuration records the duration that it took to send the PutRecords API request payload. -func (nsc *NilStatsCollector) AddPutRecordsSendDuration(time.Duration) {} +// UpdatePutRecordsSendDuration records the duration that it took to send the PutRecords API request payload. +func (nsc *NilStatsCollector) UpdatePutRecordsSendDuration(time.Duration) {} + +// UpdateProducerConcurrency records the number of concurrent workers that the producer has. +func (nsc *NilStatsCollector) UpdateProducerConcurrency(int) {} // Metric names to be exported const ( @@ -95,6 +99,7 @@ const ( MetricsPutRecordsDuration = "kinetic.producer.putrecords.duration" MetricsPutRecordsBuildDuration = "kinetic.producer.putrecords.build.duration" MetricsPutRecordsSendDuration = "kinetic.producer.putrecords.send.duration" + MetricsProducerConcurrency = "kinetic.producer.concurrency" ) // DefaultStatsCollector is a type that implements the producers's StatsCollector interface using the @@ -114,6 +119,7 @@ type DefaultStatsCollector struct { PutRecordsDuration metrics.Gauge PutRecordsBuildDuration metrics.Gauge PutRecordsSendDuration metrics.Gauge + ProducerConcurrency metrics.Gauge } // NewDefaultStatsCollector instantiates a new DefaultStatsCollector object @@ -133,6 +139,7 @@ func NewDefaultStatsCollector(r metrics.Registry) *DefaultStatsCollector { PutRecordsDuration: metrics.GetOrRegisterGauge(MetricsPutRecordsDuration, r), PutRecordsBuildDuration: metrics.GetOrRegisterGauge(MetricsPutRecordsBuildDuration, r), PutRecordsSendDuration: metrics.GetOrRegisterGauge(MetricsPutRecordsSendDuration, r), + ProducerConcurrency: metrics.GetOrRegisterGauge(MetricsProducerConcurrency, r), } } @@ -197,22 +204,27 @@ func (dsc *DefaultStatsCollector) AddPutRecordsProvisionedThroughputExceeded(cou dsc.PutRecordsProvisionedThroughputExceeded.Inc(int64(count)) } -// AddPutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls +// UpdatePutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls // are measured. -func (dsc *DefaultStatsCollector) AddPutRecordsDuration(duration time.Duration) { +func (dsc *DefaultStatsCollector) UpdatePutRecordsDuration(duration time.Duration) { dsc.PutRecordsDuration.Update(duration.Nanoseconds()) } -// AddPutRecordsBuildDuration records the duration that it took to build the PutRecords API request payload. -func (dsc *DefaultStatsCollector) AddPutRecordsBuildDuration(duration time.Duration) { +// UpdatePutRecordsBuildDuration records the duration that it took to build the PutRecords API request payload. +func (dsc *DefaultStatsCollector) UpdatePutRecordsBuildDuration(duration time.Duration) { dsc.PutRecordsBuildDuration.Update(duration.Nanoseconds()) } -// AddPutRecordsSendDuration records the duration that it took to send the PutRecords API request payload. -func (dsc *DefaultStatsCollector) AddPutRecordsSendDuration(duration time.Duration) { +// UpdatePutRecordsSendDuration records the duration that it took to send the PutRecords API request payload. +func (dsc *DefaultStatsCollector) UpdatePutRecordsSendDuration(duration time.Duration) { dsc.PutRecordsSendDuration.Update(duration.Nanoseconds()) } +// UpdateProducerConcurrency records the number of concurrent workers that the producer has. +func (dsc *DefaultStatsCollector) UpdateProducerConcurrency(count int) { + dsc.ProducerConcurrency.Update(int64(count)) +} + // PrintStats logs the stats func (dsc *DefaultStatsCollector) PrintStats() { log.Printf("Producer Stats: Sent Total: [%d]\n", dsc.SentTotal.Count()) @@ -229,4 +241,5 @@ func (dsc *DefaultStatsCollector) PrintStats() { log.Printf("Producer Stats: PutRecords Duration (ns): [%d]\n", dsc.PutRecordsDuration.Value()) log.Printf("Producer Stats: PutRecords Build Duration (ns): [%d]\n", dsc.PutRecordsBuildDuration.Value()) log.Printf("Producer Stats: PutRecords Send Duration (ns): [%d]\n", dsc.PutRecordsSendDuration.Value()) + log.Printf("Producer Stats: Producer Concurrency: [%d]\n", dsc.ProducerConcurrency.Value()) } diff --git a/producer/stats_test.go b/producer/stats_test.go index 3d6f320..f333d13 100644 --- a/producer/stats_test.go +++ b/producer/stats_test.go @@ -62,16 +62,20 @@ func TestStatsCollector(t *testing.T) { sc.AddPutRecordsTimeout(1) }) - Convey("check that AddPutRecordsDuration does not error", func() { - sc.AddPutRecordsDuration(time.Second) + Convey("check that UpdatePutRecordsDuration does not error", func() { + sc.UpdatePutRecordsDuration(time.Second) }) - Convey("check that AddPutRecordsBuildDuration does not error", func() { - sc.AddPutRecordsBuildDuration(time.Second) + Convey("check that UpdatePutRecordsBuildDuration does not error", func() { + sc.UpdatePutRecordsBuildDuration(time.Second) }) - Convey("check that AddPutRecordsSendDuration does not error", func() { - sc.AddPutRecordsSendDuration(time.Second) + Convey("check that UpdatePutRecordsSendDuration does not error", func() { + sc.UpdatePutRecordsSendDuration(time.Second) + }) + + Convey("check that UpdateProducerConcurrency does not error", func() { + sc.UpdateProducerConcurrency(5) }) }) @@ -146,19 +150,24 @@ func TestStatsCollector(t *testing.T) { So(sc.(*DefaultStatsCollector).PutRecordsTimeout.Count(), ShouldEqual, int64(count)) }) - Convey("check that AddPutRecordsDuration does not error", func() { - sc.AddPutRecordsDuration(time.Second) + Convey("check that UpdatePutRecordsDuration does not error", func() { + sc.UpdatePutRecordsDuration(time.Second) So(sc.(*DefaultStatsCollector).PutRecordsDuration.Value(), ShouldEqual, 1000000000) }) - Convey("check that AddPutRecordsBuildDuration does not error", func() { - sc.AddPutRecordsBuildDuration(time.Second) + Convey("check that UpdatePutRecordsBuildDuration does not error", func() { + sc.UpdatePutRecordsBuildDuration(time.Second) So(sc.(*DefaultStatsCollector).PutRecordsBuildDuration.Value(), ShouldEqual, 1000000000) }) - Convey("check that AddPutRecordsSendDuration does not error", func() { - sc.AddPutRecordsSendDuration(time.Second) + Convey("check that UpdatePutRecordsSendDuration does not error", func() { + sc.UpdatePutRecordsSendDuration(time.Second) So(sc.(*DefaultStatsCollector).PutRecordsSendDuration.Value(), ShouldEqual, 1000000000) }) + + Convey("check that UpdateProducerConcurrency does not error", func() { + sc.UpdateProducerConcurrency(5) + So(sc.(*DefaultStatsCollector).ProducerConcurrency.Value(), ShouldEqual, 5) + }) }) } From 13421ee4bd380bf4bef4fb34c362af0a3a53f710 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 24 Jul 2017 17:09:24 -0500 Subject: [PATCH 061/109] Added more thorough error checking around rate limiting configuration values --- errs/errors.go | 6 ++++++ producer/kinesiswriterconfig.go | 18 ++++++++++++------ 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/errs/errors.go b/errs/errors.go index bf9954d..7cfb441 100644 --- a/errs/errors.go +++ b/errs/errors.go @@ -86,4 +86,10 @@ var ( // ErrInvalidBatchSize is returned when the batchSize is invalid ErrInvalidBatchSize = errors.New("PutRecordsBatch supports batch sizes less than or equal to 500") + + // ErrInvalidMsgCountRateLimit is returned when a MsgCountRateLimit is configured incorrectly + ErrInvalidMsgCountRateLimit = errors.New("Invalid Message Count Rate Limit") + + // ErrInvalidMsgSizeRateLimit is returned when a MsgCountSizesLimit is configured incorrectly + ErrInvalidMsgSizeRateLimit = errors.New("Invalid Message Size Rate Limit") ) diff --git a/producer/kinesiswriterconfig.go b/producer/kinesiswriterconfig.go index 8b639a6..d991118 100644 --- a/producer/kinesiswriterconfig.go +++ b/producer/kinesiswriterconfig.go @@ -1,9 +1,11 @@ package producer import ( + "log" "time" "github.com/aws/aws-sdk-go/aws" + "github.com/rewardStyle/kinetic/errs" ) const ( @@ -39,18 +41,22 @@ func (c *KinesisWriterConfig) SetResponseReadTimeout(timeout time.Duration) { // SetMsgCountRateLimit configures the maximum number of messages that can be sent per second func (c *KinesisWriterConfig) SetMsgCountRateLimit(limit int) { - if limit > kinesisMsgCountRateLimit { - + if limit > 0 && limit <= kinesisMsgCountRateLimit { + c.msgCountRateLimit = limit + } else { + log.Fatal("Message Count Rate Limit must be positive and less than ", kinesisMsgCountRateLimit) + panic(errs.ErrInvalidMsgCountRateLimit) } - c.msgCountRateLimit = limit } // SetMsgSizeRateLimit configures the maximum transmission size of the messages that can be sent per second func (c *KinesisWriterConfig) SetMsgSizeRateLimit(limit int) { - if limit > kinesisMsgSizeRateLimit { - + if limit > 0 && limit <= kinesisMsgSizeRateLimit { + c.msgSizeRateLimit = limit + } else { + log.Fatal("Message Count Size Limit must be positive and less than ", kinesisMsgSizeRateLimit) + panic(errs.ErrInvalidMsgSizeRateLimit) } - c.msgSizeRateLimit = limit } // SetStatsCollector configures a listener to handle listener metrics. From c66a9e10635f9032dd45c6956bb6320ef6513d38 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 24 Jul 2017 21:23:03 -0500 Subject: [PATCH 062/109] Added a data spill callback function as a configurable parameter. --- producer/kinesiswriterconfig.go | 2 +- producer/producer.go | 37 ++++++++++++++++++++------------- producer/producerconfig.go | 7 +++++++ producer/producerconfig_test.go | 8 +++++++ testexec/main.go | 4 ++++ 5 files changed, 43 insertions(+), 15 deletions(-) diff --git a/producer/kinesiswriterconfig.go b/producer/kinesiswriterconfig.go index d991118..69a1b6e 100644 --- a/producer/kinesiswriterconfig.go +++ b/producer/kinesiswriterconfig.go @@ -41,7 +41,7 @@ func (c *KinesisWriterConfig) SetResponseReadTimeout(timeout time.Duration) { // SetMsgCountRateLimit configures the maximum number of messages that can be sent per second func (c *KinesisWriterConfig) SetMsgCountRateLimit(limit int) { - if limit > 0 && limit <= kinesisMsgCountRateLimit { + if limit > 0 && limit <= kinesisMsgCountRateLimit { c.msgCountRateLimit = limit } else { log.Fatal("Message Count Rate Limit must be positive and less than ", kinesisMsgCountRateLimit) diff --git a/producer/producer.go b/producer/producer.go index 594107d..fa74d7e 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -16,13 +16,14 @@ import ( // producerOptions holds all of the configurable settings for a Producer type producerOptions struct { - batchSize int // maximum message capacity per request - batchTimeout time.Duration // maximum time duration to wait for incoming messages - queueDepth int // maximum number of messages to enqueue in the message queue - maxRetryAttempts int // maximum number of retry attempts for failed messages - concurrency int // number of concurrent workers per shard - shardCheckFreq time.Duration // frequency (specified as a duration) with which to check the the shard size - Stats StatsCollector // stats collection mechanism + batchSize int // maximum message capacity per request + batchTimeout time.Duration // maximum time duration to wait for incoming messages + queueDepth int // maximum number of messages to enqueue in the message queue + maxRetryAttempts int // maximum number of retry attempts for failed messages + concurrency int // number of concurrent workers per shard + shardCheckFreq time.Duration // frequency (specified as a duration) with which to check the the shard size + dataSpillFn MessageHandlerAsync // callback function for handling dropped messages that the producer was unable to send to the stream + Stats StatsCollector // stats collection mechanism } // Producer sends records to AWS Kinesis or Firehose. @@ -146,7 +147,7 @@ func (p *Producer) produce() { var msgSize int if msgSize = msg.RequestEntrySize(); msgSize > p.writer.getMsgSizeRateLimit() { - // TODO: Send this huge message to data spill + p.sendToDataSpill(msg) break } batchMsgSize += msgSize @@ -214,7 +215,9 @@ func (p *Producer) shutdown() { case <-timer.C: timer.Stop() if remaining > 0 { - // TODO: Send remaining messages to the data spill + for msg := range p.messages { + p.sendToDataSpill(msg) + } } break drain } @@ -319,9 +322,7 @@ func (p *Producer) sendBatch(batch []*message.Message) []*message.Message { failed = append(failed, msg) p.Stats.AddSentRetried(1) } else { - p.Stats.AddDroppedTotal(1) - p.Stats.AddDroppedRetries(1) - // TODO: Add to data spill + p.sendToDataSpill(msg) } return nil }) @@ -355,6 +356,15 @@ func (p *Producer) sendBatch(batch []*message.Message) []*message.Message { return batch } +// sendToDataSpill is called when the producer is unable to write the message to the stream +func (p *Producer) sendToDataSpill(msg *message.Message) { + p.Stats.AddDroppedTotal(1) + p.Stats.AddDroppedCapacity(1) + if err := p.dataSpillFn(msg); err != nil { + p.LogError("Unable to call data spill function on message: ", string(msg.Data)) + } +} + // CloseWithContext initiates the graceful shutdown of the produce function, waiting for all outstanding messages and to // flush. Cancellation is supported through contexts. func (p *Producer) CloseWithContext(ctx context.Context) { @@ -392,8 +402,7 @@ func (p *Producer) TryToSend(msg *message.Message) error { case p.messages <- msg: return nil default: - p.Stats.AddDroppedTotal(1) - p.Stats.AddDroppedCapacity(1) + p.sendToDataSpill(msg) return errs.ErrDroppedMessage } } diff --git a/producer/producerconfig.go b/producer/producerconfig.go index ac54e44..4cab912 100644 --- a/producer/producerconfig.go +++ b/producer/producerconfig.go @@ -6,6 +6,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/rewardStyle/kinetic/errs" + "github.com/rewardStyle/kinetic/message" ) // Config is used to configure a Producer instance. @@ -26,6 +27,7 @@ func NewConfig(cfg *aws.Config) *Config { maxRetryAttempts: 10, concurrency: 3, shardCheckFreq: time.Minute, + dataSpillFn: func(*message.Message) error { return nil }, Stats: &NilStatsCollector{}, }, LogLevel: *cfg.LogLevel, @@ -68,6 +70,11 @@ func (c *Config) SetShardCheckFreq(duration time.Duration) { c.shardCheckFreq = duration } +// SetDataSpillFn defines a callback function to be called when a message gets dropped by the producer +func (c *Config) SetDataSpillFn(fn MessageHandlerAsync) { + c.dataSpillFn = fn +} + // SetStatsCollector configures a listener to handle producer metrics. func (c *Config) SetStatsCollector(stats StatsCollector) { c.Stats = stats diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go index 702bd7d..949a2d0 100644 --- a/producer/producerconfig_test.go +++ b/producer/producerconfig_test.go @@ -86,6 +86,14 @@ func TestNewConfig(t *testing.T) { So(cfg.concurrency, ShouldEqual, 10) }) + Convey("check that we can set the data spill callback function", func() { + fn := func(msg *message.Message) error { + return nil + } + cfg.SetDataSpillFn(fn) + So(cfg.dataSpillFn, ShouldEqual, fn) + }) + Convey("check that we can set the check shard frequency", func() { cfg.SetShardCheckFreq(time.Second) So(cfg.shardCheckFreq, ShouldEqual, time.Second) diff --git a/testexec/main.go b/testexec/main.go index 5b38064..e340bf9 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -257,6 +257,10 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce c.SetQueueDepth(10000) c.SetConcurrency(3) c.SetShardCheckFreq(time.Minute) + c.SetDataSpillFn(func(msg *message.Message) error { + //log.Printf("Message was dropped: [%s]\n", string(msg.Data)) + return nil + }) }) if err != nil { log.Fatalf("Unable to create a new producer due to: %v\n", err) From b3f5a5f5fd4dd19f3e179ee02e4a5a968c8242b5 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 25 Jul 2017 10:30:05 -0500 Subject: [PATCH 063/109] Reset rate limiters upon a detected shard count change event. --- producer/producer.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/producer/producer.go b/producer/producer.go index fa74d7e..f9b642c 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -92,9 +92,11 @@ func (p *Producer) produce() { p.LogError("Failed to call getConcurrencyMultiplier due to: ", err) } if newMult != mult && newMult > 0 { - p.resizeWorkerPool(newMult * p.concurrency) - p.Stats.UpdateProducerConcurrency(newMult * p.concurrency) mult = newMult + p.msgCountLimiter.SetLimit(rate.Limit(float64(mult * p.writer.getMsgCountRateLimit()))) + p.msgSizeLimiter.SetLimit(rate.Limit(float64(mult * p.writer.getMsgSizeRateLimit()))) + p.resizeWorkerPool(mult * p.concurrency) + p.Stats.UpdateProducerConcurrency(mult * p.concurrency) } select { From 9788d1ef4041c11019cadd4ab7df6b64af2f9be6 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 25 Jul 2017 15:57:14 -0500 Subject: [PATCH 064/109] Modified the rate limiting implementation in the produce function to fill a batch or timeout first before waiting on the count and size tokens needed to send the batch. Renamed newWorker helper function to doWork. Enhanced error handling around RequestEntrySize. Made throughput multiplier configurable for the firehosewriter. --- kinetic.go | 6 +-- message/message.go | 6 +-- producer/firehosewriter.go | 9 ++-- producer/firehosewriterconfig.go | 12 +++-- producer/producer.go | 90 +++++++++++++++++--------------- 5 files changed, 67 insertions(+), 56 deletions(-) diff --git a/kinetic.go b/kinetic.go index 957cdce..0885367 100644 --- a/kinetic.go +++ b/kinetic.go @@ -36,7 +36,7 @@ type Kinetic struct { func New(fn func(*Config)) (*Kinetic, error) { config := NewConfig() fn(config) - session, err := config.GetSession() + sess, err := config.GetSession() if err != nil { return nil, err } @@ -44,9 +44,9 @@ func New(fn func(*Config)) (*Kinetic, error) { kineticOptions: config.kineticOptions, LogHelper: &logging.LogHelper{ LogLevel: config.LogLevel, - Logger: session.Config.Logger, + Logger: sess.Config.Logger, }, - Session: session, + Session: sess, }, nil } diff --git a/message/message.go b/message/message.go index 578455a..1909956 100644 --- a/message/message.go +++ b/message/message.go @@ -42,12 +42,12 @@ func FromRecord(record *kinesis.Record) *Message { // RequestEntrySize calculates what the size (in bytes) of the message will be after calling ToRequestEntry on it and // marshalling it to json -func (m *Message) RequestEntrySize() int { +func (m *Message) RequestEntrySize() (int, error) { buf, err := json.Marshal(m.ToRequestEntry()) if err != nil { - + return 0, nil } - return len(buf) + return len(buf), nil } // ToRequestEntry creates a kinesis.PutRecordsRequestEntry to be used in the kinesis.PutRecords API call. diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go index 3885900..60cb60f 100644 --- a/producer/firehosewriter.go +++ b/producer/firehosewriter.go @@ -17,9 +17,10 @@ import ( ) type firehoseWriterOptions struct { - msgCountRateLimit int // maximum number of records to be sent per second - msgSizeRateLimit int // maximum (transmission) size of records to be sent per second - Stats StatsCollector // stats collection mechanism + msgCountRateLimit int // maximum number of records to be sent per second + msgSizeRateLimit int // maximum (transmission) size of records to be sent per second + throughputMultiplier int // integer multiplier to increase firehose throughput rate limits + Stats StatsCollector // stats collection mechanism } // FirehoseWriter handles the API to send records to Kinesis. @@ -144,5 +145,5 @@ func (w *FirehoseWriter) getMsgSizeRateLimit() int { // getConcurrencyMultiplier returns the writer's concurrency multiplier. For the firehosewriter the multiplier is 1. func (w *FirehoseWriter) getConcurrencyMultiplier() (int, error) { - return 1, nil + return w.throughputMultiplier, nil } diff --git a/producer/firehosewriterconfig.go b/producer/firehosewriterconfig.go index ef9c023..3d22692 100644 --- a/producer/firehosewriterconfig.go +++ b/producer/firehosewriterconfig.go @@ -21,9 +21,10 @@ func NewFirehoseWriterConfig(cfg *aws.Config) *FirehoseWriterConfig { return &FirehoseWriterConfig{ AwsConfig: cfg, firehoseWriterOptions: &firehoseWriterOptions{ - msgCountRateLimit: firehoseMsgCountRateLimit, - msgSizeRateLimit: firehoseMsgSizeRateLimit, - Stats: &NilStatsCollector{}, + msgCountRateLimit: firehoseMsgCountRateLimit, + msgSizeRateLimit: firehoseMsgSizeRateLimit, + throughputMultiplier: 1, + Stats: &NilStatsCollector{}, }, LogLevel: *cfg.LogLevel, } @@ -45,6 +46,11 @@ func (c *FirehoseWriterConfig) SetMsgSizeRateLimit(limit int) { c.msgSizeRateLimit = limit } +// SetThroughputMultiplier configures the throughput multiplier to increase the streaming rate limits for Firehose +func (c *FirehoseWriterConfig) SetThroughputMultiplier(count int) { + c.throughputMultiplier = count +} + // SetStatsCollector configures a listener to handle listener metrics. func (c *FirehoseWriterConfig) SetStatsCollector(stats StatsCollector) { c.Stats = stats diff --git a/producer/producer.go b/producer/producer.go index f9b642c..f4a3a74 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -125,46 +125,48 @@ func (p *Producer) produce() { return case status := <-p.status: var batch []*message.Message + timeout := time.After(p.batchTimeout) + + // Fill batch by pulling from the messages channel or flush after timeout + fillBatch: + for len(batch) < status.capacity { + select { + case <-timeout: + break fillBatch + case msg := <-p.messages: + if msg == nil { + // Drop nil message + break + } - // Start pulling from the message channel and waiting for the bucket token to - // fill at the same time + msgSize, err := msg.RequestEntrySize() + if err != nil { + p.LogError("Unable to retreive message size due to marshalling errors for: ", string(msg.Data)) + p.sendToDataSpill(msg) + break + } + if msgSize > p.writer.getMsgSizeRateLimit() { + p.LogError("Encountered a message that exceeded that message size rate limit: ", string(msg.Data)) + p.sendToDataSpill(msg) + break + } + batch = append(batch, msg) + } + } + + // Then wait (if necessary) for the required tokens before sending the batch to the worker wg := sync.WaitGroup{} wg.Add(1) go func() { defer wg.Done() - var batchMsgSize int - timeout := time.After(p.batchTimeout) - - fillBatch: - for len(batch) < status.capacity { - select { - case <-timeout: - break fillBatch - case msg := <-p.messages: - if msg == nil { - // Drop nil message - break - } - - var msgSize int - if msgSize = msg.RequestEntrySize(); msgSize > p.writer.getMsgSizeRateLimit() { - p.sendToDataSpill(msg) - break - } - batchMsgSize += msgSize - batch = append(batch, msg) - } - } - // Request and wait for the message size (transmission) rate limiter to // allow this payload ctx, cancel := context.WithTimeout(context.TODO(), p.batchTimeout) defer cancel() - if err := p.msgSizeLimiter.WaitN(ctx, batchMsgSize+status.failedSize); err != nil { - + if err := p.msgSizeLimiter.WaitN(ctx, len(batch)+status.failedSize); err != nil { + p.LogError("Error occured waiting for message size tokens") } - cancel() }() wg.Add(1) @@ -174,14 +176,12 @@ func (p *Producer) produce() { // Request and wait for the message counter rate limiter to allow this // payload ctx, cancel := context.WithTimeout(context.TODO(), p.batchTimeout) - err := p.msgCountLimiter.WaitN(ctx, status.capacity+status.failedCount) + defer cancel() + err := p.msgCountLimiter.WaitN(ctx, len(batch)+status.failedCount) if err != nil { - // TODO: handle error properly + p.LogError("Error occured waiting for message count tokens") } - cancel() }() - - // Wait until the batch is ready and the rate limiter gives the go ahead wg.Wait() // Send batch regardless if it is empty or not @@ -246,15 +246,14 @@ func (p *Producer) shutdown() { }) } -// resizeWorkerPool is called to instantiate new workers, decommission workers or recommission workers that have been -// deactivated +// resizeWorkerPool is called to spawn new go routines or send stop signals to existing workers func (p *Producer) resizeWorkerPool(desiredWorkerCount int) { p.resizeMu.Lock() defer p.resizeMu.Unlock() if p.workerCount < desiredWorkerCount { for p.workerCount < desiredWorkerCount { - go p.newWorker() + go p.doWork() p.workerCount++ } } else { @@ -265,13 +264,13 @@ func (p *Producer) resizeWorkerPool(desiredWorkerCount int) { } } -// newWorker is a (blocking) helper function that increases the number of concurrent go routines calling the -// sendBatch function. Communications between the produce function and the newWorker function occurs on the status +// doWork is a (blocking) helper function that, when called as a separate go routine, increases the number of concurrent +// sendBatch functions. Communications between the produce function and the doWork function occurs on the status // channel where the "worker" provides information regarding its previously failed message count, its capacity for // new messages and the "worker's" channel to which the produce function should send the batches. The "workers" also // listen on the dismiss channel which upon receiving a signal will continue sending previously failed messages only // until all failed messages have been sent successfully or aged out. -func (p *Producer) newWorker() { +func (p *Producer) doWork() { batches := make(chan []*message.Message) defer close(batches) @@ -296,15 +295,20 @@ func (p *Producer) newWorker() { } var failedSize int for _, msg := range retries { - failedSize += msg.RequestEntrySize() + size, err := msg.RequestEntrySize() + if err != nil { + p.LogError("Unable to retreive message size due to marshalling errors for: ", string(msg.Data)) + p.sendToDataSpill(msg) + continue + } + failedSize += size } - status := &statusReport{ + p.status <- &statusReport{ capacity: capacity, failedCount: failedCount, failedSize: failedSize, channel: batches, } - p.status <- status // Receive a batch of messages and call the producer's sendBatch function batch := <-batches From c892a14121310853c31ef4f5100ccf3415c4c857 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 26 Jul 2017 17:36:16 -0500 Subject: [PATCH 065/109] Minor refactor of listener to clean up a bit. --- listener/kclreader.go | 1 - listener/kinesisreader.go | 1 - listener/listener.go | 91 +++++++++++++++++++++++---------------- producer/producer.go | 13 ++++-- producer/producer_test.go | 2 +- 5 files changed, 64 insertions(+), 44 deletions(-) diff --git a/listener/kclreader.go b/listener/kclreader.go index 5e4f9d6..8a912b0 100644 --- a/listener/kclreader.go +++ b/listener/kclreader.go @@ -27,7 +27,6 @@ type kclReaderOptions struct { type KclReader struct { *kclReaderOptions *logging.LogHelper - throttleSem chan empty pipeOfDeath chan empty scanner *bufio.Scanner diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index dc0f5bd..e201c58 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -30,7 +30,6 @@ type kinesisReaderOptions struct { type KinesisReader struct { *kinesisReaderOptions *logging.LogHelper - stream string shard string throttleSem chan empty diff --git a/listener/listener.go b/listener/listener.go index 0140474..9c8df6a 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -45,7 +45,6 @@ type listenerOptions struct { type Listener struct { *listenerOptions *logging.LogHelper - reader StreamReader messages chan *message.Message concurrencySem chan empty @@ -108,6 +107,57 @@ func (l *Listener) stopConsuming() { l.consuming = false } +func (l *Listener) enqueueSingle(ctx context.Context) (int, error) { + n, err := l.reader.GetRecord(ctx, func(msg *message.Message, wg *sync.WaitGroup) error { + defer wg.Done() + l.messages <- msg + + return nil + }) + if err != nil { + l.handleErrorLogging(err) + return 0, err + } + return n, nil +} + +func (l *Listener) enqueueBatch(ctx context.Context) (int, error) { + n, err := l.reader.GetRecords(ctx, + func(msg *message.Message, wg *sync.WaitGroup) error { + defer wg.Done() + l.messages <- msg + + return nil + }) + if err != nil { + l.handleErrorLogging(err) + return 0, err + } + return n, nil +} + +func (l *Listener) handleErrorLogging(err error) { + switch err := err.(type) { + case net.Error: + if err.Timeout() { + l.Stats.AddGetRecordsTimeout(1) + l.LogError("Received net error:", err.Error()) + } else { + l.LogError("Received unknown net error:", err.Error()) + } + case error: + switch err { + case errs.ErrTimeoutReadResponseBody: + l.Stats.AddGetRecordsReadTimeout(1) + l.LogError("Received error:", err.Error()) + default: + l.LogError("Received error:", err.Error()) + } + default: + l.LogError("Received unknown error:", err.Error()) + } +} + // RetrieveWithContext waits for a message from the stream and returns the message. Cancellation is supported through // contexts. func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, error) { @@ -126,15 +176,7 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, e if !ok { return nil, err } - n, err := l.reader.GetRecord(childCtx, func(msg *message.Message, wg *sync.WaitGroup) error { - defer wg.Done() - l.messages <- msg - - return nil - }) - if err != nil { - return nil, err - } + n, err := l.enqueueSingle(childCtx) if n > 0 { l.Stats.AddDelivered(n) return <-l.messages, nil @@ -197,34 +239,7 @@ func (l *Listener) consume(ctx context.Context) { return } - _, err := l.reader.GetRecords(childCtx, - func(msg *message.Message, wg *sync.WaitGroup) error { - defer wg.Done() - l.messages <- msg - - return nil - }) - if err != nil { - switch err := err.(type) { - case net.Error: - if err.Timeout() { - l.Stats.AddGetRecordsTimeout(1) - l.LogError("Received net error:", err.Error()) - } else { - l.LogError("Received unknown net error:", err.Error()) - } - case error: - switch err { - case errs.ErrTimeoutReadResponseBody: - l.Stats.AddGetRecordsReadTimeout(1) - l.LogError("Received error:", err.Error()) - default: - l.LogError("Received error:", err.Error()) - } - default: - l.LogError("Received unknown error:", err.Error()) - } - } + l.enqueueBatch(childCtx) } }() } diff --git a/producer/producer.go b/producer/producer.go index f4a3a74..3537c33 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -42,7 +42,6 @@ type Producer struct { startupOnce sync.Once // used to ensure that the startup function is called once shutdownOnce sync.Once // used to ensure that the shutdown function is called once resizeMu sync.Mutex // used to prevent resizeWorkerPool from being called synchronously with itself - shardCountMu sync.Mutex // used to add thread-safe access to the shardCount variable noCopy noCopy // prevents the Producer from being copied } @@ -127,8 +126,8 @@ func (p *Producer) produce() { var batch []*message.Message timeout := time.After(p.batchTimeout) - // Fill batch by pulling from the messages channel or flush after timeout fillBatch: + // Fill a batch by pulling from the messages channel or flushing after the timeout for len(batch) < status.capacity { select { case <-timeout: @@ -239,7 +238,15 @@ func (p *Producer) shutdown() { } // Stop the running go routine in produce - p.stop <- empty{} + if p.stop != nil { + p.stop <- empty{} + close(p.stop) + } + + // Close the pipeOfDeath channel + if p.pipeOfDeath != nil { + close(p.pipeOfDeath) + } // Reset startupOnce to allow the start up sequence to happen again p.startupOnce = sync.Once{} diff --git a/producer/producer_test.go b/producer/producer_test.go index 1230680..1eab288 100644 --- a/producer/producer_test.go +++ b/producer/producer_test.go @@ -30,7 +30,7 @@ func TestProducer(t *testing.T) { So(err, ShouldBeNil) err = k.WaitUntilStreamExists(context.TODO(), stream, - request.WithWaiterDelay(request.ConstantWaiterDelay(1*time.Second))) + request.WithWaiterDelay(request.ConstantWaiterDelay(time.Second))) So(err, ShouldBeNil) shards, err := k.GetShards(stream) From e100d7d58ba41e826aec7af07c7843c773892053 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 26 Jul 2017 23:12:44 -0500 Subject: [PATCH 066/109] Added rate limiting to the listener. Added payload size as one of the return parameters for the GetRecord and GetRecords interface. --- listener/kclreader.go | 13 ++++--- listener/kinesisreader.go | 18 +++++----- listener/listener.go | 75 ++++++++++++++++++++++++++++----------- 3 files changed, 72 insertions(+), 34 deletions(-) diff --git a/listener/kclreader.go b/listener/kclreader.go index 8a912b0..7617e65 100644 --- a/listener/kclreader.go +++ b/listener/kclreader.go @@ -55,7 +55,7 @@ func NewKclReader(c *aws.Config, fn ...func(*KclReaderConfig)) (*KclReader, erro // message channel. After all the messages on the message buffer have been moved to the listener's message // channel, a message is sent (following the Multilang protocol) to acknowledge that the processRecords message // has been received / processed -func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, error) { +func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, int, error) { // Define the batchSize batchSize := 0 if len(r.msgBuffer) > 0 { @@ -67,6 +67,9 @@ func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, erro } r.Stats.AddBatchSize(batchSize) + // TODO: Define the payloadSize + var payloadSize int + // Loop through the message buffer and call the message handler function on each message var wg sync.WaitGroup for i := 0; i < batchSize; i++ { @@ -82,11 +85,11 @@ func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, erro err := r.sendMessage(multilang.NewStatusMessage(multilang.PROCESSRECORDS)) if err != nil { r.LogError(err) - return batchSize, err + return batchSize, payloadSize, err } } - return batchSize, nil + return batchSize, payloadSize, nil } func (r *KclReader) getAction() (*multilang.ActionMessage, error) { @@ -195,12 +198,12 @@ func (r *KclReader) onShutdown() error { // GetRecord calls processRecords to attempt to put one message from message buffer to the listener's message // channel -func (r *KclReader) GetRecord(ctx context.Context, fn MessageHandler) (int, error) { +func (r *KclReader) GetRecord(ctx context.Context, fn MessageHandler) (int, int, error) { return r.processRecords(fn, 1) } // GetRecords calls processRecords to attempt to put all messages on the message buffer on the listener's // message channel -func (r *KclReader) GetRecords(ctx context.Context, fn MessageHandler) (int, error) { +func (r *KclReader) GetRecords(ctx context.Context, fn MessageHandler) (int, int, error) { return r.processRecords(fn, -1) } diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index e201c58..5468797 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -127,9 +127,9 @@ func (r *KinesisReader) throttle(sem chan empty) { }) } -func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batchSize int) (int, error) { +func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batchSize int) (int, int, error) { if err := r.ensureShardIterator(); err != nil { - return 0, err + return 0, 0, err } r.throttle(r.throttleSem) @@ -176,7 +176,9 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch } }) + var payloadSize int req.Handlers.Unmarshal.PushBack(func(req *request.Request) { + payloadSize += int(req.HTTPRequest.ContentLength) r.Stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) r.LogDebug("Finished GetRecords Unmarshal, took", time.Since(start)) }) @@ -192,14 +194,14 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch default: r.LogDebug("Received AWS error:", err.Error()) } - return 0, err + return 0, 0, err } r.Stats.AddGetRecordsDuration(time.Since(start)) // Process Records r.LogDebug(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), r.shard, time.Since(start))) if resp == nil { - return 0, errs.ErrNilGetRecordsResponse + return 0, 0, errs.ErrNilGetRecordsResponse } delivered := 0 r.Stats.AddBatchSize(len(resp.Records)) @@ -211,7 +213,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch select { case <-ctx.Done(): r.LogInfo(fmt.Sprintf("getRecords received ctx.Done() while delivering messages, %d delivered, ~%d dropped", delivered, len(resp.Records)-delivered)) - return delivered, ctx.Err() + return delivered, payloadSize, ctx.Err() default: var wg sync.WaitGroup wg.Add(1) @@ -245,15 +247,15 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch // around. r.setNextShardIterator(*resp.NextShardIterator) } - return delivered, nil + return delivered, payloadSize, nil } // GetRecord calls getRecords and delivers one record into the messages channel. -func (r *KinesisReader) GetRecord(ctx context.Context, fn MessageHandler) (int, error) { +func (r *KinesisReader) GetRecord(ctx context.Context, fn MessageHandler) (int, int, error) { return r.getRecords(ctx, fn, 1) } // GetRecords calls getRecords and delivers each record into the messages channel. -func (r *KinesisReader) GetRecords(ctx context.Context, fn MessageHandler) (int, error) { +func (r *KinesisReader) GetRecords(ctx context.Context, fn MessageHandler) (int, int, error) { return r.getRecords(ctx, fn, r.batchSize) } diff --git a/listener/listener.go b/listener/listener.go index 9c8df6a..ddfa710 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -7,6 +7,7 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" + "github.com/golang/time/rate" "github.com/rewardStyle/kinetic/errs" "github.com/rewardStyle/kinetic/logging" "github.com/rewardStyle/kinetic/message" @@ -14,8 +15,8 @@ import ( // StreamReader is an interface that abstracts out a stream reader. type StreamReader interface { - GetRecord(context.Context, MessageHandler) (int, error) - GetRecords(context.Context, MessageHandler) (int, error) + GetRecord(context.Context, MessageHandler) (int, int, error) + GetRecords(context.Context, MessageHandler) (int, int, error) } // empty is used a as a dummy type for semaphore channels and the pipe of death channel. @@ -45,12 +46,14 @@ type listenerOptions struct { type Listener struct { *listenerOptions *logging.LogHelper - reader StreamReader - messages chan *message.Message - concurrencySem chan empty - pipeOfDeath chan empty - consuming bool - consumingMu sync.Mutex + reader StreamReader + txnCountRateLimiter *rate.Limiter + txSizeRateLimiter *rate.Limiter + messages chan *message.Message + concurrencySem chan empty + pipeOfDeath chan empty + consuming bool + consumingMu sync.Mutex } // NewListener creates a new Listener object for retrieving and listening to message(s) on a StreamReader. @@ -65,9 +68,7 @@ func NewListener(c *aws.Config, r StreamReader, fn ...func(*Config)) (*Listener, LogLevel: cfg.LogLevel, Logger: cfg.AwsConfig.Logger, }, - reader: r, - concurrencySem: make(chan empty, cfg.concurrency), - pipeOfDeath: make(chan empty), + reader: r, }, nil } @@ -79,6 +80,8 @@ func (l *Listener) startConsuming() bool { if !l.consuming { l.consuming = true l.messages = make(chan *message.Message, l.queueDepth) + l.concurrencySem = make(chan empty, l.concurrency) + l.pipeOfDeath = make(chan empty) return true } return false @@ -104,11 +107,14 @@ func (l *Listener) stopConsuming() { if l.consuming && l.messages != nil { close(l.messages) } + if l.concurrencySem != nil { + close(l.concurrencySem) + } l.consuming = false } -func (l *Listener) enqueueSingle(ctx context.Context) (int, error) { - n, err := l.reader.GetRecord(ctx, func(msg *message.Message, wg *sync.WaitGroup) error { +func (l *Listener) enqueueSingle(ctx context.Context) (int, int, error) { + n, m, err := l.reader.GetRecord(ctx, func(msg *message.Message, wg *sync.WaitGroup) error { defer wg.Done() l.messages <- msg @@ -116,13 +122,13 @@ func (l *Listener) enqueueSingle(ctx context.Context) (int, error) { }) if err != nil { l.handleErrorLogging(err) - return 0, err + return 0, 0, err } - return n, nil + return n, m, nil } -func (l *Listener) enqueueBatch(ctx context.Context) (int, error) { - n, err := l.reader.GetRecords(ctx, +func (l *Listener) enqueueBatch(ctx context.Context) (int, int, error) { + n, m, err := l.reader.GetRecords(ctx, func(msg *message.Message, wg *sync.WaitGroup) error { defer wg.Done() l.messages <- msg @@ -131,9 +137,9 @@ func (l *Listener) enqueueBatch(ctx context.Context) (int, error) { }) if err != nil { l.handleErrorLogging(err) - return 0, err + return 0, 0, err } - return n, nil + return n, m, nil } func (l *Listener) handleErrorLogging(err error) { @@ -176,7 +182,7 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, e if !ok { return nil, err } - n, err := l.enqueueSingle(childCtx) + n, _, err := l.enqueueSingle(childCtx) if n > 0 { l.Stats.AddDelivered(n) return <-l.messages, nil @@ -227,6 +233,10 @@ func (l *Listener) consume(ctx context.Context) { go func() { defer l.stopConsuming() + // TODO: make these parameters configurable also scale according to the shard count + l.txnCountRateLimiter = rate.NewLimiter(rate.Limit(5), 1) + l.txSizeRateLimiter = rate.NewLimiter(rate.Limit(2000000), 2000000) + childCtx, cancel := context.WithCancel(ctx) defer cancel() for { @@ -239,7 +249,30 @@ func (l *Listener) consume(ctx context.Context) { return } - l.enqueueBatch(childCtx) + _, payloadSize, err := l.enqueueBatch(childCtx) + if err != nil { + l.LogError("Encountered an error when calling enqueueBatch: ", err) + return + } + + wg := sync.WaitGroup{} + wg.Add(1) + go func() { + defer wg.Done() + + if err := l.txnCountRateLimiter.Wait(childCtx); err != nil { + l.LogError("Error occured waiting for transaction count tokens") + } + }() + + wg.Add(1) + go func() { + defer wg.Done() + if err := l.txSizeRateLimiter.WaitN(childCtx, payloadSize); err != nil { + l.LogError("Error occured waiting for transmission size tokens") + } + }() + wg.Wait() } }() } From 385c90caf1719297b4a7290ef34fecda898e9a1a Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 27 Jul 2017 10:33:39 -0500 Subject: [PATCH 067/109] Remove glide files --- glide.lock | 82 ------------------------------------------------------ glide.yaml | 67 -------------------------------------------- 2 files changed, 149 deletions(-) delete mode 100644 glide.lock delete mode 100644 glide.yaml diff --git a/glide.lock b/glide.lock deleted file mode 100644 index 00b8b30..0000000 --- a/glide.lock +++ /dev/null @@ -1,82 +0,0 @@ -hash: 478410f2430633b92c3939e14a69e21c2f95684bae1e74250640ccb998d23f48 -updated: 2017-07-20T11:57:12.855440218-05:00 -imports: -- name: github.com/aws/aws-sdk-go - version: 72e42b13da62269f68308fb6068b7ea691a416a4 - subpackages: - - aws - - aws/awserr - - aws/awsutil - - aws/client - - aws/client/metadata - - aws/corehandlers - - aws/credentials - - aws/credentials/ec2rolecreds - - aws/credentials/endpointcreds - - aws/credentials/stscreds - - aws/defaults - - aws/ec2metadata - - aws/endpoints - - aws/request - - aws/session - - aws/signer/v4 - - internal/shareddefaults - - private/protocol - - private/protocol/json/jsonutil - - private/protocol/jsonrpc - - private/protocol/query - - private/protocol/query/queryutil - - private/protocol/rest - - private/protocol/xml/xmlutil - - private/waiter - - service/firehose - - service/firehose/firehoseiface - - service/kinesis - - service/kinesis/kinesisiface - - service/sts -- name: github.com/go-ini/ini - version: 1730955e3146956d6a087861380f9b4667ed5071 -- name: github.com/golang/time - version: 8be79e1e0910c292df4e79c241bb7e8f7e725959 -- name: github.com/gopherjs/gopherjs - version: 0cc212b4ef1943b61e894a7d56490575b3216439 - subpackages: - - js -- name: github.com/jmespath/go-jmespath - version: bd40a432e4c76585ef6b72d3fd96fb9b6dc7b68d -- name: github.com/jtolds/gls - version: f1ac7f4f24f50328e6bc838ca4437d1612a0243c -- name: github.com/rcrowley/go-metrics - version: ab2277b1c5d15c3cba104e9cbddbdfc622df5ad8 -- name: github.com/smartystreets/assertions - version: c9ee7d9e9a2aeec0bee7c4a516f3e0ad7cb7e558 - subpackages: - - internal/go-render/render - - internal/oglematchers -- name: github.com/smartystreets/goconvey - version: 3bd662eac601ad6436e64776af2e112069eb2edc - subpackages: - - convey - - convey/gotest - - convey/reporting -- name: github.com/stretchr/testify - version: 69483b4bd14f5845b5a1e55bca19e954e827f1d0 - subpackages: - - assert -- name: gopkg.in/gcfg.v1 - version: 27e4946190b4a327b539185f2b5b1f7c84730728 - subpackages: - - scanner - - token - - types -- name: gopkg.in/warnings.v0 - version: 8a331561fe74dadba6edfc59f3be66c22c3b065d -testImports: -- name: github.com/davecgh/go-spew - version: 6d212800a42e8ab5c146b8ace3490ee17e5225f9 - subpackages: - - spew -- name: github.com/pmezard/go-difflib - version: d8ed2627bdf02c080bf22230dbb337003b7aba2d - subpackages: - - difflib diff --git a/glide.yaml b/glide.yaml deleted file mode 100644 index eb4b8ba..0000000 --- a/glide.yaml +++ /dev/null @@ -1,67 +0,0 @@ -package: github.com/rewardStyle/kinetic -import: -- package: github.com/aws/aws-sdk-go - version: v1.10.3 - subpackages: - - aws - - aws/awserr - - aws/awsutil - - aws/client - - aws/client/metadata - - aws/corehandlers - - aws/credentials - - aws/credentials/ec2rolecreds - - aws/credentials/endpointcreds - - aws/credentials/stscreds - - aws/defaults - - aws/ec2metadata - - aws/endpoints - - aws/request - - aws/session - - aws/signer/v4 - - private/protocol - - private/protocol/json/jsonutil - - private/protocol/jsonrpc - - private/protocol/query - - private/protocol/query/queryutil - - private/protocol/rest - - private/protocol/xml/xmlutil - - private/waiter - - service/firehose - - service/firehose/firehoseiface - - service/kinesis - - service/kinesis/kinesisiface - - service/sts -- package: github.com/go-ini/ini - version: 1730955e3146956d6a087861380f9b4667ed5071 -- package: github.com/gopherjs/gopherjs - version: 0cc212b4ef1943b61e894a7d56490575b3216439 - subpackages: - - js -- package: github.com/jmespath/go-jmespath - version: bd40a432e4c76585ef6b72d3fd96fb9b6dc7b68d -- package: github.com/jtolds/gls - version: f1ac7f4f24f50328e6bc838ca4437d1612a0243c -- package: github.com/smartystreets/assertions - version: c9ee7d9e9a2aeec0bee7c4a516f3e0ad7cb7e558 - subpackages: - - internal/go-render/render - - internal/oglematchers -- package: github.com/smartystreets/goconvey - version: 3bd662eac601ad6436e64776af2e112069eb2edc - subpackages: - - convey - - convey/gotest - - convey/reporting -- package: gopkg.in/gcfg.v1 - version: 27e4946190b4a327b539185f2b5b1f7c84730728 - subpackages: - - scanner - - token - - types -- package: gopkg.in/warnings.v0 - version: 8a331561fe74dadba6edfc59f3be66c22c3b065d -- package: github.com/stretchr/testify/assert - version: v1.1.4 -- package: github.com/rcrowley/go-metrics -- package: github.com/golang/time From efc417601cc3fb692ccc7dab1722e1640a9da88a Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 27 Jul 2017 11:07:08 -0500 Subject: [PATCH 068/109] Eliminated uncessary packages by moving files around --- README.md | 2 +- config.go | 67 ++++++++++++++++++++++++++++++-- config/config.go | 66 ------------------------------- config_test.go | 8 ++-- errs/errors.go => errors.go | 2 +- kinetic.go | 11 ++---- kinetic_integration_test.go | 11 +++--- listener/kclreader.go | 11 +++--- listener/kinesisreader.go | 22 +++++------ listener/kinesisreader_test.go | 5 +-- listener/listener.go | 38 +++++++++--------- listener/listener_test.go | 20 +++++----- listener/listenerconfig.go | 4 +- listener/listenerconfig_test.go | 7 ++-- listener/stats.go | 2 +- listener/stats_test.go | 4 +- logging/logging.go => logging.go | 8 ++-- message/message.go => message.go | 2 +- multilang/action.go | 8 ++-- producer/firehosewriter.go | 14 +++---- producer/kinesiswriter.go | 18 ++++----- producer/kinesiswriterconfig.go | 6 +-- producer/producer.go | 38 +++++++++--------- producer/producer_test.go | 7 ++-- producer/producerconfig.go | 7 ++-- producer/producerconfig_test.go | 12 +++--- producer/sharedtypes.go | 10 ++--- producer/stats.go | 2 +- producer/stats_test.go | 4 +- testexec/main.go | 11 +++--- 30 files changed, 197 insertions(+), 230 deletions(-) delete mode 100644 config/config.go rename errs/errors.go => errors.go (99%) rename logging/logging.go => logging.go (91%) rename message/message.go => message.go (99%) diff --git a/README.md b/README.md index 2ebcbdd..209fead 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,7 @@ go l.Listen(func(b []byte, fnwg *sync.WaitGroup){ }) // Send a message using the producer -p.Send(&message.Message{ +p.Send(&kinetic.Message{ Data: []byte(`{"foo":"bar"}`), }) diff --git a/config.go b/config.go index f6fb166..448778b 100644 --- a/config.go +++ b/config.go @@ -1,14 +1,17 @@ package kinetic import ( - "github.com/aws/aws-sdk-go/aws" + "net/http" + "time" - "github.com/rewardStyle/kinetic/config" + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/credentials" + "github.com/aws/aws-sdk-go/aws/session" ) // Config is used to configure a Kinetic instance type Config struct { - *config.AwsOptions + *AwsOptions *kineticOptions LogLevel aws.LogLevelType } @@ -16,7 +19,7 @@ type Config struct { // NewConfig creates a new instance of Config func NewConfig() *Config { return &Config{ - AwsOptions: config.DefaultAwsOptions(), + AwsOptions: DefaultAwsOptions(), kineticOptions: &kineticOptions{}, LogLevel: aws.LogOff, } @@ -27,3 +30,59 @@ func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { c.AwsOptions.SetLogLevel(logLevel) c.LogLevel = logLevel & 0xffff0000 } + +// AwsOptions helps configure an aws.Config and session.Session +type AwsOptions struct { + AwsConfig *aws.Config +} + +// DefaultAwsOptions initializes the default aws.Config struct +func DefaultAwsOptions() *AwsOptions { + return &AwsOptions{ + AwsConfig: aws.NewConfig().WithHTTPClient( + &http.Client{ + Timeout: 2 * time.Minute, + }, + ), + } +} + +// SetCredentials configures AWS credentials. +func (c *AwsOptions) SetCredentials(accessKey, secretKey, sessionToken string) { + c.AwsConfig.WithCredentials( + credentials.NewStaticCredentials(accessKey, secretKey, sessionToken), + ) +} + +// SetRegion configures the AWS region. +func (c *AwsOptions) SetRegion(region string) { + c.AwsConfig.WithRegion(region) +} + +// SetEndpoint sets the endpoint to be used by aws-sdk-go. +func (c *AwsOptions) SetEndpoint(endpoint string) { + c.AwsConfig.WithEndpoint(endpoint) +} + +// SetLogger configures the logger for Kinetic and the aws-sdk-go. +func (c *AwsOptions) SetLogger(logger aws.Logger) { + c.AwsConfig.WithLogger(logger) +} + +// SetLogLevel configures the log levels for the SDK. +func (c *AwsOptions) SetLogLevel(logLevel aws.LogLevelType) { + c.AwsConfig.WithLogLevel(logLevel & 0xffff) +} + +// SetHTTPClientTimeout configures the HTTP timeout for the SDK. +func (c *AwsOptions) SetHTTPClientTimeout(timeout time.Duration) { + c.AwsConfig.WithHTTPClient(&http.Client{ + Timeout: timeout, + }) +} + +// GetSession creates an instance of the session.Session to be used when creating service +// clients in aws-sdk-go. +func (c *AwsOptions) GetSession() (*session.Session, error) { + return session.NewSession(c.AwsConfig) +} diff --git a/config/config.go b/config/config.go deleted file mode 100644 index a86d1d9..0000000 --- a/config/config.go +++ /dev/null @@ -1,66 +0,0 @@ -package config - -import ( - "net/http" - "time" - - "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/aws/credentials" - "github.com/aws/aws-sdk-go/aws/session" -) - -// AwsOptions helps configure an aws.Config and session.Session -type AwsOptions struct { - AwsConfig *aws.Config -} - -// DefaultAwsOptions initializes the default aws.Config struct -func DefaultAwsOptions() *AwsOptions { - return &AwsOptions{ - AwsConfig: aws.NewConfig().WithHTTPClient( - &http.Client{ - Timeout: 2 * time.Minute, - }, - ), - } -} - -// SetCredentials configures AWS credentials. -func (c *AwsOptions) SetCredentials(accessKey, secretKey, sessionToken string) { - c.AwsConfig.WithCredentials( - credentials.NewStaticCredentials(accessKey, secretKey, sessionToken), - ) -} - -// SetRegion configures the AWS region. -func (c *AwsOptions) SetRegion(region string) { - c.AwsConfig.WithRegion(region) -} - -// SetEndpoint sets the endpoint to be used by aws-sdk-go. -func (c *AwsOptions) SetEndpoint(endpoint string) { - c.AwsConfig.WithEndpoint(endpoint) -} - -// SetLogger configures the logger for Kinetic and the aws-sdk-go. -func (c *AwsOptions) SetLogger(logger aws.Logger) { - c.AwsConfig.WithLogger(logger) -} - -// SetLogLevel configures the log levels for the SDK. -func (c *AwsOptions) SetLogLevel(logLevel aws.LogLevelType) { - c.AwsConfig.WithLogLevel(logLevel & 0xffff) -} - -// SetHTTPClientTimeout configures the HTTP timeout for the SDK. -func (c *AwsOptions) SetHTTPClientTimeout(timeout time.Duration) { - c.AwsConfig.WithHTTPClient(&http.Client{ - Timeout: timeout, - }) -} - -// GetSession creates an instance of the session.Session to be used when creating service -// clients in aws-sdk-go. -func (c *AwsOptions) GetSession() (*session.Session, error) { - return session.NewSession(c.AwsConfig) -} diff --git a/config_test.go b/config_test.go index 37b0027..392be51 100644 --- a/config_test.go +++ b/config_test.go @@ -9,8 +9,6 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/session" - - "github.com/rewardStyle/kinetic/logging" ) func getSession(config *Config) *session.Session { @@ -27,7 +25,7 @@ func TestNewConfig(t *testing.T) { Convey("check the default values for its non-zero config", func() { So(config.AwsConfig.HTTPClient.Timeout, ShouldEqual, 2*time.Minute) So(config.AwsConfig.LogLevel.Value(), ShouldEqual, aws.LogOff) - So(config.LogLevel.Value(), ShouldEqual, logging.LogOff) + So(config.LogLevel.Value(), ShouldEqual, LogOff) }) Convey("check that we can retrieve an aws.Session from it ", func() { @@ -79,12 +77,12 @@ func TestNewConfig(t *testing.T) { }) Convey("check that we can set both the sdk and kinetic log level", func() { - ll := aws.LogDebug | aws.LogDebugWithSigning | logging.LogDebug + ll := aws.LogDebug | aws.LogDebugWithSigning | LogDebug config.SetLogLevel(ll) sess := getSession(config) So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) - So(config.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) + So(config.LogLevel.AtLeast(LogDebug), ShouldBeTrue) }) Convey("check that we can set the http.Client Timeout", func() { diff --git a/errs/errors.go b/errors.go similarity index 99% rename from errs/errors.go rename to errors.go index 7cfb441..bce6888 100644 --- a/errs/errors.go +++ b/errors.go @@ -1,4 +1,4 @@ -package errs +package kinetic import ( "errors" diff --git a/kinetic.go b/kinetic.go index 0885367..00fa4da 100644 --- a/kinetic.go +++ b/kinetic.go @@ -12,9 +12,6 @@ import ( "github.com/aws/aws-sdk-go/service/firehose/firehoseiface" "github.com/aws/aws-sdk-go/service/kinesis" "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" - - "github.com/rewardStyle/kinetic/errs" - "github.com/rewardStyle/kinetic/logging" ) type kineticOptions struct { @@ -24,7 +21,7 @@ type kineticOptions struct { // methods for interacting with the AWS services. type Kinetic struct { *kineticOptions - *logging.LogHelper + *LogHelper clientMu sync.Mutex fclient firehoseiface.FirehoseAPI @@ -42,7 +39,7 @@ func New(fn func(*Config)) (*Kinetic, error) { } return &Kinetic{ kineticOptions: config.kineticOptions, - LogHelper: &logging.LogHelper{ + LogHelper: &LogHelper{ LogLevel: config.LogLevel, Logger: sess.Config.Logger, }, @@ -132,10 +129,10 @@ func (k *Kinetic) GetShards(stream string) ([]string, error) { return nil, err } if resp == nil { - return nil, errs.ErrNilDescribeStreamResponse + return nil, ErrNilDescribeStreamResponse } if resp.StreamDescription == nil { - return nil, errs.ErrNilStreamDescription + return nil, ErrNilStreamDescription } var shards []string for _, shard := range resp.StreamDescription.Shards { diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index 548e1f0..5036c5e 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -13,12 +13,11 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" "github.com/rewardStyle/kinetic/listener" - "github.com/rewardStyle/kinetic/message" "github.com/rewardStyle/kinetic/producer" "github.com/stretchr/testify/assert" ) -type Message struct { +type TestMessage struct { ID int `json:"id"` Message string `json:"message"` } @@ -172,12 +171,12 @@ func TestKineticIntegration(t *testing.T) { go func(sent *int) { defer wg.Done() for i := 0; i < numMsg; i++ { - msg := &Message{ + msg := &TestMessage{ ID: i, Message: "hello_" + strconv.Itoa(i), } jsonStr, _ := json.Marshal(msg) - if err := p.Send(&message.Message{ + if err := p.Send(&Message{ PartitionKey: aws.String("key"), Data: []byte(jsonStr), }); err == nil { @@ -188,10 +187,10 @@ func TestKineticIntegration(t *testing.T) { // Use the listener to read messages from the kinetic stream go func() { - l.Listen(func(m *message.Message, fnwg *sync.WaitGroup) error { + l.Listen(func(m *Message, fnwg *sync.WaitGroup) error { defer fnwg.Done() - msg := &Message{} + msg := &TestMessage{} json.Unmarshal(m.Data, msg) if !streamData.exists(msg.ID) { diff --git a/listener/kclreader.go b/listener/kclreader.go index 7617e65..e10d90d 100644 --- a/listener/kclreader.go +++ b/listener/kclreader.go @@ -11,8 +11,7 @@ import ( "sync" "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic/logging" - "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic" "github.com/rewardStyle/kinetic/multilang" ) @@ -26,12 +25,12 @@ type kclReaderOptions struct { // KclReader handles the KCL Multilang Protocol to read records from KCL type KclReader struct { *kclReaderOptions - *logging.LogHelper + *kinetic.LogHelper throttleSem chan empty pipeOfDeath chan empty scanner *bufio.Scanner reader *bufio.Reader - msgBuffer []message.Message + msgBuffer []kinetic.Message } // NewKclReader creates a new stream reader to read records from KCL @@ -42,12 +41,12 @@ func NewKclReader(c *aws.Config, fn ...func(*KclReaderConfig)) (*KclReader, erro } return &KclReader{ kclReaderOptions: cfg.kclReaderOptions, - LogHelper: &logging.LogHelper{ + LogHelper: &kinetic.LogHelper{ LogLevel: cfg.LogLevel, Logger: cfg.AwsConfig.Logger, }, throttleSem: make(chan empty, 5), - msgBuffer: []message.Message{}, + msgBuffer: []kinetic.Message{}, }, nil } diff --git a/listener/kinesisreader.go b/listener/kinesisreader.go index 5468797..d1fedd1 100644 --- a/listener/kinesisreader.go +++ b/listener/kinesisreader.go @@ -13,9 +13,7 @@ import ( "github.com/aws/aws-sdk-go/service/kinesis" "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" - "github.com/rewardStyle/kinetic/errs" - "github.com/rewardStyle/kinetic/logging" - "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic" ) // kinesisReaderOptions is used to hold all of the configurable settings of a KinesisReader. @@ -29,7 +27,7 @@ type kinesisReaderOptions struct { // KinesisReader handles the API to read records from Kinesis. type KinesisReader struct { *kinesisReaderOptions - *logging.LogHelper + *kinetic.LogHelper stream string shard string throttleSem chan empty @@ -50,7 +48,7 @@ func NewKinesisReader(c *aws.Config, stream string, shard string, fn ...func(*Ki } return &KinesisReader{ kinesisReaderOptions: cfg.kinesisReaderOptions, - LogHelper: &logging.LogHelper{ + LogHelper: &kinetic.LogHelper{ LogLevel: cfg.LogLevel, Logger: cfg.AwsConfig.Logger, }, @@ -83,10 +81,10 @@ func (r *KinesisReader) ensureShardIterator() error { return err } if resp == nil { - return errs.ErrNilGetShardIteratorResponse + return kinetic.ErrNilGetShardIteratorResponse } if resp.ShardIterator == nil { - return errs.ErrNilShardIterator + return kinetic.ErrNilShardIterator } return r.setNextShardIterator(*resp.ShardIterator) } @@ -97,7 +95,7 @@ func (r *KinesisReader) ensureShardIterator() error { // be taken to ensure that only one call to Listen and Retrieve/RetrieveFn can be running at a time. func (r *KinesisReader) setNextShardIterator(shardIterator string) error { if len(shardIterator) == 0 { - return errs.ErrEmptyShardIterator + return kinetic.ErrEmptyShardIterator } r.nextShardIterator = shardIterator return nil @@ -112,7 +110,7 @@ func (r *KinesisReader) setNextShardIterator(shardIterator string) error { // time. func (r *KinesisReader) setSequenceNumber(sequenceNumber string) error { if len(sequenceNumber) == 0 { - return errs.ErrEmptySequenceNumber + return kinetic.ErrEmptySequenceNumber } r.shardIterator.AtSequenceNumber(sequenceNumber) return nil @@ -147,7 +145,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch req.ApplyOptions(request.WithResponseReadTimeout(r.responseReadTimeout)) // If debug is turned on, add some handlers for GetRecords logging - if r.LogLevel.AtLeast(logging.LogDebug) { + if r.LogLevel.AtLeast(kinetic.LogDebug) { req.Handlers.Send.PushBack(func(req *request.Request) { r.LogDebug("Finished getRecords Send, took", time.Since(start)) }) @@ -201,7 +199,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch // Process Records r.LogDebug(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), r.shard, time.Since(start))) if resp == nil { - return 0, 0, errs.ErrNilGetRecordsResponse + return 0, 0, kinetic.ErrNilGetRecordsResponse } delivered := 0 r.Stats.AddBatchSize(len(resp.Records)) @@ -217,7 +215,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch default: var wg sync.WaitGroup wg.Add(1) - go fn(message.FromRecord(record), &wg) + go fn(kinetic.FromRecord(record), &wg) wg.Wait() delivered++ r.Stats.AddConsumed(1) diff --git a/listener/kinesisreader_test.go b/listener/kinesisreader_test.go index d661973..9e87549 100644 --- a/listener/kinesisreader_test.go +++ b/listener/kinesisreader_test.go @@ -9,7 +9,6 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/logging" ) func TestNewKinesisReader(t *testing.T) { @@ -42,7 +41,7 @@ func TestNewKinesisReader(t *testing.T) { Convey("check that we can create a new KinesisReader with configured values", func() { batchSize := rand.Int() respReadTimeout := time.Duration(rand.Int()) * time.Second - logLevel := aws.LogDebug | aws.LogDebugWithSigning | logging.LogDebug + logLevel := aws.LogDebug | aws.LogDebugWithSigning | kinetic.LogDebug shardIterator := NewShardIterator() myStatsCollector := &NilStatsCollector{} r, err := NewKinesisReader(k.Session.Config, stream, shard, func(krc *KinesisReaderConfig) { @@ -56,7 +55,7 @@ func TestNewKinesisReader(t *testing.T) { So(err, ShouldBeNil) So(r.batchSize, ShouldEqual, batchSize) So(r.responseReadTimeout, ShouldEqual, respReadTimeout) - So(r.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) + So(r.LogLevel.AtLeast(kinetic.LogDebug), ShouldBeTrue) So(r.shardIterator, ShouldEqual, shardIterator) So(r.Stats, ShouldEqual, myStatsCollector) }) diff --git a/listener/listener.go b/listener/listener.go index ddfa710..86dbe0a 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -7,10 +7,8 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "github.com/golang/time/rate" - "github.com/rewardStyle/kinetic/errs" - "github.com/rewardStyle/kinetic/logging" - "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic" + "golang.org/x/time/rate" ) // StreamReader is an interface that abstracts out a stream reader. @@ -25,15 +23,15 @@ type empty struct{} // MessageProcessor defines the signature of a message handler used by Listen, RetrieveFn and their associated // *WithContext functions. MessageHandler accepts a WaitGroup so the function can be run as a blocking operation as // opposed to MessageHandlerAsync. -type MessageProcessor func(*message.Message, *sync.WaitGroup) error +type MessageProcessor func(*kinetic.Message, *sync.WaitGroup) error // MessageHandler defines the signature of a message handler used by GetRecord() and GetRecords(). MessageHandler // accepts a WaitGroup so the function can be run as a blocking operation as opposed to MessageHandlerAsync. -type MessageHandler func(*message.Message, *sync.WaitGroup) error +type MessageHandler func(*kinetic.Message, *sync.WaitGroup) error // MessageHandlerAsync defines the signature of a message handler used by GetRecord() and GetRecords(). // MessageHandlerAsync is meant to be run asynchronously. -type MessageHandlerAsync func(*message.Message) error +type MessageHandlerAsync func(*kinetic.Message) error // listenerOptions is used to hold all of the configurable settings of a Listener object. type listenerOptions struct { @@ -45,11 +43,11 @@ type listenerOptions struct { // Listener polls the StreamReader for messages. type Listener struct { *listenerOptions - *logging.LogHelper + *kinetic.LogHelper reader StreamReader txnCountRateLimiter *rate.Limiter txSizeRateLimiter *rate.Limiter - messages chan *message.Message + messages chan *kinetic.Message concurrencySem chan empty pipeOfDeath chan empty consuming bool @@ -64,7 +62,7 @@ func NewListener(c *aws.Config, r StreamReader, fn ...func(*Config)) (*Listener, } return &Listener{ listenerOptions: cfg.listenerOptions, - LogHelper: &logging.LogHelper{ + LogHelper: &kinetic.LogHelper{ LogLevel: cfg.LogLevel, Logger: cfg.AwsConfig.Logger, }, @@ -79,7 +77,7 @@ func (l *Listener) startConsuming() bool { defer l.consumingMu.Unlock() if !l.consuming { l.consuming = true - l.messages = make(chan *message.Message, l.queueDepth) + l.messages = make(chan *kinetic.Message, l.queueDepth) l.concurrencySem = make(chan empty, l.concurrency) l.pipeOfDeath = make(chan empty) return true @@ -92,7 +90,7 @@ func (l *Listener) startConsuming() bool { func (l *Listener) shouldConsume(ctx context.Context) (bool, error) { select { case <-l.pipeOfDeath: - return false, errs.ErrPipeOfDeath + return false, kinetic.ErrPipeOfDeath case <-ctx.Done(): return false, ctx.Err() default: @@ -114,7 +112,7 @@ func (l *Listener) stopConsuming() { } func (l *Listener) enqueueSingle(ctx context.Context) (int, int, error) { - n, m, err := l.reader.GetRecord(ctx, func(msg *message.Message, wg *sync.WaitGroup) error { + n, m, err := l.reader.GetRecord(ctx, func(msg *kinetic.Message, wg *sync.WaitGroup) error { defer wg.Done() l.messages <- msg @@ -129,7 +127,7 @@ func (l *Listener) enqueueSingle(ctx context.Context) (int, int, error) { func (l *Listener) enqueueBatch(ctx context.Context) (int, int, error) { n, m, err := l.reader.GetRecords(ctx, - func(msg *message.Message, wg *sync.WaitGroup) error { + func(msg *kinetic.Message, wg *sync.WaitGroup) error { defer wg.Done() l.messages <- msg @@ -153,7 +151,7 @@ func (l *Listener) handleErrorLogging(err error) { } case error: switch err { - case errs.ErrTimeoutReadResponseBody: + case kinetic.ErrTimeoutReadResponseBody: l.Stats.AddGetRecordsReadTimeout(1) l.LogError("Received error:", err.Error()) default: @@ -164,11 +162,11 @@ func (l *Listener) handleErrorLogging(err error) { } } -// RetrieveWithContext waits for a message from the stream and returns the message. Cancellation is supported through +// RetrieveWithContext waits for a message from the stream and returns the kinetic. Cancellation is supported through // contexts. -func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, error) { +func (l *Listener) RetrieveWithContext(ctx context.Context) (*kinetic.Message, error) { if !l.startConsuming() { - return nil, errs.ErrAlreadyConsuming + return nil, kinetic.ErrAlreadyConsuming } defer l.stopConsuming() @@ -191,7 +189,7 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*message.Message, e } // Retrieve waits for a message from the stream and returns the value -func (l *Listener) Retrieve() (*message.Message, error) { +func (l *Listener) Retrieve() (*kinetic.Message, error) { return l.RetrieveWithContext(context.TODO()) } @@ -295,7 +293,7 @@ func (l *Listener) ListenWithContext(ctx context.Context, fn MessageProcessor) { // couple more messages (especially since select is random in which channel is read from). l.concurrencySem <- empty{} wg.Add(1) - go func(msg *message.Message) { + go func(msg *kinetic.Message) { defer func() { <-l.concurrencySem }() diff --git a/listener/listener_test.go b/listener/listener_test.go index c142de1..43757d5 100644 --- a/listener/listener_test.go +++ b/listener/listener_test.go @@ -17,8 +17,6 @@ import ( "github.com/aws/aws-sdk-go/service/kinesis" "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/errs" - "github.com/rewardStyle/kinetic/message" ) func putRecord(l *Listener, b []byte) (*string, error) { @@ -83,12 +81,12 @@ func TestListener(t *testing.T) { Convey("check that setting an empty shard iterator returns an error", func() { err := l.reader.(*KinesisReader).setNextShardIterator("") - So(err, ShouldEqual, errs.ErrEmptyShardIterator) + So(err, ShouldEqual, kinetic.ErrEmptyShardIterator) }) Convey("check that setting an empty sequence number returns an error", func() { err := l.reader.(*KinesisReader).setSequenceNumber("") - So(err, ShouldEqual, errs.ErrEmptySequenceNumber) + So(err, ShouldEqual, kinetic.ErrEmptySequenceNumber) }) Convey("check that we can get the TRIM_HORIZON shard iterator", func() { @@ -181,7 +179,7 @@ func TestListener(t *testing.T) { }() <-time.After(10 * time.Millisecond) _, err := l.Retrieve() - So(err, ShouldEqual, errs.ErrAlreadyConsuming) + So(err, ShouldEqual, kinetic.ErrAlreadyConsuming) wg.Wait() }) @@ -191,7 +189,7 @@ func TestListener(t *testing.T) { go func() { ctx, cancel := context.WithTimeout(context.TODO(), 1000*time.Millisecond) defer cancel() - l.ListenWithContext(ctx, func(msg *message.Message, wg *sync.WaitGroup) error { + l.ListenWithContext(ctx, func(msg *kinetic.Message, wg *sync.WaitGroup) error { defer wg.Done() return nil }) @@ -199,7 +197,7 @@ func TestListener(t *testing.T) { }() <-time.After(10 * time.Millisecond) _, err := l.Retrieve() - So(err, ShouldEqual, errs.ErrAlreadyConsuming) + So(err, ShouldEqual, kinetic.ErrAlreadyConsuming) wg.Wait() }) @@ -209,7 +207,7 @@ func TestListener(t *testing.T) { secs := []float64{} for i := 1; i <= 6; i++ { start := time.Now() - l.reader.GetRecord(context.TODO(), func(msg *message.Message, wg *sync.WaitGroup) error { + l.reader.GetRecord(context.TODO(), func(msg *kinetic.Message, wg *sync.WaitGroup) error { defer wg.Done() return nil @@ -226,7 +224,7 @@ func TestListener(t *testing.T) { data := "retrieved" _, err := putRecord(l, []byte(data)) So(err, ShouldBeNil) - err = l.RetrieveFn(func(msg *message.Message, wg *sync.WaitGroup) error { + err = l.RetrieveFn(func(msg *kinetic.Message, wg *sync.WaitGroup) error { defer wg.Done() called = true @@ -247,7 +245,7 @@ func TestListener(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - l.Listen(func(msg *message.Message, wg *sync.WaitGroup) error { + l.Listen(func(msg *kinetic.Message, wg *sync.WaitGroup) error { defer wg.Done() atomic.AddInt64(&count, 1) @@ -295,7 +293,7 @@ func TestListener(t *testing.T) { go func() { ctx, cancel := context.WithCancel(context.TODO()) defer wg.Done() - l.ListenWithContext(ctx, func(m *message.Message, wg *sync.WaitGroup) error { + l.ListenWithContext(ctx, func(m *kinetic.Message, wg *sync.WaitGroup) error { defer wg.Done() time.AfterFunc(time.Duration(rand.Intn(10))*time.Second, func() { n, err := strconv.Atoi(string(m.Data)) diff --git a/listener/listenerconfig.go b/listener/listenerconfig.go index 5a2ba0f..8dab8b0 100644 --- a/listener/listenerconfig.go +++ b/listener/listenerconfig.go @@ -2,7 +2,7 @@ package listener import ( "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic/logging" + "github.com/rewardStyle/kinetic" ) // Config is used to configure a Listener instance @@ -21,7 +21,7 @@ func NewConfig(cfg *aws.Config) *Config { concurrency: 10000, Stats: &NilStatsCollector{}, }, - LogLevel: logging.LogOff, + LogLevel: kinetic.LogOff, } } diff --git a/listener/listenerconfig_test.go b/listener/listenerconfig_test.go index 531056f..484ee1c 100644 --- a/listener/listenerconfig_test.go +++ b/listener/listenerconfig_test.go @@ -8,7 +8,6 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/logging" ) type DebugStatsCollector struct{} @@ -39,13 +38,13 @@ func TestNewConfig(t *testing.T) { So(config.queueDepth, ShouldEqual, 10000) So(config.concurrency, ShouldEqual, 10000) So(config.Stats, ShouldHaveSameTypeAs, &NilStatsCollector{}) - So(config.LogLevel.Value(), ShouldEqual, logging.LogOff) + So(config.LogLevel.Value(), ShouldEqual, kinetic.LogOff) }) Convey("check that we can set both the sdk and kinetic log level", func() { - ll := aws.LogDebug | aws.LogDebugWithSigning | logging.LogDebug + ll := aws.LogDebug | aws.LogDebugWithSigning | kinetic.LogDebug config.SetLogLevel(ll) - So(config.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) + So(config.LogLevel.AtLeast(kinetic.LogDebug), ShouldBeTrue) }) Convey("check that we can set the concurrency limit", func() { diff --git a/listener/stats.go b/listener/stats.go index a3947a8..3326741 100644 --- a/listener/stats.go +++ b/listener/stats.go @@ -4,7 +4,7 @@ import ( "log" "time" - "github.com/rcrowley/go-metrics" + metrics "github.com/jasonyurs/go-metrics" ) // StatsCollector allows for a collector to collect various metrics produced by diff --git a/listener/stats_test.go b/listener/stats_test.go index c5b4650..9932cbd 100644 --- a/listener/stats_test.go +++ b/listener/stats_test.go @@ -1,12 +1,12 @@ package listener import ( - . "github.com/smartystreets/goconvey/convey" "math/rand" "testing" "time" - "github.com/rcrowley/go-metrics" + metrics "github.com/jasonyurs/go-metrics" + . "github.com/smartystreets/goconvey/convey" ) func init() { diff --git a/logging/logging.go b/logging.go similarity index 91% rename from logging/logging.go rename to logging.go index 3b66722..86f44b2 100644 --- a/logging/logging.go +++ b/logging.go @@ -1,11 +1,11 @@ -package logging +package kinetic import ( "github.com/aws/aws-sdk-go/aws" ) const ( - // LogOff disables all logging. + // LogOff disables all kinetic. LogOff aws.LogLevelType = (iota * 0x1000) << 16 // LogBug enables logging of bugs in code. @@ -17,7 +17,7 @@ const ( // LogInfo enables logging of informational messages.. LogInfo - // LogDebug enables debug logging. + // LogDebug enables debug kinetic. LogDebug ) @@ -58,7 +58,7 @@ func (l *LogHelper) LogError(args ...interface{}) { l.Log(LogError, args...) } -// LogInfo logs an informational message. +// LogInfo logs an informational kinetic. func (l *LogHelper) LogInfo(args ...interface{}) { l.Log(LogInfo, args...) } diff --git a/message/message.go b/message.go similarity index 99% rename from message/message.go rename to message.go index 1909956..d157171 100644 --- a/message/message.go +++ b/message.go @@ -1,4 +1,4 @@ -package message +package kinetic import ( "encoding/json" diff --git a/multilang/action.go b/multilang/action.go index e2547a0..be85dc6 100644 --- a/multilang/action.go +++ b/multilang/action.go @@ -5,7 +5,7 @@ import ( "strconv" "time" - "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic" ) // ActionType is used as an enum for KCL Multilang protocol action message types @@ -62,14 +62,14 @@ func (t *Timestamp) UnmarshalJSON(b []byte) error { return nil } -// ToMessage is used to transform a multilang.Record struct into a message.Message struct -func (r *Record) ToMessage() *message.Message { +// ToMessage is used to transform a multilang.Record struct into a kinetic.Message struct +func (r *Record) ToMessage() *kinetic.Message { b, err := base64.StdEncoding.DecodeString(r.Data) if err != nil { panic("There was a problem decoding kcl data") } - return &message.Message{ + return &kinetic.Message{ ApproximateArrivalTimestamp: &r.ApproximateArrival.Time, Data: b, PartitionKey: &r.PartitionKey, diff --git a/producer/firehosewriter.go b/producer/firehosewriter.go index 60cb60f..9b86860 100644 --- a/producer/firehosewriter.go +++ b/producer/firehosewriter.go @@ -11,9 +11,7 @@ import ( "github.com/aws/aws-sdk-go/service/firehose" "github.com/aws/aws-sdk-go/service/firehose/firehoseiface" - "github.com/rewardStyle/kinetic/errs" - "github.com/rewardStyle/kinetic/logging" - "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic" ) type firehoseWriterOptions struct { @@ -26,7 +24,7 @@ type firehoseWriterOptions struct { // FirehoseWriter handles the API to send records to Kinesis. type FirehoseWriter struct { *firehoseWriterOptions - *logging.LogHelper + *kinetic.LogHelper stream string client firehoseiface.FirehoseAPI @@ -44,7 +42,7 @@ func NewFirehoseWriter(c *aws.Config, stream string, fn ...func(*FirehoseWriterC } return &FirehoseWriter{ firehoseWriterOptions: cfg.firehoseWriterOptions, - LogHelper: &logging.LogHelper{ + LogHelper: &kinetic.LogHelper{ LogLevel: cfg.LogLevel, Logger: cfg.AwsConfig.Logger, }, @@ -54,7 +52,7 @@ func NewFirehoseWriter(c *aws.Config, stream string, fn ...func(*FirehoseWriterC } // PutRecords sends a batch of records to Firehose and returns a list of records that need to be retried. -func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageHandlerAsync) error { +func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*kinetic.Message, fn MessageHandlerAsync) error { var startSendTime time.Time var startBuildTime time.Time @@ -99,10 +97,10 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*message.Mes w.Stats.UpdatePutRecordsDuration(time.Since(start)) if resp == nil { - return errs.ErrNilPutRecordsResponse + return kinetic.ErrNilPutRecordsResponse } if resp.FailedPutCount == nil { - return errs.ErrNilFailedRecordCount + return kinetic.ErrNilFailedRecordCount } attempted := len(messages) failed := int(aws.Int64Value(resp.FailedPutCount)) diff --git a/producer/kinesiswriter.go b/producer/kinesiswriter.go index 11fa37a..342d689 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesiswriter.go @@ -11,9 +11,7 @@ import ( "github.com/aws/aws-sdk-go/service/kinesis" "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" - "github.com/rewardStyle/kinetic/errs" - "github.com/rewardStyle/kinetic/logging" - "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic" ) type kinesisWriterOptions struct { @@ -26,7 +24,7 @@ type kinesisWriterOptions struct { // KinesisWriter handles the API to send records to Kinesis. type KinesisWriter struct { *kinesisWriterOptions - *logging.LogHelper + *kinetic.LogHelper stream string client kinesisiface.KinesisAPI @@ -44,7 +42,7 @@ func NewKinesisWriter(c *aws.Config, stream string, fn ...func(*KinesisWriterCon } return &KinesisWriter{ kinesisWriterOptions: cfg.kinesisWriterOptions, - LogHelper: &logging.LogHelper{ + LogHelper: &kinetic.LogHelper{ LogLevel: cfg.LogLevel, Logger: cfg.AwsConfig.Logger, }, @@ -54,7 +52,7 @@ func NewKinesisWriter(c *aws.Config, stream string, fn ...func(*KinesisWriterCon } // PutRecords sends a batch of records to Kinesis and returns a list of records that need to be retried. -func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Message, fn MessageHandlerAsync) error { +func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*kinetic.Message, fn MessageHandlerAsync) error { var startSendTime time.Time var startBuildTime time.Time @@ -100,10 +98,10 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*message.Mess w.Stats.UpdatePutRecordsDuration(time.Since(start)) if resp == nil { - return errs.ErrNilPutRecordsResponse + return kinetic.ErrNilPutRecordsResponse } if resp.FailedRecordCount == nil { - return errs.ErrNilFailedRecordCount + return kinetic.ErrNilFailedRecordCount } attempted := len(messages) failed := int(aws.Int64Value(resp.FailedRecordCount)) @@ -156,10 +154,10 @@ func (w *KinesisWriter) getConcurrencyMultiplier() (int, error) { return 0, err } if resp == nil { - return 0, errs.ErrNilDescribeStreamResponse + return 0, kinetic.ErrNilDescribeStreamResponse } if resp.StreamDescription == nil { - return 0, errs.ErrNilStreamDescription + return 0, kinetic.ErrNilStreamDescription } // maps shardID to a boolean that indicates whether or not the shard is a parent shard or an adjacent parent shard diff --git a/producer/kinesiswriterconfig.go b/producer/kinesiswriterconfig.go index 69a1b6e..84a720d 100644 --- a/producer/kinesiswriterconfig.go +++ b/producer/kinesiswriterconfig.go @@ -5,7 +5,7 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic/errs" + "github.com/rewardStyle/kinetic" ) const ( @@ -45,7 +45,7 @@ func (c *KinesisWriterConfig) SetMsgCountRateLimit(limit int) { c.msgCountRateLimit = limit } else { log.Fatal("Message Count Rate Limit must be positive and less than ", kinesisMsgCountRateLimit) - panic(errs.ErrInvalidMsgCountRateLimit) + panic(kinetic.ErrInvalidMsgCountRateLimit) } } @@ -55,7 +55,7 @@ func (c *KinesisWriterConfig) SetMsgSizeRateLimit(limit int) { c.msgSizeRateLimit = limit } else { log.Fatal("Message Count Size Limit must be positive and less than ", kinesisMsgSizeRateLimit) - panic(errs.ErrInvalidMsgSizeRateLimit) + panic(kinetic.ErrInvalidMsgSizeRateLimit) } } diff --git a/producer/producer.go b/producer/producer.go index 3537c33..48e8186 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -8,10 +8,8 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/awserr" - "github.com/golang/time/rate" - "github.com/rewardStyle/kinetic/errs" - "github.com/rewardStyle/kinetic/logging" - "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic" + "golang.org/x/time/rate" ) // producerOptions holds all of the configurable settings for a Producer @@ -29,12 +27,12 @@ type producerOptions struct { // Producer sends records to AWS Kinesis or Firehose. type Producer struct { *producerOptions // contains all of the configuration settings for the Producer - *logging.LogHelper // object for help with logging + *kinetic.LogHelper // object for help with logging writer StreamWriter // interface for abstracting the PutRecords call msgCountLimiter *rate.Limiter // rate limiter to limit the number of messages dispatched per second msgSizeLimiter *rate.Limiter // rate limiter to limit the total size (in bytes) of messages dispatched per second workerCount int // number of concurrent workers sending batch messages for the producer - messages chan *message.Message // channel for enqueuing messages to be put on the stream + messages chan *kinetic.Message // channel for enqueuing messages to be put on the stream status chan *statusReport // channel for workers to communicate their current status dismiss chan empty // channel for handling the decommissioning of a surplus of workers stop chan empty // channel for handling shutdown @@ -53,7 +51,7 @@ func NewProducer(c *aws.Config, w StreamWriter, fn ...func(*Config)) (*Producer, } return &Producer{ producerOptions: cfg.producerOptions, - LogHelper: &logging.LogHelper{ + LogHelper: &kinetic.LogHelper{ LogLevel: cfg.LogLevel, Logger: cfg.AwsConfig.Logger, }, @@ -72,7 +70,7 @@ func (p *Producer) produce() { p.msgSizeLimiter = rate.NewLimiter(rate.Limit(float64(p.writer.getMsgSizeRateLimit())), p.writer.getMsgSizeRateLimit()) // Create communication channels - p.messages = make(chan *message.Message, p.queueDepth) + p.messages = make(chan *kinetic.Message, p.queueDepth) p.status = make(chan *statusReport) p.dismiss = make(chan empty) p.stop = make(chan empty) @@ -123,7 +121,7 @@ func (p *Producer) produce() { case <-p.stop: return case status := <-p.status: - var batch []*message.Message + var batch []*kinetic.Message timeout := time.After(p.batchTimeout) fillBatch: @@ -278,10 +276,10 @@ func (p *Producer) resizeWorkerPool(desiredWorkerCount int) { // listen on the dismiss channel which upon receiving a signal will continue sending previously failed messages only // until all failed messages have been sent successfully or aged out. func (p *Producer) doWork() { - batches := make(chan []*message.Message) + batches := make(chan []*kinetic.Message) defer close(batches) - var retries []*message.Message + var retries []*kinetic.Message var dismissed bool for ok := true; ok; ok = !dismissed || len(retries) != 0 { // Check to see if there were any signals to dismiss workers (if eligible) @@ -328,9 +326,9 @@ func (p *Producer) doWork() { // sendBatch is the function that is called by each worker to put records on the stream. sendBatch accepts a slice of // messages to send and returns a slice of messages that failed to send -func (p *Producer) sendBatch(batch []*message.Message) []*message.Message { - var failed []*message.Message - err := p.writer.PutRecords(context.TODO(), batch, func(msg *message.Message) error { +func (p *Producer) sendBatch(batch []*kinetic.Message) []*kinetic.Message { + var failed []*kinetic.Message + err := p.writer.PutRecords(context.TODO(), batch, func(msg *kinetic.Message) error { if msg.FailCount <= p.maxRetryAttempts { failed = append(failed, msg) p.Stats.AddSentRetried(1) @@ -357,7 +355,7 @@ func (p *Producer) sendBatch(batch []*message.Message) []*message.Message { p.LogError("Received AWS error:", err.Error()) case error: switch err { - case errs.ErrRetryRecords: + case kinetic.ErrRetryRecords: break default: p.LogError("Received error:", err.Error()) @@ -370,7 +368,7 @@ func (p *Producer) sendBatch(batch []*message.Message) []*message.Message { } // sendToDataSpill is called when the producer is unable to write the message to the stream -func (p *Producer) sendToDataSpill(msg *message.Message) { +func (p *Producer) sendToDataSpill(msg *kinetic.Message) { p.Stats.AddDroppedTotal(1) p.Stats.AddDroppedCapacity(1) if err := p.dataSpillFn(msg); err != nil { @@ -392,7 +390,7 @@ func (p *Producer) Close() { } // SendWithContext sends a message to the stream. Cancellation supported through contexts. -func (p *Producer) SendWithContext(ctx context.Context, msg *message.Message) error { +func (p *Producer) SendWithContext(ctx context.Context, msg *kinetic.Message) error { p.produce() select { case p.messages <- msg: @@ -403,19 +401,19 @@ func (p *Producer) SendWithContext(ctx context.Context, msg *message.Message) er } // Send a message to the stream, waiting on the message to be put into the channel. -func (p *Producer) Send(msg *message.Message) error { +func (p *Producer) Send(msg *kinetic.Message) error { return p.SendWithContext(context.TODO(), msg) } // TryToSend will attempt to send a message to the stream if the channel has capacity for a message, or will immediately // return with an error if the channel is full. -func (p *Producer) TryToSend(msg *message.Message) error { +func (p *Producer) TryToSend(msg *kinetic.Message) error { p.produce() select { case p.messages <- msg: return nil default: p.sendToDataSpill(msg) - return errs.ErrDroppedMessage + return kinetic.ErrDroppedMessage } } diff --git a/producer/producer_test.go b/producer/producer_test.go index 1eab288..d09fc4b 100644 --- a/producer/producer_test.go +++ b/producer/producer_test.go @@ -13,7 +13,6 @@ import ( "github.com/rewardStyle/kinetic" "github.com/rewardStyle/kinetic/listener" - "github.com/rewardStyle/kinetic/message" ) func TestProducer(t *testing.T) { @@ -83,7 +82,7 @@ func TestProducer(t *testing.T) { Convey("check that we can send and receive a single message", func() { start := time.Now() data := "hello" - p.Send(&message.Message{ + p.Send(&kinetic.Message{ PartitionKey: aws.String("key"), Data: []byte(data), }) @@ -97,7 +96,7 @@ func TestProducer(t *testing.T) { Convey("check that we can send a single message after batch timeout elapses", func() { start := time.Now() data := "hello" - p.Send(&message.Message{ + p.Send(&kinetic.Message{ PartitionKey: aws.String("key"), Data: []byte(data), }) @@ -130,7 +129,7 @@ func TestProducer(t *testing.T) { go func() { defer wg.Done() for _, datum := range data { - p.Send(&message.Message{ + p.Send(&kinetic.Message{ PartitionKey: aws.String("key"), Data: []byte(datum), }) diff --git a/producer/producerconfig.go b/producer/producerconfig.go index 4cab912..05f2ce1 100644 --- a/producer/producerconfig.go +++ b/producer/producerconfig.go @@ -5,8 +5,7 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic/errs" - "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic" ) // Config is used to configure a Producer instance. @@ -27,7 +26,7 @@ func NewConfig(cfg *aws.Config) *Config { maxRetryAttempts: 10, concurrency: 3, shardCheckFreq: time.Minute, - dataSpillFn: func(*message.Message) error { return nil }, + dataSpillFn: func(*kinetic.Message) error { return nil }, Stats: &NilStatsCollector{}, }, LogLevel: *cfg.LogLevel, @@ -41,7 +40,7 @@ func (c *Config) SetBatchSize(batchSize int) { } else { // http://docs.aws.amazon.com/firehose/latest/APIReference/API_PutRecordBatch.html log.Fatal("BatchSize must be less than or equal to 500 ") - panic(errs.ErrInvalidBatchSize) + panic(kinetic.ErrInvalidBatchSize) } } diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go index 949a2d0..41a34e3 100644 --- a/producer/producerconfig_test.go +++ b/producer/producerconfig_test.go @@ -8,8 +8,6 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/logging" - "github.com/rewardStyle/kinetic/message" ) type DebugStatsCollector struct{} @@ -32,7 +30,7 @@ func (l *DebugStatsCollector) UpdateProducerConcurrency(int) {} type DebugStreamWriter struct{} -func (w *DebugStreamWriter) PutRecords(batch []*message.Message) ([]*message.Message, error) { +func (w *DebugStreamWriter) PutRecords(batch []*kinetic.Message) ([]*kinetic.Message, error) { return nil, nil } @@ -52,13 +50,13 @@ func TestNewConfig(t *testing.T) { So(cfg.concurrency, ShouldEqual, 3) So(cfg.shardCheckFreq, ShouldEqual, time.Minute) So(cfg.Stats, ShouldHaveSameTypeAs, &NilStatsCollector{}) - So(cfg.LogLevel.Value(), ShouldEqual, logging.LogOff) + So(cfg.LogLevel.Value(), ShouldEqual, kinetic.LogOff) }) Convey("check that we can set both the sdk and kinetic log level", func() { - ll := aws.LogDebug | aws.LogDebugWithSigning | logging.LogDebug + ll := aws.LogDebug | aws.LogDebugWithSigning | kinetic.LogDebug cfg.SetLogLevel(ll) - So(cfg.LogLevel.AtLeast(logging.LogDebug), ShouldBeTrue) + So(cfg.LogLevel.AtLeast(kinetic.LogDebug), ShouldBeTrue) }) Convey("check that we can set the batch size", func() { @@ -87,7 +85,7 @@ func TestNewConfig(t *testing.T) { }) Convey("check that we can set the data spill callback function", func() { - fn := func(msg *message.Message) error { + fn := func(msg *kinetic.Message) error { return nil } cfg.SetDataSpillFn(fn) diff --git a/producer/sharedtypes.go b/producer/sharedtypes.go index 7704a34..0d9a2e1 100644 --- a/producer/sharedtypes.go +++ b/producer/sharedtypes.go @@ -4,12 +4,12 @@ import ( "context" "sync" - "github.com/rewardStyle/kinetic/message" + "github.com/rewardStyle/kinetic" ) // StreamWriter is an interface that abstracts the differences in API between Kinesis and Firehose. type StreamWriter interface { - PutRecords(context.Context, []*message.Message, MessageHandlerAsync) error + PutRecords(context.Context, []*kinetic.Message, MessageHandlerAsync) error getMsgCountRateLimit() int getMsgSizeRateLimit() int getConcurrencyMultiplier() (int, error) @@ -17,18 +17,18 @@ type StreamWriter interface { // MessageHandler defines the signature of a message handler used by PutRecords(). MessageHandler accepts a WaitGroup // so the function can be run as a blocking operation as opposed to MessageHandlerAsync. -type MessageHandler func(*message.Message, *sync.WaitGroup) error +type MessageHandler func(*kinetic.Message, *sync.WaitGroup) error // MessageHandlerAsync defines the signature of a message handler used by PutRecords(). MessageHandlerAsync is meant to // be run asynchronously. -type MessageHandlerAsync func(*message.Message) error +type MessageHandlerAsync func(*kinetic.Message) error // statusReport is used to communicate a worker's capacity for new messages and to which channel they should be sent. type statusReport struct { capacity int // maximum message capacity the worker can handle failedCount int // number of previous messages that failed to send failedSize int // size in bytes of the previous messages that failed to send - channel chan []*message.Message // channel of the worker to which the batch messages should be sent + channel chan []*kinetic.Message // channel of the worker to which the batch messages should be sent } // empty is used a as a dummy type for counting semaphore channels. diff --git a/producer/stats.go b/producer/stats.go index 88d3988..8def1e2 100644 --- a/producer/stats.go +++ b/producer/stats.go @@ -4,7 +4,7 @@ import ( "log" "time" - "github.com/rcrowley/go-metrics" + metrics "github.com/jasonyurs/go-metrics" ) // StatsCollector allows for a collector to collect various metrics produced by diff --git a/producer/stats_test.go b/producer/stats_test.go index f333d13..5d159cd 100644 --- a/producer/stats_test.go +++ b/producer/stats_test.go @@ -1,12 +1,12 @@ package producer import ( - . "github.com/smartystreets/goconvey/convey" "math/rand" "testing" "time" - "github.com/rcrowley/go-metrics" + metrics "github.com/jasonyurs/go-metrics" + . "github.com/smartystreets/goconvey/convey" ) func init() { diff --git a/testexec/main.go b/testexec/main.go index e340bf9..7fdfb6c 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -17,10 +17,9 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" "github.com/aws/aws-sdk-go/aws/session" - "github.com/rcrowley/go-metrics" + metrics "github.com/jasonyurs/go-metrics" "github.com/rewardStyle/kinetic" "github.com/rewardStyle/kinetic/listener" - "github.com/rewardStyle/kinetic/message" "github.com/rewardStyle/kinetic/producer" "net/http" @@ -239,7 +238,7 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce psc := producer.NewDefaultStatsCollector(registry) w, err := producer.NewKinesisWriter(k.Session.Config, streamName, func(kwc *producer.KinesisWriterConfig) { - //kwc.SetLogLevel(logging.LogDebug) + //kwc.SetLogLevel(kinetic.LogDebug) kwc.SetResponseReadTimeout(time.Second) kwc.SetStatsCollector(psc) kwc.SetMsgCountRateLimit(1000) @@ -257,7 +256,7 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce c.SetQueueDepth(10000) c.SetConcurrency(3) c.SetShardCheckFreq(time.Minute) - c.SetDataSpillFn(func(msg *message.Message) error { + c.SetDataSpillFn(func(msg *kinetic.Message) error { //log.Printf("Message was dropped: [%s]\n", string(msg.Data)) return nil }) @@ -393,7 +392,7 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { for { <-sendSignal jsonStr, _ := json.Marshal(NewMessage()) - if err := p.TryToSend(&message.Message{ + if err := p.TryToSend(&kinetic.Message{ PartitionKey: aws.String("key"), Data: []byte(jsonStr), }); err == nil { @@ -493,7 +492,7 @@ func listen(sd *StreamData, l *listener.Listener, wg *sync.WaitGroup) { // Call Listen within a go routine go func() { - l.Listen(func(m *message.Message, wg *sync.WaitGroup) error { + l.Listen(func(m *kinetic.Message, wg *sync.WaitGroup) error { defer wg.Done() // Unmarshal data From e4a35be7124868e7f3d62c9a826e72900dab5cf9 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 27 Jul 2017 21:44:48 -0500 Subject: [PATCH 069/109] Renamed listener to consumer. Introduced functional options for configuring consumer, producer, kinesisreader, kclreader, kinesiswriter and firehosewriter objects --- README.md | 4 +- listener/listener.go => consumer/consumer.go | 113 +++++++++++------- .../stats.go => consumer/consumer_stats.go | 50 ++++---- .../consumer_stats_test.go | 2 +- .../consumer_test.go | 26 ++-- .../action.go => consumer/kcl_actions.go | 2 +- .../kclreader.go => consumer/kcl_reader.go | 95 +++++++++++---- .../kcl_reader_config.go | 6 +- .../kinesis_reader.go | 88 +++++++++++--- .../kinesis_reader_test.go | 21 ++-- .../readcloser.go => consumer/read_closer.go | 2 +- .../shard_iterator.go | 2 +- .../shard_iterator_test.go | 2 +- consumer/shared_types.go | 30 +++++ errors.go | 12 ++ kinetic_integration_test.go | 52 +++++--- listener/kinesisreaderconfig.go | 54 --------- listener/listenerconfig.go | 47 -------- listener/listenerconfig_test.go | 60 ---------- .../{firehosewriter.go => firehose_writer.go} | 106 ++++++++++++---- producer/firehosewriterconfig.go | 62 ---------- .../{kinesiswriter.go => kinesis_writer.go} | 102 ++++++++++++---- producer/kinesiswriterconfig.go | 70 ----------- producer/producer.go | 110 +++++++++++++++-- producer/{stats.go => producer_stats.go} | 0 .../{stats_test.go => producer_stats_test.go} | 0 producer/producer_test.go | 44 ++++--- producer/producerconfig.go | 85 ------------- producer/producerconfig_test.go | 105 ---------------- producer/{sharedtypes.go => shared_types.go} | 0 testexec/README.md | 2 +- testexec/main.go | 78 ++++++------ 32 files changed, 685 insertions(+), 747 deletions(-) rename listener/listener.go => consumer/consumer.go (75%) rename listener/stats.go => consumer/consumer_stats.go (88%) rename listener/stats_test.go => consumer/consumer_stats_test.go (99%) rename listener/listener_test.go => consumer/consumer_test.go (94%) rename multilang/action.go => consumer/kcl_actions.go (99%) rename listener/kclreader.go => consumer/kcl_reader.go (67%) rename listener/kclreaderconfig.go => consumer/kcl_reader_config.go (95%) rename listener/kinesisreader.go => consumer/kinesis_reader.go (80%) rename listener/kinesisreader_test.go => consumer/kinesis_reader_test.go (77%) rename listener/readcloser.go => consumer/read_closer.go (96%) rename listener/sharditerator.go => consumer/shard_iterator.go (99%) rename listener/sharditerator_test.go => consumer/shard_iterator_test.go (99%) create mode 100644 consumer/shared_types.go delete mode 100644 listener/kinesisreaderconfig.go delete mode 100644 listener/listenerconfig.go delete mode 100644 listener/listenerconfig_test.go rename producer/{firehosewriter.go => firehose_writer.go} (56%) delete mode 100644 producer/firehosewriterconfig.go rename producer/{kinesiswriter.go => kinesis_writer.go} (65%) delete mode 100644 producer/kinesiswriterconfig.go rename producer/{stats.go => producer_stats.go} (100%) rename producer/{stats_test.go => producer_stats_test.go} (100%) delete mode 100644 producer/producerconfig.go delete mode 100644 producer/producerconfig_test.go rename producer/{sharedtypes.go => shared_types.go} (100%) diff --git a/README.md b/README.md index 209fead..99af38c 100644 --- a/README.md +++ b/README.md @@ -53,9 +53,9 @@ p, _ := producer.NewProducer(func(c *producer.Config) { }) // Create a kinetic listener -l, _ := listener.NewListener(func(c *listener.Config) { +l, _ := consumer.NewListener(func(c *consumer.Config) { c.SetAwsConfig(k.Session.Config) - c.SetReader(listener.NewKinesisReader("stream-name", "shard-name")) + c.SetReader(consumer.NewKinesisReader("stream-name", "shard-name")) }) msg, err := l.Retrieve() diff --git a/listener/listener.go b/consumer/consumer.go similarity index 75% rename from listener/listener.go rename to consumer/consumer.go index 86dbe0a..dcc1a37 100644 --- a/listener/listener.go +++ b/consumer/consumer.go @@ -1,4 +1,4 @@ -package listener +package consumer import ( "context" @@ -11,38 +11,61 @@ import ( "golang.org/x/time/rate" ) -// StreamReader is an interface that abstracts out a stream reader. -type StreamReader interface { - GetRecord(context.Context, MessageHandler) (int, int, error) - GetRecords(context.Context, MessageHandler) (int, int, error) +// consumerOptions is used to hold all of the configurable settings of a Listener object. +type consumerOptions struct { + queueDepth int + concurrency int + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + Stats StatsCollector // stats collection mechanism +} + +func defaultConsumerOptions() *consumerOptions { + return &consumerOptions{ + queueDepth: 10000, + concurrency: 10, + Stats: &NilStatsCollector{}, + } } -// empty is used a as a dummy type for semaphore channels and the pipe of death channel. -type empty struct{} +type ConsumerOptionsFn func(*consumerOptions) error -// MessageProcessor defines the signature of a message handler used by Listen, RetrieveFn and their associated -// *WithContext functions. MessageHandler accepts a WaitGroup so the function can be run as a blocking operation as -// opposed to MessageHandlerAsync. -type MessageProcessor func(*kinetic.Message, *sync.WaitGroup) error +func ConsumerQueueDepth(depth int) ConsumerOptionsFn { + return func(o *consumerOptions) error { + if depth > 0 { + o.queueDepth = depth + return nil + } + return kinetic.ErrInvalidQueueDepth + } +} -// MessageHandler defines the signature of a message handler used by GetRecord() and GetRecords(). MessageHandler -// accepts a WaitGroup so the function can be run as a blocking operation as opposed to MessageHandlerAsync. -type MessageHandler func(*kinetic.Message, *sync.WaitGroup) error +func ConsumerConcurrency(count int) ConsumerOptionsFn { + return func(o *consumerOptions) error { + if count > 0 { + o.concurrency = count + return nil + } + return kinetic.ErrInvalidConcurrency + } +} -// MessageHandlerAsync defines the signature of a message handler used by GetRecord() and GetRecords(). -// MessageHandlerAsync is meant to be run asynchronously. -type MessageHandlerAsync func(*kinetic.Message) error +func ConsumerLogLevel(ll aws.LogLevelType) ConsumerOptionsFn { + return func(o *consumerOptions) error { + o.logLevel = ll & 0xffff0000 + return nil + } +} -// listenerOptions is used to hold all of the configurable settings of a Listener object. -type listenerOptions struct { - queueDepth int - concurrency int - Stats StatsCollector +func ConsumerStatsCollector(sc StatsCollector) ConsumerOptionsFn { + return func(o *consumerOptions) error { + o.Stats = sc + return nil + } } // Listener polls the StreamReader for messages. -type Listener struct { - *listenerOptions +type Consumer struct { + *consumerOptions *kinetic.LogHelper reader StreamReader txnCountRateLimiter *rate.Limiter @@ -55,16 +78,16 @@ type Listener struct { } // NewListener creates a new Listener object for retrieving and listening to message(s) on a StreamReader. -func NewListener(c *aws.Config, r StreamReader, fn ...func(*Config)) (*Listener, error) { - cfg := NewConfig(c) - for _, f := range fn { - f(cfg) +func NewConsumer(c *aws.Config, r StreamReader, optionFns ...ConsumerOptionsFn) (*Consumer, error) { + consumerOptions := defaultConsumerOptions() + for _, optionFn := range optionFns { + optionFn(consumerOptions) } - return &Listener{ - listenerOptions: cfg.listenerOptions, + return &Consumer{ + consumerOptions: consumerOptions, LogHelper: &kinetic.LogHelper{ - LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, + LogLevel: consumerOptions.logLevel, + Logger: c.Logger, }, reader: r, }, nil @@ -72,7 +95,7 @@ func NewListener(c *aws.Config, r StreamReader, fn ...func(*Config)) (*Listener, // startConsuming will initialize the message channel and set consuming to true if there is not already another consume // loop running. -func (l *Listener) startConsuming() bool { +func (l *Consumer) startConsuming() bool { l.consumingMu.Lock() defer l.consumingMu.Unlock() if !l.consuming { @@ -87,7 +110,7 @@ func (l *Listener) startConsuming() bool { // shouldConsume is a convenience function that allows functions to break their loops if the context receives a // cancellation or a pipe of death. -func (l *Listener) shouldConsume(ctx context.Context) (bool, error) { +func (l *Consumer) shouldConsume(ctx context.Context) (bool, error) { select { case <-l.pipeOfDeath: return false, kinetic.ErrPipeOfDeath @@ -99,7 +122,7 @@ func (l *Listener) shouldConsume(ctx context.Context) (bool, error) { } // stopConsuming handles any cleanup after consuming has stopped. -func (l *Listener) stopConsuming() { +func (l *Consumer) stopConsuming() { l.consumingMu.Lock() defer l.consumingMu.Unlock() if l.consuming && l.messages != nil { @@ -111,7 +134,7 @@ func (l *Listener) stopConsuming() { l.consuming = false } -func (l *Listener) enqueueSingle(ctx context.Context) (int, int, error) { +func (l *Consumer) enqueueSingle(ctx context.Context) (int, int, error) { n, m, err := l.reader.GetRecord(ctx, func(msg *kinetic.Message, wg *sync.WaitGroup) error { defer wg.Done() l.messages <- msg @@ -125,7 +148,7 @@ func (l *Listener) enqueueSingle(ctx context.Context) (int, int, error) { return n, m, nil } -func (l *Listener) enqueueBatch(ctx context.Context) (int, int, error) { +func (l *Consumer) enqueueBatch(ctx context.Context) (int, int, error) { n, m, err := l.reader.GetRecords(ctx, func(msg *kinetic.Message, wg *sync.WaitGroup) error { defer wg.Done() @@ -140,7 +163,7 @@ func (l *Listener) enqueueBatch(ctx context.Context) (int, int, error) { return n, m, nil } -func (l *Listener) handleErrorLogging(err error) { +func (l *Consumer) handleErrorLogging(err error) { switch err := err.(type) { case net.Error: if err.Timeout() { @@ -164,7 +187,7 @@ func (l *Listener) handleErrorLogging(err error) { // RetrieveWithContext waits for a message from the stream and returns the kinetic. Cancellation is supported through // contexts. -func (l *Listener) RetrieveWithContext(ctx context.Context) (*kinetic.Message, error) { +func (l *Consumer) RetrieveWithContext(ctx context.Context) (*kinetic.Message, error) { if !l.startConsuming() { return nil, kinetic.ErrAlreadyConsuming } @@ -189,13 +212,13 @@ func (l *Listener) RetrieveWithContext(ctx context.Context) (*kinetic.Message, e } // Retrieve waits for a message from the stream and returns the value -func (l *Listener) Retrieve() (*kinetic.Message, error) { +func (l *Consumer) Retrieve() (*kinetic.Message, error) { return l.RetrieveWithContext(context.TODO()) } // RetrieveFnWithContext retrieves a message from the stream and dispatches it to the supplied function. RetrieveFn // will wait until the function completes. Cancellation is supported through context. -func (l *Listener) RetrieveFnWithContext(ctx context.Context, fn MessageProcessor) error { +func (l *Consumer) RetrieveFnWithContext(ctx context.Context, fn MessageProcessor) error { msg, err := l.RetrieveWithContext(ctx) if err != nil { return err @@ -216,12 +239,12 @@ func (l *Listener) RetrieveFnWithContext(ctx context.Context, fn MessageProcesso // RetrieveFn retrieves a message from the stream and dispatches it to the supplied function. RetrieveFn will wait // until the function completes. -func (l *Listener) RetrieveFn(fn MessageProcessor) error { +func (l *Consumer) RetrieveFn(fn MessageProcessor) error { return l.RetrieveFnWithContext(context.TODO(), fn) } // consume calls getRecords with configured batch size in a loop until the listener is stopped. -func (l *Listener) consume(ctx context.Context) { +func (l *Consumer) consume(ctx context.Context) { // We need to run startConsuming to make sure that we are okay and ready to start consuming. This is mainly to // avoid a race condition where Listen() will attempt to read the messages channel prior to consume() // initializing it. We can then launch a goroutine to handle the actual consume operation. @@ -277,7 +300,7 @@ func (l *Listener) consume(ctx context.Context) { // ListenWithContext listens and delivers message to the supplied function. Upon cancellation, Listen will stop the // consumer loop and wait until the messages channel is closed and all messages are delivered. -func (l *Listener) ListenWithContext(ctx context.Context, fn MessageProcessor) { +func (l *Consumer) ListenWithContext(ctx context.Context, fn MessageProcessor) { l.consume(ctx) var wg sync.WaitGroup defer wg.Wait() @@ -314,6 +337,6 @@ func (l *Listener) ListenWithContext(ctx context.Context, fn MessageProcessor) { } // Listen listens and delivers message to the supplied function. -func (l *Listener) Listen(fn MessageProcessor) { +func (l *Consumer) Listen(fn MessageProcessor) { l.ListenWithContext(context.TODO(), fn) } diff --git a/listener/stats.go b/consumer/consumer_stats.go similarity index 88% rename from listener/stats.go rename to consumer/consumer_stats.go index 3326741..8ea9da7 100644 --- a/listener/stats.go +++ b/consumer/consumer_stats.go @@ -1,4 +1,4 @@ -package listener +package consumer import ( "log" @@ -29,29 +29,29 @@ type StatsCollector interface { type NilStatsCollector struct{} // AddConsumed records a count of the number of messages received from AWS -// Kinesis by the listener. +// Kinesis by the consumer. func (nsc *NilStatsCollector) AddConsumed(int) {} // AddDelivered records a count of the number of messages delivered to the -// application by the listener. +// application by the consumer. func (nsc *NilStatsCollector) AddDelivered(int) {} // AddProcessed records a count of the number of messages processed by the -// application by the listener. This is based on a WaitGroup that is sent to +// application by the consumer. This is based on a WaitGroup that is sent to // the RetrieveFn and Listen functions. Retrieve does not count processed // messages. func (nsc *NilStatsCollector) AddProcessed(int) {} // AddBatchSize records a count of the number of messages returned by -// GetRecords in the listener. +// GetRecords in the consumer. func (nsc *NilStatsCollector) AddBatchSize(int) {} // AddGetRecordsCalled records the number of times the GetRecords API was called -// by the listener. +// by the consumer. func (nsc *NilStatsCollector) AddGetRecordsCalled(int) {} // AddProvisionedThroughputExceeded records the number of times the GetRecords -// API returned a ErrCodeProvisionedThroughputExceededException by the listener. +// API returned a ErrCodeProvisionedThroughputExceededException by the consumer. func (nsc *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} // AddGetRecordsTimeout records the number of times the GetRecords API timed out @@ -82,18 +82,18 @@ func (nsc *NilStatsCollector) AddGetRecordsUnmarshalDuration(time.Duration) {} // Metric names to be exported const ( - MetricsConsumed = "kinetic.listener.consumed" - MetricsDelivered = "kinetic.listener.delivered" - MetricsProcessed = "kinetic.listener.processed" - MetricsBatchSize = "kinetic.listener.batchsize" - MetricsSent = "kinetic.listener.sent" - MetricsProvisionedThroughputExceeded = "kinetic.listener.provisionedthroughputexceeded" - MetricsGetRecordsTimeout = "kinetic.listener.getrecords.timeout" - MetricsGetRecordsReadTimeout = "kinetic.listener.getrecords.readtimeout" - MetricsProcessedDuration = "kinetic.listener.processed.duration" - MetricsGetRecordsDuration = "kinetic.listener.getrecords.duration" - MetricsGetRecordsReadResponseDuration = "kinetic.listener.getrecords.readresponse.duration" - MetricsGetRecordsUnmarshalDuration = "kinetic.listener.getrecords.unmarshal.duration" + MetricsConsumed = "kinetic.consumer.consumed" + MetricsDelivered = "kinetic.consumer.delivered" + MetricsProcessed = "kinetic.consumer.processed" + MetricsBatchSize = "kinetic.consumer.batchsize" + MetricsSent = "kinetic.consumer.sent" + MetricsProvisionedThroughputExceeded = "kinetic.consumer.provisionedthroughputexceeded" + MetricsGetRecordsTimeout = "kinetic.consumer.getrecords.timeout" + MetricsGetRecordsReadTimeout = "kinetic.consumer.getrecords.readtimeout" + MetricsProcessedDuration = "kinetic.consumer.processed.duration" + MetricsGetRecordsDuration = "kinetic.consumer.getrecords.duration" + MetricsGetRecordsReadResponseDuration = "kinetic.consumer.getrecords.readresponse.duration" + MetricsGetRecordsUnmarshalDuration = "kinetic.consumer.getrecords.unmarshal.duration" ) // DefaultStatsCollector is a type that implements the listener's StatsCollector interface using the @@ -132,19 +132,19 @@ func NewDefaultStatsCollector(r metrics.Registry) *DefaultStatsCollector { } // AddConsumed records a count of the number of messages received from AWS -// Kinesis by the listener. +// Kinesis by the consumer. func (dsc *DefaultStatsCollector) AddConsumed(count int) { dsc.Consumed.Inc(int64(count)) } // AddDelivered records a count of the number of messages delivered to the -// application by the listener. +// application by the consumer. func (dsc *DefaultStatsCollector) AddDelivered(count int) { dsc.Delivered.Inc(int64(count)) } // AddProcessed records a count of the number of messages processed by the -// application by the listener. This is based on a WaitGroup that is sent to +// application by the consumer. This is based on a WaitGroup that is sent to // the RetrieveFn and Listen functions. Retrieve does not count processed // messages. func (dsc *DefaultStatsCollector) AddProcessed(count int) { @@ -152,19 +152,19 @@ func (dsc *DefaultStatsCollector) AddProcessed(count int) { } // AddBatchSize records a count of the number of messages returned by -// GetRecords in the listener. +// GetRecords in the consumer. func (dsc *DefaultStatsCollector) AddBatchSize(count int) { dsc.BatchSize.Inc(int64(count)) } // AddGetRecordsCalled records the number of times the GetRecords API was called -// by the listener. +// by the consumer. func (dsc *DefaultStatsCollector) AddGetRecordsCalled(count int) { dsc.GetRecordsCalled.Inc(int64(count)) } // AddProvisionedThroughputExceeded records the number of times the GetRecords -// API returned a ErrCodeProvisionedThroughputExceededException by the listener. +// API returned a ErrCodeProvisionedThroughputExceededException by the consumer. func (dsc *DefaultStatsCollector) AddProvisionedThroughputExceeded(count int) { dsc.ProvisionedThroughputExceeded.Inc(int64(count)) } diff --git a/listener/stats_test.go b/consumer/consumer_stats_test.go similarity index 99% rename from listener/stats_test.go rename to consumer/consumer_stats_test.go index 9932cbd..d5f365c 100644 --- a/listener/stats_test.go +++ b/consumer/consumer_stats_test.go @@ -1,4 +1,4 @@ -package listener +package consumer import ( "math/rand" diff --git a/listener/listener_test.go b/consumer/consumer_test.go similarity index 94% rename from listener/listener_test.go rename to consumer/consumer_test.go index 43757d5..4e7ebc8 100644 --- a/listener/listener_test.go +++ b/consumer/consumer_test.go @@ -1,4 +1,4 @@ -package listener +package consumer import ( . "github.com/smartystreets/goconvey/convey" @@ -17,9 +17,10 @@ import ( "github.com/aws/aws-sdk-go/service/kinesis" "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/consumer" ) -func putRecord(l *Listener, b []byte) (*string, error) { +func putRecord(l *Consumer, b []byte) (*string, error) { resp, err := l.reader.(*KinesisReader).client.PutRecord(&kinesis.PutRecordInput{ Data: b, PartitionKey: aws.String("dummy"), @@ -54,16 +55,21 @@ func TestListener(t *testing.T) { So(k.Session, ShouldNotBeNil) So(k.Session.Config, ShouldNotBeNil) - r, err := NewKinesisReader(k.Session.Config, stream, shards[0], func(krc *KinesisReaderConfig) { - krc.SetResponseReadTimeout(time.Second) - }) + r, err := NewKinesisReader(k.Session.Config, stream, shards[0], + consumer.KinesisReaderBatchSize(10000), + //consumer.KinesisReaderShardIterator(), + consumer.KinesisReaderResponseReadTimeout(time.Second), + consumer.KinesisReaderLogLevel(aws.LogOff), + ) So(r, ShouldNotBeNil) So(err, ShouldBeNil) - l, err := NewListener(k.Session.Config, r, func(c *Config) { - c.SetQueueDepth(10) - c.SetConcurrency(10) - }) + l, err := NewConsumer(k.Session.Config, r, + consumer.ConsumerQueueDepth(10), + consumer.ConsumerConcurrency(10), + consumer.ConsumerLogLevel(aws.LogOff), + //consumer.ConsumerStatsCollector(lsc), + ) So(l, ShouldNotBeNil) So(err, ShouldBeNil) @@ -201,7 +207,7 @@ func TestListener(t *testing.T) { wg.Wait() }) - // TODO: Move this test to kinesisreader_test.go + // TODO: Move this test to kinesis_reader_test.go Convey("check that throttle mechanism prevents more than 5 calls to get records", func() { start := time.Now() secs := []float64{} diff --git a/multilang/action.go b/consumer/kcl_actions.go similarity index 99% rename from multilang/action.go rename to consumer/kcl_actions.go index be85dc6..5e4bd56 100644 --- a/multilang/action.go +++ b/consumer/kcl_actions.go @@ -1,4 +1,4 @@ -package multilang +package consumer import ( "encoding/base64" diff --git a/listener/kclreader.go b/consumer/kcl_reader.go similarity index 67% rename from listener/kclreader.go rename to consumer/kcl_reader.go index e10d90d..b733ef4 100644 --- a/listener/kclreader.go +++ b/consumer/kcl_reader.go @@ -1,4 +1,4 @@ -package listener +package consumer import ( "bufio" @@ -12,21 +12,67 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/multilang" ) type kclReaderOptions struct { onInitCallbackFn func() error onCheckpointCallbackFn func() error onShutdownCallbackFn func() error - Stats StatsCollector + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + stats StatsCollector // stats collection mechanism +} + +func defaultKlcReaderOptions() *kclReaderOptions { + return &kclReaderOptions{ + onInitCallbackFn: func() error { return nil }, + onCheckpointCallbackFn: func() error { return nil }, + onShutdownCallbackFn: func() error { return nil }, + logLevel: aws.LogOff, + stats: &NilStatsCollector{}, + } +} + +type KlcReaderOptionsFn func(*kclReaderOptions) error + +func KlcReaderOnInitCallbackFn(fn func() error) KlcReaderOptionsFn { + return func(o *kclReaderOptions) error { + o.onInitCallbackFn = fn + return nil + } +} + +func KlcReaderOnCheckpointCallbackFn(fn func() error) KlcReaderOptionsFn { + return func(o *kclReaderOptions) error { + o.onCheckpointCallbackFn = fn + return nil + } +} + +func KlcReaderOnShutdownCallbackFn(fn func() error) KlcReaderOptionsFn { + return func(o *kclReaderOptions) error { + o.onShutdownCallbackFn = fn + return nil + } +} + +func KlcReaderLogLevel(ll aws.LogLevelType) KlcReaderOptionsFn { + return func(o *kclReaderOptions) error { + o.logLevel = ll + return nil + } +} + +func KlcReaderStatsCollector(sc StatsCollector) KlcReaderOptionsFn { + return func(o *kclReaderOptions) error { + o.stats = sc + return nil + } } // KclReader handles the KCL Multilang Protocol to read records from KCL type KclReader struct { *kclReaderOptions *kinetic.LogHelper - throttleSem chan empty pipeOfDeath chan empty scanner *bufio.Scanner reader *bufio.Reader @@ -34,19 +80,18 @@ type KclReader struct { } // NewKclReader creates a new stream reader to read records from KCL -func NewKclReader(c *aws.Config, fn ...func(*KclReaderConfig)) (*KclReader, error) { - cfg := NewKclReaderConfig(c) - for _, f := range fn { - f(cfg) +func NewKclReader(c *aws.Config, optionFns ...KlcReaderOptionsFn) (*KclReader, error) { + kclReaderOptions := defaultKlcReaderOptions() + for _, optionFn := range optionFns { + optionFn(kclReaderOptions) } return &KclReader{ - kclReaderOptions: cfg.kclReaderOptions, + msgBuffer: []kinetic.Message{}, + kclReaderOptions: kclReaderOptions, LogHelper: &kinetic.LogHelper{ - LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, + LogLevel: kclReaderOptions.logLevel, + Logger: c.Logger, }, - throttleSem: make(chan empty, 5), - msgBuffer: []kinetic.Message{}, }, nil } @@ -64,7 +109,7 @@ func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, int, batchSize = int(math.Min(float64(len(r.msgBuffer)), float64(numRecords))) } } - r.Stats.AddBatchSize(batchSize) + r.stats.AddBatchSize(batchSize) // TODO: Define the payloadSize var payloadSize int @@ -75,13 +120,13 @@ func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, int, wg.Add(1) go fn(&r.msgBuffer[0], &wg) r.msgBuffer = r.msgBuffer[1:] - r.Stats.AddConsumed(1) + r.stats.AddConsumed(1) } wg.Wait() // Send an acknowledgement that the 'ProcessRecords' message was received/processed if len(r.msgBuffer) == 0 { - err := r.sendMessage(multilang.NewStatusMessage(multilang.PROCESSRECORDS)) + err := r.sendMessage(NewStatusMessage(PROCESSRECORDS)) if err != nil { r.LogError(err) return batchSize, payloadSize, err @@ -91,7 +136,7 @@ func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, int, return batchSize, payloadSize, nil } -func (r *KclReader) getAction() (*multilang.ActionMessage, error) { +func (r *KclReader) getAction() (*ActionMessage, error) { buffer := &bytes.Buffer{} for { line, isPrefix, err := r.reader.ReadLine() @@ -104,7 +149,7 @@ func (r *KclReader) getAction() (*multilang.ActionMessage, error) { } } - actionMsg := &multilang.ActionMessage{} + actionMsg := &ActionMessage{} err := json.Unmarshal(buffer.Bytes(), actionMsg) if err != nil { fmt.Fprintf(os.Stderr, "Could not understand line read from input: %s\n", buffer.String()) @@ -126,16 +171,16 @@ func (r *KclReader) processAction() error { } switch actionMessage.Action { - case multilang.INITIALIZE: + case INITIALIZE: r.onInit() - r.sendMessage(multilang.NewStatusMessage(multilang.INITIALIZE)) - case multilang.CHECKPOINT: + r.sendMessage(NewStatusMessage(INITIALIZE)) + case CHECKPOINT: r.onCheckpoint() - r.sendMessage(multilang.NewStatusMessage(multilang.CHECKPOINT)) - case multilang.SHUTDOWN: + r.sendMessage(NewStatusMessage(CHECKPOINT)) + case SHUTDOWN: r.onShutdown() - r.sendMessage(multilang.NewStatusMessage(multilang.SHUTDOWN)) - case multilang.PROCESSRECORDS: + r.sendMessage(NewStatusMessage(SHUTDOWN)) + case PROCESSRECORDS: go func() error { for _, msg := range actionMessage.Records { r.msgBuffer = append(r.msgBuffer, *msg.ToMessage()) diff --git a/listener/kclreaderconfig.go b/consumer/kcl_reader_config.go similarity index 95% rename from listener/kclreaderconfig.go rename to consumer/kcl_reader_config.go index a4b0f43..c9406be 100644 --- a/listener/kclreaderconfig.go +++ b/consumer/kcl_reader_config.go @@ -1,4 +1,4 @@ -package listener +package consumer import ( "github.com/aws/aws-sdk-go/aws" @@ -16,7 +16,7 @@ func NewKclReaderConfig(cfg *aws.Config) *KclReaderConfig { return &KclReaderConfig{ AwsConfig: cfg, kclReaderOptions: &kclReaderOptions{ - Stats: &NilStatsCollector{}, + stats: &NilStatsCollector{}, }, LogLevel: *cfg.LogLevel, } @@ -42,7 +42,7 @@ func (c *KclReaderConfig) SetOnShutdownCallbackFn(fn func() error) { // SetStatsCollector configures a listener to handle listener metrics. func (c *KclReaderConfig) SetStatsCollector(stats StatsCollector) { - c.Stats = stats + c.stats = stats } // SetLogLevel configures both the SDK and Kinetic log levels. diff --git a/listener/kinesisreader.go b/consumer/kinesis_reader.go similarity index 80% rename from listener/kinesisreader.go rename to consumer/kinesis_reader.go index d1fedd1..94c6caf 100644 --- a/listener/kinesisreader.go +++ b/consumer/kinesis_reader.go @@ -1,4 +1,4 @@ -package listener +package consumer import ( "context" @@ -16,12 +16,66 @@ import ( "github.com/rewardStyle/kinetic" ) +const ( + kinesisReaderBatchSize = 10000 +) + // kinesisReaderOptions is used to hold all of the configurable settings of a KinesisReader. type kinesisReaderOptions struct { batchSize int shardIterator *ShardIterator responseReadTimeout time.Duration - Stats StatsCollector + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + stats StatsCollector // stats collection mechanism +} + +func defaultKinesisReaderOptions() *kinesisReaderOptions { + return &kinesisReaderOptions{ + batchSize: kinesisReaderBatchSize, + shardIterator: NewShardIterator(), + responseReadTimeout: time.Second, + stats: &NilStatsCollector{}, + } +} + +type KinesisReaderOptionsFn func(*kinesisReaderOptions) error + +func KinesisReaderBatchSize(size int) KinesisReaderOptionsFn { + return func(o *kinesisReaderOptions) error { + if size > 0 && size <= kinesisReaderBatchSize { + o.batchSize = size + return nil + } + return kinetic.ErrInvalidBatchSize + } +} + +func KinesisReaderShardIterator(shardIterator *ShardIterator) KinesisReaderOptionsFn { + return func(o *kinesisReaderOptions) error { + o.shardIterator = shardIterator + return nil + } +} + +func KinesisReaderResponseReadTimeout(timeout time.Duration) KinesisReaderOptionsFn { + return func(o *kinesisReaderOptions) error { + o.responseReadTimeout = timeout + return nil + } +} + +func KinesisReaderLogLevel(ll aws.LogLevelType) KinesisReaderOptionsFn { + return func(o *kinesisReaderOptions) error { + o.logLevel = ll & 0xffff0000 + return nil + } +} + +func KinesisReaderStatsCollector(sc StatsCollector) KinesisReaderOptionsFn { + return func(o *kinesisReaderOptions) error { + o.stats = sc + return nil + } } // KinesisReader handles the API to read records from Kinesis. @@ -37,20 +91,20 @@ type KinesisReader struct { // NewKinesisReader creates a new KinesisReader object which implements the StreamReader interface to read records from // Kinesis. -func NewKinesisReader(c *aws.Config, stream string, shard string, fn ...func(*KinesisReaderConfig)) (*KinesisReader, error) { - cfg := NewKinesisReaderConfig(c) - for _, f := range fn { - f(cfg) +func NewKinesisReader(c *aws.Config, stream string, shard string, optionFns ...KinesisReaderOptionsFn) (*KinesisReader, error) { + kinesisReaderOptions := defaultKinesisReaderOptions() + for _, optionFn := range optionFns { + optionFn(kinesisReaderOptions) } - sess, err := session.NewSession(cfg.AwsConfig) + sess, err := session.NewSession(c) if err != nil { return nil, err } return &KinesisReader{ - kinesisReaderOptions: cfg.kinesisReaderOptions, + kinesisReaderOptions: kinesisReaderOptions, LogHelper: &kinetic.LogHelper{ - LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, + LogLevel: kinesisReaderOptions.logLevel, + Logger: c.Logger, }, stream: stream, shard: shard, @@ -167,7 +221,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch req.HTTPResponse.Body = &ReadCloserWrapper{ ReadCloser: req.HTTPResponse.Body, OnCloseFn: func() { - r.Stats.AddGetRecordsReadResponseDuration(time.Since(startReadTime)) + r.stats.AddGetRecordsReadResponseDuration(time.Since(startReadTime)) r.LogDebug("Finished GetRecords body read, took", time.Since(start)) startUnmarshalTime = time.Now() }, @@ -177,24 +231,24 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch var payloadSize int req.Handlers.Unmarshal.PushBack(func(req *request.Request) { payloadSize += int(req.HTTPRequest.ContentLength) - r.Stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) + r.stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) r.LogDebug("Finished GetRecords Unmarshal, took", time.Since(start)) }) // Send the GetRecords request r.LogDebug("Starting GetRecords Build/Sign request, took", time.Since(start)) - r.Stats.AddGetRecordsCalled(1) + r.stats.AddGetRecordsCalled(1) if err := req.Send(); err != nil { r.LogError("Error getting records:", err) switch err.(awserr.Error).Code() { case kinesis.ErrCodeProvisionedThroughputExceededException: - r.Stats.AddProvisionedThroughputExceeded(1) + r.stats.AddProvisionedThroughputExceeded(1) default: r.LogDebug("Received AWS error:", err.Error()) } return 0, 0, err } - r.Stats.AddGetRecordsDuration(time.Since(start)) + r.stats.AddGetRecordsDuration(time.Since(start)) // Process Records r.LogDebug(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), r.shard, time.Since(start))) @@ -202,7 +256,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch return 0, 0, kinetic.ErrNilGetRecordsResponse } delivered := 0 - r.Stats.AddBatchSize(len(resp.Records)) + r.stats.AddBatchSize(len(resp.Records)) for _, record := range resp.Records { if record != nil { // Allow (only) a pipeOfDeath to trigger an instance shutdown of the loop to deliver messages. @@ -218,7 +272,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch go fn(kinetic.FromRecord(record), &wg) wg.Wait() delivered++ - r.Stats.AddConsumed(1) + r.stats.AddConsumed(1) if record.SequenceNumber != nil { // We can safely ignore if this call returns error, as if we somehow receive an // empty sequence number from AWS, we will simply not set it. At worst, this diff --git a/listener/kinesisreader_test.go b/consumer/kinesis_reader_test.go similarity index 77% rename from listener/kinesisreader_test.go rename to consumer/kinesis_reader_test.go index 9e87549..78a5dd8 100644 --- a/listener/kinesisreader_test.go +++ b/consumer/kinesis_reader_test.go @@ -1,4 +1,4 @@ -package listener +package consumer import ( . "github.com/smartystreets/goconvey/convey" @@ -9,6 +9,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/rewardStyle/kinetic" + "github.com/rewardStyle/kinetic/consumer" ) func TestNewKinesisReader(t *testing.T) { @@ -35,7 +36,7 @@ func TestNewKinesisReader(t *testing.T) { So(r.batchSize, ShouldEqual, 10000) So(r.shardIterator, ShouldNotBeNil) So(r.responseReadTimeout, ShouldEqual, time.Second) - So(r.Stats, ShouldNotBeNil) + So(r.stats, ShouldNotBeNil) }) Convey("check that we can create a new KinesisReader with configured values", func() { @@ -44,20 +45,20 @@ func TestNewKinesisReader(t *testing.T) { logLevel := aws.LogDebug | aws.LogDebugWithSigning | kinetic.LogDebug shardIterator := NewShardIterator() myStatsCollector := &NilStatsCollector{} - r, err := NewKinesisReader(k.Session.Config, stream, shard, func(krc *KinesisReaderConfig) { - krc.SetBatchSize(batchSize) - krc.SetResponseReadTimeout(respReadTimeout) - krc.SetLogLevel(logLevel) - krc.SetInitialShardIterator(shardIterator) - krc.SetStatsCollector(myStatsCollector) - }) + r, err := NewKinesisReader(k.Session.Config, stream, shard, + consumer.KinesisReaderBatchSize(batchSize), + consumer.KinesisReaderShardIterator(shardIterator), + consumer.KinesisReaderResponseReadTimeout(respReadTimeout), + consumer.KinesisReaderLogLevel(logLevel), + consumer.KinesisReaderStatsCollector(myStatsCollector), + ) So(r, ShouldNotBeNil) So(err, ShouldBeNil) So(r.batchSize, ShouldEqual, batchSize) So(r.responseReadTimeout, ShouldEqual, respReadTimeout) So(r.LogLevel.AtLeast(kinetic.LogDebug), ShouldBeTrue) So(r.shardIterator, ShouldEqual, shardIterator) - So(r.Stats, ShouldEqual, myStatsCollector) + So(r.stats, ShouldEqual, myStatsCollector) }) }) } diff --git a/listener/readcloser.go b/consumer/read_closer.go similarity index 96% rename from listener/readcloser.go rename to consumer/read_closer.go index 57b011b..3d4ffb2 100644 --- a/listener/readcloser.go +++ b/consumer/read_closer.go @@ -1,4 +1,4 @@ -package listener +package consumer import ( "io" diff --git a/listener/sharditerator.go b/consumer/shard_iterator.go similarity index 99% rename from listener/sharditerator.go rename to consumer/shard_iterator.go index 3e017ee..202c96f 100644 --- a/listener/sharditerator.go +++ b/consumer/shard_iterator.go @@ -1,4 +1,4 @@ -package listener +package consumer import ( "time" diff --git a/listener/sharditerator_test.go b/consumer/shard_iterator_test.go similarity index 99% rename from listener/sharditerator_test.go rename to consumer/shard_iterator_test.go index 3efa843..c46345b 100644 --- a/listener/sharditerator_test.go +++ b/consumer/shard_iterator_test.go @@ -1,4 +1,4 @@ -package listener +package consumer import ( "github.com/aws/aws-sdk-go/aws" diff --git a/consumer/shared_types.go b/consumer/shared_types.go new file mode 100644 index 0000000..d1a07be --- /dev/null +++ b/consumer/shared_types.go @@ -0,0 +1,30 @@ +package consumer + +import ( + "context" + "sync" + + "github.com/rewardStyle/kinetic" +) + +// StreamReader is an interface that abstracts out a stream reader. +type StreamReader interface { + GetRecord(context.Context, MessageHandler) (int, int, error) + GetRecords(context.Context, MessageHandler) (int, int, error) +} + +// empty is used a as a dummy type for semaphore channels and the pipe of death channel. +type empty struct{} + +// MessageProcessor defines the signature of a message handler used by Listen, RetrieveFn and their associated +// *WithContext functions. MessageHandler accepts a WaitGroup so the function can be run as a blocking operation as +// opposed to MessageHandlerAsync. +type MessageProcessor func(*kinetic.Message, *sync.WaitGroup) error + +// MessageHandler defines the signature of a message handler used by GetRecord() and GetRecords(). MessageHandler +// accepts a WaitGroup so the function can be run as a blocking operation as opposed to MessageHandlerAsync. +type MessageHandler func(*kinetic.Message, *sync.WaitGroup) error + +// MessageHandlerAsync defines the signature of a message handler used by GetRecord() and GetRecords(). +// MessageHandlerAsync is meant to be run asynchronously. +type MessageHandlerAsync func(*kinetic.Message) error diff --git a/errors.go b/errors.go index bce6888..348f512 100644 --- a/errors.go +++ b/errors.go @@ -92,4 +92,16 @@ var ( // ErrInvalidMsgSizeRateLimit is returned when a MsgCountSizesLimit is configured incorrectly ErrInvalidMsgSizeRateLimit = errors.New("Invalid Message Size Rate Limit") + + // ErrInvalidThroughputMultiplier is returned when a ThroughputMultiplier is configured incorrectly + ErrInvalidThroughputMultiplier = errors.New("Invalid Throughput Multiplier") + + // ErrInvalidQueueDepth is returned when the Queue Depth is configured incorrectly + ErrInvalidQueueDepth = errors.New("Invalid Queue Depth") + + // ErrInvalidMaxRetryAttempts is returned when the Max Retry Attempts is configured incorrectly + ErrInvalidMaxRetryAttempts = errors.New("Invalid Max Retry Attempts") + + // ErrInvalidConcurrency is returned when the concurrency value is configured incorrectly + ErrInvalidConcurrency = errors.New("Invalid concurrency") ) diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index 5036c5e..4e1ee4b 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -12,9 +12,10 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" - "github.com/rewardStyle/kinetic/listener" + "github.com/rewardStyle/kinetic/consumer" "github.com/rewardStyle/kinetic/producer" "github.com/stretchr/testify/assert" + "github.com/rewardStyle/kinetic" ) type TestMessage struct { @@ -129,36 +130,53 @@ func TestKineticIntegration(t *testing.T) { log.Printf("Shard Name: %s\n", shards[0]) // Create a new kinesis stream writer - w, err := producer.NewKinesisWriter(k.Session.Config, stream, func(kwc *producer.KinesisWriterConfig) { - kwc.SetLogLevel(aws.LogDebug) - kwc.SetResponseReadTimeout(time.Second) - }) + w, err := producer.NewKinesisWriter(k.Session.Config, stream, + producer.KinesisWriterResponseReadTimeout(time.Second), + producer.KinesisWriterMsgCountRateLimit(1000), + producer.KinesisWriterMsgSizeRateLimit(1000000), + producer.KinesisWriterLogLevel(kinetic.LogDebug), + ) if err != nil { log.Fatalf("Unable to create a new kinesis stream writer due to: %v\n", err) } // Create a new kinetic producer - p, err := producer.NewProducer(k.Session.Config, w, func(c *producer.Config) { - c.SetBatchSize(5) - c.SetBatchTimeout(1000 * time.Millisecond) - }) + p, err := producer.NewProducer(k.Session.Config, w, + producer.ProducerBatchSize(5), + producer.ProducerBatchTimeout(time.Second), + producer.ProducerMaxRetryAttempts(3), + producer.ProducerQueueDepth(10000), + producer.ProducerConcurrency(3), + producer.ProducerShardCheckFrequency(time.Minute), + producer.ProducerDataSpillFn(func(msg *kinetic.Message) error { + //log.Printf("Message was dropped: [%s]\n", string(msg.Data)) + return nil + }), + producer.ProducerLogLevel(aws.LogOff), + //producer.ProducerStatsCollector(), + ) assert.NotNil(t, p) assert.Nil(t, err) assert.NotNil(t, k.Session) assert.NotNil(t, k.Session.Config) - r, err := listener.NewKinesisReader(k.Session.Config, stream, shards[0], - func(krc *listener.KinesisReaderConfig) { - krc.SetResponseReadTimeout(time.Second) - }) + r, err := consumer.NewKinesisReader(k.Session.Config, stream, shards[0], + //consumer.KinesisReaderBatchSize(), + //consumer.KinesisReaderShardIterator(), + consumer.KinesisReaderResponseReadTimeout(time.Second), + //consumer.KinesisReaderLogLevel(), + //consumer.KinesisReaderStatsCollector(), + ) assert.NotNil(t, r) assert.NoError(t, err) // Create a new kinetic listener - l, err := listener.NewListener(k.Session.Config, r, func(c *listener.Config) { - c.SetQueueDepth(20) - c.SetConcurrency(10) - }) + l, err := consumer.NewConsumer(k.Session.Config, r, + consumer.ConsumerQueueDepth(20), + consumer.ConsumerConcurrency(10), + consumer.ConsumerLogLevel(aws.LogOff), + //consumer.ConsumerStatsCollector(), + ) assert.NotNil(t, l) assert.Nil(t, err) diff --git a/listener/kinesisreaderconfig.go b/listener/kinesisreaderconfig.go deleted file mode 100644 index 7531313..0000000 --- a/listener/kinesisreaderconfig.go +++ /dev/null @@ -1,54 +0,0 @@ -package listener - -import ( - "time" - - "github.com/aws/aws-sdk-go/aws" -) - -// KinesisReaderConfig is used to configure a KinesisReader -type KinesisReaderConfig struct { - *kinesisReaderOptions - AwsConfig *aws.Config - LogLevel aws.LogLevelType -} - -// NewKinesisReaderConfig creates a new instance of KinesisReaderConfig -func NewKinesisReaderConfig(cfg *aws.Config) *KinesisReaderConfig { - return &KinesisReaderConfig{ - AwsConfig: cfg, - kinesisReaderOptions: &kinesisReaderOptions{ - batchSize: 10000, - shardIterator: NewShardIterator(), - responseReadTimeout: time.Second, - Stats: &NilStatsCollector{}, - }, - LogLevel: *cfg.LogLevel, - } -} - -// SetBatchSize configures the batch size of the GetRecords call. -func (c *KinesisReaderConfig) SetBatchSize(batchSize int) { - c.batchSize = batchSize -} - -// SetInitialShardIterator configures the settings used to retrieve initial shard iterator via the GetShardIterator -// call. -func (c *KinesisReaderConfig) SetInitialShardIterator(shardIterator *ShardIterator) { - c.shardIterator = shardIterator -} - -// SetResponseReadTimeout configures the time to wait for each successive Read operation on the GetRecords response payload. -func (c *KinesisReaderConfig) SetResponseReadTimeout(timeout time.Duration) { - c.responseReadTimeout = timeout -} - -// SetStatsCollector configures a listener to handle listener metrics. -func (c *KinesisReaderConfig) SetStatsCollector(stats StatsCollector) { - c.Stats = stats -} - -// SetLogLevel configures the log levels for the SDK. -func (c *KinesisReaderConfig) SetLogLevel(logLevel aws.LogLevelType) { - c.LogLevel = logLevel & 0xffff0000 -} diff --git a/listener/listenerconfig.go b/listener/listenerconfig.go deleted file mode 100644 index 8dab8b0..0000000 --- a/listener/listenerconfig.go +++ /dev/null @@ -1,47 +0,0 @@ -package listener - -import ( - "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic" -) - -// Config is used to configure a Listener instance -type Config struct { - *listenerOptions - AwsConfig *aws.Config - LogLevel aws.LogLevelType -} - -// NewConfig creates a new instance of Config -func NewConfig(cfg *aws.Config) *Config { - return &Config{ - AwsConfig: cfg, - listenerOptions: &listenerOptions{ - queueDepth: 10000, - concurrency: 10000, - Stats: &NilStatsCollector{}, - }, - LogLevel: kinetic.LogOff, - } -} - -// SetQueueDepth controls the depth of the listener queue -func (c *Config) SetQueueDepth(queueDepth int) { - c.queueDepth = queueDepth -} - -// SetConcurrency controls the number of goroutines the Listener will spawn to -// process messages. -func (c *Config) SetConcurrency(concurrency int) { - c.concurrency = concurrency -} - -// SetStatsCollector configures a listener to handle listener metrics. -func (c *Config) SetStatsCollector(stats StatsCollector) { - c.Stats = stats -} - -// SetLogLevel configures both the SDK and Kinetic log levels. -func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { - c.LogLevel = logLevel & 0xffff0000 -} diff --git a/listener/listenerconfig_test.go b/listener/listenerconfig_test.go deleted file mode 100644 index 484ee1c..0000000 --- a/listener/listenerconfig_test.go +++ /dev/null @@ -1,60 +0,0 @@ -package listener - -import ( - . "github.com/smartystreets/goconvey/convey" - - "testing" - "time" - - "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic" -) - -type DebugStatsCollector struct{} - -func (l *DebugStatsCollector) AddConsumed(int) {} -func (l *DebugStatsCollector) AddDelivered(int) {} -func (l *DebugStatsCollector) AddProcessed(int) {} -func (l *DebugStatsCollector) AddBatchSize(int) {} -func (l *DebugStatsCollector) AddGetRecordsCalled(int) {} -func (l *DebugStatsCollector) AddProvisionedThroughputExceeded(int) {} -func (l *DebugStatsCollector) AddGetRecordsTimeout(int) {} -func (l *DebugStatsCollector) AddGetRecordsReadTimeout(int) {} -func (l *DebugStatsCollector) AddProcessedDuration(time.Duration) {} -func (l *DebugStatsCollector) AddGetRecordsDuration(time.Duration) {} -func (l *DebugStatsCollector) AddGetRecordsReadResponseDuration(time.Duration) {} -func (l *DebugStatsCollector) AddGetRecordsUnmarshalDuration(time.Duration) {} - -func TestNewConfig(t *testing.T) { - Convey("given a new listener config", t, func() { - k, err := kinetic.New(func(c *kinetic.Config) { - c.SetEndpoint("bogus-endpoint") - }) - So(k, ShouldNotBeNil) - So(err, ShouldBeNil) - config := NewConfig(k.Session.Config) - - Convey("check the default values for its non-zero config", func() { - So(config.queueDepth, ShouldEqual, 10000) - So(config.concurrency, ShouldEqual, 10000) - So(config.Stats, ShouldHaveSameTypeAs, &NilStatsCollector{}) - So(config.LogLevel.Value(), ShouldEqual, kinetic.LogOff) - }) - - Convey("check that we can set both the sdk and kinetic log level", func() { - ll := aws.LogDebug | aws.LogDebugWithSigning | kinetic.LogDebug - config.SetLogLevel(ll) - So(config.LogLevel.AtLeast(kinetic.LogDebug), ShouldBeTrue) - }) - - Convey("check that we can set the concurrency limit", func() { - config.SetConcurrency(50) - So(config.concurrency, ShouldEqual, 50) - }) - - Convey("check that we can configure a stats collector", func() { - config.SetStatsCollector(&DebugStatsCollector{}) - So(config.Stats, ShouldHaveSameTypeAs, &DebugStatsCollector{}) - }) - }) -} diff --git a/producer/firehosewriter.go b/producer/firehose_writer.go similarity index 56% rename from producer/firehosewriter.go rename to producer/firehose_writer.go index 9b86860..bd172de 100644 --- a/producer/firehosewriter.go +++ b/producer/firehose_writer.go @@ -14,40 +14,101 @@ import ( "github.com/rewardStyle/kinetic" ) +const ( + firehoseMsgCountRateLimit = 5000 // AWS Firehose limit of 5000 records/sec + firehoseMsgSizeRateLimit = 5000000 // AWS Firehose limit of 5 MB/sec +) + type firehoseWriterOptions struct { - msgCountRateLimit int // maximum number of records to be sent per second - msgSizeRateLimit int // maximum (transmission) size of records to be sent per second - throughputMultiplier int // integer multiplier to increase firehose throughput rate limits - Stats StatsCollector // stats collection mechanism + msgCountRateLimit int // maximum number of records to be sent per second + msgSizeRateLimit int // maximum (transmission) size of records to be sent per second + throughputMultiplier int // integer multiplier to increase firehose throughput rate limits + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + stats StatsCollector // stats collection mechanism +} + +func defaultFirehoseWriterOptions() *firehoseWriterOptions { + return &firehoseWriterOptions{ + msgCountRateLimit: firehoseMsgCountRateLimit, + msgSizeRateLimit: firehoseMsgSizeRateLimit, + throughputMultiplier: 1, + logLevel: aws.LogOff, + stats: &NilStatsCollector{}, + } +} + +type FireHoseWriterOptionsFn func(*firehoseWriterOptions) error + +func FirehoseWriterMsgCountRateLimit(limit int) FireHoseWriterOptionsFn { + return func(o *firehoseWriterOptions) error { + if limit > 0 && limit <= firehoseMsgCountRateLimit { + o.msgCountRateLimit = limit + return nil + } + return kinetic.ErrInvalidMsgCountRateLimit + } +} + +func FirehoseWriterMsgSizeRateLimit(limit int) FireHoseWriterOptionsFn { + return func(o *firehoseWriterOptions) error { + if limit > 0 && limit <= firehoseMsgSizeRateLimit { + o.msgSizeRateLimit = limit + return nil + } + return kinetic.ErrInvalidMsgSizeRateLimit + } +} + +func FirehoseWriterThroughputMultiplier(multiplier int) FireHoseWriterOptionsFn { + return func(o *firehoseWriterOptions) error { + if multiplier > 0 { + o.throughputMultiplier = multiplier + return nil + } + return kinetic.ErrInvalidThroughputMultiplier + } +} + +func FirehoseWriterLogLevel(ll aws.LogLevelType) FireHoseWriterOptionsFn { + return func(o *firehoseWriterOptions) error { + o.logLevel = ll & 0xffff0000 + return nil + } +} + +func FirehoseWriterStatsCollector(sc StatsCollector) FireHoseWriterOptionsFn { + return func(o *firehoseWriterOptions) error { + o.stats = sc + return nil + } } // FirehoseWriter handles the API to send records to Kinesis. type FirehoseWriter struct { *firehoseWriterOptions *kinetic.LogHelper - stream string client firehoseiface.FirehoseAPI } // NewFirehoseWriter creates a new stream writer to write records to a Kinesis. -func NewFirehoseWriter(c *aws.Config, stream string, fn ...func(*FirehoseWriterConfig)) (*FirehoseWriter, error) { - cfg := NewFirehoseWriterConfig(c) - for _, f := range fn { - f(cfg) +func NewFirewhoseWriter(c *aws.Config, stream string, optionFns ...FireHoseWriterOptionsFn) (*FirehoseWriter, error) { + firehoseWriterOptions := defaultFirehoseWriterOptions() + for _, optionFn := range optionFns { + optionFn(firehoseWriterOptions) } - sess, err := session.NewSession(cfg.AwsConfig) + sess, err := session.NewSession(c) if err != nil { return nil, err } return &FirehoseWriter{ - firehoseWriterOptions: cfg.firehoseWriterOptions, - LogHelper: &kinetic.LogHelper{ - LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, - }, stream: stream, client: firehose.New(sess), + firehoseWriterOptions: firehoseWriterOptions, + LogHelper: &kinetic.LogHelper{ + LogLevel: firehoseWriterOptions.logLevel, + Logger: c.Logger, + }, }, nil } @@ -74,7 +135,7 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*kinetic.Mes }) req.Handlers.Build.PushBack(func(r *request.Request) { - w.Stats.UpdatePutRecordsBuildDuration(time.Since(startBuildTime)) + w.stats.UpdatePutRecordsBuildDuration(time.Since(startBuildTime)) w.LogDebug("Finished PutRecords Build, took", time.Since(start)) }) @@ -84,17 +145,17 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*kinetic.Mes }) req.Handlers.Send.PushBack(func(r *request.Request) { - w.Stats.UpdatePutRecordsSendDuration(time.Since(startSendTime)) + w.stats.UpdatePutRecordsSendDuration(time.Since(startSendTime)) w.LogDebug("Finished PutRecords Send, took", time.Since(start)) }) w.LogDebug("Starting PutRecords Build/Sign request, took", time.Since(start)) - w.Stats.AddPutRecordsCalled(1) + w.stats.AddPutRecordsCalled(1) if err := req.Send(); err != nil { w.LogError("Error putting records:", err.Error()) return err } - w.Stats.UpdatePutRecordsDuration(time.Since(start)) + w.stats.UpdatePutRecordsDuration(time.Since(start)) if resp == nil { return kinetic.ErrNilPutRecordsResponse @@ -111,23 +172,22 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*kinetic.Mes if record.RecordId != nil { // TODO: per-shard metrics messages[idx].RecordID = record.RecordId - w.Stats.AddSentSuccess(1) + w.stats.AddSentSuccess(1) } else { switch aws.StringValue(record.ErrorCode) { case firehose.ErrCodeLimitExceededException: - w.Stats.AddProvisionedThroughputExceeded(1) + w.stats.AddProvisionedThroughputExceeded(1) default: w.LogDebug("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(record.ErrorMessage)) } messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage messages[idx].FailCount++ - w.Stats.AddSentFailed(1) + w.stats.AddSentFailed(1) fn(messages[idx]) } } - return nil } diff --git a/producer/firehosewriterconfig.go b/producer/firehosewriterconfig.go deleted file mode 100644 index 3d22692..0000000 --- a/producer/firehosewriterconfig.go +++ /dev/null @@ -1,62 +0,0 @@ -package producer - -import ( - "github.com/aws/aws-sdk-go/aws" -) - -const ( - firehoseMsgCountRateLimit = 5000 // AWS Firehose limit of 5000 records/sec - firehoseMsgSizeRateLimit = 5000000 // AWS Firehose limit of 5 MB/sec -) - -// FirehoseWriterConfig is used to configure FirehoseWriter -type FirehoseWriterConfig struct { - *firehoseWriterOptions - AwsConfig *aws.Config - LogLevel aws.LogLevelType -} - -// NewFirehoseWriterConfig creates a new instance of FirehoseWriterConfig -func NewFirehoseWriterConfig(cfg *aws.Config) *FirehoseWriterConfig { - return &FirehoseWriterConfig{ - AwsConfig: cfg, - firehoseWriterOptions: &firehoseWriterOptions{ - msgCountRateLimit: firehoseMsgCountRateLimit, - msgSizeRateLimit: firehoseMsgSizeRateLimit, - throughputMultiplier: 1, - Stats: &NilStatsCollector{}, - }, - LogLevel: *cfg.LogLevel, - } -} - -// SetMsgCountRateLimit configures the maximum number of messages that can be sent per second -func (c *FirehoseWriterConfig) SetMsgCountRateLimit(limit int) { - if limit > firehoseMsgCountRateLimit { - - } - c.msgCountRateLimit = limit -} - -// SetMsgSizeRateLimit configures the maximum transmission size of the messages that can be sent per second -func (c *FirehoseWriterConfig) SetMsgSizeRateLimit(limit int) { - if limit > firehoseMsgSizeRateLimit { - - } - c.msgSizeRateLimit = limit -} - -// SetThroughputMultiplier configures the throughput multiplier to increase the streaming rate limits for Firehose -func (c *FirehoseWriterConfig) SetThroughputMultiplier(count int) { - c.throughputMultiplier = count -} - -// SetStatsCollector configures a listener to handle listener metrics. -func (c *FirehoseWriterConfig) SetStatsCollector(stats StatsCollector) { - c.Stats = stats -} - -// SetLogLevel configures the log levels for the SDK. -func (c *FirehoseWriterConfig) SetLogLevel(logLevel aws.LogLevelType) { - c.LogLevel = logLevel & 0xffff0000 -} diff --git a/producer/kinesiswriter.go b/producer/kinesis_writer.go similarity index 65% rename from producer/kinesiswriter.go rename to producer/kinesis_writer.go index 342d689..2ba852e 100644 --- a/producer/kinesiswriter.go +++ b/producer/kinesis_writer.go @@ -14,40 +14,98 @@ import ( "github.com/rewardStyle/kinetic" ) +const ( + kinesisMsgCountRateLimit = 1000 // AWS Kinesis limit of 1000 records/sec + kinesisMsgSizeRateLimit = 1000000 // AWS Kinesis limit of 1 MB/sec +) + type kinesisWriterOptions struct { - responseReadTimeout time.Duration // maximum time to wait for PutRecords API call before timing out - msgCountRateLimit int // maximum number of records to be sent per second - msgSizeRateLimit int // maximum (transmission) size of records to be sent per second - Stats StatsCollector // stats collection mechanism + responseReadTimeout time.Duration // maximum time to wait for PutRecords API call before timing out + msgCountRateLimit int // maximum number of records to be sent per second + msgSizeRateLimit int // maximum (transmission) size of records to be sent per second + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + stats StatsCollector // stats collection mechanism +} + +func defaultKinesisWriterOptions() *kinesisWriterOptions { + return &kinesisWriterOptions{ + responseReadTimeout: time.Second, + msgCountRateLimit: kinesisMsgCountRateLimit, + msgSizeRateLimit: kinesisMsgSizeRateLimit, + logLevel: aws.LogOff, + stats: &NilStatsCollector{}, + } +} + +type KinesisWriterOptionsFn func(*kinesisWriterOptions) error + +func KinesisWriterResponseReadTimeout(timeout time.Duration) KinesisWriterOptionsFn { + return func(o *kinesisWriterOptions) error { + o.responseReadTimeout = timeout + return nil + } +} + +func KinesisWriterMsgCountRateLimit(limit int) KinesisWriterOptionsFn { + return func(o *kinesisWriterOptions) error { + if limit > 0 && limit <= kinesisMsgCountRateLimit { + o.msgSizeRateLimit = limit + return nil + } + return kinetic.ErrInvalidMsgSizeRateLimit + } +} + +func KinesisWriterMsgSizeRateLimit(limit int) KinesisWriterOptionsFn { + return func(o *kinesisWriterOptions) error { + if limit > 0 && limit <= kinesisMsgSizeRateLimit { + o.msgSizeRateLimit = limit + return nil + } + return kinetic.ErrInvalidMsgSizeRateLimit + } +} + +func KinesisWriterLogLevel(ll aws.LogLevelType) KinesisWriterOptionsFn { + return func(o *kinesisWriterOptions) error { + o.logLevel = ll & 0xffff0000 + return nil + } +} + +func KinesisWriterStatsCollector(sc StatsCollector) KinesisWriterOptionsFn { + return func(o *kinesisWriterOptions) error { + o.stats = sc + return nil + } } // KinesisWriter handles the API to send records to Kinesis. type KinesisWriter struct { *kinesisWriterOptions *kinetic.LogHelper - stream string client kinesisiface.KinesisAPI } // NewKinesisWriter creates a new stream writer to write records to a Kinesis. -func NewKinesisWriter(c *aws.Config, stream string, fn ...func(*KinesisWriterConfig)) (*KinesisWriter, error) { - cfg := NewKinesisWriterConfig(c) - for _, f := range fn { - f(cfg) +func NewKinesisWriter(c *aws.Config, stream string, optionFns ...KinesisWriterOptionsFn) (*KinesisWriter, error) { + kinesisWriterOptions := defaultKinesisWriterOptions() + for _, option := range optionFns { + option(kinesisWriterOptions) } - sess, err := session.NewSession(cfg.AwsConfig) + sess, err := session.NewSession(c) if err != nil { return nil, err } return &KinesisWriter{ - kinesisWriterOptions: cfg.kinesisWriterOptions, - LogHelper: &kinetic.LogHelper{ - LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, - }, stream: stream, client: kinesis.New(sess), + kinesisWriterOptions: kinesisWriterOptions, + LogHelper: &kinetic.LogHelper{ + LogLevel: kinesisWriterOptions.logLevel, + Logger: c.Logger, + }, }, nil } @@ -75,7 +133,7 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*kinetic.Mess }) req.Handlers.Build.PushBack(func(r *request.Request) { - w.Stats.UpdatePutRecordsBuildDuration(time.Since(startBuildTime)) + w.stats.UpdatePutRecordsBuildDuration(time.Since(startBuildTime)) w.LogDebug("Finished PutRecords Build, took", time.Since(start)) }) @@ -85,17 +143,17 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*kinetic.Mess }) req.Handlers.Send.PushBack(func(r *request.Request) { - w.Stats.UpdatePutRecordsSendDuration(time.Since(startSendTime)) + w.stats.UpdatePutRecordsSendDuration(time.Since(startSendTime)) w.LogDebug("Finished PutRecords Send, took", time.Since(start)) }) w.LogDebug("Starting PutRecords Build/Sign request, took", time.Since(start)) - w.Stats.AddPutRecordsCalled(1) + w.stats.AddPutRecordsCalled(1) if err := req.Send(); err != nil { w.LogError("Error putting records:", err.Error()) return err } - w.Stats.UpdatePutRecordsDuration(time.Since(start)) + w.stats.UpdatePutRecordsDuration(time.Since(start)) if resp == nil { return kinetic.ErrNilPutRecordsResponse @@ -113,18 +171,18 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*kinetic.Mess // TODO: per-shard metrics messages[idx].SequenceNumber = record.SequenceNumber messages[idx].ShardID = record.ShardId - w.Stats.AddSentSuccess(1) + w.stats.AddSentSuccess(1) } else { switch aws.StringValue(record.ErrorCode) { case kinesis.ErrCodeProvisionedThroughputExceededException: - w.Stats.AddProvisionedThroughputExceeded(1) + w.stats.AddProvisionedThroughputExceeded(1) default: w.LogDebug("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(record.ErrorMessage)) } messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage messages[idx].FailCount++ - w.Stats.AddSentFailed(1) + w.stats.AddSentFailed(1) fn(messages[idx]) } diff --git a/producer/kinesiswriterconfig.go b/producer/kinesiswriterconfig.go deleted file mode 100644 index 84a720d..0000000 --- a/producer/kinesiswriterconfig.go +++ /dev/null @@ -1,70 +0,0 @@ -package producer - -import ( - "log" - "time" - - "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic" -) - -const ( - kinesisMsgCountRateLimit = 1000 // AWS Kinesis limit of 1000 records/sec - kinesisMsgSizeRateLimit = 1000000 // AWS Kinesis limit of 1 MB/sec -) - -// KinesisWriterConfig is used to configure KinesisWriter -type KinesisWriterConfig struct { - *kinesisWriterOptions - AwsConfig *aws.Config - LogLevel aws.LogLevelType -} - -// NewKinesisWriterConfig creates a new instance of KinesisWriterConfig -func NewKinesisWriterConfig(cfg *aws.Config) *KinesisWriterConfig { - return &KinesisWriterConfig{ - AwsConfig: cfg, - kinesisWriterOptions: &kinesisWriterOptions{ - responseReadTimeout: time.Second, - msgCountRateLimit: kinesisMsgCountRateLimit, - msgSizeRateLimit: kinesisMsgSizeRateLimit, - Stats: &NilStatsCollector{}, - }, - LogLevel: *cfg.LogLevel, - } -} - -// SetResponseReadTimeout configures the time to wait for each successive Read operation on the GetRecords response payload. -func (c *KinesisWriterConfig) SetResponseReadTimeout(timeout time.Duration) { - c.responseReadTimeout = timeout -} - -// SetMsgCountRateLimit configures the maximum number of messages that can be sent per second -func (c *KinesisWriterConfig) SetMsgCountRateLimit(limit int) { - if limit > 0 && limit <= kinesisMsgCountRateLimit { - c.msgCountRateLimit = limit - } else { - log.Fatal("Message Count Rate Limit must be positive and less than ", kinesisMsgCountRateLimit) - panic(kinetic.ErrInvalidMsgCountRateLimit) - } -} - -// SetMsgSizeRateLimit configures the maximum transmission size of the messages that can be sent per second -func (c *KinesisWriterConfig) SetMsgSizeRateLimit(limit int) { - if limit > 0 && limit <= kinesisMsgSizeRateLimit { - c.msgSizeRateLimit = limit - } else { - log.Fatal("Message Count Size Limit must be positive and less than ", kinesisMsgSizeRateLimit) - panic(kinetic.ErrInvalidMsgSizeRateLimit) - } -} - -// SetStatsCollector configures a listener to handle listener metrics. -func (c *KinesisWriterConfig) SetStatsCollector(stats StatsCollector) { - c.Stats = stats -} - -// SetLogLevel configures the log levels for the SDK. -func (c *KinesisWriterConfig) SetLogLevel(logLevel aws.LogLevelType) { - c.LogLevel = logLevel & 0xffff0000 -} diff --git a/producer/producer.go b/producer/producer.go index 48e8186..af7b690 100644 --- a/producer/producer.go +++ b/producer/producer.go @@ -12,6 +12,10 @@ import ( "golang.org/x/time/rate" ) +const ( + putRecordsMaxBatchSize = 500 +) + // producerOptions holds all of the configurable settings for a Producer type producerOptions struct { batchSize int // maximum message capacity per request @@ -21,9 +25,101 @@ type producerOptions struct { concurrency int // number of concurrent workers per shard shardCheckFreq time.Duration // frequency (specified as a duration) with which to check the the shard size dataSpillFn MessageHandlerAsync // callback function for handling dropped messages that the producer was unable to send to the stream + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level Stats StatsCollector // stats collection mechanism } +func defaultProducerOptions() *producerOptions { + return &producerOptions{ + batchSize: putRecordsMaxBatchSize, + batchTimeout: time.Second, + queueDepth: 10000, + maxRetryAttempts: 10, + concurrency: 3, + shardCheckFreq: time.Minute, + dataSpillFn: func(*kinetic.Message) error { return nil }, + logLevel: aws.LogOff, + Stats: &NilStatsCollector{}, + } +} + +type ProducerOptionsFn func(*producerOptions) error + +func ProducerBatchSize(size int) ProducerOptionsFn { + return func(o *producerOptions) error { + if size > 0 && size <= putRecordsMaxBatchSize { + o.batchSize = size + return nil + } + return kinetic.ErrInvalidBatchSize + } +} + +func ProducerBatchTimeout(timeout time.Duration) ProducerOptionsFn { + return func(o *producerOptions) error { + o.batchTimeout = timeout + return nil + } +} + +func ProducerQueueDepth(queueDepth int) ProducerOptionsFn { + return func(o *producerOptions) error { + if queueDepth > 0 { + o.queueDepth = queueDepth + return nil + } + return kinetic.ErrInvalidQueueDepth + } +} + +func ProducerMaxRetryAttempts(attemtps int) ProducerOptionsFn { + return func(o *producerOptions) error { + if attemtps > 0 { + o.maxRetryAttempts = attemtps + return nil + } + return kinetic.ErrInvalidMaxRetryAttempts + } +} + +func ProducerConcurrency(count int) ProducerOptionsFn { + return func(o *producerOptions) error { + if count > 0 { + o.concurrency = count + return nil + } + return kinetic.ErrInvalidConcurrency + } +} + +func ProducerShardCheckFrequency(duration time.Duration) ProducerOptionsFn { + return func(o *producerOptions) error { + o.shardCheckFreq = duration + return nil + } +} + +func ProducerDataSpillFn(fn MessageHandlerAsync) ProducerOptionsFn { + return func(o *producerOptions) error { + o.dataSpillFn = fn + return nil + } +} + +func ProducerLogLevel(ll aws.LogLevelType) ProducerOptionsFn { + return func(o *producerOptions) error { + o.logLevel = ll & 0xffff0000 + return nil + } +} + +func ProducerStatsCollector(sc StatsCollector) ProducerOptionsFn { + return func(o *producerOptions) error { + o.Stats = sc + return nil + } +} + // Producer sends records to AWS Kinesis or Firehose. type Producer struct { *producerOptions // contains all of the configuration settings for the Producer @@ -44,16 +140,16 @@ type Producer struct { } // NewProducer creates a new producer for writing records to a Kinesis or Firehose stream. -func NewProducer(c *aws.Config, w StreamWriter, fn ...func(*Config)) (*Producer, error) { - cfg := NewConfig(c) - for _, f := range fn { - f(cfg) +func NewProducer(c *aws.Config, w StreamWriter, optionFns ...ProducerOptionsFn) (*Producer, error) { + producerOptions := defaultProducerOptions() + for _, optionFn := range optionFns { + optionFn(producerOptions) } return &Producer{ - producerOptions: cfg.producerOptions, + producerOptions: producerOptions, LogHelper: &kinetic.LogHelper{ - LogLevel: cfg.LogLevel, - Logger: cfg.AwsConfig.Logger, + LogLevel: producerOptions.logLevel, + Logger: c.Logger, }, writer: w, }, nil diff --git a/producer/stats.go b/producer/producer_stats.go similarity index 100% rename from producer/stats.go rename to producer/producer_stats.go diff --git a/producer/stats_test.go b/producer/producer_stats_test.go similarity index 100% rename from producer/stats_test.go rename to producer/producer_stats_test.go diff --git a/producer/producer_test.go b/producer/producer_test.go index d09fc4b..c1b9ae9 100644 --- a/producer/producer_test.go +++ b/producer/producer_test.go @@ -12,7 +12,8 @@ import ( "github.com/aws/aws-sdk-go/aws/request" "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/listener" + "github.com/rewardStyle/kinetic/consumer" + "github.com/rewardStyle/kinetic/producer" ) func TestProducer(t *testing.T) { @@ -42,28 +43,41 @@ func TestProducer(t *testing.T) { So(w, ShouldNotBeNil) So(err, ShouldBeNil) - p, err := NewProducer(k.Session.Config, w, func(c *Config) { - c.SetBatchSize(5) - c.SetBatchTimeout(time.Second) - c.SetQueueDepth(10) - c.SetConcurrency(2) - }) + p, err := NewProducer(k.Session.Config, w, + producer.ProducerBatchSize(5), + producer.ProducerBatchTimeout(time.Second), + producer.ProducerMaxRetryAttempts(3), + producer.ProducerQueueDepth(10), + producer.ProducerConcurrency(2), + producer.ProducerShardCheckFrequency(time.Minute), + producer.ProducerDataSpillFn(func(msg *kinetic.Message) error { + //log.Printf("Message was dropped: [%s]\n", string(msg.Data)) + return nil + }), + producer.ProducerLogLevel(aws.LogOff), + //producer.ProducerStatsCollector(), + ) So(p, ShouldNotBeNil) So(err, ShouldBeNil) So(k.Session, ShouldNotBeNil) So(k.Session.Config, ShouldNotBeNil) - r, err := listener.NewKinesisReader(k.Session.Config, stream, shards[0], - func(krc *listener.KinesisReaderConfig) { - krc.SetResponseReadTimeout(1000 * time.Millisecond) - }) + r, err := consumer.NewKinesisReader(k.Session.Config, stream, shards[0], + //consumer.KinesisReaderBatchSize(), + //consumer.KinesisReaderShardIterator(), + consumer.KinesisReaderResponseReadTimeout(time.Second), + //consumer.KinesisReaderLogLevel(), + //consumer.KinesisReaderStatsCollector(), + ) So(r, ShouldNotBeNil) So(err, ShouldBeNil) - l, err := listener.NewListener(k.Session.Config, r, func(c *listener.Config) { - c.SetQueueDepth(10) - c.SetConcurrency(10) - }) + l, err := consumer.NewConsumer(k.Session.Config, r, + consumer.ConsumerQueueDepth(10), + consumer.ConsumerConcurrency(10), + consumer.ConsumerLogLevel(aws.LogOff), + //consumer.ConsumerStatsCollector(lsc), + ) So(l, ShouldNotBeNil) So(err, ShouldBeNil) diff --git a/producer/producerconfig.go b/producer/producerconfig.go deleted file mode 100644 index 05f2ce1..0000000 --- a/producer/producerconfig.go +++ /dev/null @@ -1,85 +0,0 @@ -package producer - -import ( - "log" - "time" - - "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic" -) - -// Config is used to configure a Producer instance. -type Config struct { - *producerOptions - AwsConfig *aws.Config - LogLevel aws.LogLevelType -} - -// NewConfig creates a new instance of Config. -func NewConfig(cfg *aws.Config) *Config { - return &Config{ - AwsConfig: cfg, - producerOptions: &producerOptions{ - batchSize: 500, - batchTimeout: time.Second, - queueDepth: 10000, - maxRetryAttempts: 10, - concurrency: 3, - shardCheckFreq: time.Minute, - dataSpillFn: func(*kinetic.Message) error { return nil }, - Stats: &NilStatsCollector{}, - }, - LogLevel: *cfg.LogLevel, - } -} - -// SetBatchSize configures the batch size to flush pending records to the PutRecords call. -func (c *Config) SetBatchSize(batchSize int) { - if batchSize > 0 && batchSize <= 500 { - c.batchSize = batchSize - } else { - // http://docs.aws.amazon.com/firehose/latest/APIReference/API_PutRecordBatch.html - log.Fatal("BatchSize must be less than or equal to 500 ") - panic(kinetic.ErrInvalidBatchSize) - } -} - -// SetBatchTimeout configures the timeout to flush pending records to the PutRecords call. -func (c *Config) SetBatchTimeout(timeout time.Duration) { - c.batchTimeout = timeout -} - -// SetQueueDepth controls the number of messages that can be in the channel to be processed by produce at a given time. -func (c *Config) SetQueueDepth(queueDepth int) { - c.queueDepth = queueDepth -} - -// SetMaxRetryAttempts controls the number of times a message can be retried before it is discarded. -func (c *Config) SetMaxRetryAttempts(attempts int) { - c.maxRetryAttempts = attempts -} - -// SetConcurrency defines the number of concurrent workers to run per active shard (multiplier) -func (c *Config) SetConcurrency(count int) { - c.concurrency = count -} - -// SetShardCheckFreq defines a frequency (specified as a duration) with which to check the shard size -func (c *Config) SetShardCheckFreq(duration time.Duration) { - c.shardCheckFreq = duration -} - -// SetDataSpillFn defines a callback function to be called when a message gets dropped by the producer -func (c *Config) SetDataSpillFn(fn MessageHandlerAsync) { - c.dataSpillFn = fn -} - -// SetStatsCollector configures a listener to handle producer metrics. -func (c *Config) SetStatsCollector(stats StatsCollector) { - c.Stats = stats -} - -// SetLogLevel configures both the SDK and Kinetic log levels. -func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { - c.LogLevel = logLevel & 0xffff0000 -} diff --git a/producer/producerconfig_test.go b/producer/producerconfig_test.go deleted file mode 100644 index 41a34e3..0000000 --- a/producer/producerconfig_test.go +++ /dev/null @@ -1,105 +0,0 @@ -package producer - -import ( - . "github.com/smartystreets/goconvey/convey" - - "testing" - "time" - - "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic" -) - -type DebugStatsCollector struct{} - -func (l *DebugStatsCollector) AddSentTotal(int) {} -func (l *DebugStatsCollector) AddSentSuccess(int) {} -func (l *DebugStatsCollector) AddSentFailed(int) {} -func (l *DebugStatsCollector) AddSentRetried(int) {} -func (l *DebugStatsCollector) AddDroppedTotal(int) {} -func (l *DebugStatsCollector) AddDroppedCapacity(int) {} -func (l *DebugStatsCollector) AddDroppedRetries(int) {} -func (l *DebugStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} -func (l *DebugStatsCollector) AddPutRecordsCalled(int) {} -func (l *DebugStatsCollector) AddProvisionedThroughputExceeded(int) {} -func (l *DebugStatsCollector) AddPutRecordsTimeout(int) {} -func (l *DebugStatsCollector) UpdatePutRecordsDuration(time.Duration) {} -func (l *DebugStatsCollector) UpdatePutRecordsBuildDuration(time.Duration) {} -func (l *DebugStatsCollector) UpdatePutRecordsSendDuration(time.Duration) {} -func (l *DebugStatsCollector) UpdateProducerConcurrency(int) {} - -type DebugStreamWriter struct{} - -func (w *DebugStreamWriter) PutRecords(batch []*kinetic.Message) ([]*kinetic.Message, error) { - return nil, nil -} - -func TestNewConfig(t *testing.T) { - Convey("given a new producer config", t, func() { - k, err := kinetic.New(func(c *kinetic.Config) { - c.SetEndpoint("bogus-endpoint") - }) - So(err, ShouldBeNil) - cfg := NewConfig(k.Session.Config) - - Convey("check the default values for its non-zero config", func() { - So(cfg.batchSize, ShouldEqual, 500) - So(cfg.batchTimeout, ShouldEqual, 1*time.Second) - So(cfg.queueDepth, ShouldEqual, 10000) - So(cfg.maxRetryAttempts, ShouldEqual, 10) - So(cfg.concurrency, ShouldEqual, 3) - So(cfg.shardCheckFreq, ShouldEqual, time.Minute) - So(cfg.Stats, ShouldHaveSameTypeAs, &NilStatsCollector{}) - So(cfg.LogLevel.Value(), ShouldEqual, kinetic.LogOff) - }) - - Convey("check that we can set both the sdk and kinetic log level", func() { - ll := aws.LogDebug | aws.LogDebugWithSigning | kinetic.LogDebug - cfg.SetLogLevel(ll) - So(cfg.LogLevel.AtLeast(kinetic.LogDebug), ShouldBeTrue) - }) - - Convey("check that we can set the batch size", func() { - cfg.SetBatchSize(100) - So(cfg.batchSize, ShouldEqual, 100) - }) - - Convey("check that we can set the batch timeout", func() { - cfg.SetBatchTimeout(10 * time.Second) - So(cfg.batchTimeout, ShouldEqual, 10*time.Second) - }) - - Convey("check that we can set the queue depth", func() { - cfg.SetQueueDepth(1000) - So(cfg.queueDepth, ShouldEqual, 1000) - }) - - Convey("check that we can set the max retries", func() { - cfg.SetMaxRetryAttempts(100) - So(cfg.maxRetryAttempts, ShouldEqual, 100) - }) - - Convey("check that we can set the workers per shard", func() { - cfg.SetConcurrency(10) - So(cfg.concurrency, ShouldEqual, 10) - }) - - Convey("check that we can set the data spill callback function", func() { - fn := func(msg *kinetic.Message) error { - return nil - } - cfg.SetDataSpillFn(fn) - So(cfg.dataSpillFn, ShouldEqual, fn) - }) - - Convey("check that we can set the check shard frequency", func() { - cfg.SetShardCheckFreq(time.Second) - So(cfg.shardCheckFreq, ShouldEqual, time.Second) - }) - - Convey("check that we can configure a stats collector", func() { - cfg.SetStatsCollector(&DebugStatsCollector{}) - So(cfg.Stats, ShouldHaveSameTypeAs, &DebugStatsCollector{}) - }) - }) -} diff --git a/producer/sharedtypes.go b/producer/shared_types.go similarity index 100% rename from producer/sharedtypes.go rename to producer/shared_types.go diff --git a/testexec/README.md b/testexec/README.md index 3daf1a8..a3306b9 100644 --- a/testexec/README.md +++ b/testexec/README.md @@ -1,5 +1,5 @@ # testexec -testexec is an executable CLI program to test the kinetic producer / listener. +testexec is an executable CLI program to test the kinetic producer / consumer. The executable program creates a kinetic object connecting to either a local kinesalte instance or to AWS Kinesis using local AWS credentials/config. diff --git a/testexec/main.go b/testexec/main.go index 7fdfb6c..9f6cad7 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -19,7 +19,7 @@ import ( "github.com/aws/aws-sdk-go/aws/session" metrics "github.com/jasonyurs/go-metrics" "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/listener" + "github.com/rewardStyle/kinetic/consumer" "github.com/rewardStyle/kinetic/producer" "net/http" @@ -237,30 +237,31 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce } psc := producer.NewDefaultStatsCollector(registry) - w, err := producer.NewKinesisWriter(k.Session.Config, streamName, func(kwc *producer.KinesisWriterConfig) { - //kwc.SetLogLevel(kinetic.LogDebug) - kwc.SetResponseReadTimeout(time.Second) - kwc.SetStatsCollector(psc) - kwc.SetMsgCountRateLimit(1000) - kwc.SetMsgSizeRateLimit(1000000) - }) + w, err := producer.NewKinesisWriter(k.Session.Config, streamName, + producer.KinesisWriterResponseReadTimeout(time.Second), + producer.KinesisWriterMsgCountRateLimit(1000), + producer.KinesisWriterMsgSizeRateLimit(1000000), + producer.KinesisWriterLogLevel(kinetic.LogDebug), + producer.KinesisWriterStatsCollector(psc), + ) if err != nil { log.Fatalf("Unable to create a new kinesis stream writer due to: %v\n", err) } - p, err := producer.NewProducer(k.Session.Config, w, func(c *producer.Config) { - c.SetBatchTimeout(1000 * time.Millisecond) - c.SetBatchSize(500) - c.SetMaxRetryAttempts(3) - c.SetStatsCollector(psc) - c.SetQueueDepth(10000) - c.SetConcurrency(3) - c.SetShardCheckFreq(time.Minute) - c.SetDataSpillFn(func(msg *kinetic.Message) error { + p, err := producer.NewProducer(k.Session.Config, w, + producer.ProducerBatchSize(500), + producer.ProducerBatchTimeout(time.Second), + producer.ProducerMaxRetryAttempts(3), + producer.ProducerQueueDepth(10000), + producer.ProducerConcurrency(3), + producer.ProducerShardCheckFrequency(time.Minute), + producer.ProducerDataSpillFn(func(msg *kinetic.Message) error { //log.Printf("Message was dropped: [%s]\n", string(msg.Data)) return nil - }) - }) + }), + producer.ProducerLogLevel(aws.LogOff), + producer.ProducerStatsCollector(psc), + ) if err != nil { log.Fatalf("Unable to create a new producer due to: %v\n", err) } @@ -268,7 +269,7 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce return p } -func newKineticListener(k *kinetic.Kinetic, streamName string) *listener.Listener { +func newKineticListener(k *kinetic.Kinetic, streamName string) *consumer.Consumer { if *cfg.Verbose { log.Println("Creating a kinetic listener ...") } @@ -279,21 +280,24 @@ func newKineticListener(k *kinetic.Kinetic, streamName string) *listener.Listene log.Fatalf("Unable to get shards for stream %s due to: %v\n", streamName, err) } - lsc := listener.NewDefaultStatsCollector(registry) - r, err := listener.NewKinesisReader(k.Session.Config, streamName, shards[0], - func(krc *listener.KinesisReaderConfig) { - krc.SetResponseReadTimeout(1000 * time.Millisecond) - krc.SetStatsCollector(lsc) - }) + lsc := consumer.NewDefaultStatsCollector(registry) + r, err := consumer.NewKinesisReader(k.Session.Config, streamName, shards[0], + consumer.KinesisReaderBatchSize(10000), + //consumer.KinesisReaderShardIterator(), + consumer.KinesisReaderResponseReadTimeout(time.Second), + consumer.KinesisReaderLogLevel(aws.LogOff), + consumer.KinesisReaderStatsCollector(lsc), + ) if err != nil { log.Fatalf("Unable to create a new kinesis reader due to: %v\n", err) } - l, err := listener.NewListener(k.Session.Config, r, func(c *listener.Config) { - c.SetQueueDepth(500) - c.SetConcurrency(10) - c.SetStatsCollector(lsc) - }) + l, err := consumer.NewConsumer(k.Session.Config, r, + consumer.ConsumerQueueDepth(500), + consumer.ConsumerConcurrency(10), + consumer.ConsumerLogLevel(aws.LogOff), + consumer.ConsumerStatsCollector(lsc), + ) if err != nil { log.Fatalf("Unable to create a new listener due to: %v\n", err) } @@ -314,7 +318,7 @@ func handlePoD() { } } -func display(sd *StreamData, p *producer.Producer, l *listener.Listener, wg *sync.WaitGroup) { +func display(sd *StreamData, p *producer.Producer, c *consumer.Consumer, wg *sync.WaitGroup) { defer wg.Done() for { @@ -336,7 +340,7 @@ func display(sd *StreamData, p *producer.Producer, l *listener.Listener, wg *syn p.Stats.(*producer.DefaultStatsCollector).PrintStats() } if *cfg.Mode != ModeWrite { - l.Stats.(*listener.DefaultStatsCollector).PrintStats() + c.Stats.(*consumer.DefaultStatsCollector).PrintStats() sd.printSummary() } return @@ -350,7 +354,7 @@ func display(sd *StreamData, p *producer.Producer, l *listener.Listener, wg *syn p.Stats.(*producer.DefaultStatsCollector).PrintStats() } if *cfg.Mode != ModeWrite { - l.Stats.(*listener.DefaultStatsCollector).PrintStats() + c.Stats.(*consumer.DefaultStatsCollector).PrintStats() sd.printStats() } } @@ -472,7 +476,7 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { produceWg.Wait() } -func listen(sd *StreamData, l *listener.Listener, wg *sync.WaitGroup) { +func listen(sd *StreamData, c *consumer.Consumer, wg *sync.WaitGroup) { defer wg.Done() // Return early if we don't need to produce @@ -492,7 +496,7 @@ func listen(sd *StreamData, l *listener.Listener, wg *sync.WaitGroup) { // Call Listen within a go routine go func() { - l.Listen(func(m *kinetic.Message, wg *sync.WaitGroup) error { + c.Listen(func(m *kinetic.Message, wg *sync.WaitGroup) error { defer wg.Done() // Unmarshal data @@ -544,7 +548,7 @@ func listen(sd *StreamData, l *listener.Listener, wg *sync.WaitGroup) { } return case <-time.After(time.Second): - newConsumed := l.Stats.(*listener.DefaultStatsCollector).Consumed.Count() + newConsumed := c.Stats.(*consumer.DefaultStatsCollector).Consumed.Count() if consumed != uint64(newConsumed) { staleTime.Reset(staleTimeout) consumed = uint64(newConsumed) From 56194f7e39eb05dd2c7909d429e9515df55e3f3d Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 28 Jul 2017 00:26:38 -0500 Subject: [PATCH 070/109] Consolidated the producer and consumer packages into the main kinetic package. Simplied the kinetic parameter configuration by applying the functional options pattern. --- config.go | 88 ----------- config_test.go | 93 ------------ consumer/consumer.go => consumer.go | 41 +++-- consumer/kcl_reader_config.go | 51 ------- consumer/shared_types.go | 30 ---- .../consumer_stats.go => consumer_stats.go | 140 +++++++++--------- ...er_stats_test.go => consumer_stats_test.go | 40 ++--- consumer/consumer_test.go => consumer_test.go | 49 +++--- .../firehose_writer.go => firehose_writer.go | 50 +++---- consumer/kcl_actions.go => kcl_actions.go | 10 +- consumer/kcl_reader.go => kcl_reader.go | 31 ++-- .../kinesis_reader.go => kinesis_reader.go | 48 +++--- ...s_reader_test.go => kinesis_reader_test.go | 34 ++--- .../kinesis_writer.go => kinesis_writer.go | 52 ++++--- kinetic.go | 84 ++++++++++- kinetic_integration_test.go | 65 ++++---- kinetic_test.go | 10 +- producer/producer.go => producer.go | 67 +++++---- .../producer_stats.go => producer_stats.go | 85 +++++------ ...er_stats_test.go => producer_stats_test.go | 44 +++--- producer/producer_test.go => producer_test.go | 60 ++++---- consumer/read_closer.go => read_closer.go | 2 +- .../shard_iterator.go => shard_iterator.go | 2 +- ...iterator_test.go => shard_iterator_test.go | 2 +- producer/shared_types.go => shared_types.go | 37 +++-- testexec/main.go | 100 ++++++------- 26 files changed, 546 insertions(+), 769 deletions(-) delete mode 100644 config.go delete mode 100644 config_test.go rename consumer/consumer.go => consumer.go (89%) delete mode 100644 consumer/kcl_reader_config.go delete mode 100644 consumer/shared_types.go rename consumer/consumer_stats.go => consumer_stats.go (52%) rename consumer/consumer_stats_test.go => consumer_stats_test.go (65%) rename consumer/consumer_test.go => consumer_test.go (87%) rename producer/firehose_writer.go => firehose_writer.go (80%) rename consumer/kcl_actions.go => kcl_actions.go (94%) rename consumer/kcl_reader.go => kcl_reader.go (89%) rename consumer/kinesis_reader.go => kinesis_reader.go (90%) rename consumer/kinesis_reader_test.go => kinesis_reader_test.go (61%) rename producer/kinesis_writer.go => kinesis_writer.go (83%) rename producer/producer.go => producer.go (86%) rename producer/producer_stats.go => producer_stats.go (76%) rename producer/producer_stats_test.go => producer_stats_test.go (67%) rename producer/producer_test.go => producer_test.go (72%) rename consumer/read_closer.go => read_closer.go (96%) rename consumer/shard_iterator.go => shard_iterator.go (99%) rename consumer/shard_iterator_test.go => shard_iterator_test.go (99%) rename producer/shared_types.go => shared_types.go (60%) diff --git a/config.go b/config.go deleted file mode 100644 index 448778b..0000000 --- a/config.go +++ /dev/null @@ -1,88 +0,0 @@ -package kinetic - -import ( - "net/http" - "time" - - "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/aws/credentials" - "github.com/aws/aws-sdk-go/aws/session" -) - -// Config is used to configure a Kinetic instance -type Config struct { - *AwsOptions - *kineticOptions - LogLevel aws.LogLevelType -} - -// NewConfig creates a new instance of Config -func NewConfig() *Config { - return &Config{ - AwsOptions: DefaultAwsOptions(), - kineticOptions: &kineticOptions{}, - LogLevel: aws.LogOff, - } -} - -// SetLogLevel configures both the SDK and Kinetic log levels. -func (c *Config) SetLogLevel(logLevel aws.LogLevelType) { - c.AwsOptions.SetLogLevel(logLevel) - c.LogLevel = logLevel & 0xffff0000 -} - -// AwsOptions helps configure an aws.Config and session.Session -type AwsOptions struct { - AwsConfig *aws.Config -} - -// DefaultAwsOptions initializes the default aws.Config struct -func DefaultAwsOptions() *AwsOptions { - return &AwsOptions{ - AwsConfig: aws.NewConfig().WithHTTPClient( - &http.Client{ - Timeout: 2 * time.Minute, - }, - ), - } -} - -// SetCredentials configures AWS credentials. -func (c *AwsOptions) SetCredentials(accessKey, secretKey, sessionToken string) { - c.AwsConfig.WithCredentials( - credentials.NewStaticCredentials(accessKey, secretKey, sessionToken), - ) -} - -// SetRegion configures the AWS region. -func (c *AwsOptions) SetRegion(region string) { - c.AwsConfig.WithRegion(region) -} - -// SetEndpoint sets the endpoint to be used by aws-sdk-go. -func (c *AwsOptions) SetEndpoint(endpoint string) { - c.AwsConfig.WithEndpoint(endpoint) -} - -// SetLogger configures the logger for Kinetic and the aws-sdk-go. -func (c *AwsOptions) SetLogger(logger aws.Logger) { - c.AwsConfig.WithLogger(logger) -} - -// SetLogLevel configures the log levels for the SDK. -func (c *AwsOptions) SetLogLevel(logLevel aws.LogLevelType) { - c.AwsConfig.WithLogLevel(logLevel & 0xffff) -} - -// SetHTTPClientTimeout configures the HTTP timeout for the SDK. -func (c *AwsOptions) SetHTTPClientTimeout(timeout time.Duration) { - c.AwsConfig.WithHTTPClient(&http.Client{ - Timeout: timeout, - }) -} - -// GetSession creates an instance of the session.Session to be used when creating service -// clients in aws-sdk-go. -func (c *AwsOptions) GetSession() (*session.Session, error) { - return session.NewSession(c.AwsConfig) -} diff --git a/config_test.go b/config_test.go deleted file mode 100644 index 392be51..0000000 --- a/config_test.go +++ /dev/null @@ -1,93 +0,0 @@ -package kinetic - -import ( - . "github.com/smartystreets/goconvey/convey" - - "fmt" - "testing" - "time" - - "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/aws/session" -) - -func getSession(config *Config) *session.Session { - sess, err := config.GetSession() - So(err, ShouldBeNil) - So(sess, ShouldNotBeNil) - return sess -} - -func TestNewConfig(t *testing.T) { - Convey("given a new kinetic config", t, func() { - config := NewConfig() - - Convey("check the default values for its non-zero config", func() { - So(config.AwsConfig.HTTPClient.Timeout, ShouldEqual, 2*time.Minute) - So(config.AwsConfig.LogLevel.Value(), ShouldEqual, aws.LogOff) - So(config.LogLevel.Value(), ShouldEqual, LogOff) - }) - - Convey("check that we can retrieve an aws.Session from it ", func() { - getSession(config) - }) - - Convey("check that we can set credentials", func() { - config.SetCredentials("access-key", "secret-key", "security-token") - sess := getSession(config) - creds, err := sess.Config.Credentials.Get() - So(err, ShouldBeNil) - So(creds.AccessKeyID, ShouldEqual, "access-key") - So(creds.SecretAccessKey, ShouldEqual, "secret-key") - So(creds.SessionToken, ShouldEqual, "security-token") - }) - - Convey("check that we can set the region", func() { - config.SetRegion("my-region") - sess := getSession(config) - So(aws.StringValue(sess.Config.Region), ShouldEqual, "my-region") - }) - - Convey("check that we can set the endpoint", func() { - config.SetEndpoint("my-endpoint") - sess := getSession(config) - So(aws.StringValue(sess.Config.Endpoint), ShouldEqual, "my-endpoint") - }) - - Convey("check that we can configure a logger", func() { - var logs []string - loggerFn := func(args ...interface{}) { - logs = append(logs, fmt.Sprint(args...)) - } - config.SetLogger(aws.LoggerFunc(loggerFn)) - sess := getSession(config) - - Convey("check that basic logging should work", func() { - sess.Config.Logger.Log("one") - sess.Config.Logger.Log("two") - sess.Config.Logger.Log("three") - So(len(logs), ShouldEqual, 3) - So(logs, ShouldContain, "one") - So(logs, ShouldContain, "two") - So(logs, ShouldContain, "three") - Reset(func() { - logs = nil - }) - }) - }) - - Convey("check that we can set both the sdk and kinetic log level", func() { - ll := aws.LogDebug | aws.LogDebugWithSigning | LogDebug - config.SetLogLevel(ll) - sess := getSession(config) - So(sess.Config.LogLevel.AtLeast(aws.LogDebug), ShouldBeTrue) - So(sess.Config.LogLevel.Matches(aws.LogDebugWithSigning), ShouldBeTrue) - So(config.LogLevel.AtLeast(LogDebug), ShouldBeTrue) - }) - - Convey("check that we can set the http.Client Timeout", func() { - config.SetHTTPClientTimeout(10 * time.Minute) - So(config.AwsConfig.HTTPClient.Timeout, ShouldEqual, 10*time.Minute) - }) - }) -} diff --git a/consumer/consumer.go b/consumer.go similarity index 89% rename from consumer/consumer.go rename to consumer.go index dcc1a37..3b16fa8 100644 --- a/consumer/consumer.go +++ b/consumer.go @@ -1,4 +1,4 @@ -package consumer +package kinetic import ( "context" @@ -7,7 +7,6 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic" "golang.org/x/time/rate" ) @@ -15,15 +14,15 @@ import ( type consumerOptions struct { queueDepth int concurrency int - logLevel aws.LogLevelType // log level for configuring the LogHelper's log level - Stats StatsCollector // stats collection mechanism + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + Stats ConsumerStatsCollector // stats collection mechanism } func defaultConsumerOptions() *consumerOptions { return &consumerOptions{ queueDepth: 10000, concurrency: 10, - Stats: &NilStatsCollector{}, + Stats: &NilConsumerStatsCollector{}, } } @@ -35,7 +34,7 @@ func ConsumerQueueDepth(depth int) ConsumerOptionsFn { o.queueDepth = depth return nil } - return kinetic.ErrInvalidQueueDepth + return ErrInvalidQueueDepth } } @@ -45,7 +44,7 @@ func ConsumerConcurrency(count int) ConsumerOptionsFn { o.concurrency = count return nil } - return kinetic.ErrInvalidConcurrency + return ErrInvalidConcurrency } } @@ -56,7 +55,7 @@ func ConsumerLogLevel(ll aws.LogLevelType) ConsumerOptionsFn { } } -func ConsumerStatsCollector(sc StatsCollector) ConsumerOptionsFn { +func ConsumerStats(sc ConsumerStatsCollector) ConsumerOptionsFn { return func(o *consumerOptions) error { o.Stats = sc return nil @@ -66,11 +65,11 @@ func ConsumerStatsCollector(sc StatsCollector) ConsumerOptionsFn { // Listener polls the StreamReader for messages. type Consumer struct { *consumerOptions - *kinetic.LogHelper + *LogHelper reader StreamReader txnCountRateLimiter *rate.Limiter txSizeRateLimiter *rate.Limiter - messages chan *kinetic.Message + messages chan *Message concurrencySem chan empty pipeOfDeath chan empty consuming bool @@ -85,7 +84,7 @@ func NewConsumer(c *aws.Config, r StreamReader, optionFns ...ConsumerOptionsFn) } return &Consumer{ consumerOptions: consumerOptions, - LogHelper: &kinetic.LogHelper{ + LogHelper: &LogHelper{ LogLevel: consumerOptions.logLevel, Logger: c.Logger, }, @@ -100,7 +99,7 @@ func (l *Consumer) startConsuming() bool { defer l.consumingMu.Unlock() if !l.consuming { l.consuming = true - l.messages = make(chan *kinetic.Message, l.queueDepth) + l.messages = make(chan *Message, l.queueDepth) l.concurrencySem = make(chan empty, l.concurrency) l.pipeOfDeath = make(chan empty) return true @@ -113,7 +112,7 @@ func (l *Consumer) startConsuming() bool { func (l *Consumer) shouldConsume(ctx context.Context) (bool, error) { select { case <-l.pipeOfDeath: - return false, kinetic.ErrPipeOfDeath + return false, ErrPipeOfDeath case <-ctx.Done(): return false, ctx.Err() default: @@ -135,7 +134,7 @@ func (l *Consumer) stopConsuming() { } func (l *Consumer) enqueueSingle(ctx context.Context) (int, int, error) { - n, m, err := l.reader.GetRecord(ctx, func(msg *kinetic.Message, wg *sync.WaitGroup) error { + n, m, err := l.reader.GetRecord(ctx, func(msg *Message, wg *sync.WaitGroup) error { defer wg.Done() l.messages <- msg @@ -150,7 +149,7 @@ func (l *Consumer) enqueueSingle(ctx context.Context) (int, int, error) { func (l *Consumer) enqueueBatch(ctx context.Context) (int, int, error) { n, m, err := l.reader.GetRecords(ctx, - func(msg *kinetic.Message, wg *sync.WaitGroup) error { + func(msg *Message, wg *sync.WaitGroup) error { defer wg.Done() l.messages <- msg @@ -174,7 +173,7 @@ func (l *Consumer) handleErrorLogging(err error) { } case error: switch err { - case kinetic.ErrTimeoutReadResponseBody: + case ErrTimeoutReadResponseBody: l.Stats.AddGetRecordsReadTimeout(1) l.LogError("Received error:", err.Error()) default: @@ -185,11 +184,11 @@ func (l *Consumer) handleErrorLogging(err error) { } } -// RetrieveWithContext waits for a message from the stream and returns the kinetic. Cancellation is supported through +// RetrieveWithContext waits for a message from the stream and returns the Cancellation is supported through // contexts. -func (l *Consumer) RetrieveWithContext(ctx context.Context) (*kinetic.Message, error) { +func (l *Consumer) RetrieveWithContext(ctx context.Context) (*Message, error) { if !l.startConsuming() { - return nil, kinetic.ErrAlreadyConsuming + return nil, ErrAlreadyConsuming } defer l.stopConsuming() @@ -212,7 +211,7 @@ func (l *Consumer) RetrieveWithContext(ctx context.Context) (*kinetic.Message, e } // Retrieve waits for a message from the stream and returns the value -func (l *Consumer) Retrieve() (*kinetic.Message, error) { +func (l *Consumer) Retrieve() (*Message, error) { return l.RetrieveWithContext(context.TODO()) } @@ -316,7 +315,7 @@ func (l *Consumer) ListenWithContext(ctx context.Context, fn MessageProcessor) { // couple more messages (especially since select is random in which channel is read from). l.concurrencySem <- empty{} wg.Add(1) - go func(msg *kinetic.Message) { + go func(msg *Message) { defer func() { <-l.concurrencySem }() diff --git a/consumer/kcl_reader_config.go b/consumer/kcl_reader_config.go deleted file mode 100644 index c9406be..0000000 --- a/consumer/kcl_reader_config.go +++ /dev/null @@ -1,51 +0,0 @@ -package consumer - -import ( - "github.com/aws/aws-sdk-go/aws" -) - -// KclReaderConfig is used to configure KclReader -type KclReaderConfig struct { - *kclReaderOptions - AwsConfig *aws.Config - LogLevel aws.LogLevelType -} - -// NewKclReaderConfig creates a new instance of KclReaderConfig -func NewKclReaderConfig(cfg *aws.Config) *KclReaderConfig { - return &KclReaderConfig{ - AwsConfig: cfg, - kclReaderOptions: &kclReaderOptions{ - stats: &NilStatsCollector{}, - }, - LogLevel: *cfg.LogLevel, - } -} - -// SetOnInitCallbackFn configures a callback function which is run prior to sending a status message -// acknowledging an 'initialize' message was received / processed -func (c *KclReaderConfig) SetOnInitCallbackFn(fn func() error) { - c.onInitCallbackFn = fn -} - -// SetOnCheckpointCallbackFn configures a callback function which is run prior to sending a status message -// acknowledging an 'checkpoint' message was received / processed -func (c *KclReaderConfig) SetOnCheckpointCallbackFn(fn func() error) { - c.onCheckpointCallbackFn = fn -} - -// SetOnShutdownCallbackFn configures a callback function which is run prior to sending a status message -// acknowledging a 'shutdown' message was received / processed -func (c *KclReaderConfig) SetOnShutdownCallbackFn(fn func() error) { - c.onShutdownCallbackFn = fn -} - -// SetStatsCollector configures a listener to handle listener metrics. -func (c *KclReaderConfig) SetStatsCollector(stats StatsCollector) { - c.stats = stats -} - -// SetLogLevel configures both the SDK and Kinetic log levels. -func (c *KclReaderConfig) SetLogLevel(logLevel aws.LogLevelType) { - c.LogLevel = logLevel & 0xffff0000 -} diff --git a/consumer/shared_types.go b/consumer/shared_types.go deleted file mode 100644 index d1a07be..0000000 --- a/consumer/shared_types.go +++ /dev/null @@ -1,30 +0,0 @@ -package consumer - -import ( - "context" - "sync" - - "github.com/rewardStyle/kinetic" -) - -// StreamReader is an interface that abstracts out a stream reader. -type StreamReader interface { - GetRecord(context.Context, MessageHandler) (int, int, error) - GetRecords(context.Context, MessageHandler) (int, int, error) -} - -// empty is used a as a dummy type for semaphore channels and the pipe of death channel. -type empty struct{} - -// MessageProcessor defines the signature of a message handler used by Listen, RetrieveFn and their associated -// *WithContext functions. MessageHandler accepts a WaitGroup so the function can be run as a blocking operation as -// opposed to MessageHandlerAsync. -type MessageProcessor func(*kinetic.Message, *sync.WaitGroup) error - -// MessageHandler defines the signature of a message handler used by GetRecord() and GetRecords(). MessageHandler -// accepts a WaitGroup so the function can be run as a blocking operation as opposed to MessageHandlerAsync. -type MessageHandler func(*kinetic.Message, *sync.WaitGroup) error - -// MessageHandlerAsync defines the signature of a message handler used by GetRecord() and GetRecords(). -// MessageHandlerAsync is meant to be run asynchronously. -type MessageHandlerAsync func(*kinetic.Message) error diff --git a/consumer/consumer_stats.go b/consumer_stats.go similarity index 52% rename from consumer/consumer_stats.go rename to consumer_stats.go index 8ea9da7..6e35040 100644 --- a/consumer/consumer_stats.go +++ b/consumer_stats.go @@ -1,4 +1,4 @@ -package consumer +package kinetic import ( "log" @@ -10,13 +10,13 @@ import ( // StatsCollector allows for a collector to collect various metrics produced by // the Kinetic Listener library. This was really built with rcrowley/go-metrics // in mind. -type StatsCollector interface { +type ConsumerStatsCollector interface { AddConsumed(int) AddDelivered(int) AddProcessed(int) AddBatchSize(int) AddGetRecordsCalled(int) - AddProvisionedThroughputExceeded(int) + AddGetRecordsProvisionedThroughputExceeded(int) AddGetRecordsTimeout(int) AddGetRecordsReadTimeout(int) AddProcessedDuration(time.Duration) @@ -26,120 +26,120 @@ type StatsCollector interface { } // NilStatsCollector is a stats listener that ignores all metrics. -type NilStatsCollector struct{} +type NilConsumerStatsCollector struct{} // AddConsumed records a count of the number of messages received from AWS // Kinesis by the consumer. -func (nsc *NilStatsCollector) AddConsumed(int) {} +func (nsc *NilConsumerStatsCollector) AddConsumed(int) {} // AddDelivered records a count of the number of messages delivered to the // application by the consumer. -func (nsc *NilStatsCollector) AddDelivered(int) {} +func (nsc *NilConsumerStatsCollector) AddDelivered(int) {} // AddProcessed records a count of the number of messages processed by the // application by the consumer. This is based on a WaitGroup that is sent to // the RetrieveFn and Listen functions. Retrieve does not count processed // messages. -func (nsc *NilStatsCollector) AddProcessed(int) {} +func (nsc *NilConsumerStatsCollector) AddProcessed(int) {} // AddBatchSize records a count of the number of messages returned by // GetRecords in the consumer. -func (nsc *NilStatsCollector) AddBatchSize(int) {} +func (nsc *NilConsumerStatsCollector) AddBatchSize(int) {} // AddGetRecordsCalled records the number of times the GetRecords API was called // by the consumer. -func (nsc *NilStatsCollector) AddGetRecordsCalled(int) {} +func (nsc *NilConsumerStatsCollector) AddGetRecordsCalled(int) {} // AddProvisionedThroughputExceeded records the number of times the GetRecords // API returned a ErrCodeProvisionedThroughputExceededException by the consumer. -func (nsc *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} +func (nsc *NilConsumerStatsCollector) AddGetRecordsProvisionedThroughputExceeded(int) {} // AddGetRecordsTimeout records the number of times the GetRecords API timed out // on the HTTP level. This is influenced by the WithHTTPClientTimeout // configuration. -func (nsc *NilStatsCollector) AddGetRecordsTimeout(int) {} +func (nsc *NilConsumerStatsCollector) AddGetRecordsTimeout(int) {} // AddGetRecordsReadTimeout records the number of times the GetRecords API timed // out while reading the response body. This is influenced by the // WithGetRecordsReadTimeout configuration. -func (nsc *NilStatsCollector) AddGetRecordsReadTimeout(int) {} +func (nsc *NilConsumerStatsCollector) AddGetRecordsReadTimeout(int) {} // AddProcessedDuration records the duration to process a record. See notes on // AddProcessed. -func (nsc *NilStatsCollector) AddProcessedDuration(time.Duration) {} +func (nsc *NilConsumerStatsCollector) AddProcessedDuration(time.Duration) {} // AddGetRecordsDuration records the duration that the GetRecords API request // took. Only the times of successful calls are measured. -func (nsc *NilStatsCollector) AddGetRecordsDuration(time.Duration) {} +func (nsc *NilConsumerStatsCollector) AddGetRecordsDuration(time.Duration) {} // AddGetRecordsReadResponseDuration records the duration that it took to read // the response body of a GetRecords API request. -func (nsc *NilStatsCollector) AddGetRecordsReadResponseDuration(time.Duration) {} +func (nsc *NilConsumerStatsCollector) AddGetRecordsReadResponseDuration(time.Duration) {} // AddGetRecordsUnmarshalDuration records the duration that it took to unmarshal // the response body of a GetRecords API request. -func (nsc *NilStatsCollector) AddGetRecordsUnmarshalDuration(time.Duration) {} +func (nsc *NilConsumerStatsCollector) AddGetRecordsUnmarshalDuration(time.Duration) {} // Metric names to be exported const ( - MetricsConsumed = "kinetic.consumer.consumed" - MetricsDelivered = "kinetic.consumer.delivered" - MetricsProcessed = "kinetic.consumer.processed" - MetricsBatchSize = "kinetic.consumer.batchsize" - MetricsSent = "kinetic.consumer.sent" - MetricsProvisionedThroughputExceeded = "kinetic.consumer.provisionedthroughputexceeded" - MetricsGetRecordsTimeout = "kinetic.consumer.getrecords.timeout" - MetricsGetRecordsReadTimeout = "kinetic.consumer.getrecords.readtimeout" - MetricsProcessedDuration = "kinetic.consumer.processed.duration" - MetricsGetRecordsDuration = "kinetic.consumer.getrecords.duration" - MetricsGetRecordsReadResponseDuration = "kinetic.consumer.getrecords.readresponse.duration" - MetricsGetRecordsUnmarshalDuration = "kinetic.consumer.getrecords.unmarshal.duration" + MetricsConsumed = "kinetic.consumer.consumed" + MetricsDelivered = "kinetic.consumer.delivered" + MetricsProcessed = "kinetic.consumer.processed" + MetricsBatchSize = "kinetic.consumer.batchsize" + MetricsSent = "kinetic.consumer.sent" + MetricsGetRecordsProvisionedThroughputExceeded = "kinetic.consumer.getrecords.provisionedthroughputexceeded" + MetricsGetRecordsTimeout = "kinetic.consumer.getrecords.timeout" + MetricsGetRecordsReadTimeout = "kinetic.consumer.getrecords.readtimeout" + MetricsProcessedDuration = "kinetic.consumer.processed.duration" + MetricsGetRecordsDuration = "kinetic.consumer.getrecords.duration" + MetricsGetRecordsReadResponseDuration = "kinetic.consumer.getrecords.readresponse.duration" + MetricsGetRecordsUnmarshalDuration = "kinetic.consumer.getrecords.unmarshal.duration" ) // DefaultStatsCollector is a type that implements the listener's StatsCollector interface using the // rcrowley/go-metrics library -type DefaultStatsCollector struct { - Consumed metrics.Counter - Delivered metrics.Counter - Processed metrics.Counter - BatchSize metrics.Counter - GetRecordsCalled metrics.Counter - ProvisionedThroughputExceeded metrics.Counter - GetRecordsTimeout metrics.Counter - GetRecordsReadTimeout metrics.Counter - ProcessedDuration metrics.Gauge - GetRecordsDuration metrics.Gauge - GetRecordsReadResponseDuration metrics.Gauge - GetRecordsUnmarshalDuration metrics.Gauge +type DefaultConsumerStatsCollector struct { + Consumed metrics.Counter + Delivered metrics.Counter + Processed metrics.Counter + BatchSize metrics.Counter + GetRecordsCalled metrics.Counter + GetRecordsProvisionedThroughputExceeded metrics.Counter + GetRecordsTimeout metrics.Counter + GetRecordsReadTimeout metrics.Counter + ProcessedDuration metrics.Gauge + GetRecordsDuration metrics.Gauge + GetRecordsReadResponseDuration metrics.Gauge + GetRecordsUnmarshalDuration metrics.Gauge } // NewDefaultStatsCollector instantiates a new DefaultStatsCollector object -func NewDefaultStatsCollector(r metrics.Registry) *DefaultStatsCollector { - return &DefaultStatsCollector{ - Consumed: metrics.GetOrRegisterCounter(MetricsConsumed, r), - Delivered: metrics.GetOrRegisterCounter(MetricsDelivered, r), - Processed: metrics.GetOrRegisterCounter(MetricsProcessed, r), - BatchSize: metrics.GetOrRegisterCounter(MetricsBatchSize, r), - GetRecordsCalled: metrics.GetOrRegisterCounter(MetricsSent, r), - ProvisionedThroughputExceeded: metrics.GetOrRegisterCounter(MetricsProvisionedThroughputExceeded, r), - GetRecordsTimeout: metrics.GetOrRegisterCounter(MetricsGetRecordsTimeout, r), - GetRecordsReadTimeout: metrics.GetOrRegisterCounter(MetricsGetRecordsReadTimeout, r), - ProcessedDuration: metrics.GetOrRegisterGauge(MetricsProcessedDuration, r), - GetRecordsDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsDuration, r), - GetRecordsReadResponseDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsReadResponseDuration, r), - GetRecordsUnmarshalDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsUnmarshalDuration, r), +func NewDefaultConsumerStatsCollector(r metrics.Registry) *DefaultConsumerStatsCollector { + return &DefaultConsumerStatsCollector{ + Consumed: metrics.GetOrRegisterCounter(MetricsConsumed, r), + Delivered: metrics.GetOrRegisterCounter(MetricsDelivered, r), + Processed: metrics.GetOrRegisterCounter(MetricsProcessed, r), + BatchSize: metrics.GetOrRegisterCounter(MetricsBatchSize, r), + GetRecordsCalled: metrics.GetOrRegisterCounter(MetricsSent, r), + GetRecordsProvisionedThroughputExceeded: metrics.GetOrRegisterCounter(MetricsGetRecordsProvisionedThroughputExceeded, r), + GetRecordsTimeout: metrics.GetOrRegisterCounter(MetricsGetRecordsTimeout, r), + GetRecordsReadTimeout: metrics.GetOrRegisterCounter(MetricsGetRecordsReadTimeout, r), + ProcessedDuration: metrics.GetOrRegisterGauge(MetricsProcessedDuration, r), + GetRecordsDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsDuration, r), + GetRecordsReadResponseDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsReadResponseDuration, r), + GetRecordsUnmarshalDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsUnmarshalDuration, r), } } // AddConsumed records a count of the number of messages received from AWS // Kinesis by the consumer. -func (dsc *DefaultStatsCollector) AddConsumed(count int) { +func (dsc *DefaultConsumerStatsCollector) AddConsumed(count int) { dsc.Consumed.Inc(int64(count)) } // AddDelivered records a count of the number of messages delivered to the // application by the consumer. -func (dsc *DefaultStatsCollector) AddDelivered(count int) { +func (dsc *DefaultConsumerStatsCollector) AddDelivered(count int) { dsc.Delivered.Inc(int64(count)) } @@ -147,68 +147,68 @@ func (dsc *DefaultStatsCollector) AddDelivered(count int) { // application by the consumer. This is based on a WaitGroup that is sent to // the RetrieveFn and Listen functions. Retrieve does not count processed // messages. -func (dsc *DefaultStatsCollector) AddProcessed(count int) { +func (dsc *DefaultConsumerStatsCollector) AddProcessed(count int) { dsc.Processed.Inc(int64(count)) } // AddBatchSize records a count of the number of messages returned by // GetRecords in the consumer. -func (dsc *DefaultStatsCollector) AddBatchSize(count int) { +func (dsc *DefaultConsumerStatsCollector) AddBatchSize(count int) { dsc.BatchSize.Inc(int64(count)) } // AddGetRecordsCalled records the number of times the GetRecords API was called // by the consumer. -func (dsc *DefaultStatsCollector) AddGetRecordsCalled(count int) { +func (dsc *DefaultConsumerStatsCollector) AddGetRecordsCalled(count int) { dsc.GetRecordsCalled.Inc(int64(count)) } // AddProvisionedThroughputExceeded records the number of times the GetRecords // API returned a ErrCodeProvisionedThroughputExceededException by the consumer. -func (dsc *DefaultStatsCollector) AddProvisionedThroughputExceeded(count int) { - dsc.ProvisionedThroughputExceeded.Inc(int64(count)) +func (dsc *DefaultConsumerStatsCollector) AddGetRecordsProvisionedThroughputExceeded(count int) { + dsc.GetRecordsProvisionedThroughputExceeded.Inc(int64(count)) } // AddGetRecordsTimeout records the number of times the GetRecords API timed out // on the HTTP level. This is influenced by the WithHTTPClientTimeout // configuration. -func (dsc *DefaultStatsCollector) AddGetRecordsTimeout(count int) { +func (dsc *DefaultConsumerStatsCollector) AddGetRecordsTimeout(count int) { dsc.GetRecordsTimeout.Inc(int64(count)) } // AddGetRecordsReadTimeout records the number of times the GetRecords API timed // out while reading the response body. This is influenced by the // WithGetRecordsReadTimeout configuration. -func (dsc *DefaultStatsCollector) AddGetRecordsReadTimeout(count int) { +func (dsc *DefaultConsumerStatsCollector) AddGetRecordsReadTimeout(count int) { dsc.GetRecordsReadTimeout.Inc(int64(count)) } // AddProcessedDuration records the duration to process a record. See notes on // AddProcessed. -func (dsc *DefaultStatsCollector) AddProcessedDuration(duration time.Duration) { +func (dsc *DefaultConsumerStatsCollector) AddProcessedDuration(duration time.Duration) { dsc.ProcessedDuration.Update(duration.Nanoseconds()) } // AddGetRecordsDuration records the duration that the GetRecords API request // took. Only the times of successful calls are measured. -func (dsc *DefaultStatsCollector) AddGetRecordsDuration(duration time.Duration) { +func (dsc *DefaultConsumerStatsCollector) AddGetRecordsDuration(duration time.Duration) { dsc.GetRecordsDuration.Update(duration.Nanoseconds()) } // AddGetRecordsReadResponseDuration records the duration that it took to read // the response body of a GetRecords API request. -func (dsc *DefaultStatsCollector) AddGetRecordsReadResponseDuration(duration time.Duration) { +func (dsc *DefaultConsumerStatsCollector) AddGetRecordsReadResponseDuration(duration time.Duration) { dsc.GetRecordsReadResponseDuration.Update(duration.Nanoseconds()) } // AddGetRecordsUnmarshalDuration records the duration that it took to unmarshal // the response body of a GetRecords API request. -func (dsc *DefaultStatsCollector) AddGetRecordsUnmarshalDuration(duration time.Duration) { +func (dsc *DefaultConsumerStatsCollector) AddGetRecordsUnmarshalDuration(duration time.Duration) { dsc.GetRecordsUnmarshalDuration.Update(duration.Nanoseconds()) } // PrintStats logs the stats -func (dsc *DefaultStatsCollector) PrintStats() { +func (dsc *DefaultConsumerStatsCollector) PrintStats() { log.Printf("Listener stats: Consumed: [%d]\n", dsc.Consumed.Count()) log.Printf("Listener stats: Delivered: [%d]\n", dsc.Delivered.Count()) log.Printf("Listener stats: Processed: [%d]\n", dsc.Processed.Count()) @@ -216,7 +216,7 @@ func (dsc *DefaultStatsCollector) PrintStats() { log.Printf("Listener stats: GetRecords Called: [%d]\n", dsc.GetRecordsCalled.Count()) log.Printf("Listener stats: GetRecords Timeout: [%d]\n", dsc.GetRecordsTimeout.Count()) log.Printf("Listener stats: GetRecords Read Timeout: [%d]\n", dsc.GetRecordsReadTimeout.Count()) - log.Printf("Listener stats: Provisioned Throughput Exceeded: [%d]\n", dsc.ProvisionedThroughputExceeded.Count()) + log.Printf("Listener stats: GetRecords Provisioned Throughput Exceeded: [%d]\n", dsc.GetRecordsProvisionedThroughputExceeded.Count()) log.Printf("Listener stats: Processed Duration (ns): [%d]\n", dsc.ProcessedDuration.Value()) log.Printf("Listener stats: GetRecords Duration (ns): [%d]\n", dsc.GetRecordsDuration.Value()) log.Printf("Listener stats: GetRecords Read Response Duration (ns): [%d]\n", dsc.GetRecordsReadResponseDuration.Value()) diff --git a/consumer/consumer_stats_test.go b/consumer_stats_test.go similarity index 65% rename from consumer/consumer_stats_test.go rename to consumer_stats_test.go index d5f365c..b1c44fc 100644 --- a/consumer/consumer_stats_test.go +++ b/consumer_stats_test.go @@ -1,4 +1,4 @@ -package consumer +package kinetic import ( "math/rand" @@ -13,9 +13,9 @@ func init() { rand.Seed(time.Now().UnixNano()) } -func TestStatsCollector(t *testing.T) { +func TestConsumerStatsCollector(t *testing.T) { Convey("given a NilStatsCollector", t, func() { - var sc StatsCollector = &NilStatsCollector{} + var sc ConsumerStatsCollector = &NilConsumerStatsCollector{} So(sc, ShouldNotBeNil) Convey("check that AddConsumed does not error", func() { @@ -38,8 +38,8 @@ func TestStatsCollector(t *testing.T) { sc.AddGetRecordsCalled(1) }) - Convey("check that AddProvisionedThroughputExceeded does not error", func() { - sc.AddProvisionedThroughputExceeded(1) + Convey("check that AddGetRecordsProvisionedThroughputExceeded does not error", func() { + sc.AddGetRecordsProvisionedThroughputExceeded(1) }) Convey("check that AddGetRecordsTimeout does not error", func() { @@ -69,75 +69,75 @@ func TestStatsCollector(t *testing.T) { Convey("given a DefaultStatsCollector", t, func() { r := metrics.NewRegistry() - var sc StatsCollector = NewDefaultStatsCollector(r) + var sc ConsumerStatsCollector = NewDefaultConsumerStatsCollector(r) So(sc, ShouldNotBeNil) Convey("check that AddConsumed does not error", func() { count := rand.Int() sc.AddConsumed(count) - So(sc.(*DefaultStatsCollector).Consumed.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultConsumerStatsCollector).Consumed.Count(), ShouldEqual, int64(count)) }) Convey("check that AddDelivered does not error", func() { count := rand.Int() sc.AddDelivered(count) - So(sc.(*DefaultStatsCollector).Delivered.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultConsumerStatsCollector).Delivered.Count(), ShouldEqual, int64(count)) }) Convey("check that AddProcessed does not error", func() { count := rand.Int() sc.AddProcessed(count) - So(sc.(*DefaultStatsCollector).Processed.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultConsumerStatsCollector).Processed.Count(), ShouldEqual, int64(count)) }) Convey("check that AddBatchSize does not error", func() { count := rand.Int() sc.AddBatchSize(count) - So(sc.(*DefaultStatsCollector).BatchSize.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultConsumerStatsCollector).BatchSize.Count(), ShouldEqual, int64(count)) }) Convey("check that AddGetRecordsCalled does not error", func() { count := rand.Int() sc.AddGetRecordsCalled(count) - So(sc.(*DefaultStatsCollector).GetRecordsCalled.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultConsumerStatsCollector).GetRecordsCalled.Count(), ShouldEqual, int64(count)) }) - Convey("check that AddProvisionedThroughputExceeded does not error", func() { + Convey("check that AddGetRecordsProvisionedThroughputExceeded does not error", func() { count := rand.Int() - sc.AddProvisionedThroughputExceeded(count) - So(sc.(*DefaultStatsCollector).ProvisionedThroughputExceeded.Count(), ShouldEqual, int64(count)) + sc.AddGetRecordsProvisionedThroughputExceeded(count) + So(sc.(*DefaultConsumerStatsCollector).GetRecordsProvisionedThroughputExceeded.Count(), ShouldEqual, int64(count)) }) Convey("check that AddGetRecordsTimeout does not error", func() { count := rand.Int() sc.AddGetRecordsTimeout(count) - So(sc.(*DefaultStatsCollector).GetRecordsTimeout.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultConsumerStatsCollector).GetRecordsTimeout.Count(), ShouldEqual, int64(count)) }) Convey("check that AddGetRecordsReadTimeout does not error", func() { count := rand.Int() sc.AddGetRecordsReadTimeout(count) - So(sc.(*DefaultStatsCollector).GetRecordsReadTimeout.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultConsumerStatsCollector).GetRecordsReadTimeout.Count(), ShouldEqual, int64(count)) }) Convey("check that AddProcessedDuration does not error", func() { sc.AddProcessedDuration(time.Second) - So(sc.(*DefaultStatsCollector).ProcessedDuration.Value(), ShouldEqual, 1000000000) + So(sc.(*DefaultConsumerStatsCollector).ProcessedDuration.Value(), ShouldEqual, 1000000000) }) Convey("check that AddGetRecordsDuration does not error", func() { sc.AddGetRecordsDuration(time.Second) - So(sc.(*DefaultStatsCollector).GetRecordsDuration.Value(), ShouldEqual, 1000000000) + So(sc.(*DefaultConsumerStatsCollector).GetRecordsDuration.Value(), ShouldEqual, 1000000000) }) Convey("check that AddGetRecordsReadResponseDuration does not error", func() { sc.AddGetRecordsReadResponseDuration(time.Second) - So(sc.(*DefaultStatsCollector).GetRecordsReadResponseDuration.Value(), ShouldEqual, 1000000000) + So(sc.(*DefaultConsumerStatsCollector).GetRecordsReadResponseDuration.Value(), ShouldEqual, 1000000000) }) Convey("check that AddGetRecordsUnmarshalDuration does not error", func() { sc.AddGetRecordsUnmarshalDuration(time.Second) - So(sc.(*DefaultStatsCollector).GetRecordsUnmarshalDuration.Value(), ShouldEqual, 1000000000) + So(sc.(*DefaultConsumerStatsCollector).GetRecordsUnmarshalDuration.Value(), ShouldEqual, 1000000000) }) }) } diff --git a/consumer/consumer_test.go b/consumer_test.go similarity index 87% rename from consumer/consumer_test.go rename to consumer_test.go index 4e7ebc8..6160318 100644 --- a/consumer/consumer_test.go +++ b/consumer_test.go @@ -1,4 +1,4 @@ -package consumer +package kinetic import ( . "github.com/smartystreets/goconvey/convey" @@ -15,9 +15,6 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" "github.com/aws/aws-sdk-go/service/kinesis" - - "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/consumer" ) func putRecord(l *Consumer, b []byte) (*string, error) { @@ -34,11 +31,11 @@ func putRecord(l *Consumer, b []byte) (*string, error) { func TestListener(t *testing.T) { Convey("given a listener", t, func() { - k, err := kinetic.New(func(c *kinetic.Config) { - c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") - c.SetRegion("some-region") - c.SetEndpoint("http://127.0.0.1:4567") - }) + k, err := NewKinetic( + KineticAwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + KineticAwsConfigRegion("some-region"), + KineticAwsConfigEndpoint("http://127.0.0.1:4567"), + ) stream := "some-listener-stream" @@ -56,19 +53,19 @@ func TestListener(t *testing.T) { So(k.Session, ShouldNotBeNil) So(k.Session.Config, ShouldNotBeNil) r, err := NewKinesisReader(k.Session.Config, stream, shards[0], - consumer.KinesisReaderBatchSize(10000), - //consumer.KinesisReaderShardIterator(), - consumer.KinesisReaderResponseReadTimeout(time.Second), - consumer.KinesisReaderLogLevel(aws.LogOff), + KinesisReaderBatchSize(10000), + //KinesisReaderShardIterator(), + KinesisReaderResponseReadTimeout(time.Second), + KinesisReaderLogLevel(aws.LogOff), ) So(r, ShouldNotBeNil) So(err, ShouldBeNil) l, err := NewConsumer(k.Session.Config, r, - consumer.ConsumerQueueDepth(10), - consumer.ConsumerConcurrency(10), - consumer.ConsumerLogLevel(aws.LogOff), - //consumer.ConsumerStatsCollector(lsc), + ConsumerQueueDepth(10), + ConsumerConcurrency(10), + ConsumerLogLevel(aws.LogOff), + ConsumerStats(&NilConsumerStatsCollector{}), ) So(l, ShouldNotBeNil) So(err, ShouldBeNil) @@ -87,12 +84,12 @@ func TestListener(t *testing.T) { Convey("check that setting an empty shard iterator returns an error", func() { err := l.reader.(*KinesisReader).setNextShardIterator("") - So(err, ShouldEqual, kinetic.ErrEmptyShardIterator) + So(err, ShouldEqual, ErrEmptyShardIterator) }) Convey("check that setting an empty sequence number returns an error", func() { err := l.reader.(*KinesisReader).setSequenceNumber("") - So(err, ShouldEqual, kinetic.ErrEmptySequenceNumber) + So(err, ShouldEqual, ErrEmptySequenceNumber) }) Convey("check that we can get the TRIM_HORIZON shard iterator", func() { @@ -185,7 +182,7 @@ func TestListener(t *testing.T) { }() <-time.After(10 * time.Millisecond) _, err := l.Retrieve() - So(err, ShouldEqual, kinetic.ErrAlreadyConsuming) + So(err, ShouldEqual, ErrAlreadyConsuming) wg.Wait() }) @@ -195,7 +192,7 @@ func TestListener(t *testing.T) { go func() { ctx, cancel := context.WithTimeout(context.TODO(), 1000*time.Millisecond) defer cancel() - l.ListenWithContext(ctx, func(msg *kinetic.Message, wg *sync.WaitGroup) error { + l.ListenWithContext(ctx, func(msg *Message, wg *sync.WaitGroup) error { defer wg.Done() return nil }) @@ -203,7 +200,7 @@ func TestListener(t *testing.T) { }() <-time.After(10 * time.Millisecond) _, err := l.Retrieve() - So(err, ShouldEqual, kinetic.ErrAlreadyConsuming) + So(err, ShouldEqual, ErrAlreadyConsuming) wg.Wait() }) @@ -213,7 +210,7 @@ func TestListener(t *testing.T) { secs := []float64{} for i := 1; i <= 6; i++ { start := time.Now() - l.reader.GetRecord(context.TODO(), func(msg *kinetic.Message, wg *sync.WaitGroup) error { + l.reader.GetRecord(context.TODO(), func(msg *Message, wg *sync.WaitGroup) error { defer wg.Done() return nil @@ -230,7 +227,7 @@ func TestListener(t *testing.T) { data := "retrieved" _, err := putRecord(l, []byte(data)) So(err, ShouldBeNil) - err = l.RetrieveFn(func(msg *kinetic.Message, wg *sync.WaitGroup) error { + err = l.RetrieveFn(func(msg *Message, wg *sync.WaitGroup) error { defer wg.Done() called = true @@ -251,7 +248,7 @@ func TestListener(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - l.Listen(func(msg *kinetic.Message, wg *sync.WaitGroup) error { + l.Listen(func(msg *Message, wg *sync.WaitGroup) error { defer wg.Done() atomic.AddInt64(&count, 1) @@ -299,7 +296,7 @@ func TestListener(t *testing.T) { go func() { ctx, cancel := context.WithCancel(context.TODO()) defer wg.Done() - l.ListenWithContext(ctx, func(m *kinetic.Message, wg *sync.WaitGroup) error { + l.ListenWithContext(ctx, func(m *Message, wg *sync.WaitGroup) error { defer wg.Done() time.AfterFunc(time.Duration(rand.Intn(10))*time.Second, func() { n, err := strconv.Atoi(string(m.Data)) diff --git a/producer/firehose_writer.go b/firehose_writer.go similarity index 80% rename from producer/firehose_writer.go rename to firehose_writer.go index bd172de..d3e686a 100644 --- a/producer/firehose_writer.go +++ b/firehose_writer.go @@ -1,4 +1,4 @@ -package producer +package kinetic import ( "context" @@ -10,8 +10,6 @@ import ( "github.com/aws/aws-sdk-go/aws/session" "github.com/aws/aws-sdk-go/service/firehose" "github.com/aws/aws-sdk-go/service/firehose/firehoseiface" - - "github.com/rewardStyle/kinetic" ) const ( @@ -20,11 +18,11 @@ const ( ) type firehoseWriterOptions struct { - msgCountRateLimit int // maximum number of records to be sent per second - msgSizeRateLimit int // maximum (transmission) size of records to be sent per second - throughputMultiplier int // integer multiplier to increase firehose throughput rate limits - logLevel aws.LogLevelType // log level for configuring the LogHelper's log level - stats StatsCollector // stats collection mechanism + msgCountRateLimit int // maximum number of records to be sent per second + msgSizeRateLimit int // maximum (transmission) size of records to be sent per second + throughputMultiplier int // integer multiplier to increase firehose throughput rate limits + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + Stats ProducerStatsCollector // stats collection mechanism } func defaultFirehoseWriterOptions() *firehoseWriterOptions { @@ -33,7 +31,7 @@ func defaultFirehoseWriterOptions() *firehoseWriterOptions { msgSizeRateLimit: firehoseMsgSizeRateLimit, throughputMultiplier: 1, logLevel: aws.LogOff, - stats: &NilStatsCollector{}, + Stats: &NilProducerStatsCollector{}, } } @@ -45,7 +43,7 @@ func FirehoseWriterMsgCountRateLimit(limit int) FireHoseWriterOptionsFn { o.msgCountRateLimit = limit return nil } - return kinetic.ErrInvalidMsgCountRateLimit + return ErrInvalidMsgCountRateLimit } } @@ -55,7 +53,7 @@ func FirehoseWriterMsgSizeRateLimit(limit int) FireHoseWriterOptionsFn { o.msgSizeRateLimit = limit return nil } - return kinetic.ErrInvalidMsgSizeRateLimit + return ErrInvalidMsgSizeRateLimit } } @@ -65,7 +63,7 @@ func FirehoseWriterThroughputMultiplier(multiplier int) FireHoseWriterOptionsFn o.throughputMultiplier = multiplier return nil } - return kinetic.ErrInvalidThroughputMultiplier + return ErrInvalidThroughputMultiplier } } @@ -76,9 +74,9 @@ func FirehoseWriterLogLevel(ll aws.LogLevelType) FireHoseWriterOptionsFn { } } -func FirehoseWriterStatsCollector(sc StatsCollector) FireHoseWriterOptionsFn { +func FirehoseWriterStats(sc ProducerStatsCollector) FireHoseWriterOptionsFn { return func(o *firehoseWriterOptions) error { - o.stats = sc + o.Stats = sc return nil } } @@ -86,7 +84,7 @@ func FirehoseWriterStatsCollector(sc StatsCollector) FireHoseWriterOptionsFn { // FirehoseWriter handles the API to send records to Kinesis. type FirehoseWriter struct { *firehoseWriterOptions - *kinetic.LogHelper + *LogHelper stream string client firehoseiface.FirehoseAPI } @@ -105,7 +103,7 @@ func NewFirewhoseWriter(c *aws.Config, stream string, optionFns ...FireHoseWrite stream: stream, client: firehose.New(sess), firehoseWriterOptions: firehoseWriterOptions, - LogHelper: &kinetic.LogHelper{ + LogHelper: &LogHelper{ LogLevel: firehoseWriterOptions.logLevel, Logger: c.Logger, }, @@ -113,7 +111,7 @@ func NewFirewhoseWriter(c *aws.Config, stream string, optionFns ...FireHoseWrite } // PutRecords sends a batch of records to Firehose and returns a list of records that need to be retried. -func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*kinetic.Message, fn MessageHandlerAsync) error { +func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*Message, fn MessageHandlerAsync) error { var startSendTime time.Time var startBuildTime time.Time @@ -135,7 +133,7 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*kinetic.Mes }) req.Handlers.Build.PushBack(func(r *request.Request) { - w.stats.UpdatePutRecordsBuildDuration(time.Since(startBuildTime)) + w.Stats.UpdatePutRecordsBuildDuration(time.Since(startBuildTime)) w.LogDebug("Finished PutRecords Build, took", time.Since(start)) }) @@ -145,23 +143,23 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*kinetic.Mes }) req.Handlers.Send.PushBack(func(r *request.Request) { - w.stats.UpdatePutRecordsSendDuration(time.Since(startSendTime)) + w.Stats.UpdatePutRecordsSendDuration(time.Since(startSendTime)) w.LogDebug("Finished PutRecords Send, took", time.Since(start)) }) w.LogDebug("Starting PutRecords Build/Sign request, took", time.Since(start)) - w.stats.AddPutRecordsCalled(1) + w.Stats.AddPutRecordsCalled(1) if err := req.Send(); err != nil { w.LogError("Error putting records:", err.Error()) return err } - w.stats.UpdatePutRecordsDuration(time.Since(start)) + w.Stats.UpdatePutRecordsDuration(time.Since(start)) if resp == nil { - return kinetic.ErrNilPutRecordsResponse + return ErrNilPutRecordsResponse } if resp.FailedPutCount == nil { - return kinetic.ErrNilFailedRecordCount + return ErrNilFailedRecordCount } attempted := len(messages) failed := int(aws.Int64Value(resp.FailedPutCount)) @@ -172,18 +170,18 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*kinetic.Mes if record.RecordId != nil { // TODO: per-shard metrics messages[idx].RecordID = record.RecordId - w.stats.AddSentSuccess(1) + w.Stats.AddSentSuccess(1) } else { switch aws.StringValue(record.ErrorCode) { case firehose.ErrCodeLimitExceededException: - w.stats.AddProvisionedThroughputExceeded(1) + w.Stats.AddPutRecordsProvisionedThroughputExceeded(1) default: w.LogDebug("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(record.ErrorMessage)) } messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage messages[idx].FailCount++ - w.stats.AddSentFailed(1) + w.Stats.AddSentFailed(1) fn(messages[idx]) } diff --git a/consumer/kcl_actions.go b/kcl_actions.go similarity index 94% rename from consumer/kcl_actions.go rename to kcl_actions.go index 5e4bd56..9b033d2 100644 --- a/consumer/kcl_actions.go +++ b/kcl_actions.go @@ -1,11 +1,9 @@ -package consumer +package kinetic import ( "encoding/base64" "strconv" "time" - - "github.com/rewardStyle/kinetic" ) // ActionType is used as an enum for KCL Multilang protocol action message types @@ -62,14 +60,14 @@ func (t *Timestamp) UnmarshalJSON(b []byte) error { return nil } -// ToMessage is used to transform a multilang.Record struct into a kinetic.Message struct -func (r *Record) ToMessage() *kinetic.Message { +// ToMessage is used to transform a multilang.Record struct into a Message struct +func (r *Record) ToMessage() *Message { b, err := base64.StdEncoding.DecodeString(r.Data) if err != nil { panic("There was a problem decoding kcl data") } - return &kinetic.Message{ + return &Message{ ApproximateArrivalTimestamp: &r.ApproximateArrival.Time, Data: b, PartitionKey: &r.PartitionKey, diff --git a/consumer/kcl_reader.go b/kcl_reader.go similarity index 89% rename from consumer/kcl_reader.go rename to kcl_reader.go index b733ef4..4119851 100644 --- a/consumer/kcl_reader.go +++ b/kcl_reader.go @@ -1,4 +1,4 @@ -package consumer +package kinetic import ( "bufio" @@ -11,24 +11,23 @@ import ( "sync" "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic" ) type kclReaderOptions struct { onInitCallbackFn func() error onCheckpointCallbackFn func() error onShutdownCallbackFn func() error - logLevel aws.LogLevelType // log level for configuring the LogHelper's log level - stats StatsCollector // stats collection mechanism + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + Stats ConsumerStatsCollector // stats collection mechanism } func defaultKlcReaderOptions() *kclReaderOptions { return &kclReaderOptions{ - onInitCallbackFn: func() error { return nil }, + onInitCallbackFn: func() error { return nil }, onCheckpointCallbackFn: func() error { return nil }, - onShutdownCallbackFn: func() error { return nil }, - logLevel: aws.LogOff, - stats: &NilStatsCollector{}, + onShutdownCallbackFn: func() error { return nil }, + logLevel: aws.LogOff, + Stats: &NilConsumerStatsCollector{}, } } @@ -62,9 +61,9 @@ func KlcReaderLogLevel(ll aws.LogLevelType) KlcReaderOptionsFn { } } -func KlcReaderStatsCollector(sc StatsCollector) KlcReaderOptionsFn { +func KlcReaderStats(sc ConsumerStatsCollector) KlcReaderOptionsFn { return func(o *kclReaderOptions) error { - o.stats = sc + o.Stats = sc return nil } } @@ -72,11 +71,11 @@ func KlcReaderStatsCollector(sc StatsCollector) KlcReaderOptionsFn { // KclReader handles the KCL Multilang Protocol to read records from KCL type KclReader struct { *kclReaderOptions - *kinetic.LogHelper + *LogHelper pipeOfDeath chan empty scanner *bufio.Scanner reader *bufio.Reader - msgBuffer []kinetic.Message + msgBuffer []Message } // NewKclReader creates a new stream reader to read records from KCL @@ -86,9 +85,9 @@ func NewKclReader(c *aws.Config, optionFns ...KlcReaderOptionsFn) (*KclReader, e optionFn(kclReaderOptions) } return &KclReader{ - msgBuffer: []kinetic.Message{}, + msgBuffer: []Message{}, kclReaderOptions: kclReaderOptions, - LogHelper: &kinetic.LogHelper{ + LogHelper: &LogHelper{ LogLevel: kclReaderOptions.logLevel, Logger: c.Logger, }, @@ -109,7 +108,7 @@ func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, int, batchSize = int(math.Min(float64(len(r.msgBuffer)), float64(numRecords))) } } - r.stats.AddBatchSize(batchSize) + r.Stats.AddBatchSize(batchSize) // TODO: Define the payloadSize var payloadSize int @@ -120,7 +119,7 @@ func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, int, wg.Add(1) go fn(&r.msgBuffer[0], &wg) r.msgBuffer = r.msgBuffer[1:] - r.stats.AddConsumed(1) + r.Stats.AddConsumed(1) } wg.Wait() diff --git a/consumer/kinesis_reader.go b/kinesis_reader.go similarity index 90% rename from consumer/kinesis_reader.go rename to kinesis_reader.go index 94c6caf..52e298b 100644 --- a/consumer/kinesis_reader.go +++ b/kinesis_reader.go @@ -1,4 +1,4 @@ -package consumer +package kinetic import ( "context" @@ -12,8 +12,6 @@ import ( "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" - - "github.com/rewardStyle/kinetic" ) const ( @@ -25,8 +23,8 @@ type kinesisReaderOptions struct { batchSize int shardIterator *ShardIterator responseReadTimeout time.Duration - logLevel aws.LogLevelType // log level for configuring the LogHelper's log level - stats StatsCollector // stats collection mechanism + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + Stats ConsumerStatsCollector // stats collection mechanism } func defaultKinesisReaderOptions() *kinesisReaderOptions { @@ -34,7 +32,7 @@ func defaultKinesisReaderOptions() *kinesisReaderOptions { batchSize: kinesisReaderBatchSize, shardIterator: NewShardIterator(), responseReadTimeout: time.Second, - stats: &NilStatsCollector{}, + Stats: &NilConsumerStatsCollector{}, } } @@ -46,7 +44,7 @@ func KinesisReaderBatchSize(size int) KinesisReaderOptionsFn { o.batchSize = size return nil } - return kinetic.ErrInvalidBatchSize + return ErrInvalidBatchSize } } @@ -71,9 +69,9 @@ func KinesisReaderLogLevel(ll aws.LogLevelType) KinesisReaderOptionsFn { } } -func KinesisReaderStatsCollector(sc StatsCollector) KinesisReaderOptionsFn { +func KinesisReaderStats(sc ConsumerStatsCollector) KinesisReaderOptionsFn { return func(o *kinesisReaderOptions) error { - o.stats = sc + o.Stats = sc return nil } } @@ -81,7 +79,7 @@ func KinesisReaderStatsCollector(sc StatsCollector) KinesisReaderOptionsFn { // KinesisReader handles the API to read records from Kinesis. type KinesisReader struct { *kinesisReaderOptions - *kinetic.LogHelper + *LogHelper stream string shard string throttleSem chan empty @@ -102,7 +100,7 @@ func NewKinesisReader(c *aws.Config, stream string, shard string, optionFns ...K } return &KinesisReader{ kinesisReaderOptions: kinesisReaderOptions, - LogHelper: &kinetic.LogHelper{ + LogHelper: &LogHelper{ LogLevel: kinesisReaderOptions.logLevel, Logger: c.Logger, }, @@ -135,10 +133,10 @@ func (r *KinesisReader) ensureShardIterator() error { return err } if resp == nil { - return kinetic.ErrNilGetShardIteratorResponse + return ErrNilGetShardIteratorResponse } if resp.ShardIterator == nil { - return kinetic.ErrNilShardIterator + return ErrNilShardIterator } return r.setNextShardIterator(*resp.ShardIterator) } @@ -149,7 +147,7 @@ func (r *KinesisReader) ensureShardIterator() error { // be taken to ensure that only one call to Listen and Retrieve/RetrieveFn can be running at a time. func (r *KinesisReader) setNextShardIterator(shardIterator string) error { if len(shardIterator) == 0 { - return kinetic.ErrEmptyShardIterator + return ErrEmptyShardIterator } r.nextShardIterator = shardIterator return nil @@ -164,7 +162,7 @@ func (r *KinesisReader) setNextShardIterator(shardIterator string) error { // time. func (r *KinesisReader) setSequenceNumber(sequenceNumber string) error { if len(sequenceNumber) == 0 { - return kinetic.ErrEmptySequenceNumber + return ErrEmptySequenceNumber } r.shardIterator.AtSequenceNumber(sequenceNumber) return nil @@ -199,7 +197,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch req.ApplyOptions(request.WithResponseReadTimeout(r.responseReadTimeout)) // If debug is turned on, add some handlers for GetRecords logging - if r.LogLevel.AtLeast(kinetic.LogDebug) { + if r.LogLevel.AtLeast(LogDebug) { req.Handlers.Send.PushBack(func(req *request.Request) { r.LogDebug("Finished getRecords Send, took", time.Since(start)) }) @@ -221,7 +219,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch req.HTTPResponse.Body = &ReadCloserWrapper{ ReadCloser: req.HTTPResponse.Body, OnCloseFn: func() { - r.stats.AddGetRecordsReadResponseDuration(time.Since(startReadTime)) + r.Stats.AddGetRecordsReadResponseDuration(time.Since(startReadTime)) r.LogDebug("Finished GetRecords body read, took", time.Since(start)) startUnmarshalTime = time.Now() }, @@ -231,32 +229,32 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch var payloadSize int req.Handlers.Unmarshal.PushBack(func(req *request.Request) { payloadSize += int(req.HTTPRequest.ContentLength) - r.stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) + r.Stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) r.LogDebug("Finished GetRecords Unmarshal, took", time.Since(start)) }) // Send the GetRecords request r.LogDebug("Starting GetRecords Build/Sign request, took", time.Since(start)) - r.stats.AddGetRecordsCalled(1) + r.Stats.AddGetRecordsCalled(1) if err := req.Send(); err != nil { r.LogError("Error getting records:", err) switch err.(awserr.Error).Code() { case kinesis.ErrCodeProvisionedThroughputExceededException: - r.stats.AddProvisionedThroughputExceeded(1) + r.Stats.AddGetRecordsProvisionedThroughputExceeded(1) default: r.LogDebug("Received AWS error:", err.Error()) } return 0, 0, err } - r.stats.AddGetRecordsDuration(time.Since(start)) + r.Stats.AddGetRecordsDuration(time.Since(start)) // Process Records r.LogDebug(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), r.shard, time.Since(start))) if resp == nil { - return 0, 0, kinetic.ErrNilGetRecordsResponse + return 0, 0, ErrNilGetRecordsResponse } delivered := 0 - r.stats.AddBatchSize(len(resp.Records)) + r.Stats.AddBatchSize(len(resp.Records)) for _, record := range resp.Records { if record != nil { // Allow (only) a pipeOfDeath to trigger an instance shutdown of the loop to deliver messages. @@ -269,10 +267,10 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch default: var wg sync.WaitGroup wg.Add(1) - go fn(kinetic.FromRecord(record), &wg) + go fn(FromRecord(record), &wg) wg.Wait() delivered++ - r.stats.AddConsumed(1) + r.Stats.AddConsumed(1) if record.SequenceNumber != nil { // We can safely ignore if this call returns error, as if we somehow receive an // empty sequence number from AWS, we will simply not set it. At worst, this diff --git a/consumer/kinesis_reader_test.go b/kinesis_reader_test.go similarity index 61% rename from consumer/kinesis_reader_test.go rename to kinesis_reader_test.go index 78a5dd8..a4a1f11 100644 --- a/consumer/kinesis_reader_test.go +++ b/kinesis_reader_test.go @@ -1,4 +1,4 @@ -package consumer +package kinetic import ( . "github.com/smartystreets/goconvey/convey" @@ -8,8 +8,6 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/consumer" ) func TestNewKinesisReader(t *testing.T) { @@ -18,11 +16,11 @@ func TestNewKinesisReader(t *testing.T) { Convey("given a kinetic object and kinesis stream/shard names", t, func() { // Instantiate a new kinentic object - k, err := kinetic.New(func(c *kinetic.Config) { - c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") - c.SetRegion("some-region") - c.SetEndpoint("http://127.0.0.1:4567") - }) + k, err := NewKinetic( + KineticAwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + KineticAwsConfigRegion("some-region"), + KineticAwsConfigEndpoint("http://127.0.0.1:4567"), + ) So(k, ShouldNotBeNil) So(err, ShouldBeNil) @@ -36,29 +34,29 @@ func TestNewKinesisReader(t *testing.T) { So(r.batchSize, ShouldEqual, 10000) So(r.shardIterator, ShouldNotBeNil) So(r.responseReadTimeout, ShouldEqual, time.Second) - So(r.stats, ShouldNotBeNil) + So(r.Stats, ShouldNotBeNil) }) Convey("check that we can create a new KinesisReader with configured values", func() { batchSize := rand.Int() respReadTimeout := time.Duration(rand.Int()) * time.Second - logLevel := aws.LogDebug | aws.LogDebugWithSigning | kinetic.LogDebug + logLevel := aws.LogDebug | aws.LogDebugWithSigning | LogDebug shardIterator := NewShardIterator() - myStatsCollector := &NilStatsCollector{} + myStatsCollector := &NilConsumerStatsCollector{} r, err := NewKinesisReader(k.Session.Config, stream, shard, - consumer.KinesisReaderBatchSize(batchSize), - consumer.KinesisReaderShardIterator(shardIterator), - consumer.KinesisReaderResponseReadTimeout(respReadTimeout), - consumer.KinesisReaderLogLevel(logLevel), - consumer.KinesisReaderStatsCollector(myStatsCollector), + KinesisReaderBatchSize(batchSize), + KinesisReaderShardIterator(shardIterator), + KinesisReaderResponseReadTimeout(respReadTimeout), + KinesisReaderLogLevel(logLevel), + KinesisReaderStats(myStatsCollector), ) So(r, ShouldNotBeNil) So(err, ShouldBeNil) So(r.batchSize, ShouldEqual, batchSize) So(r.responseReadTimeout, ShouldEqual, respReadTimeout) - So(r.LogLevel.AtLeast(kinetic.LogDebug), ShouldBeTrue) + So(r.LogLevel.AtLeast(LogDebug), ShouldBeTrue) So(r.shardIterator, ShouldEqual, shardIterator) - So(r.stats, ShouldEqual, myStatsCollector) + So(r.Stats, ShouldEqual, myStatsCollector) }) }) } diff --git a/producer/kinesis_writer.go b/kinesis_writer.go similarity index 83% rename from producer/kinesis_writer.go rename to kinesis_writer.go index 2ba852e..554b344 100644 --- a/producer/kinesis_writer.go +++ b/kinesis_writer.go @@ -1,4 +1,4 @@ -package producer +package kinetic import ( "context" @@ -10,8 +10,6 @@ import ( "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" - - "github.com/rewardStyle/kinetic" ) const ( @@ -20,11 +18,11 @@ const ( ) type kinesisWriterOptions struct { - responseReadTimeout time.Duration // maximum time to wait for PutRecords API call before timing out - msgCountRateLimit int // maximum number of records to be sent per second - msgSizeRateLimit int // maximum (transmission) size of records to be sent per second - logLevel aws.LogLevelType // log level for configuring the LogHelper's log level - stats StatsCollector // stats collection mechanism + responseReadTimeout time.Duration // maximum time to wait for PutRecords API call before timing out + msgCountRateLimit int // maximum number of records to be sent per second + msgSizeRateLimit int // maximum (transmission) size of records to be sent per second + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + Stats ProducerStatsCollector // stats collection mechanism } func defaultKinesisWriterOptions() *kinesisWriterOptions { @@ -33,7 +31,7 @@ func defaultKinesisWriterOptions() *kinesisWriterOptions { msgCountRateLimit: kinesisMsgCountRateLimit, msgSizeRateLimit: kinesisMsgSizeRateLimit, logLevel: aws.LogOff, - stats: &NilStatsCollector{}, + Stats: &NilProducerStatsCollector{}, } } @@ -52,7 +50,7 @@ func KinesisWriterMsgCountRateLimit(limit int) KinesisWriterOptionsFn { o.msgSizeRateLimit = limit return nil } - return kinetic.ErrInvalidMsgSizeRateLimit + return ErrInvalidMsgSizeRateLimit } } @@ -62,7 +60,7 @@ func KinesisWriterMsgSizeRateLimit(limit int) KinesisWriterOptionsFn { o.msgSizeRateLimit = limit return nil } - return kinetic.ErrInvalidMsgSizeRateLimit + return ErrInvalidMsgSizeRateLimit } } @@ -73,9 +71,9 @@ func KinesisWriterLogLevel(ll aws.LogLevelType) KinesisWriterOptionsFn { } } -func KinesisWriterStatsCollector(sc StatsCollector) KinesisWriterOptionsFn { +func KinesisWriterStats(sc ProducerStatsCollector) KinesisWriterOptionsFn { return func(o *kinesisWriterOptions) error { - o.stats = sc + o.Stats = sc return nil } } @@ -83,7 +81,7 @@ func KinesisWriterStatsCollector(sc StatsCollector) KinesisWriterOptionsFn { // KinesisWriter handles the API to send records to Kinesis. type KinesisWriter struct { *kinesisWriterOptions - *kinetic.LogHelper + *LogHelper stream string client kinesisiface.KinesisAPI } @@ -102,7 +100,7 @@ func NewKinesisWriter(c *aws.Config, stream string, optionFns ...KinesisWriterOp stream: stream, client: kinesis.New(sess), kinesisWriterOptions: kinesisWriterOptions, - LogHelper: &kinetic.LogHelper{ + LogHelper: &LogHelper{ LogLevel: kinesisWriterOptions.logLevel, Logger: c.Logger, }, @@ -110,7 +108,7 @@ func NewKinesisWriter(c *aws.Config, stream string, optionFns ...KinesisWriterOp } // PutRecords sends a batch of records to Kinesis and returns a list of records that need to be retried. -func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*kinetic.Message, fn MessageHandlerAsync) error { +func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*Message, fn MessageHandlerAsync) error { var startSendTime time.Time var startBuildTime time.Time @@ -133,7 +131,7 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*kinetic.Mess }) req.Handlers.Build.PushBack(func(r *request.Request) { - w.stats.UpdatePutRecordsBuildDuration(time.Since(startBuildTime)) + w.Stats.UpdatePutRecordsBuildDuration(time.Since(startBuildTime)) w.LogDebug("Finished PutRecords Build, took", time.Since(start)) }) @@ -143,23 +141,23 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*kinetic.Mess }) req.Handlers.Send.PushBack(func(r *request.Request) { - w.stats.UpdatePutRecordsSendDuration(time.Since(startSendTime)) + w.Stats.UpdatePutRecordsSendDuration(time.Since(startSendTime)) w.LogDebug("Finished PutRecords Send, took", time.Since(start)) }) w.LogDebug("Starting PutRecords Build/Sign request, took", time.Since(start)) - w.stats.AddPutRecordsCalled(1) + w.Stats.AddPutRecordsCalled(1) if err := req.Send(); err != nil { w.LogError("Error putting records:", err.Error()) return err } - w.stats.UpdatePutRecordsDuration(time.Since(start)) + w.Stats.UpdatePutRecordsDuration(time.Since(start)) if resp == nil { - return kinetic.ErrNilPutRecordsResponse + return ErrNilPutRecordsResponse } if resp.FailedRecordCount == nil { - return kinetic.ErrNilFailedRecordCount + return ErrNilFailedRecordCount } attempted := len(messages) failed := int(aws.Int64Value(resp.FailedRecordCount)) @@ -171,18 +169,18 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*kinetic.Mess // TODO: per-shard metrics messages[idx].SequenceNumber = record.SequenceNumber messages[idx].ShardID = record.ShardId - w.stats.AddSentSuccess(1) + w.Stats.AddSentSuccess(1) } else { switch aws.StringValue(record.ErrorCode) { case kinesis.ErrCodeProvisionedThroughputExceededException: - w.stats.AddProvisionedThroughputExceeded(1) + w.Stats.AddPutRecordsProvisionedThroughputExceeded(1) default: w.LogDebug("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(record.ErrorMessage)) } messages[idx].ErrorCode = record.ErrorCode messages[idx].ErrorMessage = record.ErrorMessage messages[idx].FailCount++ - w.stats.AddSentFailed(1) + w.Stats.AddSentFailed(1) fn(messages[idx]) } @@ -212,10 +210,10 @@ func (w *KinesisWriter) getConcurrencyMultiplier() (int, error) { return 0, err } if resp == nil { - return 0, kinetic.ErrNilDescribeStreamResponse + return 0, ErrNilDescribeStreamResponse } if resp.StreamDescription == nil { - return 0, kinetic.ErrNilStreamDescription + return 0, ErrNilStreamDescription } // maps shardID to a boolean that indicates whether or not the shard is a parent shard or an adjacent parent shard diff --git a/kinetic.go b/kinetic.go index 00fa4da..97b8aef 100644 --- a/kinetic.go +++ b/kinetic.go @@ -12,9 +12,76 @@ import ( "github.com/aws/aws-sdk-go/service/firehose/firehoseiface" "github.com/aws/aws-sdk-go/service/kinesis" "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" + "net/http" + "github.com/aws/aws-sdk-go/aws/credentials" ) type kineticOptions struct { + awsConfig *aws.Config // + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level +} + +func defaultKineticOptions() *kineticOptions { + return &kineticOptions{ + awsConfig: aws.NewConfig().WithHTTPClient( + &http.Client{ + Timeout: 2 * time.Minute, + }), + logLevel: aws.LogOff, + } +} + +type KineticOptionsFn func(*kineticOptions) error + +func KineticAwsConfigCredentials(accessKey, secretKey, sessionToken string) KineticOptionsFn { + return func(o *kineticOptions) error { + o.awsConfig.WithCredentials(credentials.NewStaticCredentials(accessKey, secretKey, sessionToken)) + return nil + } +} + +func KineticAwsConfigRegion(region string) KineticOptionsFn { + return func(o *kineticOptions) error { + o.awsConfig.WithRegion(region) + return nil + } +} + +func KineticAwsConfigEndpoint(endpoint string) KineticOptionsFn { + return func(o *kineticOptions) error { + o.awsConfig.WithEndpoint(endpoint) + return nil + } +} + +func KineticAwsConfigLogger(logger aws.Logger) KineticOptionsFn { + return func(o *kineticOptions) error { + o.awsConfig.WithLogger(logger) + return nil + } +} + +func KineticAwsConfigLogLevel(logLevel aws.LogLevelType) KineticOptionsFn { + return func(o *kineticOptions) error { + o.awsConfig.WithLogLevel(logLevel) + return nil + } +} + +func KineticAwsConfigHttpClientTimeout(timeout time.Duration) KineticOptionsFn { + return func(o *kineticOptions) error { + o.awsConfig.WithHTTPClient(&http.Client{ + Timeout: timeout, + }) + return nil + } +} + +func KineticLogLevel(logLevel aws.LogLevelType) KineticOptionsFn { + return func(o *kineticOptions) error { + o.logLevel = logLevel & 0xffff0000 + return nil + } } // Kinetic represents a kinesis and firehose client and provides some utility @@ -22,25 +89,26 @@ type kineticOptions struct { type Kinetic struct { *kineticOptions *LogHelper - clientMu sync.Mutex fclient firehoseiface.FirehoseAPI kclient kinesisiface.KinesisAPI Session *session.Session } -// New creates a new instance of Kientic. -func New(fn func(*Config)) (*Kinetic, error) { - config := NewConfig() - fn(config) - sess, err := config.GetSession() +// New creates a new instance of Kinetic. +func NewKinetic(optionFns ...KineticOptionsFn) (*Kinetic, error) { + kineticOptions := defaultKineticOptions() + for _, optionFn := range optionFns { + optionFn(kineticOptions) + } + sess, err := session.NewSession(kineticOptions.awsConfig) if err != nil { return nil, err } return &Kinetic{ - kineticOptions: config.kineticOptions, + kineticOptions: kineticOptions, LogHelper: &LogHelper{ - LogLevel: config.LogLevel, + LogLevel: kineticOptions.logLevel, Logger: sess.Config.Logger, }, Session: sess, diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index 4e1ee4b..903c0f7 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -12,10 +12,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" - "github.com/rewardStyle/kinetic/consumer" - "github.com/rewardStyle/kinetic/producer" "github.com/stretchr/testify/assert" - "github.com/rewardStyle/kinetic" ) type TestMessage struct { @@ -94,11 +91,11 @@ func TestKineticIntegration(t *testing.T) { streamData := NewStreamData() // Instantiate a new kinentic object - k, err := New(func(c *Config) { - c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") - c.SetRegion("some-region") - c.SetEndpoint("http://127.0.0.1:4567") - }) + k, err := NewKinetic( + KineticAwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + KineticAwsConfigRegion("some-region"), + KineticAwsConfigEndpoint("http://127.0.0.1:4567"), + ) assert.NotNil(t, k) assert.Nil(t, err) @@ -130,52 +127,52 @@ func TestKineticIntegration(t *testing.T) { log.Printf("Shard Name: %s\n", shards[0]) // Create a new kinesis stream writer - w, err := producer.NewKinesisWriter(k.Session.Config, stream, - producer.KinesisWriterResponseReadTimeout(time.Second), - producer.KinesisWriterMsgCountRateLimit(1000), - producer.KinesisWriterMsgSizeRateLimit(1000000), - producer.KinesisWriterLogLevel(kinetic.LogDebug), + w, err := NewKinesisWriter(k.Session.Config, stream, + KinesisWriterResponseReadTimeout(time.Second), + KinesisWriterMsgCountRateLimit(1000), + KinesisWriterMsgSizeRateLimit(1000000), + KinesisWriterLogLevel(LogDebug), ) if err != nil { log.Fatalf("Unable to create a new kinesis stream writer due to: %v\n", err) } // Create a new kinetic producer - p, err := producer.NewProducer(k.Session.Config, w, - producer.ProducerBatchSize(5), - producer.ProducerBatchTimeout(time.Second), - producer.ProducerMaxRetryAttempts(3), - producer.ProducerQueueDepth(10000), - producer.ProducerConcurrency(3), - producer.ProducerShardCheckFrequency(time.Minute), - producer.ProducerDataSpillFn(func(msg *kinetic.Message) error { + p, err := NewProducer(k.Session.Config, w, + ProducerBatchSize(5), + ProducerBatchTimeout(time.Second), + ProducerMaxRetryAttempts(3), + ProducerQueueDepth(10000), + ProducerConcurrency(3), + ProducerShardCheckFrequency(time.Minute), + ProducerDataSpillFn(func(msg *Message) error { //log.Printf("Message was dropped: [%s]\n", string(msg.Data)) return nil }), - producer.ProducerLogLevel(aws.LogOff), - //producer.ProducerStatsCollector(), + ProducerLogLevel(aws.LogOff), + ProducerStats(&NilProducerStatsCollector{}), ) assert.NotNil(t, p) assert.Nil(t, err) assert.NotNil(t, k.Session) assert.NotNil(t, k.Session.Config) - r, err := consumer.NewKinesisReader(k.Session.Config, stream, shards[0], - //consumer.KinesisReaderBatchSize(), - //consumer.KinesisReaderShardIterator(), - consumer.KinesisReaderResponseReadTimeout(time.Second), - //consumer.KinesisReaderLogLevel(), - //consumer.KinesisReaderStatsCollector(), + r, err := NewKinesisReader(k.Session.Config, stream, shards[0], + //KinesisReaderBatchSize(), + //KinesisReaderShardIterator(), + KinesisReaderResponseReadTimeout(time.Second), + //KinesisReaderLogLevel(), + //KinesisReaderStatsCollector(), ) assert.NotNil(t, r) assert.NoError(t, err) // Create a new kinetic listener - l, err := consumer.NewConsumer(k.Session.Config, r, - consumer.ConsumerQueueDepth(20), - consumer.ConsumerConcurrency(10), - consumer.ConsumerLogLevel(aws.LogOff), - //consumer.ConsumerStatsCollector(), + l, err := NewConsumer(k.Session.Config, r, + ConsumerQueueDepth(20), + ConsumerConcurrency(10), + ConsumerLogLevel(aws.LogOff), + ConsumerStats(&NilConsumerStatsCollector{}), ) assert.NotNil(t, l) assert.Nil(t, err) diff --git a/kinetic_test.go b/kinetic_test.go index a59948c..0ace2e4 100644 --- a/kinetic_test.go +++ b/kinetic_test.go @@ -14,11 +14,11 @@ import ( func TestKinetic(t *testing.T) { Convey("given a kinetic object", t, func() { - k, err := New(func(c *Config) { - c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") - c.SetRegion("some-region") - c.SetEndpoint("http://127.0.0.1:4567") - }) + k, err := NewKinetic( + KineticAwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + KineticAwsConfigRegion("some-region"), + KineticAwsConfigEndpoint("http://127.0.0.1:4567"), + ) So(k, ShouldNotBeNil) So(err, ShouldBeNil) diff --git a/producer/producer.go b/producer.go similarity index 86% rename from producer/producer.go rename to producer.go index af7b690..9f1649d 100644 --- a/producer/producer.go +++ b/producer.go @@ -1,4 +1,4 @@ -package producer +package kinetic import ( "context" @@ -8,7 +8,6 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/awserr" - "github.com/rewardStyle/kinetic" "golang.org/x/time/rate" ) @@ -18,15 +17,15 @@ const ( // producerOptions holds all of the configurable settings for a Producer type producerOptions struct { - batchSize int // maximum message capacity per request - batchTimeout time.Duration // maximum time duration to wait for incoming messages - queueDepth int // maximum number of messages to enqueue in the message queue - maxRetryAttempts int // maximum number of retry attempts for failed messages - concurrency int // number of concurrent workers per shard - shardCheckFreq time.Duration // frequency (specified as a duration) with which to check the the shard size - dataSpillFn MessageHandlerAsync // callback function for handling dropped messages that the producer was unable to send to the stream - logLevel aws.LogLevelType // log level for configuring the LogHelper's log level - Stats StatsCollector // stats collection mechanism + batchSize int // maximum message capacity per request + batchTimeout time.Duration // maximum time duration to wait for incoming messages + queueDepth int // maximum number of messages to enqueue in the message queue + maxRetryAttempts int // maximum number of retry attempts for failed messages + concurrency int // number of concurrent workers per shard + shardCheckFreq time.Duration // frequency (specified as a duration) with which to check the the shard size + dataSpillFn MessageHandlerAsync // callback function for handling dropped messages that the producer was unable to send to the stream + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + Stats ProducerStatsCollector // stats collection mechanism } func defaultProducerOptions() *producerOptions { @@ -37,9 +36,9 @@ func defaultProducerOptions() *producerOptions { maxRetryAttempts: 10, concurrency: 3, shardCheckFreq: time.Minute, - dataSpillFn: func(*kinetic.Message) error { return nil }, + dataSpillFn: func(*Message) error { return nil }, logLevel: aws.LogOff, - Stats: &NilStatsCollector{}, + Stats: &NilProducerStatsCollector{}, } } @@ -51,7 +50,7 @@ func ProducerBatchSize(size int) ProducerOptionsFn { o.batchSize = size return nil } - return kinetic.ErrInvalidBatchSize + return ErrInvalidBatchSize } } @@ -68,7 +67,7 @@ func ProducerQueueDepth(queueDepth int) ProducerOptionsFn { o.queueDepth = queueDepth return nil } - return kinetic.ErrInvalidQueueDepth + return ErrInvalidQueueDepth } } @@ -78,7 +77,7 @@ func ProducerMaxRetryAttempts(attemtps int) ProducerOptionsFn { o.maxRetryAttempts = attemtps return nil } - return kinetic.ErrInvalidMaxRetryAttempts + return ErrInvalidMaxRetryAttempts } } @@ -88,7 +87,7 @@ func ProducerConcurrency(count int) ProducerOptionsFn { o.concurrency = count return nil } - return kinetic.ErrInvalidConcurrency + return ErrInvalidConcurrency } } @@ -113,7 +112,7 @@ func ProducerLogLevel(ll aws.LogLevelType) ProducerOptionsFn { } } -func ProducerStatsCollector(sc StatsCollector) ProducerOptionsFn { +func ProducerStats(sc ProducerStatsCollector) ProducerOptionsFn { return func(o *producerOptions) error { o.Stats = sc return nil @@ -123,12 +122,12 @@ func ProducerStatsCollector(sc StatsCollector) ProducerOptionsFn { // Producer sends records to AWS Kinesis or Firehose. type Producer struct { *producerOptions // contains all of the configuration settings for the Producer - *kinetic.LogHelper // object for help with logging + *LogHelper // object for help with logging writer StreamWriter // interface for abstracting the PutRecords call msgCountLimiter *rate.Limiter // rate limiter to limit the number of messages dispatched per second msgSizeLimiter *rate.Limiter // rate limiter to limit the total size (in bytes) of messages dispatched per second workerCount int // number of concurrent workers sending batch messages for the producer - messages chan *kinetic.Message // channel for enqueuing messages to be put on the stream + messages chan *Message // channel for enqueuing messages to be put on the stream status chan *statusReport // channel for workers to communicate their current status dismiss chan empty // channel for handling the decommissioning of a surplus of workers stop chan empty // channel for handling shutdown @@ -147,7 +146,7 @@ func NewProducer(c *aws.Config, w StreamWriter, optionFns ...ProducerOptionsFn) } return &Producer{ producerOptions: producerOptions, - LogHelper: &kinetic.LogHelper{ + LogHelper: &LogHelper{ LogLevel: producerOptions.logLevel, Logger: c.Logger, }, @@ -166,7 +165,7 @@ func (p *Producer) produce() { p.msgSizeLimiter = rate.NewLimiter(rate.Limit(float64(p.writer.getMsgSizeRateLimit())), p.writer.getMsgSizeRateLimit()) // Create communication channels - p.messages = make(chan *kinetic.Message, p.queueDepth) + p.messages = make(chan *Message, p.queueDepth) p.status = make(chan *statusReport) p.dismiss = make(chan empty) p.stop = make(chan empty) @@ -217,7 +216,7 @@ func (p *Producer) produce() { case <-p.stop: return case status := <-p.status: - var batch []*kinetic.Message + var batch []*Message timeout := time.After(p.batchTimeout) fillBatch: @@ -372,10 +371,10 @@ func (p *Producer) resizeWorkerPool(desiredWorkerCount int) { // listen on the dismiss channel which upon receiving a signal will continue sending previously failed messages only // until all failed messages have been sent successfully or aged out. func (p *Producer) doWork() { - batches := make(chan []*kinetic.Message) + batches := make(chan []*Message) defer close(batches) - var retries []*kinetic.Message + var retries []*Message var dismissed bool for ok := true; ok; ok = !dismissed || len(retries) != 0 { // Check to see if there were any signals to dismiss workers (if eligible) @@ -422,9 +421,9 @@ func (p *Producer) doWork() { // sendBatch is the function that is called by each worker to put records on the stream. sendBatch accepts a slice of // messages to send and returns a slice of messages that failed to send -func (p *Producer) sendBatch(batch []*kinetic.Message) []*kinetic.Message { - var failed []*kinetic.Message - err := p.writer.PutRecords(context.TODO(), batch, func(msg *kinetic.Message) error { +func (p *Producer) sendBatch(batch []*Message) []*Message { + var failed []*Message + err := p.writer.PutRecords(context.TODO(), batch, func(msg *Message) error { if msg.FailCount <= p.maxRetryAttempts { failed = append(failed, msg) p.Stats.AddSentRetried(1) @@ -451,7 +450,7 @@ func (p *Producer) sendBatch(batch []*kinetic.Message) []*kinetic.Message { p.LogError("Received AWS error:", err.Error()) case error: switch err { - case kinetic.ErrRetryRecords: + case ErrRetryRecords: break default: p.LogError("Received error:", err.Error()) @@ -464,7 +463,7 @@ func (p *Producer) sendBatch(batch []*kinetic.Message) []*kinetic.Message { } // sendToDataSpill is called when the producer is unable to write the message to the stream -func (p *Producer) sendToDataSpill(msg *kinetic.Message) { +func (p *Producer) sendToDataSpill(msg *Message) { p.Stats.AddDroppedTotal(1) p.Stats.AddDroppedCapacity(1) if err := p.dataSpillFn(msg); err != nil { @@ -486,7 +485,7 @@ func (p *Producer) Close() { } // SendWithContext sends a message to the stream. Cancellation supported through contexts. -func (p *Producer) SendWithContext(ctx context.Context, msg *kinetic.Message) error { +func (p *Producer) SendWithContext(ctx context.Context, msg *Message) error { p.produce() select { case p.messages <- msg: @@ -497,19 +496,19 @@ func (p *Producer) SendWithContext(ctx context.Context, msg *kinetic.Message) er } // Send a message to the stream, waiting on the message to be put into the channel. -func (p *Producer) Send(msg *kinetic.Message) error { +func (p *Producer) Send(msg *Message) error { return p.SendWithContext(context.TODO(), msg) } // TryToSend will attempt to send a message to the stream if the channel has capacity for a message, or will immediately // return with an error if the channel is full. -func (p *Producer) TryToSend(msg *kinetic.Message) error { +func (p *Producer) TryToSend(msg *Message) error { p.produce() select { case p.messages <- msg: return nil default: p.sendToDataSpill(msg) - return kinetic.ErrDroppedMessage + return ErrDroppedMessage } } diff --git a/producer/producer_stats.go b/producer_stats.go similarity index 76% rename from producer/producer_stats.go rename to producer_stats.go index 8def1e2..84034c8 100644 --- a/producer/producer_stats.go +++ b/producer_stats.go @@ -1,4 +1,4 @@ -package producer +package kinetic import ( "log" @@ -10,7 +10,7 @@ import ( // StatsCollector allows for a collector to collect various metrics produced by // the Kinetic producer library. This was really built with rcrowley/go-metrics // in mind. -type StatsCollector interface { +type ProducerStatsCollector interface { AddSentTotal(int) AddSentSuccess(int) AddSentFailed(int) @@ -20,7 +20,6 @@ type StatsCollector interface { AddDroppedRetries(int) AddPutRecordsCalled(int) AddPutRecordsTimeout(int) - AddProvisionedThroughputExceeded(int) AddPutRecordsProvisionedThroughputExceeded(int) UpdatePutRecordsDuration(time.Duration) UpdatePutRecordsBuildDuration(time.Duration) @@ -29,59 +28,55 @@ type StatsCollector interface { } // NilStatsCollector is a stats listener that ignores all metrics. -type NilStatsCollector struct{} +type NilProducerStatsCollector struct{} // AddSentTotal records a count of the total number of messages attempted by PutRecords in the producer. -func (nsc *NilStatsCollector) AddSentTotal(int) {} +func (nsc *NilProducerStatsCollector) AddSentTotal(int) {} // AddSentSuccess records a count of the number of messages sent successfully to AWS Kinesis by the producer. -func (nsc *NilStatsCollector) AddSentSuccess(int) {} +func (nsc *NilProducerStatsCollector) AddSentSuccess(int) {} // AddSentFailed records a count of the number of messages that failed to be sent to AWS Kinesis by the producer. -func (nsc *NilStatsCollector) AddSentFailed(int) {} +func (nsc *NilProducerStatsCollector) AddSentFailed(int) {} // AddSentRetried records a count of the number of messages that were retried after some error occurred when sending // to AWS Kinesis by the producer. -func (nsc *NilStatsCollector) AddSentRetried(int) {} +func (nsc *NilProducerStatsCollector) AddSentRetried(int) {} // AddDroppedTotal records a count of the total number of messages dropped by the application after multiple failures. -func (nsc *NilStatsCollector) AddDroppedTotal(int) {} +func (nsc *NilProducerStatsCollector) AddDroppedTotal(int) {} // AddDroppedCapacity records a count of the number of messages that were dropped by the application due to the stream // writer being at capacity. -func (nsc *NilStatsCollector) AddDroppedCapacity(int) {} +func (nsc *NilProducerStatsCollector) AddDroppedCapacity(int) {} // AddDroppedRetries records a count of the number of retry messages dropped by the application after the max number of // retries was exceeded. -func (nsc *NilStatsCollector) AddDroppedRetries(int) {} +func (nsc *NilProducerStatsCollector) AddDroppedRetries(int) {} // AddPutRecordsCalled records the number of times the PutRecords API was called by the producer. -func (nsc *NilStatsCollector) AddPutRecordsCalled(int) {} +func (nsc *NilProducerStatsCollector) AddPutRecordsCalled(int) {} // AddPutRecordsTimeout records the number of times the PutRecords API timed out on the HTTP level. This is influenced // by the WithHTTPClientTimeout configuration. -func (nsc *NilStatsCollector) AddPutRecordsTimeout(int) {} - -// AddProvisionedThroughputExceeded records the number of times the PutRecords API response contained a record which -// contained an ErrCodeProvisionedThroughputExceededException error. -func (nsc *NilStatsCollector) AddProvisionedThroughputExceeded(int) {} +func (nsc *NilProducerStatsCollector) AddPutRecordsTimeout(int) {} // AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords API returned a // ErrCodeProvisionedThroughputExceededException by the producer. -func (nsc *NilStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} +func (nsc *NilProducerStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} // UpdatePutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls // are measured. -func (nsc *NilStatsCollector) UpdatePutRecordsDuration(time.Duration) {} +func (nsc *NilProducerStatsCollector) UpdatePutRecordsDuration(time.Duration) {} // UpdatePutRecordsBuildDuration records the duration that it took to build the PutRecords API request payload. -func (nsc *NilStatsCollector) UpdatePutRecordsBuildDuration(time.Duration) {} +func (nsc *NilProducerStatsCollector) UpdatePutRecordsBuildDuration(time.Duration) {} // UpdatePutRecordsSendDuration records the duration that it took to send the PutRecords API request payload. -func (nsc *NilStatsCollector) UpdatePutRecordsSendDuration(time.Duration) {} +func (nsc *NilProducerStatsCollector) UpdatePutRecordsSendDuration(time.Duration) {} // UpdateProducerConcurrency records the number of concurrent workers that the producer has. -func (nsc *NilStatsCollector) UpdateProducerConcurrency(int) {} +func (nsc *NilProducerStatsCollector) UpdateProducerConcurrency(int) {} // Metric names to be exported const ( @@ -94,7 +89,6 @@ const ( MetricsDroppedRetries = "kinetic.producer.dropped.retries" MetricsPutRecordsCalled = "kinetic.producer.putrecords.called" MetricsPutRecordsTimeout = "kinetic.producer.putrecords.timeout" - MetricsProvisionedThroughputExceeded = "kinetic.producer.provisionedthroughputexceeded" MetricsPutRecordsProvisionedThroughputExceeded = "kinetic.producer.putrecords.provisionedthroughputexceeded" MetricsPutRecordsDuration = "kinetic.producer.putrecords.duration" MetricsPutRecordsBuildDuration = "kinetic.producer.putrecords.build.duration" @@ -104,7 +98,7 @@ const ( // DefaultStatsCollector is a type that implements the producers's StatsCollector interface using the // rcrowley/go-metrics library -type DefaultStatsCollector struct { +type DefaultProducerStatsCollector struct { SentTotal metrics.Counter SentSuccess metrics.Counter SentFailed metrics.Counter @@ -114,7 +108,6 @@ type DefaultStatsCollector struct { DroppedRetries metrics.Counter PutRecordsCalled metrics.Counter PutRecordsTimeout metrics.Counter - ProvisionedThroughputExceeded metrics.Counter PutRecordsProvisionedThroughputExceeded metrics.Counter PutRecordsDuration metrics.Gauge PutRecordsBuildDuration metrics.Gauge @@ -123,8 +116,8 @@ type DefaultStatsCollector struct { } // NewDefaultStatsCollector instantiates a new DefaultStatsCollector object -func NewDefaultStatsCollector(r metrics.Registry) *DefaultStatsCollector { - return &DefaultStatsCollector{ +func NewDefaultProducerStatsCollector(r metrics.Registry) *DefaultProducerStatsCollector { + return &DefaultProducerStatsCollector{ SentTotal: metrics.GetOrRegisterCounter(MetricsSentTotal, r), SentSuccess: metrics.GetOrRegisterCounter(MetricsSentSuccess, r), SentFailed: metrics.GetOrRegisterCounter(MetricsSentFailed, r), @@ -134,7 +127,6 @@ func NewDefaultStatsCollector(r metrics.Registry) *DefaultStatsCollector { DroppedRetries: metrics.GetOrRegisterCounter(MetricsDroppedRetries, r), PutRecordsCalled: metrics.GetOrRegisterCounter(MetricsPutRecordsCalled, r), PutRecordsTimeout: metrics.GetOrRegisterCounter(MetricsPutRecordsTimeout, r), - ProvisionedThroughputExceeded: metrics.GetOrRegisterCounter(MetricsProvisionedThroughputExceeded, r), PutRecordsProvisionedThroughputExceeded: metrics.GetOrRegisterCounter(MetricsPutRecordsProvisionedThroughputExceeded, r), PutRecordsDuration: metrics.GetOrRegisterGauge(MetricsPutRecordsDuration, r), PutRecordsBuildDuration: metrics.GetOrRegisterGauge(MetricsPutRecordsBuildDuration, r), @@ -144,89 +136,83 @@ func NewDefaultStatsCollector(r metrics.Registry) *DefaultStatsCollector { } // AddSentTotal records a count of the total number of messages attempted by PutRecords in the producer. -func (dsc *DefaultStatsCollector) AddSentTotal(count int) { +func (dsc *DefaultProducerStatsCollector) AddSentTotal(count int) { dsc.SentTotal.Inc(int64(count)) } // AddSentSuccess records a count of the number of messages sent successfully to AWS Kinesis by the producer. -func (dsc *DefaultStatsCollector) AddSentSuccess(count int) { +func (dsc *DefaultProducerStatsCollector) AddSentSuccess(count int) { dsc.SentSuccess.Inc(int64(count)) } // AddSentFailed records a count of the number of messages that failed to be sent to AWS Kinesis by the producer. -func (dsc *DefaultStatsCollector) AddSentFailed(count int) { +func (dsc *DefaultProducerStatsCollector) AddSentFailed(count int) { dsc.SentFailed.Inc(int64(count)) } // AddSentRetried records a count of the number of messages that were retried after some error occurred when sending // to AWS Kinesis by the producer. -func (dsc *DefaultStatsCollector) AddSentRetried(count int) { +func (dsc *DefaultProducerStatsCollector) AddSentRetried(count int) { dsc.SentRetried.Inc(int64(count)) } // AddDroppedTotal records a count of the total number of messages dropped by the application after multiple failures. -func (dsc *DefaultStatsCollector) AddDroppedTotal(count int) { +func (dsc *DefaultProducerStatsCollector) AddDroppedTotal(count int) { dsc.DroppedTotal.Inc(int64(count)) } // AddDroppedCapacity records a count of the number of messages that were dropped by the application due to the stream // writer being at capacity. -func (dsc *DefaultStatsCollector) AddDroppedCapacity(count int) { +func (dsc *DefaultProducerStatsCollector) AddDroppedCapacity(count int) { dsc.DroppedCapacity.Inc(int64(count)) } // AddDroppedRetries records a count of the number of retry messages dropped by the application after the max number of // retries was exceeded. -func (dsc *DefaultStatsCollector) AddDroppedRetries(count int) { +func (dsc *DefaultProducerStatsCollector) AddDroppedRetries(count int) { dsc.DroppedRetries.Inc(int64(count)) } // AddPutRecordsCalled records the number of times the PutRecords API was called by the producer. -func (dsc *DefaultStatsCollector) AddPutRecordsCalled(count int) { +func (dsc *DefaultProducerStatsCollector) AddPutRecordsCalled(count int) { dsc.PutRecordsCalled.Inc(int64(count)) } // AddPutRecordsTimeout records the number of times the PutRecords API timed out on the HTTP level. This is influenced // by the WithHTTPClientTimeout configuration. -func (dsc *DefaultStatsCollector) AddPutRecordsTimeout(count int) { +func (dsc *DefaultProducerStatsCollector) AddPutRecordsTimeout(count int) { dsc.PutRecordsTimeout.Inc(int64(count)) } -// AddProvisionedThroughputExceeded records the number of times the PutRecords API response contained a record which -// contained an ErrCodeProvisionedThroughputExceededException error. -func (dsc *DefaultStatsCollector) AddProvisionedThroughputExceeded(count int) { - dsc.ProvisionedThroughputExceeded.Inc(int64(count)) -} - // AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords API returned a // ErrCodeProvisionedThroughputExceededException by the producer. -func (dsc *DefaultStatsCollector) AddPutRecordsProvisionedThroughputExceeded(count int) { +func (dsc *DefaultProducerStatsCollector) AddPutRecordsProvisionedThroughputExceeded(count int) { dsc.PutRecordsProvisionedThroughputExceeded.Inc(int64(count)) } // UpdatePutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls // are measured. -func (dsc *DefaultStatsCollector) UpdatePutRecordsDuration(duration time.Duration) { +func (dsc *DefaultProducerStatsCollector) UpdatePutRecordsDuration(duration time.Duration) { dsc.PutRecordsDuration.Update(duration.Nanoseconds()) } // UpdatePutRecordsBuildDuration records the duration that it took to build the PutRecords API request payload. -func (dsc *DefaultStatsCollector) UpdatePutRecordsBuildDuration(duration time.Duration) { +func (dsc *DefaultProducerStatsCollector) UpdatePutRecordsBuildDuration(duration time.Duration) { dsc.PutRecordsBuildDuration.Update(duration.Nanoseconds()) } // UpdatePutRecordsSendDuration records the duration that it took to send the PutRecords API request payload. -func (dsc *DefaultStatsCollector) UpdatePutRecordsSendDuration(duration time.Duration) { +func (dsc *DefaultProducerStatsCollector) UpdatePutRecordsSendDuration(duration time.Duration) { dsc.PutRecordsSendDuration.Update(duration.Nanoseconds()) } // UpdateProducerConcurrency records the number of concurrent workers that the producer has. -func (dsc *DefaultStatsCollector) UpdateProducerConcurrency(count int) { +func (dsc *DefaultProducerStatsCollector) UpdateProducerConcurrency(count int) { dsc.ProducerConcurrency.Update(int64(count)) } // PrintStats logs the stats -func (dsc *DefaultStatsCollector) PrintStats() { +func (dsc *DefaultProducerStatsCollector) PrintStats() { log.Printf("Producer Stats: Sent Total: [%d]\n", dsc.SentTotal.Count()) log.Printf("Producer Stats: Sent Success: [%d]\n", dsc.SentSuccess.Count()) log.Printf("Producer Stats: Sent Failed: [%d]\n", dsc.SentFailed.Count()) @@ -236,7 +222,6 @@ func (dsc *DefaultStatsCollector) PrintStats() { log.Printf("Producer Stats: Dropped Capacity: [%d]\n", dsc.DroppedCapacity.Count()) log.Printf("Producer Stats: PutRecords Called: [%d]\n", dsc.PutRecordsCalled.Count()) log.Printf("Producer Stats: PutRecords Timeout: [%d]\n", dsc.PutRecordsTimeout.Count()) - log.Printf("Producer Stats: Provisioned Throughput Exceeded: [%d]\n", dsc.ProvisionedThroughputExceeded.Count()) log.Printf("Producer Stats: PutRecords Provisioned Throughput Exceeded: [%d]\n", dsc.PutRecordsProvisionedThroughputExceeded.Count()) log.Printf("Producer Stats: PutRecords Duration (ns): [%d]\n", dsc.PutRecordsDuration.Value()) log.Printf("Producer Stats: PutRecords Build Duration (ns): [%d]\n", dsc.PutRecordsBuildDuration.Value()) diff --git a/producer/producer_stats_test.go b/producer_stats_test.go similarity index 67% rename from producer/producer_stats_test.go rename to producer_stats_test.go index 5d159cd..fd2e592 100644 --- a/producer/producer_stats_test.go +++ b/producer_stats_test.go @@ -1,4 +1,4 @@ -package producer +package kinetic import ( "math/rand" @@ -13,9 +13,9 @@ func init() { rand.Seed(time.Now().UnixNano()) } -func TestStatsCollector(t *testing.T) { +func TestProducerStatsCollector(t *testing.T) { Convey("given a NilStatsCollector", t, func() { - var sc StatsCollector = &NilStatsCollector{} + var sc ProducerStatsCollector = &NilProducerStatsCollector{} So(sc, ShouldNotBeNil) Convey("check that AddSentTotal does not error", func() { @@ -55,7 +55,7 @@ func TestStatsCollector(t *testing.T) { }) Convey("check that AddProvisionedThroughputExceeded does not erro", func() { - sc.AddProvisionedThroughputExceeded(1) + sc.AddPutRecordsProvisionedThroughputExceeded(1) }) Convey("check that AddPutRecordsTimeout does not error", func() { @@ -81,93 +81,93 @@ func TestStatsCollector(t *testing.T) { Convey("given a DefaulStatsCollector", t, func() { r := metrics.NewRegistry() - var sc StatsCollector = NewDefaultStatsCollector(r) + var sc ProducerStatsCollector = NewDefaultProducerStatsCollector(r) So(sc, ShouldNotBeNil) Convey("check that AddSentTotal does not error", func() { count := rand.Int() sc.AddSentTotal(count) - So(sc.(*DefaultStatsCollector).SentTotal.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultProducerStatsCollector).SentTotal.Count(), ShouldEqual, int64(count)) }) Convey("check that AddSentSuccess does not error", func() { count := rand.Int() sc.AddSentSuccess(count) - So(sc.(*DefaultStatsCollector).SentSuccess.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultProducerStatsCollector).SentSuccess.Count(), ShouldEqual, int64(count)) }) Convey("check that AddSentFailed does not error", func() { count := rand.Int() sc.AddSentFailed(count) - So(sc.(*DefaultStatsCollector).SentFailed.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultProducerStatsCollector).SentFailed.Count(), ShouldEqual, int64(count)) }) Convey("check that AddSentRetried does not error", func() { count := rand.Int() sc.AddSentRetried(count) - So(sc.(*DefaultStatsCollector).SentRetried.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultProducerStatsCollector).SentRetried.Count(), ShouldEqual, int64(count)) }) Convey("check that AddDroppedTotal does not error", func() { count := rand.Int() sc.AddDroppedTotal(count) - So(sc.(*DefaultStatsCollector).DroppedTotal.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultProducerStatsCollector).DroppedTotal.Count(), ShouldEqual, int64(count)) }) Convey("check that AddDroppedCapacity does not error", func() { count := rand.Int() sc.AddDroppedCapacity(count) - So(sc.(*DefaultStatsCollector).DroppedCapacity.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultProducerStatsCollector).DroppedCapacity.Count(), ShouldEqual, int64(count)) }) Convey("check that AddDroppedRetries does not error", func() { count := rand.Int() sc.AddDroppedRetries(count) - So(sc.(*DefaultStatsCollector).DroppedRetries.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultProducerStatsCollector).DroppedRetries.Count(), ShouldEqual, int64(count)) }) Convey("check that AddPutRecordsProvisionedThroughputExceeded does not error", func() { count := rand.Int() sc.AddPutRecordsProvisionedThroughputExceeded(count) - So(sc.(*DefaultStatsCollector).PutRecordsProvisionedThroughputExceeded.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultProducerStatsCollector).PutRecordsProvisionedThroughputExceeded.Count(), ShouldEqual, int64(count)) }) Convey("check that AddPutRecordsCalled does not eroror", func() { count := rand.Int() sc.AddPutRecordsCalled(count) - So(sc.(*DefaultStatsCollector).PutRecordsCalled.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultProducerStatsCollector).PutRecordsCalled.Count(), ShouldEqual, int64(count)) }) - Convey("check that AddProvisionedThroughputExceeded does not erro", func() { + Convey("check that AddPutRecordsProvisionedThroughputExceeded does not erro", func() { count := rand.Int() - sc.AddProvisionedThroughputExceeded(count) - So(sc.(*DefaultStatsCollector).ProvisionedThroughputExceeded.Count(), ShouldEqual, int64(count)) + sc.AddPutRecordsProvisionedThroughputExceeded(count) + So(sc.(*DefaultProducerStatsCollector).PutRecordsProvisionedThroughputExceeded.Count(), ShouldEqual, int64(count)) }) Convey("check that AddPutRecordsTimeout does not error", func() { count := rand.Int() sc.AddPutRecordsTimeout(count) - So(sc.(*DefaultStatsCollector).PutRecordsTimeout.Count(), ShouldEqual, int64(count)) + So(sc.(*DefaultProducerStatsCollector).PutRecordsTimeout.Count(), ShouldEqual, int64(count)) }) Convey("check that UpdatePutRecordsDuration does not error", func() { sc.UpdatePutRecordsDuration(time.Second) - So(sc.(*DefaultStatsCollector).PutRecordsDuration.Value(), ShouldEqual, 1000000000) + So(sc.(*DefaultProducerStatsCollector).PutRecordsDuration.Value(), ShouldEqual, 1000000000) }) Convey("check that UpdatePutRecordsBuildDuration does not error", func() { sc.UpdatePutRecordsBuildDuration(time.Second) - So(sc.(*DefaultStatsCollector).PutRecordsBuildDuration.Value(), ShouldEqual, 1000000000) + So(sc.(*DefaultProducerStatsCollector).PutRecordsBuildDuration.Value(), ShouldEqual, 1000000000) }) Convey("check that UpdatePutRecordsSendDuration does not error", func() { sc.UpdatePutRecordsSendDuration(time.Second) - So(sc.(*DefaultStatsCollector).PutRecordsSendDuration.Value(), ShouldEqual, 1000000000) + So(sc.(*DefaultProducerStatsCollector).PutRecordsSendDuration.Value(), ShouldEqual, 1000000000) }) Convey("check that UpdateProducerConcurrency does not error", func() { sc.UpdateProducerConcurrency(5) - So(sc.(*DefaultStatsCollector).ProducerConcurrency.Value(), ShouldEqual, 5) + So(sc.(*DefaultProducerStatsCollector).ProducerConcurrency.Value(), ShouldEqual, 5) }) }) } diff --git a/producer/producer_test.go b/producer_test.go similarity index 72% rename from producer/producer_test.go rename to producer_test.go index c1b9ae9..dabbbf6 100644 --- a/producer/producer_test.go +++ b/producer_test.go @@ -1,4 +1,4 @@ -package producer +package kinetic import ( . "github.com/smartystreets/goconvey/convey" @@ -12,17 +12,15 @@ import ( "github.com/aws/aws-sdk-go/aws/request" "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/consumer" - "github.com/rewardStyle/kinetic/producer" ) func TestProducer(t *testing.T) { Convey("given a producer", t, func() { - k, err := kinetic.New(func(c *kinetic.Config) { - c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") - c.SetRegion("some-region") - c.SetEndpoint("http://127.0.0.1:4567") - }) + k, err := NewKinetic( + KineticAwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + KineticAwsConfigRegion("some-region"), + KineticAwsConfigEndpoint("http://127.0.0.1:4567"), + ) stream := "some-producer-stream" @@ -44,39 +42,39 @@ func TestProducer(t *testing.T) { So(err, ShouldBeNil) p, err := NewProducer(k.Session.Config, w, - producer.ProducerBatchSize(5), - producer.ProducerBatchTimeout(time.Second), - producer.ProducerMaxRetryAttempts(3), - producer.ProducerQueueDepth(10), - producer.ProducerConcurrency(2), - producer.ProducerShardCheckFrequency(time.Minute), - producer.ProducerDataSpillFn(func(msg *kinetic.Message) error { + kinetic.ProducerBatchSize(5), + kinetic.ProducerBatchTimeout(time.Second), + kinetic.ProducerMaxRetryAttempts(3), + kinetic.ProducerQueueDepth(10), + kinetic.ProducerConcurrency(2), + kinetic.ProducerShardCheckFrequency(time.Minute), + kinetic.ProducerDataSpillFn(func(msg *kinetic.Message) error { //log.Printf("Message was dropped: [%s]\n", string(msg.Data)) return nil }), - producer.ProducerLogLevel(aws.LogOff), - //producer.ProducerStatsCollector(), + kinetic.ProducerLogLevel(aws.LogOff), + //kinetic.ProducerStats(), ) So(p, ShouldNotBeNil) So(err, ShouldBeNil) So(k.Session, ShouldNotBeNil) So(k.Session.Config, ShouldNotBeNil) - r, err := consumer.NewKinesisReader(k.Session.Config, stream, shards[0], - //consumer.KinesisReaderBatchSize(), - //consumer.KinesisReaderShardIterator(), - consumer.KinesisReaderResponseReadTimeout(time.Second), - //consumer.KinesisReaderLogLevel(), - //consumer.KinesisReaderStatsCollector(), + r, err := NewKinesisReader(k.Session.Config, stream, shards[0], + //KinesisReaderBatchSize(), + //KinesisReaderShardIterator(), + KinesisReaderResponseReadTimeout(time.Second), + //KinesisReaderLogLevel(), + //KinesisReaderStats(), ) So(r, ShouldNotBeNil) So(err, ShouldBeNil) - l, err := consumer.NewConsumer(k.Session.Config, r, - consumer.ConsumerQueueDepth(10), - consumer.ConsumerConcurrency(10), - consumer.ConsumerLogLevel(aws.LogOff), - //consumer.ConsumerStatsCollector(lsc), + l, err := NewConsumer(k.Session.Config, r, + ConsumerQueueDepth(10), + ConsumerConcurrency(10), + ConsumerLogLevel(aws.LogOff), + ConsumerStats(&NilConsumerStatsCollector{}), ) So(l, ShouldNotBeNil) So(err, ShouldBeNil) @@ -96,7 +94,7 @@ func TestProducer(t *testing.T) { Convey("check that we can send and receive a single message", func() { start := time.Now() data := "hello" - p.Send(&kinetic.Message{ + p.Send(&Message{ PartitionKey: aws.String("key"), Data: []byte(data), }) @@ -110,7 +108,7 @@ func TestProducer(t *testing.T) { Convey("check that we can send a single message after batch timeout elapses", func() { start := time.Now() data := "hello" - p.Send(&kinetic.Message{ + p.Send(&Message{ PartitionKey: aws.String("key"), Data: []byte(data), }) @@ -143,7 +141,7 @@ func TestProducer(t *testing.T) { go func() { defer wg.Done() for _, datum := range data { - p.Send(&kinetic.Message{ + p.Send(&Message{ PartitionKey: aws.String("key"), Data: []byte(datum), }) diff --git a/consumer/read_closer.go b/read_closer.go similarity index 96% rename from consumer/read_closer.go rename to read_closer.go index 3d4ffb2..77c58fa 100644 --- a/consumer/read_closer.go +++ b/read_closer.go @@ -1,4 +1,4 @@ -package consumer +package kinetic import ( "io" diff --git a/consumer/shard_iterator.go b/shard_iterator.go similarity index 99% rename from consumer/shard_iterator.go rename to shard_iterator.go index 202c96f..10a4089 100644 --- a/consumer/shard_iterator.go +++ b/shard_iterator.go @@ -1,4 +1,4 @@ -package consumer +package kinetic import ( "time" diff --git a/consumer/shard_iterator_test.go b/shard_iterator_test.go similarity index 99% rename from consumer/shard_iterator_test.go rename to shard_iterator_test.go index c46345b..df1f2ef 100644 --- a/consumer/shard_iterator_test.go +++ b/shard_iterator_test.go @@ -1,4 +1,4 @@ -package consumer +package kinetic import ( "github.com/aws/aws-sdk-go/aws" diff --git a/producer/shared_types.go b/shared_types.go similarity index 60% rename from producer/shared_types.go rename to shared_types.go index 0d9a2e1..bdda51b 100644 --- a/producer/shared_types.go +++ b/shared_types.go @@ -1,35 +1,36 @@ -package producer +package kinetic import ( "context" "sync" - - "github.com/rewardStyle/kinetic" ) // StreamWriter is an interface that abstracts the differences in API between Kinesis and Firehose. type StreamWriter interface { - PutRecords(context.Context, []*kinetic.Message, MessageHandlerAsync) error + PutRecords(context.Context, []*Message, MessageHandlerAsync) error getMsgCountRateLimit() int getMsgSizeRateLimit() int getConcurrencyMultiplier() (int, error) } +// StreamReader is an interface that abstracts out a stream reader. +type StreamReader interface { + GetRecord(context.Context, MessageHandler) (int, int, error) + GetRecords(context.Context, MessageHandler) (int, int, error) +} + +// MessageProcessor defines the signature of a message handler used by Listen, RetrieveFn and their associated +// *WithContext functions. MessageHandler accepts a WaitGroup so the function can be run as a blocking operation as +// opposed to MessageHandlerAsync. +type MessageProcessor func(*Message, *sync.WaitGroup) error + // MessageHandler defines the signature of a message handler used by PutRecords(). MessageHandler accepts a WaitGroup // so the function can be run as a blocking operation as opposed to MessageHandlerAsync. -type MessageHandler func(*kinetic.Message, *sync.WaitGroup) error +type MessageHandler func(*Message, *sync.WaitGroup) error // MessageHandlerAsync defines the signature of a message handler used by PutRecords(). MessageHandlerAsync is meant to // be run asynchronously. -type MessageHandlerAsync func(*kinetic.Message) error - -// statusReport is used to communicate a worker's capacity for new messages and to which channel they should be sent. -type statusReport struct { - capacity int // maximum message capacity the worker can handle - failedCount int // number of previous messages that failed to send - failedSize int // size in bytes of the previous messages that failed to send - channel chan []*kinetic.Message // channel of the worker to which the batch messages should be sent -} +type MessageHandlerAsync func(*Message) error // empty is used a as a dummy type for counting semaphore channels. type empty struct{} @@ -39,3 +40,11 @@ type noCopy struct{} // Lock is a no-op used by -copylocks checker from `go vet`. func (*noCopy) Lock() {} + +// statusReport is used to communicate a worker's capacity for new messages and to which channel they should be sent. +type statusReport struct { + capacity int // maximum message capacity the worker can handle + failedCount int // number of previous messages that failed to send + failedSize int // size in bytes of the previous messages that failed to send + channel chan []*Message // channel of the worker to which the batch messages should be sent +} diff --git a/testexec/main.go b/testexec/main.go index 9f6cad7..25e2878 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -19,8 +19,6 @@ import ( "github.com/aws/aws-sdk-go/aws/session" metrics "github.com/jasonyurs/go-metrics" "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/consumer" - "github.com/rewardStyle/kinetic/producer" "net/http" _ "net/http/pprof" @@ -141,11 +139,11 @@ func main() { } func newDefaultKinetic() *kinetic.Kinetic { - k, err := kinetic.New(func(c *kinetic.Config) { - c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") - c.SetRegion("some-region") - c.SetEndpoint("http://127.0.0.1:4567") - }) + k, err := kinetic.NewKinetic( + kinetic.KineticAwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + kinetic.KineticAwsConfigRegion("some-region"), + kinetic.KineticAwsConfigEndpoint("http://127.0.0.1:4567"), + ) if err != nil { log.Fatalf("Unable to create new default kinetic object due to: %v\n", err) } @@ -181,10 +179,10 @@ func newAwsKinetic() *kinetic.Kinetic { creds, _ := sess.Config.Credentials.Get() // Instantiate a new kinetic object configured with appropriate configs - k, err := kinetic.New(func(c *kinetic.Config) { - c.SetCredentials(creds.AccessKeyID, creds.SecretAccessKey, creds.SessionToken) - c.SetRegion(*sess.Config.Region) - }) + k, err := kinetic.NewKinetic( + kinetic.KineticAwsConfigCredentials(creds.AccessKeyID, creds.SecretAccessKey, creds.SessionToken), + kinetic.KineticAwsConfigRegion(*sess.Config.Region), + ) if err != nil { log.Fatalf("Unable to create new aws kinetic object due to: %v\n", err) } @@ -231,36 +229,36 @@ func prepareKinesisStream(k *kinetic.Kinetic, streamName string) { } } -func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Producer { +func newKineticProducer(k *kinetic.Kinetic, streamName string) *kinetic.Producer { if *cfg.Verbose { log.Println("Creating a kinetic producer ...") } - psc := producer.NewDefaultStatsCollector(registry) - w, err := producer.NewKinesisWriter(k.Session.Config, streamName, - producer.KinesisWriterResponseReadTimeout(time.Second), - producer.KinesisWriterMsgCountRateLimit(1000), - producer.KinesisWriterMsgSizeRateLimit(1000000), - producer.KinesisWriterLogLevel(kinetic.LogDebug), - producer.KinesisWriterStatsCollector(psc), + psc := kinetic.NewDefaultProducerStatsCollector(registry) + w, err := kinetic.NewKinesisWriter(k.Session.Config, streamName, + kinetic.KinesisWriterResponseReadTimeout(time.Second), + kinetic.KinesisWriterMsgCountRateLimit(1000), + kinetic.KinesisWriterMsgSizeRateLimit(1000000), + kinetic.KinesisWriterLogLevel(kinetic.LogDebug), + kinetic.KinesisWriterStats(psc), ) if err != nil { log.Fatalf("Unable to create a new kinesis stream writer due to: %v\n", err) } - p, err := producer.NewProducer(k.Session.Config, w, - producer.ProducerBatchSize(500), - producer.ProducerBatchTimeout(time.Second), - producer.ProducerMaxRetryAttempts(3), - producer.ProducerQueueDepth(10000), - producer.ProducerConcurrency(3), - producer.ProducerShardCheckFrequency(time.Minute), - producer.ProducerDataSpillFn(func(msg *kinetic.Message) error { + p, err := kinetic.NewProducer(k.Session.Config, w, + kinetic.ProducerBatchSize(500), + kinetic.ProducerBatchTimeout(time.Second), + kinetic.ProducerMaxRetryAttempts(3), + kinetic.ProducerQueueDepth(10000), + kinetic.ProducerConcurrency(3), + kinetic.ProducerShardCheckFrequency(time.Minute), + kinetic.ProducerDataSpillFn(func(msg *kinetic.Message) error { //log.Printf("Message was dropped: [%s]\n", string(msg.Data)) return nil }), - producer.ProducerLogLevel(aws.LogOff), - producer.ProducerStatsCollector(psc), + kinetic.ProducerLogLevel(aws.LogOff), + kinetic.ProducerStats(psc), ) if err != nil { log.Fatalf("Unable to create a new producer due to: %v\n", err) @@ -269,7 +267,7 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *producer.Produce return p } -func newKineticListener(k *kinetic.Kinetic, streamName string) *consumer.Consumer { +func newKineticListener(k *kinetic.Kinetic, streamName string) *kinetic.Consumer { if *cfg.Verbose { log.Println("Creating a kinetic listener ...") } @@ -280,23 +278,23 @@ func newKineticListener(k *kinetic.Kinetic, streamName string) *consumer.Consume log.Fatalf("Unable to get shards for stream %s due to: %v\n", streamName, err) } - lsc := consumer.NewDefaultStatsCollector(registry) - r, err := consumer.NewKinesisReader(k.Session.Config, streamName, shards[0], - consumer.KinesisReaderBatchSize(10000), - //consumer.KinesisReaderShardIterator(), - consumer.KinesisReaderResponseReadTimeout(time.Second), - consumer.KinesisReaderLogLevel(aws.LogOff), - consumer.KinesisReaderStatsCollector(lsc), + lsc := kinetic.NewDefaultConsumerStatsCollector(registry) + r, err := kinetic.NewKinesisReader(k.Session.Config, streamName, shards[0], + kinetic.KinesisReaderBatchSize(10000), + //kinetic.KinesisReaderShardIterator(), + kinetic.KinesisReaderResponseReadTimeout(time.Second), + kinetic.KinesisReaderLogLevel(aws.LogOff), + kinetic.KinesisReaderStats(lsc), ) if err != nil { log.Fatalf("Unable to create a new kinesis reader due to: %v\n", err) } - l, err := consumer.NewConsumer(k.Session.Config, r, - consumer.ConsumerQueueDepth(500), - consumer.ConsumerConcurrency(10), - consumer.ConsumerLogLevel(aws.LogOff), - consumer.ConsumerStatsCollector(lsc), + l, err := kinetic.NewConsumer(k.Session.Config, r, + kinetic.ConsumerQueueDepth(500), + kinetic.ConsumerConcurrency(10), + kinetic.ConsumerLogLevel(aws.LogOff), + kinetic.ConsumerStats(lsc), ) if err != nil { log.Fatalf("Unable to create a new listener due to: %v\n", err) @@ -318,7 +316,7 @@ func handlePoD() { } } -func display(sd *StreamData, p *producer.Producer, c *consumer.Consumer, wg *sync.WaitGroup) { +func display(sd *StreamData, p *kinetic.Producer, c *kinetic.Consumer, wg *sync.WaitGroup) { defer wg.Done() for { @@ -337,10 +335,10 @@ func display(sd *StreamData, p *producer.Producer, c *consumer.Consumer, wg *syn log.Println() log.Println("***** Stream Data Summary *****") if *cfg.Mode != ModeRead { - p.Stats.(*producer.DefaultStatsCollector).PrintStats() + p.Stats.(*kinetic.DefaultProducerStatsCollector).PrintStats() } if *cfg.Mode != ModeWrite { - c.Stats.(*consumer.DefaultStatsCollector).PrintStats() + c.Stats.(*kinetic.DefaultConsumerStatsCollector).PrintStats() sd.printSummary() } return @@ -351,17 +349,17 @@ func display(sd *StreamData, p *producer.Producer, c *consumer.Consumer, wg *syn log.Println() log.Println("***** Stream Data Stats *****") if *cfg.Mode != ModeRead { - p.Stats.(*producer.DefaultStatsCollector).PrintStats() + p.Stats.(*kinetic.DefaultProducerStatsCollector).PrintStats() } if *cfg.Mode != ModeWrite { - c.Stats.(*consumer.DefaultStatsCollector).PrintStats() + c.Stats.(*kinetic.DefaultConsumerStatsCollector).PrintStats() sd.printStats() } } } } -func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { +func produce(sd *StreamData, p *kinetic.Producer, wg *sync.WaitGroup) { defer wg.Done() // Return early if we don't need to produce @@ -465,7 +463,7 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { } return case <-time.After(time.Second): - newSent := p.Stats.(*producer.DefaultStatsCollector).SentSuccess.Count() + newSent := p.Stats.(*kinetic.DefaultProducerStatsCollector).SentSuccess.Count() if sent != uint64(newSent) { staleTime.Reset(staleTimeout) sent = uint64(newSent) @@ -476,7 +474,7 @@ func produce(sd *StreamData, p *producer.Producer, wg *sync.WaitGroup) { produceWg.Wait() } -func listen(sd *StreamData, c *consumer.Consumer, wg *sync.WaitGroup) { +func listen(sd *StreamData, c *kinetic.Consumer, wg *sync.WaitGroup) { defer wg.Done() // Return early if we don't need to produce @@ -548,7 +546,7 @@ func listen(sd *StreamData, c *consumer.Consumer, wg *sync.WaitGroup) { } return case <-time.After(time.Second): - newConsumed := c.Stats.(*consumer.DefaultStatsCollector).Consumed.Count() + newConsumed := c.Stats.(*kinetic.DefaultConsumerStatsCollector).Consumed.Count() if consumed != uint64(newConsumed) { staleTime.Reset(staleTimeout) consumed = uint64(newConsumed) From a61c7322bd5059f5d4b56a2f12c0d676c90c9d75 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 28 Jul 2017 10:30:23 -0500 Subject: [PATCH 071/109] Documented new data types and functions to fix golint errors. Fixed broken unit tests. --- README.md | 8 ++++---- consumer.go | 24 ++++++++++++++++------- consumer_stats.go | 38 ++++++++++++++++++------------------- consumer_stats_test.go | 4 ++-- consumer_test.go | 12 ++++++------ errors.go | 16 ++++++++-------- firehose_writer.go | 14 +++++++++++++- kcl_reader.go | 25 +++++++++++++++++------- kinesis_reader.go | 13 ++++++++++++- kinesis_reader_test.go | 8 ++++---- kinesis_writer.go | 12 ++++++++++++ kinetic.go | 31 ++++++++++++++++++++---------- kinetic_integration_test.go | 12 ++++++------ kinetic_test.go | 6 +++--- producer.go | 13 ++++++++++++- producer_stats.go | 8 ++++---- producer_stats_test.go | 4 ++-- producer_test.go | 26 ++++++++++++------------- testexec/README.md | 2 +- testexec/main.go | 30 ++++++++++++++--------------- 20 files changed, 191 insertions(+), 115 deletions(-) diff --git a/README.md b/README.md index 99af38c..20121c3 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@ Kinetic is an MIT-licensed high-performance AWS Kinesis Client for Go Kinetic wraps [sendgridlabs go-kinesis library](https://github.com/sendgridlabs/go-kinesis) to provide maximum throughput for AWS Kinesis producers and consumers. -An instance of a Kinetic listener/producer is meant to be used for each shard, so please use it accordingly. If you use more than one instance per-shard then you will +An instance of a Kinetic consumer/producer is meant to be used for each shard, so please use it accordingly. If you use more than one instance per-shard then you will hit the AWS Kinesis throughput [limits](http://docs.aws.amazon.com/kinesis/latest/dev/service-sizes-and-limits.html). ### Getting Started @@ -33,7 +33,7 @@ Kinetic can be used to interface with kinesis like so: ```go import ( "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/listener" + "github.com/rewardStyle/kinetic/consumer" "github.com/rewardStyle/kinetic/message" "github.com/rewardStyle/kinetic/producer" "sync" @@ -52,8 +52,8 @@ p, _ := producer.NewProducer(func(c *producer.Config) { c.SetKinesisStream("stream-name") }) -// Create a kinetic listener -l, _ := consumer.NewListener(func(c *consumer.Config) { +// Create a kinetic consumer +l, _ := consumer.NewConsumer(func(c *consumer.Config) { c.SetAwsConfig(k.Session.Config) c.SetReader(consumer.NewKinesisReader("stream-name", "shard-name")) }) diff --git a/consumer.go b/consumer.go index 3b16fa8..12e0502 100644 --- a/consumer.go +++ b/consumer.go @@ -10,14 +10,15 @@ import ( "golang.org/x/time/rate" ) -// consumerOptions is used to hold all of the configurable settings of a Listener object. +// consumerOptions is used to hold all of the configurable settings of a Consumer. type consumerOptions struct { - queueDepth int - concurrency int + queueDepth int // size of the consumer's message channel + concurrency int // number of concurrent routines processing messages off of the message channel logLevel aws.LogLevelType // log level for configuring the LogHelper's log level Stats ConsumerStatsCollector // stats collection mechanism } +// defaultConsumerOptions instantiates a consumerOptions with default values. func defaultConsumerOptions() *consumerOptions { return &consumerOptions{ queueDepth: 10000, @@ -26,8 +27,10 @@ func defaultConsumerOptions() *consumerOptions { } } +// ConsumerOptionsFn is a method signature for defining functional option methods for configuring the Consumer. type ConsumerOptionsFn func(*consumerOptions) error +// ConsumerQueueDepth is a functional option method for configuring the consumer's queueDepth. func ConsumerQueueDepth(depth int) ConsumerOptionsFn { return func(o *consumerOptions) error { if depth > 0 { @@ -38,6 +41,7 @@ func ConsumerQueueDepth(depth int) ConsumerOptionsFn { } } +// ConsumerConcurrency is a functional option method for configuring the consumer's concurrency. func ConsumerConcurrency(count int) ConsumerOptionsFn { return func(o *consumerOptions) error { if count > 0 { @@ -48,6 +52,7 @@ func ConsumerConcurrency(count int) ConsumerOptionsFn { } } +// ConsumerLogLevel is a functional option method for configuring the consumer's log level. func ConsumerLogLevel(ll aws.LogLevelType) ConsumerOptionsFn { return func(o *consumerOptions) error { o.logLevel = ll & 0xffff0000 @@ -55,6 +60,7 @@ func ConsumerLogLevel(ll aws.LogLevelType) ConsumerOptionsFn { } } +// ConsumerStats is a functional option method for configuring the consumer's stats collector. func ConsumerStats(sc ConsumerStatsCollector) ConsumerOptionsFn { return func(o *consumerOptions) error { o.Stats = sc @@ -62,7 +68,7 @@ func ConsumerStats(sc ConsumerStatsCollector) ConsumerOptionsFn { } } -// Listener polls the StreamReader for messages. +// Consumer polls the StreamReader for messages. type Consumer struct { *consumerOptions *LogHelper @@ -76,7 +82,7 @@ type Consumer struct { consumingMu sync.Mutex } -// NewListener creates a new Listener object for retrieving and listening to message(s) on a StreamReader. +// NewConsumer creates a new Consumer object for retrieving and listening to message(s) on a StreamReader. func NewConsumer(c *aws.Config, r StreamReader, optionFns ...ConsumerOptionsFn) (*Consumer, error) { consumerOptions := defaultConsumerOptions() for _, optionFn := range optionFns { @@ -133,6 +139,7 @@ func (l *Consumer) stopConsuming() { l.consuming = false } +// enqueueSingle calls the readers's GetRecord method and enqueus a single message on the message channel. func (l *Consumer) enqueueSingle(ctx context.Context) (int, int, error) { n, m, err := l.reader.GetRecord(ctx, func(msg *Message, wg *sync.WaitGroup) error { defer wg.Done() @@ -147,6 +154,7 @@ func (l *Consumer) enqueueSingle(ctx context.Context) (int, int, error) { return n, m, nil } +// enqueueBatch calls the reader's GetRecords method and enqueues a batch of messages on the message chanel. func (l *Consumer) enqueueBatch(ctx context.Context) (int, int, error) { n, m, err := l.reader.GetRecords(ctx, func(msg *Message, wg *sync.WaitGroup) error { @@ -162,6 +170,8 @@ func (l *Consumer) enqueueBatch(ctx context.Context) (int, int, error) { return n, m, nil } +// handleErrorLogging is a helper method for handling and logging errors from calling the reader's +// GetRecord and GetRecords method. func (l *Consumer) handleErrorLogging(err error) { switch err := err.(type) { case net.Error: @@ -210,7 +220,7 @@ func (l *Consumer) RetrieveWithContext(ctx context.Context) (*Message, error) { } } -// Retrieve waits for a message from the stream and returns the value +// Retrieve waits for a message from the stream and returns the value. func (l *Consumer) Retrieve() (*Message, error) { return l.RetrieveWithContext(context.TODO()) } @@ -242,7 +252,7 @@ func (l *Consumer) RetrieveFn(fn MessageProcessor) error { return l.RetrieveFnWithContext(context.TODO(), fn) } -// consume calls getRecords with configured batch size in a loop until the listener is stopped. +// consume calls getRecords with configured batch size in a loop until the consumer is stopped. func (l *Consumer) consume(ctx context.Context) { // We need to run startConsuming to make sure that we are okay and ready to start consuming. This is mainly to // avoid a race condition where Listen() will attempt to read the messages channel prior to consume() diff --git a/consumer_stats.go b/consumer_stats.go index 6e35040..0a30941 100644 --- a/consumer_stats.go +++ b/consumer_stats.go @@ -7,8 +7,8 @@ import ( metrics "github.com/jasonyurs/go-metrics" ) -// StatsCollector allows for a collector to collect various metrics produced by -// the Kinetic Listener library. This was really built with rcrowley/go-metrics +// ConsumerStatsCollector allows for a collector to collect various metrics produced by +// the Kinetic Consumer library. This was really built with rcrowley/go-metrics // in mind. type ConsumerStatsCollector interface { AddConsumed(int) @@ -25,7 +25,7 @@ type ConsumerStatsCollector interface { AddGetRecordsUnmarshalDuration(time.Duration) } -// NilStatsCollector is a stats listener that ignores all metrics. +// NilConsumerStatsCollector is a stats consumer that ignores all metrics. type NilConsumerStatsCollector struct{} // AddConsumed records a count of the number of messages received from AWS @@ -50,7 +50,7 @@ func (nsc *NilConsumerStatsCollector) AddBatchSize(int) {} // by the consumer. func (nsc *NilConsumerStatsCollector) AddGetRecordsCalled(int) {} -// AddProvisionedThroughputExceeded records the number of times the GetRecords +// AddGetRecordsProvisionedThroughputExceeded records the number of times the GetRecords // API returned a ErrCodeProvisionedThroughputExceededException by the consumer. func (nsc *NilConsumerStatsCollector) AddGetRecordsProvisionedThroughputExceeded(int) {} @@ -96,7 +96,7 @@ const ( MetricsGetRecordsUnmarshalDuration = "kinetic.consumer.getrecords.unmarshal.duration" ) -// DefaultStatsCollector is a type that implements the listener's StatsCollector interface using the +// DefaultConsumerStatsCollector is a type that implements the consumer's StatsCollector interface using the // rcrowley/go-metrics library type DefaultConsumerStatsCollector struct { Consumed metrics.Counter @@ -113,7 +113,7 @@ type DefaultConsumerStatsCollector struct { GetRecordsUnmarshalDuration metrics.Gauge } -// NewDefaultStatsCollector instantiates a new DefaultStatsCollector object +// NewDefaultConsumerStatsCollector instantiates a new DefaultStatsCollector object func NewDefaultConsumerStatsCollector(r metrics.Registry) *DefaultConsumerStatsCollector { return &DefaultConsumerStatsCollector{ Consumed: metrics.GetOrRegisterCounter(MetricsConsumed, r), @@ -163,7 +163,7 @@ func (dsc *DefaultConsumerStatsCollector) AddGetRecordsCalled(count int) { dsc.GetRecordsCalled.Inc(int64(count)) } -// AddProvisionedThroughputExceeded records the number of times the GetRecords +// AddGetRecordsProvisionedThroughputExceeded records the number of times the GetRecords // API returned a ErrCodeProvisionedThroughputExceededException by the consumer. func (dsc *DefaultConsumerStatsCollector) AddGetRecordsProvisionedThroughputExceeded(count int) { dsc.GetRecordsProvisionedThroughputExceeded.Inc(int64(count)) @@ -209,16 +209,16 @@ func (dsc *DefaultConsumerStatsCollector) AddGetRecordsUnmarshalDuration(duratio // PrintStats logs the stats func (dsc *DefaultConsumerStatsCollector) PrintStats() { - log.Printf("Listener stats: Consumed: [%d]\n", dsc.Consumed.Count()) - log.Printf("Listener stats: Delivered: [%d]\n", dsc.Delivered.Count()) - log.Printf("Listener stats: Processed: [%d]\n", dsc.Processed.Count()) - log.Printf("Listener stats: Batch Size: [%d]\n", dsc.BatchSize.Count()) - log.Printf("Listener stats: GetRecords Called: [%d]\n", dsc.GetRecordsCalled.Count()) - log.Printf("Listener stats: GetRecords Timeout: [%d]\n", dsc.GetRecordsTimeout.Count()) - log.Printf("Listener stats: GetRecords Read Timeout: [%d]\n", dsc.GetRecordsReadTimeout.Count()) - log.Printf("Listener stats: GetRecords Provisioned Throughput Exceeded: [%d]\n", dsc.GetRecordsProvisionedThroughputExceeded.Count()) - log.Printf("Listener stats: Processed Duration (ns): [%d]\n", dsc.ProcessedDuration.Value()) - log.Printf("Listener stats: GetRecords Duration (ns): [%d]\n", dsc.GetRecordsDuration.Value()) - log.Printf("Listener stats: GetRecords Read Response Duration (ns): [%d]\n", dsc.GetRecordsReadResponseDuration.Value()) - log.Printf("Listener stats: GetRecords Unmarshal Duration (ns): [%d]\n", dsc.GetRecordsUnmarshalDuration.Value()) + log.Printf("Consumer stats: Consumed: [%d]\n", dsc.Consumed.Count()) + log.Printf("Consumer stats: Delivered: [%d]\n", dsc.Delivered.Count()) + log.Printf("Consumer stats: Processed: [%d]\n", dsc.Processed.Count()) + log.Printf("Consumer stats: Batch Size: [%d]\n", dsc.BatchSize.Count()) + log.Printf("Consumer stats: GetRecords Called: [%d]\n", dsc.GetRecordsCalled.Count()) + log.Printf("Consumer stats: GetRecords Timeout: [%d]\n", dsc.GetRecordsTimeout.Count()) + log.Printf("Consumer stats: GetRecords Read Timeout: [%d]\n", dsc.GetRecordsReadTimeout.Count()) + log.Printf("Consumer stats: GetRecords Provisioned Throughput Exceeded: [%d]\n", dsc.GetRecordsProvisionedThroughputExceeded.Count()) + log.Printf("Consumer stats: Processed Duration (ns): [%d]\n", dsc.ProcessedDuration.Value()) + log.Printf("Consumer stats: GetRecords Duration (ns): [%d]\n", dsc.GetRecordsDuration.Value()) + log.Printf("Consumer stats: GetRecords Read Response Duration (ns): [%d]\n", dsc.GetRecordsReadResponseDuration.Value()) + log.Printf("Consumer stats: GetRecords Unmarshal Duration (ns): [%d]\n", dsc.GetRecordsUnmarshalDuration.Value()) } diff --git a/consumer_stats_test.go b/consumer_stats_test.go index b1c44fc..cb1446d 100644 --- a/consumer_stats_test.go +++ b/consumer_stats_test.go @@ -14,7 +14,7 @@ func init() { } func TestConsumerStatsCollector(t *testing.T) { - Convey("given a NilStatsCollector", t, func() { + Convey("given a NilConsumerStatsCollector", t, func() { var sc ConsumerStatsCollector = &NilConsumerStatsCollector{} So(sc, ShouldNotBeNil) @@ -67,7 +67,7 @@ func TestConsumerStatsCollector(t *testing.T) { }) }) - Convey("given a DefaultStatsCollector", t, func() { + Convey("given a DefaultConsumerStatsCollector", t, func() { r := metrics.NewRegistry() var sc ConsumerStatsCollector = NewDefaultConsumerStatsCollector(r) So(sc, ShouldNotBeNil) diff --git a/consumer_test.go b/consumer_test.go index 6160318..48264e6 100644 --- a/consumer_test.go +++ b/consumer_test.go @@ -29,15 +29,15 @@ func putRecord(l *Consumer, b []byte) (*string, error) { return resp.SequenceNumber, nil } -func TestListener(t *testing.T) { - Convey("given a listener", t, func() { +func TestConsumer(t *testing.T) { + Convey("given a consumer", t, func() { k, err := NewKinetic( - KineticAwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), - KineticAwsConfigRegion("some-region"), - KineticAwsConfigEndpoint("http://127.0.0.1:4567"), + AwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + AwsConfigRegion("some-region"), + AwsConfigEndpoint("http://127.0.0.1:4567"), ) - stream := "some-listener-stream" + stream := "some-consumer-stream" err = k.CreateStream(stream, 1) So(err, ShouldBeNil) diff --git a/errors.go b/errors.go index 348f512..2792ecf 100644 --- a/errors.go +++ b/errors.go @@ -19,19 +19,19 @@ var ( ) var ( - // ErrNilListener is returned when the Listener is nil when it is not + // ErrNilConsumer is returned when the Consumer is nil when it is not // supposed to be - ErrNilListener = errors.New("StreamReader not associated with a listener") + ErrNilConsumer = errors.New("StreamReader not associated with a consumer") - // ErrListenerAlreadyAssociated is returned by a StreamReader attempting - // to associate it with a Listener when it already has an association - // with a listener - ErrListenerAlreadyAssociated = errors.New("StreamReader already associated with a listener") + // ErrConsumerAlreadyAssociated is returned by a StreamReader attempting + // to associate it with a Consumer when it already has an association + // with a consumer + ErrConsumerAlreadyAssociated = errors.New("StreamReader already associated with a consumer") // ErrAlreadyConsuming is returned when attempting to consume when the - // Listener is already consuming. May be returned by + // Consumer is already consuming. May be returned by // Retrieve/RetrieveFn. - ErrAlreadyConsuming = errors.New("Listener already consuming. Only one Listen, Retrieve, or RetrieveFn may be active at a time") + ErrAlreadyConsuming = errors.New("Consumer already consuming. Only one Listen, Retrieve, or RetrieveFn may be active at a time") // ErrEmptySequenceNumber is returned when attempting to set an empty // sequence number. diff --git a/firehose_writer.go b/firehose_writer.go index d3e686a..681dfce 100644 --- a/firehose_writer.go +++ b/firehose_writer.go @@ -17,6 +17,7 @@ const ( firehoseMsgSizeRateLimit = 5000000 // AWS Firehose limit of 5 MB/sec ) +// firehoseWriterOptions is a struct that holds all of the FirehoseWriter's configurable parameters. type firehoseWriterOptions struct { msgCountRateLimit int // maximum number of records to be sent per second msgSizeRateLimit int // maximum (transmission) size of records to be sent per second @@ -25,6 +26,7 @@ type firehoseWriterOptions struct { Stats ProducerStatsCollector // stats collection mechanism } +// defaultFirehoseWriterOptions instantiates a firehoseWriterOptions with default values. func defaultFirehoseWriterOptions() *firehoseWriterOptions { return &firehoseWriterOptions{ msgCountRateLimit: firehoseMsgCountRateLimit, @@ -35,8 +37,12 @@ func defaultFirehoseWriterOptions() *firehoseWriterOptions { } } +// FireHoseWriterOptionsFn is a method signature for defining functional option methods for configuring +// the FirehoseWriter. type FireHoseWriterOptionsFn func(*firehoseWriterOptions) error +// FirehoseWriterMsgCountRateLimit is a functional option method for configuring the FirehoseWriter's +// message count rate limit. func FirehoseWriterMsgCountRateLimit(limit int) FireHoseWriterOptionsFn { return func(o *firehoseWriterOptions) error { if limit > 0 && limit <= firehoseMsgCountRateLimit { @@ -47,6 +53,8 @@ func FirehoseWriterMsgCountRateLimit(limit int) FireHoseWriterOptionsFn { } } +// FirehoseWriterMsgSizeRateLimit is a functional option method for configuring the FirehoseWriter's +// messsage size rate limit. func FirehoseWriterMsgSizeRateLimit(limit int) FireHoseWriterOptionsFn { return func(o *firehoseWriterOptions) error { if limit > 0 && limit <= firehoseMsgSizeRateLimit { @@ -57,6 +65,8 @@ func FirehoseWriterMsgSizeRateLimit(limit int) FireHoseWriterOptionsFn { } } +// FirehoseWriterThroughputMultiplier is a functional option method for configuring the FirehoseWriter's +// throughput multiplier. func FirehoseWriterThroughputMultiplier(multiplier int) FireHoseWriterOptionsFn { return func(o *firehoseWriterOptions) error { if multiplier > 0 { @@ -67,6 +77,7 @@ func FirehoseWriterThroughputMultiplier(multiplier int) FireHoseWriterOptionsFn } } +// FirehoseWriterLogLevel is a functional option method for configuring the FirehoseWriter's log level. func FirehoseWriterLogLevel(ll aws.LogLevelType) FireHoseWriterOptionsFn { return func(o *firehoseWriterOptions) error { o.logLevel = ll & 0xffff0000 @@ -74,6 +85,7 @@ func FirehoseWriterLogLevel(ll aws.LogLevelType) FireHoseWriterOptionsFn { } } +// FirehoseWriterStats is a functional option method for configuring the FirehoseWriter's stats collector. func FirehoseWriterStats(sc ProducerStatsCollector) FireHoseWriterOptionsFn { return func(o *firehoseWriterOptions) error { o.Stats = sc @@ -90,7 +102,7 @@ type FirehoseWriter struct { } // NewFirehoseWriter creates a new stream writer to write records to a Kinesis. -func NewFirewhoseWriter(c *aws.Config, stream string, optionFns ...FireHoseWriterOptionsFn) (*FirehoseWriter, error) { +func NewFirehoseWriter(c *aws.Config, stream string, optionFns ...FireHoseWriterOptionsFn) (*FirehoseWriter, error) { firehoseWriterOptions := defaultFirehoseWriterOptions() for _, optionFn := range optionFns { optionFn(firehoseWriterOptions) diff --git a/kcl_reader.go b/kcl_reader.go index 4119851..14d1c61 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -13,14 +13,16 @@ import ( "github.com/aws/aws-sdk-go/aws" ) +// kclReaderOptions is a struct that holds all of the KclReader's configurable parameters. type kclReaderOptions struct { - onInitCallbackFn func() error - onCheckpointCallbackFn func() error - onShutdownCallbackFn func() error + onInitCallbackFn func() error // callback function that gets called after initialization + onCheckpointCallbackFn func() error // callback function that gets called after checkpointing + onShutdownCallbackFn func() error // callback function that gets called after shutdown logLevel aws.LogLevelType // log level for configuring the LogHelper's log level Stats ConsumerStatsCollector // stats collection mechanism } +// defaultKlcReaderOptions instantiates a kclReaderOptions with default values. func defaultKlcReaderOptions() *kclReaderOptions { return &kclReaderOptions{ onInitCallbackFn: func() error { return nil }, @@ -31,8 +33,11 @@ func defaultKlcReaderOptions() *kclReaderOptions { } } +// KlcReaderOptionsFn is a method signature for defining functional option methods for configuring the KclReader. type KlcReaderOptionsFn func(*kclReaderOptions) error +// KlcReaderOnInitCallbackFn is a functional option method for configuring the KclReader's +// onInitCallbackFn. func KlcReaderOnInitCallbackFn(fn func() error) KlcReaderOptionsFn { return func(o *kclReaderOptions) error { o.onInitCallbackFn = fn @@ -40,6 +45,8 @@ func KlcReaderOnInitCallbackFn(fn func() error) KlcReaderOptionsFn { } } +// KlcReaderOnCheckpointCallbackFn is a functional option method for configuring the KclReader's +// onCheckpointCallbackFn. func KlcReaderOnCheckpointCallbackFn(fn func() error) KlcReaderOptionsFn { return func(o *kclReaderOptions) error { o.onCheckpointCallbackFn = fn @@ -47,6 +54,8 @@ func KlcReaderOnCheckpointCallbackFn(fn func() error) KlcReaderOptionsFn { } } +// KlcReaderOnShutdownCallbackFn is a functional option method for configuring the KclReader's +// onShutdownCallbackFn. func KlcReaderOnShutdownCallbackFn(fn func() error) KlcReaderOptionsFn { return func(o *kclReaderOptions) error { o.onShutdownCallbackFn = fn @@ -54,6 +63,7 @@ func KlcReaderOnShutdownCallbackFn(fn func() error) KlcReaderOptionsFn { } } +// KlcReaderLogLevel is a functional option method for configuring the KclReader's log level. func KlcReaderLogLevel(ll aws.LogLevelType) KlcReaderOptionsFn { return func(o *kclReaderOptions) error { o.logLevel = ll @@ -61,6 +71,7 @@ func KlcReaderLogLevel(ll aws.LogLevelType) KlcReaderOptionsFn { } } +// KlcReaderStats is a functional option method for configuring the KclReader's stats collector. func KlcReaderStats(sc ConsumerStatsCollector) KlcReaderOptionsFn { return func(o *kclReaderOptions) error { o.Stats = sc @@ -94,8 +105,8 @@ func NewKclReader(c *aws.Config, optionFns ...KlcReaderOptionsFn) (*KclReader, e }, nil } -// processRecords is a helper method which loops through the message buffer and puts messages on the listener's -// message channel. After all the messages on the message buffer have been moved to the listener's message +// processRecords is a helper method which loops through the message buffer and puts messages on the consumer's +// message channel. After all the messages on the message buffer have been moved to the consumer's message // channel, a message is sent (following the Multilang protocol) to acknowledge that the processRecords message // has been received / processed func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, int, error) { @@ -239,13 +250,13 @@ func (r *KclReader) onShutdown() error { return nil } -// GetRecord calls processRecords to attempt to put one message from message buffer to the listener's message +// GetRecord calls processRecords to attempt to put one message from message buffer to the consumer's message // channel func (r *KclReader) GetRecord(ctx context.Context, fn MessageHandler) (int, int, error) { return r.processRecords(fn, 1) } -// GetRecords calls processRecords to attempt to put all messages on the message buffer on the listener's +// GetRecords calls processRecords to attempt to put all messages on the message buffer on the consumer's // message channel func (r *KclReader) GetRecords(ctx context.Context, fn MessageHandler) (int, int, error) { return r.processRecords(fn, -1) diff --git a/kinesis_reader.go b/kinesis_reader.go index 52e298b..a531fbe 100644 --- a/kinesis_reader.go +++ b/kinesis_reader.go @@ -18,7 +18,7 @@ const ( kinesisReaderBatchSize = 10000 ) -// kinesisReaderOptions is used to hold all of the configurable settings of a KinesisReader. +// kinesisReaderOptions a struct that holds all of the KinesisReader's configurable parameters. type kinesisReaderOptions struct { batchSize int shardIterator *ShardIterator @@ -27,6 +27,7 @@ type kinesisReaderOptions struct { Stats ConsumerStatsCollector // stats collection mechanism } +// defaultKinesisReaderOptions instantiates a kinesisReaderOptions with default values. func defaultKinesisReaderOptions() *kinesisReaderOptions { return &kinesisReaderOptions{ batchSize: kinesisReaderBatchSize, @@ -36,8 +37,12 @@ func defaultKinesisReaderOptions() *kinesisReaderOptions { } } +// KinesisReaderOptionsFn is a method signature for defining functional option methods for configuring +// the KinesisReader. type KinesisReaderOptionsFn func(*kinesisReaderOptions) error +// KinesisReaderBatchSize is a functional option method for configuring the KinesisReader's +// batch size. func KinesisReaderBatchSize(size int) KinesisReaderOptionsFn { return func(o *kinesisReaderOptions) error { if size > 0 && size <= kinesisReaderBatchSize { @@ -48,6 +53,8 @@ func KinesisReaderBatchSize(size int) KinesisReaderOptionsFn { } } +// KinesisReaderShardIterator is a functional option method for configuring the KinesisReader's +// shard iterator. func KinesisReaderShardIterator(shardIterator *ShardIterator) KinesisReaderOptionsFn { return func(o *kinesisReaderOptions) error { o.shardIterator = shardIterator @@ -55,6 +62,8 @@ func KinesisReaderShardIterator(shardIterator *ShardIterator) KinesisReaderOptio } } +// KinesisReaderResponseReadTimeout is a functional option method for configuring the KinesisReader's +// response read timeout. func KinesisReaderResponseReadTimeout(timeout time.Duration) KinesisReaderOptionsFn { return func(o *kinesisReaderOptions) error { o.responseReadTimeout = timeout @@ -62,6 +71,7 @@ func KinesisReaderResponseReadTimeout(timeout time.Duration) KinesisReaderOption } } +// KinesisReaderLogLevel is a functional option method for configuring the KinesisReader's log level. func KinesisReaderLogLevel(ll aws.LogLevelType) KinesisReaderOptionsFn { return func(o *kinesisReaderOptions) error { o.logLevel = ll & 0xffff0000 @@ -69,6 +79,7 @@ func KinesisReaderLogLevel(ll aws.LogLevelType) KinesisReaderOptionsFn { } } +// KinesisReaderStats is a functional option method for configuring the KinesisReader's stats collector. func KinesisReaderStats(sc ConsumerStatsCollector) KinesisReaderOptionsFn { return func(o *kinesisReaderOptions) error { o.Stats = sc diff --git a/kinesis_reader_test.go b/kinesis_reader_test.go index a4a1f11..dba0061 100644 --- a/kinesis_reader_test.go +++ b/kinesis_reader_test.go @@ -17,9 +17,9 @@ func TestNewKinesisReader(t *testing.T) { Convey("given a kinetic object and kinesis stream/shard names", t, func() { // Instantiate a new kinentic object k, err := NewKinetic( - KineticAwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), - KineticAwsConfigRegion("some-region"), - KineticAwsConfigEndpoint("http://127.0.0.1:4567"), + AwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + AwsConfigRegion("some-region"), + AwsConfigEndpoint("http://127.0.0.1:4567"), ) So(k, ShouldNotBeNil) So(err, ShouldBeNil) @@ -38,7 +38,7 @@ func TestNewKinesisReader(t *testing.T) { }) Convey("check that we can create a new KinesisReader with configured values", func() { - batchSize := rand.Int() + batchSize := rand.Intn(kinesisReaderBatchSize) respReadTimeout := time.Duration(rand.Int()) * time.Second logLevel := aws.LogDebug | aws.LogDebugWithSigning | LogDebug shardIterator := NewShardIterator() diff --git a/kinesis_writer.go b/kinesis_writer.go index 554b344..205e67c 100644 --- a/kinesis_writer.go +++ b/kinesis_writer.go @@ -17,6 +17,7 @@ const ( kinesisMsgSizeRateLimit = 1000000 // AWS Kinesis limit of 1 MB/sec ) +// kinesisWriterOptions is a struct that holds all of the KinesisWriter's configurable parameters. type kinesisWriterOptions struct { responseReadTimeout time.Duration // maximum time to wait for PutRecords API call before timing out msgCountRateLimit int // maximum number of records to be sent per second @@ -25,6 +26,7 @@ type kinesisWriterOptions struct { Stats ProducerStatsCollector // stats collection mechanism } +// defaultKinesisWriterOptions instantiates a kinesisWriterOptions with default values. func defaultKinesisWriterOptions() *kinesisWriterOptions { return &kinesisWriterOptions{ responseReadTimeout: time.Second, @@ -35,8 +37,12 @@ func defaultKinesisWriterOptions() *kinesisWriterOptions { } } +// KinesisWriterOptionsFn is a method signature for defining functional option methods for configuring +// the KinesisWriter. type KinesisWriterOptionsFn func(*kinesisWriterOptions) error +// KinesisWriterResponseReadTimeout is a functional option method for configuring the KinesisWriter's +// response read timeout func KinesisWriterResponseReadTimeout(timeout time.Duration) KinesisWriterOptionsFn { return func(o *kinesisWriterOptions) error { o.responseReadTimeout = timeout @@ -44,6 +50,8 @@ func KinesisWriterResponseReadTimeout(timeout time.Duration) KinesisWriterOption } } +// KinesisWriterMsgCountRateLimit is a functional option method for configuring the KinesisWriter's +// message count rate limit func KinesisWriterMsgCountRateLimit(limit int) KinesisWriterOptionsFn { return func(o *kinesisWriterOptions) error { if limit > 0 && limit <= kinesisMsgCountRateLimit { @@ -54,6 +62,8 @@ func KinesisWriterMsgCountRateLimit(limit int) KinesisWriterOptionsFn { } } +// KinesisWriterMsgSizeRateLimit is a functional option method for configuring the KinesisWriter's +// message size rate limit func KinesisWriterMsgSizeRateLimit(limit int) KinesisWriterOptionsFn { return func(o *kinesisWriterOptions) error { if limit > 0 && limit <= kinesisMsgSizeRateLimit { @@ -64,6 +74,7 @@ func KinesisWriterMsgSizeRateLimit(limit int) KinesisWriterOptionsFn { } } +// KinesisWriterLogLevel is a functional option method for configuring the KinesisWriter's log level func KinesisWriterLogLevel(ll aws.LogLevelType) KinesisWriterOptionsFn { return func(o *kinesisWriterOptions) error { o.logLevel = ll & 0xffff0000 @@ -71,6 +82,7 @@ func KinesisWriterLogLevel(ll aws.LogLevelType) KinesisWriterOptionsFn { } } +// KinesisWriterStats is a functional option method for configuring the KinesisWriter's stats collector func KinesisWriterStats(sc ProducerStatsCollector) KinesisWriterOptionsFn { return func(o *kinesisWriterOptions) error { o.Stats = sc diff --git a/kinetic.go b/kinetic.go index 97b8aef..c0398e1 100644 --- a/kinetic.go +++ b/kinetic.go @@ -16,11 +16,13 @@ import ( "github.com/aws/aws-sdk-go/aws/credentials" ) +// kineticOptions is a struct that holds all of Kinetic's configurable parameters. type kineticOptions struct { awsConfig *aws.Config // logLevel aws.LogLevelType // log level for configuring the LogHelper's log level } +// defaultKineticOptions instantiates a kineticOptions with default values. func defaultKineticOptions() *kineticOptions { return &kineticOptions{ awsConfig: aws.NewConfig().WithHTTPClient( @@ -31,44 +33,52 @@ func defaultKineticOptions() *kineticOptions { } } -type KineticOptionsFn func(*kineticOptions) error +// OptionsFn is a method signature for defining functional option methods for configuring Kinetic. +type OptionsFn func(*kineticOptions) error -func KineticAwsConfigCredentials(accessKey, secretKey, sessionToken string) KineticOptionsFn { +// AwsConfigCredentials is a functional option method for configuring Kinetic's AwsConfig credentials. +func AwsConfigCredentials(accessKey, secretKey, sessionToken string) OptionsFn { return func(o *kineticOptions) error { o.awsConfig.WithCredentials(credentials.NewStaticCredentials(accessKey, secretKey, sessionToken)) return nil } } -func KineticAwsConfigRegion(region string) KineticOptionsFn { +// AwsConfigRegion is a functional option method for configuring Kinetic's AwsConfig region. +func AwsConfigRegion(region string) OptionsFn { return func(o *kineticOptions) error { o.awsConfig.WithRegion(region) return nil } } -func KineticAwsConfigEndpoint(endpoint string) KineticOptionsFn { +// AwsConfigEndpoint is a functional option method for configuring Kinetic's AwsConfig endpoint. +func AwsConfigEndpoint(endpoint string) OptionsFn { return func(o *kineticOptions) error { o.awsConfig.WithEndpoint(endpoint) return nil } } -func KineticAwsConfigLogger(logger aws.Logger) KineticOptionsFn { +// AwsConfigLogger is a functional option method for configuring Kinetic's AwsConfig logger. +func AwsConfigLogger(logger aws.Logger) OptionsFn { return func(o *kineticOptions) error { o.awsConfig.WithLogger(logger) return nil } } -func KineticAwsConfigLogLevel(logLevel aws.LogLevelType) KineticOptionsFn { +// AwsConfigLogLevel is a functional option method for configuring Kinetic's AwsConfig log level. +func AwsConfigLogLevel(logLevel aws.LogLevelType) OptionsFn { return func(o *kineticOptions) error { o.awsConfig.WithLogLevel(logLevel) return nil } } -func KineticAwsConfigHttpClientTimeout(timeout time.Duration) KineticOptionsFn { +// AwsConfigHTTPClientTimeout is a functional option method for configuring Kinetic's +// AwsConfig HTTP client timeout. +func AwsConfigHTTPClientTimeout(timeout time.Duration) OptionsFn { return func(o *kineticOptions) error { o.awsConfig.WithHTTPClient(&http.Client{ Timeout: timeout, @@ -77,7 +87,8 @@ func KineticAwsConfigHttpClientTimeout(timeout time.Duration) KineticOptionsFn { } } -func KineticLogLevel(logLevel aws.LogLevelType) KineticOptionsFn { +// LogLevel is a functional option method for configuring Kinetic's log level. +func LogLevel(logLevel aws.LogLevelType) OptionsFn { return func(o *kineticOptions) error { o.logLevel = logLevel & 0xffff0000 return nil @@ -95,8 +106,8 @@ type Kinetic struct { Session *session.Session } -// New creates a new instance of Kinetic. -func NewKinetic(optionFns ...KineticOptionsFn) (*Kinetic, error) { +// NewKinetic creates a new instance of Kinetic. +func NewKinetic(optionFns ...OptionsFn) (*Kinetic, error) { kineticOptions := defaultKineticOptions() for _, optionFn := range optionFns { optionFn(kineticOptions) diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index 903c0f7..c2a52f2 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -92,9 +92,9 @@ func TestKineticIntegration(t *testing.T) { // Instantiate a new kinentic object k, err := NewKinetic( - KineticAwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), - KineticAwsConfigRegion("some-region"), - KineticAwsConfigEndpoint("http://127.0.0.1:4567"), + AwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + AwsConfigRegion("some-region"), + AwsConfigEndpoint("http://127.0.0.1:4567"), ) assert.NotNil(t, k) assert.Nil(t, err) @@ -131,7 +131,7 @@ func TestKineticIntegration(t *testing.T) { KinesisWriterResponseReadTimeout(time.Second), KinesisWriterMsgCountRateLimit(1000), KinesisWriterMsgSizeRateLimit(1000000), - KinesisWriterLogLevel(LogDebug), + KinesisWriterLogLevel(aws.LogOff), ) if err != nil { log.Fatalf("Unable to create a new kinesis stream writer due to: %v\n", err) @@ -167,7 +167,7 @@ func TestKineticIntegration(t *testing.T) { assert.NotNil(t, r) assert.NoError(t, err) - // Create a new kinetic listener + // Create a new kinetic consumer l, err := NewConsumer(k.Session.Config, r, ConsumerQueueDepth(20), ConsumerConcurrency(10), @@ -200,7 +200,7 @@ func TestKineticIntegration(t *testing.T) { } }(&numSent) - // Use the listener to read messages from the kinetic stream + // Use the consumer to read messages from the kinetic stream go func() { l.Listen(func(m *Message, fnwg *sync.WaitGroup) error { defer fnwg.Done() diff --git a/kinetic_test.go b/kinetic_test.go index 0ace2e4..03f3270 100644 --- a/kinetic_test.go +++ b/kinetic_test.go @@ -15,9 +15,9 @@ import ( func TestKinetic(t *testing.T) { Convey("given a kinetic object", t, func() { k, err := NewKinetic( - KineticAwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), - KineticAwsConfigRegion("some-region"), - KineticAwsConfigEndpoint("http://127.0.0.1:4567"), + AwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + AwsConfigRegion("some-region"), + AwsConfigEndpoint("http://127.0.0.1:4567"), ) So(k, ShouldNotBeNil) So(err, ShouldBeNil) diff --git a/producer.go b/producer.go index 9f1649d..75becb8 100644 --- a/producer.go +++ b/producer.go @@ -15,7 +15,7 @@ const ( putRecordsMaxBatchSize = 500 ) -// producerOptions holds all of the configurable settings for a Producer +// producerOptions holds all of the configurable settings for a Producer. type producerOptions struct { batchSize int // maximum message capacity per request batchTimeout time.Duration // maximum time duration to wait for incoming messages @@ -28,6 +28,7 @@ type producerOptions struct { Stats ProducerStatsCollector // stats collection mechanism } +// defaultProducerOptions instantiates a producerOptions with default values. func defaultProducerOptions() *producerOptions { return &producerOptions{ batchSize: putRecordsMaxBatchSize, @@ -42,8 +43,10 @@ func defaultProducerOptions() *producerOptions { } } +// ProducerOptionsFn is a method signature for defining functional option methods for configuring the Producer. type ProducerOptionsFn func(*producerOptions) error +// ProducerBatchSize is a functional option method for configuing the producer's batch size. func ProducerBatchSize(size int) ProducerOptionsFn { return func(o *producerOptions) error { if size > 0 && size <= putRecordsMaxBatchSize { @@ -54,6 +57,7 @@ func ProducerBatchSize(size int) ProducerOptionsFn { } } +// ProducerBatchTimeout is a functional option method for configuing the producer's batch timeout. func ProducerBatchTimeout(timeout time.Duration) ProducerOptionsFn { return func(o *producerOptions) error { o.batchTimeout = timeout @@ -61,6 +65,7 @@ func ProducerBatchTimeout(timeout time.Duration) ProducerOptionsFn { } } +// ProducerQueueDepth is a functional option method for configuing the producer's queue depth. func ProducerQueueDepth(queueDepth int) ProducerOptionsFn { return func(o *producerOptions) error { if queueDepth > 0 { @@ -71,6 +76,7 @@ func ProducerQueueDepth(queueDepth int) ProducerOptionsFn { } } +// ProducerMaxRetryAttempts is a functional option method for configuing the producer's max retry attempts. func ProducerMaxRetryAttempts(attemtps int) ProducerOptionsFn { return func(o *producerOptions) error { if attemtps > 0 { @@ -81,6 +87,7 @@ func ProducerMaxRetryAttempts(attemtps int) ProducerOptionsFn { } } +// ProducerConcurrency is a functional option method for configuing the producer's concurrency. func ProducerConcurrency(count int) ProducerOptionsFn { return func(o *producerOptions) error { if count > 0 { @@ -91,6 +98,7 @@ func ProducerConcurrency(count int) ProducerOptionsFn { } } +// ProducerShardCheckFrequency is a functional option method for configuing the producer's shard check frequency. func ProducerShardCheckFrequency(duration time.Duration) ProducerOptionsFn { return func(o *producerOptions) error { o.shardCheckFreq = duration @@ -98,6 +106,7 @@ func ProducerShardCheckFrequency(duration time.Duration) ProducerOptionsFn { } } +// ProducerDataSpillFn is a functional option method for configuing the producer's data spill callback function. func ProducerDataSpillFn(fn MessageHandlerAsync) ProducerOptionsFn { return func(o *producerOptions) error { o.dataSpillFn = fn @@ -105,6 +114,7 @@ func ProducerDataSpillFn(fn MessageHandlerAsync) ProducerOptionsFn { } } +// ProducerLogLevel is a functional option method for configuing the producer's log level. func ProducerLogLevel(ll aws.LogLevelType) ProducerOptionsFn { return func(o *producerOptions) error { o.logLevel = ll & 0xffff0000 @@ -112,6 +122,7 @@ func ProducerLogLevel(ll aws.LogLevelType) ProducerOptionsFn { } } +// ProducerStats is a functional option method for configuing the producer's stats collector. func ProducerStats(sc ProducerStatsCollector) ProducerOptionsFn { return func(o *producerOptions) error { o.Stats = sc diff --git a/producer_stats.go b/producer_stats.go index 84034c8..331230f 100644 --- a/producer_stats.go +++ b/producer_stats.go @@ -7,7 +7,7 @@ import ( metrics "github.com/jasonyurs/go-metrics" ) -// StatsCollector allows for a collector to collect various metrics produced by +// ProducerStatsCollector allows for a collector to collect various metrics produced by // the Kinetic producer library. This was really built with rcrowley/go-metrics // in mind. type ProducerStatsCollector interface { @@ -27,7 +27,7 @@ type ProducerStatsCollector interface { UpdateProducerConcurrency(int) } -// NilStatsCollector is a stats listener that ignores all metrics. +// NilProducerStatsCollector is a stats consumer that ignores all metrics. type NilProducerStatsCollector struct{} // AddSentTotal records a count of the total number of messages attempted by PutRecords in the producer. @@ -96,7 +96,7 @@ const ( MetricsProducerConcurrency = "kinetic.producer.concurrency" ) -// DefaultStatsCollector is a type that implements the producers's StatsCollector interface using the +// DefaultProducerStatsCollector is a type that implements the producers's StatsCollector interface using the // rcrowley/go-metrics library type DefaultProducerStatsCollector struct { SentTotal metrics.Counter @@ -115,7 +115,7 @@ type DefaultProducerStatsCollector struct { ProducerConcurrency metrics.Gauge } -// NewDefaultStatsCollector instantiates a new DefaultStatsCollector object +// NewDefaultProducerStatsCollector instantiates a new DefaultStatsCollector object func NewDefaultProducerStatsCollector(r metrics.Registry) *DefaultProducerStatsCollector { return &DefaultProducerStatsCollector{ SentTotal: metrics.GetOrRegisterCounter(MetricsSentTotal, r), diff --git a/producer_stats_test.go b/producer_stats_test.go index fd2e592..20dd8ef 100644 --- a/producer_stats_test.go +++ b/producer_stats_test.go @@ -14,7 +14,7 @@ func init() { } func TestProducerStatsCollector(t *testing.T) { - Convey("given a NilStatsCollector", t, func() { + Convey("given a NilProducerStatsCollector", t, func() { var sc ProducerStatsCollector = &NilProducerStatsCollector{} So(sc, ShouldNotBeNil) @@ -79,7 +79,7 @@ func TestProducerStatsCollector(t *testing.T) { }) }) - Convey("given a DefaulStatsCollector", t, func() { + Convey("given a DefaultProdcuerStatsCollector", t, func() { r := metrics.NewRegistry() var sc ProducerStatsCollector = NewDefaultProducerStatsCollector(r) So(sc, ShouldNotBeNil) diff --git a/producer_test.go b/producer_test.go index dabbbf6..b3f25dd 100644 --- a/producer_test.go +++ b/producer_test.go @@ -10,16 +10,14 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" - - "github.com/rewardStyle/kinetic" ) func TestProducer(t *testing.T) { Convey("given a producer", t, func() { k, err := NewKinetic( - KineticAwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), - KineticAwsConfigRegion("some-region"), - KineticAwsConfigEndpoint("http://127.0.0.1:4567"), + AwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + AwsConfigRegion("some-region"), + AwsConfigEndpoint("http://127.0.0.1:4567"), ) stream := "some-producer-stream" @@ -42,18 +40,18 @@ func TestProducer(t *testing.T) { So(err, ShouldBeNil) p, err := NewProducer(k.Session.Config, w, - kinetic.ProducerBatchSize(5), - kinetic.ProducerBatchTimeout(time.Second), - kinetic.ProducerMaxRetryAttempts(3), - kinetic.ProducerQueueDepth(10), - kinetic.ProducerConcurrency(2), - kinetic.ProducerShardCheckFrequency(time.Minute), - kinetic.ProducerDataSpillFn(func(msg *kinetic.Message) error { + ProducerBatchSize(5), + ProducerBatchTimeout(time.Second), + ProducerMaxRetryAttempts(3), + ProducerQueueDepth(10), + ProducerConcurrency(2), + ProducerShardCheckFrequency(time.Minute), + ProducerDataSpillFn(func(msg *Message) error { //log.Printf("Message was dropped: [%s]\n", string(msg.Data)) return nil }), - kinetic.ProducerLogLevel(aws.LogOff), - //kinetic.ProducerStats(), + ProducerLogLevel(aws.LogOff), + //ProducerStats(), ) So(p, ShouldNotBeNil) So(err, ShouldBeNil) diff --git a/testexec/README.md b/testexec/README.md index a3306b9..00c45b9 100644 --- a/testexec/README.md +++ b/testexec/README.md @@ -7,7 +7,7 @@ A new stream will be created (with a random stream name) unless a stream name is There are three modes of operations permitted: 'read', 'write' and 'readwrite'. Write mode is produce only where as read mode is consume only. Readwrite mode is both produce and consume concurrently. -So depending on the mode of operation, a kinetic producer will stream dummy data to the kinesis stream and a kinetic listener will stream from the kinesis stream. +So depending on the mode of operation, a kinetic producer will stream dummy data to the kinesis stream and a kinetic consumer will stream from the kinesis stream. The program logs stream data stats to the console log periodically. diff --git a/testexec/main.go b/testexec/main.go index 25e2878..899624d 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -122,8 +122,8 @@ func main() { // Create a new kinetic producer p := newKineticProducer(k, streamName) - // Create a new kinetic listener - l := newKineticListener(k, streamName) + // Create a new kinetic consumer + l := newKineticConsumer(k, streamName) // Instantiate StreamData Object to keep stats streamData := NewStreamData() @@ -140,9 +140,9 @@ func main() { func newDefaultKinetic() *kinetic.Kinetic { k, err := kinetic.NewKinetic( - kinetic.KineticAwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), - kinetic.KineticAwsConfigRegion("some-region"), - kinetic.KineticAwsConfigEndpoint("http://127.0.0.1:4567"), + kinetic.AwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + kinetic.AwsConfigRegion("some-region"), + kinetic.AwsConfigEndpoint("http://127.0.0.1:4567"), ) if err != nil { log.Fatalf("Unable to create new default kinetic object due to: %v\n", err) @@ -180,8 +180,8 @@ func newAwsKinetic() *kinetic.Kinetic { // Instantiate a new kinetic object configured with appropriate configs k, err := kinetic.NewKinetic( - kinetic.KineticAwsConfigCredentials(creds.AccessKeyID, creds.SecretAccessKey, creds.SessionToken), - kinetic.KineticAwsConfigRegion(*sess.Config.Region), + kinetic.AwsConfigCredentials(creds.AccessKeyID, creds.SecretAccessKey, creds.SessionToken), + kinetic.AwsConfigRegion(*sess.Config.Region), ) if err != nil { log.Fatalf("Unable to create new aws kinetic object due to: %v\n", err) @@ -267,9 +267,9 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *kinetic.Producer return p } -func newKineticListener(k *kinetic.Kinetic, streamName string) *kinetic.Consumer { +func newKineticConsumer(k *kinetic.Kinetic, streamName string) *kinetic.Consumer { if *cfg.Verbose { - log.Println("Creating a kinetic listener ...") + log.Println("Creating a kinetic consumer ...") } // Determine the shard name @@ -297,7 +297,7 @@ func newKineticListener(k *kinetic.Kinetic, streamName string) *kinetic.Consumer kinetic.ConsumerStats(lsc), ) if err != nil { - log.Fatalf("Unable to create a new listener due to: %v\n", err) + log.Fatalf("Unable to create a new consumer due to: %v\n", err) } return l @@ -485,10 +485,10 @@ func listen(sd *StreamData, c *kinetic.Consumer, wg *sync.WaitGroup) { defer func() { if *cfg.Verbose { log.Println() - log.Println("listener: Exiting listen ...") + log.Println("consumer: Exiting listen ...") } - // In read and readwrite mode the listener controls when to stop displaying + // In read and readwrite mode the consumer controls when to stop displaying stopDisplay <- struct{}{} }() @@ -504,7 +504,7 @@ func listen(sd *StreamData, c *kinetic.Consumer, wg *sync.WaitGroup) { // Only mark "done" if the message isn't a duplicate if sd.exists(msg.ID) { if *cfg.Verbose { - log.Printf("listener: Duplicate message: %v\n", msg) + log.Printf("consumer: Duplicate message: %v\n", msg) } } @@ -536,13 +536,13 @@ func listen(sd *StreamData, c *kinetic.Consumer, wg *sync.WaitGroup) { case <-stopListen: if *cfg.Verbose { log.Println() - log.Println("listener: Received stop listen ...") + log.Println("consumer: Received stop listen ...") } return case <-staleTime.C: if *cfg.Verbose { log.Println() - log.Println("listener: No more incoming messages from listener ...") + log.Println("consumer: No more incoming messages from consumer ...") } return case <-time.After(time.Second): From a6c3ba5ab8fd0166a247a408eab14dd3efe83833 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 28 Jul 2017 10:44:07 -0500 Subject: [PATCH 072/109] Regenerated glide.yaml and glide.lock files. --- glide.lock | 83 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ glide.yaml | 25 ++++++++++++++++ 2 files changed, 108 insertions(+) create mode 100644 glide.lock create mode 100644 glide.yaml diff --git a/glide.lock b/glide.lock new file mode 100644 index 0000000..33e6407 --- /dev/null +++ b/glide.lock @@ -0,0 +1,83 @@ +hash: 67b38976bd434e7c37b4cb0446ba002edb0203a0964c82fb2932ef6597395b3a +updated: 2017-07-28T10:43:09.939381844-05:00 +imports: +- name: github.com/aws/aws-sdk-go + version: 72e42b13da62269f68308fb6068b7ea691a416a4 + subpackages: + - aws + - aws/awserr + - aws/awsutil + - aws/client + - aws/client/metadata + - aws/corehandlers + - aws/credentials + - aws/credentials/ec2rolecreds + - aws/credentials/endpointcreds + - aws/credentials/stscreds + - aws/defaults + - aws/ec2metadata + - aws/endpoints + - aws/request + - aws/session + - aws/signer/v4 + - internal/shareddefaults + - private/protocol + - private/protocol/json/jsonutil + - private/protocol/jsonrpc + - private/protocol/query + - private/protocol/query/queryutil + - private/protocol/rest + - private/protocol/xml/xmlutil + - service/firehose + - service/firehose/firehoseiface + - service/kinesis + - service/kinesis/kinesisiface + - service/sts +- name: github.com/go-ini/ini + version: 1730955e3146956d6a087861380f9b4667ed5071 +- name: github.com/jasonyurs/go-metrics + version: 63f9066b204fd9a315fc7691094d05aff455ecf5 +- name: github.com/jmespath/go-jmespath + version: bd40a432e4c76585ef6b72d3fd96fb9b6dc7b68d +- name: github.com/rcrowley/go-metrics + version: ab2277b1c5d15c3cba104e9cbddbdfc622df5ad8 +- name: github.com/rewardStyle/kinetic + version: 1d4a236f3ee6c01b3a69d77d3fb0e64f35361b3e + subpackages: + - config + - errs + - logging +- name: golang.org/x/net + version: 0e2717dc3cc05907dc23096ef3a9086ea93f567f + subpackages: + - context +- name: golang.org/x/time + version: 8be79e1e0910c292df4e79c241bb7e8f7e725959 + subpackages: + - rate +testImports: +- name: github.com/davecgh/go-spew + version: 6d212800a42e8ab5c146b8ace3490ee17e5225f9 + subpackages: + - spew +- name: github.com/jtolds/gls + version: f1ac7f4f24f50328e6bc838ca4437d1612a0243c +- name: github.com/pmezard/go-difflib + version: d8ed2627bdf02c080bf22230dbb337003b7aba2d + subpackages: + - difflib +- name: github.com/smartystreets/assertions + version: c9ee7d9e9a2aeec0bee7c4a516f3e0ad7cb7e558 + subpackages: + - internal/go-render/render + - internal/oglematchers +- name: github.com/smartystreets/goconvey + version: 3bd662eac601ad6436e64776af2e112069eb2edc + subpackages: + - convey + - convey/gotest + - convey/reporting +- name: github.com/stretchr/testify + version: 69483b4bd14f5845b5a1e55bca19e954e827f1d0 + subpackages: + - assert diff --git a/glide.yaml b/glide.yaml new file mode 100644 index 0000000..7ef334e --- /dev/null +++ b/glide.yaml @@ -0,0 +1,25 @@ +package: . +import: +- package: github.com/aws/aws-sdk-go + subpackages: + - aws + - aws/awserr + - aws/credentials + - aws/request + - aws/session + - service/firehose + - service/firehose/firehoseiface + - service/kinesis + - service/kinesis/kinesisiface +- package: github.com/rcrowley/go-metrics +- package: github.com/rewardStyle/kinetic +- package: golang.org/x/time + subpackages: + - rate +testImport: +- package: github.com/smartystreets/goconvey + subpackages: + - convey +- package: github.com/stretchr/testify + subpackages: + - assert From 83b9b32007145e80d38369092725bb8e6db37731 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 28 Jul 2017 11:05:32 -0500 Subject: [PATCH 073/109] Fix vendoring issues. --- consumer_stats.go | 2 +- consumer_stats_test.go | 2 +- glide.lock | 6 ++---- glide.yaml | 1 - producer_stats.go | 2 +- producer_stats_test.go | 2 +- testexec/main.go | 2 +- 7 files changed, 7 insertions(+), 10 deletions(-) diff --git a/consumer_stats.go b/consumer_stats.go index 0a30941..1849f56 100644 --- a/consumer_stats.go +++ b/consumer_stats.go @@ -4,7 +4,7 @@ import ( "log" "time" - metrics "github.com/jasonyurs/go-metrics" + metrics "github.com/rcrowley/go-metrics" ) // ConsumerStatsCollector allows for a collector to collect various metrics produced by diff --git a/consumer_stats_test.go b/consumer_stats_test.go index cb1446d..a3ab7bb 100644 --- a/consumer_stats_test.go +++ b/consumer_stats_test.go @@ -5,7 +5,7 @@ import ( "testing" "time" - metrics "github.com/jasonyurs/go-metrics" + metrics "github.com/rcrowley/go-metrics" . "github.com/smartystreets/goconvey/convey" ) diff --git a/glide.lock b/glide.lock index 33e6407..2f592a4 100644 --- a/glide.lock +++ b/glide.lock @@ -1,5 +1,5 @@ -hash: 67b38976bd434e7c37b4cb0446ba002edb0203a0964c82fb2932ef6597395b3a -updated: 2017-07-28T10:43:09.939381844-05:00 +hash: 88992f22cbe218586fcad11169719591a709cd6b17de91e7718ee7b6b87c67f5 +updated: 2017-07-28T11:04:02.656139446-05:00 imports: - name: github.com/aws/aws-sdk-go version: 72e42b13da62269f68308fb6068b7ea691a416a4 @@ -35,8 +35,6 @@ imports: - service/sts - name: github.com/go-ini/ini version: 1730955e3146956d6a087861380f9b4667ed5071 -- name: github.com/jasonyurs/go-metrics - version: 63f9066b204fd9a315fc7691094d05aff455ecf5 - name: github.com/jmespath/go-jmespath version: bd40a432e4c76585ef6b72d3fd96fb9b6dc7b68d - name: github.com/rcrowley/go-metrics diff --git a/glide.yaml b/glide.yaml index 7ef334e..3387fbd 100644 --- a/glide.yaml +++ b/glide.yaml @@ -12,7 +12,6 @@ import: - service/kinesis - service/kinesis/kinesisiface - package: github.com/rcrowley/go-metrics -- package: github.com/rewardStyle/kinetic - package: golang.org/x/time subpackages: - rate diff --git a/producer_stats.go b/producer_stats.go index 331230f..11f5532 100644 --- a/producer_stats.go +++ b/producer_stats.go @@ -4,7 +4,7 @@ import ( "log" "time" - metrics "github.com/jasonyurs/go-metrics" + metrics "github.com/rcrowley/go-metrics" ) // ProducerStatsCollector allows for a collector to collect various metrics produced by diff --git a/producer_stats_test.go b/producer_stats_test.go index 20dd8ef..e3dea49 100644 --- a/producer_stats_test.go +++ b/producer_stats_test.go @@ -5,7 +5,7 @@ import ( "testing" "time" - metrics "github.com/jasonyurs/go-metrics" + metrics "github.com/rcrowley/go-metrics" . "github.com/smartystreets/goconvey/convey" ) diff --git a/testexec/main.go b/testexec/main.go index 899624d..c180194 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -17,7 +17,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/request" "github.com/aws/aws-sdk-go/aws/session" - metrics "github.com/jasonyurs/go-metrics" + metrics "github.com/rcrowley/go-metrics" "github.com/rewardStyle/kinetic" "net/http" From adcd33a4972965d871bc291c5286ee1223758568 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 28 Jul 2017 11:24:07 -0500 Subject: [PATCH 074/109] Additional changes to glide.yaml file. --- glide.lock | 12 +++--------- glide.yaml | 5 ++++- 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/glide.lock b/glide.lock index 2f592a4..fb83edb 100644 --- a/glide.lock +++ b/glide.lock @@ -1,5 +1,5 @@ -hash: 88992f22cbe218586fcad11169719591a709cd6b17de91e7718ee7b6b87c67f5 -updated: 2017-07-28T11:04:02.656139446-05:00 +hash: 7203d745bb3c303ab9257a7d20b250ce919248e1251502ebe6d56a5af086126c +updated: 2017-07-28T11:21:23.082612845-05:00 imports: - name: github.com/aws/aws-sdk-go version: 72e42b13da62269f68308fb6068b7ea691a416a4 @@ -39,12 +39,6 @@ imports: version: bd40a432e4c76585ef6b72d3fd96fb9b6dc7b68d - name: github.com/rcrowley/go-metrics version: ab2277b1c5d15c3cba104e9cbddbdfc622df5ad8 -- name: github.com/rewardStyle/kinetic - version: 1d4a236f3ee6c01b3a69d77d3fb0e64f35361b3e - subpackages: - - config - - errs - - logging - name: golang.org/x/net version: 0e2717dc3cc05907dc23096ef3a9086ea93f567f subpackages: @@ -70,7 +64,7 @@ testImports: - internal/go-render/render - internal/oglematchers - name: github.com/smartystreets/goconvey - version: 3bd662eac601ad6436e64776af2e112069eb2edc + version: d4c757aa9afd1e2fc1832aaab209b5794eb336e1 subpackages: - convey - convey/gotest diff --git a/glide.yaml b/glide.yaml index 3387fbd..9ce7ec7 100644 --- a/glide.yaml +++ b/glide.yaml @@ -1,6 +1,7 @@ -package: . +package: github.com/rewardStyle/kinetic import: - package: github.com/aws/aws-sdk-go + version: v1.10.3 subpackages: - aws - aws/awserr @@ -17,8 +18,10 @@ import: - rate testImport: - package: github.com/smartystreets/goconvey + version: v1.6.2 subpackages: - convey - package: github.com/stretchr/testify + version: v1.1.4 subpackages: - assert From a352c0b0cfb1f4758b5b92499c6530171e84803f Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 28 Jul 2017 13:42:36 -0500 Subject: [PATCH 075/109] Refactored the produce/consumer constructor methods to not require a StreamWriter/StreamReader. Converted the writer/reader to configurable options which can be configured using a functional option method. --- consumer.go | 38 ++++++++++++++++++++++----------- consumer_test.go | 12 ++++++++++- kinetic_integration_test.go | 6 ++++-- producer.go | 42 ++++++++++++++++++++++++------------- producer_test.go | 6 ++++-- testexec/main.go | 6 ++++-- 6 files changed, 77 insertions(+), 33 deletions(-) diff --git a/consumer.go b/consumer.go index 12e0502..081d618 100644 --- a/consumer.go +++ b/consumer.go @@ -12,6 +12,7 @@ import ( // consumerOptions is used to hold all of the configurable settings of a Consumer. type consumerOptions struct { + reader StreamReader // interface for abstracting the GetRecord/GetRecords call queueDepth int // size of the consumer's message channel concurrency int // number of concurrent routines processing messages off of the message channel logLevel aws.LogLevelType // log level for configuring the LogHelper's log level @@ -30,6 +31,14 @@ func defaultConsumerOptions() *consumerOptions { // ConsumerOptionsFn is a method signature for defining functional option methods for configuring the Consumer. type ConsumerOptionsFn func(*consumerOptions) error +// ConsumerReader is a functional option method for configuring the consumer's stream reader. +func ConsumerReader(r StreamReader) ConsumerOptionsFn { + return func(o *consumerOptions) error { + o.reader = r + return nil + } +} + // ConsumerQueueDepth is a functional option method for configuring the consumer's queueDepth. func ConsumerQueueDepth(depth int) ConsumerOptionsFn { return func(o *consumerOptions) error { @@ -70,31 +79,36 @@ func ConsumerStats(sc ConsumerStatsCollector) ConsumerOptionsFn { // Consumer polls the StreamReader for messages. type Consumer struct { - *consumerOptions - *LogHelper - reader StreamReader - txnCountRateLimiter *rate.Limiter - txSizeRateLimiter *rate.Limiter - messages chan *Message - concurrencySem chan empty - pipeOfDeath chan empty - consuming bool - consumingMu sync.Mutex + *consumerOptions // contains all of the configuration settings for the Consumer + *LogHelper // object for help with logging + txnCountRateLimiter *rate.Limiter // rate limiter to limit the number of transactions per second + txSizeRateLimiter *rate.Limiter // rate limiter to limit the transmission size per seccond + messages chan *Message // channel for storing messages that have been retrieved from the stream + concurrencySem chan empty // channel for controlling the number of concurrent workers processing messages from the message channel + pipeOfDeath chan empty // channel for handling pipe of death + consuming bool // flag for indicating whether or not the consumer is consuming + consumingMu sync.Mutex // mutex for making the consuming flag thread safe } // NewConsumer creates a new Consumer object for retrieving and listening to message(s) on a StreamReader. -func NewConsumer(c *aws.Config, r StreamReader, optionFns ...ConsumerOptionsFn) (*Consumer, error) { +func NewConsumer(c *aws.Config, stream string, shard string, optionFns ...ConsumerOptionsFn) (*Consumer, error) { consumerOptions := defaultConsumerOptions() for _, optionFn := range optionFns { optionFn(consumerOptions) } + if consumerOptions.reader == nil { + r, err := NewKinesisReader(c, stream, shard) + if err != nil { + return nil, err + } + consumerOptions.reader = r + } return &Consumer{ consumerOptions: consumerOptions, LogHelper: &LogHelper{ LogLevel: consumerOptions.logLevel, Logger: c.Logger, }, - reader: r, }, nil } diff --git a/consumer_test.go b/consumer_test.go index 48264e6..69df0a1 100644 --- a/consumer_test.go +++ b/consumer_test.go @@ -61,7 +61,8 @@ func TestConsumer(t *testing.T) { So(r, ShouldNotBeNil) So(err, ShouldBeNil) - l, err := NewConsumer(k.Session.Config, r, + l, err := NewConsumer(k.Session.Config, stream, shards[0], + ConsumerReader(r), ConsumerQueueDepth(10), ConsumerConcurrency(10), ConsumerLogLevel(aws.LogOff), @@ -71,12 +72,21 @@ func TestConsumer(t *testing.T) { So(err, ShouldBeNil) Convey("given a kinesis reader", func() { + Convey("check that the reader was initialized correctly", func() { + So(l.reader, ShouldNotBeNil) + So(l.reader, ShouldEqual, r) + }) + r := l.reader.(*KinesisReader) Convey("check that the reader was initialized with the correct stream name", func() { So(r.stream, ShouldEqual, stream) }) + Convey("check that the reader was initialized with the correct shard", func() { + So(r.shard, ShouldEqual, shards[0]) + }) + Convey("check that the kinesis client was initialized correctly", func() { So(r.client, ShouldNotBeNil) }) diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index c2a52f2..8b33d98 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -138,7 +138,8 @@ func TestKineticIntegration(t *testing.T) { } // Create a new kinetic producer - p, err := NewProducer(k.Session.Config, w, + p, err := NewProducer(k.Session.Config, stream, + ProducerWriter(w), ProducerBatchSize(5), ProducerBatchTimeout(time.Second), ProducerMaxRetryAttempts(3), @@ -168,7 +169,8 @@ func TestKineticIntegration(t *testing.T) { assert.NoError(t, err) // Create a new kinetic consumer - l, err := NewConsumer(k.Session.Config, r, + l, err := NewConsumer(k.Session.Config, stream, shards[0], + ConsumerReader(r), ConsumerQueueDepth(20), ConsumerConcurrency(10), ConsumerLogLevel(aws.LogOff), diff --git a/producer.go b/producer.go index 75becb8..2820402 100644 --- a/producer.go +++ b/producer.go @@ -17,15 +17,16 @@ const ( // producerOptions holds all of the configurable settings for a Producer. type producerOptions struct { - batchSize int // maximum message capacity per request - batchTimeout time.Duration // maximum time duration to wait for incoming messages - queueDepth int // maximum number of messages to enqueue in the message queue - maxRetryAttempts int // maximum number of retry attempts for failed messages - concurrency int // number of concurrent workers per shard - shardCheckFreq time.Duration // frequency (specified as a duration) with which to check the the shard size - dataSpillFn MessageHandlerAsync // callback function for handling dropped messages that the producer was unable to send to the stream - logLevel aws.LogLevelType // log level for configuring the LogHelper's log level - Stats ProducerStatsCollector // stats collection mechanism + writer StreamWriter // interface for abstracting the PutRecords call + batchSize int // maximum message capacity per request + batchTimeout time.Duration // maximum time duration to wait for incoming messages + queueDepth int // maximum number of messages to enqueue in the message queue + maxRetryAttempts int // maximum number of retry attempts for failed messages + concurrency int // number of concurrent workers per shard + shardCheckFreq time.Duration // frequency (specified as a duration) with which to check the the shard size + dataSpillFn MessageHandlerAsync // callback function for handling dropped messages that the producer was unable to send to the stream + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + Stats ProducerStatsCollector // stats collection mechanism } // defaultProducerOptions instantiates a producerOptions with default values. @@ -46,6 +47,14 @@ func defaultProducerOptions() *producerOptions { // ProducerOptionsFn is a method signature for defining functional option methods for configuring the Producer. type ProducerOptionsFn func(*producerOptions) error +// ProducerWriter is a functional option method for configuing the producer's stream writer. +func ProducerWriter(w StreamWriter) ProducerOptionsFn { + return func(o *producerOptions) error { + o.writer = w + return nil + } +} + // ProducerBatchSize is a functional option method for configuing the producer's batch size. func ProducerBatchSize(size int) ProducerOptionsFn { return func(o *producerOptions) error { @@ -133,12 +142,11 @@ func ProducerStats(sc ProducerStatsCollector) ProducerOptionsFn { // Producer sends records to AWS Kinesis or Firehose. type Producer struct { *producerOptions // contains all of the configuration settings for the Producer - *LogHelper // object for help with logging - writer StreamWriter // interface for abstracting the PutRecords call + *LogHelper // object for help with logging msgCountLimiter *rate.Limiter // rate limiter to limit the number of messages dispatched per second msgSizeLimiter *rate.Limiter // rate limiter to limit the total size (in bytes) of messages dispatched per second workerCount int // number of concurrent workers sending batch messages for the producer - messages chan *Message // channel for enqueuing messages to be put on the stream + messages chan *Message // channel for enqueuing messages to be put on the stream status chan *statusReport // channel for workers to communicate their current status dismiss chan empty // channel for handling the decommissioning of a surplus of workers stop chan empty // channel for handling shutdown @@ -150,18 +158,24 @@ type Producer struct { } // NewProducer creates a new producer for writing records to a Kinesis or Firehose stream. -func NewProducer(c *aws.Config, w StreamWriter, optionFns ...ProducerOptionsFn) (*Producer, error) { +func NewProducer(c *aws.Config, stream string, optionFns ...ProducerOptionsFn) (*Producer, error) { producerOptions := defaultProducerOptions() for _, optionFn := range optionFns { optionFn(producerOptions) } + if producerOptions.writer == nil { + w, err := NewKinesisWriter(c, stream) + if err != nil { + return nil, err + } + producerOptions.writer = w + } return &Producer{ producerOptions: producerOptions, LogHelper: &LogHelper{ LogLevel: producerOptions.logLevel, Logger: c.Logger, }, - writer: w, }, nil } diff --git a/producer_test.go b/producer_test.go index b3f25dd..74a301a 100644 --- a/producer_test.go +++ b/producer_test.go @@ -39,7 +39,8 @@ func TestProducer(t *testing.T) { So(w, ShouldNotBeNil) So(err, ShouldBeNil) - p, err := NewProducer(k.Session.Config, w, + p, err := NewProducer(k.Session.Config, stream, + ProducerWriter(w), ProducerBatchSize(5), ProducerBatchTimeout(time.Second), ProducerMaxRetryAttempts(3), @@ -68,7 +69,8 @@ func TestProducer(t *testing.T) { So(r, ShouldNotBeNil) So(err, ShouldBeNil) - l, err := NewConsumer(k.Session.Config, r, + l, err := NewConsumer(k.Session.Config, stream, shards[0], + ConsumerReader(r), ConsumerQueueDepth(10), ConsumerConcurrency(10), ConsumerLogLevel(aws.LogOff), diff --git a/testexec/main.go b/testexec/main.go index c180194..44b376a 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -246,7 +246,8 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *kinetic.Producer log.Fatalf("Unable to create a new kinesis stream writer due to: %v\n", err) } - p, err := kinetic.NewProducer(k.Session.Config, w, + p, err := kinetic.NewProducer(k.Session.Config, streamName, + kinetic.ProducerWriter(w), kinetic.ProducerBatchSize(500), kinetic.ProducerBatchTimeout(time.Second), kinetic.ProducerMaxRetryAttempts(3), @@ -290,7 +291,8 @@ func newKineticConsumer(k *kinetic.Kinetic, streamName string) *kinetic.Consumer log.Fatalf("Unable to create a new kinesis reader due to: %v\n", err) } - l, err := kinetic.NewConsumer(k.Session.Config, r, + l, err := kinetic.NewConsumer(k.Session.Config, streamName, shards[0], + kinetic.ConsumerReader(r), kinetic.ConsumerQueueDepth(500), kinetic.ConsumerConcurrency(10), kinetic.ConsumerLogLevel(aws.LogOff), From f7a09a8ff18f7f4a75a4b5e1ca7f108a3d133dc1 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 28 Jul 2017 14:48:41 -0500 Subject: [PATCH 076/109] Fix broken unit test --- consumer_test.go | 32 +++++++++++++------------------- 1 file changed, 13 insertions(+), 19 deletions(-) diff --git a/consumer_test.go b/consumer_test.go index 69df0a1..365c146 100644 --- a/consumer_test.go +++ b/consumer_test.go @@ -53,7 +53,7 @@ func TestConsumer(t *testing.T) { So(k.Session, ShouldNotBeNil) So(k.Session.Config, ShouldNotBeNil) r, err := NewKinesisReader(k.Session.Config, stream, shards[0], - KinesisReaderBatchSize(10000), + KinesisReaderBatchSize(5), //KinesisReaderShardIterator(), KinesisReaderResponseReadTimeout(time.Second), KinesisReaderLogLevel(aws.LogOff), @@ -301,26 +301,20 @@ func TestConsumer(t *testing.T) { So(err, ShouldBeNil) } var count int64 - var wg sync.WaitGroup - wg.Add(1) - go func() { - ctx, cancel := context.WithCancel(context.TODO()) + ctx, cancel := context.WithCancel(context.TODO()) + l.ListenWithContext(ctx, func(m *Message, wg *sync.WaitGroup) error { defer wg.Done() - l.ListenWithContext(ctx, func(m *Message, wg *sync.WaitGroup) error { - defer wg.Done() - time.AfterFunc(time.Duration(rand.Intn(10))*time.Second, func() { - n, err := strconv.Atoi(string(m.Data)) - c.So(err, ShouldBeNil) - atomic.AddInt64(&count, 1) - if n == 15 { - cancel() - } - }) - - return nil + time.AfterFunc(time.Duration(rand.Intn(10))*time.Second, func() { + n, err := strconv.Atoi(string(m.Data)) + c.So(err, ShouldBeNil) + atomic.AddInt64(&count, 1) + if n >= 15 { + cancel() + } }) - }() - wg.Wait() + + return nil + }) So(atomic.LoadInt64(&count), ShouldBeBetweenOrEqual, 1, 20) Printf("(count was %d)", atomic.LoadInt64(&count)) }) From 58b6aa256d39b444f7b7ab0d31f09011704579df Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 31 Jul 2017 16:37:37 -0500 Subject: [PATCH 077/109] Fix some megacheck errors and updated README.md files. --- README.md | 38 ++++------ consumer.go | 146 ++++++++++++++++++------------------ consumer_test.go | 29 ++++--- kinetic_integration_test.go | 6 +- producer.go | 2 +- producer_test.go | 2 + testexec/main.go | 14 ++-- 7 files changed, 117 insertions(+), 120 deletions(-) diff --git a/README.md b/README.md index 20121c3..944e9d0 100644 --- a/README.md +++ b/README.md @@ -4,14 +4,12 @@ # kinetic Kinetic is an MIT-licensed high-performance AWS Kinesis Client for Go -Kinetic wraps [sendgridlabs go-kinesis library](https://github.com/sendgridlabs/go-kinesis) to provide maximum throughput for AWS Kinesis producers and consumers. -An instance of a Kinetic consumer/producer is meant to be used for each shard, so please use it accordingly. If you use more than one instance per-shard then you will -hit the AWS Kinesis throughput [limits](http://docs.aws.amazon.com/kinesis/latest/dev/service-sizes-and-limits.html). +Kinetic wraps [aws-sdk-go](https://github.com/aws/aws-sdk-go.git) to provide maximum throughput with built-in fault tolerance and retry logic for AWS Kinesis producers and consumers. +The Kinetic producer can write to Kinesis or Firehose and the Kinetic listener can consume stream data from Kinesis using the aws-go-sdk or using the Kinesis client library (written in Java). ### Getting Started Before using kinetic, you should make sure you have a created a Kinesis stream and your configuration file has the credentails necessary to read and write to the stream. Once this stream exists in AWS, kinetic will ensure it is in the "ACTIVE" state before running. - ## Testing Tests are written using [goconvey](http://goconvey.co/) and [kinesalite](https://github.com/mhart/kinesalite). Make sure you have kinesalite running locally before attempting to run the tests. They can be run either via the comamnd line: @@ -33,40 +31,34 @@ Kinetic can be used to interface with kinesis like so: ```go import ( "github.com/rewardStyle/kinetic" - "github.com/rewardStyle/kinetic/consumer" - "github.com/rewardStyle/kinetic/message" - "github.com/rewardStyle/kinetic/producer" "sync" ) // Create a kinetic object associated with a local kinesalite stream -k, _ := kinetic.New(func(c *kinetic.Config) { - c.SetCredentials("some-access-key", "some-secret-key", "some-security-token") - c.SetRegion("some-region") - c.SetEndpoint("http://127.0.0.1:4567") -}) +k, _ := kinetic.NewKinetic( + kinetic.AwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + kinetic.AwsConfigRegion("some-region"), + kinetic.AwsConfigEndpoint(""http://127.0.0.1:4567""), +) // Create a kinetic producer -p, _ := producer.NewProducer(func(c *producer.Config) { - c.SetAwsConfig(k.Session.Config) - c.SetKinesisStream("stream-name") -}) +p, _ := kinetic.NewProducer(k.Session.Config, "stream-name") // Create a kinetic consumer -l, _ := consumer.NewConsumer(func(c *consumer.Config) { - c.SetAwsConfig(k.Session.Config) - c.SetReader(consumer.NewKinesisReader("stream-name", "shard-name")) -}) +c, err := kinetic.NewConsumer(k.Session.Config, "stream-name", "shard-name") + -msg, err := l.Retrieve() +// Retrieve one message using the consumer's Retrieve function +msg, err := c.Retrieve() if err != nil { println(err) } // Using Listen - will block unless sent in goroutine -go l.Listen(func(b []byte, fnwg *sync.WaitGroup){ +go c.Listen(func(b []byte, wg *sync.WaitGroup){ + defer wg.Done() + println(string(b)) - fnwg.Done() }) // Send a message using the producer diff --git a/consumer.go b/consumer.go index 081d618..3569f67 100644 --- a/consumer.go +++ b/consumer.go @@ -88,6 +88,7 @@ type Consumer struct { pipeOfDeath chan empty // channel for handling pipe of death consuming bool // flag for indicating whether or not the consumer is consuming consumingMu sync.Mutex // mutex for making the consuming flag thread safe + noCopy noCopy // prevents the Consumer from being copied } // NewConsumer creates a new Consumer object for retrieving and listening to message(s) on a StreamReader. @@ -114,14 +115,14 @@ func NewConsumer(c *aws.Config, stream string, shard string, optionFns ...Consum // startConsuming will initialize the message channel and set consuming to true if there is not already another consume // loop running. -func (l *Consumer) startConsuming() bool { - l.consumingMu.Lock() - defer l.consumingMu.Unlock() - if !l.consuming { - l.consuming = true - l.messages = make(chan *Message, l.queueDepth) - l.concurrencySem = make(chan empty, l.concurrency) - l.pipeOfDeath = make(chan empty) +func (c *Consumer) startConsuming() bool { + c.consumingMu.Lock() + defer c.consumingMu.Unlock() + if !c.consuming { + c.consuming = true + c.messages = make(chan *Message, c.queueDepth) + c.concurrencySem = make(chan empty, c.concurrency) + c.pipeOfDeath = make(chan empty) return true } return false @@ -129,9 +130,9 @@ func (l *Consumer) startConsuming() bool { // shouldConsume is a convenience function that allows functions to break their loops if the context receives a // cancellation or a pipe of death. -func (l *Consumer) shouldConsume(ctx context.Context) (bool, error) { +func (c *Consumer) shouldConsume(ctx context.Context) (bool, error) { select { - case <-l.pipeOfDeath: + case <-c.pipeOfDeath: return false, ErrPipeOfDeath case <-ctx.Done(): return false, ctx.Err() @@ -141,44 +142,41 @@ func (l *Consumer) shouldConsume(ctx context.Context) (bool, error) { } // stopConsuming handles any cleanup after consuming has stopped. -func (l *Consumer) stopConsuming() { - l.consumingMu.Lock() - defer l.consumingMu.Unlock() - if l.consuming && l.messages != nil { - close(l.messages) +func (c *Consumer) stopConsuming() { + c.consumingMu.Lock() + defer c.consumingMu.Unlock() + if c.consuming && c.messages != nil { + close(c.messages) } - if l.concurrencySem != nil { - close(l.concurrencySem) - } - l.consuming = false + c.consuming = false } // enqueueSingle calls the readers's GetRecord method and enqueus a single message on the message channel. -func (l *Consumer) enqueueSingle(ctx context.Context) (int, int, error) { - n, m, err := l.reader.GetRecord(ctx, func(msg *Message, wg *sync.WaitGroup) error { +func (c *Consumer) enqueueSingle(ctx context.Context) (int, int, error) { + n, m, err := c.reader.GetRecord(ctx, func(msg *Message, wg *sync.WaitGroup) error { defer wg.Done() - l.messages <- msg + c.messages <- msg return nil }) if err != nil { - l.handleErrorLogging(err) + c.handleErrorLogging(err) return 0, 0, err } return n, m, nil } // enqueueBatch calls the reader's GetRecords method and enqueues a batch of messages on the message chanel. -func (l *Consumer) enqueueBatch(ctx context.Context) (int, int, error) { - n, m, err := l.reader.GetRecords(ctx, +func (c *Consumer) enqueueBatch(ctx context.Context) (int, int, error) { + n, m, err := c.reader.GetRecords(ctx, func(msg *Message, wg *sync.WaitGroup) error { defer wg.Done() - l.messages <- msg + c.messages <- msg return nil }) if err != nil { - l.handleErrorLogging(err) + c.handleErrorLogging(err) return 0, 0, err } return n, m, nil @@ -186,35 +184,35 @@ func (l *Consumer) enqueueBatch(ctx context.Context) (int, int, error) { // handleErrorLogging is a helper method for handling and logging errors from calling the reader's // GetRecord and GetRecords method. -func (l *Consumer) handleErrorLogging(err error) { +func (c *Consumer) handleErrorLogging(err error) { switch err := err.(type) { case net.Error: if err.Timeout() { - l.Stats.AddGetRecordsTimeout(1) - l.LogError("Received net error:", err.Error()) + c.Stats.AddGetRecordsTimeout(1) + c.LogError("Received net error:", err.Error()) } else { - l.LogError("Received unknown net error:", err.Error()) + c.LogError("Received unknown net error:", err.Error()) } case error: switch err { case ErrTimeoutReadResponseBody: - l.Stats.AddGetRecordsReadTimeout(1) - l.LogError("Received error:", err.Error()) + c.Stats.AddGetRecordsReadTimeout(1) + c.LogError("Received error:", err.Error()) default: - l.LogError("Received error:", err.Error()) + c.LogError("Received error:", err.Error()) } default: - l.LogError("Received unknown error:", err.Error()) + c.LogError("Received unknown error:", err.Error()) } } // RetrieveWithContext waits for a message from the stream and returns the Cancellation is supported through // contexts. -func (l *Consumer) RetrieveWithContext(ctx context.Context) (*Message, error) { - if !l.startConsuming() { +func (c *Consumer) RetrieveWithContext(ctx context.Context) (*Message, error) { + if !c.startConsuming() { return nil, ErrAlreadyConsuming } - defer l.stopConsuming() + defer c.stopConsuming() childCtx, cancel := context.WithCancel(ctx) defer cancel() @@ -222,27 +220,27 @@ func (l *Consumer) RetrieveWithContext(ctx context.Context) (*Message, error) { // A cancellation or closing the pipe of death will cause Retrieve (and related functions) to abort in // between getRecord calls. Note, that this would only occur when there are no new records to retrieve. // Otherwise, getRecords will be allowed to run to completion and deliver one record. - ok, err := l.shouldConsume(ctx) + ok, err := c.shouldConsume(ctx) if !ok { return nil, err } - n, _, err := l.enqueueSingle(childCtx) + n, _, _ := c.enqueueSingle(childCtx) if n > 0 { - l.Stats.AddDelivered(n) - return <-l.messages, nil + c.Stats.AddDelivered(n) + return <-c.messages, nil } } } // Retrieve waits for a message from the stream and returns the value. -func (l *Consumer) Retrieve() (*Message, error) { - return l.RetrieveWithContext(context.TODO()) +func (c *Consumer) Retrieve() (*Message, error) { + return c.RetrieveWithContext(context.TODO()) } // RetrieveFnWithContext retrieves a message from the stream and dispatches it to the supplied function. RetrieveFn // will wait until the function completes. Cancellation is supported through context. -func (l *Consumer) RetrieveFnWithContext(ctx context.Context, fn MessageProcessor) error { - msg, err := l.RetrieveWithContext(ctx) +func (c *Consumer) RetrieveFnWithContext(ctx context.Context, fn MessageProcessor) error { + msg, err := c.RetrieveWithContext(ctx) if err != nil { return err } @@ -253,8 +251,8 @@ func (l *Consumer) RetrieveFnWithContext(ctx context.Context, fn MessageProcesso go func() { start := time.Now() fn(msg, &wg) - l.Stats.AddProcessedDuration(time.Since(start)) - l.Stats.AddProcessed(1) + c.Stats.AddProcessedDuration(time.Since(start)) + c.Stats.AddProcessed(1) }() return nil @@ -262,24 +260,24 @@ func (l *Consumer) RetrieveFnWithContext(ctx context.Context, fn MessageProcesso // RetrieveFn retrieves a message from the stream and dispatches it to the supplied function. RetrieveFn will wait // until the function completes. -func (l *Consumer) RetrieveFn(fn MessageProcessor) error { - return l.RetrieveFnWithContext(context.TODO(), fn) +func (c *Consumer) RetrieveFn(fn MessageProcessor) error { + return c.RetrieveFnWithContext(context.TODO(), fn) } // consume calls getRecords with configured batch size in a loop until the consumer is stopped. -func (l *Consumer) consume(ctx context.Context) { +func (c *Consumer) consume(ctx context.Context) { // We need to run startConsuming to make sure that we are okay and ready to start consuming. This is mainly to // avoid a race condition where Listen() will attempt to read the messages channel prior to consume() // initializing it. We can then launch a goroutine to handle the actual consume operation. - if !l.startConsuming() { + if !c.startConsuming() { return } go func() { - defer l.stopConsuming() + defer c.stopConsuming() // TODO: make these parameters configurable also scale according to the shard count - l.txnCountRateLimiter = rate.NewLimiter(rate.Limit(5), 1) - l.txSizeRateLimiter = rate.NewLimiter(rate.Limit(2000000), 2000000) + c.txnCountRateLimiter = rate.NewLimiter(rate.Limit(5), 1) + c.txSizeRateLimiter = rate.NewLimiter(rate.Limit(2000000), 2000000) childCtx, cancel := context.WithCancel(ctx) defer cancel() @@ -289,13 +287,13 @@ func (l *Consumer) consume(ctx context.Context) { // call below will be allowed to complete (as getRecords does not regard context cancellation). // In the case of cancellation by pipe of death, however, the getRecords will immediately abort // and allow the consume function to immediately abort as well. - if ok, _ := l.shouldConsume(ctx); !ok { + if ok, _ := c.shouldConsume(ctx); !ok { return } - _, payloadSize, err := l.enqueueBatch(childCtx) + _, payloadSize, err := c.enqueueBatch(childCtx) if err != nil { - l.LogError("Encountered an error when calling enqueueBatch: ", err) + c.LogError("Encountered an error when calling enqueueBatch: ", err) return } @@ -304,16 +302,16 @@ func (l *Consumer) consume(ctx context.Context) { go func() { defer wg.Done() - if err := l.txnCountRateLimiter.Wait(childCtx); err != nil { - l.LogError("Error occured waiting for transaction count tokens") + if err := c.txnCountRateLimiter.Wait(childCtx); err != nil { + c.LogError("Error occured waiting for transaction count tokens") } }() wg.Add(1) go func() { defer wg.Done() - if err := l.txSizeRateLimiter.WaitN(childCtx, payloadSize); err != nil { - l.LogError("Error occured waiting for transmission size tokens") + if err := c.txSizeRateLimiter.WaitN(childCtx, payloadSize); err != nil { + c.LogError("Error occured waiting for transmission size tokens") } }() wg.Wait() @@ -323,43 +321,43 @@ func (l *Consumer) consume(ctx context.Context) { // ListenWithContext listens and delivers message to the supplied function. Upon cancellation, Listen will stop the // consumer loop and wait until the messages channel is closed and all messages are delivered. -func (l *Consumer) ListenWithContext(ctx context.Context, fn MessageProcessor) { - l.consume(ctx) +func (c *Consumer) ListenWithContext(ctx context.Context, fn MessageProcessor) { + c.consume(ctx) var wg sync.WaitGroup defer wg.Wait() for { select { - case msg, ok := <-l.messages: + case msg, ok := <-c.messages: if !ok { return } - l.Stats.AddDelivered(1) + c.Stats.AddDelivered(1) // For simplicity, did not do the pipe of death here. If POD is received, we may deliver a // couple more messages (especially since select is random in which channel is read from). - l.concurrencySem <- empty{} + c.concurrencySem <- empty{} wg.Add(1) go func(msg *Message) { defer func() { - <-l.concurrencySem + <-c.concurrencySem }() var fnWg sync.WaitGroup fnWg.Add(1) start := time.Now() fn(msg, &fnWg) fnWg.Wait() - l.Stats.AddProcessedDuration(time.Since(start)) - l.Stats.AddProcessed(1) + c.Stats.AddProcessedDuration(time.Since(start)) + c.Stats.AddProcessed(1) wg.Done() }(msg) - case <-l.pipeOfDeath: - l.LogInfo("ListenWithContext received pipe of death") + case <-c.pipeOfDeath: + c.LogInfo("ListenWithContext received pipe of death") return } } } // Listen listens and delivers message to the supplied function. -func (l *Consumer) Listen(fn MessageProcessor) { - l.ListenWithContext(context.TODO(), fn) +func (c *Consumer) Listen(fn MessageProcessor) { + c.ListenWithContext(context.TODO(), fn) } diff --git a/consumer_test.go b/consumer_test.go index 365c146..c05a9d9 100644 --- a/consumer_test.go +++ b/consumer_test.go @@ -36,6 +36,9 @@ func TestConsumer(t *testing.T) { AwsConfigRegion("some-region"), AwsConfigEndpoint("http://127.0.0.1:4567"), ) + So(k, ShouldNotBeNil) + So(err, ShouldBeNil) + stream := "some-consumer-stream" @@ -302,19 +305,21 @@ func TestConsumer(t *testing.T) { } var count int64 ctx, cancel := context.WithCancel(context.TODO()) - l.ListenWithContext(ctx, func(m *Message, wg *sync.WaitGroup) error { - defer wg.Done() - time.AfterFunc(time.Duration(rand.Intn(10))*time.Second, func() { - n, err := strconv.Atoi(string(m.Data)) - c.So(err, ShouldBeNil) - atomic.AddInt64(&count, 1) - if n >= 15 { - cancel() - } - }) + go func() { + l.ListenWithContext(ctx, func(m *Message, wg *sync.WaitGroup) error { + defer wg.Done() + time.AfterFunc(time.Duration(rand.Intn(3))*time.Second, func() { + n, err := strconv.Atoi(string(m.Data)) + c.So(n, ShouldBeBetweenOrEqual, 0, 19) + c.So(err, ShouldBeNil) + atomic.AddInt64(&count, 1) + }) - return nil - }) + return nil + }) + }() + <-time.After(1 * time.Second) + cancel() So(atomic.LoadInt64(&count), ShouldBeBetweenOrEqual, 1, 20) Printf("(count was %d)", atomic.LoadInt64(&count)) }) diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index 8b33d98..950bb84 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -169,14 +169,14 @@ func TestKineticIntegration(t *testing.T) { assert.NoError(t, err) // Create a new kinetic consumer - l, err := NewConsumer(k.Session.Config, stream, shards[0], + c, err := NewConsumer(k.Session.Config, stream, shards[0], ConsumerReader(r), ConsumerQueueDepth(20), ConsumerConcurrency(10), ConsumerLogLevel(aws.LogOff), ConsumerStats(&NilConsumerStatsCollector{}), ) - assert.NotNil(t, l) + assert.NotNil(t, c) assert.Nil(t, err) numMsg := 1000 @@ -204,7 +204,7 @@ func TestKineticIntegration(t *testing.T) { // Use the consumer to read messages from the kinetic stream go func() { - l.Listen(func(m *Message, fnwg *sync.WaitGroup) error { + c.Listen(func(m *Message, fnwg *sync.WaitGroup) error { defer fnwg.Done() msg := &TestMessage{} diff --git a/producer.go b/producer.go index 2820402..2082913 100644 --- a/producer.go +++ b/producer.go @@ -221,7 +221,7 @@ func (p *Producer) produce() { timer.Stop() return case <-timer.C: - break + continue } } }() diff --git a/producer_test.go b/producer_test.go index 74a301a..6396cdc 100644 --- a/producer_test.go +++ b/producer_test.go @@ -19,6 +19,8 @@ func TestProducer(t *testing.T) { AwsConfigRegion("some-region"), AwsConfigEndpoint("http://127.0.0.1:4567"), ) + So(k, ShouldNotBeNil) + So(err, ShouldBeNil) stream := "some-producer-stream" diff --git a/testexec/main.go b/testexec/main.go index 44b376a..27f2532 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -239,7 +239,7 @@ func newKineticProducer(k *kinetic.Kinetic, streamName string) *kinetic.Producer kinetic.KinesisWriterResponseReadTimeout(time.Second), kinetic.KinesisWriterMsgCountRateLimit(1000), kinetic.KinesisWriterMsgSizeRateLimit(1000000), - kinetic.KinesisWriterLogLevel(kinetic.LogDebug), + kinetic.KinesisWriterLogLevel(aws.LogOff), kinetic.KinesisWriterStats(psc), ) if err != nil { @@ -279,30 +279,30 @@ func newKineticConsumer(k *kinetic.Kinetic, streamName string) *kinetic.Consumer log.Fatalf("Unable to get shards for stream %s due to: %v\n", streamName, err) } - lsc := kinetic.NewDefaultConsumerStatsCollector(registry) + csc := kinetic.NewDefaultConsumerStatsCollector(registry) r, err := kinetic.NewKinesisReader(k.Session.Config, streamName, shards[0], kinetic.KinesisReaderBatchSize(10000), //kinetic.KinesisReaderShardIterator(), kinetic.KinesisReaderResponseReadTimeout(time.Second), kinetic.KinesisReaderLogLevel(aws.LogOff), - kinetic.KinesisReaderStats(lsc), + kinetic.KinesisReaderStats(csc), ) if err != nil { log.Fatalf("Unable to create a new kinesis reader due to: %v\n", err) } - l, err := kinetic.NewConsumer(k.Session.Config, streamName, shards[0], + c, err := kinetic.NewConsumer(k.Session.Config, streamName, shards[0], kinetic.ConsumerReader(r), kinetic.ConsumerQueueDepth(500), kinetic.ConsumerConcurrency(10), kinetic.ConsumerLogLevel(aws.LogOff), - kinetic.ConsumerStats(lsc), + kinetic.ConsumerStats(csc), ) if err != nil { log.Fatalf("Unable to create a new consumer due to: %v\n", err) } - return l + return c } func handlePoD() { @@ -416,7 +416,7 @@ func produce(sd *StreamData, p *kinetic.Producer, wg *sync.WaitGroup) { if *cfg.Blast { sendTicker = time.NewTicker(time.Nanosecond) } else { - sendTicker = time.NewTicker(time.Millisecond) + sendTicker = time.NewTicker(time.Duration(rand.Intn(1) + 1) * time.Millisecond) } produce: for { From 66a064ff66f8d154eea226f7303e0a7611aba493 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 1 Aug 2017 22:46:33 -0500 Subject: [PATCH 078/109] Removed the WaitGroup argument from the messageHandler and MessageProcessor callback functions and refactored the APIs to call the callback function asynchronously. Made KCL action types private. Refactored the KCL reader implementation. Fix typos and inconsistencies in data type names. --- consumer.go | 62 ++++------- consumer_test.go | 18 +-- firehose_writer.go | 18 +-- kcl_actions.go | 74 ++++++------- kcl_reader.go | 211 +++++++++++++----------------------- kinesis_reader.go | 53 +++++---- kinesis_reader_test.go | 2 +- kinesis_writer.go | 2 +- kinetic_integration_test.go | 12 +- producer.go | 4 +- shared_types.go | 27 ++--- testexec/main.go | 6 +- 12 files changed, 196 insertions(+), 293 deletions(-) diff --git a/consumer.go b/consumer.go index 3569f67..aabb38e 100644 --- a/consumer.go +++ b/consumer.go @@ -152,34 +152,25 @@ func (c *Consumer) stopConsuming() { } // enqueueSingle calls the readers's GetRecord method and enqueus a single message on the message channel. -func (c *Consumer) enqueueSingle(ctx context.Context) (int, int, error) { - n, m, err := c.reader.GetRecord(ctx, func(msg *Message, wg *sync.WaitGroup) error { - defer wg.Done() - c.messages <- msg +func (c *Consumer) enqueueSingle(ctx context.Context) (count int, size int) { + count, size, _ = c.reader.GetRecord(ctx, + func(msg *Message) error { + c.messages <- msg + return nil + }) - return nil - }) - if err != nil { - c.handleErrorLogging(err) - return 0, 0, err - } - return n, m, nil + return count, size } // enqueueBatch calls the reader's GetRecords method and enqueues a batch of messages on the message chanel. -func (c *Consumer) enqueueBatch(ctx context.Context) (int, int, error) { - n, m, err := c.reader.GetRecords(ctx, - func(msg *Message, wg *sync.WaitGroup) error { - defer wg.Done() +func (c *Consumer) enqueueBatch(ctx context.Context) (count, size int) { + count, size, _ = c.reader.GetRecords(ctx, + func(msg *Message) error { c.messages <- msg - return nil }) - if err != nil { - c.handleErrorLogging(err) - return 0, 0, err - } - return n, m, nil + + return count, size } // handleErrorLogging is a helper method for handling and logging errors from calling the reader's @@ -224,7 +215,7 @@ func (c *Consumer) RetrieveWithContext(ctx context.Context) (*Message, error) { if !ok { return nil, err } - n, _, _ := c.enqueueSingle(childCtx) + n, _ := c.enqueueSingle(childCtx) if n > 0 { c.Stats.AddDelivered(n) return <-c.messages, nil @@ -245,16 +236,10 @@ func (c *Consumer) RetrieveFnWithContext(ctx context.Context, fn MessageProcesso return err } - var wg sync.WaitGroup - defer wg.Wait() - wg.Add(1) - go func() { - start := time.Now() - fn(msg, &wg) - c.Stats.AddProcessedDuration(time.Since(start)) - c.Stats.AddProcessed(1) - }() - + start := time.Now() + fn(msg) + c.Stats.AddProcessedDuration(time.Since(start)) + c.Stats.AddProcessed(1) return nil } @@ -291,11 +276,7 @@ func (c *Consumer) consume(ctx context.Context) { return } - _, payloadSize, err := c.enqueueBatch(childCtx) - if err != nil { - c.LogError("Encountered an error when calling enqueueBatch: ", err) - return - } + _, size := c.enqueueBatch(childCtx) wg := sync.WaitGroup{} wg.Add(1) @@ -310,7 +291,7 @@ func (c *Consumer) consume(ctx context.Context) { wg.Add(1) go func() { defer wg.Done() - if err := c.txSizeRateLimiter.WaitN(childCtx, payloadSize); err != nil { + if err := c.txSizeRateLimiter.WaitN(childCtx, size); err != nil { c.LogError("Error occured waiting for transmission size tokens") } }() @@ -341,11 +322,8 @@ func (c *Consumer) ListenWithContext(ctx context.Context, fn MessageProcessor) { defer func() { <-c.concurrencySem }() - var fnWg sync.WaitGroup - fnWg.Add(1) start := time.Now() - fn(msg, &fnWg) - fnWg.Wait() + fn(msg) c.Stats.AddProcessedDuration(time.Since(start)) c.Stats.AddProcessed(1) wg.Done() diff --git a/consumer_test.go b/consumer_test.go index c05a9d9..64f926e 100644 --- a/consumer_test.go +++ b/consumer_test.go @@ -205,8 +205,7 @@ func TestConsumer(t *testing.T) { go func() { ctx, cancel := context.WithTimeout(context.TODO(), 1000*time.Millisecond) defer cancel() - l.ListenWithContext(ctx, func(msg *Message, wg *sync.WaitGroup) error { - defer wg.Done() + l.ListenWithContext(ctx, func(msg *Message) error { return nil }) wg.Done() @@ -217,15 +216,12 @@ func TestConsumer(t *testing.T) { wg.Wait() }) - // TODO: Move this test to kinesis_reader_test.go Convey("check that throttle mechanism prevents more than 5 calls to get records", func() { start := time.Now() secs := []float64{} for i := 1; i <= 6; i++ { start := time.Now() - l.reader.GetRecord(context.TODO(), func(msg *Message, wg *sync.WaitGroup) error { - defer wg.Done() - + l.reader.GetRecord(context.TODO(), func(msg *Message) error { return nil }) secs = append(secs, time.Since(start).Seconds()) @@ -240,9 +236,7 @@ func TestConsumer(t *testing.T) { data := "retrieved" _, err := putRecord(l, []byte(data)) So(err, ShouldBeNil) - err = l.RetrieveFn(func(msg *Message, wg *sync.WaitGroup) error { - defer wg.Done() - + err = l.RetrieveFn(func(msg *Message) error { called = true // Note that because this is called in a goroutine, we have to use // the goconvey context @@ -261,8 +255,7 @@ func TestConsumer(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - l.Listen(func(msg *Message, wg *sync.WaitGroup) error { - defer wg.Done() + l.Listen(func(msg *Message) error { atomic.AddInt64(&count, 1) return nil @@ -306,8 +299,7 @@ func TestConsumer(t *testing.T) { var count int64 ctx, cancel := context.WithCancel(context.TODO()) go func() { - l.ListenWithContext(ctx, func(m *Message, wg *sync.WaitGroup) error { - defer wg.Done() + l.ListenWithContext(ctx, func(m *Message) error { time.AfterFunc(time.Duration(rand.Intn(3))*time.Second, func() { n, err := strconv.Atoi(string(m.Data)) c.So(n, ShouldBeBetweenOrEqual, 0, 19) diff --git a/firehose_writer.go b/firehose_writer.go index 681dfce..c4e8669 100644 --- a/firehose_writer.go +++ b/firehose_writer.go @@ -37,13 +37,13 @@ func defaultFirehoseWriterOptions() *firehoseWriterOptions { } } -// FireHoseWriterOptionsFn is a method signature for defining functional option methods for configuring +// FirehoseWriterOptionsFn is a method signature for defining functional option methods for configuring // the FirehoseWriter. -type FireHoseWriterOptionsFn func(*firehoseWriterOptions) error +type FirehoseWriterOptionsFn func(*firehoseWriterOptions) error // FirehoseWriterMsgCountRateLimit is a functional option method for configuring the FirehoseWriter's // message count rate limit. -func FirehoseWriterMsgCountRateLimit(limit int) FireHoseWriterOptionsFn { +func FirehoseWriterMsgCountRateLimit(limit int) FirehoseWriterOptionsFn { return func(o *firehoseWriterOptions) error { if limit > 0 && limit <= firehoseMsgCountRateLimit { o.msgCountRateLimit = limit @@ -55,7 +55,7 @@ func FirehoseWriterMsgCountRateLimit(limit int) FireHoseWriterOptionsFn { // FirehoseWriterMsgSizeRateLimit is a functional option method for configuring the FirehoseWriter's // messsage size rate limit. -func FirehoseWriterMsgSizeRateLimit(limit int) FireHoseWriterOptionsFn { +func FirehoseWriterMsgSizeRateLimit(limit int) FirehoseWriterOptionsFn { return func(o *firehoseWriterOptions) error { if limit > 0 && limit <= firehoseMsgSizeRateLimit { o.msgSizeRateLimit = limit @@ -67,7 +67,7 @@ func FirehoseWriterMsgSizeRateLimit(limit int) FireHoseWriterOptionsFn { // FirehoseWriterThroughputMultiplier is a functional option method for configuring the FirehoseWriter's // throughput multiplier. -func FirehoseWriterThroughputMultiplier(multiplier int) FireHoseWriterOptionsFn { +func FirehoseWriterThroughputMultiplier(multiplier int) FirehoseWriterOptionsFn { return func(o *firehoseWriterOptions) error { if multiplier > 0 { o.throughputMultiplier = multiplier @@ -78,7 +78,7 @@ func FirehoseWriterThroughputMultiplier(multiplier int) FireHoseWriterOptionsFn } // FirehoseWriterLogLevel is a functional option method for configuring the FirehoseWriter's log level. -func FirehoseWriterLogLevel(ll aws.LogLevelType) FireHoseWriterOptionsFn { +func FirehoseWriterLogLevel(ll aws.LogLevelType) FirehoseWriterOptionsFn { return func(o *firehoseWriterOptions) error { o.logLevel = ll & 0xffff0000 return nil @@ -86,7 +86,7 @@ func FirehoseWriterLogLevel(ll aws.LogLevelType) FireHoseWriterOptionsFn { } // FirehoseWriterStats is a functional option method for configuring the FirehoseWriter's stats collector. -func FirehoseWriterStats(sc ProducerStatsCollector) FireHoseWriterOptionsFn { +func FirehoseWriterStats(sc ProducerStatsCollector) FirehoseWriterOptionsFn { return func(o *firehoseWriterOptions) error { o.Stats = sc return nil @@ -102,7 +102,7 @@ type FirehoseWriter struct { } // NewFirehoseWriter creates a new stream writer to write records to a Kinesis. -func NewFirehoseWriter(c *aws.Config, stream string, optionFns ...FireHoseWriterOptionsFn) (*FirehoseWriter, error) { +func NewFirehoseWriter(c *aws.Config, stream string, optionFns ...FirehoseWriterOptionsFn) (*FirehoseWriter, error) { firehoseWriterOptions := defaultFirehoseWriterOptions() for _, optionFn := range optionFns { optionFn(firehoseWriterOptions) @@ -123,7 +123,7 @@ func NewFirehoseWriter(c *aws.Config, stream string, optionFns ...FireHoseWriter } // PutRecords sends a batch of records to Firehose and returns a list of records that need to be retried. -func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*Message, fn MessageHandlerAsync) error { +func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*Message, fn messageHandler) error { var startSendTime time.Time var startBuildTime time.Time diff --git a/kcl_actions.go b/kcl_actions.go index 9b033d2..ef7a3d0 100644 --- a/kcl_actions.go +++ b/kcl_actions.go @@ -7,47 +7,47 @@ import ( ) // ActionType is used as an enum for KCL Multilang protocol action message types -type ActionType string +type kclActionType string // These are the enumerated KCL Multilang protocol action message types const ( - INITIALIZE ActionType = "initialize" - PROCESSRECORDS ActionType = "processRecords" - RECORD ActionType = "record" - CHECKPOINT ActionType = "checkpoint" - SHUTDOWN ActionType = "shutdown" - STATUS ActionType = "status" + kclActionTypeInitialize kclActionType = "initialize" + kclActionTypeProcessRecords kclActionType = "processRecords" + kclActionTypeRecord kclActionType = "record" + kclActionTypeCheckpoint kclActionType = "checkpoint" + kcActionTypeShutdown kclActionType = "shutdown" + KclActionTypeStatus kclActionType = "status" ) -// ActionMessage is a struct used to marshal / unmarshal KCL Multilang protocol action messages -type ActionMessage struct { - Action ActionType `json:"action"` - ShardID string `json:"shardId,omitempty"` - SequenceNumber string `json:"sequenceNumber,omitempty"` - Records []Record `json:"records,omitempty"` - Checkpoint string `json:"checkpoint,omitempty"` - Error string `json:"error,omitempty"` - Reason string `json:"reason,omitempty"` - ResponseFor ActionType `json:"responseFor,omitempty"` +// actionMessage is a struct used to marshal / unmarshal KCL Multilang protocol action messages +type actionMessage struct { + Action kclActionType `json:"action"` + ShardID string `json:"shardId,omitempty"` + SequenceNumber string `json:"sequenceNumber,omitempty"` + Records []record `json:"records,omitempty"` + Checkpoint string `json:"checkpoint,omitempty"` + Error string `json:"error,omitempty"` + Reason string `json:"reason,omitempty"` + ResponseFor kclActionType `json:"responseFor,omitempty"` } -// Record is a struct used to marshal / unmarshal kinesis records from KCL Multilang protocol -type Record struct { - Action ActionType `json:"action"` - ApproximateArrival Timestamp `json:"approximateArrivalTimestamp"` - Data string `json:"data,omitempty"` - PartitionKey string `json:"partitionKey,omitempty"` - SequenceNumber string `json:"sequenceNumber,omitempty"` - SubSequenceNumber int `json:"subSequenceNumber,omitempty"` +// record is a struct used to marshal / unmarshal kinesis records from KCL Multilang protocol +type record struct { + Action kclActionType `json:"action"` + ApproximateArrival timestamp `json:"approximateArrivalTimestamp"` + Data string `json:"data,omitempty"` + PartitionKey string `json:"partitionKey,omitempty"` + SequenceNumber string `json:"sequenceNumber,omitempty"` + SubSequenceNumber int `json:"subSequenceNumber,omitempty"` } -// Timestamp is a time.Time type -type Timestamp struct { +// timestamp is a time.Time type +type timestamp struct { time.Time } // UnmarshalJSON is used as a custom unmarshaller unmarshal unix time stamps -func (t *Timestamp) UnmarshalJSON(b []byte) error { +func (t *timestamp) UnmarshalJSON(b []byte) error { ts, err := strconv.Atoi(string(b)) if err != nil { return err @@ -61,7 +61,7 @@ func (t *Timestamp) UnmarshalJSON(b []byte) error { } // ToMessage is used to transform a multilang.Record struct into a Message struct -func (r *Record) ToMessage() *Message { +func (r *record) ToMessage() *Message { b, err := base64.StdEncoding.DecodeString(r.Data) if err != nil { panic("There was a problem decoding kcl data") @@ -75,18 +75,18 @@ func (r *Record) ToMessage() *Message { } } -// NewCheckpointMessage is used to create a new checkpoint message -func NewCheckpointMessage(seqNum string) *ActionMessage { - return &ActionMessage{ - Action: CHECKPOINT, +// newCheckpointMessage is used to create a new checkpoint message +func newCheckpointMessage(seqNum string) *actionMessage { + return &actionMessage{ + Action: kclActionTypeCheckpoint, Checkpoint: seqNum, } } -// NewStatusMessage is used to create a new status message -func NewStatusMessage(actionType ActionType) *ActionMessage { - return &ActionMessage{ - Action: STATUS, +// newStatusMessage is used to create a new status message +func newStatusMessage(actionType kclActionType) *actionMessage { + return &actionMessage{ + Action: KclActionTypeStatus, ResponseFor: actionType, } } diff --git a/kcl_reader.go b/kcl_reader.go index 14d1c61..b6b7445 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -6,73 +6,55 @@ import ( "context" "encoding/json" "fmt" - "math" "os" - "sync" "github.com/aws/aws-sdk-go/aws" ) +const ( + kclReaderMaxBatchSize = 10000 +) + // kclReaderOptions is a struct that holds all of the KclReader's configurable parameters. type kclReaderOptions struct { - onInitCallbackFn func() error // callback function that gets called after initialization - onCheckpointCallbackFn func() error // callback function that gets called after checkpointing - onShutdownCallbackFn func() error // callback function that gets called after shutdown + batchSize int // maximum records per GetRecordsRequest call // callback function that gets called after shutdown logLevel aws.LogLevelType // log level for configuring the LogHelper's log level Stats ConsumerStatsCollector // stats collection mechanism } -// defaultKlcReaderOptions instantiates a kclReaderOptions with default values. -func defaultKlcReaderOptions() *kclReaderOptions { +// defaultkclReaderOptions instantiates a kclReaderOptions with default values. +func defaultKclReaderOptions() *kclReaderOptions { return &kclReaderOptions{ - onInitCallbackFn: func() error { return nil }, - onCheckpointCallbackFn: func() error { return nil }, - onShutdownCallbackFn: func() error { return nil }, + batchSize: kclReaderMaxBatchSize, logLevel: aws.LogOff, Stats: &NilConsumerStatsCollector{}, } } -// KlcReaderOptionsFn is a method signature for defining functional option methods for configuring the KclReader. -type KlcReaderOptionsFn func(*kclReaderOptions) error - -// KlcReaderOnInitCallbackFn is a functional option method for configuring the KclReader's -// onInitCallbackFn. -func KlcReaderOnInitCallbackFn(fn func() error) KlcReaderOptionsFn { - return func(o *kclReaderOptions) error { - o.onInitCallbackFn = fn - return nil - } -} - -// KlcReaderOnCheckpointCallbackFn is a functional option method for configuring the KclReader's -// onCheckpointCallbackFn. -func KlcReaderOnCheckpointCallbackFn(fn func() error) KlcReaderOptionsFn { - return func(o *kclReaderOptions) error { - o.onCheckpointCallbackFn = fn - return nil - } -} +// kclReaderOptionsFn is a method signature for defining functional option methods for configuring the KclReader. +type kclReaderOptionsFn func(*kclReaderOptions) error -// KlcReaderOnShutdownCallbackFn is a functional option method for configuring the KclReader's -// onShutdownCallbackFn. -func KlcReaderOnShutdownCallbackFn(fn func() error) KlcReaderOptionsFn { +// kclReaderBatchSize is a functional option method for configuring the KclReader's batch size +func kclReaderBatchSize(size int) kclReaderOptionsFn { return func(o *kclReaderOptions) error { - o.onShutdownCallbackFn = fn - return nil + if size >= 0 && size <= kclReaderMaxBatchSize { + o.batchSize = size + return nil + } + return ErrInvalidBatchSize } } -// KlcReaderLogLevel is a functional option method for configuring the KclReader's log level. -func KlcReaderLogLevel(ll aws.LogLevelType) KlcReaderOptionsFn { +// kclReaderLogLevel is a functional option method for configuring the KclReader's log level. +func kclReaderLogLevel(ll aws.LogLevelType) kclReaderOptionsFn { return func(o *kclReaderOptions) error { o.logLevel = ll return nil } } -// KlcReaderStats is a functional option method for configuring the KclReader's stats collector. -func KlcReaderStats(sc ConsumerStatsCollector) KlcReaderOptionsFn { +// kclReaderStats is a functional option method for configuring the KclReader's stats collector. +func kclReaderStats(sc ConsumerStatsCollector) kclReaderOptionsFn { return func(o *kclReaderOptions) error { o.Stats = sc return nil @@ -84,19 +66,20 @@ type KclReader struct { *kclReaderOptions *LogHelper pipeOfDeath chan empty + stop chan empty scanner *bufio.Scanner reader *bufio.Reader - msgBuffer []Message + messages chan *Message } // NewKclReader creates a new stream reader to read records from KCL -func NewKclReader(c *aws.Config, optionFns ...KlcReaderOptionsFn) (*KclReader, error) { - kclReaderOptions := defaultKlcReaderOptions() +func NewKclReader(c *aws.Config, optionFns ...kclReaderOptionsFn) (*KclReader, error) { + kclReaderOptions := defaultKclReaderOptions() for _, optionFn := range optionFns { optionFn(kclReaderOptions) } return &KclReader{ - msgBuffer: []Message{}, + messages: make(chan *Message), kclReaderOptions: kclReaderOptions, LogHelper: &LogHelper{ LogLevel: kclReaderOptions.logLevel, @@ -105,48 +88,29 @@ func NewKclReader(c *aws.Config, optionFns ...KlcReaderOptionsFn) (*KclReader, e }, nil } -// processRecords is a helper method which loops through the message buffer and puts messages on the consumer's -// message channel. After all the messages on the message buffer have been moved to the consumer's message -// channel, a message is sent (following the Multilang protocol) to acknowledge that the processRecords message -// has been received / processed -func (r *KclReader) processRecords(fn MessageHandler, numRecords int) (int, int, error) { - // Define the batchSize - batchSize := 0 - if len(r.msgBuffer) > 0 { - if numRecords < 0 { - batchSize = len(r.msgBuffer) - } else { - batchSize = int(math.Min(float64(len(r.msgBuffer)), float64(numRecords))) - } - } - r.Stats.AddBatchSize(batchSize) - - // TODO: Define the payloadSize - var payloadSize int - - // Loop through the message buffer and call the message handler function on each message - var wg sync.WaitGroup +// processRecords is a helper method which pulls from the reader's message channel and calls the callback function +func (r *KclReader) processRecords(batchSize int, fn messageHandler) (count int, size int, err error) { for i := 0; i < batchSize; i++ { - wg.Add(1) - go fn(&r.msgBuffer[0], &wg) - r.msgBuffer = r.msgBuffer[1:] - r.Stats.AddConsumed(1) - } - wg.Wait() - - // Send an acknowledgement that the 'ProcessRecords' message was received/processed - if len(r.msgBuffer) == 0 { - err := r.sendMessage(NewStatusMessage(PROCESSRECORDS)) + msg := <-r.messages + err = fn(msg) if err != nil { - r.LogError(err) - return batchSize, payloadSize, err + r.LogError("messageHandler resulted in an error: ", err) + } else { + r.Stats.AddConsumed(1) + count++ + b, err := json.Marshal(msg) + if err != nil { + r.LogError("Unable to marshal message: ", err) + } else { + size += len(b) + } } } - return batchSize, payloadSize, nil + return count, size, nil } -func (r *KclReader) getAction() (*ActionMessage, error) { +func (r *KclReader) getAction() (*actionMessage, error) { buffer := &bytes.Buffer{} for { line, isPrefix, err := r.reader.ReadLine() @@ -159,7 +123,7 @@ func (r *KclReader) getAction() (*ActionMessage, error) { } } - actionMsg := &ActionMessage{} + actionMsg := &actionMessage{} err := json.Unmarshal(buffer.Bytes(), actionMsg) if err != nil { fmt.Fprintf(os.Stderr, "Could not understand line read from input: %s\n", buffer.String()) @@ -170,7 +134,6 @@ func (r *KclReader) getAction() (*ActionMessage, error) { // processAction listens to STDIN and processes action messages based on the Multilang protocol from KCL func (r *KclReader) processAction() error { for { - // Retrieve action message actionMessage, err := r.getAction() if err != nil { @@ -181,30 +144,41 @@ func (r *KclReader) processAction() error { } switch actionMessage.Action { - case INITIALIZE: - r.onInit() - r.sendMessage(NewStatusMessage(INITIALIZE)) - case CHECKPOINT: - r.onCheckpoint() - r.sendMessage(NewStatusMessage(CHECKPOINT)) - case SHUTDOWN: - r.onShutdown() - r.sendMessage(NewStatusMessage(SHUTDOWN)) - case PROCESSRECORDS: - go func() error { - for _, msg := range actionMessage.Records { - r.msgBuffer = append(r.msgBuffer, *msg.ToMessage()) - } - - return nil - }() + case kclActionTypeInitialize: + err := r.sendMessage(newStatusMessage(kclActionTypeInitialize)) + if err != nil { + r.LogError("Unable to send Initialize acknowledgement due to: ", err) + } + case kclActionTypeCheckpoint: + err := r.sendMessage(newStatusMessage(kclActionTypeCheckpoint)) + if err != nil { + r.LogError("Unable to send Checkpoint acknowledgement due to: ", err) + } + case kcActionTypeShutdown: + err := r.sendMessage(newStatusMessage(kcActionTypeShutdown)) + if err != nil { + r.LogError("Unable to send Shutdown acknowledgement due to: ", err) + } + case kclActionTypeProcessRecords: + // Put all the messages on the reader's message channel + for _, msg := range actionMessage.Records { + r.messages <-msg.ToMessage() + } + + // Send an acknowledgement that all the messages were received + err := r.sendMessage(newStatusMessage(kclActionTypeProcessRecords)) + if err != nil { + r.LogError("Unable to send ProcessRecords acknowledgement due to: ", err) + } default: + r.LogError("processAction received an invalid action: ", actionMessage.Action) } } return nil } +// sendMessage func (r *KclReader) sendMessage(msg interface{}) error { b, err := json.Marshal(msg) if err != nil { @@ -217,47 +191,16 @@ func (r *KclReader) sendMessage(msg interface{}) error { return nil } -func (r *KclReader) onInit() error { - if r.onInitCallbackFn != nil { - err := r.onInitCallbackFn() - if err != nil { - r.LogError(err) - return err - } - } - return nil -} - -func (r *KclReader) onCheckpoint() error { - if r.onCheckpointCallbackFn != nil { - err := r.onCheckpointCallbackFn() - if err != nil { - r.LogError(err) - return err - } - } - return nil -} - -func (r *KclReader) onShutdown() error { - if r.onShutdownCallbackFn != nil { - err := r.onShutdownCallbackFn() - if err != nil { - r.LogError(err) - return err - } - } - return nil -} - // GetRecord calls processRecords to attempt to put one message from message buffer to the consumer's message // channel -func (r *KclReader) GetRecord(ctx context.Context, fn MessageHandler) (int, int, error) { - return r.processRecords(fn, 1) +func (r *KclReader) GetRecord(ctx context.Context, fn messageHandler) (count int, size int, err error) { + count, size, err = r.processRecords(1, fn) + return count, size, err } // GetRecords calls processRecords to attempt to put all messages on the message buffer on the consumer's // message channel -func (r *KclReader) GetRecords(ctx context.Context, fn MessageHandler) (int, int, error) { - return r.processRecords(fn, -1) +func (r *KclReader) GetRecords(ctx context.Context, fn messageHandler) (count int, size int, err error) { + count, size, err = r.processRecords(r.batchSize, fn) + return count, size, err } diff --git a/kinesis_reader.go b/kinesis_reader.go index a531fbe..48d9d01 100644 --- a/kinesis_reader.go +++ b/kinesis_reader.go @@ -3,7 +3,6 @@ package kinetic import ( "context" "fmt" - "sync" "time" "github.com/aws/aws-sdk-go/aws" @@ -15,14 +14,14 @@ import ( ) const ( - kinesisReaderBatchSize = 10000 + kinesisReaderMaxBatchSize = 10000 ) // kinesisReaderOptions a struct that holds all of the KinesisReader's configurable parameters. type kinesisReaderOptions struct { - batchSize int - shardIterator *ShardIterator - responseReadTimeout time.Duration + batchSize int // maximum records per GetRecordsRequest call + shardIterator *ShardIterator // shard iterator for Kinesis stream + responseReadTimeout time.Duration // response read time out for GetRecordsRequest API call logLevel aws.LogLevelType // log level for configuring the LogHelper's log level Stats ConsumerStatsCollector // stats collection mechanism } @@ -30,7 +29,7 @@ type kinesisReaderOptions struct { // defaultKinesisReaderOptions instantiates a kinesisReaderOptions with default values. func defaultKinesisReaderOptions() *kinesisReaderOptions { return &kinesisReaderOptions{ - batchSize: kinesisReaderBatchSize, + batchSize: kinesisReaderMaxBatchSize, shardIterator: NewShardIterator(), responseReadTimeout: time.Second, Stats: &NilConsumerStatsCollector{}, @@ -45,7 +44,7 @@ type KinesisReaderOptionsFn func(*kinesisReaderOptions) error // batch size. func KinesisReaderBatchSize(size int) KinesisReaderOptionsFn { return func(o *kinesisReaderOptions) error { - if size > 0 && size <= kinesisReaderBatchSize { + if size > 0 && size <= kinesisReaderMaxBatchSize { o.batchSize = size return nil } @@ -188,9 +187,10 @@ func (r *KinesisReader) throttle(sem chan empty) { }) } -func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batchSize int) (int, int, error) { - if err := r.ensureShardIterator(); err != nil { - return 0, 0, err +func (r *KinesisReader) getRecords(ctx context.Context, fn messageHandler, batchSize int) (count int, size int, err error) { + if err = r.ensureShardIterator(); err != nil { + r.LogError("Error calling ensureShardIterator(): ", err) + return count, size, err } r.throttle(r.throttleSem) @@ -237,9 +237,8 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch } }) - var payloadSize int req.Handlers.Unmarshal.PushBack(func(req *request.Request) { - payloadSize += int(req.HTTPRequest.ContentLength) + size += int(req.HTTPRequest.ContentLength) r.Stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) r.LogDebug("Finished GetRecords Unmarshal, took", time.Since(start)) }) @@ -247,7 +246,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch // Send the GetRecords request r.LogDebug("Starting GetRecords Build/Sign request, took", time.Since(start)) r.Stats.AddGetRecordsCalled(1) - if err := req.Send(); err != nil { + if err = req.Send(); err != nil { r.LogError("Error getting records:", err) switch err.(awserr.Error).Code() { case kinesis.ErrCodeProvisionedThroughputExceededException: @@ -255,16 +254,15 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch default: r.LogDebug("Received AWS error:", err.Error()) } - return 0, 0, err + return count, size, err } r.Stats.AddGetRecordsDuration(time.Since(start)) // Process Records r.LogDebug(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), r.shard, time.Since(start))) if resp == nil { - return 0, 0, ErrNilGetRecordsResponse + return count, size, ErrNilGetRecordsResponse } - delivered := 0 r.Stats.AddBatchSize(len(resp.Records)) for _, record := range resp.Records { if record != nil { @@ -273,14 +271,11 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch // the current batch of records. select { case <-ctx.Done(): - r.LogInfo(fmt.Sprintf("getRecords received ctx.Done() while delivering messages, %d delivered, ~%d dropped", delivered, len(resp.Records)-delivered)) - return delivered, payloadSize, ctx.Err() + r.LogInfo(fmt.Sprintf("getRecords received ctx.Done() while delivering messages, %d delivered, ~%d dropped", count, len(resp.Records)-count)) + return count, size, ctx.Err() default: - var wg sync.WaitGroup - wg.Add(1) - go fn(FromRecord(record), &wg) - wg.Wait() - delivered++ + fn(FromRecord(record)) + count++ r.Stats.AddConsumed(1) if record.SequenceNumber != nil { // We can safely ignore if this call returns error, as if we somehow receive an @@ -308,15 +303,17 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn MessageHandler, batch // around. r.setNextShardIterator(*resp.NextShardIterator) } - return delivered, payloadSize, nil + return count, size, nil } // GetRecord calls getRecords and delivers one record into the messages channel. -func (r *KinesisReader) GetRecord(ctx context.Context, fn MessageHandler) (int, int, error) { - return r.getRecords(ctx, fn, 1) +func (r *KinesisReader) GetRecord(ctx context.Context, fn messageHandler) (count int, size int, err error) { + count, size, err = r.getRecords(ctx, fn, 1) + return count, size, err } // GetRecords calls getRecords and delivers each record into the messages channel. -func (r *KinesisReader) GetRecords(ctx context.Context, fn MessageHandler) (int, int, error) { - return r.getRecords(ctx, fn, r.batchSize) +func (r *KinesisReader) GetRecords(ctx context.Context, fn messageHandler) (count int, size int, err error) { + count, size, err = r.getRecords(ctx, fn, r.batchSize) + return count, size, err } diff --git a/kinesis_reader_test.go b/kinesis_reader_test.go index dba0061..60ca651 100644 --- a/kinesis_reader_test.go +++ b/kinesis_reader_test.go @@ -38,7 +38,7 @@ func TestNewKinesisReader(t *testing.T) { }) Convey("check that we can create a new KinesisReader with configured values", func() { - batchSize := rand.Intn(kinesisReaderBatchSize) + batchSize := rand.Intn(kinesisReaderMaxBatchSize) respReadTimeout := time.Duration(rand.Int()) * time.Second logLevel := aws.LogDebug | aws.LogDebugWithSigning | LogDebug shardIterator := NewShardIterator() diff --git a/kinesis_writer.go b/kinesis_writer.go index 205e67c..b2b05d7 100644 --- a/kinesis_writer.go +++ b/kinesis_writer.go @@ -120,7 +120,7 @@ func NewKinesisWriter(c *aws.Config, stream string, optionFns ...KinesisWriterOp } // PutRecords sends a batch of records to Kinesis and returns a list of records that need to be retried. -func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*Message, fn MessageHandlerAsync) error { +func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*Message, fn messageHandler) error { var startSendTime time.Time var startBuildTime time.Time diff --git a/kinetic_integration_test.go b/kinetic_integration_test.go index 950bb84..91decf4 100644 --- a/kinetic_integration_test.go +++ b/kinetic_integration_test.go @@ -185,7 +185,7 @@ func TestKineticIntegration(t *testing.T) { // Use the producer to write messages to the kinetic stream wg := sync.WaitGroup{} wg.Add(numMsg + 1) - go func(sent *int) { + go func() { defer wg.Done() for i := 0; i < numMsg; i++ { msg := &TestMessage{ @@ -197,16 +197,14 @@ func TestKineticIntegration(t *testing.T) { PartitionKey: aws.String("key"), Data: []byte(jsonStr), }); err == nil { - *sent++ + numSent++ } } - }(&numSent) + }() // Use the consumer to read messages from the kinetic stream go func() { - c.Listen(func(m *Message, fnwg *sync.WaitGroup) error { - defer fnwg.Done() - + c.Listen(func(m *Message) error { msg := &TestMessage{} json.Unmarshal(m.Data, msg) @@ -221,8 +219,8 @@ func TestKineticIntegration(t *testing.T) { return nil }) }() - wg.Wait() + assert.Equal(t, int(numSent), numMsg, "Number of message sent should equal the number of messages") assert.Equal(t, streamData.size(), numMsg, "Number of messages") diff --git a/producer.go b/producer.go index 2082913..b444148 100644 --- a/producer.go +++ b/producer.go @@ -24,7 +24,7 @@ type producerOptions struct { maxRetryAttempts int // maximum number of retry attempts for failed messages concurrency int // number of concurrent workers per shard shardCheckFreq time.Duration // frequency (specified as a duration) with which to check the the shard size - dataSpillFn MessageHandlerAsync // callback function for handling dropped messages that the producer was unable to send to the stream + dataSpillFn MessageProcessor // callback function for handling dropped messages that the producer was unable to send to the stream logLevel aws.LogLevelType // log level for configuring the LogHelper's log level Stats ProducerStatsCollector // stats collection mechanism } @@ -116,7 +116,7 @@ func ProducerShardCheckFrequency(duration time.Duration) ProducerOptionsFn { } // ProducerDataSpillFn is a functional option method for configuing the producer's data spill callback function. -func ProducerDataSpillFn(fn MessageHandlerAsync) ProducerOptionsFn { +func ProducerDataSpillFn(fn MessageProcessor) ProducerOptionsFn { return func(o *producerOptions) error { o.dataSpillFn = fn return nil diff --git a/shared_types.go b/shared_types.go index bdda51b..11681e2 100644 --- a/shared_types.go +++ b/shared_types.go @@ -2,12 +2,11 @@ package kinetic import ( "context" - "sync" ) // StreamWriter is an interface that abstracts the differences in API between Kinesis and Firehose. type StreamWriter interface { - PutRecords(context.Context, []*Message, MessageHandlerAsync) error + PutRecords(context.Context, []*Message, messageHandler) error getMsgCountRateLimit() int getMsgSizeRateLimit() int getConcurrencyMultiplier() (int, error) @@ -15,22 +14,20 @@ type StreamWriter interface { // StreamReader is an interface that abstracts out a stream reader. type StreamReader interface { - GetRecord(context.Context, MessageHandler) (int, int, error) - GetRecords(context.Context, MessageHandler) (int, int, error) + GetRecord(context.Context, messageHandler) (count int, size int, err error) + GetRecords(context.Context, messageHandler) (count int, size int, err error) } -// MessageProcessor defines the signature of a message handler used by Listen, RetrieveFn and their associated -// *WithContext functions. MessageHandler accepts a WaitGroup so the function can be run as a blocking operation as -// opposed to MessageHandlerAsync. -type MessageProcessor func(*Message, *sync.WaitGroup) error +// MessageProcessor defines the signature of a (asynchronous) callback function used by Listen, RetrieveFn and +// their associated *WithContext functions. MessageProcessor is a user-defined callback function for processing +// messages after they have been pulled off of the consumer's message channel or for processing the producer's +// dropped message. +type MessageProcessor func(*Message) error -// MessageHandler defines the signature of a message handler used by PutRecords(). MessageHandler accepts a WaitGroup -// so the function can be run as a blocking operation as opposed to MessageHandlerAsync. -type MessageHandler func(*Message, *sync.WaitGroup) error - -// MessageHandlerAsync defines the signature of a message handler used by PutRecords(). MessageHandlerAsync is meant to -// be run asynchronously. -type MessageHandlerAsync func(*Message) error +// messageHandler defines the signature of a message handler used by PutRecords(), GetRecord() and GetRecords(). +// The messageHandler is used as a callback function defined by the producer/consumer so the writers/readers +// don't need to know about the producer's/consumer's message channels. +type messageHandler func(*Message) error // empty is used a as a dummy type for counting semaphore channels. type empty struct{} diff --git a/testexec/main.go b/testexec/main.go index 27f2532..3d6b485 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -416,7 +416,7 @@ func produce(sd *StreamData, p *kinetic.Producer, wg *sync.WaitGroup) { if *cfg.Blast { sendTicker = time.NewTicker(time.Nanosecond) } else { - sendTicker = time.NewTicker(time.Duration(rand.Intn(1) + 1) * time.Millisecond) + sendTicker = time.NewTicker(time.Duration(rand.Intn(1)+1) * time.Millisecond) } produce: for { @@ -496,9 +496,7 @@ func listen(sd *StreamData, c *kinetic.Consumer, wg *sync.WaitGroup) { // Call Listen within a go routine go func() { - c.Listen(func(m *kinetic.Message, wg *sync.WaitGroup) error { - defer wg.Done() - + c.Listen(func(m *kinetic.Message) error { // Unmarshal data msg := &Message{} json.Unmarshal(m.Data, msg) From 40157db7ef7ceb2fe7cfd1ec7678c4b8b01bd85d Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 2 Aug 2017 15:31:05 -0500 Subject: [PATCH 079/109] Refactor KclReader implementation to add startup and shutdown functions. Renamed provisioned through exceeded stats for producer and consumer. Added randomization to dummy messages and fixed partion key error for test executable. --- consumer.go | 12 +- consumer_stats.go | 20 ++-- consumer_stats_test.go | 10 +- firehose_writer.go | 2 +- kcl_reader.go | 246 ++++++++++++++++++++++++++--------------- kinesis_reader.go | 2 +- kinesis_writer.go | 2 +- producer.go | 14 ++- producer_stats.go | 10 +- producer_stats_test.go | 14 +-- testexec/main.go | 6 +- testexec/messages.go | 10 +- 12 files changed, 218 insertions(+), 130 deletions(-) diff --git a/consumer.go b/consumer.go index aabb38e..d658e49 100644 --- a/consumer.go +++ b/consumer.go @@ -153,22 +153,30 @@ func (c *Consumer) stopConsuming() { // enqueueSingle calls the readers's GetRecord method and enqueus a single message on the message channel. func (c *Consumer) enqueueSingle(ctx context.Context) (count int, size int) { - count, size, _ = c.reader.GetRecord(ctx, + var err error + count, size, err = c.reader.GetRecord(ctx, func(msg *Message) error { c.messages <- msg return nil }) + if err != nil { + c.handleErrorLogging(err) + } return count, size } // enqueueBatch calls the reader's GetRecords method and enqueues a batch of messages on the message chanel. func (c *Consumer) enqueueBatch(ctx context.Context) (count, size int) { - count, size, _ = c.reader.GetRecords(ctx, + var err error + count, size, err = c.reader.GetRecords(ctx, func(msg *Message) error { c.messages <- msg return nil }) + if err != nil { + c.handleErrorLogging(err) + } return count, size } diff --git a/consumer_stats.go b/consumer_stats.go index 1849f56..e940edf 100644 --- a/consumer_stats.go +++ b/consumer_stats.go @@ -16,7 +16,7 @@ type ConsumerStatsCollector interface { AddProcessed(int) AddBatchSize(int) AddGetRecordsCalled(int) - AddGetRecordsProvisionedThroughputExceeded(int) + AddReadProvisionedThroughputExceeded(int) AddGetRecordsTimeout(int) AddGetRecordsReadTimeout(int) AddProcessedDuration(time.Duration) @@ -50,9 +50,9 @@ func (nsc *NilConsumerStatsCollector) AddBatchSize(int) {} // by the consumer. func (nsc *NilConsumerStatsCollector) AddGetRecordsCalled(int) {} -// AddGetRecordsProvisionedThroughputExceeded records the number of times the GetRecords +// AddReadProvisionedThroughputExceeded records the number of times the GetRecords // API returned a ErrCodeProvisionedThroughputExceededException by the consumer. -func (nsc *NilConsumerStatsCollector) AddGetRecordsProvisionedThroughputExceeded(int) {} +func (nsc *NilConsumerStatsCollector) AddReadProvisionedThroughputExceeded(int) {} // AddGetRecordsTimeout records the number of times the GetRecords API timed out // on the HTTP level. This is influenced by the WithHTTPClientTimeout @@ -87,7 +87,7 @@ const ( MetricsProcessed = "kinetic.consumer.processed" MetricsBatchSize = "kinetic.consumer.batchsize" MetricsSent = "kinetic.consumer.sent" - MetricsGetRecordsProvisionedThroughputExceeded = "kinetic.consumer.getrecords.provisionedthroughputexceeded" + MetricsReadProvisionedThroughputExceeded = "kinetic.consumer.getrecords.provisionedthroughputexceeded" MetricsGetRecordsTimeout = "kinetic.consumer.getrecords.timeout" MetricsGetRecordsReadTimeout = "kinetic.consumer.getrecords.readtimeout" MetricsProcessedDuration = "kinetic.consumer.processed.duration" @@ -104,7 +104,7 @@ type DefaultConsumerStatsCollector struct { Processed metrics.Counter BatchSize metrics.Counter GetRecordsCalled metrics.Counter - GetRecordsProvisionedThroughputExceeded metrics.Counter + ReadProvisionedThroughputExceeded metrics.Counter GetRecordsTimeout metrics.Counter GetRecordsReadTimeout metrics.Counter ProcessedDuration metrics.Gauge @@ -121,7 +121,7 @@ func NewDefaultConsumerStatsCollector(r metrics.Registry) *DefaultConsumerStatsC Processed: metrics.GetOrRegisterCounter(MetricsProcessed, r), BatchSize: metrics.GetOrRegisterCounter(MetricsBatchSize, r), GetRecordsCalled: metrics.GetOrRegisterCounter(MetricsSent, r), - GetRecordsProvisionedThroughputExceeded: metrics.GetOrRegisterCounter(MetricsGetRecordsProvisionedThroughputExceeded, r), + ReadProvisionedThroughputExceeded: metrics.GetOrRegisterCounter(MetricsReadProvisionedThroughputExceeded, r), GetRecordsTimeout: metrics.GetOrRegisterCounter(MetricsGetRecordsTimeout, r), GetRecordsReadTimeout: metrics.GetOrRegisterCounter(MetricsGetRecordsReadTimeout, r), ProcessedDuration: metrics.GetOrRegisterGauge(MetricsProcessedDuration, r), @@ -163,10 +163,10 @@ func (dsc *DefaultConsumerStatsCollector) AddGetRecordsCalled(count int) { dsc.GetRecordsCalled.Inc(int64(count)) } -// AddGetRecordsProvisionedThroughputExceeded records the number of times the GetRecords +// AddReadProvisionedThroughputExceeded records the number of times the GetRecords // API returned a ErrCodeProvisionedThroughputExceededException by the consumer. -func (dsc *DefaultConsumerStatsCollector) AddGetRecordsProvisionedThroughputExceeded(count int) { - dsc.GetRecordsProvisionedThroughputExceeded.Inc(int64(count)) +func (dsc *DefaultConsumerStatsCollector) AddReadProvisionedThroughputExceeded(count int) { + dsc.ReadProvisionedThroughputExceeded.Inc(int64(count)) } // AddGetRecordsTimeout records the number of times the GetRecords API timed out @@ -216,7 +216,7 @@ func (dsc *DefaultConsumerStatsCollector) PrintStats() { log.Printf("Consumer stats: GetRecords Called: [%d]\n", dsc.GetRecordsCalled.Count()) log.Printf("Consumer stats: GetRecords Timeout: [%d]\n", dsc.GetRecordsTimeout.Count()) log.Printf("Consumer stats: GetRecords Read Timeout: [%d]\n", dsc.GetRecordsReadTimeout.Count()) - log.Printf("Consumer stats: GetRecords Provisioned Throughput Exceeded: [%d]\n", dsc.GetRecordsProvisionedThroughputExceeded.Count()) + log.Printf("Consumer stats: GetRecords Provisioned Throughput Exceeded: [%d]\n", dsc.ReadProvisionedThroughputExceeded.Count()) log.Printf("Consumer stats: Processed Duration (ns): [%d]\n", dsc.ProcessedDuration.Value()) log.Printf("Consumer stats: GetRecords Duration (ns): [%d]\n", dsc.GetRecordsDuration.Value()) log.Printf("Consumer stats: GetRecords Read Response Duration (ns): [%d]\n", dsc.GetRecordsReadResponseDuration.Value()) diff --git a/consumer_stats_test.go b/consumer_stats_test.go index a3ab7bb..4cc75b3 100644 --- a/consumer_stats_test.go +++ b/consumer_stats_test.go @@ -38,8 +38,8 @@ func TestConsumerStatsCollector(t *testing.T) { sc.AddGetRecordsCalled(1) }) - Convey("check that AddGetRecordsProvisionedThroughputExceeded does not error", func() { - sc.AddGetRecordsProvisionedThroughputExceeded(1) + Convey("check that AddReadProvisionedThroughputExceeded does not error", func() { + sc.AddReadProvisionedThroughputExceeded(1) }) Convey("check that AddGetRecordsTimeout does not error", func() { @@ -102,10 +102,10 @@ func TestConsumerStatsCollector(t *testing.T) { So(sc.(*DefaultConsumerStatsCollector).GetRecordsCalled.Count(), ShouldEqual, int64(count)) }) - Convey("check that AddGetRecordsProvisionedThroughputExceeded does not error", func() { + Convey("check that AddReadProvisionedThroughputExceeded does not error", func() { count := rand.Int() - sc.AddGetRecordsProvisionedThroughputExceeded(count) - So(sc.(*DefaultConsumerStatsCollector).GetRecordsProvisionedThroughputExceeded.Count(), ShouldEqual, int64(count)) + sc.AddReadProvisionedThroughputExceeded(count) + So(sc.(*DefaultConsumerStatsCollector).ReadProvisionedThroughputExceeded.Count(), ShouldEqual, int64(count)) }) Convey("check that AddGetRecordsTimeout does not error", func() { diff --git a/firehose_writer.go b/firehose_writer.go index c4e8669..42e1d2b 100644 --- a/firehose_writer.go +++ b/firehose_writer.go @@ -186,7 +186,7 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*Message, fn } else { switch aws.StringValue(record.ErrorCode) { case firehose.ErrCodeLimitExceededException: - w.Stats.AddPutRecordsProvisionedThroughputExceeded(1) + w.Stats.AddWriteProvisionedThroughputExceeded(1) default: w.LogDebug("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(record.ErrorMessage)) } diff --git a/kcl_reader.go b/kcl_reader.go index b6b7445..68dd990 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -7,6 +7,7 @@ import ( "encoding/json" "fmt" "os" + "sync" "github.com/aws/aws-sdk-go/aws" ) @@ -17,7 +18,10 @@ const ( // kclReaderOptions is a struct that holds all of the KclReader's configurable parameters. type kclReaderOptions struct { - batchSize int // maximum records per GetRecordsRequest call // callback function that gets called after shutdown + batchSize int // maximum records per GetRecordsRequest call + onInitCallbackFn func() error // callback function that gets called after initialization + onCheckpointCallbackFn func() error // callback function that gets called after checkpointing + onShutdownCallbackFn func() error // callback function that gets called after shutdown logLevel aws.LogLevelType // log level for configuring the LogHelper's log level Stats ConsumerStatsCollector // stats collection mechanism } @@ -25,19 +29,19 @@ type kclReaderOptions struct { // defaultkclReaderOptions instantiates a kclReaderOptions with default values. func defaultKclReaderOptions() *kclReaderOptions { return &kclReaderOptions{ - batchSize: kclReaderMaxBatchSize, - logLevel: aws.LogOff, - Stats: &NilConsumerStatsCollector{}, + batchSize: kclReaderMaxBatchSize, + logLevel: aws.LogOff, + Stats: &NilConsumerStatsCollector{}, } } -// kclReaderOptionsFn is a method signature for defining functional option methods for configuring the KclReader. -type kclReaderOptionsFn func(*kclReaderOptions) error +// KclReaderOptionsFn is a method signature for defining functional option methods for configuring the KclReader. +type KclReaderOptionsFn func(*kclReaderOptions) error // kclReaderBatchSize is a functional option method for configuring the KclReader's batch size -func kclReaderBatchSize(size int) kclReaderOptionsFn { +func kclReaderBatchSize(size int) KclReaderOptionsFn { return func(o *kclReaderOptions) error { - if size >= 0 && size <= kclReaderMaxBatchSize { + if size > 0 && size <= kclReaderMaxBatchSize { o.batchSize = size return nil } @@ -45,8 +49,35 @@ func kclReaderBatchSize(size int) kclReaderOptionsFn { } } +// KclReaderOnInitCallbackFn is a functional option method for configuring the KclReader's +// onInitCallbackFn. +func KclReaderOnInitCallbackFn(fn func() error) KclReaderOptionsFn { + return func(o *kclReaderOptions) error { + o.onInitCallbackFn = fn + return nil + } +} + +// KclReaderOnCheckpointCallbackFn is a functional option method for configuring the KclReader's +// onCheckpointCallbackFn. +func KclReaderOnCheckpointCallbackFn(fn func() error) KclReaderOptionsFn { + return func(o *kclReaderOptions) error { + o.onCheckpointCallbackFn = fn + return nil + } +} + +// KclReaderOnShutdownCallbackFn is a functional option method for configuring the KclReader's +// onShutdownCallbackFn. +func KclReaderOnShutdownCallbackFn(fn func() error) KclReaderOptionsFn { + return func(o *kclReaderOptions) error { + o.onShutdownCallbackFn = fn + return nil + } +} + // kclReaderLogLevel is a functional option method for configuring the KclReader's log level. -func kclReaderLogLevel(ll aws.LogLevelType) kclReaderOptionsFn { +func kclReaderLogLevel(ll aws.LogLevelType) KclReaderOptionsFn { return func(o *kclReaderOptions) error { o.logLevel = ll return nil @@ -54,7 +85,7 @@ func kclReaderLogLevel(ll aws.LogLevelType) kclReaderOptionsFn { } // kclReaderStats is a functional option method for configuring the KclReader's stats collector. -func kclReaderStats(sc ConsumerStatsCollector) kclReaderOptionsFn { +func kclReaderStats(sc ConsumerStatsCollector) KclReaderOptionsFn { return func(o *kclReaderOptions) error { o.Stats = sc return nil @@ -63,23 +94,21 @@ func kclReaderStats(sc ConsumerStatsCollector) kclReaderOptionsFn { // KclReader handles the KCL Multilang Protocol to read records from KCL type KclReader struct { - *kclReaderOptions - *LogHelper - pipeOfDeath chan empty - stop chan empty - scanner *bufio.Scanner - reader *bufio.Reader - messages chan *Message + *kclReaderOptions // contains all of the configuration settings for the KclReader + *LogHelper // object for help with logging + reader *bufio.Reader // io reader to read from STDIN + messages chan *Message // unbuffered message channel used to throttle the record processing from KCL + startupOnce sync.Once // used to ensure that the startup function is called once + shutdownOnce sync.Once // used to ensure that the shutdown function is called once } // NewKclReader creates a new stream reader to read records from KCL -func NewKclReader(c *aws.Config, optionFns ...kclReaderOptionsFn) (*KclReader, error) { +func NewKclReader(c *aws.Config, optionFns ...KclReaderOptionsFn) (*KclReader, error) { kclReaderOptions := defaultKclReaderOptions() for _, optionFn := range optionFns { optionFn(kclReaderOptions) } return &KclReader{ - messages: make(chan *Message), kclReaderOptions: kclReaderOptions, LogHelper: &LogHelper{ LogLevel: kclReaderOptions.logLevel, @@ -88,34 +117,96 @@ func NewKclReader(c *aws.Config, optionFns ...kclReaderOptionsFn) (*KclReader, e }, nil } -// processRecords is a helper method which pulls from the reader's message channel and calls the callback function -func (r *KclReader) processRecords(batchSize int, fn messageHandler) (count int, size int, err error) { - for i := 0; i < batchSize; i++ { - msg := <-r.messages - err = fn(msg) - if err != nil { - r.LogError("messageHandler resulted in an error: ", err) - } else { - r.Stats.AddConsumed(1) - count++ - b, err := json.Marshal(msg) - if err != nil { - r.LogError("Unable to marshal message: ", err) - } else { - size += len(b) +func (r *KclReader) process(ctx context.Context) { + r.startupOnce.Do(func (){ + defer func(){ + // Reset shutdownOnce to allow the shut down sequence to happen again + r.shutdownOnce = sync.Once{} + }() + + // create communication channels + r.messages = make(chan *Message) + + // listen to STDIN and processes action messages based on the Multilang protocol from KCL + go func(){ + defer r.shutdown() + + for { + select { + case <-ctx.Done(): + r.LogInfo("KclReader received ctx.Done() while processing messages from KCL") + return + default: + } + + // Retrieve action message + actionMessage, err := r.getAction() + if err != nil || actionMessage == nil { + return + } + + switch actionMessage.Action { + case kclActionTypeInitialize: + r.LogDebug("Receieved Initialize action from KCL") + r.onInitCallbackFn() + err := r.sendMessage(newStatusMessage(kclActionTypeInitialize)) + if err != nil { + r.LogError("Unable to send Initialize acknowledgement due to: ", err) + } + case kclActionTypeCheckpoint: + r.LogDebug("Receieved Checkpoint action from KCL") + r.onCheckpointCallbackFn() + err := r.sendMessage(newStatusMessage(kclActionTypeCheckpoint)) + if err != nil { + r.LogError("Unable to send Checkpoint acknowledgement due to: ", err) + } + case kcActionTypeShutdown: + r.LogDebug("Receieved Shutdown action from KCL") + r.onShutdownCallbackFn() + err := r.sendMessage(newStatusMessage(kcActionTypeShutdown)) + if err != nil { + r.LogError("Unable to send Shutdown acknowledgement due to: ", err) + } + return + case kclActionTypeProcessRecords: + r.LogDebug("Receieved ProcessRecords action from KCL") + // Put all the messages on the reader's message channel + for _, msg := range actionMessage.Records { + r.messages <-msg.ToMessage() + } + + // Send an acknowledgement that all the messages were received + err := r.sendMessage(newStatusMessage(kclActionTypeProcessRecords)) + if err != nil { + r.LogError("Unable to send ProcessRecords acknowledgement due to: ", err) + } + default: + r.LogError("processAction received an invalid action: ", actionMessage.Action) + } } - } - } + }() + }) +} - return count, size, nil +func (r *KclReader) shutdown() { + defer func(){ + // Reset startupOnce to allow the start up sequence to happen again + r.startupOnce = sync.Once{} + }() + + if r.messages != nil { + close(r.messages) + } } +// getAction reads messages from STDIN based on the Multilang Daemon protocol from KCL func (r *KclReader) getAction() (*actionMessage, error) { buffer := &bytes.Buffer{} for { line, isPrefix, err := r.reader.ReadLine() if err != nil { - panic("Unable to read line from stdin " + err.Error()) + r.LogError("Unable to read line from stdin ", err) + return nil, err } buffer.Write(line) if !isPrefix { @@ -126,59 +217,13 @@ func (r *KclReader) getAction() (*actionMessage, error) { actionMsg := &actionMessage{} err := json.Unmarshal(buffer.Bytes(), actionMsg) if err != nil { - fmt.Fprintf(os.Stderr, "Could not understand line read from input: %s\n", buffer.String()) + r.LogError("Unable to unmarshal line read from input: ", buffer.String()) + return nil, err } return actionMsg, nil } -// processAction listens to STDIN and processes action messages based on the Multilang protocol from KCL -func (r *KclReader) processAction() error { - for { - // Retrieve action message - actionMessage, err := r.getAction() - if err != nil { - return err - } - if actionMessage == nil { - break - } - - switch actionMessage.Action { - case kclActionTypeInitialize: - err := r.sendMessage(newStatusMessage(kclActionTypeInitialize)) - if err != nil { - r.LogError("Unable to send Initialize acknowledgement due to: ", err) - } - case kclActionTypeCheckpoint: - err := r.sendMessage(newStatusMessage(kclActionTypeCheckpoint)) - if err != nil { - r.LogError("Unable to send Checkpoint acknowledgement due to: ", err) - } - case kcActionTypeShutdown: - err := r.sendMessage(newStatusMessage(kcActionTypeShutdown)) - if err != nil { - r.LogError("Unable to send Shutdown acknowledgement due to: ", err) - } - case kclActionTypeProcessRecords: - // Put all the messages on the reader's message channel - for _, msg := range actionMessage.Records { - r.messages <-msg.ToMessage() - } - - // Send an acknowledgement that all the messages were received - err := r.sendMessage(newStatusMessage(kclActionTypeProcessRecords)) - if err != nil { - r.LogError("Unable to send ProcessRecords acknowledgement due to: ", err) - } - default: - r.LogError("processAction received an invalid action: ", actionMessage.Action) - } - } - - return nil -} - -// sendMessage +// sendMessage writes messages to STDOUT based on the Multilang Daemon protocol from KCL func (r *KclReader) sendMessage(msg interface{}) error { b, err := json.Marshal(msg) if err != nil { @@ -191,16 +236,39 @@ func (r *KclReader) sendMessage(msg interface{}) error { return nil } +// processRecords is a helper method which pulls from the reader's message channel and calls the callback function +func (r *KclReader) processRecords(ctx context.Context, batchSize int, fn messageHandler) (count int, size int, err error) { + r.process(ctx) + for i := 0; i < batchSize; i++ { + msg := <-r.messages + err = fn(msg) + if err != nil { + r.LogError("messageHandler resulted in an error: ", err) + } else { + r.Stats.AddConsumed(1) + count++ + b, err := json.Marshal(msg) + if err != nil { + r.LogError("Unable to marshal message: ", err) + } else { + size += len(b) + } + } + } + + return count, size, nil +} + // GetRecord calls processRecords to attempt to put one message from message buffer to the consumer's message // channel func (r *KclReader) GetRecord(ctx context.Context, fn messageHandler) (count int, size int, err error) { - count, size, err = r.processRecords(1, fn) + count, size, err = r.processRecords(ctx, 1, fn) return count, size, err } // GetRecords calls processRecords to attempt to put all messages on the message buffer on the consumer's // message channel func (r *KclReader) GetRecords(ctx context.Context, fn messageHandler) (count int, size int, err error) { - count, size, err = r.processRecords(r.batchSize, fn) + count, size, err = r.processRecords(ctx, r.batchSize, fn) return count, size, err } diff --git a/kinesis_reader.go b/kinesis_reader.go index 48d9d01..8591114 100644 --- a/kinesis_reader.go +++ b/kinesis_reader.go @@ -250,7 +250,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn messageHandler, batch r.LogError("Error getting records:", err) switch err.(awserr.Error).Code() { case kinesis.ErrCodeProvisionedThroughputExceededException: - r.Stats.AddGetRecordsProvisionedThroughputExceeded(1) + r.Stats.AddReadProvisionedThroughputExceeded(1) default: r.LogDebug("Received AWS error:", err.Error()) } diff --git a/kinesis_writer.go b/kinesis_writer.go index b2b05d7..f25a1c6 100644 --- a/kinesis_writer.go +++ b/kinesis_writer.go @@ -185,7 +185,7 @@ func (w *KinesisWriter) PutRecords(ctx context.Context, messages []*Message, fn } else { switch aws.StringValue(record.ErrorCode) { case kinesis.ErrCodeProvisionedThroughputExceededException: - w.Stats.AddPutRecordsProvisionedThroughputExceeded(1) + w.Stats.AddWriteProvisionedThroughputExceeded(1) default: w.LogDebug("PutRecords record failed with error:", aws.StringValue(record.ErrorCode), aws.StringValue(record.ErrorMessage)) } diff --git a/producer.go b/producer.go index b444148..92161e8 100644 --- a/producer.go +++ b/producer.go @@ -182,8 +182,10 @@ func NewProducer(c *aws.Config, stream string, optionFns ...ProducerOptionsFn) ( // produce is called once to initialize a pool of workers which send batches of messages concurrently func (p *Producer) produce() { p.startupOnce.Do(func() { - // Reset shutdownOnce to allow the shut down sequence to happen again - p.shutdownOnce = sync.Once{} + defer func() { + // Reset shutdownOnce to allow the shut down sequence to happen again + p.shutdownOnce = sync.Once{} + }() // Instantiate rate limiters p.msgCountLimiter = rate.NewLimiter(rate.Limit(float64(p.writer.getMsgCountRateLimit())), p.batchSize) @@ -312,6 +314,11 @@ func (p *Producer) produce() { // shutdown is called once to handle the graceful shutdown of the produce function func (p *Producer) shutdown() { p.shutdownOnce.Do(func() { + defer func() { + // Reset startupOnce to allow the start up sequence to happen again + p.startupOnce = sync.Once{} + }() + // Close the messages channel to prevent any more incoming messages if p.messages != nil { close(p.messages) @@ -365,9 +372,6 @@ func (p *Producer) shutdown() { if p.pipeOfDeath != nil { close(p.pipeOfDeath) } - - // Reset startupOnce to allow the start up sequence to happen again - p.startupOnce = sync.Once{} }) } diff --git a/producer_stats.go b/producer_stats.go index 11f5532..02bac07 100644 --- a/producer_stats.go +++ b/producer_stats.go @@ -20,7 +20,7 @@ type ProducerStatsCollector interface { AddDroppedRetries(int) AddPutRecordsCalled(int) AddPutRecordsTimeout(int) - AddPutRecordsProvisionedThroughputExceeded(int) + AddWriteProvisionedThroughputExceeded(int) UpdatePutRecordsDuration(time.Duration) UpdatePutRecordsBuildDuration(time.Duration) UpdatePutRecordsSendDuration(time.Duration) @@ -61,9 +61,9 @@ func (nsc *NilProducerStatsCollector) AddPutRecordsCalled(int) {} // by the WithHTTPClientTimeout configuration. func (nsc *NilProducerStatsCollector) AddPutRecordsTimeout(int) {} -// AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords API returned a +// AddWriteProvisionedThroughputExceeded records the number of times the PutRecords API returned a // ErrCodeProvisionedThroughputExceededException by the producer. -func (nsc *NilProducerStatsCollector) AddPutRecordsProvisionedThroughputExceeded(int) {} +func (nsc *NilProducerStatsCollector) AddWriteProvisionedThroughputExceeded(int) {} // UpdatePutRecordsDuration records the duration that the PutRecords API request took. Only the times of successful calls // are measured. @@ -184,9 +184,9 @@ func (dsc *DefaultProducerStatsCollector) AddPutRecordsTimeout(count int) { dsc.PutRecordsTimeout.Inc(int64(count)) } -// AddPutRecordsProvisionedThroughputExceeded records the number of times the PutRecords API returned a +// AddWriteProvisionedThroughputExceeded records the number of times the PutRecords API returned a // ErrCodeProvisionedThroughputExceededException by the producer. -func (dsc *DefaultProducerStatsCollector) AddPutRecordsProvisionedThroughputExceeded(count int) { +func (dsc *DefaultProducerStatsCollector) AddWriteProvisionedThroughputExceeded(count int) { dsc.PutRecordsProvisionedThroughputExceeded.Inc(int64(count)) } diff --git a/producer_stats_test.go b/producer_stats_test.go index e3dea49..4acb8b5 100644 --- a/producer_stats_test.go +++ b/producer_stats_test.go @@ -46,8 +46,8 @@ func TestProducerStatsCollector(t *testing.T) { sc.AddDroppedRetries(1) }) - Convey("check that AddPutRecordsProvisionedThroughputExceeded does not error", func() { - sc.AddPutRecordsProvisionedThroughputExceeded(1) + Convey("check that WriteProvisionedThroughputExceeded does not error", func() { + sc.AddWriteProvisionedThroughputExceeded(1) }) Convey("check that AddPutRecordsCalled does not eroror", func() { @@ -55,7 +55,7 @@ func TestProducerStatsCollector(t *testing.T) { }) Convey("check that AddProvisionedThroughputExceeded does not erro", func() { - sc.AddPutRecordsProvisionedThroughputExceeded(1) + sc.AddWriteProvisionedThroughputExceeded(1) }) Convey("check that AddPutRecordsTimeout does not error", func() { @@ -126,9 +126,9 @@ func TestProducerStatsCollector(t *testing.T) { So(sc.(*DefaultProducerStatsCollector).DroppedRetries.Count(), ShouldEqual, int64(count)) }) - Convey("check that AddPutRecordsProvisionedThroughputExceeded does not error", func() { + Convey("check that WriteProvisionedThroughputExceeded does not error", func() { count := rand.Int() - sc.AddPutRecordsProvisionedThroughputExceeded(count) + sc.AddWriteProvisionedThroughputExceeded(count) So(sc.(*DefaultProducerStatsCollector).PutRecordsProvisionedThroughputExceeded.Count(), ShouldEqual, int64(count)) }) @@ -138,9 +138,9 @@ func TestProducerStatsCollector(t *testing.T) { So(sc.(*DefaultProducerStatsCollector).PutRecordsCalled.Count(), ShouldEqual, int64(count)) }) - Convey("check that AddPutRecordsProvisionedThroughputExceeded does not erro", func() { + Convey("check that WriteProvisionedThroughputExceeded does not erro", func() { count := rand.Int() - sc.AddPutRecordsProvisionedThroughputExceeded(count) + sc.AddWriteProvisionedThroughputExceeded(count) So(sc.(*DefaultProducerStatsCollector).PutRecordsProvisionedThroughputExceeded.Count(), ShouldEqual, int64(count)) }) diff --git a/testexec/main.go b/testexec/main.go index 3d6b485..adc6743 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -395,9 +395,11 @@ func produce(sd *StreamData, p *kinetic.Producer, wg *sync.WaitGroup) { go func() { for { <-sendSignal + key := make([]byte, 16) + rand.Read(key) jsonStr, _ := json.Marshal(NewMessage()) if err := p.TryToSend(&kinetic.Message{ - PartitionKey: aws.String("key"), + PartitionKey: aws.String(string(key)), Data: []byte(jsonStr), }); err == nil { sd.incrementMsgCount() @@ -416,7 +418,7 @@ func produce(sd *StreamData, p *kinetic.Producer, wg *sync.WaitGroup) { if *cfg.Blast { sendTicker = time.NewTicker(time.Nanosecond) } else { - sendTicker = time.NewTicker(time.Duration(rand.Intn(1)+1) * time.Millisecond) + sendTicker = time.NewTicker(time.Millisecond) } produce: for { diff --git a/testexec/messages.go b/testexec/messages.go index b963dad..b77f279 100644 --- a/testexec/messages.go +++ b/testexec/messages.go @@ -1,7 +1,7 @@ package main import ( - "strconv" + "math/rand" "sync/atomic" ) @@ -15,10 +15,16 @@ type Message struct { // NewMessage creates a new Message struct with a unique identifier func NewMessage() *Message { + // Increment ID atomic.AddUint64(&msgCount, 1) id := atomic.LoadUint64(&msgCount) + + // Create random string as message + message := make([]byte, 100) + rand.Read(message) + return &Message{ ID: id, - Message: "hello_" + strconv.Itoa(int(id)), + Message: string(message), } } From 8c11d7246efd2dc3e629f6b0052e21bafeda1073 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 10 Aug 2017 22:47:32 -0500 Subject: [PATCH 080/109] Added checkpointing in KclReader and exposed new functions through the StreamReader interface. Refactored KclReader process method to allow checkpoint calls to KCL in between long-poll processRecord handling. --- checkpoint.go | 575 +++++++++++++++++++++++++ checkpoint_test.go | 1024 ++++++++++++++++++++++++++++++++++++++++++++ consumer.go | 10 +- consumer_stats.go | 68 +-- consumer_test.go | 1 - firehose_writer.go | 4 +- kcl_reader.go | 160 +++++-- kinesis_reader.go | 18 + kinesis_writer.go | 4 +- kinetic.go | 5 +- producer.go | 50 +-- shared_types.go | 9 +- 12 files changed, 1816 insertions(+), 112 deletions(-) create mode 100644 checkpoint.go create mode 100644 checkpoint_test.go diff --git a/checkpoint.go b/checkpoint.go new file mode 100644 index 0000000..4af57eb --- /dev/null +++ b/checkpoint.go @@ -0,0 +1,575 @@ +package kinetic + +import ( + "context" + "errors" + "math" + "strconv" + "sync" + "sync/atomic" + "time" +) + +// btreeNode is a binary tree node used to encapsulate checkpointing information in a binary search tree. +type btreeNode struct { + seqNum int // key to store the sequence number in the checkpointing system + done bool // flag to indicate whether or not processing of a record is complete + age time.Time // time stamp used to expire btreeNodes from the checkpointing system + left *btreeNode // pointer to a btreeNode whose sequence number is less than own + right *btreeNode // pointer to a btreeNode whose sequence number is greater than own +} + +// insert is a (recursive) method used to insert sequence numbers into the binary search tree. +func (n *btreeNode) insert(seqNum int) error { + if n == nil { + return errors.New("Cannot call insert on a nil node") + } + + switch { + case seqNum == n.seqNum: + return nil + case seqNum < n.seqNum: + if n.left == nil { + n.left = &btreeNode{ + seqNum: seqNum, + age: time.Now(), + } + return nil + } + return n.left.insert(seqNum) + case seqNum > n.seqNum: + if n.right == nil { + n.right = &btreeNode{ + seqNum: seqNum, + age: time.Now(), + } + return nil + } + return n.right.insert(seqNum) + } + return nil +} + +// find is a (recursive) method used to retrieve a pointer to the btreeNode containing the sequence number or 'false' +// if the sequence number provided is not found. +func (n *btreeNode) find(seqNum int) (*btreeNode, bool) { + if n == nil { + return nil, false + } + + switch { + case seqNum == n.seqNum: + return n, true + case seqNum < n.seqNum: + return n.left.find(seqNum) + default: + return n.right.find(seqNum) + } +} + +// findMin is a helper function used to find the minimum element in the subtree of the given parent node. +func (n *btreeNode) findMin(parent *btreeNode) (*btreeNode, *btreeNode) { + if n.left == nil { + return n, parent + } + return n.left.findMin(n) +} + +// findMax is a helper function used to find the maximum element in the subtree of the given parent node. +func (n *btreeNode) findMax(parent *btreeNode) (*btreeNode, *btreeNode) { + if n.right == nil { + return n, parent + } + return n.right.findMax(n) +} + +// replaceNode is a helper method used to replace btreeNodes with a binary search tree. +func (n *btreeNode) replaceNode(parent, replacement *btreeNode) error { + if n == nil { + return errors.New("Cannot call replaceNode on a nil node") + } + + if n == parent.left { + parent.left = replacement + return nil + } + parent.right = replacement + return nil +} + +// delete is a (recursive) method used remove (and re-organize if necessary) a btreeNode from the binary search tree. +func (n *btreeNode) delete(seqNum int, parent *btreeNode) error { + if n == nil { + return errors.New("Cannot call delete on a nil node") + } + + switch { + case seqNum < n.seqNum: + return n.left.delete(seqNum, n) + case seqNum > n.seqNum: + return n.right.delete(seqNum, n) + default: + if n.left == nil && n.right == nil { + n.replaceNode(parent, nil) + return nil + } + if n.left == nil { + n.replaceNode(parent, n.right) + return nil + } + if n.right == nil { + n.replaceNode(parent, n.left) + return nil + } + + replacement, replacementParent := n.left.findMax(n) + n.seqNum = replacement.seqNum + n.done = replacement.done + n.age = replacement.age + + return replacement.delete(replacement.seqNum, replacementParent) + } +} + +// markDone is a method used to change the 'done' flag of a btreeNode to 'true'. +func (n *btreeNode) markDone() { + n.done = true +} + +// btree is a binary search tree. +type btree struct { + root *btreeNode // the root node for the binary search tree +} + +// insert is a method used to insert sequence numbers into the binary search tree of the checkpointing system. +func (t *btree) insert(seqNum int) error { + if t.root == nil { + t.root = &btreeNode{ + seqNum: seqNum, + age: time.Now(), + } + return nil + } + return t.root.insert(seqNum) +} + +// find is a method used retrieve a pointer to a btreeNode of the given sequence number or 'false' if the sequence +// number was not found. +func (t *btree) find(seqNum int) (*btreeNode, bool) { + if t.root == nil { + return nil, false + } + return t.root.find(seqNum) +} + +// delete is a method used to remove a btreeNode with the given sequence number from the binary search tree. +func (t *btree) delete(seqNum int) error { + if t.root == nil { + return errors.New("Cannot call delete on a nil tree") + } + + if t.root.seqNum == seqNum { + fakeParent := &btreeNode{right: t.root} + if err := t.root.delete(seqNum, fakeParent); err != nil { + return err + } + t.root = fakeParent.right + return nil + } + return t.root.delete(seqNum, nil) +} + +// traverseOrdered is a (recursive) method used to apply a callback function on all of the nodes in the binary +// search tree ordered by sequence number from smallest to largest. +func (t *btree) traverseOrdered(n *btreeNode, fn func(*btreeNode)) { + if n == nil { + return + } + + t.traverseOrdered(n.left, fn) + fn(n) + t.traverseOrdered(n.right, fn) +} + +// traverseChildren is a (recursive) method used to apply a callback function on all of children nodes in the +// binary search tree. +func (t *btree) traverseChildren(n *btreeNode, fn func(*btreeNode)) { + if n == nil { + return + } + + t.traverseChildren(n.left, fn) + t.traverseChildren(n.right, fn) + fn(n) +} + +// size is a method used to determine the number of btreeNodes in the binary search tree. +func (t *btree) size() int { + var size int + if t.root != nil { + t.traverseOrdered(t.root, func(n *btreeNode) { + size++ + }) + } + return size +} + +// minNode is a method used to retrieve a pointer to the btreeNode containing the smallest sequence number in the +// binary search tree. +func (t *btree) minNode() *btreeNode { + if t.root == nil { + return nil + } + node, _ := t.root.findMin(t.root) + return node +} + +// maxNode is a method used to retrieve a pointer to the btreeNode containing the largest sequence number in the +// binary search tree. +func (t *btree) maxNode() *btreeNode { + if t.root == nil { + return nil + } + node, _ := t.root.findMax(t.root) + return node +} + +// trim is a method used to trim the binary search tree so that the sequence number of all of the remaining nodes +// fall within (non-inclusive) the minSeqNum and maxSeqNum values. +func (t *btree) trim(minSeqNum, maxSeqNum *int) error { + var min, max int + if minSeqNum == nil { + min = math.MinInt64 + } else { + min = *minSeqNum + } + if maxSeqNum == nil { + max = math.MaxInt64 + } else { + max = *maxSeqNum + } + + fakeParent := &btreeNode{right: t.root} + t.traverseChildren(fakeParent, func(n *btreeNode) { + if n.left != nil && (n.left.seqNum <= min || n.left.seqNum >= max) { + n.delete(n.left.seqNum, n) + } + if n.right != nil && (n.right.seqNum <= min || n.right.seqNum >= max) { + n.delete(n.right.seqNum, n) + } + }) + t.root = fakeParent.right + + return nil +} + +// expire is a method used to delete btreeNodes from the binary search tree whose age is older than the given age. +func (t *btree) expire(age time.Duration) error { + fakeParent := &btreeNode{right: t.root} + t.traverseChildren(fakeParent, func(n *btreeNode) { + if n.left != nil && time.Since(n.left.age) > age { + n.delete(n.left.seqNum, n) + } + if n.right != nil && time.Since(n.right.age) > age { + n.delete(n.right.seqNum, n) + } + }) + t.root = fakeParent.right + + return nil +} + +// checkpointOptions is a struct containing all of the configurable options for a checkpoint object. +type checkpointOptions struct { + autoCheckpointCount int // count of newly inserted messages before triggering an automatic checkpoint call + autoCheckpointFreq time.Duration // frequency with which to automatically call checkpoint + maxAge time.Duration // maximum duration for an inserted sequence number to live before being expired + maxSize int // maximum capacity (number of btreeNodes) of the checkpoint system + checkpointFn func(checkpoint string) error // callback function to call on a sequence number when a checkpoint is discovered with the checkpoint call + expireFn func(checkpoint string) error // callback function to call when a sequence number is aged out + capacityFn func(checkpoint string) error // callback function to call when the checkpoint system has reached max capacity + countCheckFreq time.Duration // frequency with which to check the insert count +} + +// defaultCheckpointOptions is a function that returns a pointer to a checkpointOptions object with default values. +func defaultCheckpointOptions() *checkpointOptions { + return &checkpointOptions{ + autoCheckpointCount: 10000, + autoCheckpointFreq: time.Minute, + maxAge: time.Hour, + maxSize: 1000000, + checkpointFn: func(string) error { return nil }, + expireFn: func(string) error { return nil }, + capacityFn: func(string) error { return nil }, + countCheckFreq: time.Second, + } +} + +// checkpointOptionsFn is a function signature used to define function options for configuring all of the +// configurable options of a checkpoint object. +type checkpointOptionsFn func(*checkpointOptions) error + +// checkpointAutoCheckpointCount is a functional option method for configuring the checkpoint's auto checkpoint count. +func checkpointAutoCheckpointCount(count int) checkpointOptionsFn { + return func(o *checkpointOptions) error { + o.autoCheckpointCount = count + return nil + } +} + +// checkpointAutoCheckpointFreq is a functional option method for configuring the checkpoint's auto checkpoint +// frequency. +func checkpointAutoCheckpointFreq(freq time.Duration) checkpointOptionsFn { + return func(o *checkpointOptions) error { + o.autoCheckpointFreq = freq + return nil + } +} + +// checkpointMaxAge is a functional option method for configuring the checkpoint's maximum age. +func checkpointMaxAge(age time.Duration) checkpointOptionsFn { + return func(o *checkpointOptions) error { + o.maxAge = age + return nil + } +} + +// checkpointMaxSize is a functional option method for configuring the checkpoint's maximum size. +func checkpointMaxSize(size int) checkpointOptionsFn { + return func(o *checkpointOptions) error { + o.maxSize = size + return nil + } +} + +// checkpointCheckpointFn is a functional option method for configuring the checkpoint's checkpoint callback function. +func checkpointCheckpointFn(fn func(string) error) checkpointOptionsFn { + return func(o *checkpointOptions) error { + o.checkpointFn = fn + return nil + } +} + +// checkpointExpireFn is a functional option method for configuring the checkpoint's expire callback function. +func checkpointExpireFn(fn func(string) error) checkpointOptionsFn { + return func(o *checkpointOptions) error { + o.expireFn = fn + return nil + } +} + +// checkpointCapacityFn is a functional option method for configuring the checkpoint's capacity callback function. +func checkpointCapacityFn(fn func(string) error) checkpointOptionsFn { + return func(o *checkpointOptions) error { + o.capacityFn = fn + return nil + } +} + +// checkpointCountCheckFreq is a functional option method for configuring the checkpoint's count check frequency. +func checkpointCountCheckFreq(freq time.Duration) checkpointOptionsFn { + return func(o *checkpointOptions) error { + o.countCheckFreq = freq + return nil + } +} + +// checkpoint is a data structure that is used as a bookkeeping system to track the state of data processing +// for records pulled off of the consumer's message channel. The check pointing system uses a binary search +// tree keyed off of the Kinesis sequence number of the record. The sequence numbers should be inserted into the +// binary search tree using the insert() function after messages are pulled off of the consumer's message +// channel and should be marked done using the markDone() function after data processing is completed. The +// checkpoint() can be called periodically which may trigger a checkpoint call to KCL if the oldest sequence +// numbers have been marked complete. Call startup() to enable automatic checkpointing and expiration. +type checkpoint struct { + *checkpointOptions // contains all of the configuration settings for the checkpoint object + keys *btree // binary search tree used to store Kinesis record sequence numbers + keysMu sync.Mutex // mutex to make keys thread safe + counter uint64 // counter to track the number of messages inserted since the last checkpoint + checkpointCh chan struct{} // channel with which to communicate / coordinate checkpointing + startupOnce sync.Once // used to ensure that the startup function is called once + shutdownOnce sync.Once // used to ensure that the shutdown function is called once +} + +// newCheckpoint instantiates a new checkpoint object with default configuration settings unless the function option +// methods are provided to change the default values. +func newCheckpoint(optionFns ...checkpointOptionsFn) *checkpoint { + checkpointOptions := defaultCheckpointOptions() + for _, optionFn := range optionFns { + optionFn(checkpointOptions) + } + return &checkpoint{ + checkpointOptions: checkpointOptions, + keys: &btree{}, + } +} + +// startup is a method used to enable automatic checkpointing and expiration of btreeNodes from the +// checkpointing system. +func (c *checkpoint) startup(ctx context.Context) { + c.startupOnce.Do(func() { + defer func() { + c.shutdownOnce = sync.Once{} + }() + + c.checkpointCh = make(chan struct{}) + go func() { + defer c.shutdown() + + autoCheckpointTimer := time.NewTimer(c.autoCheckpointFreq) + expirationTicker := time.NewTicker(c.maxAge) + counterCheckTicker := time.NewTicker(c.countCheckFreq) + + for { + wait: + for atomic.LoadUint64(&c.counter) < uint64(c.autoCheckpointCount) { + select { + case <-ctx.Done(): + autoCheckpointTimer.Stop() + expirationTicker.Stop() + counterCheckTicker.Stop() + return + case <-c.checkpointCh: + break wait + case <-autoCheckpointTimer.C: + break wait + case <-expirationTicker.C: + c.expire(c.maxAge) + break + case <-counterCheckTicker.C: + break + } + } + + // call check to obtain the checkpoint + cp, found := c.check() + if found { + c.checkpointFn(cp) + c.trim(cp) + } + + // reset counter and timer + atomic.StoreUint64(&c.counter, 0) + autoCheckpointTimer.Reset(c.autoCheckpointFreq) + } + }() + }) +} + +// shutdown is a method used to clean up the checkpoint object +func (c *checkpoint) shutdown() { + c.shutdownOnce.Do(func() { + defer func() { + c.startupOnce = sync.Once{} + }() + + if c.checkpointCh != nil { + close(c.checkpointCh) + } + }) +} + +// insert safely inserts a sequence number into the binary search tree. +func (c *checkpoint) insert(seqNumStr string) error { + c.keysMu.Lock() + defer c.keysMu.Unlock() + + if c.keys.size() >= c.maxSize { + c.capacityFn(seqNumStr) + return errors.New("Unable to insert due to capacity") + } + + if seqNum, err := strconv.Atoi(seqNumStr); err == nil { + err := c.keys.insert(seqNum) + if err != nil { + return err + } + } + + atomic.AddUint64(&c.counter, 1) + return nil +} + +// markDone safely marks the given sequence number as done. +func (c *checkpoint) markDone(seqNumStr string) error { + c.keysMu.Lock() + defer c.keysMu.Unlock() + + seqNum, err := strconv.Atoi(seqNumStr) + if err != nil { + return err + } + + node, ok := c.keys.find(seqNum) + if !ok { + return errors.New("Sequence number not found") + } + + node.markDone() + return nil +} + +// check returns the largest sequence number marked as done where all smaller sequence numbers have +// also been marked as done. +func (c *checkpoint) check() (checkpoint string, found bool) { + c.keysMu.Lock() + defer c.keysMu.Unlock() + + if c.keys == nil { + return "", found + } + + var stop bool // flag to stop applying function in traverse + c.keys.traverseOrdered(c.keys.root, func(n *btreeNode) { + if stop { + return + } + + if n.left == nil && !n.done { + stop = true + return + } else if !n.done { + stop = true + return + } + checkpoint = strconv.Itoa(n.seqNum) + found = true + }) + return checkpoint, found +} + +// trim safely simplifies the binary tree housing the sequence numbers so that only the smallest element is +// marked as done. +func (c *checkpoint) trim(checkpointStr string) error { + c.keysMu.Lock() + defer c.keysMu.Unlock() + + checkpoint, err := strconv.Atoi(checkpointStr) + if err != nil { + return err + } + + if _, ok := c.keys.find(checkpoint); !ok { + return errors.New("Sequence number not found") + } + + cp := int(checkpoint) + return c.keys.trim(&cp, nil) +} + +// expire safely removes sequence numberse from the checkpointing system that are older than the given age. +func (c *checkpoint) expire(age time.Duration) error { + c.keysMu.Lock() + defer c.keysMu.Unlock() + + return c.keys.expire(age) +} + +// checkpoint sends a signal to the checkpointCh channel to enable a call to checkpoint. +func (c *checkpoint) checkpoint() error { + if c.checkpointCh == nil { + return errors.New("Nil checkpoint channel") + } + c.checkpointCh <- struct{}{} + return nil +} diff --git a/checkpoint_test.go b/checkpoint_test.go new file mode 100644 index 0000000..e3424da --- /dev/null +++ b/checkpoint_test.go @@ -0,0 +1,1024 @@ +package kinetic + +import ( + . "github.com/smartystreets/goconvey/convey" + + "context" + "fmt" + "math" + "math/rand" + "sort" + "strconv" + "sync/atomic" + "testing" + "time" +) + +// initialize the randomization seed for the random number generator. +func init() { + rand.Seed(time.Now().UnixNano()) +} + +// sortByKeys is an alias to a slice array of int that implements the sort interface for sorting. +type sortByKeys []int + +// Len implements the Len function for hte sort interface. +func (k sortByKeys) Len() int { + return len(k) +} + +// Less implements the Less function for the sort interface. +func (k sortByKeys) Less(i, j int) bool { + return k[i] < k[j] +} + +// Swap implements the Swap function for the sort interface. +func (k sortByKeys) Swap(i, j int) { + k[i], k[j] = k[j], k[i] +} + +// insertKeys is a helper function for inserting multiple keys into the binary search tree. +func insertKeys(t *btree, keys []int) bool { + for _, key := range keys { + err := t.insert(key) + if err != nil { + return false + } + } + return true +} + +// findKeys is a helper function for confirm that all of the given keys were found. +func findKeys(t *btree, keys []int) bool { + for _, key := range keys { + node, ok := t.find(key) + if !ok { + return false + } + if node == nil || key != node.seqNum { + return false + } + } + return true +} + +// sortUnique is a helper method used to return a unique and sorted slice of keys. +func sortUnique(keys []int) sortByKeys { + keyMap := make(map[int]bool) + for _, key := range keys { + if _, ok := keyMap[key]; !ok { + keyMap[key] = true + } + } + + var sortedKeys sortByKeys + for key := range keyMap { + sortedKeys = append(sortedKeys, key) + } + sort.Sort(sortedKeys) + + return sortedKeys +} + +// randomKeys is helper method used to generate a random set of keys for testing. +func randomKeys(size, max int) (keys []int) { + for i := 0; i < size; i++ { + keys = append(keys, int(rand.Intn(max))) + } + return keys +} + +// intPointer is a helper method used to return the address of an int +func intPointer(x int) *int { + return &x +} + +func TestCheckpointBTreeInsert(t *testing.T) { + var tests = []struct { + scenario string + keys []int + }{ + { + scenario: "empty tree", + keys: []int{}, + }, + { + scenario: "single key", + keys: []int{1}, + }, + { + scenario: "duplicate key", + keys: []int{1, 1}, + }, + { + scenario: "keys in ascending order", + keys: []int{1, 2, 3}, + }, + { + scenario: "keys in descendeing order", + keys: []int{4, 3, 2, 1}, + }, + { + scenario: "keys in random order", + keys: randomKeys(100, 1000), + }, + } + + for _, test := range tests { + var b btree + sortedKeys := sortUnique(test.keys) + Convey(fmt.Sprintf("running binary tree insert test suit on scenario: [%s]", test.scenario), t, func() { + Convey("should be able to insert keys without errors", func() { + ok := insertKeys(&b, test.keys) + So(ok, ShouldBeTrue) + + if len(sortedKeys) > 0 { + Convey("should be able to find all keys without errors", func() { + ok := findKeys(&b, sortedKeys) + So(ok, ShouldBeTrue) + }) + } + + Convey("should be able to determine size correctly", func() { + So(b.size(), ShouldEqual, len(sortedKeys)) + }) + + Convey("should be able to confirm sorting order", func() { + var keys sortByKeys + b.traverseOrdered(b.root, func(n *btreeNode) { + keys = append(keys, n.seqNum) + }) + So(b.size(), ShouldEqual, len(keys)) + So(sort.IsSorted(keys), ShouldBeTrue) + }) + + Convey("should be able to find the min node correctly", func() { + minNode := b.minNode() + if minNode == nil { + So(b.root, ShouldBeNil) + } else { + So(b.minNode().seqNum, ShouldEqual, sortedKeys[0]) + } + }) + + Convey("should be able to find the max node correctly", func() { + maxNode := b.maxNode() + if maxNode == nil { + So(b.root, ShouldBeNil) + } else { + So(b.maxNode().seqNum, ShouldEqual, sortedKeys[len(sortedKeys)-1]) + } + }) + }) + }) + } +} + +func TestCheckpointBTreeFind(t *testing.T) { + var tests = []struct { + scenario string + keys []int + findKey int + shouldFind bool + }{ + { + scenario: "key is on a root node", + keys: []int{2, 3, 1}, + findKey: 2, + shouldFind: true, + }, + { + scenario: "key is on a right child node", + keys: []int{2, 3, 1}, + findKey: 3, + shouldFind: true, + }, + { + scenario: "key is on a left child node", + keys: []int{2, 3, 1}, + findKey: 1, + shouldFind: true, + }, + { + scenario: "key does not exist", + keys: []int{2, 3, 1}, + findKey: 0, + shouldFind: false, + }, + { + scenario: "empty tree / key does not exist", + keys: []int{}, + findKey: 0, + shouldFind: false, + }, + } + + var b btree + for _, test := range tests { + Convey(fmt.Sprintf("running binary tree find test suit on scenario: [%s]", test.scenario), t, func() { + ok := insertKeys(&b, test.keys) + So(ok, ShouldBeTrue) + + Convey("calling find should behave accordingly", func() { + node, ok := b.find(test.findKey) + So(ok, ShouldEqual, test.shouldFind) + if test.shouldFind { + Convey("the node return should match the key", func() { + So(node.seqNum, ShouldEqual, test.findKey) + }) + } + }) + }) + } +} + +func TestCheckpointBTreeDelete(t *testing.T) { + var tests = []struct { + scenario string + keys []int + deleteKey int + shouldError bool + }{ + { + scenario: "delete on empty tree", + keys: []int{}, + deleteKey: 1, + shouldError: true, + }, + { + scenario: "delete a key that doesn't exist", + keys: []int{2, 3, 1}, + deleteKey: 0, + shouldError: true, + }, + { + scenario: "delete a node with no children", + keys: []int{5, 2, 18, -4, 3}, + deleteKey: -4, + shouldError: false, + }, + { + scenario: "delete a node with one child on the left", + keys: []int{5, 2, 18, -4, 3, 16, 15, 17}, + deleteKey: 18, + shouldError: false, + }, + { + scenario: "delete a node with one child on the right", + keys: []int{5, 2, 18, -4, 3, 21, 19, 25}, + deleteKey: 18, + shouldError: false, + }, + { + scenario: "delete a node with two children", + keys: []int{5, 2, 19, 16, 25}, + deleteKey: 19, + shouldError: false, + }, + { + scenario: "delete a root node with no children", + keys: []int{1}, + deleteKey: 1, + shouldError: false, + }, + { + scenario: "delete a root node with one child on the left", + keys: []int{2, 1}, + deleteKey: 2, + shouldError: false, + }, + { + scenario: "delete a root node with one child on the right", + keys: []int{2, 3}, + deleteKey: 2, + shouldError: false, + }, + { + scenario: "delete a root node with two children", + keys: []int{2, 1, 3}, + deleteKey: 3, + shouldError: false, + }, + } + + for _, test := range tests { + var b btree + sortedKeys := sortUnique(test.keys) + Convey(fmt.Sprintf("running binary tree delete test suit on scenario: [%s]", test.scenario), t, func() { + ok := insertKeys(&b, test.keys) + So(ok, ShouldBeTrue) + + Convey("should be able to call delete and respond correctly", func() { + err := b.delete(test.deleteKey) + if test.shouldError { + So(err, ShouldNotBeNil) + } else { + So(err, ShouldBeNil) + } + + if test.shouldError { + if len(sortedKeys) > 0 { + Convey("should be able to find all keys without errors", func() { + ok := findKeys(&b, sortedKeys) + So(ok, ShouldBeTrue) + }) + } + + Convey("should not be able to find delete key", func() { + _, ok := b.find(test.deleteKey) + So(ok, ShouldBeFalse) + }) + + Convey("should be able to determine size correctly", func() { + So(b.size(), ShouldEqual, len(sortedKeys)) + }) + + Convey("should be able to confirm sorting order", func() { + var keys sortByKeys + b.traverseOrdered(b.root, func(n *btreeNode) { + keys = append(keys, n.seqNum) + }) + So(b.size(), ShouldEqual, len(keys)) + So(sort.IsSorted(keys), ShouldBeTrue) + }) + + Convey("should be able to find the min node correctly", func() { + minNode := b.minNode() + if minNode == nil { + So(b.root, ShouldBeNil) + } else { + So(b.minNode().seqNum, ShouldEqual, sortedKeys[0]) + } + }) + + Convey("should be able to find the max node correctly", func() { + maxNode := b.maxNode() + if maxNode == nil { + So(b.root, ShouldBeNil) + } else { + So(b.maxNode().seqNum, ShouldEqual, sortedKeys[len(sortedKeys)-1]) + } + }) + + } else { + // determine the new key set (sorted and unique) after removing the first element + var newKeys sortByKeys + for _, k := range test.keys { + if k != test.deleteKey { + newKeys = append(newKeys, k) + } + } + sort.Sort(newKeys) + + if len(newKeys) > 0 { + Convey("should be able to find remaining keys after a delete", func() { + ok := findKeys(&b, newKeys) + So(ok, ShouldBeTrue) + }) + } + + Convey("should be able to determine size correctly after a delete", func() { + So(b.size(), ShouldEqual, len(newKeys)) + }) + + Convey("should be able to confirm sorting order after a delete", func() { + var keys sortByKeys + b.traverseOrdered(b.root, func(n *btreeNode) { + keys = append(keys, n.seqNum) + }) + So(b.size(), ShouldEqual, len(keys)) + So(sort.IsSorted(keys), ShouldBeTrue) + }) + + Convey("should be able to find the min node correctly after a delete", func() { + minNode := b.minNode() + if minNode == nil { + So(b.root, ShouldBeNil) + } else { + So(b.minNode().seqNum, ShouldEqual, newKeys[0]) + } + }) + + Convey("should be able to find the max node correctly after a delete", func() { + maxNode := b.maxNode() + if maxNode == nil { + So(b.root, ShouldBeNil) + } else { + So(b.maxNode().seqNum, ShouldEqual, newKeys[len(newKeys)-1]) + } + }) + } + }) + }) + } +} + +func TestCheckpointBTreeTrim(t *testing.T) { + testKeys := []int{50, 20, 60, 40, 0, 10, 30} + var tests = []struct { + scenario string + keys []int + trimMin *int + trimMax *int + }{ + { + scenario: "trim on empty tree", + keys: []int{}, + trimMin: intPointer(0), + trimMax: intPointer(10), + }, + { + scenario: "trim left - none", + keys: testKeys, + trimMin: intPointer(-1), + trimMax: nil, + }, + { + scenario: "trim left - some", + keys: testKeys, + trimMin: intPointer(20), + trimMax: nil, + }, + { + scenario: "trim left - all", + keys: testKeys, + trimMin: intPointer(60), + trimMax: nil, + }, + { + scenario: "trim right - none", + keys: testKeys, + trimMin: nil, + trimMax: intPointer(70), + }, + { + scenario: "trim right - some", + keys: testKeys, + trimMin: nil, + trimMax: intPointer(40), + }, + { + scenario: "trim right - all", + keys: testKeys, + trimMin: nil, + trimMax: intPointer(0), + }, + { + scenario: "trim - none", + keys: testKeys, + trimMin: intPointer(0), + trimMax: intPointer(60), + }, + { + scenario: "trim - some", + keys: testKeys, + trimMin: intPointer(10), + trimMax: intPointer(50), + }, + { + scenario: "trim - all", + keys: testKeys, + trimMin: intPointer(1), + trimMax: intPointer(9), + }, + } + + for _, test := range tests { + var b btree + sortedKeys := sortUnique(test.keys) + Convey(fmt.Sprintf("running binary tree trim test suite on [%s]", test.scenario), t, func() { + Convey("should be able to insert keys without errors", func() { + ok := insertKeys(&b, test.keys) + So(ok, ShouldBeTrue) + + if len(sortedKeys) > 0 { + Convey("should be able to find all keys without errors", func() { + ok := findKeys(&b, sortedKeys) + So(ok, ShouldBeTrue) + }) + } + + uniqueKeys := sortUnique(test.keys) + var remainingKeys sortByKeys + var trimmedKeys sortByKeys + var min, max int + if test.trimMin == nil { + min = math.MinInt64 + } else { + min = *test.trimMin + } + if test.trimMax == nil { + max = math.MaxInt64 + } else { + max = *test.trimMax + } + + for _, key := range uniqueKeys { + if key > min && key < max { + remainingKeys = append(remainingKeys, key) + } else { + trimmedKeys = append(trimmedKeys, key) + } + } + + Convey("calling trim on tree", func() { + err := b.trim(test.trimMin, test.trimMax) + So(err, ShouldBeNil) + + if len(remainingKeys) > 0 { + Convey("should be able to find all remaining keys", func() { + So(findKeys(&b, remainingKeys), ShouldBeTrue) + }) + } + + if len(trimmedKeys) > 0 { + Convey("should not be able to find any trimmed keys", func() { + for _, key := range trimmedKeys { + _, ok := b.find(key) + So(ok, ShouldBeFalse) + } + }) + } + }) + }) + }) + } +} + +func TestCheckpointBTreeExpire(t *testing.T) { + testKeys1 := []int{20, 30, 10} + testKeys2 := []int{15, 25, -5} + var tests = []struct { + scenario string + keys1 []int + keys2 []int + pause1 time.Duration + pause2 time.Duration + expiration time.Duration + expectedSize int + }{ + { + scenario: "expire on an empty tree", + keys1: []int{}, + keys2: []int{}, + pause1: 0 * time.Second, + pause2: 0 * time.Second, + expiration: time.Second, + expectedSize: 0, + }, + { + scenario: "expire none", + keys1: testKeys1, + keys2: testKeys2, + pause1: 0 * time.Second, + pause2: 0 * time.Second, + expiration: time.Second, + expectedSize: len(testKeys1) + len(testKeys2), + }, + { + scenario: "expire some", + keys1: testKeys1, + keys2: testKeys2, + pause1: time.Second, + pause2: 0 * time.Second, + expiration: time.Second, + expectedSize: len(testKeys2), + }, + { + scenario: "expire all", + keys1: testKeys1, + keys2: testKeys2, + pause1: 0 * time.Second, + pause2: time.Second, + expiration: time.Second, + expectedSize: 0, + }, + } + + for _, test := range tests { + var b btree + sortedKeys := sortUnique(append(test.keys1, test.keys2...)) + Convey(fmt.Sprintf("running binary tree expire test suit on scenario: [%s]", test.scenario), t, func() { + ok1 := insertKeys(&b, test.keys1) + So(ok1, ShouldBeTrue) + <-time.After(test.pause1) + + ok2 := insertKeys(&b, test.keys2) + So(ok2, ShouldBeTrue) + <-time.After(test.pause2) + + Convey("should be able to find all inserted keys", func() { + ok := findKeys(&b, sortedKeys) + So(ok, ShouldBeTrue) + + Convey("calling trim with expiration", func() { + err := b.expire(test.expiration) + So(err, ShouldBeNil) + + Convey("size should match expected size", func() { + So(b.size(), ShouldEqual, test.expectedSize) + + if test.expectedSize > 0 { + Convey("remaining nodes should be under the age limit", func() { + var failed bool + b.traverseOrdered(b.root, func(n *btreeNode) { + var stop bool + if stop { + return + } + if time.Since(n.age) > test.expiration { + failed = false + stop = true + } + }) + So(failed, ShouldBeFalse) + }) + } + }) + }) + }) + }) + } +} + +func TestCheckpointOffNominal(t *testing.T) { + Convey("off nominal operations on an empty tree", t, func() { + cp := newCheckpoint() + Convey("markDone on an empty tree should error", func() { + err := cp.markDone("1") + So(err, ShouldNotBeNil) + }) + + Convey("trim on an empty tree should error", func() { + err := cp.trim("1") + So(err, ShouldNotBeNil) + + }) + }) + + Convey("off nominal operations on a non-empty tree", t, func() { + cp := newCheckpoint() + cp.insert("2") + cp.insert("3") + cp.insert("1") + Convey("markDone on a key that does not exist should error", func() { + err := cp.markDone("0") + So(err, ShouldNotBeNil) + }) + + Convey("trim on a key that does not exist should error", func() { + err := cp.trim("0") + So(err, ShouldNotBeNil) + + }) + }) +} + +func TestCheckpointConfigs(t *testing.T) { + testKeys := []int{1, -2, 3, -4, 5} + Convey("test auto checkpoint using autoCheckpointCount", t, func() { + autoCheckpointCount := 5 + Convey("instantiating a new checkpoint with options", func() { + var checkpointFnCalled uint64 + cp := newCheckpoint( + checkpointCountCheckFreq(time.Millisecond), + checkpointAutoCheckpointCount(autoCheckpointCount), + checkpointCheckpointFn(func(seqNum string) error { + atomic.AddUint64(&checkpointFnCalled, 1) + return nil + }), + ) + So(cp, ShouldNotBeNil) + So(cp.autoCheckpointCount, ShouldEqual, autoCheckpointCount) + + ctx, cancel := context.WithCancel(context.TODO()) + defer cancel() + cp.startup(ctx) + + Convey("inserting and marking all keys as done", func() { + var failedCount int + for _, testKey := range testKeys { + key := strconv.Itoa(testKey) + if err := cp.insert(key); err != nil { + failedCount++ + } + if err := cp.markDone(key); err != nil { + failedCount++ + } + } + So(failedCount, ShouldEqual, 0) + + <-time.After(time.Millisecond) + Convey("confirming that checkpoint and trim was called", func() { + So(atomic.LoadUint64(&checkpointFnCalled), ShouldBeGreaterThan, 0) + cp.keysMu.Lock() + So(cp.keys.size(), ShouldEqual, 0) + cp.keysMu.Unlock() + }) + }) + }) + }) + + Convey("test auto checkpoint using autoCheckpointFreq", t, func() { + autoCheckpointFreq := 100 * time.Millisecond + Convey("instantiating a new checkpoint with options", func() { + var checkpointFnCalled uint64 + cp := newCheckpoint( + checkpointAutoCheckpointFreq(autoCheckpointFreq), + checkpointCheckpointFn(func(seqNum string) error { + atomic.AddUint64(&checkpointFnCalled, 1) + return nil + }), + ) + So(cp, ShouldNotBeNil) + So(cp.autoCheckpointFreq, ShouldEqual, autoCheckpointFreq) + + ctx, cancel := context.WithCancel(context.TODO()) + defer cancel() + cp.startup(ctx) + + Convey("inserting and marking all keys as done", func() { + var failedCount int + for _, testKey := range testKeys { + key := strconv.Itoa(testKey) + if err := cp.insert(key); err != nil { + failedCount++ + } + if err := cp.markDone(key); err != nil { + failedCount++ + } + } + So(failedCount, ShouldEqual, 0) + + <-time.After(time.Duration(2) * autoCheckpointFreq) + Convey("confirming that checkpoint and trim was called", func() { + So(atomic.LoadUint64(&checkpointFnCalled), ShouldBeGreaterThan, 0) + cp.keysMu.Lock() + So(cp.keys.size(), ShouldEqual, 0) + cp.keysMu.Unlock() + }) + }) + }) + + }) + + Convey("test auto expiration using maxAge", t, func() { + maxAge := 500 * time.Millisecond + Convey("instantiating a new checkpoint with options", func() { + cp := newCheckpoint( + checkpointMaxAge(maxAge), + ) + So(cp, ShouldNotBeNil) + So(cp.maxAge, ShouldEqual, maxAge) + + ctx, cancel := context.WithCancel(context.TODO()) + defer cancel() + cp.startup(ctx) + + Convey("inserting and marking all keys as done", func() { + var failedCount int + for _, testKey := range testKeys { + key := strconv.Itoa(testKey) + if err := cp.insert(key); err != nil { + failedCount++ + } + if err := cp.markDone(key); err != nil { + failedCount++ + } + } + So(failedCount, ShouldEqual, 0) + + <-time.After(time.Duration(2) * maxAge) + Convey("confirming that expire was called", func() { + cp.keysMu.Lock() + So(cp.keys.size(), ShouldEqual, 0) + cp.keysMu.Unlock() + }) + }) + }) + }) + + Convey("test preventing inserts after maxSize capacity is reached", t, func() { + Convey("instantiating a new checkpoint with max size set", func() { + testKeys := []int{1, -23, 45, -67, 89} + maxSize := 4 + var capacityFnCalled bool + cp := newCheckpoint( + checkpointMaxSize(maxSize), + checkpointCapacityFn(func(seqNum string) error { + capacityFnCalled = true + return nil + }), + ) + So(cp, ShouldNotBeNil) + So(cp.maxSize, ShouldEqual, maxSize) + + ctx, cancel := context.WithCancel(context.TODO()) + defer cancel() + cp.startup(ctx) + + Convey("inserting one too many keys than allowed", func() { + var failCount int + for _, testKey := range testKeys { + if err := cp.insert(strconv.Itoa(testKey)); err != nil { + failCount++ + } + } + So(failCount, ShouldBeGreaterThan, 0) + So(capacityFnCalled, ShouldBeTrue) + }) + }) + }) + + Convey("test calling checkpoint function", t, func() { + testKeys := []int{1, -2, 3, -4, 5} + var checkpointFnCalled uint64 + cp := newCheckpoint( + checkpointCheckpointFn(func(seqNum string) error { + atomic.AddUint64(&checkpointFnCalled, 1) + return nil + }), + ) + So(cp, ShouldNotBeNil) + + ctx, cancel := context.WithCancel(context.TODO()) + defer cancel() + cp.startup(ctx) + + Convey("inserting and marking all keys as done", func() { + var failedCount int + for _, testKey := range testKeys { + key := strconv.Itoa(testKey) + if err := cp.insert(key); err != nil { + failedCount++ + } + if err := cp.markDone(key); err != nil { + failedCount++ + } + } + So(failedCount, ShouldEqual, 0) + + Convey("calling checkpoint", func() { + err := cp.checkpoint() + So(err, ShouldBeNil) + + <-time.After(500 * time.Millisecond) + Convey("confirming that checkpointFn was called", func() { + So(atomic.LoadUint64(&checkpointFnCalled), ShouldBeGreaterThan, 0) + cp.keysMu.Lock() + So(cp.keys.size(), ShouldEqual, 0) + cp.keysMu.Unlock() + }) + }) + }) + }) +} + +func TestCheckpointNominal(t *testing.T) { + testKeys := []int{5, 2, 6, 4, 0, 1, 3} + var tests = []struct { + scenario string + keys []int + doneKeys []int + shouldBeOK bool + }{ + { + scenario: "checkpoint/trim when there are no keys inserted", + keys: []int{}, + doneKeys: []int{}, + shouldBeOK: false, + }, + { + scenario: "checkpoint/trim when no keys are marked done", + keys: testKeys, + doneKeys: []int{}, + shouldBeOK: false, + }, + { + scenario: "checkpoint/trim when minimum node is not marked done", + keys: testKeys, + doneKeys: []int{1, 2, 3}, + shouldBeOK: false, + }, + { + scenario: "checkpoint/trim when minimum node is marked done", + keys: testKeys, + doneKeys: []int{2, 3, 1, 0}, + shouldBeOK: true, + }, + { + scenario: "checkpoint/trim when all keys are marked done", + keys: testKeys, + doneKeys: testKeys, + shouldBeOK: true, + }, + { + scenario: "checkpoint/trim only one key exists", + keys: []int{1}, + doneKeys: []int{1}, + shouldBeOK: true, + }, + } + + for _, test := range tests { + cp := newCheckpoint() + Convey(fmt.Sprintf("running checkpoint test suit on scenario: [%s]", test.scenario), t, func() { + Convey("inserting the test keys", func() { + var errCount int + for _, key := range test.keys { + if err := cp.insert(strconv.Itoa(key)); err != nil { + errCount++ + } + } + So(errCount, ShouldEqual, 0) + + Convey("confirming keys were inserted correctly", func() { + So(findKeys(cp.keys, test.keys), ShouldBeTrue) + }) + + uniqueDoneKeys := sortUnique(test.doneKeys) + Convey("marking test keys as done", func() { + if len(test.doneKeys) > 0 { + var errCount int + for _, doneKey := range test.doneKeys { + if err := cp.markDone(strconv.Itoa(doneKey)); err != nil { + errCount++ + } + } + So(errCount, ShouldEqual, 0) + + Convey("confirm that only the test keys were marked done", func() { + var markedDone sortByKeys + cp.keys.traverseOrdered(cp.keys.root, func(n *btreeNode) { + if n.done { + markedDone = append(markedDone, n.seqNum) + } + }) + So(len(markedDone), ShouldEqual, len(uniqueDoneKeys)) + + for _, doneKey := range test.doneKeys { + node, ok := cp.keys.find(doneKey) + So(node.done, ShouldBeTrue) + So(ok, ShouldBeTrue) + } + }) + } + + Convey("calling checkpoint should yeild the expected behavior", func() { + checkpoint, ok := cp.check() + So(checkpoint, ShouldNotBeNil) + So(ok, ShouldEqual, test.shouldBeOK) + //So(cp.lastCheckpoint, ShouldEqual, checkpoint) + + if test.shouldBeOK { + uniqueKeys := sortUnique(test.keys) + Convey("and should also result in the correct checkpoint key", func() { + var expectedCheckpoint int + for i := 0; i < len(uniqueDoneKeys); i++ { + if uniqueKeys[i] == uniqueDoneKeys[i] { + expectedCheckpoint = uniqueDoneKeys[i] + } else { + break + } + } + So(checkpoint, ShouldEqual, strconv.Itoa(expectedCheckpoint)) + }) + + Convey("calling trim using the checkpoint key", func() { + err := cp.trim(checkpoint) + So(err, ShouldBeNil) + + Convey("should not be able to find any keys marked done", func() { + var foundDone bool + var stop bool + cp.keys.traverseOrdered(cp.keys.root, func(n *btreeNode) { + if stop { + return + } + if n.done { + foundDone = true + stop = true + } + }) + So(foundDone, ShouldBeFalse) + }) + + if len(uniqueKeys)-len(uniqueDoneKeys) > 0 { + Convey("should still be able to find remaining keys", func() { + remainingKeys := uniqueKeys[len(uniqueDoneKeys):] + for _, key := range remainingKeys { + node, ok := cp.keys.find(key) + So(ok, ShouldBeTrue) + So(node, ShouldNotBeNil) + So(node.done, ShouldBeFalse) + } + }) + } + }) + } + }) + }) + }) + }) + } +} diff --git a/consumer.go b/consumer.go index d658e49..23289db 100644 --- a/consumer.go +++ b/consumer.go @@ -13,8 +13,8 @@ import ( // consumerOptions is used to hold all of the configurable settings of a Consumer. type consumerOptions struct { reader StreamReader // interface for abstracting the GetRecord/GetRecords call - queueDepth int // size of the consumer's message channel - concurrency int // number of concurrent routines processing messages off of the message channel + queueDepth int // size of the consumer's message channel + concurrency int // number of concurrent routines processing messages off of the message channel logLevel aws.LogLevelType // log level for configuring the LogHelper's log level Stats ConsumerStatsCollector // stats collection mechanism } @@ -79,14 +79,14 @@ func ConsumerStats(sc ConsumerStatsCollector) ConsumerOptionsFn { // Consumer polls the StreamReader for messages. type Consumer struct { - *consumerOptions // contains all of the configuration settings for the Consumer - *LogHelper // object for help with logging + *consumerOptions // contains all of the configuration settings for the Consumer + *LogHelper // object for help with logging txnCountRateLimiter *rate.Limiter // rate limiter to limit the number of transactions per second txSizeRateLimiter *rate.Limiter // rate limiter to limit the transmission size per seccond messages chan *Message // channel for storing messages that have been retrieved from the stream concurrencySem chan empty // channel for controlling the number of concurrent workers processing messages from the message channel pipeOfDeath chan empty // channel for handling pipe of death - consuming bool // flag for indicating whether or not the consumer is consuming + consuming bool // flag for indicating whether or not the consumer is consuming consumingMu sync.Mutex // mutex for making the consuming flag thread safe noCopy noCopy // prevents the Consumer from being copied } diff --git a/consumer_stats.go b/consumer_stats.go index e940edf..f806ef7 100644 --- a/consumer_stats.go +++ b/consumer_stats.go @@ -82,52 +82,52 @@ func (nsc *NilConsumerStatsCollector) AddGetRecordsUnmarshalDuration(time.Durati // Metric names to be exported const ( - MetricsConsumed = "kinetic.consumer.consumed" - MetricsDelivered = "kinetic.consumer.delivered" - MetricsProcessed = "kinetic.consumer.processed" - MetricsBatchSize = "kinetic.consumer.batchsize" - MetricsSent = "kinetic.consumer.sent" + MetricsConsumed = "kinetic.consumer.consumed" + MetricsDelivered = "kinetic.consumer.delivered" + MetricsProcessed = "kinetic.consumer.processed" + MetricsBatchSize = "kinetic.consumer.batchsize" + MetricsSent = "kinetic.consumer.sent" MetricsReadProvisionedThroughputExceeded = "kinetic.consumer.getrecords.provisionedthroughputexceeded" - MetricsGetRecordsTimeout = "kinetic.consumer.getrecords.timeout" - MetricsGetRecordsReadTimeout = "kinetic.consumer.getrecords.readtimeout" - MetricsProcessedDuration = "kinetic.consumer.processed.duration" - MetricsGetRecordsDuration = "kinetic.consumer.getrecords.duration" - MetricsGetRecordsReadResponseDuration = "kinetic.consumer.getrecords.readresponse.duration" - MetricsGetRecordsUnmarshalDuration = "kinetic.consumer.getrecords.unmarshal.duration" + MetricsGetRecordsTimeout = "kinetic.consumer.getrecords.timeout" + MetricsGetRecordsReadTimeout = "kinetic.consumer.getrecords.readtimeout" + MetricsProcessedDuration = "kinetic.consumer.processed.duration" + MetricsGetRecordsDuration = "kinetic.consumer.getrecords.duration" + MetricsGetRecordsReadResponseDuration = "kinetic.consumer.getrecords.readresponse.duration" + MetricsGetRecordsUnmarshalDuration = "kinetic.consumer.getrecords.unmarshal.duration" ) // DefaultConsumerStatsCollector is a type that implements the consumer's StatsCollector interface using the // rcrowley/go-metrics library type DefaultConsumerStatsCollector struct { - Consumed metrics.Counter - Delivered metrics.Counter - Processed metrics.Counter - BatchSize metrics.Counter - GetRecordsCalled metrics.Counter + Consumed metrics.Counter + Delivered metrics.Counter + Processed metrics.Counter + BatchSize metrics.Counter + GetRecordsCalled metrics.Counter ReadProvisionedThroughputExceeded metrics.Counter - GetRecordsTimeout metrics.Counter - GetRecordsReadTimeout metrics.Counter - ProcessedDuration metrics.Gauge - GetRecordsDuration metrics.Gauge - GetRecordsReadResponseDuration metrics.Gauge - GetRecordsUnmarshalDuration metrics.Gauge + GetRecordsTimeout metrics.Counter + GetRecordsReadTimeout metrics.Counter + ProcessedDuration metrics.Gauge + GetRecordsDuration metrics.Gauge + GetRecordsReadResponseDuration metrics.Gauge + GetRecordsUnmarshalDuration metrics.Gauge } // NewDefaultConsumerStatsCollector instantiates a new DefaultStatsCollector object func NewDefaultConsumerStatsCollector(r metrics.Registry) *DefaultConsumerStatsCollector { return &DefaultConsumerStatsCollector{ - Consumed: metrics.GetOrRegisterCounter(MetricsConsumed, r), - Delivered: metrics.GetOrRegisterCounter(MetricsDelivered, r), - Processed: metrics.GetOrRegisterCounter(MetricsProcessed, r), - BatchSize: metrics.GetOrRegisterCounter(MetricsBatchSize, r), - GetRecordsCalled: metrics.GetOrRegisterCounter(MetricsSent, r), - ReadProvisionedThroughputExceeded: metrics.GetOrRegisterCounter(MetricsReadProvisionedThroughputExceeded, r), - GetRecordsTimeout: metrics.GetOrRegisterCounter(MetricsGetRecordsTimeout, r), - GetRecordsReadTimeout: metrics.GetOrRegisterCounter(MetricsGetRecordsReadTimeout, r), - ProcessedDuration: metrics.GetOrRegisterGauge(MetricsProcessedDuration, r), - GetRecordsDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsDuration, r), - GetRecordsReadResponseDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsReadResponseDuration, r), - GetRecordsUnmarshalDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsUnmarshalDuration, r), + Consumed: metrics.GetOrRegisterCounter(MetricsConsumed, r), + Delivered: metrics.GetOrRegisterCounter(MetricsDelivered, r), + Processed: metrics.GetOrRegisterCounter(MetricsProcessed, r), + BatchSize: metrics.GetOrRegisterCounter(MetricsBatchSize, r), + GetRecordsCalled: metrics.GetOrRegisterCounter(MetricsSent, r), + ReadProvisionedThroughputExceeded: metrics.GetOrRegisterCounter(MetricsReadProvisionedThroughputExceeded, r), + GetRecordsTimeout: metrics.GetOrRegisterCounter(MetricsGetRecordsTimeout, r), + GetRecordsReadTimeout: metrics.GetOrRegisterCounter(MetricsGetRecordsReadTimeout, r), + ProcessedDuration: metrics.GetOrRegisterGauge(MetricsProcessedDuration, r), + GetRecordsDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsDuration, r), + GetRecordsReadResponseDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsReadResponseDuration, r), + GetRecordsUnmarshalDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsUnmarshalDuration, r), } } diff --git a/consumer_test.go b/consumer_test.go index 64f926e..ed48d42 100644 --- a/consumer_test.go +++ b/consumer_test.go @@ -39,7 +39,6 @@ func TestConsumer(t *testing.T) { So(k, ShouldNotBeNil) So(err, ShouldBeNil) - stream := "some-consumer-stream" err = k.CreateStream(stream, 1) diff --git a/firehose_writer.go b/firehose_writer.go index 42e1d2b..0ef38f9 100644 --- a/firehose_writer.go +++ b/firehose_writer.go @@ -112,8 +112,8 @@ func NewFirehoseWriter(c *aws.Config, stream string, optionFns ...FirehoseWriter return nil, err } return &FirehoseWriter{ - stream: stream, - client: firehose.New(sess), + stream: stream, + client: firehose.New(sess), firehoseWriterOptions: firehoseWriterOptions, LogHelper: &LogHelper{ LogLevel: firehoseWriterOptions.logLevel, diff --git a/kcl_reader.go b/kcl_reader.go index 68dd990..8a48713 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -8,6 +8,7 @@ import ( "fmt" "os" "sync" + "time" "github.com/aws/aws-sdk-go/aws" ) @@ -18,7 +19,11 @@ const ( // kclReaderOptions is a struct that holds all of the KclReader's configurable parameters. type kclReaderOptions struct { - batchSize int // maximum records per GetRecordsRequest call + batchSize int + autoCheckpointCount int // maximum number of messages pulled off the message queue before triggering an auto checkpoint + autoCheckpointFreq time.Duration // frequency with which to auto checkpoint + checkpointMaxAge time.Duration // maximum duration for which a sequence number lives in the checkpoint system + checkpointMaxSize int // maximum records per GetRecordsRequest call onInitCallbackFn func() error // callback function that gets called after initialization onCheckpointCallbackFn func() error // callback function that gets called after checkpointing onShutdownCallbackFn func() error // callback function that gets called after shutdown @@ -26,12 +31,16 @@ type kclReaderOptions struct { Stats ConsumerStatsCollector // stats collection mechanism } -// defaultkclReaderOptions instantiates a kclReaderOptions with default values. +// defaultKclReaderOptions instantiates a kclReaderOptions with default values. func defaultKclReaderOptions() *kclReaderOptions { return &kclReaderOptions{ - batchSize: kclReaderMaxBatchSize, - logLevel: aws.LogOff, - Stats: &NilConsumerStatsCollector{}, + batchSize: kclReaderMaxBatchSize, + autoCheckpointCount: 10000, + autoCheckpointFreq: time.Minute, + checkpointMaxSize: 1000000, + checkpointMaxAge: time.Hour, + logLevel: aws.LogOff, + Stats: &NilConsumerStatsCollector{}, } } @@ -49,6 +58,38 @@ func kclReaderBatchSize(size int) KclReaderOptionsFn { } } +// KclReaderAutoCheckpointFreq is a functional option method for configuring the KclReader's checkpoint frequency +func KclReaderAutoCheckpointFreq(freq time.Duration) KclReaderOptionsFn { + return func(o *kclReaderOptions) error { + o.autoCheckpointFreq = freq + return nil + } +} + +// KclReaderAutoCheckpointCount is a functional option method for configuring the KclReader's checkpoint count +func KclReaderAutoCheckpointCount(count int) KclReaderOptionsFn { + return func(o *kclReaderOptions) error { + o.autoCheckpointCount = count + return nil + } +} + +// KclReaderCheckpointMaxAge is a functional option method for configuring the KclReader's checkpoint max age +func KclReaderCheckpointMaxAge(age time.Duration) KclReaderOptionsFn { + return func(o *kclReaderOptions) error { + o.checkpointMaxAge = age + return nil + } +} + +// KclReaderCheckpointMaxSize is a functional option method for configuring the KclReader's checkpoint max age +func KclReaderCheckpointMaxSize(size int) KclReaderOptionsFn { + return func(o *kclReaderOptions) error { + o.checkpointMaxSize = size + return nil + } +} + // KclReaderOnInitCallbackFn is a functional option method for configuring the KclReader's // onInitCallbackFn. func KclReaderOnInitCallbackFn(fn func() error) KclReaderOptionsFn { @@ -94,12 +135,14 @@ func kclReaderStats(sc ConsumerStatsCollector) KclReaderOptionsFn { // KclReader handles the KCL Multilang Protocol to read records from KCL type KclReader struct { - *kclReaderOptions // contains all of the configuration settings for the KclReader - *LogHelper // object for help with logging - reader *bufio.Reader // io reader to read from STDIN - messages chan *Message // unbuffered message channel used to throttle the record processing from KCL - startupOnce sync.Once // used to ensure that the startup function is called once - shutdownOnce sync.Once // used to ensure that the shutdown function is called once + *kclReaderOptions // contains all of the configuration settings for the KclReader + *LogHelper // object for help with logging + reader *bufio.Reader // io reader to read from STDIN + checkpoint *checkpoint // data structure used to manage checkpointing + messages chan *Message // unbuffered message channel used to throttle the record processing from KCL + actions chan *actionMessage // unbuffered action message channel used internally to coordinate sending action messages to KCL + startupOnce sync.Once // used to ensure that the startup function is called once + shutdownOnce sync.Once // used to ensure that the shutdown function is called once } // NewKclReader creates a new stream reader to read records from KCL @@ -118,17 +161,52 @@ func NewKclReader(c *aws.Config, optionFns ...KclReaderOptionsFn) (*KclReader, e } func (r *KclReader) process(ctx context.Context) { - r.startupOnce.Do(func (){ - defer func(){ + r.startupOnce.Do(func() { + defer func() { // Reset shutdownOnce to allow the shut down sequence to happen again r.shutdownOnce = sync.Once{} }() // create communication channels r.messages = make(chan *Message) + r.actions = make(chan *actionMessage) + + // instantiate and start the checkpointing system + r.checkpoint = newCheckpoint( + checkpointAutoCheckpointCount(r.autoCheckpointCount), + checkpointAutoCheckpointFreq(r.autoCheckpointFreq), + checkpointMaxAge(r.checkpointMaxAge), + checkpointMaxSize(r.checkpointMaxSize), + checkpointCheckpointFn(func(checkpoint string) error { + r.actions <- newCheckpointMessage(checkpoint) + return nil + }), + checkpointExpireFn(func(checkpoint string) error { + r.LogError(fmt.Sprintf("Checkpoint: Sequence number [%s] exceeded max age", checkpoint)) + return nil + }), + checkpointCapacityFn(func(checkpoint string) error { + r.LogError(fmt.Sprintf("Checkpoint: Sequence number [%s] exceeded max size", checkpoint)) + return nil + }), + ) + r.checkpoint.startup(ctx) + + // send messages to KCL + go func() { + for { + actionMessage := <-r.actions + r.LogInfo(fmt.Sprintf("Sending a %s action message to KCL", actionMessage.Action)) + err := r.sendMessage(actionMessage) + if err != nil { + r.LogError(fmt.Sprintf("Unable to send %s action message due to: ", + actionMessage.Action), err) + } + } + }() // listen to STDIN and processes action messages based on the Multilang protocol from KCL - go func(){ + go func() { defer r.shutdown() for { @@ -149,37 +227,25 @@ func (r *KclReader) process(ctx context.Context) { case kclActionTypeInitialize: r.LogDebug("Receieved Initialize action from KCL") r.onInitCallbackFn() - err := r.sendMessage(newStatusMessage(kclActionTypeInitialize)) - if err != nil { - r.LogError("Unable to send Initialize acknowledgement due to: ", err) - } + r.actions <- newStatusMessage(actionMessage.Action) case kclActionTypeCheckpoint: r.LogDebug("Receieved Checkpoint action from KCL") r.onCheckpointCallbackFn() - err := r.sendMessage(newStatusMessage(kclActionTypeCheckpoint)) - if err != nil { - r.LogError("Unable to send Checkpoint acknowledgement due to: ", err) - } + r.actions <- newStatusMessage(actionMessage.Action) case kcActionTypeShutdown: r.LogDebug("Receieved Shutdown action from KCL") r.onShutdownCallbackFn() - err := r.sendMessage(newStatusMessage(kcActionTypeShutdown)) - if err != nil { - r.LogError("Unable to send Shutdown acknowledgement due to: ", err) - } - return + r.actions <- newStatusMessage(actionMessage.Action) case kclActionTypeProcessRecords: r.LogDebug("Receieved ProcessRecords action from KCL") - // Put all the messages on the reader's message channel - for _, msg := range actionMessage.Records { - r.messages <-msg.ToMessage() - } - - // Send an acknowledgement that all the messages were received - err := r.sendMessage(newStatusMessage(kclActionTypeProcessRecords)) - if err != nil { - r.LogError("Unable to send ProcessRecords acknowledgement due to: ", err) - } + // Put the messages on the reader's message channel + // (one by one as they are pulled off and processed by the reader) + go func() { + for _, msg := range actionMessage.Records { + r.messages <- msg.ToMessage() + } + r.actions <- newStatusMessage(actionMessage.Action) + }() default: r.LogError("processAction received an invalid action: ", actionMessage.Action) } @@ -189,7 +255,7 @@ func (r *KclReader) process(ctx context.Context) { } func (r *KclReader) shutdown() { - defer func(){ + defer func() { // Reset startupOnce to allow the start up sequence to happen again r.startupOnce = sync.Once{} }() @@ -240,7 +306,10 @@ func (r *KclReader) sendMessage(msg interface{}) error { func (r *KclReader) processRecords(ctx context.Context, batchSize int, fn messageHandler) (count int, size int, err error) { r.process(ctx) for i := 0; i < batchSize; i++ { + // pull a message off of the reader's message channel msg := <-r.messages + + // execute the caller's callback function for processing the message err = fn(msg) if err != nil { r.LogError("messageHandler resulted in an error: ", err) @@ -272,3 +341,18 @@ func (r *KclReader) GetRecords(ctx context.Context, fn messageHandler) (count in count, size, err = r.processRecords(ctx, r.batchSize, fn) return count, size, err } + +// Checkpoint sends a message to KCL if there is sequence number that can be checkpointed +func (r *KclReader) Checkpoint() error { + return r.checkpoint.checkpoint() +} + +// CheckpointInsert registers a sequence number with the checkpointing system +func (r *KclReader) CheckpointInsert(seqNum string) error { + return r.checkpoint.insert(seqNum) +} + +// CheckpointDone marks the given sequence number as done in the checkpointing system +func (r *KclReader) CheckpointDone(seqNum string) error { + return r.checkpoint.markDone(seqNum) +} diff --git a/kinesis_reader.go b/kinesis_reader.go index 8591114..dd9c184 100644 --- a/kinesis_reader.go +++ b/kinesis_reader.go @@ -317,3 +317,21 @@ func (r *KinesisReader) GetRecords(ctx context.Context, fn messageHandler) (coun count, size, err = r.getRecords(ctx, fn, r.batchSize) return count, size, err } + +// Checkpoint sends a message to KCL if there is sequence number that can be checkpointed +func (r *KinesisReader) Checkpoint() error { + // No-op (only applicable to KclReader) + return nil +} + +// CheckpointInsert registers a sequence number with the checkpointing system +func (r *KinesisReader) CheckpointInsert(seqNum string) error { + // No-op (only applicable to KclReader) + return nil +} + +// CheckpointDone marks the given sequence number as done in the checkpointing system +func (r *KinesisReader) CheckpointDone(seqNum string) error { + // No-op (only applicable to KclReader) + return nil +} diff --git a/kinesis_writer.go b/kinesis_writer.go index f25a1c6..b990e5d 100644 --- a/kinesis_writer.go +++ b/kinesis_writer.go @@ -109,8 +109,8 @@ func NewKinesisWriter(c *aws.Config, stream string, optionFns ...KinesisWriterOp return nil, err } return &KinesisWriter{ - stream: stream, - client: kinesis.New(sess), + stream: stream, + client: kinesis.New(sess), kinesisWriterOptions: kinesisWriterOptions, LogHelper: &LogHelper{ LogLevel: kinesisWriterOptions.logLevel, diff --git a/kinetic.go b/kinetic.go index c0398e1..0d7445f 100644 --- a/kinetic.go +++ b/kinetic.go @@ -9,11 +9,12 @@ import ( "github.com/aws/aws-sdk-go/aws/request" "github.com/aws/aws-sdk-go/aws/session" //"github.com/aws/aws-sdk-go/service/firehose" + "net/http" + + "github.com/aws/aws-sdk-go/aws/credentials" "github.com/aws/aws-sdk-go/service/firehose/firehoseiface" "github.com/aws/aws-sdk-go/service/kinesis" "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" - "net/http" - "github.com/aws/aws-sdk-go/aws/credentials" ) // kineticOptions is a struct that holds all of Kinetic's configurable parameters. diff --git a/producer.go b/producer.go index 92161e8..f2072f9 100644 --- a/producer.go +++ b/producer.go @@ -47,7 +47,7 @@ func defaultProducerOptions() *producerOptions { // ProducerOptionsFn is a method signature for defining functional option methods for configuring the Producer. type ProducerOptionsFn func(*producerOptions) error -// ProducerWriter is a functional option method for configuing the producer's stream writer. +// ProducerWriter is a functional option method for configuring the producer's stream writer. func ProducerWriter(w StreamWriter) ProducerOptionsFn { return func(o *producerOptions) error { o.writer = w @@ -55,7 +55,7 @@ func ProducerWriter(w StreamWriter) ProducerOptionsFn { } } -// ProducerBatchSize is a functional option method for configuing the producer's batch size. +// ProducerBatchSize is a functional option method for configuring the producer's batch size. func ProducerBatchSize(size int) ProducerOptionsFn { return func(o *producerOptions) error { if size > 0 && size <= putRecordsMaxBatchSize { @@ -66,7 +66,7 @@ func ProducerBatchSize(size int) ProducerOptionsFn { } } -// ProducerBatchTimeout is a functional option method for configuing the producer's batch timeout. +// ProducerBatchTimeout is a functional option method for configuring the producer's batch timeout. func ProducerBatchTimeout(timeout time.Duration) ProducerOptionsFn { return func(o *producerOptions) error { o.batchTimeout = timeout @@ -74,7 +74,7 @@ func ProducerBatchTimeout(timeout time.Duration) ProducerOptionsFn { } } -// ProducerQueueDepth is a functional option method for configuing the producer's queue depth. +// ProducerQueueDepth is a functional option method for configuring the producer's queue depth. func ProducerQueueDepth(queueDepth int) ProducerOptionsFn { return func(o *producerOptions) error { if queueDepth > 0 { @@ -85,10 +85,10 @@ func ProducerQueueDepth(queueDepth int) ProducerOptionsFn { } } -// ProducerMaxRetryAttempts is a functional option method for configuing the producer's max retry attempts. +// ProducerMaxRetryAttempts is a functional option method for configuring the producer's max retry attempts. func ProducerMaxRetryAttempts(attemtps int) ProducerOptionsFn { return func(o *producerOptions) error { - if attemtps > 0 { + if attemtps > 0 { o.maxRetryAttempts = attemtps return nil } @@ -96,7 +96,7 @@ func ProducerMaxRetryAttempts(attemtps int) ProducerOptionsFn { } } -// ProducerConcurrency is a functional option method for configuing the producer's concurrency. +// ProducerConcurrency is a functional option method for configuring the producer's concurrency. func ProducerConcurrency(count int) ProducerOptionsFn { return func(o *producerOptions) error { if count > 0 { @@ -107,7 +107,7 @@ func ProducerConcurrency(count int) ProducerOptionsFn { } } -// ProducerShardCheckFrequency is a functional option method for configuing the producer's shard check frequency. +// ProducerShardCheckFrequency is a functional option method for configuring the producer's shard check frequency. func ProducerShardCheckFrequency(duration time.Duration) ProducerOptionsFn { return func(o *producerOptions) error { o.shardCheckFreq = duration @@ -115,7 +115,7 @@ func ProducerShardCheckFrequency(duration time.Duration) ProducerOptionsFn { } } -// ProducerDataSpillFn is a functional option method for configuing the producer's data spill callback function. +// ProducerDataSpillFn is a functional option method for configuring the producer's data spill callback function. func ProducerDataSpillFn(fn MessageProcessor) ProducerOptionsFn { return func(o *producerOptions) error { o.dataSpillFn = fn @@ -123,7 +123,7 @@ func ProducerDataSpillFn(fn MessageProcessor) ProducerOptionsFn { } } -// ProducerLogLevel is a functional option method for configuing the producer's log level. +// ProducerLogLevel is a functional option method for configuring the producer's log level. func ProducerLogLevel(ll aws.LogLevelType) ProducerOptionsFn { return func(o *producerOptions) error { o.logLevel = ll & 0xffff0000 @@ -131,7 +131,7 @@ func ProducerLogLevel(ll aws.LogLevelType) ProducerOptionsFn { } } -// ProducerStats is a functional option method for configuing the producer's stats collector. +// ProducerStats is a functional option method for configuring the producer's stats collector. func ProducerStats(sc ProducerStatsCollector) ProducerOptionsFn { return func(o *producerOptions) error { o.Stats = sc @@ -141,20 +141,20 @@ func ProducerStats(sc ProducerStatsCollector) ProducerOptionsFn { // Producer sends records to AWS Kinesis or Firehose. type Producer struct { - *producerOptions // contains all of the configuration settings for the Producer - *LogHelper // object for help with logging - msgCountLimiter *rate.Limiter // rate limiter to limit the number of messages dispatched per second - msgSizeLimiter *rate.Limiter // rate limiter to limit the total size (in bytes) of messages dispatched per second - workerCount int // number of concurrent workers sending batch messages for the producer - messages chan *Message // channel for enqueuing messages to be put on the stream - status chan *statusReport // channel for workers to communicate their current status - dismiss chan empty // channel for handling the decommissioning of a surplus of workers - stop chan empty // channel for handling shutdown - pipeOfDeath chan empty // channel for handling pipe of death - startupOnce sync.Once // used to ensure that the startup function is called once - shutdownOnce sync.Once // used to ensure that the shutdown function is called once - resizeMu sync.Mutex // used to prevent resizeWorkerPool from being called synchronously with itself - noCopy noCopy // prevents the Producer from being copied + *producerOptions // contains all of the configuration settings for the Producer + *LogHelper // object for help with logging + msgCountLimiter *rate.Limiter // rate limiter to limit the number of messages dispatched per second + msgSizeLimiter *rate.Limiter // rate limiter to limit the total size (in bytes) of messages dispatched per second + workerCount int // number of concurrent workers sending batch messages for the producer + messages chan *Message // channel for enqueuing messages to be put on the stream + status chan *statusReport // channel for workers to communicate their current status + dismiss chan empty // channel for handling the decommissioning of a surplus of workers + stop chan empty // channel for handling shutdown + pipeOfDeath chan empty // channel for handling pipe of death + startupOnce sync.Once // used to ensure that the startup function is called once + shutdownOnce sync.Once // used to ensure that the shutdown function is called once + resizeMu sync.Mutex // used to prevent resizeWorkerPool from being called synchronously with itself + noCopy noCopy // prevents the Producer from being copied } // NewProducer creates a new producer for writing records to a Kinesis or Firehose stream. diff --git a/shared_types.go b/shared_types.go index 11681e2..be65a11 100644 --- a/shared_types.go +++ b/shared_types.go @@ -16,6 +16,9 @@ type StreamWriter interface { type StreamReader interface { GetRecord(context.Context, messageHandler) (count int, size int, err error) GetRecords(context.Context, messageHandler) (count int, size int, err error) + Checkpoint() error // only applicable for KclReader + CheckpointInsert(seqNum string) error // only applicable for KclReader + CheckpointDone(seqNum string) error // only applicable for KclReader } // MessageProcessor defines the signature of a (asynchronous) callback function used by Listen, RetrieveFn and @@ -40,8 +43,8 @@ func (*noCopy) Lock() {} // statusReport is used to communicate a worker's capacity for new messages and to which channel they should be sent. type statusReport struct { - capacity int // maximum message capacity the worker can handle - failedCount int // number of previous messages that failed to send - failedSize int // size in bytes of the previous messages that failed to send + capacity int // maximum message capacity the worker can handle + failedCount int // number of previous messages that failed to send + failedSize int // size in bytes of the previous messages that failed to send channel chan []*Message // channel of the worker to which the batch messages should be sent } From 9856250aed8b42fabcb4565f8a8f132559f79980 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 14 Aug 2017 13:23:17 -0500 Subject: [PATCH 081/109] Expose function option methods (KclReaderLogLevel and KclReaderStats) --- checkpoint.go | 2 +- kcl_reader.go | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/checkpoint.go b/checkpoint.go index 4af57eb..d08ca10 100644 --- a/checkpoint.go +++ b/checkpoint.go @@ -557,7 +557,7 @@ func (c *checkpoint) trim(checkpointStr string) error { return c.keys.trim(&cp, nil) } -// expire safely removes sequence numberse from the checkpointing system that are older than the given age. +// expire safely removes sequence numbers from the checkpointing system that are older than the given age. func (c *checkpoint) expire(age time.Duration) error { c.keysMu.Lock() defer c.keysMu.Unlock() diff --git a/kcl_reader.go b/kcl_reader.go index 8a48713..e9544ce 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -117,16 +117,16 @@ func KclReaderOnShutdownCallbackFn(fn func() error) KclReaderOptionsFn { } } -// kclReaderLogLevel is a functional option method for configuring the KclReader's log level. -func kclReaderLogLevel(ll aws.LogLevelType) KclReaderOptionsFn { +// KclReaderLogLevel is a functional option method for configuring the KclReader's log level. +func KclReaderLogLevel(ll aws.LogLevelType) KclReaderOptionsFn { return func(o *kclReaderOptions) error { o.logLevel = ll return nil } } -// kclReaderStats is a functional option method for configuring the KclReader's stats collector. -func kclReaderStats(sc ConsumerStatsCollector) KclReaderOptionsFn { +// KclReaderStats is a functional option method for configuring the KclReader's stats collector. +func KclReaderStats(sc ConsumerStatsCollector) KclReaderOptionsFn { return func(o *kclReaderOptions) error { o.Stats = sc return nil From 80ea7c338b601789df1bb48e06386a410bb09c3f Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 18 Aug 2017 11:52:45 -0500 Subject: [PATCH 082/109] Renamed checkpoint to checkpointer. Removed capacity constraint on checkpointer. Removed node expirary on checkpointer. Added checkpoint stats in consumer's stats collector. Added periodic update of checkpointer size in KclReader. --- checkpoint.go => checkpointer.go | 138 ++++----------- checkpoint_test.go => checkpointer_test.go | 190 ++------------------- consumer.go | 4 +- consumer_stats.go | 140 +++++++++++---- consumer_stats_test.go | 102 ++++++++--- glide.lock | 10 +- glide.yaml | 2 + kcl_reader.go | 145 +++++++++------- kinesis_reader.go | 12 +- 9 files changed, 337 insertions(+), 406 deletions(-) rename checkpoint.go => checkpointer.go (74%) rename checkpoint_test.go => checkpointer_test.go (81%) diff --git a/checkpoint.go b/checkpointer.go similarity index 74% rename from checkpoint.go rename to checkpointer.go index d08ca10..d09eb92 100644 --- a/checkpoint.go +++ b/checkpointer.go @@ -12,9 +12,8 @@ import ( // btreeNode is a binary tree node used to encapsulate checkpointing information in a binary search tree. type btreeNode struct { - seqNum int // key to store the sequence number in the checkpointing system + seqNum int // key to store the sequence number in the checkpointer done bool // flag to indicate whether or not processing of a record is complete - age time.Time // time stamp used to expire btreeNodes from the checkpointing system left *btreeNode // pointer to a btreeNode whose sequence number is less than own right *btreeNode // pointer to a btreeNode whose sequence number is greater than own } @@ -32,7 +31,6 @@ func (n *btreeNode) insert(seqNum int) error { if n.left == nil { n.left = &btreeNode{ seqNum: seqNum, - age: time.Now(), } return nil } @@ -41,7 +39,6 @@ func (n *btreeNode) insert(seqNum int) error { if n.right == nil { n.right = &btreeNode{ seqNum: seqNum, - age: time.Now(), } return nil } @@ -125,7 +122,6 @@ func (n *btreeNode) delete(seqNum int, parent *btreeNode) error { replacement, replacementParent := n.left.findMax(n) n.seqNum = replacement.seqNum n.done = replacement.done - n.age = replacement.age return replacement.delete(replacement.seqNum, replacementParent) } @@ -141,12 +137,11 @@ type btree struct { root *btreeNode // the root node for the binary search tree } -// insert is a method used to insert sequence numbers into the binary search tree of the checkpointing system. +// insert is a method used to insert sequence numbers into the binary search tree of the checkpointer. func (t *btree) insert(seqNum int) error { if t.root == nil { t.root = &btreeNode{ seqNum: seqNum, - age: time.Now(), } return nil } @@ -263,31 +258,11 @@ func (t *btree) trim(minSeqNum, maxSeqNum *int) error { return nil } -// expire is a method used to delete btreeNodes from the binary search tree whose age is older than the given age. -func (t *btree) expire(age time.Duration) error { - fakeParent := &btreeNode{right: t.root} - t.traverseChildren(fakeParent, func(n *btreeNode) { - if n.left != nil && time.Since(n.left.age) > age { - n.delete(n.left.seqNum, n) - } - if n.right != nil && time.Since(n.right.age) > age { - n.delete(n.right.seqNum, n) - } - }) - t.root = fakeParent.right - - return nil -} - // checkpointOptions is a struct containing all of the configurable options for a checkpoint object. type checkpointOptions struct { autoCheckpointCount int // count of newly inserted messages before triggering an automatic checkpoint call autoCheckpointFreq time.Duration // frequency with which to automatically call checkpoint - maxAge time.Duration // maximum duration for an inserted sequence number to live before being expired - maxSize int // maximum capacity (number of btreeNodes) of the checkpoint system checkpointFn func(checkpoint string) error // callback function to call on a sequence number when a checkpoint is discovered with the checkpoint call - expireFn func(checkpoint string) error // callback function to call when a sequence number is aged out - capacityFn func(checkpoint string) error // callback function to call when the checkpoint system has reached max capacity countCheckFreq time.Duration // frequency with which to check the insert count } @@ -296,11 +271,7 @@ func defaultCheckpointOptions() *checkpointOptions { return &checkpointOptions{ autoCheckpointCount: 10000, autoCheckpointFreq: time.Minute, - maxAge: time.Hour, - maxSize: 1000000, checkpointFn: func(string) error { return nil }, - expireFn: func(string) error { return nil }, - capacityFn: func(string) error { return nil }, countCheckFreq: time.Second, } } @@ -326,22 +297,6 @@ func checkpointAutoCheckpointFreq(freq time.Duration) checkpointOptionsFn { } } -// checkpointMaxAge is a functional option method for configuring the checkpoint's maximum age. -func checkpointMaxAge(age time.Duration) checkpointOptionsFn { - return func(o *checkpointOptions) error { - o.maxAge = age - return nil - } -} - -// checkpointMaxSize is a functional option method for configuring the checkpoint's maximum size. -func checkpointMaxSize(size int) checkpointOptionsFn { - return func(o *checkpointOptions) error { - o.maxSize = size - return nil - } -} - // checkpointCheckpointFn is a functional option method for configuring the checkpoint's checkpoint callback function. func checkpointCheckpointFn(fn func(string) error) checkpointOptionsFn { return func(o *checkpointOptions) error { @@ -350,22 +305,6 @@ func checkpointCheckpointFn(fn func(string) error) checkpointOptionsFn { } } -// checkpointExpireFn is a functional option method for configuring the checkpoint's expire callback function. -func checkpointExpireFn(fn func(string) error) checkpointOptionsFn { - return func(o *checkpointOptions) error { - o.expireFn = fn - return nil - } -} - -// checkpointCapacityFn is a functional option method for configuring the checkpoint's capacity callback function. -func checkpointCapacityFn(fn func(string) error) checkpointOptionsFn { - return func(o *checkpointOptions) error { - o.capacityFn = fn - return nil - } -} - // checkpointCountCheckFreq is a functional option method for configuring the checkpoint's count check frequency. func checkpointCountCheckFreq(freq time.Duration) checkpointOptionsFn { return func(o *checkpointOptions) error { @@ -381,43 +320,41 @@ func checkpointCountCheckFreq(freq time.Duration) checkpointOptionsFn { // channel and should be marked done using the markDone() function after data processing is completed. The // checkpoint() can be called periodically which may trigger a checkpoint call to KCL if the oldest sequence // numbers have been marked complete. Call startup() to enable automatic checkpointing and expiration. -type checkpoint struct { - *checkpointOptions // contains all of the configuration settings for the checkpoint object - keys *btree // binary search tree used to store Kinesis record sequence numbers - keysMu sync.Mutex // mutex to make keys thread safe - counter uint64 // counter to track the number of messages inserted since the last checkpoint - checkpointCh chan struct{} // channel with which to communicate / coordinate checkpointing - startupOnce sync.Once // used to ensure that the startup function is called once - shutdownOnce sync.Once // used to ensure that the shutdown function is called once +type checkpointer struct { + *checkpointOptions // contains all of the configuration settings for the checkpoint object + keys *btree // binary search tree used to store Kinesis record sequence numbers + keysMu sync.Mutex // mutex to make keys thread safe + counter uint64 // counter to track the number of messages inserted since the last checkpoint + checkpointCh chan empty // channel with which to communicate / coordinate checkpointing + startupOnce sync.Once // used to ensure that the startup function is called once + shutdownOnce sync.Once // used to ensure that the shutdown function is called once } // newCheckpoint instantiates a new checkpoint object with default configuration settings unless the function option // methods are provided to change the default values. -func newCheckpoint(optionFns ...checkpointOptionsFn) *checkpoint { +func newCheckpointer(optionFns ...checkpointOptionsFn) *checkpointer { checkpointOptions := defaultCheckpointOptions() for _, optionFn := range optionFns { optionFn(checkpointOptions) } - return &checkpoint{ + return &checkpointer{ checkpointOptions: checkpointOptions, keys: &btree{}, } } -// startup is a method used to enable automatic checkpointing and expiration of btreeNodes from the -// checkpointing system. -func (c *checkpoint) startup(ctx context.Context) { +// startup is a method used to enable automatic checkpointing. +func (c *checkpointer) startup(ctx context.Context) { c.startupOnce.Do(func() { defer func() { c.shutdownOnce = sync.Once{} }() - c.checkpointCh = make(chan struct{}) + c.checkpointCh = make(chan empty) go func() { defer c.shutdown() autoCheckpointTimer := time.NewTimer(c.autoCheckpointFreq) - expirationTicker := time.NewTicker(c.maxAge) counterCheckTicker := time.NewTicker(c.countCheckFreq) for { @@ -426,16 +363,12 @@ func (c *checkpoint) startup(ctx context.Context) { select { case <-ctx.Done(): autoCheckpointTimer.Stop() - expirationTicker.Stop() counterCheckTicker.Stop() return case <-c.checkpointCh: break wait case <-autoCheckpointTimer.C: break wait - case <-expirationTicker.C: - c.expire(c.maxAge) - break case <-counterCheckTicker.C: break } @@ -457,7 +390,7 @@ func (c *checkpoint) startup(ctx context.Context) { } // shutdown is a method used to clean up the checkpoint object -func (c *checkpoint) shutdown() { +func (c *checkpointer) shutdown() { c.shutdownOnce.Do(func() { defer func() { c.startupOnce = sync.Once{} @@ -469,15 +402,18 @@ func (c *checkpoint) shutdown() { }) } -// insert safely inserts a sequence number into the binary search tree. -func (c *checkpoint) insert(seqNumStr string) error { +// size safely determines the number of nodes in the checkpointer. +func (c *checkpointer) size() int { c.keysMu.Lock() defer c.keysMu.Unlock() - if c.keys.size() >= c.maxSize { - c.capacityFn(seqNumStr) - return errors.New("Unable to insert due to capacity") - } + return c.keys.size() +} + +// insert safely inserts a sequence number into the binary search tree. +func (c *checkpointer) insert(seqNumStr string) error { + c.keysMu.Lock() + defer c.keysMu.Unlock() if seqNum, err := strconv.Atoi(seqNumStr); err == nil { err := c.keys.insert(seqNum) @@ -491,7 +427,7 @@ func (c *checkpoint) insert(seqNumStr string) error { } // markDone safely marks the given sequence number as done. -func (c *checkpoint) markDone(seqNumStr string) error { +func (c *checkpointer) markDone(seqNumStr string) error { c.keysMu.Lock() defer c.keysMu.Unlock() @@ -511,7 +447,7 @@ func (c *checkpoint) markDone(seqNumStr string) error { // check returns the largest sequence number marked as done where all smaller sequence numbers have // also been marked as done. -func (c *checkpoint) check() (checkpoint string, found bool) { +func (c *checkpointer) check() (checkpoint string, found bool) { c.keysMu.Lock() defer c.keysMu.Unlock() @@ -538,9 +474,9 @@ func (c *checkpoint) check() (checkpoint string, found bool) { return checkpoint, found } -// trim safely simplifies the binary tree housing the sequence numbers so that only the smallest element is -// marked as done. -func (c *checkpoint) trim(checkpointStr string) error { +// trim safely simplifies the binary tree housing the sequence numbers so that oldest sequence numbers +// chained and marked as done are removed. +func (c *checkpointer) trim(checkpointStr string) error { c.keysMu.Lock() defer c.keysMu.Unlock() @@ -557,19 +493,11 @@ func (c *checkpoint) trim(checkpointStr string) error { return c.keys.trim(&cp, nil) } -// expire safely removes sequence numbers from the checkpointing system that are older than the given age. -func (c *checkpoint) expire(age time.Duration) error { - c.keysMu.Lock() - defer c.keysMu.Unlock() - - return c.keys.expire(age) -} - -// checkpoint sends a signal to the checkpointCh channel to enable a call to checkpoint. -func (c *checkpoint) checkpoint() error { +// checkpoint sends a signal to the checkpointCh channel to trigger a call to checkpoint (potentially). +func (c *checkpointer) checkpoint() error { if c.checkpointCh == nil { return errors.New("Nil checkpoint channel") } - c.checkpointCh <- struct{}{} + c.checkpointCh <- empty{} return nil } diff --git a/checkpoint_test.go b/checkpointer_test.go similarity index 81% rename from checkpoint_test.go rename to checkpointer_test.go index e3424da..1032f03 100644 --- a/checkpoint_test.go +++ b/checkpointer_test.go @@ -93,7 +93,7 @@ func intPointer(x int) *int { return &x } -func TestCheckpointBTreeInsert(t *testing.T) { +func TestCheckpointerBTreeInsert(t *testing.T) { var tests = []struct { scenario string keys []int @@ -174,7 +174,7 @@ func TestCheckpointBTreeInsert(t *testing.T) { } } -func TestCheckpointBTreeFind(t *testing.T) { +func TestCheckpointerBTreeFind(t *testing.T) { var tests = []struct { scenario string keys []int @@ -232,7 +232,7 @@ func TestCheckpointBTreeFind(t *testing.T) { } } -func TestCheckpointBTreeDelete(t *testing.T) { +func TestCheckpointerBTreeDelete(t *testing.T) { var tests = []struct { scenario string keys []int @@ -413,7 +413,7 @@ func TestCheckpointBTreeDelete(t *testing.T) { } } -func TestCheckpointBTreeTrim(t *testing.T) { +func TestCheckpointerBTreeTrim(t *testing.T) { testKeys := []int{50, 20, 60, 40, 0, 10, 30} var tests = []struct { scenario string @@ -545,105 +545,9 @@ func TestCheckpointBTreeTrim(t *testing.T) { } } -func TestCheckpointBTreeExpire(t *testing.T) { - testKeys1 := []int{20, 30, 10} - testKeys2 := []int{15, 25, -5} - var tests = []struct { - scenario string - keys1 []int - keys2 []int - pause1 time.Duration - pause2 time.Duration - expiration time.Duration - expectedSize int - }{ - { - scenario: "expire on an empty tree", - keys1: []int{}, - keys2: []int{}, - pause1: 0 * time.Second, - pause2: 0 * time.Second, - expiration: time.Second, - expectedSize: 0, - }, - { - scenario: "expire none", - keys1: testKeys1, - keys2: testKeys2, - pause1: 0 * time.Second, - pause2: 0 * time.Second, - expiration: time.Second, - expectedSize: len(testKeys1) + len(testKeys2), - }, - { - scenario: "expire some", - keys1: testKeys1, - keys2: testKeys2, - pause1: time.Second, - pause2: 0 * time.Second, - expiration: time.Second, - expectedSize: len(testKeys2), - }, - { - scenario: "expire all", - keys1: testKeys1, - keys2: testKeys2, - pause1: 0 * time.Second, - pause2: time.Second, - expiration: time.Second, - expectedSize: 0, - }, - } - - for _, test := range tests { - var b btree - sortedKeys := sortUnique(append(test.keys1, test.keys2...)) - Convey(fmt.Sprintf("running binary tree expire test suit on scenario: [%s]", test.scenario), t, func() { - ok1 := insertKeys(&b, test.keys1) - So(ok1, ShouldBeTrue) - <-time.After(test.pause1) - - ok2 := insertKeys(&b, test.keys2) - So(ok2, ShouldBeTrue) - <-time.After(test.pause2) - - Convey("should be able to find all inserted keys", func() { - ok := findKeys(&b, sortedKeys) - So(ok, ShouldBeTrue) - - Convey("calling trim with expiration", func() { - err := b.expire(test.expiration) - So(err, ShouldBeNil) - - Convey("size should match expected size", func() { - So(b.size(), ShouldEqual, test.expectedSize) - - if test.expectedSize > 0 { - Convey("remaining nodes should be under the age limit", func() { - var failed bool - b.traverseOrdered(b.root, func(n *btreeNode) { - var stop bool - if stop { - return - } - if time.Since(n.age) > test.expiration { - failed = false - stop = true - } - }) - So(failed, ShouldBeFalse) - }) - } - }) - }) - }) - }) - } -} - -func TestCheckpointOffNominal(t *testing.T) { +func TestCheckpointerOffNominal(t *testing.T) { Convey("off nominal operations on an empty tree", t, func() { - cp := newCheckpoint() + cp := newCheckpointer() Convey("markDone on an empty tree should error", func() { err := cp.markDone("1") So(err, ShouldNotBeNil) @@ -657,7 +561,7 @@ func TestCheckpointOffNominal(t *testing.T) { }) Convey("off nominal operations on a non-empty tree", t, func() { - cp := newCheckpoint() + cp := newCheckpointer() cp.insert("2") cp.insert("3") cp.insert("1") @@ -674,13 +578,13 @@ func TestCheckpointOffNominal(t *testing.T) { }) } -func TestCheckpointConfigs(t *testing.T) { +func TestCheckpointerConfigs(t *testing.T) { testKeys := []int{1, -2, 3, -4, 5} Convey("test auto checkpoint using autoCheckpointCount", t, func() { autoCheckpointCount := 5 Convey("instantiating a new checkpoint with options", func() { var checkpointFnCalled uint64 - cp := newCheckpoint( + cp := newCheckpointer( checkpointCountCheckFreq(time.Millisecond), checkpointAutoCheckpointCount(autoCheckpointCount), checkpointCheckpointFn(func(seqNum string) error { @@ -723,7 +627,7 @@ func TestCheckpointConfigs(t *testing.T) { autoCheckpointFreq := 100 * time.Millisecond Convey("instantiating a new checkpoint with options", func() { var checkpointFnCalled uint64 - cp := newCheckpoint( + cp := newCheckpointer( checkpointAutoCheckpointFreq(autoCheckpointFreq), checkpointCheckpointFn(func(seqNum string) error { atomic.AddUint64(&checkpointFnCalled, 1) @@ -762,78 +666,10 @@ func TestCheckpointConfigs(t *testing.T) { }) - Convey("test auto expiration using maxAge", t, func() { - maxAge := 500 * time.Millisecond - Convey("instantiating a new checkpoint with options", func() { - cp := newCheckpoint( - checkpointMaxAge(maxAge), - ) - So(cp, ShouldNotBeNil) - So(cp.maxAge, ShouldEqual, maxAge) - - ctx, cancel := context.WithCancel(context.TODO()) - defer cancel() - cp.startup(ctx) - - Convey("inserting and marking all keys as done", func() { - var failedCount int - for _, testKey := range testKeys { - key := strconv.Itoa(testKey) - if err := cp.insert(key); err != nil { - failedCount++ - } - if err := cp.markDone(key); err != nil { - failedCount++ - } - } - So(failedCount, ShouldEqual, 0) - - <-time.After(time.Duration(2) * maxAge) - Convey("confirming that expire was called", func() { - cp.keysMu.Lock() - So(cp.keys.size(), ShouldEqual, 0) - cp.keysMu.Unlock() - }) - }) - }) - }) - - Convey("test preventing inserts after maxSize capacity is reached", t, func() { - Convey("instantiating a new checkpoint with max size set", func() { - testKeys := []int{1, -23, 45, -67, 89} - maxSize := 4 - var capacityFnCalled bool - cp := newCheckpoint( - checkpointMaxSize(maxSize), - checkpointCapacityFn(func(seqNum string) error { - capacityFnCalled = true - return nil - }), - ) - So(cp, ShouldNotBeNil) - So(cp.maxSize, ShouldEqual, maxSize) - - ctx, cancel := context.WithCancel(context.TODO()) - defer cancel() - cp.startup(ctx) - - Convey("inserting one too many keys than allowed", func() { - var failCount int - for _, testKey := range testKeys { - if err := cp.insert(strconv.Itoa(testKey)); err != nil { - failCount++ - } - } - So(failCount, ShouldBeGreaterThan, 0) - So(capacityFnCalled, ShouldBeTrue) - }) - }) - }) - Convey("test calling checkpoint function", t, func() { testKeys := []int{1, -2, 3, -4, 5} var checkpointFnCalled uint64 - cp := newCheckpoint( + cp := newCheckpointer( checkpointCheckpointFn(func(seqNum string) error { atomic.AddUint64(&checkpointFnCalled, 1) return nil @@ -874,7 +710,7 @@ func TestCheckpointConfigs(t *testing.T) { }) } -func TestCheckpointNominal(t *testing.T) { +func TestCheckpointerNominal(t *testing.T) { testKeys := []int{5, 2, 6, 4, 0, 1, 3} var tests = []struct { scenario string @@ -921,7 +757,7 @@ func TestCheckpointNominal(t *testing.T) { } for _, test := range tests { - cp := newCheckpoint() + cp := newCheckpointer() Convey(fmt.Sprintf("running checkpoint test suit on scenario: [%s]", test.scenario), t, func() { Convey("inserting the test keys", func() { var errCount int diff --git a/consumer.go b/consumer.go index 23289db..6a5d048 100644 --- a/consumer.go +++ b/consumer.go @@ -246,7 +246,7 @@ func (c *Consumer) RetrieveFnWithContext(ctx context.Context, fn MessageProcesso start := time.Now() fn(msg) - c.Stats.AddProcessedDuration(time.Since(start)) + c.Stats.UpdateProcessedDuration(time.Since(start)) c.Stats.AddProcessed(1) return nil } @@ -332,7 +332,7 @@ func (c *Consumer) ListenWithContext(ctx context.Context, fn MessageProcessor) { }() start := time.Now() fn(msg) - c.Stats.AddProcessedDuration(time.Since(start)) + c.Stats.UpdateProcessedDuration(time.Since(start)) c.Stats.AddProcessed(1) wg.Done() }(msg) diff --git a/consumer_stats.go b/consumer_stats.go index f806ef7..1e40a2e 100644 --- a/consumer_stats.go +++ b/consumer_stats.go @@ -14,15 +14,21 @@ type ConsumerStatsCollector interface { AddConsumed(int) AddDelivered(int) AddProcessed(int) - AddBatchSize(int) + UpdateBatchSize(int) AddGetRecordsCalled(int) AddReadProvisionedThroughputExceeded(int) AddGetRecordsTimeout(int) AddGetRecordsReadTimeout(int) - AddProcessedDuration(time.Duration) - AddGetRecordsDuration(time.Duration) - AddGetRecordsReadResponseDuration(time.Duration) - AddGetRecordsUnmarshalDuration(time.Duration) + UpdateProcessedDuration(time.Duration) + UpdateGetRecordsDuration(time.Duration) + UpdateGetRecordsReadResponseDuration(time.Duration) + UpdateGetRecordsUnmarshalDuration(time.Duration) + AddCheckpointInsert(int) + AddCheckpointDone(int) + UpdateCheckpointSize(int) + AddCheckpointSent(int) + AddCheckpointSuccess(int) + AddCheckpointError(int) } // NilConsumerStatsCollector is a stats consumer that ignores all metrics. @@ -42,9 +48,9 @@ func (nsc *NilConsumerStatsCollector) AddDelivered(int) {} // messages. func (nsc *NilConsumerStatsCollector) AddProcessed(int) {} -// AddBatchSize records a count of the number of messages returned by +// UpdateBatchSize records a count of the number of messages returned by // GetRecords in the consumer. -func (nsc *NilConsumerStatsCollector) AddBatchSize(int) {} +func (nsc *NilConsumerStatsCollector) UpdateBatchSize(int) {} // AddGetRecordsCalled records the number of times the GetRecords API was called // by the consumer. @@ -64,21 +70,41 @@ func (nsc *NilConsumerStatsCollector) AddGetRecordsTimeout(int) {} // WithGetRecordsReadTimeout configuration. func (nsc *NilConsumerStatsCollector) AddGetRecordsReadTimeout(int) {} -// AddProcessedDuration records the duration to process a record. See notes on +// UpdateProcessedDuration records the duration to process a record. See notes on // AddProcessed. -func (nsc *NilConsumerStatsCollector) AddProcessedDuration(time.Duration) {} +func (nsc *NilConsumerStatsCollector) UpdateProcessedDuration(time.Duration) {} -// AddGetRecordsDuration records the duration that the GetRecords API request +// UpdateGetRecordsDuration records the duration that the GetRecords API request // took. Only the times of successful calls are measured. -func (nsc *NilConsumerStatsCollector) AddGetRecordsDuration(time.Duration) {} +func (nsc *NilConsumerStatsCollector) UpdateGetRecordsDuration(time.Duration) {} -// AddGetRecordsReadResponseDuration records the duration that it took to read +// UpdateGetRecordsReadResponseDuration records the duration that it took to read // the response body of a GetRecords API request. -func (nsc *NilConsumerStatsCollector) AddGetRecordsReadResponseDuration(time.Duration) {} +func (nsc *NilConsumerStatsCollector) UpdateGetRecordsReadResponseDuration(time.Duration) {} -// AddGetRecordsUnmarshalDuration records the duration that it took to unmarshal +// UpdateGetRecordsUnmarshalDuration records the duration that it took to unmarshal // the response body of a GetRecords API request. -func (nsc *NilConsumerStatsCollector) AddGetRecordsUnmarshalDuration(time.Duration) {} +func (nsc *NilConsumerStatsCollector) UpdateGetRecordsUnmarshalDuration(time.Duration) {} + +// AddCheckpointInsert records the number of times the CheckpointInsert API was called. +func (nsc *NilConsumerStatsCollector) AddCheckpointInsert(int) {} + +// AddCheckpointDone records the number of times the CheckpointDone API was called. +func (nsc *NilConsumerStatsCollector) AddCheckpointDone(int) {} + +// UpdateCheckpointSize records the current size of the checkpointer. +func (nsc *NilConsumerStatsCollector) UpdateCheckpointSize(int) {} + +// AddCheckpointSent records the number of times a checkpoint action message was sent to KCL. +func (nsc *NilConsumerStatsCollector) AddCheckpointSent(int) {} + +// AddCheckpointSuccess records the number of times KCL send a checkpoint acknowledgement indicating that +// checkpointing was successful +func (nsc *NilConsumerStatsCollector) AddCheckpointSuccess(int) {} + +// AddCheckpointError records the number of times KCL send a checkpoint acknowledgement indicating that +// checkpointing was not successful +func (nsc *NilConsumerStatsCollector) AddCheckpointError(int) {} // Metric names to be exported const ( @@ -94,6 +120,12 @@ const ( MetricsGetRecordsDuration = "kinetic.consumer.getrecords.duration" MetricsGetRecordsReadResponseDuration = "kinetic.consumer.getrecords.readresponse.duration" MetricsGetRecordsUnmarshalDuration = "kinetic.consumer.getrecords.unmarshal.duration" + MetricsCheckpointInsert = "kinetic.consumer.checkpoint.insert" + MetricsCheckpointDone = "kinetic.consumer.checkpoint.done" + MetricsCheckpointSize = "kinetic.consumer.checkpoint.size" + MetricsCheckpointSent = "kinetic.consumer.checkpoint.sent" + MetricsCheckpointSuccess = "kinetic.consumer.checkpoint.success" + MetricsCheckpointError = "kinetic.consumer.checkpoint.error" ) // DefaultConsumerStatsCollector is a type that implements the consumer's StatsCollector interface using the @@ -102,7 +134,7 @@ type DefaultConsumerStatsCollector struct { Consumed metrics.Counter Delivered metrics.Counter Processed metrics.Counter - BatchSize metrics.Counter + BatchSize metrics.Gauge GetRecordsCalled metrics.Counter ReadProvisionedThroughputExceeded metrics.Counter GetRecordsTimeout metrics.Counter @@ -111,6 +143,12 @@ type DefaultConsumerStatsCollector struct { GetRecordsDuration metrics.Gauge GetRecordsReadResponseDuration metrics.Gauge GetRecordsUnmarshalDuration metrics.Gauge + CheckpointInsert metrics.Counter + CheckpointDone metrics.Counter + CheckpointSize metrics.Gauge + CheckpointSent metrics.Counter + CheckpointSuccess metrics.Counter + CheckpointError metrics.Counter } // NewDefaultConsumerStatsCollector instantiates a new DefaultStatsCollector object @@ -119,7 +157,7 @@ func NewDefaultConsumerStatsCollector(r metrics.Registry) *DefaultConsumerStatsC Consumed: metrics.GetOrRegisterCounter(MetricsConsumed, r), Delivered: metrics.GetOrRegisterCounter(MetricsDelivered, r), Processed: metrics.GetOrRegisterCounter(MetricsProcessed, r), - BatchSize: metrics.GetOrRegisterCounter(MetricsBatchSize, r), + BatchSize: metrics.GetOrRegisterGauge(MetricsBatchSize, r), GetRecordsCalled: metrics.GetOrRegisterCounter(MetricsSent, r), ReadProvisionedThroughputExceeded: metrics.GetOrRegisterCounter(MetricsReadProvisionedThroughputExceeded, r), GetRecordsTimeout: metrics.GetOrRegisterCounter(MetricsGetRecordsTimeout, r), @@ -128,6 +166,12 @@ func NewDefaultConsumerStatsCollector(r metrics.Registry) *DefaultConsumerStatsC GetRecordsDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsDuration, r), GetRecordsReadResponseDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsReadResponseDuration, r), GetRecordsUnmarshalDuration: metrics.GetOrRegisterGauge(MetricsGetRecordsUnmarshalDuration, r), + CheckpointInsert: metrics.GetOrRegisterCounter(MetricsCheckpointInsert, r), + CheckpointDone: metrics.GetOrRegisterCounter(MetricsCheckpointDone, r), + CheckpointSize: metrics.GetOrRegisterGauge(MetricsCheckpointSize, r), + CheckpointSent: metrics.GetOrRegisterCounter(MetricsCheckpointSent, r), + CheckpointSuccess: metrics.GetOrRegisterCounter(MetricsCheckpointSuccess, r), + CheckpointError: metrics.GetOrRegisterCounter(MetricsCheckpointError, r), } } @@ -151,10 +195,10 @@ func (dsc *DefaultConsumerStatsCollector) AddProcessed(count int) { dsc.Processed.Inc(int64(count)) } -// AddBatchSize records a count of the number of messages returned by +// UpdateBatchSize records a count of the number of messages returned by // GetRecords in the consumer. -func (dsc *DefaultConsumerStatsCollector) AddBatchSize(count int) { - dsc.BatchSize.Inc(int64(count)) +func (dsc *DefaultConsumerStatsCollector) UpdateBatchSize(count int) { + dsc.BatchSize.Update(int64(count)) } // AddGetRecordsCalled records the number of times the GetRecords API was called @@ -183,36 +227,68 @@ func (dsc *DefaultConsumerStatsCollector) AddGetRecordsReadTimeout(count int) { dsc.GetRecordsReadTimeout.Inc(int64(count)) } -// AddProcessedDuration records the duration to process a record. See notes on +// UpdateProcessedDuration records the duration to process a record. See notes on // AddProcessed. -func (dsc *DefaultConsumerStatsCollector) AddProcessedDuration(duration time.Duration) { +func (dsc *DefaultConsumerStatsCollector) UpdateProcessedDuration(duration time.Duration) { dsc.ProcessedDuration.Update(duration.Nanoseconds()) } -// AddGetRecordsDuration records the duration that the GetRecords API request +// UpdateGetRecordsDuration records the duration that the GetRecords API request // took. Only the times of successful calls are measured. -func (dsc *DefaultConsumerStatsCollector) AddGetRecordsDuration(duration time.Duration) { +func (dsc *DefaultConsumerStatsCollector) UpdateGetRecordsDuration(duration time.Duration) { dsc.GetRecordsDuration.Update(duration.Nanoseconds()) } -// AddGetRecordsReadResponseDuration records the duration that it took to read +// UpdateGetRecordsReadResponseDuration records the duration that it took to read // the response body of a GetRecords API request. -func (dsc *DefaultConsumerStatsCollector) AddGetRecordsReadResponseDuration(duration time.Duration) { +func (dsc *DefaultConsumerStatsCollector) UpdateGetRecordsReadResponseDuration(duration time.Duration) { dsc.GetRecordsReadResponseDuration.Update(duration.Nanoseconds()) } -// AddGetRecordsUnmarshalDuration records the duration that it took to unmarshal +// UpdateGetRecordsUnmarshalDuration records the duration that it took to unmarshal // the response body of a GetRecords API request. -func (dsc *DefaultConsumerStatsCollector) AddGetRecordsUnmarshalDuration(duration time.Duration) { +func (dsc *DefaultConsumerStatsCollector) UpdateGetRecordsUnmarshalDuration(duration time.Duration) { dsc.GetRecordsUnmarshalDuration.Update(duration.Nanoseconds()) } +// AddCheckpointInsert records the number of times the CheckpointInsert API was called. +func (dsc *DefaultConsumerStatsCollector) AddCheckpointInsert(count int) { + dsc.CheckpointInsert.Inc(int64(count)) +} + +// AddCheckpointDone records the number of times the CheckpointDone API was called. +func (dsc *DefaultConsumerStatsCollector) AddCheckpointDone(count int) { + dsc.CheckpointDone.Inc(int64(count)) +} + +// UpdateCheckpointSize records the current size of the checkpointer. +func (dsc *DefaultConsumerStatsCollector) UpdateCheckpointSize(size int) { + dsc.CheckpointSize.Update(int64(size)) +} + +// AddCheckpointSent records the number of times a checkpoint action message was sent to KCL. +func (dsc *DefaultConsumerStatsCollector) AddCheckpointSent(count int) { + dsc.CheckpointSent.Inc(int64(count)) +} + +// AddCheckpointSuccess records the number of times KCL send a checkpoint acknowledgement indicating that +// checkpointing was successful +func (dsc *DefaultConsumerStatsCollector) AddCheckpointSuccess(count int) { + dsc.CheckpointSuccess.Inc(int64(count)) +} + +// AddCheckpointError records the number of times KCL send a checkpoint acknowledgement indicating that +// checkpointing was not successful +func (dsc *DefaultConsumerStatsCollector) AddCheckpointError(count int) { + dsc.CheckpointError.Inc(int64(count)) +} + // PrintStats logs the stats func (dsc *DefaultConsumerStatsCollector) PrintStats() { log.Printf("Consumer stats: Consumed: [%d]\n", dsc.Consumed.Count()) log.Printf("Consumer stats: Delivered: [%d]\n", dsc.Delivered.Count()) log.Printf("Consumer stats: Processed: [%d]\n", dsc.Processed.Count()) - log.Printf("Consumer stats: Batch Size: [%d]\n", dsc.BatchSize.Count()) + log.Printf("Consumer stats: Batch Size: [%d]\n", dsc.BatchSize.Value()) log.Printf("Consumer stats: GetRecords Called: [%d]\n", dsc.GetRecordsCalled.Count()) log.Printf("Consumer stats: GetRecords Timeout: [%d]\n", dsc.GetRecordsTimeout.Count()) log.Printf("Consumer stats: GetRecords Read Timeout: [%d]\n", dsc.GetRecordsReadTimeout.Count()) @@ -221,4 +297,10 @@ func (dsc *DefaultConsumerStatsCollector) PrintStats() { log.Printf("Consumer stats: GetRecords Duration (ns): [%d]\n", dsc.GetRecordsDuration.Value()) log.Printf("Consumer stats: GetRecords Read Response Duration (ns): [%d]\n", dsc.GetRecordsReadResponseDuration.Value()) log.Printf("Consumer stats: GetRecords Unmarshal Duration (ns): [%d]\n", dsc.GetRecordsUnmarshalDuration.Value()) + log.Printf("Consumer stats: Checkpoint Insert: [%d]\n", dsc.CheckpointInsert.Count()) + log.Printf("Consumer stats: Checkpoint Done: [%d]\n", dsc.CheckpointDone.Count()) + log.Printf("Consumer stats: Checkpoint Size: [%d]\n", dsc.CheckpointSize.Value()) + log.Printf("Consumer stats: Checkpoint Sent: [%d]\n", dsc.CheckpointSent.Count()) + log.Printf("Consumer stats: Checkpoint Success: [%d]\n", dsc.CheckpointSuccess.Count()) + log.Printf("Consumer stats: Checkpoint Error: [%d]\n", dsc.CheckpointError.Count()) } diff --git a/consumer_stats_test.go b/consumer_stats_test.go index 4cc75b3..6d135e1 100644 --- a/consumer_stats_test.go +++ b/consumer_stats_test.go @@ -30,8 +30,8 @@ func TestConsumerStatsCollector(t *testing.T) { sc.AddProcessed(1) }) - Convey("check that AddBatchSize does not error", func() { - sc.AddBatchSize(1) + Convey("check that UpdateBatchSize does not error", func() { + sc.UpdateBatchSize(1) }) Convey("check that AddGetRecordsCalled does not error", func() { @@ -50,20 +50,44 @@ func TestConsumerStatsCollector(t *testing.T) { sc.AddGetRecordsReadTimeout(1) }) - Convey("check that AddProcessedDuration does not error", func() { - sc.AddProcessedDuration(1) + Convey("check that UpdateProcessedDuration does not error", func() { + sc.UpdateProcessedDuration(1) }) - Convey("check that AddGetRecordsDuration does not error", func() { - sc.AddGetRecordsDuration(1) + Convey("check that UpdateGetRecordsDuration does not error", func() { + sc.UpdateGetRecordsDuration(1) }) - Convey("check that AddGetRecordsReadResponseDuration does not error", func() { - sc.AddGetRecordsReadResponseDuration(1) + Convey("check that UpdateGetRecordsReadResponseDuration does not error", func() { + sc.UpdateGetRecordsReadResponseDuration(1) }) - Convey("check that AddGetRecordsUnmarshalDuration does not error", func() { - sc.AddGetRecordsUnmarshalDuration(1) + Convey("check that UpdateGetRecordsUnmarshalDuration does not error", func() { + sc.UpdateGetRecordsUnmarshalDuration(1) + }) + + Convey("check that AddCheckpointInsert does not error", func() { + sc.AddCheckpointInsert(1) + }) + + Convey("check that AddCheckpointDone does not error", func() { + sc.AddCheckpointDone(1) + }) + + Convey("check that UpdateCheckpointSize does not error", func() { + sc.UpdateCheckpointSize(1) + }) + + Convey("check that AddCheckpointSent does not error", func() { + sc.AddCheckpointSent(1) + }) + + Convey("check that AddCheckpointSuccess does not error", func() { + sc.AddCheckpointSuccess(1) + }) + + Convey("check that AddCheckpointError does not error", func() { + sc.AddCheckpointError(1) }) }) @@ -90,10 +114,10 @@ func TestConsumerStatsCollector(t *testing.T) { So(sc.(*DefaultConsumerStatsCollector).Processed.Count(), ShouldEqual, int64(count)) }) - Convey("check that AddBatchSize does not error", func() { + Convey("check that UpdateBatchSize does not error", func() { count := rand.Int() - sc.AddBatchSize(count) - So(sc.(*DefaultConsumerStatsCollector).BatchSize.Count(), ShouldEqual, int64(count)) + sc.UpdateBatchSize(count) + So(sc.(*DefaultConsumerStatsCollector).BatchSize.Value(), ShouldEqual, int64(count)) }) Convey("check that AddGetRecordsCalled does not error", func() { @@ -120,24 +144,60 @@ func TestConsumerStatsCollector(t *testing.T) { So(sc.(*DefaultConsumerStatsCollector).GetRecordsReadTimeout.Count(), ShouldEqual, int64(count)) }) - Convey("check that AddProcessedDuration does not error", func() { - sc.AddProcessedDuration(time.Second) + Convey("check that UpdateProcessedDuration does not error", func() { + sc.UpdateProcessedDuration(time.Second) So(sc.(*DefaultConsumerStatsCollector).ProcessedDuration.Value(), ShouldEqual, 1000000000) }) - Convey("check that AddGetRecordsDuration does not error", func() { - sc.AddGetRecordsDuration(time.Second) + Convey("check that UpdateGetRecordsDuration does not error", func() { + sc.UpdateGetRecordsDuration(time.Second) So(sc.(*DefaultConsumerStatsCollector).GetRecordsDuration.Value(), ShouldEqual, 1000000000) }) - Convey("check that AddGetRecordsReadResponseDuration does not error", func() { - sc.AddGetRecordsReadResponseDuration(time.Second) + Convey("check that UpdateGetRecordsReadResponseDuration does not error", func() { + sc.UpdateGetRecordsReadResponseDuration(time.Second) So(sc.(*DefaultConsumerStatsCollector).GetRecordsReadResponseDuration.Value(), ShouldEqual, 1000000000) }) - Convey("check that AddGetRecordsUnmarshalDuration does not error", func() { - sc.AddGetRecordsUnmarshalDuration(time.Second) + Convey("check that UpdateGetRecordsUnmarshalDuration does not error", func() { + sc.UpdateGetRecordsUnmarshalDuration(time.Second) So(sc.(*DefaultConsumerStatsCollector).GetRecordsUnmarshalDuration.Value(), ShouldEqual, 1000000000) }) + + Convey("check that AddCheckpointInsert does not error", func() { + count := rand.Int() + sc.AddCheckpointInsert(count) + So(sc.(*DefaultConsumerStatsCollector).CheckpointInsert.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddCheckpointDone does not error", func() { + count := rand.Int() + sc.AddCheckpointDone(count) + So(sc.(*DefaultConsumerStatsCollector).CheckpointDone.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that UpdateCheckpointSize does not error", func() { + count := rand.Int() + sc.UpdateCheckpointSize(count) + So(sc.(*DefaultConsumerStatsCollector).CheckpointSize.Value(), ShouldEqual, int64(count)) + }) + + Convey("check that AddCheckpointSent does not error", func() { + count := rand.Int() + sc.AddCheckpointSent(count) + So(sc.(*DefaultConsumerStatsCollector).CheckpointSent.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddCheckpointSuccess does not error", func() { + count := rand.Int() + sc.AddCheckpointSuccess(count) + So(sc.(*DefaultConsumerStatsCollector).CheckpointSuccess.Count(), ShouldEqual, int64(count)) + }) + + Convey("check that AddCheckpointError does not error", func() { + count := rand.Int() + sc.AddCheckpointError(count) + So(sc.(*DefaultConsumerStatsCollector).CheckpointError.Count(), ShouldEqual, int64(count)) + }) }) } diff --git a/glide.lock b/glide.lock index fb83edb..220600a 100644 --- a/glide.lock +++ b/glide.lock @@ -1,5 +1,5 @@ -hash: 7203d745bb3c303ab9257a7d20b250ce919248e1251502ebe6d56a5af086126c -updated: 2017-07-28T11:21:23.082612845-05:00 +hash: f6ee7c06e2bb454e7b5cd139b007893526a80c6e5856e0a1c24ad03a086949c5 +updated: 2017-08-16T16:46:40.529793303-05:00 imports: - name: github.com/aws/aws-sdk-go version: 72e42b13da62269f68308fb6068b7ea691a416a4 @@ -34,13 +34,13 @@ imports: - service/kinesis/kinesisiface - service/sts - name: github.com/go-ini/ini - version: 1730955e3146956d6a087861380f9b4667ed5071 + version: c787282c39ac1fc618827141a1f762240def08a3 - name: github.com/jmespath/go-jmespath version: bd40a432e4c76585ef6b72d3fd96fb9b6dc7b68d - name: github.com/rcrowley/go-metrics version: ab2277b1c5d15c3cba104e9cbddbdfc622df5ad8 - name: golang.org/x/net - version: 0e2717dc3cc05907dc23096ef3a9086ea93f567f + version: 1c05540f6879653db88113bc4a2b70aec4bd491f subpackages: - context - name: golang.org/x/time @@ -59,7 +59,7 @@ testImports: subpackages: - difflib - name: github.com/smartystreets/assertions - version: c9ee7d9e9a2aeec0bee7c4a516f3e0ad7cb7e558 + version: e9a2f6771d97c709b2f7609fdf89ebde8b64c926 subpackages: - internal/go-render/render - internal/oglematchers diff --git a/glide.yaml b/glide.yaml index 9ce7ec7..dfe125f 100644 --- a/glide.yaml +++ b/glide.yaml @@ -21,6 +21,8 @@ testImport: version: v1.6.2 subpackages: - convey +- package: github.com/smartystreets/assertions + version: v1.6.0 - package: github.com/stretchr/testify version: v1.1.4 subpackages: diff --git a/kcl_reader.go b/kcl_reader.go index e9544ce..960cff1 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -19,28 +19,26 @@ const ( // kclReaderOptions is a struct that holds all of the KclReader's configurable parameters. type kclReaderOptions struct { - batchSize int - autoCheckpointCount int // maximum number of messages pulled off the message queue before triggering an auto checkpoint - autoCheckpointFreq time.Duration // frequency with which to auto checkpoint - checkpointMaxAge time.Duration // maximum duration for which a sequence number lives in the checkpoint system - checkpointMaxSize int // maximum records per GetRecordsRequest call - onInitCallbackFn func() error // callback function that gets called after initialization - onCheckpointCallbackFn func() error // callback function that gets called after checkpointing - onShutdownCallbackFn func() error // callback function that gets called after shutdown - logLevel aws.LogLevelType // log level for configuring the LogHelper's log level - Stats ConsumerStatsCollector // stats collection mechanism + batchSize int + autoCheckpointCount int // maximum number of messages pulled off the message queue before triggering an auto checkpoint + autoCheckpointFreq time.Duration // frequency with which to auto checkpoint + updateCheckpointSizeFreq time.Duration // frequency with which to update the CheckpointSize stats + onInitCallbackFn func() error // callback function that gets called after an initialize action message is sent from KCL + onCheckpointCallbackFn func(string, string) error // callback function that gets called after a checkpoint action message is sent from KCL + onShutdownCallbackFn func() error // callback function that gets called after shutdown action message is sent from KCL + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + Stats ConsumerStatsCollector // stats collection mechanism } // defaultKclReaderOptions instantiates a kclReaderOptions with default values. func defaultKclReaderOptions() *kclReaderOptions { return &kclReaderOptions{ - batchSize: kclReaderMaxBatchSize, - autoCheckpointCount: 10000, - autoCheckpointFreq: time.Minute, - checkpointMaxSize: 1000000, - checkpointMaxAge: time.Hour, - logLevel: aws.LogOff, - Stats: &NilConsumerStatsCollector{}, + batchSize: kclReaderMaxBatchSize, + autoCheckpointCount: 10000, + autoCheckpointFreq: time.Minute, + updateCheckpointSizeFreq: time.Minute, + logLevel: aws.LogOff, + Stats: &NilConsumerStatsCollector{}, } } @@ -58,14 +56,6 @@ func kclReaderBatchSize(size int) KclReaderOptionsFn { } } -// KclReaderAutoCheckpointFreq is a functional option method for configuring the KclReader's checkpoint frequency -func KclReaderAutoCheckpointFreq(freq time.Duration) KclReaderOptionsFn { - return func(o *kclReaderOptions) error { - o.autoCheckpointFreq = freq - return nil - } -} - // KclReaderAutoCheckpointCount is a functional option method for configuring the KclReader's checkpoint count func KclReaderAutoCheckpointCount(count int) KclReaderOptionsFn { return func(o *kclReaderOptions) error { @@ -74,18 +64,19 @@ func KclReaderAutoCheckpointCount(count int) KclReaderOptionsFn { } } -// KclReaderCheckpointMaxAge is a functional option method for configuring the KclReader's checkpoint max age -func KclReaderCheckpointMaxAge(age time.Duration) KclReaderOptionsFn { +// KclReaderAutoCheckpointFreq is a functional option method for configuring the KclReader's checkpoint frequency +func KclReaderAutoCheckpointFreq(freq time.Duration) KclReaderOptionsFn { return func(o *kclReaderOptions) error { - o.checkpointMaxAge = age + o.autoCheckpointFreq = freq return nil } } -// KclReaderCheckpointMaxSize is a functional option method for configuring the KclReader's checkpoint max age -func KclReaderCheckpointMaxSize(size int) KclReaderOptionsFn { +// KclReaderUpdateCheckpointSizeFreq is a functional option method for configuring the KclReader's +// update checkpoint size stats frequency +func KclReaderUpdateCheckpointSizeFreq(freq time.Duration) KclReaderOptionsFn { return func(o *kclReaderOptions) error { - o.checkpointMaxSize = size + o.updateCheckpointSizeFreq = freq return nil } } @@ -101,7 +92,7 @@ func KclReaderOnInitCallbackFn(fn func() error) KclReaderOptionsFn { // KclReaderOnCheckpointCallbackFn is a functional option method for configuring the KclReader's // onCheckpointCallbackFn. -func KclReaderOnCheckpointCallbackFn(fn func() error) KclReaderOptionsFn { +func KclReaderOnCheckpointCallbackFn(fn func(seqNum string, err string) error) KclReaderOptionsFn { return func(o *kclReaderOptions) error { o.onCheckpointCallbackFn = fn return nil @@ -138,7 +129,9 @@ type KclReader struct { *kclReaderOptions // contains all of the configuration settings for the KclReader *LogHelper // object for help with logging reader *bufio.Reader // io reader to read from STDIN - checkpoint *checkpoint // data structure used to manage checkpointing + checkpointer *checkpointer // data structure used to manage checkpointing + ticker *time.Ticker // a ticker with which to update the CheckpointSize stats + tickerDone chan empty // a channel used to communicate when to stop updating the CheckpointSize stats messages chan *Message // unbuffered message channel used to throttle the record processing from KCL actions chan *actionMessage // unbuffered action message channel used internally to coordinate sending action messages to KCL startupOnce sync.Once // used to ensure that the startup function is called once @@ -170,34 +163,25 @@ func (r *KclReader) process(ctx context.Context) { // create communication channels r.messages = make(chan *Message) r.actions = make(chan *actionMessage) + r.tickerDone = make(chan empty) - // instantiate and start the checkpointing system - r.checkpoint = newCheckpoint( + // instantiate and start the checkpointer + r.checkpointer = newCheckpointer( checkpointAutoCheckpointCount(r.autoCheckpointCount), checkpointAutoCheckpointFreq(r.autoCheckpointFreq), - checkpointMaxAge(r.checkpointMaxAge), - checkpointMaxSize(r.checkpointMaxSize), checkpointCheckpointFn(func(checkpoint string) error { r.actions <- newCheckpointMessage(checkpoint) - return nil - }), - checkpointExpireFn(func(checkpoint string) error { - r.LogError(fmt.Sprintf("Checkpoint: Sequence number [%s] exceeded max age", checkpoint)) - return nil - }), - checkpointCapacityFn(func(checkpoint string) error { - r.LogError(fmt.Sprintf("Checkpoint: Sequence number [%s] exceeded max size", checkpoint)) + r.Stats.AddCheckpointSent(1) return nil }), ) - r.checkpoint.startup(ctx) + r.checkpointer.startup(ctx) // send messages to KCL go func() { - for { - actionMessage := <-r.actions + for actionMessage := range r.actions { r.LogInfo(fmt.Sprintf("Sending a %s action message to KCL", actionMessage.Action)) - err := r.sendMessage(actionMessage) + err := r.sendToStdOut(actionMessage) if err != nil { r.LogError(fmt.Sprintf("Unable to send %s action message due to: ", actionMessage.Action), err) @@ -218,7 +202,7 @@ func (r *KclReader) process(ctx context.Context) { } // Retrieve action message - actionMessage, err := r.getAction() + actionMessage, err := r.receiveFromStdIn() if err != nil || actionMessage == nil { return } @@ -230,8 +214,16 @@ func (r *KclReader) process(ctx context.Context) { r.actions <- newStatusMessage(actionMessage.Action) case kclActionTypeCheckpoint: r.LogDebug("Receieved Checkpoint action from KCL") - r.onCheckpointCallbackFn() - r.actions <- newStatusMessage(actionMessage.Action) + if actionMessage.Error == "" { + r.LogInfo("Checkpoint acknowledgement from KCL was successful "+ + "for sequence number: ", actionMessage.SequenceNumber) + r.Stats.AddCheckpointSuccess(1) + } else { + r.LogError("Checkpoint acknowledgement from KCL had an error: ", + actionMessage.Error) + r.Stats.AddCheckpointError(1) + } + r.onCheckpointCallbackFn(actionMessage.SequenceNumber, actionMessage.Error) case kcActionTypeShutdown: r.LogDebug("Receieved Shutdown action from KCL") r.onShutdownCallbackFn() @@ -251,6 +243,19 @@ func (r *KclReader) process(ctx context.Context) { } } }() + + // Periodically update the CheckpointSize stats + go func() { + r.ticker = time.NewTicker(r.updateCheckpointSizeFreq) + defer r.ticker.Stop() + + select { + case <-r.ticker.C: + r.Stats.UpdateCheckpointSize(r.checkpointer.size()) + case <-r.tickerDone: + return + } + }() }) } @@ -263,10 +268,18 @@ func (r *KclReader) shutdown() { if r.messages != nil { close(r.messages) } + + if r.actions != nil { + close(r.actions) + } + + if r.tickerDone != nil { + close(r.tickerDone) + } } -// getAction reads messages from STDIN based on the Multilang Daemon protocol from KCL -func (r *KclReader) getAction() (*actionMessage, error) { +// receiveFromStdIn reads messages from STDIN based on the Multilang Daemon protocol from KCL +func (r *KclReader) receiveFromStdIn() (*actionMessage, error) { buffer := &bytes.Buffer{} for { line, isPrefix, err := r.reader.ReadLine() @@ -289,8 +302,8 @@ func (r *KclReader) getAction() (*actionMessage, error) { return actionMsg, nil } -// sendMessage writes messages to STDOUT based on the Multilang Daemon protocol from KCL -func (r *KclReader) sendMessage(msg interface{}) error { +// sendToStdOut writes messages to STDOUT based on the Multilang Daemon protocol from KCL +func (r *KclReader) sendToStdOut(msg interface{}) error { b, err := json.Marshal(msg) if err != nil { r.LogError(err) @@ -344,15 +357,25 @@ func (r *KclReader) GetRecords(ctx context.Context, fn messageHandler) (count in // Checkpoint sends a message to KCL if there is sequence number that can be checkpointed func (r *KclReader) Checkpoint() error { - return r.checkpoint.checkpoint() + return r.checkpointer.checkpoint() } -// CheckpointInsert registers a sequence number with the checkpointing system +// CheckpointInsert registers a sequence number with the checkpointer func (r *KclReader) CheckpointInsert(seqNum string) error { - return r.checkpoint.insert(seqNum) + err := r.checkpointer.insert(seqNum) + if err != nil { + return err + } + r.Stats.AddCheckpointInsert(1) + return nil } -// CheckpointDone marks the given sequence number as done in the checkpointing system +// CheckpointDone marks the given sequence number as done in the checkpointer func (r *KclReader) CheckpointDone(seqNum string) error { - return r.checkpoint.markDone(seqNum) + err := r.checkpointer.markDone(seqNum) + if err != nil { + return err + } + r.Stats.AddCheckpointDone(1) + return nil } diff --git a/kinesis_reader.go b/kinesis_reader.go index dd9c184..ff9b53f 100644 --- a/kinesis_reader.go +++ b/kinesis_reader.go @@ -230,7 +230,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn messageHandler, batch req.HTTPResponse.Body = &ReadCloserWrapper{ ReadCloser: req.HTTPResponse.Body, OnCloseFn: func() { - r.Stats.AddGetRecordsReadResponseDuration(time.Since(startReadTime)) + r.Stats.UpdateGetRecordsReadResponseDuration(time.Since(startReadTime)) r.LogDebug("Finished GetRecords body read, took", time.Since(start)) startUnmarshalTime = time.Now() }, @@ -239,7 +239,7 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn messageHandler, batch req.Handlers.Unmarshal.PushBack(func(req *request.Request) { size += int(req.HTTPRequest.ContentLength) - r.Stats.AddGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) + r.Stats.UpdateGetRecordsUnmarshalDuration(time.Since(startUnmarshalTime)) r.LogDebug("Finished GetRecords Unmarshal, took", time.Since(start)) }) @@ -256,14 +256,14 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn messageHandler, batch } return count, size, err } - r.Stats.AddGetRecordsDuration(time.Since(start)) + r.Stats.UpdateGetRecordsDuration(time.Since(start)) // Process Records r.LogDebug(fmt.Sprintf("Finished GetRecords request, %d records from shard %s, took %v\n", len(resp.Records), r.shard, time.Since(start))) if resp == nil { return count, size, ErrNilGetRecordsResponse } - r.Stats.AddBatchSize(len(resp.Records)) + r.Stats.UpdateBatchSize(len(resp.Records)) for _, record := range resp.Records { if record != nil { // Allow (only) a pipeOfDeath to trigger an instance shutdown of the loop to deliver messages. @@ -324,13 +324,13 @@ func (r *KinesisReader) Checkpoint() error { return nil } -// CheckpointInsert registers a sequence number with the checkpointing system +// CheckpointInsert registers a sequence number with the checkpointer func (r *KinesisReader) CheckpointInsert(seqNum string) error { // No-op (only applicable to KclReader) return nil } -// CheckpointDone marks the given sequence number as done in the checkpointing system +// CheckpointDone marks the given sequence number as done in the checkpointer func (r *KinesisReader) CheckpointDone(seqNum string) error { // No-op (only applicable to KclReader) return nil From da9aabc74fb2a62b7432b2db2f4f4f1aca688833 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 22 Aug 2017 16:55:37 -0500 Subject: [PATCH 083/109] Refactored the checkpointer implementation to use a doubly linked list (instead of a binary search tree). Modified markDone so delete elements from the list upon forming a chain of sequence numbers marked done in order to address the potential for unchecked memory growth. Updated all unit tests to reflect the new expected behavior. --- checkpointer.go | 375 ++++--------------- checkpointer_test.go | 867 ++++++++++++------------------------------- 2 files changed, 324 insertions(+), 918 deletions(-) diff --git a/checkpointer.go b/checkpointer.go index d09eb92..d9c0303 100644 --- a/checkpointer.go +++ b/checkpointer.go @@ -1,261 +1,66 @@ package kinetic import ( + "container/list" "context" "errors" - "math" "strconv" "sync" "sync/atomic" "time" ) -// btreeNode is a binary tree node used to encapsulate checkpointing information in a binary search tree. -type btreeNode struct { - seqNum int // key to store the sequence number in the checkpointer - done bool // flag to indicate whether or not processing of a record is complete - left *btreeNode // pointer to a btreeNode whose sequence number is less than own - right *btreeNode // pointer to a btreeNode whose sequence number is greater than own +// checkpointElement is the struct used to store checkpointing information. +type checkpointElement struct { + seqNum int + done bool } -// insert is a (recursive) method used to insert sequence numbers into the binary search tree. -func (n *btreeNode) insert(seqNum int) error { - if n == nil { - return errors.New("Cannot call insert on a nil node") - } - - switch { - case seqNum == n.seqNum: - return nil - case seqNum < n.seqNum: - if n.left == nil { - n.left = &btreeNode{ - seqNum: seqNum, - } - return nil - } - return n.left.insert(seqNum) - case seqNum > n.seqNum: - if n.right == nil { - n.right = &btreeNode{ - seqNum: seqNum, - } - return nil - } - return n.right.insert(seqNum) - } - return nil -} - -// find is a (recursive) method used to retrieve a pointer to the btreeNode containing the sequence number or 'false' -// if the sequence number provided is not found. -func (n *btreeNode) find(seqNum int) (*btreeNode, bool) { - if n == nil { - return nil, false - } - - switch { - case seqNum == n.seqNum: - return n, true - case seqNum < n.seqNum: - return n.left.find(seqNum) - default: - return n.right.find(seqNum) - } -} - -// findMin is a helper function used to find the minimum element in the subtree of the given parent node. -func (n *btreeNode) findMin(parent *btreeNode) (*btreeNode, *btreeNode) { - if n.left == nil { - return n, parent - } - return n.left.findMin(n) -} - -// findMax is a helper function used to find the maximum element in the subtree of the given parent node. -func (n *btreeNode) findMax(parent *btreeNode) (*btreeNode, *btreeNode) { - if n.right == nil { - return n, parent - } - return n.right.findMax(n) -} - -// replaceNode is a helper method used to replace btreeNodes with a binary search tree. -func (n *btreeNode) replaceNode(parent, replacement *btreeNode) error { - if n == nil { - return errors.New("Cannot call replaceNode on a nil node") - } - - if n == parent.left { - parent.left = replacement - return nil - } - parent.right = replacement - return nil +// checkpointList is a doubly linked list used to track sequence numbers. Sequence numbers are linked in ascending +// order. +type checkpointList struct { + *list.List } -// delete is a (recursive) method used remove (and re-organize if necessary) a btreeNode from the binary search tree. -func (n *btreeNode) delete(seqNum int, parent *btreeNode) error { - if n == nil { - return errors.New("Cannot call delete on a nil node") - } - - switch { - case seqNum < n.seqNum: - return n.left.delete(seqNum, n) - case seqNum > n.seqNum: - return n.right.delete(seqNum, n) - default: - if n.left == nil && n.right == nil { - n.replaceNode(parent, nil) +// insert is a method used to insert sequence numbers (ordered) into the doubly linked list. insert is optimized to +// insert from the back, so inserting in sequential or ascending order is faster than reverse / descending order. +func (c *checkpointList) insert(seqNum int) error { + value := &checkpointElement{seqNum: seqNum} + for e := c.Back(); e != nil; e = e.Prev() { + element := e.Value.(*checkpointElement) + switch { + case element.seqNum < seqNum: + c.InsertAfter(value, e) return nil - } - if n.left == nil { - n.replaceNode(parent, n.right) - return nil - } - if n.right == nil { - n.replaceNode(parent, n.left) + case element.seqNum == seqNum: return nil + case element.seqNum > seqNum: + if e.Prev() != nil && e.Prev().Value.(*checkpointElement).seqNum < seqNum { + c.InsertBefore(value, e) + return nil + } } - - replacement, replacementParent := n.left.findMax(n) - n.seqNum = replacement.seqNum - n.done = replacement.done - - return replacement.delete(replacement.seqNum, replacementParent) - } -} - -// markDone is a method used to change the 'done' flag of a btreeNode to 'true'. -func (n *btreeNode) markDone() { - n.done = true -} - -// btree is a binary search tree. -type btree struct { - root *btreeNode // the root node for the binary search tree -} - -// insert is a method used to insert sequence numbers into the binary search tree of the checkpointer. -func (t *btree) insert(seqNum int) error { - if t.root == nil { - t.root = &btreeNode{ - seqNum: seqNum, - } - return nil } - return t.root.insert(seqNum) -} -// find is a method used retrieve a pointer to a btreeNode of the given sequence number or 'false' if the sequence -// number was not found. -func (t *btree) find(seqNum int) (*btreeNode, bool) { - if t.root == nil { - return nil, false - } - return t.root.find(seqNum) + c.PushFront(value) + return nil } -// delete is a method used to remove a btreeNode with the given sequence number from the binary search tree. -func (t *btree) delete(seqNum int) error { - if t.root == nil { - return errors.New("Cannot call delete on a nil tree") - } - - if t.root.seqNum == seqNum { - fakeParent := &btreeNode{right: t.root} - if err := t.root.delete(seqNum, fakeParent); err != nil { - return err +// find is a method used to retreive the element in the doubly linked list for a given sequence number. find is +// optimized for searching oldest numbers first as it traverse the linked list starting from the beginning. +func (c *checkpointList) find(seqNum int) (*list.Element, bool) { + for e := c.Front(); e != nil; e = e.Next() { + element := e.Value.(*checkpointElement) + switch { + case element.seqNum == seqNum: + return e, true + case element.seqNum < seqNum: + break + default: + return nil, false } - t.root = fakeParent.right - return nil - } - return t.root.delete(seqNum, nil) -} - -// traverseOrdered is a (recursive) method used to apply a callback function on all of the nodes in the binary -// search tree ordered by sequence number from smallest to largest. -func (t *btree) traverseOrdered(n *btreeNode, fn func(*btreeNode)) { - if n == nil { - return - } - - t.traverseOrdered(n.left, fn) - fn(n) - t.traverseOrdered(n.right, fn) -} - -// traverseChildren is a (recursive) method used to apply a callback function on all of children nodes in the -// binary search tree. -func (t *btree) traverseChildren(n *btreeNode, fn func(*btreeNode)) { - if n == nil { - return - } - - t.traverseChildren(n.left, fn) - t.traverseChildren(n.right, fn) - fn(n) -} - -// size is a method used to determine the number of btreeNodes in the binary search tree. -func (t *btree) size() int { - var size int - if t.root != nil { - t.traverseOrdered(t.root, func(n *btreeNode) { - size++ - }) } - return size -} - -// minNode is a method used to retrieve a pointer to the btreeNode containing the smallest sequence number in the -// binary search tree. -func (t *btree) minNode() *btreeNode { - if t.root == nil { - return nil - } - node, _ := t.root.findMin(t.root) - return node -} - -// maxNode is a method used to retrieve a pointer to the btreeNode containing the largest sequence number in the -// binary search tree. -func (t *btree) maxNode() *btreeNode { - if t.root == nil { - return nil - } - node, _ := t.root.findMax(t.root) - return node -} - -// trim is a method used to trim the binary search tree so that the sequence number of all of the remaining nodes -// fall within (non-inclusive) the minSeqNum and maxSeqNum values. -func (t *btree) trim(minSeqNum, maxSeqNum *int) error { - var min, max int - if minSeqNum == nil { - min = math.MinInt64 - } else { - min = *minSeqNum - } - if maxSeqNum == nil { - max = math.MaxInt64 - } else { - max = *maxSeqNum - } - - fakeParent := &btreeNode{right: t.root} - t.traverseChildren(fakeParent, func(n *btreeNode) { - if n.left != nil && (n.left.seqNum <= min || n.left.seqNum >= max) { - n.delete(n.left.seqNum, n) - } - if n.right != nil && (n.right.seqNum <= min || n.right.seqNum >= max) { - n.delete(n.right.seqNum, n) - } - }) - t.root = fakeParent.right - - return nil + return nil, false } // checkpointOptions is a struct containing all of the configurable options for a checkpoint object. @@ -321,9 +126,9 @@ func checkpointCountCheckFreq(freq time.Duration) checkpointOptionsFn { // checkpoint() can be called periodically which may trigger a checkpoint call to KCL if the oldest sequence // numbers have been marked complete. Call startup() to enable automatic checkpointing and expiration. type checkpointer struct { - *checkpointOptions // contains all of the configuration settings for the checkpoint object - keys *btree // binary search tree used to store Kinesis record sequence numbers - keysMu sync.Mutex // mutex to make keys thread safe + *checkpointOptions // contains all of the configuration settings for the checkpoint object + list *checkpointList + listMu sync.Mutex counter uint64 // counter to track the number of messages inserted since the last checkpoint checkpointCh chan empty // channel with which to communicate / coordinate checkpointing startupOnce sync.Once // used to ensure that the startup function is called once @@ -339,7 +144,7 @@ func newCheckpointer(optionFns ...checkpointOptionsFn) *checkpointer { } return &checkpointer{ checkpointOptions: checkpointOptions, - keys: &btree{}, + list: &checkpointList{list.New()}, } } @@ -375,10 +180,8 @@ func (c *checkpointer) startup(ctx context.Context) { } // call check to obtain the checkpoint - cp, found := c.check() - if found { + if cp, found := c.check(); found { c.checkpointFn(cp) - c.trim(cp) } // reset counter and timer @@ -404,22 +207,24 @@ func (c *checkpointer) shutdown() { // size safely determines the number of nodes in the checkpointer. func (c *checkpointer) size() int { - c.keysMu.Lock() - defer c.keysMu.Unlock() + c.listMu.Lock() + defer c.listMu.Unlock() - return c.keys.size() + return c.list.Len() } // insert safely inserts a sequence number into the binary search tree. func (c *checkpointer) insert(seqNumStr string) error { - c.keysMu.Lock() - defer c.keysMu.Unlock() + c.listMu.Lock() + defer c.listMu.Unlock() - if seqNum, err := strconv.Atoi(seqNumStr); err == nil { - err := c.keys.insert(seqNum) - if err != nil { - return err - } + seqNum, err := strconv.Atoi(seqNumStr) + if err != nil { + return err + } + + if err := c.list.insert(seqNum); err != nil { + return err } atomic.AddUint64(&c.counter, 1) @@ -428,69 +233,49 @@ func (c *checkpointer) insert(seqNumStr string) error { // markDone safely marks the given sequence number as done. func (c *checkpointer) markDone(seqNumStr string) error { - c.keysMu.Lock() - defer c.keysMu.Unlock() + c.listMu.Lock() + defer c.listMu.Unlock() seqNum, err := strconv.Atoi(seqNumStr) if err != nil { return err } - node, ok := c.keys.find(seqNum) + e, ok := c.list.find(seqNum) if !ok { return errors.New("Sequence number not found") } - node.markDone() + e.Value.(*checkpointElement).done = true + if e.Prev() != nil && e.Prev().Value.(*checkpointElement).done { + c.list.Remove(e.Prev()) + } + if e.Next() != nil && e.Next().Value.(*checkpointElement).done { + c.list.Remove(e) + } + return nil } // check returns the largest sequence number marked as done where all smaller sequence numbers have // also been marked as done. func (c *checkpointer) check() (checkpoint string, found bool) { - c.keysMu.Lock() - defer c.keysMu.Unlock() - - if c.keys == nil { - return "", found - } - - var stop bool // flag to stop applying function in traverse - c.keys.traverseOrdered(c.keys.root, func(n *btreeNode) { - if stop { - return - } - - if n.left == nil && !n.done { - stop = true - return - } else if !n.done { - stop = true - return + c.listMu.Lock() + defer c.listMu.Unlock() + + if c.list != nil { + for e := c.list.Front(); e != nil; e = e.Next() { + checkpointElement := e.Value.(*checkpointElement) + if !checkpointElement.done { + break + } else { + checkpoint = strconv.Itoa(checkpointElement.seqNum) + found = true + } } - checkpoint = strconv.Itoa(n.seqNum) - found = true - }) - return checkpoint, found -} - -// trim safely simplifies the binary tree housing the sequence numbers so that oldest sequence numbers -// chained and marked as done are removed. -func (c *checkpointer) trim(checkpointStr string) error { - c.keysMu.Lock() - defer c.keysMu.Unlock() - - checkpoint, err := strconv.Atoi(checkpointStr) - if err != nil { - return err - } - - if _, ok := c.keys.find(checkpoint); !ok { - return errors.New("Sequence number not found") } - cp := int(checkpoint) - return c.keys.trim(&cp, nil) + return checkpoint, found } // checkpoint sends a signal to the checkpointCh channel to trigger a call to checkpoint (potentially). diff --git a/checkpointer_test.go b/checkpointer_test.go index 1032f03..9d25f87 100644 --- a/checkpointer_test.go +++ b/checkpointer_test.go @@ -3,9 +3,9 @@ package kinetic import ( . "github.com/smartystreets/goconvey/convey" + "container/list" "context" "fmt" - "math" "math/rand" "sort" "strconv" @@ -19,51 +19,26 @@ func init() { rand.Seed(time.Now().UnixNano()) } -// sortByKeys is an alias to a slice array of int that implements the sort interface for sorting. -type sortByKeys []int +// sortInt is an alias to a slice array of int that implements the sort interface for sorting. +type sortInt []int // Len implements the Len function for hte sort interface. -func (k sortByKeys) Len() int { +func (k sortInt) Len() int { return len(k) } // Less implements the Less function for the sort interface. -func (k sortByKeys) Less(i, j int) bool { +func (k sortInt) Less(i, j int) bool { return k[i] < k[j] } // Swap implements the Swap function for the sort interface. -func (k sortByKeys) Swap(i, j int) { +func (k sortInt) Swap(i, j int) { k[i], k[j] = k[j], k[i] } -// insertKeys is a helper function for inserting multiple keys into the binary search tree. -func insertKeys(t *btree, keys []int) bool { - for _, key := range keys { - err := t.insert(key) - if err != nil { - return false - } - } - return true -} - -// findKeys is a helper function for confirm that all of the given keys were found. -func findKeys(t *btree, keys []int) bool { - for _, key := range keys { - node, ok := t.find(key) - if !ok { - return false - } - if node == nil || key != node.seqNum { - return false - } - } - return true -} - // sortUnique is a helper method used to return a unique and sorted slice of keys. -func sortUnique(keys []int) sortByKeys { +func sortUnique(keys sortInt) sortInt { keyMap := make(map[int]bool) for _, key := range keys { if _, ok := keyMap[key]; !ok { @@ -71,7 +46,7 @@ func sortUnique(keys []int) sortByKeys { } } - var sortedKeys sortByKeys + var sortedKeys sortInt for key := range keyMap { sortedKeys = append(sortedKeys, key) } @@ -80,12 +55,12 @@ func sortUnique(keys []int) sortByKeys { return sortedKeys } -// randomKeys is helper method used to generate a random set of keys for testing. -func randomKeys(size, max int) (keys []int) { +// randomInts is helper method used to generate a random set of keys for testing. +func randomInts(size, max int) (ints []int) { for i := 0; i < size; i++ { - keys = append(keys, int(rand.Intn(max))) + ints = append(ints, int(rand.Intn(max))) } - return keys + return ints } // intPointer is a helper method used to return the address of an int @@ -93,493 +68,289 @@ func intPointer(x int) *int { return &x } -func TestCheckpointerBTreeInsert(t *testing.T) { +func TestCheckpointListInsert(t *testing.T) { var tests = []struct { scenario string - keys []int + numbers []int }{ - { - scenario: "empty tree", - keys: []int{}, - }, - { - scenario: "single key", - keys: []int{1}, - }, - { - scenario: "duplicate key", - keys: []int{1, 1}, - }, - { - scenario: "keys in ascending order", - keys: []int{1, 2, 3}, - }, - { - scenario: "keys in descendeing order", - keys: []int{4, 3, 2, 1}, - }, - { - scenario: "keys in random order", - keys: randomKeys(100, 1000), - }, + {"insert into an empty list", []int{1}}, + {"insert duplicate element", []int{1, 1}}, + {"insert in ascending order", []int{1, 2, 3, 4, 5}}, + {"insert in descending order", []int{5, 4, 3, 2, 1}}, + {"insert in random order", randomInts(100, 1000)}, } for _, test := range tests { - var b btree - sortedKeys := sortUnique(test.keys) - Convey(fmt.Sprintf("running binary tree insert test suit on scenario: [%s]", test.scenario), t, func() { - Convey("should be able to insert keys without errors", func() { - ok := insertKeys(&b, test.keys) - So(ok, ShouldBeTrue) - - if len(sortedKeys) > 0 { - Convey("should be able to find all keys without errors", func() { - ok := findKeys(&b, sortedKeys) - So(ok, ShouldBeTrue) - }) - } - - Convey("should be able to determine size correctly", func() { - So(b.size(), ShouldEqual, len(sortedKeys)) - }) - - Convey("should be able to confirm sorting order", func() { - var keys sortByKeys - b.traverseOrdered(b.root, func(n *btreeNode) { - keys = append(keys, n.seqNum) - }) - So(b.size(), ShouldEqual, len(keys)) - So(sort.IsSorted(keys), ShouldBeTrue) - }) - - Convey("should be able to find the min node correctly", func() { - minNode := b.minNode() - if minNode == nil { - So(b.root, ShouldBeNil) - } else { - So(b.minNode().seqNum, ShouldEqual, sortedKeys[0]) + c := &checkpointList{list.New()} + Convey(fmt.Sprintf("running checkpointList test suite on scenario: [%s]", test.scenario), t, func() { + Convey("inserting numbers into list", func() { + var errCount int + for _, number := range test.numbers { + if err := c.insert(number); err != nil { + errCount++ } - }) + } + So(errCount, ShouldEqual, 0) - Convey("should be able to find the max node correctly", func() { - maxNode := b.maxNode() - if maxNode == nil { - So(b.root, ShouldBeNil) - } else { - So(b.maxNode().seqNum, ShouldEqual, sortedKeys[len(sortedKeys)-1]) + Convey("confirm the length and order of the list", func() { + sortedUniqueNumbers := sortUnique(test.numbers) + var actualNumbers sortInt + for e := c.Front(); e != nil; e = e.Next() { + actualNumbers = append(actualNumbers, e.Value.(*checkpointElement).seqNum) } + So(len(actualNumbers), ShouldEqual, len(sortedUniqueNumbers)) + So(sort.IsSorted(actualNumbers), ShouldBeTrue) }) }) }) } } -func TestCheckpointerBTreeFind(t *testing.T) { +func TestCheckpointListFind(t *testing.T) { + testNumbers := []int{3, 2, 4, 5, 1} var tests = []struct { scenario string - keys []int - findKey int + numbers []int + find int shouldFind bool }{ - { - scenario: "key is on a root node", - keys: []int{2, 3, 1}, - findKey: 2, - shouldFind: true, - }, - { - scenario: "key is on a right child node", - keys: []int{2, 3, 1}, - findKey: 3, - shouldFind: true, - }, - { - scenario: "key is on a left child node", - keys: []int{2, 3, 1}, - findKey: 1, - shouldFind: true, - }, - { - scenario: "key does not exist", - keys: []int{2, 3, 1}, - findKey: 0, - shouldFind: false, - }, - { - scenario: "empty tree / key does not exist", - keys: []int{}, - findKey: 0, - shouldFind: false, - }, + {"find on an empty list", []int{}, 0, false}, + {"find on an element that does not exist", testNumbers, 0, false}, + {"find the biggeset element", testNumbers, 5, true}, + {"find the smallest element", testNumbers, 1, true}, + {"find a middle element", testNumbers, 3, true}, } - var b btree for _, test := range tests { - Convey(fmt.Sprintf("running binary tree find test suit on scenario: [%s]", test.scenario), t, func() { - ok := insertKeys(&b, test.keys) - So(ok, ShouldBeTrue) - - Convey("calling find should behave accordingly", func() { - node, ok := b.find(test.findKey) - So(ok, ShouldEqual, test.shouldFind) - if test.shouldFind { - Convey("the node return should match the key", func() { - So(node.seqNum, ShouldEqual, test.findKey) - }) + c := &checkpointList{list.New()} + Convey(fmt.Sprintf("running find test suite on scenario: [%s]", test.scenario), t, func() { + Convey("inserting numbers into list", func() { + var errCount int + for _, number := range test.numbers { + if err := c.insert(number); err != nil { + errCount++ + } } + So(errCount, ShouldEqual, 0) + + Convey("find should behave as expected", func() { + el, found := c.find(test.find) + So(found, ShouldEqual, test.shouldFind) + if found { + So(el.Value.(*checkpointElement).seqNum, ShouldEqual, test.find) + } + }) }) }) } } -func TestCheckpointerBTreeDelete(t *testing.T) { +func TestCheckpointerInsertAndSize(t *testing.T) { var tests = []struct { - scenario string - keys []int - deleteKey int - shouldError bool + scenario string + numbers []int }{ - { - scenario: "delete on empty tree", - keys: []int{}, - deleteKey: 1, - shouldError: true, - }, - { - scenario: "delete a key that doesn't exist", - keys: []int{2, 3, 1}, - deleteKey: 0, - shouldError: true, - }, - { - scenario: "delete a node with no children", - keys: []int{5, 2, 18, -4, 3}, - deleteKey: -4, - shouldError: false, - }, - { - scenario: "delete a node with one child on the left", - keys: []int{5, 2, 18, -4, 3, 16, 15, 17}, - deleteKey: 18, - shouldError: false, - }, - { - scenario: "delete a node with one child on the right", - keys: []int{5, 2, 18, -4, 3, 21, 19, 25}, - deleteKey: 18, - shouldError: false, - }, - { - scenario: "delete a node with two children", - keys: []int{5, 2, 19, 16, 25}, - deleteKey: 19, - shouldError: false, - }, - { - scenario: "delete a root node with no children", - keys: []int{1}, - deleteKey: 1, - shouldError: false, - }, - { - scenario: "delete a root node with one child on the left", - keys: []int{2, 1}, - deleteKey: 2, - shouldError: false, - }, - { - scenario: "delete a root node with one child on the right", - keys: []int{2, 3}, - deleteKey: 2, - shouldError: false, - }, - { - scenario: "delete a root node with two children", - keys: []int{2, 1, 3}, - deleteKey: 3, - shouldError: false, - }, + {"insert into an empty list", []int{1}}, + {"insert duplicate elements", []int{1, 1}}, + {"insert into an empty list", []int{1}}, + {"insert in ascending order", []int{11, 22, 33, 44, 55}}, + {"insert in descending order", []int{55, 44, 33, 22, 11}}, + {"insert in random order", randomInts(100, 1000)}, } for _, test := range tests { - var b btree - sortedKeys := sortUnique(test.keys) - Convey(fmt.Sprintf("running binary tree delete test suit on scenario: [%s]", test.scenario), t, func() { - ok := insertKeys(&b, test.keys) - So(ok, ShouldBeTrue) - - Convey("should be able to call delete and respond correctly", func() { - err := b.delete(test.deleteKey) - if test.shouldError { - So(err, ShouldNotBeNil) - } else { - So(err, ShouldBeNil) - } - - if test.shouldError { - if len(sortedKeys) > 0 { - Convey("should be able to find all keys without errors", func() { - ok := findKeys(&b, sortedKeys) - So(ok, ShouldBeTrue) - }) - } - - Convey("should not be able to find delete key", func() { - _, ok := b.find(test.deleteKey) - So(ok, ShouldBeFalse) - }) - - Convey("should be able to determine size correctly", func() { - So(b.size(), ShouldEqual, len(sortedKeys)) - }) - - Convey("should be able to confirm sorting order", func() { - var keys sortByKeys - b.traverseOrdered(b.root, func(n *btreeNode) { - keys = append(keys, n.seqNum) - }) - So(b.size(), ShouldEqual, len(keys)) - So(sort.IsSorted(keys), ShouldBeTrue) - }) - - Convey("should be able to find the min node correctly", func() { - minNode := b.minNode() - if minNode == nil { - So(b.root, ShouldBeNil) - } else { - So(b.minNode().seqNum, ShouldEqual, sortedKeys[0]) - } - }) - - Convey("should be able to find the max node correctly", func() { - maxNode := b.maxNode() - if maxNode == nil { - So(b.root, ShouldBeNil) - } else { - So(b.maxNode().seqNum, ShouldEqual, sortedKeys[len(sortedKeys)-1]) - } - }) - - } else { - // determine the new key set (sorted and unique) after removing the first element - var newKeys sortByKeys - for _, k := range test.keys { - if k != test.deleteKey { - newKeys = append(newKeys, k) - } + c := newCheckpointer() + Convey(fmt.Sprintf("running insert and size test suite on scenario: [%s]", test.scenario), t, func() { + Convey("inserting numbers into list", func() { + var errCount int + for _, number := range test.numbers { + if err := c.insert(strconv.Itoa(number)); err != nil { + errCount++ } - sort.Sort(newKeys) + } + So(errCount, ShouldEqual, 0) - if len(newKeys) > 0 { - Convey("should be able to find remaining keys after a delete", func() { - ok := findKeys(&b, newKeys) - So(ok, ShouldBeTrue) - }) + Convey("confirm the size and order of the list", func() { + sortedUniqueNumbers := sortUnique(test.numbers) + var actualNumbers sortInt + for e := c.list.Front(); e != nil; e = e.Next() { + actualNumbers = append(actualNumbers, e.Value.(*checkpointElement).seqNum) } - - Convey("should be able to determine size correctly after a delete", func() { - So(b.size(), ShouldEqual, len(newKeys)) - }) - - Convey("should be able to confirm sorting order after a delete", func() { - var keys sortByKeys - b.traverseOrdered(b.root, func(n *btreeNode) { - keys = append(keys, n.seqNum) - }) - So(b.size(), ShouldEqual, len(keys)) - So(sort.IsSorted(keys), ShouldBeTrue) - }) - - Convey("should be able to find the min node correctly after a delete", func() { - minNode := b.minNode() - if minNode == nil { - So(b.root, ShouldBeNil) - } else { - So(b.minNode().seqNum, ShouldEqual, newKeys[0]) - } - }) - - Convey("should be able to find the max node correctly after a delete", func() { - maxNode := b.maxNode() - if maxNode == nil { - So(b.root, ShouldBeNil) - } else { - So(b.maxNode().seqNum, ShouldEqual, newKeys[len(newKeys)-1]) - } - }) - } + So(c.size(), ShouldEqual, len(sortedUniqueNumbers)) + So(len(actualNumbers), ShouldEqual, len(sortedUniqueNumbers)) + So(sort.IsSorted(actualNumbers), ShouldBeTrue) + }) }) }) } } -func TestCheckpointerBTreeTrim(t *testing.T) { - testKeys := []int{50, 20, 60, 40, 0, 10, 30} +func TestCheckpointerMarkDoneAndCheck(t *testing.T) { + testEmpty := []int{} + testNumbers := []int{5, 2, 6, 4, 0, 1, 3} var tests = []struct { - scenario string - keys []int - trimMin *int - trimMax *int + scenario string + numbers []int + doneNumbers []int + shouldErrorOnMarkDone bool + expectedSize *int + shouldFindCheckpoint bool + expectedCheckpoint *int }{ { - scenario: "trim on empty tree", - keys: []int{}, - trimMin: intPointer(0), - trimMax: intPointer(10), - }, - { - scenario: "trim left - none", - keys: testKeys, - trimMin: intPointer(-1), - trimMax: nil, - }, - { - scenario: "trim left - some", - keys: testKeys, - trimMin: intPointer(20), - trimMax: nil, - }, - { - scenario: "trim left - all", - keys: testKeys, - trimMin: intPointer(60), - trimMax: nil, - }, - { - scenario: "trim right - none", - keys: testKeys, - trimMin: nil, - trimMax: intPointer(70), - }, - { - scenario: "trim right - some", - keys: testKeys, - trimMin: nil, - trimMax: intPointer(40), - }, - { - scenario: "trim right - all", - keys: testKeys, - trimMin: nil, - trimMax: intPointer(0), - }, - { - scenario: "trim - none", - keys: testKeys, - trimMin: intPointer(0), - trimMax: intPointer(60), - }, - { - scenario: "trim - some", - keys: testKeys, - trimMin: intPointer(10), - trimMax: intPointer(50), - }, - { - scenario: "trim - all", - keys: testKeys, - trimMin: intPointer(1), - trimMax: intPointer(9), + scenario: "markDone on an empty list", + numbers: testEmpty, + doneNumbers: []int{1}, + shouldErrorOnMarkDone: true, + expectedSize: nil, + shouldFindCheckpoint: false, + expectedCheckpoint: nil, + }, + { + scenario: "markDone on non existent number", + numbers: testNumbers, + doneNumbers: []int{-1}, + shouldErrorOnMarkDone: true, + expectedSize: nil, + shouldFindCheckpoint: false, + expectedCheckpoint: nil, + }, + { + scenario: "check when no numbers are inserted", + numbers: testEmpty, + doneNumbers: testEmpty, + shouldErrorOnMarkDone: false, + expectedSize: intPointer(len(testEmpty)), + shouldFindCheckpoint: false, + expectedCheckpoint: nil, + }, + { + scenario: "check when no numbers are marked done", + numbers: testNumbers, + doneNumbers: testEmpty, + shouldErrorOnMarkDone: false, + expectedSize: intPointer(len(testNumbers)), + shouldFindCheckpoint: false, + expectedCheckpoint: nil, + }, + { + scenario: "markDone on an existing number (not smallest number) and check", + numbers: testNumbers, + doneNumbers: []int{1}, + shouldErrorOnMarkDone: false, + expectedSize: intPointer(len(testNumbers)), + shouldFindCheckpoint: false, + expectedCheckpoint: nil, + }, + { + scenario: "markDone twice on an existing number (not smallest number) and check ", + numbers: testNumbers, + doneNumbers: []int{1, 1}, + shouldErrorOnMarkDone: false, + expectedSize: intPointer(len(testNumbers)), + shouldFindCheckpoint: false, + expectedCheckpoint: nil, + }, + { + scenario: "markDone on two sequential numbers (not smallest number, ascending) and check", + numbers: testNumbers, + doneNumbers: []int{1, 2}, + shouldErrorOnMarkDone: false, + expectedSize: intPointer(len(testNumbers) - 1), + shouldFindCheckpoint: false, + expectedCheckpoint: nil, + }, + { + scenario: "markDone on two sequential numbers (not smallest number, descending) and check", + numbers: testNumbers, + doneNumbers: []int{2, 1}, + shouldErrorOnMarkDone: false, + expectedSize: intPointer(len(testNumbers) - 1), + shouldFindCheckpoint: false, + expectedCheckpoint: nil, + }, + { + scenario: "markDone on three sequential numbers (not smallest number) and check", + numbers: testNumbers, + doneNumbers: []int{3, 1, 2}, + shouldErrorOnMarkDone: false, + expectedSize: intPointer(len(testNumbers) - 2), + shouldFindCheckpoint: false, + expectedCheckpoint: nil, + }, + { + scenario: "check when the smallest number is not marked done", + numbers: testNumbers, + doneNumbers: []int{1, 2, 3, 4, 5, 6}, + shouldErrorOnMarkDone: false, + expectedSize: intPointer(2), + shouldFindCheckpoint: false, + expectedCheckpoint: nil, + }, + { + scenario: "check when a chain (including the smallest number) is marked done", + numbers: testNumbers, + doneNumbers: []int{2, 1, 3, 0}, + shouldErrorOnMarkDone: false, + expectedSize: intPointer(4), + shouldFindCheckpoint: true, + expectedCheckpoint: intPointer(3), + }, + { + scenario: "markDone on all numbers and check", + numbers: testNumbers, + doneNumbers: testNumbers, + shouldErrorOnMarkDone: false, + expectedSize: intPointer(1), + shouldFindCheckpoint: true, + expectedCheckpoint: intPointer(6), }, } for _, test := range tests { - var b btree - sortedKeys := sortUnique(test.keys) - Convey(fmt.Sprintf("running binary tree trim test suite on [%s]", test.scenario), t, func() { - Convey("should be able to insert keys without errors", func() { - ok := insertKeys(&b, test.keys) - So(ok, ShouldBeTrue) - - if len(sortedKeys) > 0 { - Convey("should be able to find all keys without errors", func() { - ok := findKeys(&b, sortedKeys) - So(ok, ShouldBeTrue) - }) - } - - uniqueKeys := sortUnique(test.keys) - var remainingKeys sortByKeys - var trimmedKeys sortByKeys - var min, max int - if test.trimMin == nil { - min = math.MinInt64 - } else { - min = *test.trimMin - } - if test.trimMax == nil { - max = math.MaxInt64 - } else { - max = *test.trimMax - } - - for _, key := range uniqueKeys { - if key > min && key < max { - remainingKeys = append(remainingKeys, key) - } else { - trimmedKeys = append(trimmedKeys, key) + c := newCheckpointer() + Convey(fmt.Sprintf("running markDone and check test suite on scenario: [%s]", test.scenario), t, func() { + Convey("inserting numbers into list", func() { + var errCount int + for _, number := range test.numbers { + if err := c.insert(strconv.Itoa(number)); err != nil { + errCount++ } } + So(errCount, ShouldEqual, 0) - Convey("calling trim on tree", func() { - err := b.trim(test.trimMin, test.trimMax) - So(err, ShouldBeNil) - - if len(remainingKeys) > 0 { - Convey("should be able to find all remaining keys", func() { - So(findKeys(&b, remainingKeys), ShouldBeTrue) - }) + Convey("mark numbers done should behave as expected", func() { + var doneErrCount int + for _, doneNumber := range test.doneNumbers { + if err := c.markDone(strconv.Itoa(doneNumber)); err != nil { + doneErrCount++ + } } - - if len(trimmedKeys) > 0 { - Convey("should not be able to find any trimmed keys", func() { - for _, key := range trimmedKeys { - _, ok := b.find(key) - So(ok, ShouldBeFalse) - } - }) + if test.shouldErrorOnMarkDone { + So(doneErrCount, ShouldBeGreaterThan, 0) + return } + + So(doneErrCount, ShouldEqual, 0) + So(c.size(), ShouldEqual, *test.expectedSize) + Convey("call check should behave as expected", func() { + cp, found := c.check() + if test.shouldFindCheckpoint { + So(found, ShouldBeTrue) + So(cp, ShouldEqual, strconv.Itoa(*test.expectedCheckpoint)) + } else { + So(found, ShouldBeFalse) + + } + }) }) }) }) } } -func TestCheckpointerOffNominal(t *testing.T) { - Convey("off nominal operations on an empty tree", t, func() { - cp := newCheckpointer() - Convey("markDone on an empty tree should error", func() { - err := cp.markDone("1") - So(err, ShouldNotBeNil) - }) - - Convey("trim on an empty tree should error", func() { - err := cp.trim("1") - So(err, ShouldNotBeNil) - - }) - }) - - Convey("off nominal operations on a non-empty tree", t, func() { - cp := newCheckpointer() - cp.insert("2") - cp.insert("3") - cp.insert("1") - Convey("markDone on a key that does not exist should error", func() { - err := cp.markDone("0") - So(err, ShouldNotBeNil) - }) - - Convey("trim on a key that does not exist should error", func() { - err := cp.trim("0") - So(err, ShouldNotBeNil) - - }) - }) -} - -func TestCheckpointerConfigs(t *testing.T) { - testKeys := []int{1, -2, 3, -4, 5} +func TestCheckpointerAutoCheckpointing(t *testing.T) { + testNumbers := []int{1, -2, 3, -4, 5} Convey("test auto checkpoint using autoCheckpointCount", t, func() { autoCheckpointCount := 5 Convey("instantiating a new checkpoint with options", func() { @@ -601,7 +372,7 @@ func TestCheckpointerConfigs(t *testing.T) { Convey("inserting and marking all keys as done", func() { var failedCount int - for _, testKey := range testKeys { + for _, testKey := range testNumbers { key := strconv.Itoa(testKey) if err := cp.insert(key); err != nil { failedCount++ @@ -615,9 +386,6 @@ func TestCheckpointerConfigs(t *testing.T) { <-time.After(time.Millisecond) Convey("confirming that checkpoint and trim was called", func() { So(atomic.LoadUint64(&checkpointFnCalled), ShouldBeGreaterThan, 0) - cp.keysMu.Lock() - So(cp.keys.size(), ShouldEqual, 0) - cp.keysMu.Unlock() }) }) }) @@ -643,7 +411,7 @@ func TestCheckpointerConfigs(t *testing.T) { Convey("inserting and marking all keys as done", func() { var failedCount int - for _, testKey := range testKeys { + for _, testKey := range testNumbers { key := strconv.Itoa(testKey) if err := cp.insert(key); err != nil { failedCount++ @@ -657,9 +425,6 @@ func TestCheckpointerConfigs(t *testing.T) { <-time.After(time.Duration(2) * autoCheckpointFreq) Convey("confirming that checkpoint and trim was called", func() { So(atomic.LoadUint64(&checkpointFnCalled), ShouldBeGreaterThan, 0) - cp.keysMu.Lock() - So(cp.keys.size(), ShouldEqual, 0) - cp.keysMu.Unlock() }) }) }) @@ -701,160 +466,16 @@ func TestCheckpointerConfigs(t *testing.T) { <-time.After(500 * time.Millisecond) Convey("confirming that checkpointFn was called", func() { So(atomic.LoadUint64(&checkpointFnCalled), ShouldBeGreaterThan, 0) - cp.keysMu.Lock() - So(cp.keys.size(), ShouldEqual, 0) - cp.keysMu.Unlock() }) }) }) }) } -func TestCheckpointerNominal(t *testing.T) { - testKeys := []int{5, 2, 6, 4, 0, 1, 3} - var tests = []struct { - scenario string - keys []int - doneKeys []int - shouldBeOK bool - }{ - { - scenario: "checkpoint/trim when there are no keys inserted", - keys: []int{}, - doneKeys: []int{}, - shouldBeOK: false, - }, - { - scenario: "checkpoint/trim when no keys are marked done", - keys: testKeys, - doneKeys: []int{}, - shouldBeOK: false, - }, - { - scenario: "checkpoint/trim when minimum node is not marked done", - keys: testKeys, - doneKeys: []int{1, 2, 3}, - shouldBeOK: false, - }, - { - scenario: "checkpoint/trim when minimum node is marked done", - keys: testKeys, - doneKeys: []int{2, 3, 1, 0}, - shouldBeOK: true, - }, - { - scenario: "checkpoint/trim when all keys are marked done", - keys: testKeys, - doneKeys: testKeys, - shouldBeOK: true, - }, - { - scenario: "checkpoint/trim only one key exists", - keys: []int{1}, - doneKeys: []int{1}, - shouldBeOK: true, - }, - } - - for _, test := range tests { - cp := newCheckpointer() - Convey(fmt.Sprintf("running checkpoint test suit on scenario: [%s]", test.scenario), t, func() { - Convey("inserting the test keys", func() { - var errCount int - for _, key := range test.keys { - if err := cp.insert(strconv.Itoa(key)); err != nil { - errCount++ - } - } - So(errCount, ShouldEqual, 0) - - Convey("confirming keys were inserted correctly", func() { - So(findKeys(cp.keys, test.keys), ShouldBeTrue) - }) - - uniqueDoneKeys := sortUnique(test.doneKeys) - Convey("marking test keys as done", func() { - if len(test.doneKeys) > 0 { - var errCount int - for _, doneKey := range test.doneKeys { - if err := cp.markDone(strconv.Itoa(doneKey)); err != nil { - errCount++ - } - } - So(errCount, ShouldEqual, 0) - - Convey("confirm that only the test keys were marked done", func() { - var markedDone sortByKeys - cp.keys.traverseOrdered(cp.keys.root, func(n *btreeNode) { - if n.done { - markedDone = append(markedDone, n.seqNum) - } - }) - So(len(markedDone), ShouldEqual, len(uniqueDoneKeys)) - - for _, doneKey := range test.doneKeys { - node, ok := cp.keys.find(doneKey) - So(node.done, ShouldBeTrue) - So(ok, ShouldBeTrue) - } - }) - } - - Convey("calling checkpoint should yeild the expected behavior", func() { - checkpoint, ok := cp.check() - So(checkpoint, ShouldNotBeNil) - So(ok, ShouldEqual, test.shouldBeOK) - //So(cp.lastCheckpoint, ShouldEqual, checkpoint) - - if test.shouldBeOK { - uniqueKeys := sortUnique(test.keys) - Convey("and should also result in the correct checkpoint key", func() { - var expectedCheckpoint int - for i := 0; i < len(uniqueDoneKeys); i++ { - if uniqueKeys[i] == uniqueDoneKeys[i] { - expectedCheckpoint = uniqueDoneKeys[i] - } else { - break - } - } - So(checkpoint, ShouldEqual, strconv.Itoa(expectedCheckpoint)) - }) - - Convey("calling trim using the checkpoint key", func() { - err := cp.trim(checkpoint) - So(err, ShouldBeNil) - - Convey("should not be able to find any keys marked done", func() { - var foundDone bool - var stop bool - cp.keys.traverseOrdered(cp.keys.root, func(n *btreeNode) { - if stop { - return - } - if n.done { - foundDone = true - stop = true - } - }) - So(foundDone, ShouldBeFalse) - }) - - if len(uniqueKeys)-len(uniqueDoneKeys) > 0 { - Convey("should still be able to find remaining keys", func() { - remainingKeys := uniqueKeys[len(uniqueDoneKeys):] - for _, key := range remainingKeys { - node, ok := cp.keys.find(key) - So(ok, ShouldBeTrue) - So(node, ShouldNotBeNil) - So(node.done, ShouldBeFalse) - } - }) - } - }) - } - }) - }) - }) - }) - } +func TestCheckpointerOffNominal(t *testing.T) { + Convey("insert a bogus number into the list", t, func() { + c := newCheckpointer() + err := c.insert("bogus data") + So(err, ShouldNotBeNil) + }) } From 739e25222f8b812f2ea64a9e0b5a0dde9c0af987 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 23 Aug 2017 08:07:59 -0500 Subject: [PATCH 084/109] Refactored the signature of the function option methods to operate on the main (public) objects instead of the respective (private) option configuration structs. --- checkpointer.go | 30 ++++++++++++++++-------------- consumer.go | 35 ++++++++++++++++++----------------- firehose_writer.go | 37 ++++++++++++++++++++----------------- kcl_reader.go | 38 +++++++++++++++++++------------------- kinesis_reader.go | 45 ++++++++++++++++++++++++--------------------- kinesis_writer.go | 37 ++++++++++++++++++++----------------- producer.go | 45 +++++++++++++++++++++++---------------------- 7 files changed, 140 insertions(+), 127 deletions(-) diff --git a/checkpointer.go b/checkpointer.go index d9c0303..3d65e86 100644 --- a/checkpointer.go +++ b/checkpointer.go @@ -10,7 +10,7 @@ import ( "time" ) -// checkpointElement is the struct used to store checkpointing information. +// checkpointElement is the struct used to store checkpointing information in the doubly linked list. type checkpointElement struct { seqNum int done bool @@ -46,7 +46,7 @@ func (c *checkpointList) insert(seqNum int) error { return nil } -// find is a method used to retreive the element in the doubly linked list for a given sequence number. find is +// find is a method used to retrieve the element in the doubly linked list for a given sequence number. find is // optimized for searching oldest numbers first as it traverse the linked list starting from the beginning. func (c *checkpointList) find(seqNum int) (*list.Element, bool) { for e := c.Front(); e != nil; e = e.Next() { @@ -83,11 +83,11 @@ func defaultCheckpointOptions() *checkpointOptions { // checkpointOptionsFn is a function signature used to define function options for configuring all of the // configurable options of a checkpoint object. -type checkpointOptionsFn func(*checkpointOptions) error +type checkpointOptionsFn func(*checkpointer) error // checkpointAutoCheckpointCount is a functional option method for configuring the checkpoint's auto checkpoint count. func checkpointAutoCheckpointCount(count int) checkpointOptionsFn { - return func(o *checkpointOptions) error { + return func(o *checkpointer) error { o.autoCheckpointCount = count return nil } @@ -96,7 +96,7 @@ func checkpointAutoCheckpointCount(count int) checkpointOptionsFn { // checkpointAutoCheckpointFreq is a functional option method for configuring the checkpoint's auto checkpoint // frequency. func checkpointAutoCheckpointFreq(freq time.Duration) checkpointOptionsFn { - return func(o *checkpointOptions) error { + return func(o *checkpointer) error { o.autoCheckpointFreq = freq return nil } @@ -104,7 +104,7 @@ func checkpointAutoCheckpointFreq(freq time.Duration) checkpointOptionsFn { // checkpointCheckpointFn is a functional option method for configuring the checkpoint's checkpoint callback function. func checkpointCheckpointFn(fn func(string) error) checkpointOptionsFn { - return func(o *checkpointOptions) error { + return func(o *checkpointer) error { o.checkpointFn = fn return nil } @@ -112,7 +112,7 @@ func checkpointCheckpointFn(fn func(string) error) checkpointOptionsFn { // checkpointCountCheckFreq is a functional option method for configuring the checkpoint's count check frequency. func checkpointCountCheckFreq(freq time.Duration) checkpointOptionsFn { - return func(o *checkpointOptions) error { + return func(o *checkpointer) error { o.countCheckFreq = freq return nil } @@ -138,14 +138,15 @@ type checkpointer struct { // newCheckpoint instantiates a new checkpoint object with default configuration settings unless the function option // methods are provided to change the default values. func newCheckpointer(optionFns ...checkpointOptionsFn) *checkpointer { - checkpointOptions := defaultCheckpointOptions() - for _, optionFn := range optionFns { - optionFn(checkpointOptions) - } - return &checkpointer{ - checkpointOptions: checkpointOptions, + checkpointer := &checkpointer{ + checkpointOptions: defaultCheckpointOptions(), list: &checkpointList{list.New()}, } + for _, optionFn := range optionFns { + optionFn(checkpointer) + } + + return checkpointer } // startup is a method used to enable automatic checkpointing. @@ -231,7 +232,8 @@ func (c *checkpointer) insert(seqNumStr string) error { return nil } -// markDone safely marks the given sequence number as done. +// markDone safely marks the given sequence number as done and attempts to remove it's previous element if the +// previous element is also marked done or attempts to remove itself if it's next element is also marked done func (c *checkpointer) markDone(seqNumStr string) error { c.listMu.Lock() defer c.listMu.Unlock() diff --git a/consumer.go b/consumer.go index 6a5d048..cfb7e9f 100644 --- a/consumer.go +++ b/consumer.go @@ -29,11 +29,11 @@ func defaultConsumerOptions() *consumerOptions { } // ConsumerOptionsFn is a method signature for defining functional option methods for configuring the Consumer. -type ConsumerOptionsFn func(*consumerOptions) error +type ConsumerOptionsFn func(*Consumer) error // ConsumerReader is a functional option method for configuring the consumer's stream reader. func ConsumerReader(r StreamReader) ConsumerOptionsFn { - return func(o *consumerOptions) error { + return func(o *Consumer) error { o.reader = r return nil } @@ -41,7 +41,7 @@ func ConsumerReader(r StreamReader) ConsumerOptionsFn { // ConsumerQueueDepth is a functional option method for configuring the consumer's queueDepth. func ConsumerQueueDepth(depth int) ConsumerOptionsFn { - return func(o *consumerOptions) error { + return func(o *Consumer) error { if depth > 0 { o.queueDepth = depth return nil @@ -52,7 +52,7 @@ func ConsumerQueueDepth(depth int) ConsumerOptionsFn { // ConsumerConcurrency is a functional option method for configuring the consumer's concurrency. func ConsumerConcurrency(count int) ConsumerOptionsFn { - return func(o *consumerOptions) error { + return func(o *Consumer) error { if count > 0 { o.concurrency = count return nil @@ -63,7 +63,7 @@ func ConsumerConcurrency(count int) ConsumerOptionsFn { // ConsumerLogLevel is a functional option method for configuring the consumer's log level. func ConsumerLogLevel(ll aws.LogLevelType) ConsumerOptionsFn { - return func(o *consumerOptions) error { + return func(o *Consumer) error { o.logLevel = ll & 0xffff0000 return nil } @@ -71,7 +71,7 @@ func ConsumerLogLevel(ll aws.LogLevelType) ConsumerOptionsFn { // ConsumerStats is a functional option method for configuring the consumer's stats collector. func ConsumerStats(sc ConsumerStatsCollector) ConsumerOptionsFn { - return func(o *consumerOptions) error { + return func(o *Consumer) error { o.Stats = sc return nil } @@ -93,24 +93,25 @@ type Consumer struct { // NewConsumer creates a new Consumer object for retrieving and listening to message(s) on a StreamReader. func NewConsumer(c *aws.Config, stream string, shard string, optionFns ...ConsumerOptionsFn) (*Consumer, error) { - consumerOptions := defaultConsumerOptions() + consumer := &Consumer{consumerOptions: defaultConsumerOptions()} for _, optionFn := range optionFns { - optionFn(consumerOptions) + optionFn(consumer) } - if consumerOptions.reader == nil { + + if consumer.reader == nil { r, err := NewKinesisReader(c, stream, shard) if err != nil { return nil, err } - consumerOptions.reader = r + consumer.reader = r + } + + consumer.LogHelper = &LogHelper{ + LogLevel: consumer.logLevel, + Logger: c.Logger, } - return &Consumer{ - consumerOptions: consumerOptions, - LogHelper: &LogHelper{ - LogLevel: consumerOptions.logLevel, - Logger: c.Logger, - }, - }, nil + + return consumer, nil } // startConsuming will initialize the message channel and set consuming to true if there is not already another consume diff --git a/firehose_writer.go b/firehose_writer.go index 0ef38f9..cd2584d 100644 --- a/firehose_writer.go +++ b/firehose_writer.go @@ -39,12 +39,12 @@ func defaultFirehoseWriterOptions() *firehoseWriterOptions { // FirehoseWriterOptionsFn is a method signature for defining functional option methods for configuring // the FirehoseWriter. -type FirehoseWriterOptionsFn func(*firehoseWriterOptions) error +type FirehoseWriterOptionsFn func(*FirehoseWriter) error // FirehoseWriterMsgCountRateLimit is a functional option method for configuring the FirehoseWriter's // message count rate limit. func FirehoseWriterMsgCountRateLimit(limit int) FirehoseWriterOptionsFn { - return func(o *firehoseWriterOptions) error { + return func(o *FirehoseWriter) error { if limit > 0 && limit <= firehoseMsgCountRateLimit { o.msgCountRateLimit = limit return nil @@ -56,7 +56,7 @@ func FirehoseWriterMsgCountRateLimit(limit int) FirehoseWriterOptionsFn { // FirehoseWriterMsgSizeRateLimit is a functional option method for configuring the FirehoseWriter's // messsage size rate limit. func FirehoseWriterMsgSizeRateLimit(limit int) FirehoseWriterOptionsFn { - return func(o *firehoseWriterOptions) error { + return func(o *FirehoseWriter) error { if limit > 0 && limit <= firehoseMsgSizeRateLimit { o.msgSizeRateLimit = limit return nil @@ -68,7 +68,7 @@ func FirehoseWriterMsgSizeRateLimit(limit int) FirehoseWriterOptionsFn { // FirehoseWriterThroughputMultiplier is a functional option method for configuring the FirehoseWriter's // throughput multiplier. func FirehoseWriterThroughputMultiplier(multiplier int) FirehoseWriterOptionsFn { - return func(o *firehoseWriterOptions) error { + return func(o *FirehoseWriter) error { if multiplier > 0 { o.throughputMultiplier = multiplier return nil @@ -79,7 +79,7 @@ func FirehoseWriterThroughputMultiplier(multiplier int) FirehoseWriterOptionsFn // FirehoseWriterLogLevel is a functional option method for configuring the FirehoseWriter's log level. func FirehoseWriterLogLevel(ll aws.LogLevelType) FirehoseWriterOptionsFn { - return func(o *firehoseWriterOptions) error { + return func(o *FirehoseWriter) error { o.logLevel = ll & 0xffff0000 return nil } @@ -87,7 +87,7 @@ func FirehoseWriterLogLevel(ll aws.LogLevelType) FirehoseWriterOptionsFn { // FirehoseWriterStats is a functional option method for configuring the FirehoseWriter's stats collector. func FirehoseWriterStats(sc ProducerStatsCollector) FirehoseWriterOptionsFn { - return func(o *firehoseWriterOptions) error { + return func(o *FirehoseWriter) error { o.Stats = sc return nil } @@ -103,23 +103,26 @@ type FirehoseWriter struct { // NewFirehoseWriter creates a new stream writer to write records to a Kinesis. func NewFirehoseWriter(c *aws.Config, stream string, optionFns ...FirehoseWriterOptionsFn) (*FirehoseWriter, error) { - firehoseWriterOptions := defaultFirehoseWriterOptions() - for _, optionFn := range optionFns { - optionFn(firehoseWriterOptions) - } sess, err := session.NewSession(c) if err != nil { return nil, err } - return &FirehoseWriter{ + + firehoseWriter := &FirehoseWriter{ + firehoseWriterOptions: defaultFirehoseWriterOptions(), stream: stream, client: firehose.New(sess), - firehoseWriterOptions: firehoseWriterOptions, - LogHelper: &LogHelper{ - LogLevel: firehoseWriterOptions.logLevel, - Logger: c.Logger, - }, - }, nil + } + for _, optionFn := range optionFns { + optionFn(firehoseWriter) + } + + firehoseWriter.LogHelper = &LogHelper{ + LogLevel: firehoseWriter.logLevel, + Logger: c.Logger, + } + + return firehoseWriter, nil } // PutRecords sends a batch of records to Firehose and returns a list of records that need to be retried. diff --git a/kcl_reader.go b/kcl_reader.go index 960cff1..37a6e37 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -43,11 +43,11 @@ func defaultKclReaderOptions() *kclReaderOptions { } // KclReaderOptionsFn is a method signature for defining functional option methods for configuring the KclReader. -type KclReaderOptionsFn func(*kclReaderOptions) error +type KclReaderOptionsFn func(*KclReader) error // kclReaderBatchSize is a functional option method for configuring the KclReader's batch size func kclReaderBatchSize(size int) KclReaderOptionsFn { - return func(o *kclReaderOptions) error { + return func(o *KclReader) error { if size > 0 && size <= kclReaderMaxBatchSize { o.batchSize = size return nil @@ -58,7 +58,7 @@ func kclReaderBatchSize(size int) KclReaderOptionsFn { // KclReaderAutoCheckpointCount is a functional option method for configuring the KclReader's checkpoint count func KclReaderAutoCheckpointCount(count int) KclReaderOptionsFn { - return func(o *kclReaderOptions) error { + return func(o *KclReader) error { o.autoCheckpointCount = count return nil } @@ -66,7 +66,7 @@ func KclReaderAutoCheckpointCount(count int) KclReaderOptionsFn { // KclReaderAutoCheckpointFreq is a functional option method for configuring the KclReader's checkpoint frequency func KclReaderAutoCheckpointFreq(freq time.Duration) KclReaderOptionsFn { - return func(o *kclReaderOptions) error { + return func(o *KclReader) error { o.autoCheckpointFreq = freq return nil } @@ -75,7 +75,7 @@ func KclReaderAutoCheckpointFreq(freq time.Duration) KclReaderOptionsFn { // KclReaderUpdateCheckpointSizeFreq is a functional option method for configuring the KclReader's // update checkpoint size stats frequency func KclReaderUpdateCheckpointSizeFreq(freq time.Duration) KclReaderOptionsFn { - return func(o *kclReaderOptions) error { + return func(o *KclReader) error { o.updateCheckpointSizeFreq = freq return nil } @@ -84,7 +84,7 @@ func KclReaderUpdateCheckpointSizeFreq(freq time.Duration) KclReaderOptionsFn { // KclReaderOnInitCallbackFn is a functional option method for configuring the KclReader's // onInitCallbackFn. func KclReaderOnInitCallbackFn(fn func() error) KclReaderOptionsFn { - return func(o *kclReaderOptions) error { + return func(o *KclReader) error { o.onInitCallbackFn = fn return nil } @@ -93,7 +93,7 @@ func KclReaderOnInitCallbackFn(fn func() error) KclReaderOptionsFn { // KclReaderOnCheckpointCallbackFn is a functional option method for configuring the KclReader's // onCheckpointCallbackFn. func KclReaderOnCheckpointCallbackFn(fn func(seqNum string, err string) error) KclReaderOptionsFn { - return func(o *kclReaderOptions) error { + return func(o *KclReader) error { o.onCheckpointCallbackFn = fn return nil } @@ -102,7 +102,7 @@ func KclReaderOnCheckpointCallbackFn(fn func(seqNum string, err string) error) K // KclReaderOnShutdownCallbackFn is a functional option method for configuring the KclReader's // onShutdownCallbackFn. func KclReaderOnShutdownCallbackFn(fn func() error) KclReaderOptionsFn { - return func(o *kclReaderOptions) error { + return func(o *KclReader) error { o.onShutdownCallbackFn = fn return nil } @@ -110,7 +110,7 @@ func KclReaderOnShutdownCallbackFn(fn func() error) KclReaderOptionsFn { // KclReaderLogLevel is a functional option method for configuring the KclReader's log level. func KclReaderLogLevel(ll aws.LogLevelType) KclReaderOptionsFn { - return func(o *kclReaderOptions) error { + return func(o *KclReader) error { o.logLevel = ll return nil } @@ -118,7 +118,7 @@ func KclReaderLogLevel(ll aws.LogLevelType) KclReaderOptionsFn { // KclReaderStats is a functional option method for configuring the KclReader's stats collector. func KclReaderStats(sc ConsumerStatsCollector) KclReaderOptionsFn { - return func(o *kclReaderOptions) error { + return func(o *KclReader) error { o.Stats = sc return nil } @@ -140,17 +140,17 @@ type KclReader struct { // NewKclReader creates a new stream reader to read records from KCL func NewKclReader(c *aws.Config, optionFns ...KclReaderOptionsFn) (*KclReader, error) { - kclReaderOptions := defaultKclReaderOptions() + kclReader := &KclReader{kclReaderOptions: defaultKclReaderOptions()} for _, optionFn := range optionFns { - optionFn(kclReaderOptions) + optionFn(kclReader) } - return &KclReader{ - kclReaderOptions: kclReaderOptions, - LogHelper: &LogHelper{ - LogLevel: kclReaderOptions.logLevel, - Logger: c.Logger, - }, - }, nil + + kclReader.LogHelper = &LogHelper{ + LogLevel: kclReader.logLevel, + Logger: c.Logger, + } + + return kclReader, nil } func (r *KclReader) process(ctx context.Context) { diff --git a/kinesis_reader.go b/kinesis_reader.go index ff9b53f..f710173 100644 --- a/kinesis_reader.go +++ b/kinesis_reader.go @@ -38,12 +38,12 @@ func defaultKinesisReaderOptions() *kinesisReaderOptions { // KinesisReaderOptionsFn is a method signature for defining functional option methods for configuring // the KinesisReader. -type KinesisReaderOptionsFn func(*kinesisReaderOptions) error +type KinesisReaderOptionsFn func(*KinesisReader) error // KinesisReaderBatchSize is a functional option method for configuring the KinesisReader's // batch size. func KinesisReaderBatchSize(size int) KinesisReaderOptionsFn { - return func(o *kinesisReaderOptions) error { + return func(o *KinesisReader) error { if size > 0 && size <= kinesisReaderMaxBatchSize { o.batchSize = size return nil @@ -55,7 +55,7 @@ func KinesisReaderBatchSize(size int) KinesisReaderOptionsFn { // KinesisReaderShardIterator is a functional option method for configuring the KinesisReader's // shard iterator. func KinesisReaderShardIterator(shardIterator *ShardIterator) KinesisReaderOptionsFn { - return func(o *kinesisReaderOptions) error { + return func(o *KinesisReader) error { o.shardIterator = shardIterator return nil } @@ -64,7 +64,7 @@ func KinesisReaderShardIterator(shardIterator *ShardIterator) KinesisReaderOptio // KinesisReaderResponseReadTimeout is a functional option method for configuring the KinesisReader's // response read timeout. func KinesisReaderResponseReadTimeout(timeout time.Duration) KinesisReaderOptionsFn { - return func(o *kinesisReaderOptions) error { + return func(o *KinesisReader) error { o.responseReadTimeout = timeout return nil } @@ -72,7 +72,7 @@ func KinesisReaderResponseReadTimeout(timeout time.Duration) KinesisReaderOption // KinesisReaderLogLevel is a functional option method for configuring the KinesisReader's log level. func KinesisReaderLogLevel(ll aws.LogLevelType) KinesisReaderOptionsFn { - return func(o *kinesisReaderOptions) error { + return func(o *KinesisReader) error { o.logLevel = ll & 0xffff0000 return nil } @@ -80,7 +80,7 @@ func KinesisReaderLogLevel(ll aws.LogLevelType) KinesisReaderOptionsFn { // KinesisReaderStats is a functional option method for configuring the KinesisReader's stats collector. func KinesisReaderStats(sc ConsumerStatsCollector) KinesisReaderOptionsFn { - return func(o *kinesisReaderOptions) error { + return func(o *KinesisReader) error { o.Stats = sc return nil } @@ -100,25 +100,28 @@ type KinesisReader struct { // NewKinesisReader creates a new KinesisReader object which implements the StreamReader interface to read records from // Kinesis. func NewKinesisReader(c *aws.Config, stream string, shard string, optionFns ...KinesisReaderOptionsFn) (*KinesisReader, error) { - kinesisReaderOptions := defaultKinesisReaderOptions() - for _, optionFn := range optionFns { - optionFn(kinesisReaderOptions) - } sess, err := session.NewSession(c) if err != nil { return nil, err } - return &KinesisReader{ - kinesisReaderOptions: kinesisReaderOptions, - LogHelper: &LogHelper{ - LogLevel: kinesisReaderOptions.logLevel, - Logger: c.Logger, - }, - stream: stream, - shard: shard, - throttleSem: make(chan empty, 5), - client: kinesis.New(sess), - }, nil + + kinesisReader := &KinesisReader{ + kinesisReaderOptions: defaultKinesisReaderOptions(), + stream: stream, + shard: shard, + throttleSem: make(chan empty, 5), + client: kinesis.New(sess), + } + for _, optionFn := range optionFns { + optionFn(kinesisReader) + } + + kinesisReader.LogHelper = &LogHelper{ + LogLevel: kinesisReader.logLevel, + Logger: c.Logger, + } + + return kinesisReader, nil } // ensureShardIterator will lazily make sure that we have a valid ShardIterator, calling the GetShardIterator API with diff --git a/kinesis_writer.go b/kinesis_writer.go index b990e5d..d47f7f1 100644 --- a/kinesis_writer.go +++ b/kinesis_writer.go @@ -39,12 +39,12 @@ func defaultKinesisWriterOptions() *kinesisWriterOptions { // KinesisWriterOptionsFn is a method signature for defining functional option methods for configuring // the KinesisWriter. -type KinesisWriterOptionsFn func(*kinesisWriterOptions) error +type KinesisWriterOptionsFn func(*KinesisWriter) error // KinesisWriterResponseReadTimeout is a functional option method for configuring the KinesisWriter's // response read timeout func KinesisWriterResponseReadTimeout(timeout time.Duration) KinesisWriterOptionsFn { - return func(o *kinesisWriterOptions) error { + return func(o *KinesisWriter) error { o.responseReadTimeout = timeout return nil } @@ -53,7 +53,7 @@ func KinesisWriterResponseReadTimeout(timeout time.Duration) KinesisWriterOption // KinesisWriterMsgCountRateLimit is a functional option method for configuring the KinesisWriter's // message count rate limit func KinesisWriterMsgCountRateLimit(limit int) KinesisWriterOptionsFn { - return func(o *kinesisWriterOptions) error { + return func(o *KinesisWriter) error { if limit > 0 && limit <= kinesisMsgCountRateLimit { o.msgSizeRateLimit = limit return nil @@ -65,7 +65,7 @@ func KinesisWriterMsgCountRateLimit(limit int) KinesisWriterOptionsFn { // KinesisWriterMsgSizeRateLimit is a functional option method for configuring the KinesisWriter's // message size rate limit func KinesisWriterMsgSizeRateLimit(limit int) KinesisWriterOptionsFn { - return func(o *kinesisWriterOptions) error { + return func(o *KinesisWriter) error { if limit > 0 && limit <= kinesisMsgSizeRateLimit { o.msgSizeRateLimit = limit return nil @@ -76,7 +76,7 @@ func KinesisWriterMsgSizeRateLimit(limit int) KinesisWriterOptionsFn { // KinesisWriterLogLevel is a functional option method for configuring the KinesisWriter's log level func KinesisWriterLogLevel(ll aws.LogLevelType) KinesisWriterOptionsFn { - return func(o *kinesisWriterOptions) error { + return func(o *KinesisWriter) error { o.logLevel = ll & 0xffff0000 return nil } @@ -84,7 +84,7 @@ func KinesisWriterLogLevel(ll aws.LogLevelType) KinesisWriterOptionsFn { // KinesisWriterStats is a functional option method for configuring the KinesisWriter's stats collector func KinesisWriterStats(sc ProducerStatsCollector) KinesisWriterOptionsFn { - return func(o *kinesisWriterOptions) error { + return func(o *KinesisWriter) error { o.Stats = sc return nil } @@ -100,23 +100,26 @@ type KinesisWriter struct { // NewKinesisWriter creates a new stream writer to write records to a Kinesis. func NewKinesisWriter(c *aws.Config, stream string, optionFns ...KinesisWriterOptionsFn) (*KinesisWriter, error) { - kinesisWriterOptions := defaultKinesisWriterOptions() - for _, option := range optionFns { - option(kinesisWriterOptions) - } sess, err := session.NewSession(c) if err != nil { return nil, err } - return &KinesisWriter{ + + kinesisWriter := &KinesisWriter{ + kinesisWriterOptions: defaultKinesisWriterOptions(), stream: stream, client: kinesis.New(sess), - kinesisWriterOptions: kinesisWriterOptions, - LogHelper: &LogHelper{ - LogLevel: kinesisWriterOptions.logLevel, - Logger: c.Logger, - }, - }, nil + } + for _, option := range optionFns { + option(kinesisWriter) + } + + kinesisWriter.LogHelper = &LogHelper{ + LogLevel: kinesisWriter.logLevel, + Logger: c.Logger, + } + + return kinesisWriter, nil } // PutRecords sends a batch of records to Kinesis and returns a list of records that need to be retried. diff --git a/producer.go b/producer.go index f2072f9..e60609a 100644 --- a/producer.go +++ b/producer.go @@ -45,11 +45,11 @@ func defaultProducerOptions() *producerOptions { } // ProducerOptionsFn is a method signature for defining functional option methods for configuring the Producer. -type ProducerOptionsFn func(*producerOptions) error +type ProducerOptionsFn func(*Producer) error // ProducerWriter is a functional option method for configuring the producer's stream writer. func ProducerWriter(w StreamWriter) ProducerOptionsFn { - return func(o *producerOptions) error { + return func(o *Producer) error { o.writer = w return nil } @@ -57,7 +57,7 @@ func ProducerWriter(w StreamWriter) ProducerOptionsFn { // ProducerBatchSize is a functional option method for configuring the producer's batch size. func ProducerBatchSize(size int) ProducerOptionsFn { - return func(o *producerOptions) error { + return func(o *Producer) error { if size > 0 && size <= putRecordsMaxBatchSize { o.batchSize = size return nil @@ -68,7 +68,7 @@ func ProducerBatchSize(size int) ProducerOptionsFn { // ProducerBatchTimeout is a functional option method for configuring the producer's batch timeout. func ProducerBatchTimeout(timeout time.Duration) ProducerOptionsFn { - return func(o *producerOptions) error { + return func(o *Producer) error { o.batchTimeout = timeout return nil } @@ -76,7 +76,7 @@ func ProducerBatchTimeout(timeout time.Duration) ProducerOptionsFn { // ProducerQueueDepth is a functional option method for configuring the producer's queue depth. func ProducerQueueDepth(queueDepth int) ProducerOptionsFn { - return func(o *producerOptions) error { + return func(o *Producer) error { if queueDepth > 0 { o.queueDepth = queueDepth return nil @@ -87,7 +87,7 @@ func ProducerQueueDepth(queueDepth int) ProducerOptionsFn { // ProducerMaxRetryAttempts is a functional option method for configuring the producer's max retry attempts. func ProducerMaxRetryAttempts(attemtps int) ProducerOptionsFn { - return func(o *producerOptions) error { + return func(o *Producer) error { if attemtps > 0 { o.maxRetryAttempts = attemtps return nil @@ -98,7 +98,7 @@ func ProducerMaxRetryAttempts(attemtps int) ProducerOptionsFn { // ProducerConcurrency is a functional option method for configuring the producer's concurrency. func ProducerConcurrency(count int) ProducerOptionsFn { - return func(o *producerOptions) error { + return func(o *Producer) error { if count > 0 { o.concurrency = count return nil @@ -109,7 +109,7 @@ func ProducerConcurrency(count int) ProducerOptionsFn { // ProducerShardCheckFrequency is a functional option method for configuring the producer's shard check frequency. func ProducerShardCheckFrequency(duration time.Duration) ProducerOptionsFn { - return func(o *producerOptions) error { + return func(o *Producer) error { o.shardCheckFreq = duration return nil } @@ -117,7 +117,7 @@ func ProducerShardCheckFrequency(duration time.Duration) ProducerOptionsFn { // ProducerDataSpillFn is a functional option method for configuring the producer's data spill callback function. func ProducerDataSpillFn(fn MessageProcessor) ProducerOptionsFn { - return func(o *producerOptions) error { + return func(o *Producer) error { o.dataSpillFn = fn return nil } @@ -125,7 +125,7 @@ func ProducerDataSpillFn(fn MessageProcessor) ProducerOptionsFn { // ProducerLogLevel is a functional option method for configuring the producer's log level. func ProducerLogLevel(ll aws.LogLevelType) ProducerOptionsFn { - return func(o *producerOptions) error { + return func(o *Producer) error { o.logLevel = ll & 0xffff0000 return nil } @@ -133,7 +133,7 @@ func ProducerLogLevel(ll aws.LogLevelType) ProducerOptionsFn { // ProducerStats is a functional option method for configuring the producer's stats collector. func ProducerStats(sc ProducerStatsCollector) ProducerOptionsFn { - return func(o *producerOptions) error { + return func(o *Producer) error { o.Stats = sc return nil } @@ -159,24 +159,25 @@ type Producer struct { // NewProducer creates a new producer for writing records to a Kinesis or Firehose stream. func NewProducer(c *aws.Config, stream string, optionFns ...ProducerOptionsFn) (*Producer, error) { - producerOptions := defaultProducerOptions() + producer := &Producer{producerOptions: defaultProducerOptions()} for _, optionFn := range optionFns { - optionFn(producerOptions) + optionFn(producer) } - if producerOptions.writer == nil { + + if producer.writer == nil { w, err := NewKinesisWriter(c, stream) if err != nil { return nil, err } - producerOptions.writer = w + producer.writer = w + } + + producer.LogHelper = &LogHelper{ + LogLevel: producer.logLevel, + Logger: c.Logger, } - return &Producer{ - producerOptions: producerOptions, - LogHelper: &LogHelper{ - LogLevel: producerOptions.logLevel, - Logger: c.Logger, - }, - }, nil + + return producer, nil } // produce is called once to initialize a pool of workers which send batches of messages concurrently From c37399eb8aaf2ffff9ab82dfde0e1015dbf35eec Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 23 Aug 2017 09:42:08 -0500 Subject: [PATCH 085/109] Added the ability to configure the KinesisReader's GetRecords concurrency by adding a new concurrency field to the kinesisReaderOptions struct. Also added a new function method option to be able to adjust this parameter. --- kinesis_reader.go | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/kinesis_reader.go b/kinesis_reader.go index f710173..262f453 100644 --- a/kinesis_reader.go +++ b/kinesis_reader.go @@ -15,11 +15,13 @@ import ( const ( kinesisReaderMaxBatchSize = 10000 + kinesisReaderDefaultConcurrency = 5 ) // kinesisReaderOptions a struct that holds all of the KinesisReader's configurable parameters. type kinesisReaderOptions struct { batchSize int // maximum records per GetRecordsRequest call + concurrency int // maximum number of concurrent GetRecord or GetRecords calls allowed shardIterator *ShardIterator // shard iterator for Kinesis stream responseReadTimeout time.Duration // response read time out for GetRecordsRequest API call logLevel aws.LogLevelType // log level for configuring the LogHelper's log level @@ -30,6 +32,7 @@ type kinesisReaderOptions struct { func defaultKinesisReaderOptions() *kinesisReaderOptions { return &kinesisReaderOptions{ batchSize: kinesisReaderMaxBatchSize, + concurrency: kinesisReaderDefaultConcurrency, shardIterator: NewShardIterator(), responseReadTimeout: time.Second, Stats: &NilConsumerStatsCollector{}, @@ -52,6 +55,18 @@ func KinesisReaderBatchSize(size int) KinesisReaderOptionsFn { } } +// KinesisReaderConcurrency is a functional option method for configuring the KinesisReader's +// concurrency. +func KinesisReaderConcurrency(count int) KinesisReaderOptionsFn { + return func(o *KinesisReader) error { + if count > 0 { + o.concurrency = count + return nil + } + return ErrInvalidConcurrency + } +} + // KinesisReaderShardIterator is a functional option method for configuring the KinesisReader's // shard iterator. func KinesisReaderShardIterator(shardIterator *ShardIterator) KinesisReaderOptionsFn { @@ -109,13 +124,13 @@ func NewKinesisReader(c *aws.Config, stream string, shard string, optionFns ...K kinesisReaderOptions: defaultKinesisReaderOptions(), stream: stream, shard: shard, - throttleSem: make(chan empty, 5), client: kinesis.New(sess), } for _, optionFn := range optionFns { optionFn(kinesisReader) } + kinesisReader.throttleSem = make(chan empty, kinesisReader.concurrency) kinesisReader.LogHelper = &LogHelper{ LogLevel: kinesisReader.logLevel, Logger: c.Logger, From 819161dc53a15a7862e7ae67c4d7a0c8b0e654c3 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 23 Aug 2017 09:43:13 -0500 Subject: [PATCH 086/109] gofmt. --- kinesis_reader.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kinesis_reader.go b/kinesis_reader.go index 262f453..01e8c28 100644 --- a/kinesis_reader.go +++ b/kinesis_reader.go @@ -14,7 +14,7 @@ import ( ) const ( - kinesisReaderMaxBatchSize = 10000 + kinesisReaderMaxBatchSize = 10000 kinesisReaderDefaultConcurrency = 5 ) From fcac48be8706b754acdd81f74c0d5c7b8a8007a0 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 23 Aug 2017 15:33:40 -0500 Subject: [PATCH 087/109] Remove Checkpoint methods from StreamReader interface. --- kinesis_reader.go | 18 ------------------ shared_types.go | 3 --- 2 files changed, 21 deletions(-) diff --git a/kinesis_reader.go b/kinesis_reader.go index 01e8c28..201c5f7 100644 --- a/kinesis_reader.go +++ b/kinesis_reader.go @@ -335,21 +335,3 @@ func (r *KinesisReader) GetRecords(ctx context.Context, fn messageHandler) (coun count, size, err = r.getRecords(ctx, fn, r.batchSize) return count, size, err } - -// Checkpoint sends a message to KCL if there is sequence number that can be checkpointed -func (r *KinesisReader) Checkpoint() error { - // No-op (only applicable to KclReader) - return nil -} - -// CheckpointInsert registers a sequence number with the checkpointer -func (r *KinesisReader) CheckpointInsert(seqNum string) error { - // No-op (only applicable to KclReader) - return nil -} - -// CheckpointDone marks the given sequence number as done in the checkpointer -func (r *KinesisReader) CheckpointDone(seqNum string) error { - // No-op (only applicable to KclReader) - return nil -} diff --git a/shared_types.go b/shared_types.go index be65a11..d356714 100644 --- a/shared_types.go +++ b/shared_types.go @@ -16,9 +16,6 @@ type StreamWriter interface { type StreamReader interface { GetRecord(context.Context, messageHandler) (count int, size int, err error) GetRecords(context.Context, messageHandler) (count int, size int, err error) - Checkpoint() error // only applicable for KclReader - CheckpointInsert(seqNum string) error // only applicable for KclReader - CheckpointDone(seqNum string) error // only applicable for KclReader } // MessageProcessor defines the signature of a (asynchronous) callback function used by Listen, RetrieveFn and From 1cd9bfb544609a06c966f6331050dbaf871b7ee7 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 25 Aug 2017 10:59:32 -0500 Subject: [PATCH 088/109] Added documentation in doc.go and updated README.md file. Fixed a logical bug in the producer's doWork exit condition. Added read response timeout to FirehoseWriter. --- README.md | 5 +- doc.go | 390 +++++++++++++++++++++++++++++++++++++++++++++ firehose_writer.go | 11 ++ kinetic.go | 26 ++- producer.go | 3 +- 5 files changed, 417 insertions(+), 18 deletions(-) create mode 100644 doc.go diff --git a/README.md b/README.md index 944e9d0..876781a 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@ Kinetic is an MIT-licensed high-performance AWS Kinesis Client for Go Kinetic wraps [aws-sdk-go](https://github.com/aws/aws-sdk-go.git) to provide maximum throughput with built-in fault tolerance and retry logic for AWS Kinesis producers and consumers. -The Kinetic producer can write to Kinesis or Firehose and the Kinetic listener can consume stream data from Kinesis using the aws-go-sdk or using the Kinesis client library (written in Java). +The Kinetic producer can write to Kinesis or Firehose and the Kinetic consumer can consume stream data from Kinesis using the aws-sdk-go or using the Kinesis client library (written in Java). ### Getting Started Before using kinetic, you should make sure you have a created a Kinesis stream and your configuration file has the credentails necessary to read and write to the stream. Once this stream exists in AWS, kinetic will ensure it is in the "ACTIVE" state before running. @@ -68,4 +68,5 @@ p.Send(&kinetic.Message{ ``` -For more examples take a look at the tests or the test program in the `testexec` directory. API documentation can be found [here](https://godoc.org/github.com/rewardStyle/kinetic). +For more examples take a look at the tests or the test program in the `testexec` directory. For +additional information see the kinetic package documentation [here](https://godoc.org/github.com/rewardStyle/kinetic). diff --git a/doc.go b/doc.go new file mode 100644 index 0000000..edb1648 --- /dev/null +++ b/doc.go @@ -0,0 +1,390 @@ +/* +Package kinetic is a library written in Go that provides a scalable, high-performance and +fault-tolerant means to write to and read from an Amazon Kinesis stream. + +The kinetic library is distributed under an MIT License (see LICENSE.txt file). + +The kinetic library wraps the aws-sdk-go library (http://github.com/aws/aws-sdk-go.git) to provide +maximum throughput streaming to/from an Amazon Kinesis stream with built-in fault tolerance and +retry logic via the kinetic Producer and Consumer. The kinetic Producer can write to Kinesis +via the KinesisWriter or to Firehose via the FirehoseWriter. The kinetic Consumer can stream +data from Kinesis using the aws-sdk-go library via the KinesisReader or by using the Kinesis +Client Library +(http://docs.aws.amazon.com/streams/latest/dev/developing-consumers-with-kcl.html) via the +KclReader. + +Kinetic + +The Kinetic object provides access to a Kinesis client via the kinesisiface API +(https://docs.aws.amazon.com/sdk-for-go/api/service/kinesis/kinesisiface/) in +addition to providing utility methods for working with Kinesis streams such as creating and +deleting streams and retrieving shard information. When instantiating a Kinetic object, a new +AWS session is created using your AWS configuration (credentials, region and endpoint). Kinetic +also supports streaming to/from a local kinesalite (https://hub.docker.com/r/dlsniper/kinesalite/) +instance. + +A new Kinetic object can be instantiate using the NewKinetic() function as such: + + k, err := kinetic.NewKinetic( + kinetic.AwsConfigCredentials("some-access-key", "some-secret-key", "some-security-token"), + kinetic.AwsConfigRegion("some-region"), + kinetic.AwsConfigEndpoint("http://127.0.0.1:4567"), + ) + +which accepts a variable number of functional option methods. + +Producer + +The Producer object is used to stream data to an Amazon Kinesis or Firehose stream in batches via +the Send, SendWithContext or TryToSend functions. + +The Producer achieves optimal throughput by implementing a dispatcher/worker model to pull +messages off of a queue and send batches concurrently (within rate limits) based on the number +active shards, which is automatically adjusted after a re-sharding operation occurs externally. + +Usage: + +To create a Producer with default values, call NewProducer() with a pointer to the +aws.Config struct and the stream name as such: + + p, err := kinetic.NewProducer(k.Session.Config, "some-stream-name") + +where k is the Kinetic object created (above). + +To instantiate a Producer with custom parameters, pass in functional option methods like this: + + psc := kinetic.NewDefaultProducerStatsCollector(registry) + w, err := kinetic.NewKinesisWriter(k.Session.Config, "some-stream-name", + kinetic.KinesisWriterResponseReadTimeout(time.Second), + kinetic.KinesisWriterMsgCountRateLimit(1000), + kinetic.KinesisWriterMsgSizeRateLimit(1000000), + kinetic.KinesisWriterLogLevel(aws.LogOff), + kinetic.KinesisWriterStats(psc), + ) + + p, err := kinetic.NewProducer(k.Session.Config, "some-stream-name", + kinetic.ProducerWriter(w), + kinetic.ProducerBatchSize(500), + kinetic.ProducerBatchTimeout(time.Second), + kinetic.ProducerMaxRetryAttempts(3), + kinetic.ProducerQueueDepth(10000), + kinetic.ProducerConcurrency(3), + kinetic.ProducerShardCheckFrequency(time.Minute), + kinetic.ProducerDataSpillFn(func(msg *kinetic.Message) error { + log.Printf("Message was dropped: [%s]\n", string(msg.Data)) + return nil + }), + kinetic.ProducerLogLevel(aws.LogOff), + kinetic.ProducerStats(psc), + ) + +Then to send messages: + + jsonStr, err := json.Marshal(yourDataStruct) + msg := &kinetic.Message{ + PartitionKey: aws.String("unique-id-per-message"), + Data: []byte(jsonStr), + } + + // Drops message to data spill function if message channel is full + p.TryToSend(msg) + +OR + + // Blocks until message gets enqueued onto the message channel + p.Send(msg) + +How it works: + +The Producer is started by calling one of the Send APIs (Send / SendWithContext / TryToSend) which +kicks off two goroutines: 1) one to periodically check the number of active shards and adjust the +number of concurrent workers accordingly and 2) one to dispatch messages to the workers and +throttle the number of messages per second and the transmission (memory) size of the batch per +second based on the stream limits set by AWS +(http://docs.aws.amazon.com/streams/latest/dev/service-sizes-and-limits.html). Messages sent +via the Send APIs are enqueued in a message channel which get batched and sent via the +dispatcher/worker model. + +If the shard monitoring determines that there are too few workers available based on the number of +active shards, it will spawn additional workers which run in separate goroutines. If the shard +monitoring determines that there are too many workers available based on the number of active +shards, it will send a signal to a dismiss channel to which workers actively listen. The rate +limiters get reset after a change is detected. + +After being spawned, a worker runs indefinitely doing two things: 1) communicating to the +dispatcher (via a status channel) how many messages it can send in its next batch (based on the +batch size less the number of previously failed messages) and 2) calling the Producer's sendBatch +function AFTER receiving a new batch from the dispatcher. If any messages failed to send in the +sendBatch call, those messages are retried by the worker in subsequent batches until they reach +the maximum number of retry attempts, after which they are sent to a data spill callback function. +Workers run indefinitely until they receive a signal to dismiss AND they have no retry messages +left to send. The status channel is an unbuffered channel, so workers will idle while trying to +communicate to the dispatcher on this channel when the message volume is low. + +The dispatcher does the following: 1) listens to the status channel for workers signaling their +availability and capacity 2) pulls messages off of the message channel and 3) throttles the +outgoing messages using a token bucket based on the allowed rate limits. Two concurrent +goroutinues are started in parallel (one for the message count of the batch and one for the +memory transmission size of the batch) to wait until the respective tokens are available. + +In a high volume scenario, the rate limiters prevent the Producer from sending too many messages +in a short period of time which would cause provisioned throughput exceeded exceptions. +In a low volume scenario, a batch timeout causes batches to be sent with fewer messages than the +maximum batch limit so the rate limiters do not throttle the outgoing messages. + +KinesisWriter + +The KinesisWriter is the default writer used by the Producer to stream data. To stream +messages to Amazon Firehose use the FirehoseWriter. + +The KinesisWriter implements the StreamWriter interface using the kinesisiface API to make +PutRecordsRequests. Successfully sent messages are updated with meta data containing the +sequence number (http://docs.aws.amazon.com/streams/latest/dev/key-concepts.html#sequence-number) +and the shardId. Failed messages are updated with meta data which tracks the error code, error +message and fail count. + +To create a custom Producer with a KinesisWriter, see the example in Producer (above). + +FirehoseWriter + +The FirehoseWriter implements the StreamWriter interface using the firehoseiface API to make +PutRecordBatchRequests. Successfully sent messages are updated with meta data containing the +sequence number (http://docs.aws.amazon.com/streams/latest/dev/key-concepts.html#sequence-number) +and the shardId. Failed messages are updated with meta data which tracks the error code, error +message and fail count. + +To create a Producer with a FirehoseWriter, use the functional option methods like this: + + psc := kinetic.NewDefaultProducerStatsCollector(registry) + w, err := kinetic.NewFirehoseWriter(k.Session.Config, "some-stream-name", + kinetic.FirehoseWriterResponseReadTimeout(time.Second), + kinetic.FirehoseWriterMsgCountRateLimit(1000), + kinetic.FirehoseWriterMsgSizeRateLimit(1000000), + kinetic.FirehoseWriterLogLevel(aws.LogOff), + kinetic.FirehoseWriterStats(psc), + ) + + p, err := kinetic.NewProducer(k.Session.Config, "some-stream-name", + kinetic.ProducerWriter(w), + kinetic.ProducerBatchSize(500), + kinetic.ProducerBatchTimeout(time.Second), + kinetic.ProducerMaxRetryAttempts(3), + kinetic.ProducerQueueDepth(10000), + kinetic.ProducerConcurrency(3), + kinetic.ProducerShardCheckFrequency(time.Minute), + kinetic.ProducerDataSpillFn(func(msg *kinetic.Message) error { + log.Printf("Message was dropped: [%s]\n", string(msg.Data)) + return nil + }), + kinetic.ProducerLogLevel(aws.LogOff), + kinetic.ProducerStats(psc), + ) + +Consumer + +The Consumer object is used to stream data from an Amazon Kinesis stream in batches via the +Retrieve or Listen functions. + +The Consumer achieves optimal throughput by implementing a dispatcher/worker model to pull +messages off of a queue and send batches concurrently (within rate limits) based on the number +active shards, which is automatically adjusted after a re-sharding operation occurs externally. + +Usage: + +To create a Consumer with default values, call the NewConsumer function like this: + + c, err := kinetic.NewConsumer(k.Session.Config, "some-stream-name", "some-shard-name") + +where k is the Kinetic object created (above) + +To create a Consumer with custom parameters, pass in functional option methods like this: + + csc := kinetic.NewDefaultConsumerStatsCollector(registry) + r, err := kinetic.NewKinesisReader(k.Session.Config, "some-stream-name", "some-shard-id", + kinetic.KinesisReaderBatchSize(10000), + kinetic.KinesisReaderResponseReadTimeout(time.Second), + kinetic.KinesisReaderLogLevel(aws.LogOff), + kinetic.KinesisReaderStats(csc), + ) + + c, err := kinetic.NewConsumer(k.Session.Config, "some-stream-name", "some-shard-id", + kinetic.ConsumerReader(r), + kinetic.ConsumerQueueDepth(500), + kinetic.ConsumerConcurrency(10), + kinetic.ConsumerLogLevel(aws.LogOff), + kinetic.ConsumerStats(csc), + ) + +Then to retrieve messages: + + // To retrieve a single message off of the message queue + message, err := c.Retrieve() + msg := &YourDataStruct{} + json.Unmarshal(message.Data, msg) + + // Do something with msg + +OR + + // To listen to the message channel indefinitely + c.Listen(func(m *kinetic.Message) error { + // Unmarshal data + msg := &YourDataStruct{} + err := json.Unmarshal(m.Data, msg) + + // Do something with msg + + return err + }) + +How it works: + +The Consumer is started by calling one of the Retrieve / Listen APIs which kicks off a goroutine +that does two things in an infinite loop: 1) invokes a GetRecords call and 2) enqueues the batch +of messages from the GetRecords call to a message queue. The GetRecords calls are throttled by +a rate limiter which utilizes a token bucket system for the number of GetRecords transactions per +second and the transmission (memory) size of the batch per second. Because the message queue is +a buffered message channel, this goroutine becomes blocked once the message queue is full. + +The Retrieve APIs (Retrieve, RetrieveWithContext, RetrieveFn and RetrieveFnWithContext) pull one +message off of the message queue while the Listen APIs (Listen and ListenWithContext) pull +messages off of the message queue concurrently based on the Consumer's concurrency setting. The +supplied callback function for the Listen APIs is run (asynchronously) on the messages as the are +pulled off. + +KinesisReader + +The KinesisReader is the default reader used by the Consumer to stream data. + +The KinesisReader implements the StreamReader interface using the kinesisiface API to make +GetRecordsRequest, the results of which are enqueued to the Consumer's message queue. + +To create a custom Consumer with a KinesisReader, see the example for Consumer (above). + +KclReader + +The KclReader implements the StreamReader interface and adheres to the MultiLangDaemon protocol +to communicate with the Kinetic Client Libray (KCL) over STDIN and STDOUT as documented here: +https://github.com/awslabs/amazon-kinesis-client/blob/master/src/main/java/com/amazonawservices/kinesis/multilang/package-info.java + +The KCL is a Java library which requires a Java Runtime Environment to run +(http://docs.aws.amazon.com/streams/latest/dev/developing-consumers-with-kcl.html). If you +choose to use the KclReader to stream your data from Amazon Kinesis (instead of the +KinesisReader), your application must have a JRE and must run the MultiLangDaemon as a background +process. + +To create a custom Consumer with a KclReader, use the functional option methods as such: + + csc := kinetic.NewDefaultConsumerStatsCollector(registry) + r, err := kinetic.NewKclReader(k.Session.Config, "some-stream-name", "some-shard-id", + kinetic.KclReaderAutoCheckpointCount(1000), + kinetic.KclReaderAutoCheckpointFreq(time.Minute), + kinetic.KclReaderUpdateCheckpointSizeFreq(time.Minute), + kinetic.KclReaderOnInitCallbackFn(func () {return nil}), + kinetic.KclReaderOnCheckpointCallbackFn(func () {return nil}), + kinetic.KclReaderOnShutdownCallbackFn(func () {return nil}), + kinetic.KclReaderLogLevel(aws.LogOff), + kinetic.KclReaderStats(csc), + ) + + c, err := kinetic.NewConsumer(k.Session.Config, "some-stream-name", "some-shard-id", + kinetic.ConsumerReader(r), + kinetic.ConsumerQueueDepth(500), + kinetic.ConsumerConcurrency(10), + kinetic.ConsumerLogLevel(aws.LogOff), + kinetic.ConsumerStats(csc), + ) + +How it works: + +The KclReader is started by calling one of the Retrieve / Listen APIs of the Consumer which kicks +off separate goroutines that write to STDOUT and read from STDIN in accordance with the +MultiLangDaemon protocol. + + The multi language protocol defines a system for communication between a KCL multi-lang + application and another process (referred to as the "child process") over STDIN and + STDOUT of the child process. + +The sequence of events that occurs between the KCL process and the KclReader goes something like +this: + + * an "initialize" action message is sent to the child process which the KclReader reads + from STDIN. + * KclReader sends an acknowledgement status message to STDOUT. + * a "processRecord" message is sent to the child process which the KclReader reads from + STDIN. + - In a separate goroutine, the KclReader ranges through each record enqueueing + the message to an internal message channel (unbuffered) + - messages are pulled off the queue as needed by the Consumer's GetRecord and + GetRecords APIs + - upon completion of the batch of messages sent by KCL... + * KclReader sends an acknowledgement stats message to STDOUT. + * At any point after the KclReader sends acknowledgement status message for the + "initialize" status message, the KclReader can send a "checkpoint" status message to KCL + over STDOUT. + * KCL then sends an acknowledgement to the "checkpoint" status message + * When KCL sends a "shutdown" action message, KclReader invokes an OnShutdown callback + function and sends an acknowledgement to KCL over STDOUT. + +It is important to note that checkpointing can occur in between receiving a "processRecord" +message from KCL and sending an acknowledgement message. + +Checkpointing: + +Checkpointing is a feature that is only available for the KclReader. Checkpointing allows your +application to indicate to the KclReader that all records up to and including a given sequence +number have been successfully received and processed by your application so that if your +application were to be restarted (for whatever reason) the streaming would start with the next +sequence number instead of the TRIM_HORIZON. Checkpointing is not necessary if your application +consuming Kinesis data is idempotent, but starting from the TRIM_HORIZON on restarts will create +unnecessary (and avoidable) processing. + +The KclReader provides three methods for checkpointing: (CheckpointInsert, CheckpointDelete and +Checkpoint). + +Usage of the Checkpoint APIs goes something like this: + + r, _ := kinetic.NewKclReader( + kinetic.KclReaderAutoCheckpointCount(1000), + kinetic.KclReaderAutoCheckpointFreq(time.Hour), + kinetic.KclReaderOnCheckpointCallbackFn( + func (seqNum string, err error) error { + return nil + }, + ), + // Additional functional option methods here + ) + + c, err := kinetic.NewConsumer(k.Session.Config, "some-stream-name", "some-shard-id", + kinetic.ConsumerReader(r), + // Additional functional option methods here + ) + + // Stream data from Kinesis + go func() { + c.Listen(func (m *kinetic.Message) error { + r.CheckpointInsert(m.SequenceNumber) + defer r.CheckpointDone(m.SequenceNumber) + + // Unmarshal data + msg := &Message{} + json.Unmarshal(m.Data, msg) + + // Do stuff with your data + }) + }() + + // Periodically call checkpoint + go func() { + ticker := time.NewTicker(time.Minute) + defer r.ticker.Stop() + + select { + case <-ticker.C: + r.Checkpoint() + case <-stop: + return + } + }() +*/ +package kinetic \ No newline at end of file diff --git a/firehose_writer.go b/firehose_writer.go index cd2584d..ac2d604 100644 --- a/firehose_writer.go +++ b/firehose_writer.go @@ -19,6 +19,7 @@ const ( // firehoseWriterOptions is a struct that holds all of the FirehoseWriter's configurable parameters. type firehoseWriterOptions struct { + responseReadTimeout time.Duration // maximum time to wait for PutRecords API call before timing out msgCountRateLimit int // maximum number of records to be sent per second msgSizeRateLimit int // maximum (transmission) size of records to be sent per second throughputMultiplier int // integer multiplier to increase firehose throughput rate limits @@ -41,6 +42,15 @@ func defaultFirehoseWriterOptions() *firehoseWriterOptions { // the FirehoseWriter. type FirehoseWriterOptionsFn func(*FirehoseWriter) error +// FirehoseWriterResponseReadTimeout is a functional option method for configuring the +// FirehoseWriter's response read timeout +func FirehoseWriterResponseReadTimeout(timeout time.Duration) KinesisWriterOptionsFn { + return func(o *KinesisWriter) error { + o.responseReadTimeout = timeout + return nil + } +} + // FirehoseWriterMsgCountRateLimit is a functional option method for configuring the FirehoseWriter's // message count rate limit. func FirehoseWriterMsgCountRateLimit(limit int) FirehoseWriterOptionsFn { @@ -141,6 +151,7 @@ func (w *FirehoseWriter) PutRecords(ctx context.Context, messages []*Message, fn DeliveryStreamName: aws.String(w.stream), Records: records, }) + req.ApplyOptions(request.WithResponseReadTimeout(w.responseReadTimeout)) req.Handlers.Build.PushFront(func(r *request.Request) { startBuildTime = time.Now() diff --git a/kinetic.go b/kinetic.go index 0d7445f..a24e567 100644 --- a/kinetic.go +++ b/kinetic.go @@ -12,7 +12,6 @@ import ( "net/http" "github.com/aws/aws-sdk-go/aws/credentials" - "github.com/aws/aws-sdk-go/service/firehose/firehoseiface" "github.com/aws/aws-sdk-go/service/kinesis" "github.com/aws/aws-sdk-go/service/kinesis/kinesisiface" ) @@ -34,11 +33,11 @@ func defaultKineticOptions() *kineticOptions { } } -// OptionsFn is a method signature for defining functional option methods for configuring Kinetic. -type OptionsFn func(*kineticOptions) error +// AwsConfigOptionsFn is a method signature for defining functional option methods for configuring Kinetic. +type AwsConfigOptionsFn func(*kineticOptions) error // AwsConfigCredentials is a functional option method for configuring Kinetic's AwsConfig credentials. -func AwsConfigCredentials(accessKey, secretKey, sessionToken string) OptionsFn { +func AwsConfigCredentials(accessKey, secretKey, sessionToken string) AwsConfigOptionsFn { return func(o *kineticOptions) error { o.awsConfig.WithCredentials(credentials.NewStaticCredentials(accessKey, secretKey, sessionToken)) return nil @@ -46,7 +45,7 @@ func AwsConfigCredentials(accessKey, secretKey, sessionToken string) OptionsFn { } // AwsConfigRegion is a functional option method for configuring Kinetic's AwsConfig region. -func AwsConfigRegion(region string) OptionsFn { +func AwsConfigRegion(region string) AwsConfigOptionsFn { return func(o *kineticOptions) error { o.awsConfig.WithRegion(region) return nil @@ -54,7 +53,7 @@ func AwsConfigRegion(region string) OptionsFn { } // AwsConfigEndpoint is a functional option method for configuring Kinetic's AwsConfig endpoint. -func AwsConfigEndpoint(endpoint string) OptionsFn { +func AwsConfigEndpoint(endpoint string) AwsConfigOptionsFn { return func(o *kineticOptions) error { o.awsConfig.WithEndpoint(endpoint) return nil @@ -62,7 +61,7 @@ func AwsConfigEndpoint(endpoint string) OptionsFn { } // AwsConfigLogger is a functional option method for configuring Kinetic's AwsConfig logger. -func AwsConfigLogger(logger aws.Logger) OptionsFn { +func AwsConfigLogger(logger aws.Logger) AwsConfigOptionsFn { return func(o *kineticOptions) error { o.awsConfig.WithLogger(logger) return nil @@ -70,7 +69,7 @@ func AwsConfigLogger(logger aws.Logger) OptionsFn { } // AwsConfigLogLevel is a functional option method for configuring Kinetic's AwsConfig log level. -func AwsConfigLogLevel(logLevel aws.LogLevelType) OptionsFn { +func AwsConfigLogLevel(logLevel aws.LogLevelType) AwsConfigOptionsFn { return func(o *kineticOptions) error { o.awsConfig.WithLogLevel(logLevel) return nil @@ -79,7 +78,7 @@ func AwsConfigLogLevel(logLevel aws.LogLevelType) OptionsFn { // AwsConfigHTTPClientTimeout is a functional option method for configuring Kinetic's // AwsConfig HTTP client timeout. -func AwsConfigHTTPClientTimeout(timeout time.Duration) OptionsFn { +func AwsConfigHTTPClientTimeout(timeout time.Duration) AwsConfigOptionsFn { return func(o *kineticOptions) error { o.awsConfig.WithHTTPClient(&http.Client{ Timeout: timeout, @@ -89,26 +88,25 @@ func AwsConfigHTTPClientTimeout(timeout time.Duration) OptionsFn { } // LogLevel is a functional option method for configuring Kinetic's log level. -func LogLevel(logLevel aws.LogLevelType) OptionsFn { +func LogLevel(logLevel aws.LogLevelType) AwsConfigOptionsFn { return func(o *kineticOptions) error { o.logLevel = logLevel & 0xffff0000 return nil } } -// Kinetic represents a kinesis and firehose client and provides some utility -// methods for interacting with the AWS services. +// Kinetic provides access to a Kinesis client and provides some utility methods for interacting +// with the AWS Kinesis service. type Kinetic struct { *kineticOptions *LogHelper clientMu sync.Mutex - fclient firehoseiface.FirehoseAPI kclient kinesisiface.KinesisAPI Session *session.Session } // NewKinetic creates a new instance of Kinetic. -func NewKinetic(optionFns ...OptionsFn) (*Kinetic, error) { +func NewKinetic(optionFns ...AwsConfigOptionsFn) (*Kinetic, error) { kineticOptions := defaultKineticOptions() for _, optionFn := range optionFns { optionFn(kineticOptions) diff --git a/producer.go b/producer.go index e60609a..cc559d5 100644 --- a/producer.go +++ b/producer.go @@ -406,7 +406,7 @@ func (p *Producer) doWork() { var retries []*Message var dismissed bool - for ok := true; ok; ok = !dismissed || len(retries) != 0 { + for ok := true; ok; ok = !(dismissed && len(retries) == 0) { // Check to see if there were any signals to dismiss workers (if eligible) if !dismissed { select { @@ -445,7 +445,6 @@ func (p *Producer) doWork() { if len(batch)+len(retries) > 0 { retries = p.sendBatch(append(retries, batch...)) } - } } From ed36ce97b3174977979a5c4fbfe34334baf758ea Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 25 Aug 2017 11:16:02 -0500 Subject: [PATCH 089/109] Fix race condition in checkpoint_test.go that was causing tests to fail. --- checkpointer_test.go | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/checkpointer_test.go b/checkpointer_test.go index 9d25f87..c32c33e 100644 --- a/checkpointer_test.go +++ b/checkpointer_test.go @@ -357,6 +357,7 @@ func TestCheckpointerAutoCheckpointing(t *testing.T) { var checkpointFnCalled uint64 cp := newCheckpointer( checkpointCountCheckFreq(time.Millisecond), + checkpointAutoCheckpointFreq(time.Minute), checkpointAutoCheckpointCount(autoCheckpointCount), checkpointCheckpointFn(func(seqNum string) error { atomic.AddUint64(&checkpointFnCalled, 1) @@ -383,8 +384,8 @@ func TestCheckpointerAutoCheckpointing(t *testing.T) { } So(failedCount, ShouldEqual, 0) - <-time.After(time.Millisecond) - Convey("confirming that checkpoint and trim was called", func() { + <-time.After(2*time.Millisecond) + Convey("confirming that checkpoint was called", func() { So(atomic.LoadUint64(&checkpointFnCalled), ShouldBeGreaterThan, 0) }) }) @@ -397,6 +398,7 @@ func TestCheckpointerAutoCheckpointing(t *testing.T) { var checkpointFnCalled uint64 cp := newCheckpointer( checkpointAutoCheckpointFreq(autoCheckpointFreq), + checkpointAutoCheckpointCount(1000), checkpointCheckpointFn(func(seqNum string) error { atomic.AddUint64(&checkpointFnCalled, 1) return nil @@ -423,7 +425,7 @@ func TestCheckpointerAutoCheckpointing(t *testing.T) { So(failedCount, ShouldEqual, 0) <-time.After(time.Duration(2) * autoCheckpointFreq) - Convey("confirming that checkpoint and trim was called", func() { + Convey("confirming that checkpoint was called", func() { So(atomic.LoadUint64(&checkpointFnCalled), ShouldBeGreaterThan, 0) }) }) From 6a5bfe27eac65ae6a8fae05d9835c91729d53dc6 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Tue, 29 Aug 2017 11:21:59 -0500 Subject: [PATCH 090/109] Moved the rate limiting logic from the Consumer to the KinesisReader. Simplified the GetRecord / GetRecords API to only return an error (instead of returning count and size as well). Updated documentation on the Consumer / KinesisReader around rate limiting. --- checkpointer_test.go | 2 +- consumer.go | 83 +++++++++++++--------------------- doc.go | 44 ++++++++++-------- errors.go | 6 +++ kcl_reader.go | 12 ++--- kinesis_reader.go | 104 +++++++++++++++++++++++++++++++++---------- shared_types.go | 4 +- testexec/main.go | 3 ++ 8 files changed, 156 insertions(+), 102 deletions(-) diff --git a/checkpointer_test.go b/checkpointer_test.go index c32c33e..c47c73c 100644 --- a/checkpointer_test.go +++ b/checkpointer_test.go @@ -384,7 +384,7 @@ func TestCheckpointerAutoCheckpointing(t *testing.T) { } So(failedCount, ShouldEqual, 0) - <-time.After(2*time.Millisecond) + <-time.After(2 * time.Millisecond) Convey("confirming that checkpoint was called", func() { So(atomic.LoadUint64(&checkpointFnCalled), ShouldBeGreaterThan, 0) }) diff --git a/consumer.go b/consumer.go index cfb7e9f..fb6c9f7 100644 --- a/consumer.go +++ b/consumer.go @@ -7,7 +7,6 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "golang.org/x/time/rate" ) // consumerOptions is used to hold all of the configurable settings of a Consumer. @@ -79,16 +78,14 @@ func ConsumerStats(sc ConsumerStatsCollector) ConsumerOptionsFn { // Consumer polls the StreamReader for messages. type Consumer struct { - *consumerOptions // contains all of the configuration settings for the Consumer - *LogHelper // object for help with logging - txnCountRateLimiter *rate.Limiter // rate limiter to limit the number of transactions per second - txSizeRateLimiter *rate.Limiter // rate limiter to limit the transmission size per seccond - messages chan *Message // channel for storing messages that have been retrieved from the stream - concurrencySem chan empty // channel for controlling the number of concurrent workers processing messages from the message channel - pipeOfDeath chan empty // channel for handling pipe of death - consuming bool // flag for indicating whether or not the consumer is consuming - consumingMu sync.Mutex // mutex for making the consuming flag thread safe - noCopy noCopy // prevents the Consumer from being copied + *consumerOptions // contains all of the configuration settings for the Consumer + *LogHelper // object for help with logging + messages chan *Message // channel for storing messages that have been retrieved from the stream + concurrencySem chan empty // channel for controlling the number of concurrent workers processing messages from the message channel + pipeOfDeath chan empty // channel for handling pipe of death + consuming bool // flag for indicating whether or not the consumer is consuming + consumingMu sync.Mutex // mutex for making the consuming flag thread safe + noCopy noCopy // prevents the Consumer from being copied } // NewConsumer creates a new Consumer object for retrieving and listening to message(s) on a StreamReader. @@ -153,33 +150,33 @@ func (c *Consumer) stopConsuming() { } // enqueueSingle calls the readers's GetRecord method and enqueus a single message on the message channel. -func (c *Consumer) enqueueSingle(ctx context.Context) (count int, size int) { - var err error - count, size, err = c.reader.GetRecord(ctx, +func (c *Consumer) enqueueSingle(ctx context.Context) error { + err := c.reader.GetRecord(ctx, func(msg *Message) error { c.messages <- msg return nil }) if err != nil { c.handleErrorLogging(err) + return err } - return count, size + return nil } // enqueueBatch calls the reader's GetRecords method and enqueues a batch of messages on the message chanel. -func (c *Consumer) enqueueBatch(ctx context.Context) (count, size int) { - var err error - count, size, err = c.reader.GetRecords(ctx, +func (c *Consumer) enqueueBatch(ctx context.Context) error { + err := c.reader.GetRecords(ctx, func(msg *Message) error { c.messages <- msg return nil }) if err != nil { c.handleErrorLogging(err) + return err } - return count, size + return nil } // handleErrorLogging is a helper method for handling and logging errors from calling the reader's @@ -224,10 +221,12 @@ func (c *Consumer) RetrieveWithContext(ctx context.Context) (*Message, error) { if !ok { return nil, err } - n, _ := c.enqueueSingle(childCtx) - if n > 0 { - c.Stats.AddDelivered(n) - return <-c.messages, nil + + c.enqueueSingle(childCtx) + select { + case msg := <-c.messages: + return msg, nil + default: } } } @@ -245,10 +244,15 @@ func (c *Consumer) RetrieveFnWithContext(ctx context.Context, fn MessageProcesso return err } - start := time.Now() - fn(msg) - c.Stats.UpdateProcessedDuration(time.Since(start)) - c.Stats.AddProcessed(1) + if fn != nil { + start := time.Now() + if err := fn(msg); err != nil { + return err + } + c.Stats.UpdateProcessedDuration(time.Since(start)) + c.Stats.AddProcessed(1) + + } return nil } @@ -269,10 +273,6 @@ func (c *Consumer) consume(ctx context.Context) { go func() { defer c.stopConsuming() - // TODO: make these parameters configurable also scale according to the shard count - c.txnCountRateLimiter = rate.NewLimiter(rate.Limit(5), 1) - c.txSizeRateLimiter = rate.NewLimiter(rate.Limit(2000000), 2000000) - childCtx, cancel := context.WithCancel(ctx) defer cancel() for { @@ -285,26 +285,7 @@ func (c *Consumer) consume(ctx context.Context) { return } - _, size := c.enqueueBatch(childCtx) - - wg := sync.WaitGroup{} - wg.Add(1) - go func() { - defer wg.Done() - - if err := c.txnCountRateLimiter.Wait(childCtx); err != nil { - c.LogError("Error occured waiting for transaction count tokens") - } - }() - - wg.Add(1) - go func() { - defer wg.Done() - if err := c.txSizeRateLimiter.WaitN(childCtx, size); err != nil { - c.LogError("Error occured waiting for transmission size tokens") - } - }() - wg.Wait() + c.enqueueBatch(childCtx) } }() } diff --git a/doc.go b/doc.go index edb1648..835b1d3 100644 --- a/doc.go +++ b/doc.go @@ -183,11 +183,11 @@ To create a Producer with a FirehoseWriter, use the functional option methods li Consumer The Consumer object is used to stream data from an Amazon Kinesis stream in batches via the -Retrieve or Listen functions. - -The Consumer achieves optimal throughput by implementing a dispatcher/worker model to pull -messages off of a queue and send batches concurrently (within rate limits) based on the number -active shards, which is automatically adjusted after a re-sharding operation occurs externally. +Retrieve and Listen APIs which ultimately call the StreamReader's GetRecord / GetRecords functions. +All records returned by the GetRecord / GetRecords calls are placed onto the Consumer's buffered +message queue from which the Consumer's Retrieve and Listen functions pull and process +concurrently. The Consumer's concurrency setting can be adjusted to modify the number of +concurrent goroutines processing messages from the message queue. Usage: @@ -202,6 +202,9 @@ To create a Consumer with custom parameters, pass in functional option methods l csc := kinetic.NewDefaultConsumerStatsCollector(registry) r, err := kinetic.NewKinesisReader(k.Session.Config, "some-stream-name", "some-shard-id", kinetic.KinesisReaderBatchSize(10000), + kinetic.KinesisReaderConcurrency(5), + kinetic.KinesisReaderTransactionCountLimit(5), + kinetic.KinesisReaderTransmissionSizeLimit(2000000), kinetic.KinesisReaderResponseReadTimeout(time.Second), kinetic.KinesisReaderLogLevel(aws.LogOff), kinetic.KinesisReaderStats(csc), @@ -239,25 +242,30 @@ OR How it works: -The Consumer is started by calling one of the Retrieve / Listen APIs which kicks off a goroutine -that does two things in an infinite loop: 1) invokes a GetRecords call and 2) enqueues the batch -of messages from the GetRecords call to a message queue. The GetRecords calls are throttled by -a rate limiter which utilizes a token bucket system for the number of GetRecords transactions per -second and the transmission (memory) size of the batch per second. Because the message queue is -a buffered message channel, this goroutine becomes blocked once the message queue is full. +The Retrieve APIs (Retrieve, RetrieveWithContext, RetrieveFn and RetrieveFnWithContext) work +by starting the Consumer if not already started, which involves instantiating the necessary +communications channels, and then indefinitely calling GetRecord until a single message is pulled +off of the Consumer's message queue or the supplied context is cancelled. + +The Listen APIs (Listen and ListenWithContext) work by starting the Consumer if not already +started, which involves instantiating the necessary communications channels, and then +indefinitely calling GetRecords and enqueueing the returned records onto the Consumer's message +channel (indefinitely in a separate go routine until the supplied context is cancelled). In the +Listen function, the messages are pulled off of the message queue and the user-supplied callback +function is applied to it. -The Retrieve APIs (Retrieve, RetrieveWithContext, RetrieveFn and RetrieveFnWithContext) pull one -message off of the message queue while the Listen APIs (Listen and ListenWithContext) pull -messages off of the message queue concurrently based on the Consumer's concurrency setting. The -supplied callback function for the Listen APIs is run (asynchronously) on the messages as the are -pulled off. +If the Retrieve/Listen function(s) are called concurrently with each other, the first call +will block while the subsequent Retrieve/Listen calls will fail with an ErrAlreadyConsuming error. KinesisReader The KinesisReader is the default reader used by the Consumer to stream data. The KinesisReader implements the StreamReader interface using the kinesisiface API to make -GetRecordsRequest, the results of which are enqueued to the Consumer's message queue. +GetRecordsRequest, the results of which are enqueued to the Consumer's message queue. The +GetRecords calls are throttled by a rate limiter which utilizes a token bucket system for the +number of GetRecords transactions per second and the transmission (memory) size of the batch per +second. To create a custom Consumer with a KinesisReader, see the example for Consumer (above). @@ -387,4 +395,4 @@ Usage of the Checkpoint APIs goes something like this: } }() */ -package kinetic \ No newline at end of file +package kinetic diff --git a/errors.go b/errors.go index 2792ecf..76f8b09 100644 --- a/errors.go +++ b/errors.go @@ -104,4 +104,10 @@ var ( // ErrInvalidConcurrency is returned when the concurrency value is configured incorrectly ErrInvalidConcurrency = errors.New("Invalid concurrency") + + // ErrInvalidTransactionCountLimit is returned when the transaction count limit is configured incorrectly + ErrInvalidTransactionCountLimit = errors.New("Invalid transaction count limit") + + // ErrInvalidTransmissionSizeLimit is returned when the transmission size limit is configured incorrectly + ErrInvalidTransmissionSizeLimit = errors.New("Invalid transmission size limit") ) diff --git a/kcl_reader.go b/kcl_reader.go index 37a6e37..b49a9a6 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -343,16 +343,16 @@ func (r *KclReader) processRecords(ctx context.Context, batchSize int, fn messag // GetRecord calls processRecords to attempt to put one message from message buffer to the consumer's message // channel -func (r *KclReader) GetRecord(ctx context.Context, fn messageHandler) (count int, size int, err error) { - count, size, err = r.processRecords(ctx, 1, fn) - return count, size, err +func (r *KclReader) GetRecord(ctx context.Context, fn messageHandler) error { + _, _, err := r.processRecords(ctx, 1, fn) + return err } // GetRecords calls processRecords to attempt to put all messages on the message buffer on the consumer's // message channel -func (r *KclReader) GetRecords(ctx context.Context, fn messageHandler) (count int, size int, err error) { - count, size, err = r.processRecords(ctx, r.batchSize, fn) - return count, size, err +func (r *KclReader) GetRecords(ctx context.Context, fn messageHandler) error { + _, _, err := r.processRecords(ctx, r.batchSize, fn) + return err } // Checkpoint sends a message to KCL if there is sequence number that can be checkpointed diff --git a/kinesis_reader.go b/kinesis_reader.go index 201c5f7..1dd4ab5 100644 --- a/kinesis_reader.go +++ b/kinesis_reader.go @@ -3,6 +3,7 @@ package kinetic import ( "context" "fmt" + "sync" "time" "github.com/aws/aws-sdk-go/aws" @@ -11,31 +12,38 @@ import ( "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" + "golang.org/x/time/rate" ) const ( - kinesisReaderMaxBatchSize = 10000 - kinesisReaderDefaultConcurrency = 5 + kinesisReaderMaxBatchSize = 10000 + kinesisReaderDefaultConcurrency = 5 + kinesisReaderDefaultTransactionCountLimit = 5 + kinesisReaderDefaultTransmissionSizeLimit = 2000000 ) // kinesisReaderOptions a struct that holds all of the KinesisReader's configurable parameters. type kinesisReaderOptions struct { - batchSize int // maximum records per GetRecordsRequest call - concurrency int // maximum number of concurrent GetRecord or GetRecords calls allowed - shardIterator *ShardIterator // shard iterator for Kinesis stream - responseReadTimeout time.Duration // response read time out for GetRecordsRequest API call - logLevel aws.LogLevelType // log level for configuring the LogHelper's log level - Stats ConsumerStatsCollector // stats collection mechanism + batchSize int // maximum records per GetRecordsRequest call + concurrency int // maximum number of concurrent GetRecord or GetRecords calls allowed + transactionCountLimit int // maximum transactions per second for GetRecords calls + transmissionSizeLimit int // maximum transmission size per second for GetRecords calls + shardIterator *ShardIterator // shard iterator for Kinesis GetRecords API calls + responseReadTimeout time.Duration // response read time out for GetRecordsRequest API call + logLevel aws.LogLevelType // log level for configuring the LogHelper's log level + Stats ConsumerStatsCollector // stats collection mechanism } // defaultKinesisReaderOptions instantiates a kinesisReaderOptions with default values. func defaultKinesisReaderOptions() *kinesisReaderOptions { return &kinesisReaderOptions{ - batchSize: kinesisReaderMaxBatchSize, - concurrency: kinesisReaderDefaultConcurrency, - shardIterator: NewShardIterator(), - responseReadTimeout: time.Second, - Stats: &NilConsumerStatsCollector{}, + batchSize: kinesisReaderMaxBatchSize, + concurrency: kinesisReaderDefaultConcurrency, + transactionCountLimit: kinesisReaderDefaultTransactionCountLimit, + transmissionSizeLimit: kinesisReaderDefaultTransmissionSizeLimit, + shardIterator: NewShardIterator(), + responseReadTimeout: time.Second, + Stats: &NilConsumerStatsCollector{}, } } @@ -67,6 +75,30 @@ func KinesisReaderConcurrency(count int) KinesisReaderOptionsFn { } } +// KinesisReaderTransactionCountLimit is a functional option method for configuring the +// KinesisReader's transaction count limit +func KinesisReaderTransactionCountLimit(count int) KinesisReaderOptionsFn { + return func(o *KinesisReader) error { + if count > 0 { + o.transactionCountLimit = count + return nil + } + return ErrInvalidTransactionCountLimit + } +} + +// KinesisReaderTransmissionSizeLimit is a functional option method for configuring the +// KinesisReader's transmission size limit +func KinesisReaderTransmissionSizeLimit(size int) KinesisReaderOptionsFn { + return func(o *KinesisReader) error { + if size > 0 { + o.transmissionSizeLimit = size + return nil + } + return ErrInvalidTransmissionSizeLimit + } +} + // KinesisReaderShardIterator is a functional option method for configuring the KinesisReader's // shard iterator. func KinesisReaderShardIterator(shardIterator *ShardIterator) KinesisReaderOptionsFn { @@ -105,11 +137,13 @@ func KinesisReaderStats(sc ConsumerStatsCollector) KinesisReaderOptionsFn { type KinesisReader struct { *kinesisReaderOptions *LogHelper - stream string - shard string - throttleSem chan empty - nextShardIterator string - client kinesisiface.KinesisAPI + stream string // name of AWS Kinesis Stream to stream from + shard string // shardID of AWS Kinesis Stream to stream from + throttleSem chan empty // channel used to throttle concurrent GetRecord calls + txnCountRateLimiter *rate.Limiter // rate limiter to limit the number of transactions per second + txSizeRateLimiter *rate.Limiter // rate limiter to limit the transmission size per seccond + nextShardIterator string // shardIterator to start with with GetRecord request + client kinesisiface.KinesisAPI // client to Kinesis API } // NewKinesisReader creates a new KinesisReader object which implements the StreamReader interface to read records from @@ -131,6 +165,8 @@ func NewKinesisReader(c *aws.Config, stream string, shard string, optionFns ...K } kinesisReader.throttleSem = make(chan empty, kinesisReader.concurrency) + kinesisReader.txnCountRateLimiter = rate.NewLimiter(rate.Limit(kinesisReader.transactionCountLimit), 1) + kinesisReader.txSizeRateLimiter = rate.NewLimiter(rate.Limit(kinesisReader.transmissionSizeLimit), kinesisReader.transmissionSizeLimit) kinesisReader.LogHelper = &LogHelper{ LogLevel: kinesisReader.logLevel, Logger: c.Logger, @@ -321,17 +357,37 @@ func (r *KinesisReader) getRecords(ctx context.Context, fn messageHandler, batch // around. r.setNextShardIterator(*resp.NextShardIterator) } + + wg := sync.WaitGroup{} + wg.Add(1) + go func() { + defer wg.Done() + + if err := r.txnCountRateLimiter.Wait(ctx); err != nil { + r.LogError("Error occured waiting for transaction count tokens") + } + }() + + wg.Add(1) + go func() { + defer wg.Done() + if err := r.txSizeRateLimiter.WaitN(ctx, size); err != nil { + r.LogError("Error occured waiting for transmission size tokens") + } + }() + wg.Wait() + return count, size, nil } // GetRecord calls getRecords and delivers one record into the messages channel. -func (r *KinesisReader) GetRecord(ctx context.Context, fn messageHandler) (count int, size int, err error) { - count, size, err = r.getRecords(ctx, fn, 1) - return count, size, err +func (r *KinesisReader) GetRecord(ctx context.Context, fn messageHandler) error { + _, _, err := r.getRecords(ctx, fn, 1) + return err } // GetRecords calls getRecords and delivers each record into the messages channel. -func (r *KinesisReader) GetRecords(ctx context.Context, fn messageHandler) (count int, size int, err error) { - count, size, err = r.getRecords(ctx, fn, r.batchSize) - return count, size, err +func (r *KinesisReader) GetRecords(ctx context.Context, fn messageHandler) error { + _, _, err := r.getRecords(ctx, fn, r.batchSize) + return err } diff --git a/shared_types.go b/shared_types.go index d356714..dfb0a58 100644 --- a/shared_types.go +++ b/shared_types.go @@ -14,8 +14,8 @@ type StreamWriter interface { // StreamReader is an interface that abstracts out a stream reader. type StreamReader interface { - GetRecord(context.Context, messageHandler) (count int, size int, err error) - GetRecords(context.Context, messageHandler) (count int, size int, err error) + GetRecord(context.Context, messageHandler) error + GetRecords(context.Context, messageHandler) error } // MessageProcessor defines the signature of a (asynchronous) callback function used by Listen, RetrieveFn and diff --git a/testexec/main.go b/testexec/main.go index adc6743..7b4632c 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -282,6 +282,9 @@ func newKineticConsumer(k *kinetic.Kinetic, streamName string) *kinetic.Consumer csc := kinetic.NewDefaultConsumerStatsCollector(registry) r, err := kinetic.NewKinesisReader(k.Session.Config, streamName, shards[0], kinetic.KinesisReaderBatchSize(10000), + kinetic.KinesisReaderConcurrency(5), + kinetic.KinesisReaderTransactionCountLimit(5), + kinetic.KinesisReaderTransmissionSizeLimit(2000000), //kinetic.KinesisReaderShardIterator(), kinetic.KinesisReaderResponseReadTimeout(time.Second), kinetic.KinesisReaderLogLevel(aws.LogOff), From 0176c832db4fb022b1d7c710bf9f333f68469569 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 30 Aug 2017 13:36:04 -0500 Subject: [PATCH 091/109] Remove throttle function from KinesisReader as it is now redundant with the count/size rate limiter. --- kinesis_reader.go | 28 ---------------------------- testexec/main.go | 1 - 2 files changed, 29 deletions(-) diff --git a/kinesis_reader.go b/kinesis_reader.go index 1dd4ab5..9bfc2e4 100644 --- a/kinesis_reader.go +++ b/kinesis_reader.go @@ -17,7 +17,6 @@ import ( const ( kinesisReaderMaxBatchSize = 10000 - kinesisReaderDefaultConcurrency = 5 kinesisReaderDefaultTransactionCountLimit = 5 kinesisReaderDefaultTransmissionSizeLimit = 2000000 ) @@ -25,7 +24,6 @@ const ( // kinesisReaderOptions a struct that holds all of the KinesisReader's configurable parameters. type kinesisReaderOptions struct { batchSize int // maximum records per GetRecordsRequest call - concurrency int // maximum number of concurrent GetRecord or GetRecords calls allowed transactionCountLimit int // maximum transactions per second for GetRecords calls transmissionSizeLimit int // maximum transmission size per second for GetRecords calls shardIterator *ShardIterator // shard iterator for Kinesis GetRecords API calls @@ -38,7 +36,6 @@ type kinesisReaderOptions struct { func defaultKinesisReaderOptions() *kinesisReaderOptions { return &kinesisReaderOptions{ batchSize: kinesisReaderMaxBatchSize, - concurrency: kinesisReaderDefaultConcurrency, transactionCountLimit: kinesisReaderDefaultTransactionCountLimit, transmissionSizeLimit: kinesisReaderDefaultTransmissionSizeLimit, shardIterator: NewShardIterator(), @@ -63,18 +60,6 @@ func KinesisReaderBatchSize(size int) KinesisReaderOptionsFn { } } -// KinesisReaderConcurrency is a functional option method for configuring the KinesisReader's -// concurrency. -func KinesisReaderConcurrency(count int) KinesisReaderOptionsFn { - return func(o *KinesisReader) error { - if count > 0 { - o.concurrency = count - return nil - } - return ErrInvalidConcurrency - } -} - // KinesisReaderTransactionCountLimit is a functional option method for configuring the // KinesisReader's transaction count limit func KinesisReaderTransactionCountLimit(count int) KinesisReaderOptionsFn { @@ -139,7 +124,6 @@ type KinesisReader struct { *LogHelper stream string // name of AWS Kinesis Stream to stream from shard string // shardID of AWS Kinesis Stream to stream from - throttleSem chan empty // channel used to throttle concurrent GetRecord calls txnCountRateLimiter *rate.Limiter // rate limiter to limit the number of transactions per second txSizeRateLimiter *rate.Limiter // rate limiter to limit the transmission size per seccond nextShardIterator string // shardIterator to start with with GetRecord request @@ -164,7 +148,6 @@ func NewKinesisReader(c *aws.Config, stream string, shard string, optionFns ...K optionFn(kinesisReader) } - kinesisReader.throttleSem = make(chan empty, kinesisReader.concurrency) kinesisReader.txnCountRateLimiter = rate.NewLimiter(rate.Limit(kinesisReader.transactionCountLimit), 1) kinesisReader.txSizeRateLimiter = rate.NewLimiter(rate.Limit(kinesisReader.transmissionSizeLimit), kinesisReader.transmissionSizeLimit) kinesisReader.LogHelper = &LogHelper{ @@ -232,23 +215,12 @@ func (r *KinesisReader) setSequenceNumber(sequenceNumber string) error { return nil } -// Kinesis allows five read ops per second per shard. -// http://docs.aws.amazon.com/kinesis/latest/dev/service-sizes-and-limits.html -func (r *KinesisReader) throttle(sem chan empty) { - sem <- empty{} - time.AfterFunc(1*time.Second, func() { - <-sem - }) -} - func (r *KinesisReader) getRecords(ctx context.Context, fn messageHandler, batchSize int) (count int, size int, err error) { if err = r.ensureShardIterator(); err != nil { r.LogError("Error calling ensureShardIterator(): ", err) return count, size, err } - r.throttle(r.throttleSem) - // We use the GetRecordsRequest method of creating requests to allow for registering custom handlers for better // control over the API request. var startReadTime time.Time diff --git a/testexec/main.go b/testexec/main.go index 7b4632c..2b10b2c 100644 --- a/testexec/main.go +++ b/testexec/main.go @@ -282,7 +282,6 @@ func newKineticConsumer(k *kinetic.Kinetic, streamName string) *kinetic.Consumer csc := kinetic.NewDefaultConsumerStatsCollector(registry) r, err := kinetic.NewKinesisReader(k.Session.Config, streamName, shards[0], kinetic.KinesisReaderBatchSize(10000), - kinetic.KinesisReaderConcurrency(5), kinetic.KinesisReaderTransactionCountLimit(5), kinetic.KinesisReaderTransmissionSizeLimit(2000000), //kinetic.KinesisReaderShardIterator(), From e8ffae22028c2957c7dd8a719944714f87f65cbd Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 2 Oct 2017 10:53:55 -0500 Subject: [PATCH 092/109] Removed unnecessary CloseWithContext function. --- producer.go | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/producer.go b/producer.go index cc559d5..58d0d4d 100644 --- a/producer.go +++ b/producer.go @@ -500,17 +500,9 @@ func (p *Producer) sendToDataSpill(msg *Message) { } } -// CloseWithContext initiates the graceful shutdown of the produce function, waiting for all outstanding messages and to -// flush. Cancellation is supported through contexts. -func (p *Producer) CloseWithContext(ctx context.Context) { - p.shutdown() - <-ctx.Done() - close(p.pipeOfDeath) -} - // Close initiates the graceful shutdown of the produce function, waiting for all outstanding messages and to flush. func (p *Producer) Close() { - p.CloseWithContext(context.TODO()) + p.shutdown() } // SendWithContext sends a message to the stream. Cancellation supported through contexts. From 116fb8c29cff61e32ed54abad95c0b94c0431f64 Mon Sep 17 00:00:00 2001 From: Ian Chiles Date: Mon, 2 Oct 2017 13:39:13 -0500 Subject: [PATCH 093/109] strip out pipe of death --- consumer.go | 9 ++------- consumer_test.go | 10 ++++++---- errors.go | 3 --- producer.go | 9 --------- 4 files changed, 8 insertions(+), 23 deletions(-) diff --git a/consumer.go b/consumer.go index fb6c9f7..b7dd78e 100644 --- a/consumer.go +++ b/consumer.go @@ -82,7 +82,6 @@ type Consumer struct { *LogHelper // object for help with logging messages chan *Message // channel for storing messages that have been retrieved from the stream concurrencySem chan empty // channel for controlling the number of concurrent workers processing messages from the message channel - pipeOfDeath chan empty // channel for handling pipe of death consuming bool // flag for indicating whether or not the consumer is consuming consumingMu sync.Mutex // mutex for making the consuming flag thread safe noCopy noCopy // prevents the Consumer from being copied @@ -120,7 +119,6 @@ func (c *Consumer) startConsuming() bool { c.consuming = true c.messages = make(chan *Message, c.queueDepth) c.concurrencySem = make(chan empty, c.concurrency) - c.pipeOfDeath = make(chan empty) return true } return false @@ -130,8 +128,6 @@ func (c *Consumer) startConsuming() bool { // cancellation or a pipe of death. func (c *Consumer) shouldConsume(ctx context.Context) (bool, error) { select { - case <-c.pipeOfDeath: - return false, ErrPipeOfDeath case <-ctx.Done(): return false, ctx.Err() default: @@ -299,6 +295,8 @@ func (c *Consumer) ListenWithContext(ctx context.Context, fn MessageProcessor) { for { select { + case <-ctx.Done(): + return case msg, ok := <-c.messages: if !ok { return @@ -318,9 +316,6 @@ func (c *Consumer) ListenWithContext(ctx context.Context, fn MessageProcessor) { c.Stats.AddProcessed(1) wg.Done() }(msg) - case <-c.pipeOfDeath: - c.LogInfo("ListenWithContext received pipe of death") - return } } } diff --git a/consumer_test.go b/consumer_test.go index ed48d42..4456cfb 100644 --- a/consumer_test.go +++ b/consumer_test.go @@ -201,15 +201,15 @@ func TestConsumer(t *testing.T) { Convey("check that listen and retrieve can not be called concurrently", func(c C) { var wg sync.WaitGroup wg.Add(1) + ctx, cancel := context.WithTimeout(context.TODO(), 1000*time.Millisecond) go func() { - ctx, cancel := context.WithTimeout(context.TODO(), 1000*time.Millisecond) - defer cancel() l.ListenWithContext(ctx, func(msg *Message) error { return nil }) wg.Done() }() <-time.After(10 * time.Millisecond) + cancel() _, err := l.Retrieve() So(err, ShouldEqual, ErrAlreadyConsuming) wg.Wait() @@ -252,9 +252,11 @@ func TestConsumer(t *testing.T) { var count int64 var wg sync.WaitGroup wg.Add(1) + + ctx, cancel := context.WithCancel(context.TODO()) go func() { defer wg.Done() - l.Listen(func(msg *Message) error { + l.ListenWithContext(ctx, func(msg *Message) error { atomic.AddInt64(&count, 1) return nil @@ -285,7 +287,7 @@ func TestConsumer(t *testing.T) { } // FIXME: probably a race condition here as consume may // not have grabbed all data from the channel yet. - close(l.pipeOfDeath) + cancel() wg.Wait() So(atomic.LoadInt64(&count), ShouldEqual, len(planets)) }) diff --git a/errors.go b/errors.go index 76f8b09..e975014 100644 --- a/errors.go +++ b/errors.go @@ -13,9 +13,6 @@ var ( // DescribeStream request returns a response with a nil // StreamDescription ErrNilStreamDescription = errors.New("DescribeStream returned a nil StreamDescription") - - // ErrPipeOfDeath returns when the pipe of death is closed. - ErrPipeOfDeath = errors.New("Received pipe of death") ) var ( diff --git a/producer.go b/producer.go index 58d0d4d..cb2729d 100644 --- a/producer.go +++ b/producer.go @@ -150,7 +150,6 @@ type Producer struct { status chan *statusReport // channel for workers to communicate their current status dismiss chan empty // channel for handling the decommissioning of a surplus of workers stop chan empty // channel for handling shutdown - pipeOfDeath chan empty // channel for handling pipe of death startupOnce sync.Once // used to ensure that the startup function is called once shutdownOnce sync.Once // used to ensure that the shutdown function is called once resizeMu sync.Mutex // used to prevent resizeWorkerPool from being called synchronously with itself @@ -197,7 +196,6 @@ func (p *Producer) produce() { p.status = make(chan *statusReport) p.dismiss = make(chan empty) p.stop = make(chan empty) - p.pipeOfDeath = make(chan empty) // Run a separate go routine to check the shard size (throughput multiplier) and resize the worker pool // periodically if needed @@ -239,8 +237,6 @@ func (p *Producer) produce() { for { select { - case <-p.pipeOfDeath: - return case <-p.stop: return case status := <-p.status: @@ -368,11 +364,6 @@ func (p *Producer) shutdown() { p.stop <- empty{} close(p.stop) } - - // Close the pipeOfDeath channel - if p.pipeOfDeath != nil { - close(p.pipeOfDeath) - } }) } From d48871e40c71279e616de8879a159fb6a2f1f513 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 4 Oct 2017 15:41:51 -0500 Subject: [PATCH 094/109] Define reader as a bufio.NewReader that reads from STDIN --- checkpointer.go | 4 ++-- kcl_reader.go | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/checkpointer.go b/checkpointer.go index 3d65e86..65d1655 100644 --- a/checkpointer.go +++ b/checkpointer.go @@ -55,7 +55,7 @@ func (c *checkpointList) find(seqNum int) (*list.Element, bool) { case element.seqNum == seqNum: return e, true case element.seqNum < seqNum: - break + continue default: return nil, false } @@ -176,7 +176,7 @@ func (c *checkpointer) startup(ctx context.Context) { case <-autoCheckpointTimer.C: break wait case <-counterCheckTicker.C: - break + continue } } diff --git a/kcl_reader.go b/kcl_reader.go index b49a9a6..d5c67bd 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -145,6 +145,8 @@ func NewKclReader(c *aws.Config, optionFns ...KclReaderOptionsFn) (*KclReader, e optionFn(kclReader) } + kclReader.reader = bufio.NewReader(os.Stdin) + kclReader.LogHelper = &LogHelper{ LogLevel: kclReader.logLevel, Logger: c.Logger, From 83c97093926d8f45732284250f54ed9d7206ed75 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 4 Oct 2017 15:55:02 -0500 Subject: [PATCH 095/109] FU staticcheck --- circle.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/circle.yml b/circle.yml index 4ffb89a..c0baca0 100644 --- a/circle.yml +++ b/circle.yml @@ -10,7 +10,7 @@ dependencies: pre: - bash < <(curl -s -S -L https://raw.githubusercontent.com/moovweb/gvm/1.0.22/binscripts/gvm-installer) - source "$HOME/.gvm/scripts/gvm" && gvm install go$GO_VERSION -B --name=stable && gvm use stable --default - - go get honnef.co/go/tools/cmd/staticcheck +# - go get honnef.co/go/tools/cmd/staticcheck - go get honnef.co/go/tools/cmd/gosimple - go get -u github.com/golang/lint/golint - docker run -d -t -p 4567:4567 dlsniper/kinesalite:1.8.0 From bfbb8a2ba384deb7b4370733a87dcbd537ae8460 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Wed, 4 Oct 2017 16:03:53 -0500 Subject: [PATCH 096/109] staticcheck and gosimple site is down --- circle.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/circle.yml b/circle.yml index c0baca0..6d14763 100644 --- a/circle.yml +++ b/circle.yml @@ -11,7 +11,7 @@ dependencies: - bash < <(curl -s -S -L https://raw.githubusercontent.com/moovweb/gvm/1.0.22/binscripts/gvm-installer) - source "$HOME/.gvm/scripts/gvm" && gvm install go$GO_VERSION -B --name=stable && gvm use stable --default # - go get honnef.co/go/tools/cmd/staticcheck - - go get honnef.co/go/tools/cmd/gosimple +# - go get honnef.co/go/tools/cmd/gosimple - go get -u github.com/golang/lint/golint - docker run -d -t -p 4567:4567 dlsniper/kinesalite:1.8.0 override: @@ -25,7 +25,7 @@ test: pre: - cd "$GOPATH/src/$IMPORT_PATH" && go vet $(glide nv) - cd "$GOPATH/src/$IMPORT_PATH" && golint -set_exit_status $(glide nv) - - cd "$GOPATH/src/$IMPORT_PATH" && gosimple $(glide nv) +# - cd "$GOPATH/src/$IMPORT_PATH" && gosimple $(glide nv) override: - cd $GOPATH/src/$IMPORT_PATH && go test -v -race -cover $(glide nv) From 3c41824404a2e00913d1ea4a1ff5e69d42ea7451 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 5 Oct 2017 14:35:27 -0500 Subject: [PATCH 097/109] Added TeeReader for debugging --- kcl_reader.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/kcl_reader.go b/kcl_reader.go index d5c67bd..1f645b3 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -11,6 +11,7 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" + "io" ) const ( @@ -145,7 +146,9 @@ func NewKclReader(c *aws.Config, optionFns ...KclReaderOptionsFn) (*KclReader, e optionFn(kclReader) } - kclReader.reader = bufio.NewReader(os.Stdin) + f, _ := os.Open("this_file.txt") + r := io.TeeReader(os.Stdin, f) + kclReader.reader = bufio.NewReader(r) kclReader.LogHelper = &LogHelper{ LogLevel: kclReader.logLevel, From fde947e5a7c7ef53cbb038dfac35965a77db376a Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 5 Oct 2017 15:41:49 -0500 Subject: [PATCH 098/109] Added TeeReader for debugging --- kcl_reader.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kcl_reader.go b/kcl_reader.go index 1f645b3..5fe0799 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -146,7 +146,7 @@ func NewKclReader(c *aws.Config, optionFns ...KclReaderOptionsFn) (*KclReader, e optionFn(kclReader) } - f, _ := os.Open("this_file.txt") + f, _ := os.Create("this_file.txt") r := io.TeeReader(os.Stdin, f) kclReader.reader = bufio.NewReader(r) From e79cb77dc8402f03090d12aafcae18900d4d91bb Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 5 Oct 2017 15:49:11 -0500 Subject: [PATCH 099/109] Print errors with file creation --- kcl_reader.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/kcl_reader.go b/kcl_reader.go index 5fe0799..549e50d 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -146,7 +146,10 @@ func NewKclReader(c *aws.Config, optionFns ...KclReaderOptionsFn) (*KclReader, e optionFn(kclReader) } - f, _ := os.Create("this_file.txt") + f, err := os.Create("this_file.txt") + if err != nil { + panic(err) + } r := io.TeeReader(os.Stdin, f) kclReader.reader = bufio.NewReader(r) From c08d8e4bcb19ecc72772155cf3d59c5ac841c548 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 5 Oct 2017 15:54:44 -0500 Subject: [PATCH 100/109] Panic --- kcl_reader.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kcl_reader.go b/kcl_reader.go index 549e50d..09e4446 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -212,7 +212,7 @@ func (r *KclReader) process(ctx context.Context) { // Retrieve action message actionMessage, err := r.receiveFromStdIn() if err != nil || actionMessage == nil { - return + panic(err) } switch actionMessage.Action { From 854c5f60d6293b4da2696f6bd52240350846a3ae Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 6 Oct 2017 11:08:29 -0500 Subject: [PATCH 101/109] Revert TeeReader stuff (for debugging) and fix megacheck errors. --- consumer.go | 2 +- glide.lock | 6 +++--- kcl_actions.go | 1 - kcl_reader.go | 12 +++--------- producer.go | 2 +- 5 files changed, 8 insertions(+), 15 deletions(-) diff --git a/consumer.go b/consumer.go index b7dd78e..7d96249 100644 --- a/consumer.go +++ b/consumer.go @@ -84,7 +84,7 @@ type Consumer struct { concurrencySem chan empty // channel for controlling the number of concurrent workers processing messages from the message channel consuming bool // flag for indicating whether or not the consumer is consuming consumingMu sync.Mutex // mutex for making the consuming flag thread safe - noCopy noCopy // prevents the Consumer from being copied + _ noCopy // prevents the Consumer from being copied } // NewConsumer creates a new Consumer object for retrieving and listening to message(s) on a StreamReader. diff --git a/glide.lock b/glide.lock index 220600a..f0e8d7c 100644 --- a/glide.lock +++ b/glide.lock @@ -1,5 +1,5 @@ hash: f6ee7c06e2bb454e7b5cd139b007893526a80c6e5856e0a1c24ad03a086949c5 -updated: 2017-08-16T16:46:40.529793303-05:00 +updated: 2017-10-05T14:33:54.979115329-05:00 imports: - name: github.com/aws/aws-sdk-go version: 72e42b13da62269f68308fb6068b7ea691a416a4 @@ -40,11 +40,11 @@ imports: - name: github.com/rcrowley/go-metrics version: ab2277b1c5d15c3cba104e9cbddbdfc622df5ad8 - name: golang.org/x/net - version: 1c05540f6879653db88113bc4a2b70aec4bd491f + version: a04bdaca5b32abe1c069418fb7088ae607de5bd0 subpackages: - context - name: golang.org/x/time - version: 8be79e1e0910c292df4e79c241bb7e8f7e725959 + version: 6dc17368e09b0e8634d71cac8168d853e869a0c7 subpackages: - rate testImports: diff --git a/kcl_actions.go b/kcl_actions.go index ef7a3d0..a788aa8 100644 --- a/kcl_actions.go +++ b/kcl_actions.go @@ -13,7 +13,6 @@ type kclActionType string const ( kclActionTypeInitialize kclActionType = "initialize" kclActionTypeProcessRecords kclActionType = "processRecords" - kclActionTypeRecord kclActionType = "record" kclActionTypeCheckpoint kclActionType = "checkpoint" kcActionTypeShutdown kclActionType = "shutdown" KclActionTypeStatus kclActionType = "status" diff --git a/kcl_reader.go b/kcl_reader.go index 09e4446..562bbaf 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -11,7 +11,6 @@ import ( "time" "github.com/aws/aws-sdk-go/aws" - "io" ) const ( @@ -46,8 +45,8 @@ func defaultKclReaderOptions() *kclReaderOptions { // KclReaderOptionsFn is a method signature for defining functional option methods for configuring the KclReader. type KclReaderOptionsFn func(*KclReader) error -// kclReaderBatchSize is a functional option method for configuring the KclReader's batch size -func kclReaderBatchSize(size int) KclReaderOptionsFn { +// KclReaderBatchSize is a functional option method for configuring the KclReader's batch size +func KclReaderBatchSize(size int) KclReaderOptionsFn { return func(o *KclReader) error { if size > 0 && size <= kclReaderMaxBatchSize { o.batchSize = size @@ -146,12 +145,7 @@ func NewKclReader(c *aws.Config, optionFns ...KclReaderOptionsFn) (*KclReader, e optionFn(kclReader) } - f, err := os.Create("this_file.txt") - if err != nil { - panic(err) - } - r := io.TeeReader(os.Stdin, f) - kclReader.reader = bufio.NewReader(r) + kclReader.reader = bufio.NewReader(os.Stdin) kclReader.LogHelper = &LogHelper{ LogLevel: kclReader.logLevel, diff --git a/producer.go b/producer.go index cb2729d..dd40c8c 100644 --- a/producer.go +++ b/producer.go @@ -153,7 +153,7 @@ type Producer struct { startupOnce sync.Once // used to ensure that the startup function is called once shutdownOnce sync.Once // used to ensure that the shutdown function is called once resizeMu sync.Mutex // used to prevent resizeWorkerPool from being called synchronously with itself - noCopy noCopy // prevents the Producer from being copied + _ noCopy // prevents the Producer from being copied } // NewProducer creates a new producer for writing records to a Kinesis or Firehose stream. From 4eaca887f160df8c998488957b1e5cb6e68f91f7 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 6 Oct 2017 11:21:37 -0500 Subject: [PATCH 102/109] Simplified Consumer.ListenWithContext by eliminatin the select statement --- consumer.go | 38 +++++++++++++++++--------------------- 1 file changed, 17 insertions(+), 21 deletions(-) diff --git a/consumer.go b/consumer.go index 7d96249..d5103f6 100644 --- a/consumer.go +++ b/consumer.go @@ -294,29 +294,25 @@ func (c *Consumer) ListenWithContext(ctx context.Context, fn MessageProcessor) { defer wg.Wait() for { - select { - case <-ctx.Done(): + msg, ok := <-c.messages + if !ok { return - case msg, ok := <-c.messages: - if !ok { - return - } - c.Stats.AddDelivered(1) - // For simplicity, did not do the pipe of death here. If POD is received, we may deliver a - // couple more messages (especially since select is random in which channel is read from). - c.concurrencySem <- empty{} - wg.Add(1) - go func(msg *Message) { - defer func() { - <-c.concurrencySem - }() - start := time.Now() - fn(msg) - c.Stats.UpdateProcessedDuration(time.Since(start)) - c.Stats.AddProcessed(1) - wg.Done() - }(msg) } + c.Stats.AddDelivered(1) + // For simplicity, did not do the pipe of death here. If POD is received, we may deliver a + // couple more messages (especially since select is random in which channel is read from). + c.concurrencySem <- empty{} + wg.Add(1) + go func(msg *Message) { + defer func() { + <-c.concurrencySem + }() + start := time.Now() + fn(msg) + c.Stats.UpdateProcessedDuration(time.Since(start)) + c.Stats.AddProcessed(1) + wg.Done() + }(msg) } } From 94e63ce93638a761196436aa9c7e052cdcfa87ef Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 9 Oct 2017 10:53:49 -0500 Subject: [PATCH 103/109] Uncomment out staticcheck and gosimple dependencies in circle.yml file. --- circle.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/circle.yml b/circle.yml index 6d14763..4ffb89a 100644 --- a/circle.yml +++ b/circle.yml @@ -10,8 +10,8 @@ dependencies: pre: - bash < <(curl -s -S -L https://raw.githubusercontent.com/moovweb/gvm/1.0.22/binscripts/gvm-installer) - source "$HOME/.gvm/scripts/gvm" && gvm install go$GO_VERSION -B --name=stable && gvm use stable --default -# - go get honnef.co/go/tools/cmd/staticcheck -# - go get honnef.co/go/tools/cmd/gosimple + - go get honnef.co/go/tools/cmd/staticcheck + - go get honnef.co/go/tools/cmd/gosimple - go get -u github.com/golang/lint/golint - docker run -d -t -p 4567:4567 dlsniper/kinesalite:1.8.0 override: @@ -25,7 +25,7 @@ test: pre: - cd "$GOPATH/src/$IMPORT_PATH" && go vet $(glide nv) - cd "$GOPATH/src/$IMPORT_PATH" && golint -set_exit_status $(glide nv) -# - cd "$GOPATH/src/$IMPORT_PATH" && gosimple $(glide nv) + - cd "$GOPATH/src/$IMPORT_PATH" && gosimple $(glide nv) override: - cd $GOPATH/src/$IMPORT_PATH && go test -v -race -cover $(glide nv) From e266d0a229c252c395488992b6aa4f9df981c177 Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Mon, 9 Oct 2017 15:52:31 -0500 Subject: [PATCH 104/109] Add default KclReader callback functions --- kcl_reader.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/kcl_reader.go b/kcl_reader.go index 562bbaf..2fd8a63 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -37,6 +37,9 @@ func defaultKclReaderOptions() *kclReaderOptions { autoCheckpointCount: 10000, autoCheckpointFreq: time.Minute, updateCheckpointSizeFreq: time.Minute, + onInitCallbackFn: func() error { return nil }, + onCheckpointCallbackFn: func(a, b string) error { return nil }, + onShutdownCallbackFn: func() error { return nil }, logLevel: aws.LogOff, Stats: &NilConsumerStatsCollector{}, } From 3b561fac9229012005193aace65b8c3869641f0c Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Thu, 26 Oct 2017 23:30:18 -0500 Subject: [PATCH 105/109] Added a bufio writer for KclReader to write to Stdout. --- kcl_reader.go | 27 ++++++++++++++++----------- kinesis_reader.go | 2 +- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/kcl_reader.go b/kcl_reader.go index 2fd8a63..dc77aad 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -129,16 +129,17 @@ func KclReaderStats(sc ConsumerStatsCollector) KclReaderOptionsFn { // KclReader handles the KCL Multilang Protocol to read records from KCL type KclReader struct { - *kclReaderOptions // contains all of the configuration settings for the KclReader - *LogHelper // object for help with logging - reader *bufio.Reader // io reader to read from STDIN - checkpointer *checkpointer // data structure used to manage checkpointing - ticker *time.Ticker // a ticker with which to update the CheckpointSize stats - tickerDone chan empty // a channel used to communicate when to stop updating the CheckpointSize stats - messages chan *Message // unbuffered message channel used to throttle the record processing from KCL - actions chan *actionMessage // unbuffered action message channel used internally to coordinate sending action messages to KCL - startupOnce sync.Once // used to ensure that the startup function is called once - shutdownOnce sync.Once // used to ensure that the shutdown function is called once + *kclReaderOptions // contains all of the configuration settings for the KclReader + *LogHelper // object for help with logging + reader *bufio.Reader // buffered reader to read messages from KCL + writer *bufio.Writer // buffered writer to write messages to KCL + checkpointer *checkpointer // data structure used to manage checkpointing + ticker *time.Ticker // a ticker with which to update the CheckpointSize stats + tickerDone chan empty // a channel used to communicate when to stop updating the CheckpointSize stats + messages chan *Message // unbuffered message channel used to throttle the record processing from KCL + actions chan *actionMessage // unbuffered action message channel used internally to coordinate sending action messages to KCL + startupOnce sync.Once // used to ensure that the startup function is called once + shutdownOnce sync.Once // used to ensure that the shutdown function is called once } // NewKclReader creates a new stream reader to read records from KCL @@ -148,7 +149,9 @@ func NewKclReader(c *aws.Config, optionFns ...KclReaderOptionsFn) (*KclReader, e optionFn(kclReader) } + // Setup a buffered reader/writer from the io reader/writer for communicating via the Multilang Daemon Protocol kclReader.reader = bufio.NewReader(os.Stdin) + kclReader.writer = bufio.NewWriter(os.Stdout) kclReader.LogHelper = &LogHelper{ LogLevel: kclReader.logLevel, @@ -315,7 +318,9 @@ func (r *KclReader) sendToStdOut(msg interface{}) error { return err } - fmt.Fprintln(os.Stdout, string(b)) + r.writer.Write(b) + fmt.Println(r.writer, string(b)) + r.writer.Flush() return nil } diff --git a/kinesis_reader.go b/kinesis_reader.go index 9bfc2e4..511c226 100644 --- a/kinesis_reader.go +++ b/kinesis_reader.go @@ -23,7 +23,7 @@ const ( // kinesisReaderOptions a struct that holds all of the KinesisReader's configurable parameters. type kinesisReaderOptions struct { - batchSize int // maximum records per GetRecordsRequest call + batchSize int // maximum records per GetRecords call transactionCountLimit int // maximum transactions per second for GetRecords calls transmissionSizeLimit int // maximum transmission size per second for GetRecords calls shardIterator *ShardIterator // shard iterator for Kinesis GetRecords API calls From e0d36894cc2b30fe63ec30c8f91d23437dbc38ea Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 27 Oct 2017 00:26:32 -0500 Subject: [PATCH 106/109] Added function option methods for configuring the KclReader's reader/writer. --- kcl_reader.go | 39 ++++++++++++++++++++++++++++++--------- 1 file changed, 30 insertions(+), 9 deletions(-) diff --git a/kcl_reader.go b/kcl_reader.go index dc77aad..e44c6fc 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -6,6 +6,7 @@ import ( "context" "encoding/json" "fmt" + "io" "os" "sync" "time" @@ -19,7 +20,9 @@ const ( // kclReaderOptions is a struct that holds all of the KclReader's configurable parameters. type kclReaderOptions struct { - batchSize int + batchSize int // maximum records per GetRecords call + reader io.Reader // reader for reading from KCL + writer io.Writer // writer for writing to KCL autoCheckpointCount int // maximum number of messages pulled off the message queue before triggering an auto checkpoint autoCheckpointFreq time.Duration // frequency with which to auto checkpoint updateCheckpointSizeFreq time.Duration // frequency with which to update the CheckpointSize stats @@ -34,6 +37,8 @@ type kclReaderOptions struct { func defaultKclReaderOptions() *kclReaderOptions { return &kclReaderOptions{ batchSize: kclReaderMaxBatchSize, + reader: os.Stdin, + writer: os.Stdout, autoCheckpointCount: 10000, autoCheckpointFreq: time.Minute, updateCheckpointSizeFreq: time.Minute, @@ -59,6 +64,22 @@ func KclReaderBatchSize(size int) KclReaderOptionsFn { } } +// KclReaderReader is a functional option method for configuring the KclReader's reader (defaults to os.Stdin). +func KclReaderReader(reader io.Reader) KclReaderOptionsFn { + return func(o *KclReader) error { + o.reader = reader + return nil + } +} + +// KclReaderWriter is a functional option method for configuring the KclReader's writer (defaults to os.Stdout). +func KclReaderWriter(writer io.Writer) KclReaderOptionsFn { + return func(o *KclReader) error { + o.writer = writer + return nil + } +} + // KclReaderAutoCheckpointCount is a functional option method for configuring the KclReader's checkpoint count func KclReaderAutoCheckpointCount(count int) KclReaderOptionsFn { return func(o *KclReader) error { @@ -131,8 +152,8 @@ func KclReaderStats(sc ConsumerStatsCollector) KclReaderOptionsFn { type KclReader struct { *kclReaderOptions // contains all of the configuration settings for the KclReader *LogHelper // object for help with logging - reader *bufio.Reader // buffered reader to read messages from KCL - writer *bufio.Writer // buffered writer to write messages to KCL + bufReader *bufio.Reader // buffered reader to read messages from KCL + burWriter *bufio.Writer // buffered writer to write messages to KCL checkpointer *checkpointer // data structure used to manage checkpointing ticker *time.Ticker // a ticker with which to update the CheckpointSize stats tickerDone chan empty // a channel used to communicate when to stop updating the CheckpointSize stats @@ -150,8 +171,8 @@ func NewKclReader(c *aws.Config, optionFns ...KclReaderOptionsFn) (*KclReader, e } // Setup a buffered reader/writer from the io reader/writer for communicating via the Multilang Daemon Protocol - kclReader.reader = bufio.NewReader(os.Stdin) - kclReader.writer = bufio.NewWriter(os.Stdout) + kclReader.bufReader = bufio.NewReader(kclReader.reader) + kclReader.burWriter = bufio.NewWriter(kclReader.writer) kclReader.LogHelper = &LogHelper{ LogLevel: kclReader.logLevel, @@ -290,7 +311,7 @@ func (r *KclReader) shutdown() { func (r *KclReader) receiveFromStdIn() (*actionMessage, error) { buffer := &bytes.Buffer{} for { - line, isPrefix, err := r.reader.ReadLine() + line, isPrefix, err := r.bufReader.ReadLine() if err != nil { r.LogError("Unable to read line from stdin ", err) return nil, err @@ -318,9 +339,9 @@ func (r *KclReader) sendToStdOut(msg interface{}) error { return err } - r.writer.Write(b) - fmt.Println(r.writer, string(b)) - r.writer.Flush() + r.burWriter.Write(b) + fmt.Println(r.burWriter, string(b)) + r.burWriter.Flush() return nil } From f3225b0db452222fda593c4c89a95bf9cb96b2ca Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 27 Oct 2017 10:20:00 -0500 Subject: [PATCH 107/109] Change from fmt.Println to fmt.FPrintln to write to the correct io.Writer --- kcl_reader.go | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/kcl_reader.go b/kcl_reader.go index e44c6fc..3eb9d42 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -153,7 +153,7 @@ type KclReader struct { *kclReaderOptions // contains all of the configuration settings for the KclReader *LogHelper // object for help with logging bufReader *bufio.Reader // buffered reader to read messages from KCL - burWriter *bufio.Writer // buffered writer to write messages to KCL + bufWriter *bufio.Writer // buffered writer to write messages to KCL checkpointer *checkpointer // data structure used to manage checkpointing ticker *time.Ticker // a ticker with which to update the CheckpointSize stats tickerDone chan empty // a channel used to communicate when to stop updating the CheckpointSize stats @@ -172,7 +172,7 @@ func NewKclReader(c *aws.Config, optionFns ...KclReaderOptionsFn) (*KclReader, e // Setup a buffered reader/writer from the io reader/writer for communicating via the Multilang Daemon Protocol kclReader.bufReader = bufio.NewReader(kclReader.reader) - kclReader.burWriter = bufio.NewWriter(kclReader.writer) + kclReader.bufWriter = bufio.NewWriter(kclReader.writer) kclReader.LogHelper = &LogHelper{ LogLevel: kclReader.logLevel, @@ -339,9 +339,9 @@ func (r *KclReader) sendToStdOut(msg interface{}) error { return err } - r.burWriter.Write(b) - fmt.Println(r.burWriter, string(b)) - r.burWriter.Flush() + r.bufWriter.Write(b) + fmt.Fprintln(r.bufWriter, string(b)) + r.bufWriter.Flush() return nil } From 3efae98b6a6d7a037e8ecaa316de24c1333433ae Mon Sep 17 00:00:00 2001 From: Jason Yu Date: Fri, 27 Oct 2017 15:47:41 -0500 Subject: [PATCH 108/109] Fix a bug where kcl_reader was writing to KCL twice. --- kcl_reader.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kcl_reader.go b/kcl_reader.go index 3eb9d42..5a649e5 100644 --- a/kcl_reader.go +++ b/kcl_reader.go @@ -340,7 +340,7 @@ func (r *KclReader) sendToStdOut(msg interface{}) error { } r.bufWriter.Write(b) - fmt.Fprintln(r.bufWriter, string(b)) + r.bufWriter.WriteByte('\n') r.bufWriter.Flush() return nil From 46d513eec0d3775fc8ef5ffe155bfe7393a6b2dd Mon Sep 17 00:00:00 2001 From: Dan Travis Date: Thu, 1 Mar 2018 09:05:31 -0600 Subject: [PATCH 109/109] [LM-594] Fixed firehose writer timeout type --- firehose_writer.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/firehose_writer.go b/firehose_writer.go index ac2d604..95453b8 100644 --- a/firehose_writer.go +++ b/firehose_writer.go @@ -44,8 +44,8 @@ type FirehoseWriterOptionsFn func(*FirehoseWriter) error // FirehoseWriterResponseReadTimeout is a functional option method for configuring the // FirehoseWriter's response read timeout -func FirehoseWriterResponseReadTimeout(timeout time.Duration) KinesisWriterOptionsFn { - return func(o *KinesisWriter) error { +func FirehoseWriterResponseReadTimeout(timeout time.Duration) FirehoseWriterOptionsFn { + return func(o *FirehoseWriter) error { o.responseReadTimeout = timeout return nil }