Release/1.7.0

config.go

@@ -6,6 +6,26 @@ import (
 	"time"
 )
 
+// MetricsConfig controls which metrics are sent to the StatReceiver
+type MetricsConfig struct {
+	EnableCheckpoint           bool
+	EnableEventToClient        bool
+	EnableEventsFromKinesis    bool
+	EnableRecordsInMemory      bool
+	EnableRecordsInMemoryBytes bool
+}
+
+// NewMetricsConfig returns a MetricsConfig with all metrics enabled (backward compatibility)
+func NewMetricsConfig() MetricsConfig {
+	return MetricsConfig{
+		EnableCheckpoint:           true,
+		EnableEventToClient:        true,
+		EnableEventsFromKinesis:    true,
+		EnableRecordsInMemory:      true,
+		EnableRecordsInMemoryBytes: true,
+	}
+}
+
 //TODO: Update documentation to include the defaults
 //TODO: Update the 'with' methods' comments to be less ridiculous
 
@@ -52,6 +72,20 @@ type Config struct {
 
 	// use ListShards to avoid LimitExceedException from DescribeStream
 	useListShardsForKinesisStreamReady bool
+
+	// Maximum number of records to fetch per GetRecords request
+	// AWS Kinesis allows up to 10,000 records per request (default)
+	// Reducing this value helps control memory usage at the cost of increased API calls
+	getRecordsLimit int
+	// ---------- [ Memory Management ] ----------
+	// Maximum number of shards that can fetch records from Kinesis concurrently
+	// 0 (default) means unlimited concurrent shards
+	// Setting this helps prevent memory exhaustion during scaling events when a single
+	// client temporarily handles many shards. With random data distribution across shards,
+	// this provides natural fairness over time while controlling peak memory usage.
+	// Example: Setting to 20 limits memory to ~20 * 10k records * avg_record_size
+	maxConcurrentShards int
+	metricsConfig       MetricsConfig
 }
 
 // NewConfig returns a default Config struct
@@ -67,6 +101,8 @@ func NewConfig() Config {
 		dynamoWriteCapacity:   10,
 		dynamoWaiterDelay:     3 * time.Second,
 		logger:                &DefaultLogger{},
+		getRecordsLimit:       10000,
+		metricsConfig:         NewMetricsConfig(),
 	}
 }
 
@@ -115,7 +151,7 @@ func (c Config) WithBufferSize(bufferSize int) Config {
 	return c
 }
 
