producer

package module
v0.3.3 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Nov 14, 2025 License: MIT-0 Imports: 20 Imported by: 0

README

Amazon kinesis producer Build status License GoDoc

A KPL-like batch producer for Amazon Kinesis built on top of the official Go AWS SDK
and using the same aggregation format that KPL use.

Example
package main

import (
	"context"
	"log"
	"time"

	"github.com/a8m/kinesis-producer"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/kinesis"
)

func main() {
	// Load AWS SDK v2 configuration
	cfg, err := config.LoadDefaultConfig(context.TODO())
	if err != nil {
		log.Fatalf("failed to load config: %v", err)
	}
	client := kinesis.NewFromConfig(cfg)

	pr, err := producer.New(&producer.Config{
		StreamName:   "test",
		BacklogCount: 2000,
		Client:       client,
		// RateLimit:  150,  // Optional: default is 150 (150% of shard limits)
		// ShardMapRefreshTimeout: 30 * time.Second,  // Optional: default is 30s
	})
	if err != nil {
		log.Fatalf("failed to create producer: %v", err)
	}

	if err := pr.Start(); err != nil {
		log.Fatalf("failed to start producer: %v", err)
	}

	// Handle failures
	go func() {
		for r := range pr.NotifyFailures() {
			// r contains `Data`, `PartitionKey` and `Error()`
			log.Printf("failed to put record: %v", r)
		}
	}()

	go func() {
		for i := 0; i < 5000; i++ {
			err := pr.Put([]byte("foo"), "bar")
			if err != nil {
				log.Fatalf("error producing: %v", err)
			}
		}
	}()

	time.Sleep(3 * time.Second)
	pr.Stop()
}
Specifying logger implementation

producer.Config takes an optional logging.Logger implementation.

Using a custom logger
customLogger := &CustomLogger{}

&producer.Config{
  StreamName:   "test",
  BacklogCount: 2000,
  Client:       client,
  Logger:       customLogger,
}
Using logrus
import (
	"github.com/sirupsen/logrus"
	producer "github.com/a8m/kinesis-producer"
	"github.com/a8m/kinesis-producer/loggers"
)

log := logrus.New()

&producer.Config{
  StreamName:   "test",
  BacklogCount: 2000,
  Client:       client,
  Logger:       loggers.Logrus(log),
}

kinesis-producer ships with three logger implementations.

  • producer.Standard used the standard library logger
  • loggers.Logrus uses logrus logger
  • loggers.Zap uses zap logger
License

MIT

Documentation

Overview

Amazon kinesis producer A KPL-like batch producer for Amazon Kinesis built on top of the official Go AWS SDK and using the same aggregation format that KPL use.

Note: this project start as a fork of `tj/go-kinesis`. if you are not intersting in the KPL aggregation logic, you probably want to check it out.

Example
logger := &StdLogger{log.New(os.Stdout, "", log.LstdFlags)}
cfg, _ := config.LoadDefaultConfig(context.TODO())
client := kinesis.NewFromConfig(cfg)
pr, err := New(&Config{
	StreamName:   "test",
	BacklogCount: 2000,
	Client:       client,
	Logger:       logger,
})
if err != nil {
	logger.Error("failed to create producer", err)
	return
}

pr.Start()

// Handle failures
go func() {
	for r := range pr.NotifyFailures() {
		// r contains `Data`, `PartitionKey` and `Error()`
		logger.Error("detected put failure", r.error)
	}
}()

go func() {
	for i := 0; i < 5000; i++ {
		err := pr.Put([]byte("foo"), "bar")
		if err != nil {
			logger.Error("error producing", err)
		}
	}
}()

time.Sleep(3 * time.Second)
pr.Stop()

Index

Examples

Constants

This section is empty.

Variables

View Source
var (
	ErrStoppedProducer     = errors.New("Unable to Put record. Producer is already stopped")
	ErrIllegalPartitionKey = errors.New("Invalid partition key. Length must be at least 1 and at most 256")
	ErrRecordSizeExceeded  = errors.New("Data must be less than or equal to 1MB in size")
)

Errors

View Source
var File_messages_proto protoreflect.FileDescriptor

Functions

This section is empty.

Types

type AggregatedRecord

