Feat/buffered data metric

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
 
@@ -65,6 +85,7 @@ type Config struct {
 	// 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
@@ -81,6 +102,7 @@ func NewConfig() Config {
 		dynamoWaiterDelay:     3 * time.Second,
 		logger:                &DefaultLogger{},
 		getRecordsLimit:       10000,
+		metricsConfig:         NewMetricsConfig(),
 	}
 }
 
@@ -129,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
@@ -188,6 +210,13 @@ func (c Config) WithMaxConcurrentShards(maxConcurrentShards int) Config {
 	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 {

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
@@ -65,6 +171,9 @@ type Kinsumer struct {
 	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
@@ -116,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,
@@ -133,6 +248,7 @@ func NewWithInterfaces(
 		config:                config,
 		maxAgeForClientRecord: *config.clientRecordMaxAge,
 		maxAgeForLeaderRecord: config.leaderActionFrequency * 5,
+		metricsManager:        newMetricsManager(config.logger),
 	}
 
 	// Initialize semaphore for limiting concurrent shard record fetching
@@ -398,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()
@@ -413,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
@@ -424,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)
@@ -454,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)
@@ -480,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)
 			}
 		}
 	}()

noopstatreceiver.go

@@ -17,3 +17,9 @@ func (*NoopStatReceiver) EventToClient(inserted, retrieved time.Time) {}
 
 // EventsFromKinesis implementation that doesn't do anything
 func (*NoopStatReceiver) EventsFromKinesis(num int, shardID string, lag time.Duration) {}
+
+// RecordsInMemory implementation that doesn't do anything
+func (*NoopStatReceiver) RecordsInMemory(count int64) {}
+
+// RecordsInMemoryBytes implementation that doesn't do anything
+func (*NoopStatReceiver) RecordsInMemoryBytes(bytes int64) {}

shard_consumer.go

@@ -6,16 +6,16 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/twitchscience/kinsumer/kinsumeriface"
 	"time"
 
+	"github.com/twitchscience/kinsumer/kinsumeriface"
+
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/kinesis"
 	ktypes "github.com/aws/aws-sdk-go-v2/service/kinesis/types"
 	smithy "github.com/aws/smithy-go"
 )
 
-
 // getShardIterator gets a shard iterator after the last sequence number we read or at the start of the stream
 func getShardIterator(k kinsumeriface.KinesisAPI, streamName string, shardID string, sequenceNumber string, iteratorStartTimestamp *time.Time) (string, error) {
 	shardIteratorType := ktypes.ShardIteratorTypeAfterSequenceNumber
@@ -227,11 +227,34 @@ mainloop:
 
 		// Put all the records we got onto the channel
 		k.config.stats.EventsFromKinesis(len(records), shardID, lag)
+		k.metricsManager.updateMetric(RecordsIncrement, int64(len(records)))
+
+		// Calculate total payload bytes in the batch
+		totalBytes := int64(0)
+		for _, record := range records {
+			totalBytes += int64(len(record.Data))
+		}
+		k.metricsManager.updateMetric(BytesIncrement, totalBytes)
+
+		// Track records for cleanup in case of early return
+		recordsToCleanup := int64(len(records))
+		bytesToCleanup := totalBytes
+		defer func() {
+			// Decrement any records that weren't successfully processed
+			if recordsToCleanup > 0 {
+				k.metricsManager.updateMetric(RecordsDecrement, recordsToCleanup)
+			}
+			if bytesToCleanup > 0 {
+				k.metricsManager.updateMetric(BytesDecrement, bytesToCleanup)
+			}
+		}()
+
 		if len(records) > 0 {
 			retrievedAt := time.Now()
 			for _, record := range records {
-			RecordLoop:
 				// Loop until we stop or the record is consumed, checkpointing if necessary.
+				recordPayloadBytes := int64(len(record.Data))
+			RecordLoop:
 				for {
 					select {
 					case <-commitTicker.C:
@@ -249,7 +272,10 @@ mainloop:
 						record:       &record,
 						checkpointer: checkpointer,
 						retrievedAt:  retrievedAt,
+						payloadBytes: recordPayloadBytes,
 					}:
+						recordsToCleanup--                         // Record successfully sent to channel
+						bytesToCleanup -= recordPayloadBytes       // Decrement bytes for successfully sent record
 						checkpointer.lastRecordPassed = time.Now() // Mark the time so we don't retain shards when we're too slow to do so
 						lastSeqToCheckp = aws.ToString(record.SequenceNumber)
 						break RecordLoop
@@ -260,7 +286,7 @@ mainloop:
 			// Update the last sequence number we saw, in case we reached the end of the stream.
 			lastSeqNum = aws.ToString(records[len(records)-1].SequenceNumber)
 		}
-		
+
 		// Release semaphore after successfully processing all records from this batch
 		releaseSemaphore()
 		iterator = next