-// WithStats returns a Config with a modified stats
+// WithStats returns a Config with a modified stats receiver
 func (c Config) WithStats(stats StatReceiver) Config {
 	c.stats = stats
 	return c
@@ -157,6 +193,30 @@ func (c Config) WithUseListShardsForKinesisStreamReady(shouldUse bool) Config {
 	return c
 }
 
+// WithGetRecordsLimit returns a Config with a modified maximum records per GetRecords request
+// This controls how many records to fetch per GetRecords API call.
+// AWS Kinesis allows up to 10,000 records per request. Reducing this helps control memory usage.
+func (c Config) WithGetRecordsLimit(getRecordsLimit int) Config {
+	c.getRecordsLimit = getRecordsLimit
+	return c
+}
+
+// WithMaxConcurrentShards returns a Config with a modified maximum concurrent shard limit
+// This controls how many shards can fetch records from Kinesis simultaneously.
+// Setting this helps prevent memory pressure during scaling events.
+// 0 (default) means unlimited concurrent shards.
+func (c Config) WithMaxConcurrentShards(maxConcurrentShards int) Config {
+	c.maxConcurrentShards = maxConcurrentShards
+	return c
+}
+
+// WithMetricsConfig returns a Config with a modified metrics configuration
+// This controls which metrics are sent to the StatReceiver
+func (c Config) WithMetricsConfig(metricsConfig MetricsConfig) Config {
+	c.metricsConfig = metricsConfig
+	return c
+}
+
 // Verify that a config struct has sane and valid values
 func validateConfig(c *Config) error {
 	if c.throttleDelay < 200*time.Millisecond {
@@ -199,5 +259,12 @@ func validateConfig(c *Config) error {
 		return ErrConfigInvalidLogger
 	}
 
+	if c.getRecordsLimit <= 0 || c.getRecordsLimit > 10000 {
+		return ErrConfigInvalidGetRecordsLimit
+	}
+	if c.maxConcurrentShards < 0 {
+		return ErrConfigInvalidMaxConcurrentShards
+	}
+
 	return nil
 }

config_test.go

@@ -52,6 +52,10 @@ func TestConfigErrors(t *testing.T) {
 	config = NewConfig().WithStats(nil)
 	err = validateConfig(&config)
 	require.EqualError(t, err, ErrConfigInvalidStats.Error())
+
+	config = NewConfig().WithMaxConcurrentShards(-1)
+	err = validateConfig(&config)
+	require.EqualError(t, err, ErrConfigInvalidMaxConcurrentShards.Error())
 }
 
 func TestConfigWithMethods(t *testing.T) {
@@ -65,7 +69,8 @@ func TestConfigWithMethods(t *testing.T) {
 		WithThrottleDelay(1 * time.Second).
 		WithStats(stats).
 		WithIteratorStartTimestamp(&tstamp).
-		WithUseListShardsForKinesisStreamReady(true)
+		WithUseListShardsForKinesisStreamReady(true).
+		WithMaxConcurrentShards(10)
 
 	err := validateConfig(&config)
 	require.NoError(t, err)
@@ -78,4 +83,5 @@ func TestConfigWithMethods(t *testing.T) {
 	require.Equal(t, stats, config.stats)
 	require.Equal(t, &tstamp, config.iteratorStartTimestamp)
 	require.Equal(t, true, config.useListShardsForKinesisStreamReady)
+	require.Equal(t, 10, config.maxConcurrentShards)
 }

errors.go

@@ -38,6 +38,10 @@ var (
 	ErrConfigInvalidDynamoCapacity = errors.New("dynamo read/write capacity cannot be 0")
 	// ErrConfigInvalidLogger - Logger cannot be nil
 	ErrConfigInvalidLogger = errors.New("logger cannot be nil")
+	// ErrConfigInvalidGetRecordsLimit - getRecordsLimit must be between 1 and 10000
+	ErrConfigInvalidGetRecordsLimit = errors.New("getRecordsLimit must be between 1 and 10000")
+	// ErrConfigInvalidMaxConcurrentShards - maxConcurrentShards must be >= 0
+	ErrConfigInvalidMaxConcurrentShards = errors.New("maxConcurrentShards must be >= 0")
 
 	// ErrStreamBusy - Stream is busy
 	ErrStreamBusy = errors.New("stream is busy")

kinsumer.go

@@ -5,11 +5,12 @@ package kinsumer
 import (
 	"context"
 	"fmt"
-	"github.com/twitchscience/kinsumer/kinsumeriface"
 	"sync"
 	"sync/atomic"
 	"time"
 
+	"github.com/twitchscience/kinsumer/kinsumeriface"
+
 	"github.com/aws/aws-sdk-go-v2/aws"
 	awsconfig "github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
@@ -31,6 +32,111 @@ type consumedRecord struct {
 	record       *ktypes.Record // Record retrieved from kinesis
 	checkpointer *checkpointer  // Object that will store the checkpoint back to the database
 	retrievedAt  time.Time      // Time the record was retrieved from Kinesis
+	payloadBytes int64          // Size of record.Data payload in bytes
+}
+
+// MetricType represents the type of metric update
+type MetricType int
+
+const (
+	RecordsIncrement MetricType = iota
+	RecordsDecrement
+	BytesIncrement
+	BytesDecrement
+	CombinedDecrement
+)
+
+// MetricUpdate represents a single metric update to be processed
+type MetricUpdate struct {
+	Type         MetricType
+	Value        int64
+	RecordsDelta int64 // For combined updates: change in records count
+	BytesDelta   int64 // For combined updates: change in bytes count
+}
+
+// MetricsManager handles metrics updates through channels to avoid atomic contention
+type MetricsManager struct {
+	updates      chan MetricUpdate
+	stop         chan struct{}
+	recordsCount int64  // Current records count - only written by metrics goroutine
+	bytesCount   int64  // Current bytes count - only written by metrics goroutine
+	logger       Logger // Logger for warnings and errors
+}
+
+// newMetricsManager creates a new MetricsManager with buffered channel
+func newMetricsManager(logger Logger) *MetricsManager {
+	return &MetricsManager{
+		updates: make(chan MetricUpdate, 5000), // Buffer sized for batched writes
+		stop:    make(chan struct{}),
+		logger:  logger,
+	}
+}
+
+// updateMetric sends a non-blocking metric update
+func (mm *MetricsManager) updateMetric(t MetricType, value int64) {
+	select {
+	case mm.updates <- MetricUpdate{Type: t, Value: value}:
+		// Successfully queued
+	default:
+		// Channel full - log immediately since drops should be very rare with batching
+		mm.logger.Log("Warning: metrics channel full, dropping single metric update (type: %d, value: %d)", t, value)
+	}
+}
+
+// decrementCombined sends a combined decrement update for both records and bytes
+func (mm *MetricsManager) decrementCombined(recordsDelta, bytesDelta int64) {
+	select {
+	case mm.updates <- MetricUpdate{
+		Type:         CombinedDecrement,
+		RecordsDelta: recordsDelta,
+		BytesDelta:   bytesDelta,
+	}:
+		// Successfully queued
+	default:
+		// Channel full - log immediately since drops should be very rare with batching
+		mm.logger.Log("Warning: metrics channel full, dropping update with %d records, %d bytes", recordsDelta, bytesDelta)
+	}
+}
+
+// getCurrentMetrics returns current metric values using atomic loads
+func (mm *MetricsManager) getCurrentMetrics() (int64, int64) {
+	return atomic.LoadInt64(&mm.recordsCount), atomic.LoadInt64(&mm.bytesCount)
+}
+
+// run processes metric updates in a separate goroutine
+func (mm *MetricsManager) run() {
+	var recordsCount, bytesCount int64
+
+	for {
+		select {
+		case update := <-mm.updates:
+			switch update.Type {
+			case RecordsIncrement:
+				recordsCount += update.Value
+			case RecordsDecrement:
+				recordsCount -= update.Value
+			case BytesIncrement:
+				bytesCount += update.Value
+			case BytesDecrement:
+				bytesCount -= update.Value
+			case CombinedDecrement:
+				recordsCount -= update.RecordsDelta
+				bytesCount -= update.BytesDelta
+			}
+
+			// Update shared state - single writer, no contention
+			atomic.StoreInt64(&mm.recordsCount, recordsCount)
+			atomic.StoreInt64(&mm.bytesCount, bytesCount)
+
+		case <-mm.stop:
+			return
+		}
+	}
+}
+
+// shutdown stops the metrics goroutine
+func (mm *MetricsManager) shutdown() {
+	close(mm.stop)
 }
 
 // Kinsumer is a Kinesis Consumer that tries to reduce duplicate reads while allowing for multiple
@@ -64,6 +170,10 @@ type Kinsumer struct {
 	leaderWG              sync.WaitGroup          // waitGroup for the leader loop
 	maxAgeForClientRecord time.Duration           // Cutoff for client/checkpoint records we read from dynamodb before we assume the record is stale
 	maxAgeForLeaderRecord time.Duration           // Cutoff for leader/shard cache records we read from dynamodb before we assume the record is stale
+	shardSemaphore        chan struct{}           // Semaphore to limit concurrent shard record fetching
+	metricsManager        *MetricsManager         // Channel-based metrics manager to avoid atomic contention
+	pendingRecords        int64                   // Accumulated record decrements for batching
+	pendingBytes          int64                   // Accumulated byte decrements for batching
 }
 
 // New returns a Kinsumer Interface with default kinesis and dynamodb instances, to be used in ec2 instances to get default auth and config
@@ -115,6 +225,12 @@ func NewWithInterfaces(
 		config.clientRecordMaxAge = &maxAge
 	}
 
+	// Wrap the stats receiver with filtering based on metrics configuration
+	// This happens after config validation to ensure we have the final configuration
+	if config.stats != nil {
+		config.stats = newFilteredStatReceiver(config.stats, config.metricsConfig)
+	}
+
 	consumer := &Kinsumer{
 		streamName:            streamName,
 		kinesis:               kinesis,
@@ -132,6 +248,12 @@ func NewWithInterfaces(
 		config:                config,
 		maxAgeForClientRecord: *config.clientRecordMaxAge,
 		maxAgeForLeaderRecord: config.leaderActionFrequency * 5,
+		metricsManager:        newMetricsManager(config.logger),
+	}
+
+	// Initialize semaphore for limiting concurrent shard record fetching
+	if config.maxConcurrentShards > 0 {
+		consumer.shardSemaphore = make(chan struct{}, config.maxConcurrentShards)
 	}
 	return consumer, nil
 }
@@ -392,6 +514,9 @@ func (k *Kinsumer) Run() error {
 		return fmt.Errorf("error in kinsumer Run initial refreshShards: %v", err)
 	}
 
+	// Start metrics goroutine
+	go k.metricsManager.run()
+
 	k.mainWG.Add(1)
 	go func() {
 		defer k.mainWG.Done()
@@ -407,6 +532,12 @@ func (k *Kinsumer) Run() error {
 			// Do this outside the k.isLeader check in case k.isLeader was false because
 			// we lost leadership but haven't had time to shutdown the goroutine yet.
 			k.leaderWG.Wait()
+			// Flush any remaining pending metrics before shutdown
+			if k.pendingRecords > 0 || k.pendingBytes > 0 {
+				k.metricsManager.decrementCombined(k.pendingRecords, k.pendingBytes)
+			}
+			// Shutdown metrics goroutine
+			k.metricsManager.shutdown()
 		}()
 
 		// We close k.output so that Next() stops, this is also the reason
@@ -418,6 +549,12 @@ func (k *Kinsumer) Run() error {
 			shardChangeTicker.Stop()
 		}()
 
+		// Report buffer size every 1 seconds
+		bufferReportTicker := time.NewTicker(1 * time.Second)
+		defer func() {
+			bufferReportTicker.Stop()
+		}()
+
 		var record *consumedRecord
 		if err := k.startConsumers(); err != nil {
 			k.errors <- fmt.Errorf("error starting consumers: %s", err)
@@ -448,6 +585,13 @@ func (k *Kinsumer) Run() error {
 				if !k.config.manualCheckpointing {
 					record.checkpointer.update(aws.ToString(record.record.SequenceNumber))
 				}
+				k.pendingRecords++
+				k.pendingBytes += record.payloadBytes
+				if k.pendingRecords >= 50 {
+					k.metricsManager.decrementCombined(k.pendingRecords, k.pendingBytes)
+					k.pendingRecords = 0
+					k.pendingBytes = 0
+				}
 				record = nil
 			case se := <-k.shardErrors:
 				k.errors <- fmt.Errorf("shard error (%s) in %s: %s", se.shardID, se.action, se.err)
@@ -474,6 +618,12 @@ func (k *Kinsumer) Run() error {
 
 					k.isRestartingConsumers = false
 				}
+			case <-bufferReportTicker.C:
+				// Report the current number of records pulled from Kinesis but not yet delivered to client
+				recordsCount, bytesCount := k.metricsManager.getCurrentMetrics()
+				k.config.stats.RecordsInMemory(recordsCount)
+				// Report the current total bytes of record payloads pulled from Kinesis but not yet delivered to client
+				k.config.stats.RecordsInMemoryBytes(bytesCount)
 			}
 		}
 	}()