type AggregatedRecord struct {
	PartitionKeyTable    []string  `protobuf:"bytes,1,rep,name=partition_key_table,json=partitionKeyTable" json:"partition_key_table,omitempty"`
	ExplicitHashKeyTable []string  `protobuf:"bytes,2,rep,name=explicit_hash_key_table,json=explicitHashKeyTable" json:"explicit_hash_key_table,omitempty"`
	Records              []*Record `protobuf:"bytes,3,rep,name=records" json:"records,omitempty"`
	// contains filtered or unexported fields
}

func (*AggregatedRecord) Descriptor deprecated

func (*AggregatedRecord) Descriptor() ([]byte, []int)

Deprecated: Use AggregatedRecord.ProtoReflect.Descriptor instead.

func (*AggregatedRecord) GetExplicitHashKeyTable

func (x *AggregatedRecord) GetExplicitHashKeyTable() []string

func (*AggregatedRecord) GetPartitionKeyTable

func (x *AggregatedRecord) GetPartitionKeyTable() []string

func (*AggregatedRecord) GetRecords

func (x *AggregatedRecord) GetRecords() []*Record

func (*AggregatedRecord) ProtoMessage

func (*AggregatedRecord) ProtoMessage()

func (*AggregatedRecord) ProtoReflect

func (x *AggregatedRecord) ProtoReflect() protoreflect.Message

func (*AggregatedRecord) Reset

func (x *AggregatedRecord) Reset()

func (*AggregatedRecord) String

func (x *AggregatedRecord) String() string

type Aggregator

type Aggregator struct {
	// contains filtered or unexported fields
}

Aggregator coordinates multiple Reducers for per-shard aggregation.

Routing Strategy:

Routes records by target shard ID (computed from partition key MD5 hash).
ShardMap is required to determine the target shard for each partition key.

This matches Java KPL behavior where per-shard aggregation is always enabled and routes records based on shard boundaries.

Example:

shardMap, _ := NewShardMap(streamName, client, logger, false, 30*time.Second)
aggregator := NewAggregator(shardMap, config)
aggregator.Put(data, "random-uuid-123")  // Routes to shard based on hash

func NewAggregator

func NewAggregator(shardMap *ShardMap, config *ReducerConfig) (*Aggregator, error)

NewAggregator creates a new Aggregator with a required ShardMap. ShardMap must not be nil as it is required for per-shard aggregation (matching Java KPL behavior). Returns an error if ShardMap is nil.

func (*Aggregator) Count

func (a *Aggregator) Count() int

Count returns the total number of records across all reducers.

func (*Aggregator) Drain

func (a *Aggregator) Drain() ([]*ktypes.PutRecordsRequestEntry, error)

Drain returns all records from all reducers, regardless of flush conditions. This is useful for flushing on shutdown or periodic flushes. After draining, empty reducers are removed.

func (*Aggregator) DrainShard added in v0.3.3

func (a *Aggregator) DrainShard(shardID string) (*ktypes.PutRecordsRequestEntry, error)

DrainShard drains a specific shard's reducer by shardID. Returns the aggregated record if the shard exists and has records, nil otherwise. This is used for event-driven flushing when FlushInterval expires.

func (*Aggregator) GetShardCount

func (a *Aggregator) GetShardCount() int

GetShardCount returns the number of distinct shards/keys being aggregated.

func (*Aggregator) Put

func (a *Aggregator) Put(data []byte, partitionKey string) (*ktypes.PutRecordsRequestEntry, error)

Put adds a record to the appropriate reducer. Routes by shard ID (computed from partition key hash via ShardMap). Returns flushed aggregated record if flush conditions are met, nil otherwise. This matches Java KPL behavior where records are immediately returned when ready.

func (*Aggregator) Size

func (a *Aggregator) Size() int

Size returns the total bytes across all reducers.

type Config

