-
Notifications
You must be signed in to change notification settings - Fork 26
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* feat: Kinesis Data Streams support Makes following changes to resolve #5 - Add `streams` package for the Kinesis Data Streams output plugin - Updated `plugins/firehose/main.go` to also include the brand-new Streams plugin - Rename `plugins/firehose` to `plugins/kinesis`, and `firehose.so` to `kinesis.so` according to the scope of the plugin - Update README accordingly * Try to fix the integration test on Travis
- Loading branch information
Showing
13 changed files
with
361 additions
and
16 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,8 +1,8 @@ | ||
package firehose | ||
|
||
import ( | ||
"time" | ||
"errors" | ||
"time" | ||
) | ||
|
||
type FirehoseConfig struct { | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,133 @@ | ||
package streams | ||
|
||
import ( | ||
"fmt" | ||
"github.com/aws/aws-sdk-go/aws" | ||
"github.com/aws/aws-sdk-go/aws/session" | ||
"github.com/aws/aws-sdk-go/service/kinesis" | ||
"github.com/elastic/beats/libbeat/beat" | ||
"github.com/elastic/beats/libbeat/logp" | ||
"github.com/elastic/beats/libbeat/outputs" | ||
"github.com/elastic/beats/libbeat/outputs/codec" | ||
"github.com/elastic/beats/libbeat/outputs/codec/json" | ||
"github.com/elastic/beats/libbeat/publisher" | ||
"time" | ||
) | ||
|
||
type client struct { | ||
streams *kinesis.Kinesis | ||
streamName string | ||
partitionKey string | ||
beatName string | ||
encoder codec.Codec | ||
timeout time.Duration | ||
observer outputs.Observer | ||
} | ||
|
||
func newClient(sess *session.Session, config *StreamsConfig, observer outputs.Observer, beat beat.Info) (*client, error) { | ||
client := &client{ | ||
streams: kinesis.New(sess), | ||
streamName: config.DeliveryStreamName, | ||
partitionKey: config.PartitionKey, | ||
beatName: beat.Beat, | ||
encoder: json.New(false, beat.Version), | ||
timeout: config.Timeout, | ||
observer: observer, | ||
} | ||
|
||
return client, nil | ||
} | ||
|
||
func (client *client) Close() error { | ||
return nil | ||
} | ||
|
||
func (client *client) Connect() error { | ||
return nil | ||
} | ||
|
||
func (client *client) Publish(batch publisher.Batch) error { | ||
events := batch.Events() | ||
observer := client.observer | ||
observer.NewBatch(len(events)) | ||
|
||
records, dropped := client.mapEvents(events) | ||
res, err := client.sendRecords(records) | ||
if err != nil { | ||
logp.Critical("Unable to send batch: %v", err) | ||
observer.Dropped(len(events)) | ||
return err | ||
} | ||
|
||
processFailedDeliveries(res, batch) | ||
batch.ACK() | ||
logp.Debug("kinesis", "Sent %d records", len(events)) | ||
observer.Dropped(dropped) | ||
observer.Acked(len(events) - dropped) | ||
return nil | ||
} | ||
|
||
func (client *client) mapEvents(events []publisher.Event) ([]*kinesis.PutRecordsRequestEntry, int) { | ||
dropped := 0 | ||
records := make([]*kinesis.PutRecordsRequestEntry, 0, len(events)) | ||
for _, event := range events { | ||
record, err := client.mapEvent(&event) | ||
if err != nil { | ||
dropped++ | ||
} else { | ||
records = append(records, record) | ||
} | ||
} | ||
|
||
return records, dropped | ||
} | ||
|
||
func (client *client) mapEvent(event *publisher.Event) (*kinesis.PutRecordsRequestEntry, error) { | ||
serializedEvent, err := client.encoder.Encode(client.beatName, &event.Content) | ||
if err != nil { | ||
if !event.Guaranteed() { | ||
return nil, err | ||
} | ||
|
||
logp.Critical("Unable to encode event: %v", err) | ||
return nil, err | ||
} | ||
|
||
rawPartitionKey, err := event.Content.GetValue(client.partitionKey) | ||
if err != nil { | ||
return nil, fmt.Errorf("failed to get parition key: %v", err) | ||
} | ||
|
||
partitionKey, ok := rawPartitionKey.(string) | ||
if !ok { | ||
return nil, fmt.Errorf("failed to get partition key: %s(=%v) is found, but not a string", client.partitionKey, rawPartitionKey) | ||
} | ||
|
||
return &kinesis.PutRecordsRequestEntry{Data: serializedEvent, PartitionKey: aws.String(partitionKey)}, nil | ||
} | ||
func (client *client) sendRecords(records []*kinesis.PutRecordsRequestEntry) (*kinesis.PutRecordsOutput, error) { | ||
request := kinesis.PutRecordsInput{ | ||
StreamName: &client.streamName, | ||
Records: records, | ||
} | ||
return client.streams.PutRecords(&request) | ||
} | ||
|
||
func processFailedDeliveries(res *kinesis.PutRecordsOutput, batch publisher.Batch) { | ||
if *res.FailedRecordCount > 0 { | ||
events := batch.Events() | ||
failedEvents := make([]publisher.Event, 0) | ||
records := res.Records | ||
for i, r := range records { | ||
if *r.ErrorCode != "" { | ||
failedEvents = append(failedEvents, events[i]) | ||
} | ||
} | ||
|
||
if len(failedEvents) > 0 { | ||
logp.Warn("Retrying %d events", len(failedEvents)) | ||
batch.RetryEvents(failedEvents) | ||
return | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,54 @@ | ||
package streams | ||
|
||
import ( | ||
"github.com/aws/aws-sdk-go/aws" | ||
"github.com/elastic/beats/libbeat/beat" | ||
"github.com/elastic/beats/libbeat/common" | ||
"github.com/elastic/beats/libbeat/publisher" | ||
"testing" | ||
) | ||
|
||
type MockCodec struct { | ||
} | ||
|
||
func (mock MockCodec) Encode(index string, event *beat.Event) ([]byte, error) { | ||
return []byte("boom"), nil | ||
} | ||
|
||
func TestMapEvent(t *testing.T) { | ||
fieldForPartitionKey := "mypartitionkey" | ||
expectedPartitionKey := "foobar" | ||
client := client{encoder: MockCodec{}, partitionKey: fieldForPartitionKey} | ||
event := &publisher.Event{Content: beat.Event{Fields: common.MapStr{fieldForPartitionKey: expectedPartitionKey}}} | ||
record, err := client.mapEvent(event) | ||
|
||
if err != nil { | ||
t.Fatalf("uenxpected error: %v", err) | ||
} | ||
|
||
if string(record.Data) != "boom" { | ||
t.Errorf("Unexpected data: %s", record.Data) | ||
} | ||
|
||
actualPartitionKey := aws.StringValue(record.PartitionKey) | ||
if actualPartitionKey != expectedPartitionKey { | ||
t.Errorf("unexpected partition key: %s", actualPartitionKey) | ||
} | ||
} | ||
|
||
func TestMapEvents(t *testing.T) { | ||
fieldForPartitionKey := "mypartitionkey" | ||
expectedPartitionKey := "foobar" | ||
client := client{encoder: MockCodec{}, partitionKey: fieldForPartitionKey} | ||
event := publisher.Event{Content: beat.Event{Fields: common.MapStr{fieldForPartitionKey: expectedPartitionKey}}} | ||
events := []publisher.Event{event} | ||
records, _ := client.mapEvents(events) | ||
|
||
if len(records) != 1 { | ||
t.Errorf("Expected 1 records, got %v", len(records)) | ||
} | ||
|
||
if string(records[0].Data) != "boom" { | ||
t.Errorf("Unexpected data %s", records[0].Data) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,54 @@ | ||
package streams | ||
|
||
import ( | ||
"errors" | ||
"time" | ||
) | ||
|
||
type StreamsConfig struct { | ||
Region string `config:"region"` | ||
DeliveryStreamName string `config:"stream_name"` | ||
PartitionKey string `config:"partition_key"` | ||
BatchSize int `config:"batch_size"` | ||
MaxRetries int `config:"max_retries"` | ||
Timeout time.Duration `config:"timeout"` | ||
Backoff backoff `config:"backoff"` | ||
} | ||
|
||
type backoff struct { | ||
Init time.Duration | ||
Max time.Duration | ||
} | ||
|
||
const ( | ||
defaultBatchSize = 50 | ||
// As per https://docs.aws.amazon.com/sdk-for-go/api/service/kinesis/#Kinesis.PutRecords | ||
maxBatchSize = 500 | ||
) | ||
|
||
var ( | ||
defaultConfig = StreamsConfig{ | ||
Timeout: 90 * time.Second, | ||
MaxRetries: 3, | ||
Backoff: backoff{ | ||
Init: 1 * time.Second, | ||
Max: 60 * time.Second, | ||
}, | ||
} | ||
) | ||
|
||
func (c *StreamsConfig) Validate() error { | ||
if c.Region == "" { | ||
return errors.New("region is not defined") | ||
} | ||
|
||
if c.DeliveryStreamName == "" { | ||
return errors.New("stream_name is not defined") | ||
} | ||
|
||
if c.BatchSize > maxBatchSize || c.BatchSize < 1 { | ||
return errors.New("invalid batch size") | ||
} | ||
|
||
return nil | ||
} |
Oops, something went wrong.