type Config struct {
	// StreamName is the Kinesis stream.
	StreamName string

	// FlushInterval is the maximum amount of time a record can stay in the aggregation buffer.
	// After this time, records will be flushed regardless of size or count limits.
	// This matches Java KPL's RecordMaxBufferedTime behavior.
	// Default: 100ms (matching Java KPL)
	FlushInterval time.Duration

	// BatchCount determine the maximum number of items to pack in batch.
	// Must not exceed length. Defaults to 500.
	BatchCount int

	// BatchSize determine the maximum number of bytes to send with a PutRecords request.
	// Must not exceed 5MiB; Default to 5MiB.
	BatchSize int

	// AggregateBatchCount determine the maximum number of items to pack into an aggregated record.
	AggregateBatchCount int

	// AggregationBatchSize determine the maximum number of bytes to pack into an aggregated record. User records larger
	// than this will bypass aggregation.
	AggregateBatchSize int

	// BacklogCount determines the channel capacity before Put() will begin blocking. Default to `BatchCount`.
	BacklogCount int

	// Number of requests to sent concurrently. Default to 24.
	MaxConnections int

	// Logger is the logger used. Default to producer.Logger.
	Logger Logger

	// Enabling verbose logging. Default to false.
	Verbose bool

	// Client is the Putter interface implementation.
	Client Putter

	// RateLimit sets the maximum allowed put rate for a shard, as a percentage of backend limits.
	// Like Java KPL, rate limiting is always enabled and cannot be disabled.
	// The rate limit is implemented using a token bucket algorithm per shard.
	//
	// Backend limits per shard:
	//   - 1000 records per second
	//   - 1 MB per second
	//
	// With RateLimit=150 (default):
	//   - 1500 records per second per shard (1000 × 1.5)
	//   - 1.5 MB per second per shard (1 MB × 1.5)
	//
	// The default of 150% allows a single producer to completely saturate a shard.
	// Set to 100 to match exact shard limits, or lower to reduce throttling risk.
	// Minimum value: 1, Maximum value: 9223372036854775807
	//
	// Default: 150 (matches Java KPL)
	// Note: ShardMap and RateLimit are always enabled (matching Java KPL behavior)
	RateLimit int

	// ShardMap is an optional pre-configured ShardMap for testing purposes.
	// If provided, the producer will use this instead of creating a new one.
	// Leave nil for normal operation where ShardMap will be automatically initialized.
	ShardMap *ShardMap

	// ShardMapRefreshTimeout is the maximum time to wait for a ShardMap refresh operation.
	// This timeout is used when calling DescribeStream API.
	// Default: 30 seconds
	ShardMapRefreshTimeout time.Duration

	// BufferFlushTimeout is the maximum time Producer.buf records wait before flush.
	// This matches Java KPL's putrecords_buffer_duration (min(50ms, RecordMaxBufferedTime * 0.2)).
	// Default: 20ms (matching Java KPL's typical value: 100ms * 0.2 = 20ms)
	BufferFlushTimeout time.Duration
}

Config is the Producer configuration.

type FailureRecord

type FailureRecord struct {
	Data         []byte
	PartitionKey string
	// contains filtered or unexported fields
}

Failure record type

type Limiter

type Limiter struct {
	// contains filtered or unexported fields
}

Limiter implements per-shard rate limiting using token bucket algorithm. Each shard has independent rate limits for records per second and bytes per second.

func NewLimiter

func NewLimiter(config *LimiterConfig) *Limiter

NewLimiter creates a new Limiter with the given configuration.

func (*Limiter) Allow

func (l *Limiter) Allow(shardID string, records int, bytes int) bool

Allow checks if the operation is allowed under rate limits without blocking. Returns true if allowed, false if rate limited.

func (*Limiter) GetShardCount

func (l *Limiter) GetShardCount() int

GetShardCount returns the number of shards being tracked.

func (*Limiter) Reset

func (l *Limiter) Reset()

Reset clears all shard limiters (useful for testing).

func (*Limiter) Wait

func (l *Limiter) Wait(shardID string, records int, bytes int)

Wait blocks until the operation is allowed under rate limits. This ensures the operation will proceed without exceeding rate limits.

type LimiterConfig

type LimiterConfig struct {
	// RateLimit is the maximum allowed put rate as a percentage of backend limits.
	// Backend limits per shard: 1000 records/sec, 1 MB/sec
	// Default: 150 (150% = 1500 records/sec, 1.5 MB/sec)
	RateLimit int
}

LimiterConfig holds rate limiting configuration.

type LogValue

type LogValue struct {
	Name  string
	Value interface{}
}

LogValue represents a key:value pair used by the Logger interface

func (LogValue) String

func (v LogValue) String() string

type Logger

type Logger interface {
	Info(msg string, values ...LogValue)
	Error(msg string, err error, values ...LogValue)
}

Logger represents a simple interface used by kinesis-producer to handle logging

type Producer

type Producer struct {
	sync.RWMutex
	*Config
	// contains filtered or unexported fields
}

Producer batches records.

func New

func New(config *Config) (*Producer, error)

New creates new producer with the given config. Returns an error if the configuration is invalid.

func (*Producer) NotifyFailures

func (p *Producer) NotifyFailures() <-chan *FailureRecord

NotifyFailures registers and return listener to handle undeliverable messages. The incoming struct has a copy of the Data and the PartitionKey along with some error information about why the publishing failed.

func (*Producer) Put

func (p *Producer) Put(data []byte, partitionKey string) error

Put `data` using `partitionKey` asynchronously. This method is thread-safe.

Under the covers, the Producer will automatically re-attempt puts in case of transient errors. When unrecoverable error has detected(e.g: trying to put to in a stream that doesn't exist), the message will returned by the Producer. Add a listener with `Producer.NotifyFailures` to handle undeliverable messages.

func (*Producer) Start

func (p *Producer) Start() error

Start the producer

func (*Producer) Stop

func (p *Producer) Stop()

Stop the producer gracefully. Flushes any in-flight data.

type Putter

type Putter interface {
	PutRecords(context.Context, *k.PutRecordsInput, ...func(*k.Options)) (*k.PutRecordsOutput, error)
}

Putter is the interface that wraps the KinesisAPI.PutRecords method.

type Record

type Record struct {
	PartitionKeyIndex    *uint64 `protobuf:"varint,1,req,name=partition_key_index,json=partitionKeyIndex" json:"partition_key_index,omitempty"`
	ExplicitHashKeyIndex *uint64 `protobuf:"varint,2,opt,name=explicit_hash_key_index,json=explicitHashKeyIndex" json:"explicit_hash_key_index,omitempty"`
	Data                 []byte  `protobuf:"bytes,3,req,name=data" json:"data,omitempty"`
	Tags                 []*Tag  `protobuf:"bytes,4,rep,name=tags" json:"tags,omitempty"`
	// contains filtered or unexported fields
}

func (*Record) Descriptor deprecated

func (*Record) Descriptor() ([]byte, []int)

Deprecated: Use Record.ProtoReflect.Descriptor instead.

func (*Record) GetData

func (x *Record) GetData() []byte

func (*Record) GetExplicitHashKeyIndex

func (x *Record) GetExplicitHashKeyIndex() uint64

func (*Record) GetPartitionKeyIndex

func (x *Record) GetPartitionKeyIndex() uint64

func (*Record) GetTags

func (x *Record) GetTags() []*Tag

func (*Record) ProtoMessage

func (*Record) ProtoMessage()

func (*Record) ProtoReflect

func (x *Record) ProtoReflect() protoreflect.Message

func (*Record) Reset

func (x *Record) Reset()

func (*Record) String

func (x *Record) String() string

type Reducer added in v0.3.3

type Reducer struct {
	// contains filtered or unexported fields
}

Reducer aggregates records for a single shard. This matches the KPL Reducer<UserRecord, KinesisRecord> behavior. All records in this reducer will be sent to the same shard.

func NewReducer added in v0.3.3

func NewReducer(shardID string, config *ReducerConfig) *Reducer

NewReducer creates a new shard-specific reducer (for testing without notification).

func (*Reducer) Count added in v0.3.3

func (r *Reducer) Count() int

Count returns the number of records.

func (*Reducer) Drain added in v0.3.3

func (r *Reducer) Drain() (*ktypes.PutRecordsRequestEntry, error)

Drain creates an aggregated PutRecordsRequestEntry and resets the reducer. Returns nil if there are no records to drain.

func (*Reducer) IsEmpty added in v0.3.3

func (r *Reducer) IsEmpty() bool

IsEmpty returns true if the reducer has no records.

func (*Reducer) Put added in v0.3.3

func (r *Reducer) Put(data []byte, partitionKey string) (*ktypes.PutRecordsRequestEntry, error)

Put adds a record to the reducer with partition key deduplication. Returns flushed aggregated record if flush conditions are met, nil otherwise. This matches Java KPL behavior where Reducer.add() returns flushed container.

func (*Reducer) ShardID added in v0.3.3

func (r *Reducer) ShardID() string

ShardID returns the shard ID this reducer is for.

func (*Reducer) Size added in v0.3.3

func (r *Reducer) Size() int

Size returns the current size in bytes.

type ReducerConfig added in v0.3.3

type ReducerConfig struct {
	MaxSize       int           // max bytes per aggregated record (default 50KB)
	MaxCount      int           // max records per aggregated record (default 1000)
	FlushInterval time.Duration // max time records can stay buffered (default 100ms, matching Java KPL)
}

ReducerConfig holds configuration for reducers (per-shard aggregation).

type ShardLimiter

type ShardLimiter struct {
	// contains filtered or unexported fields
}

ShardLimiter implements token bucket rate limiting for a single shard.

type ShardMap

type ShardMap struct {
	// contains filtered or unexported fields
}

ShardMap maintains a mapping from partition key hashes to shard IDs. It uses an event-driven approach to handle shard splits and merges.

Event-Driven Updates (matching Java KPL behavior): The shard map is invalidated when a mismatch is detected between the predicted shard and the actual shard where a record was sent. This triggers an immediate refresh with exponential backoff retry logic (1s -> 30s).

Unlike the periodic refresh approach, this matches Java KPL's reactive strategy where updates only occur when needed.

func NewShardMap

func NewShardMap(streamName string, client *k.Client, logger Logger, verbose bool, refreshTimeout time.Duration) (*ShardMap, error)

NewShardMap creates a new ShardMap with event-driven updates. The shard map uses an event-driven approach (matching Java KPL) where updates only occur when a mismatch is detected. refreshTimeout specifies the maximum time to wait for DescribeStream API calls.

func (*ShardMap) GetShardCount

func (sm *ShardMap) GetShardCount() int

GetShardCount returns the current number of open shards.

func (*ShardMap) GetShardID

func (sm *ShardMap) GetShardID(partitionKey string) (string, error)

GetShardID returns the shard ID for a given partition key. It computes the MD5 hash of the partition key and performs a binary search through the sorted shard ranges.

Thread-safety: Uses RLock to allow concurrent reads. The Copy-on-Write pattern in refresh() ensures readers never see partial updates and are only briefly blocked during the atomic pointer swap.

func (*ShardMap) GetShardIDByHash

func (sm *ShardMap) GetShardIDByHash(hashKeyStr string) (string, error)

GetShardIDByHash returns the shard ID for a given hash key (decimal string). This is used to compute the predicted shard ID from ExplicitHashKey.

func (*ShardMap) Invalidate

func (sm *ShardMap) Invalidate(seenAt time.Time, predictedShardID string)

Invalidate marks the shard map as invalid and triggers an immediate refresh. This is called when a mismatch is detected between the predicted shard and the actual shard where a record was sent.

seenAt: The time when the mismatch was detected predictedShardID: The shard ID that was predicted (optional)

func (*ShardMap) Stop

func (sm *ShardMap) Stop()

Stop cancels any scheduled updates.

type ShardMapState

type ShardMapState int

ShardMapState represents the current state of the shard map.

const (
	ShardMapStateReady    ShardMapState = iota // Shard map is ready for use
	ShardMapStateInvalid                       // Shard map is invalid and needs refresh
	ShardMapStateUpdating                      // Shard map is currently being updated
)

type ShardRange

type ShardRange struct {
	ShardID         string
	StartingHashKey *big.Int // 128-bit hash key
	EndingHashKey   *big.Int // 128-bit hash key
}

ShardRange represents a shard's hash key range.

type StdLogger

type StdLogger struct {
	Logger *log.Logger
}

StdLogger implements the Logger interface using standard library loggers

func (*StdLogger) Error

func (l *StdLogger) Error(msg string, err error, values ...LogValue)

Error prints log message

func (*StdLogger) Info

func (l *StdLogger) Info(msg string, values ...LogValue)

Info prints log message

type Tag

type Tag struct {
	Key   *string `protobuf:"bytes,1,req,name=key" json:"key,omitempty"`
	Value *string `protobuf:"bytes,2,opt,name=value" json:"value,omitempty"`
	// contains filtered or unexported fields
}

func (*Tag) Descriptor deprecated

func (*Tag) Descriptor() ([]byte, []int)

Deprecated: Use Tag.ProtoReflect.Descriptor instead.

func (*Tag) GetKey

func (x *Tag) GetKey() string

func (*Tag) GetValue

func (x *Tag) GetValue() string

func (*Tag) ProtoMessage

func (*Tag) ProtoMessage()

func (*Tag) ProtoReflect

func (x *Tag) ProtoReflect() protoreflect.Message

func (*Tag) Reset

func (x *Tag) Reset()

func (*Tag) String

func (x *Tag) String() string

Directories

Path Synopsis
loggers

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL