diff --git a/config.go b/config.go index 697f362..ec4af17 100644 --- a/config.go +++ b/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 { diff --git a/kinsumer.go b/kinsumer.go index 8e4cb9c..6564d7a 100644 --- a/kinsumer.go +++ b/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) } } }() diff --git a/noopstatreceiver.go b/noopstatreceiver.go index 1f972ad..000c8d4 100644 --- a/noopstatreceiver.go +++ b/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) {} diff --git a/shard_consumer.go b/shard_consumer.go index 5a87b5b..527e59d 100644 --- a/shard_consumer.go +++ b/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 diff --git a/statreceiver.go b/statreceiver.go index 9eb954b..99fa8b3 100644 --- a/statreceiver.go +++ b/statreceiver.go @@ -27,4 +27,61 @@ type StatReceiver interface { // `shardID` ID of the shard that the records were retrieved from // `lag` How far the records are from the tip of the stream. EventsFromKinesis(num int, shardID string, lag time.Duration) + + // RecordsInMemory is called periodically to report the current number of records + // that have been pulled from Kinesis and are buffered in memory, waiting to be + // delivered to the client. + // `count` Current number of records in the internal buffer + RecordsInMemory(count int64) + + // RecordsInMemoryBytes is called periodically to report the current total bytes + // of record payloads that have been pulled from Kinesis and are buffered in memory, + // waiting to be delivered to the client. + // `bytes` Current total payload bytes in the internal buffer + RecordsInMemoryBytes(bytes int64) +} + +// filteredStatReceiver wraps a StatReceiver and filters method calls based on MetricsConfig +// This allows selective enabling/disabling of metrics at the configuration level +type filteredStatReceiver struct { + underlying StatReceiver + config MetricsConfig +} + +// newFilteredStatReceiver creates a new filteredStatReceiver that wraps the underlying StatReceiver +func newFilteredStatReceiver(underlying StatReceiver, config MetricsConfig) StatReceiver { + return &filteredStatReceiver{ + underlying: underlying, + config: config, + } +} + +func (f *filteredStatReceiver) Checkpoint() { + if f.config.EnableCheckpoint { + f.underlying.Checkpoint() + } +} + +func (f *filteredStatReceiver) EventToClient(inserted, retrieved time.Time) { + if f.config.EnableEventToClient { + f.underlying.EventToClient(inserted, retrieved) + } +} + +func (f *filteredStatReceiver) EventsFromKinesis(num int, shardID string, lag time.Duration) { + if f.config.EnableEventsFromKinesis { + f.underlying.EventsFromKinesis(num, shardID, lag) + } +} + +func (f *filteredStatReceiver) RecordsInMemory(count int64) { + if f.config.EnableRecordsInMemory { + f.underlying.RecordsInMemory(count) + } +} + +func (f *filteredStatReceiver) RecordsInMemoryBytes(bytes int64) { + if f.config.EnableRecordsInMemoryBytes { + f.underlying.RecordsInMemoryBytes(bytes) + } } diff --git a/statreceiver_test.go b/statreceiver_test.go new file mode 100644 index 0000000..3109c61 --- /dev/null +++ b/statreceiver_test.go @@ -0,0 +1,338 @@ +// Copyright (c) 2016 Twitch Interactive + +package kinsumer + +import ( + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestStatReceiver is a StatReceiver implementation that captures all method calls +// for testing purposes. It is thread-safe since StatReceiver methods can be called +// from multiple goroutines. +type TestStatReceiver struct { + mu sync.Mutex + recordsInMemoryCalls []int64 + recordsInMemoryBytesCalls []int64 + checkpointCalls int + eventToClientCalls int + eventsFromKinesisCalls []EventsFromKinesisCall +} + +// EventsFromKinesisCall captures the parameters of EventsFromKinesis calls +type EventsFromKinesisCall struct { + Num int + ShardID string + Lag time.Duration +} + +// Checkpoint captures checkpoint calls +func (t *TestStatReceiver) Checkpoint() { + t.mu.Lock() + defer t.mu.Unlock() + t.checkpointCalls++ +} + +// EventToClient captures event to client calls +func (t *TestStatReceiver) EventToClient(inserted, retrieved time.Time) { + t.mu.Lock() + defer t.mu.Unlock() + t.eventToClientCalls++ +} + +// EventsFromKinesis captures events from kinesis calls +func (t *TestStatReceiver) EventsFromKinesis(num int, shardID string, lag time.Duration) { + t.mu.Lock() + defer t.mu.Unlock() + t.eventsFromKinesisCalls = append(t.eventsFromKinesisCalls, EventsFromKinesisCall{ + Num: num, + ShardID: shardID, + Lag: lag, + }) +} + +// RecordsInMemory captures records in memory calls +func (t *TestStatReceiver) RecordsInMemory(count int64) { + t.mu.Lock() + defer t.mu.Unlock() + t.recordsInMemoryCalls = append(t.recordsInMemoryCalls, count) +} + +// RecordsInMemoryBytes captures records in memory bytes calls +func (t *TestStatReceiver) RecordsInMemoryBytes(bytes int64) { + t.mu.Lock() + defer t.mu.Unlock() + t.recordsInMemoryBytesCalls = append(t.recordsInMemoryBytesCalls, bytes) +} + +// Helper methods for testing + +// GetRecordsInMemoryCalls returns a copy of all RecordsInMemory calls +func (t *TestStatReceiver) GetRecordsInMemoryCalls() []int64 { + t.mu.Lock() + defer t.mu.Unlock() + calls := make([]int64, len(t.recordsInMemoryCalls)) + copy(calls, t.recordsInMemoryCalls) + return calls +} + +// GetRecordsInMemoryBytesCalls returns a copy of all RecordsInMemoryBytes calls +func (t *TestStatReceiver) GetRecordsInMemoryBytesCalls() []int64 { + t.mu.Lock() + defer t.mu.Unlock() + calls := make([]int64, len(t.recordsInMemoryBytesCalls)) + copy(calls, t.recordsInMemoryBytesCalls) + return calls +} + +// GetLastRecordsInMemory returns the last RecordsInMemory value, or -1 if no calls +func (t *TestStatReceiver) GetLastRecordsInMemory() int64 { + t.mu.Lock() + defer t.mu.Unlock() + if len(t.recordsInMemoryCalls) == 0 { + return -1 + } + return t.recordsInMemoryCalls[len(t.recordsInMemoryCalls)-1] +} + +// GetLastRecordsInMemoryBytes returns the last RecordsInMemoryBytes value, or -1 if no calls +func (t *TestStatReceiver) GetLastRecordsInMemoryBytes() int64 { + t.mu.Lock() + defer t.mu.Unlock() + if len(t.recordsInMemoryBytesCalls) == 0 { + return -1 + } + return t.recordsInMemoryBytesCalls[len(t.recordsInMemoryBytesCalls)-1] +} + +// GetCallCounts returns counts of various method calls +func (t *TestStatReceiver) GetCallCounts() (recordsInMemory, recordsInMemoryBytes, checkpoint, eventToClient, eventsFromKinesis int) { + t.mu.Lock() + defer t.mu.Unlock() + return len(t.recordsInMemoryCalls), len(t.recordsInMemoryBytesCalls), t.checkpointCalls, t.eventToClientCalls, len(t.eventsFromKinesisCalls) +} + +// Reset clears all captured calls +func (t *TestStatReceiver) Reset() { + t.mu.Lock() + defer t.mu.Unlock() + t.recordsInMemoryCalls = nil + t.recordsInMemoryBytesCalls = nil + t.checkpointCalls = 0 + t.eventToClientCalls = 0 + t.eventsFromKinesisCalls = nil +} + +// WaitForRecordsInMemoryCalls waits until at least expectedCalls RecordsInMemory calls are made +// Returns true if the expected calls were made within the timeout +func (t *TestStatReceiver) WaitForRecordsInMemoryCalls(expectedCalls int, timeout time.Duration) bool { + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if len(t.GetRecordsInMemoryCalls()) >= expectedCalls { + return true + } + time.Sleep(10 * time.Millisecond) + } + return false +} + +// WaitForRecordsInMemoryBytesCalls waits until at least expectedCalls RecordsInMemoryBytes calls are made +// Returns true if the expected calls were made within the timeout +func (t *TestStatReceiver) WaitForRecordsInMemoryBytesCalls(expectedCalls int, timeout time.Duration) bool { + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if len(t.GetRecordsInMemoryBytesCalls()) >= expectedCalls { + return true + } + time.Sleep(10 * time.Millisecond) + } + return false +} + +// TestMetricsBasic tests that the new metrics are reported to StatReceiver as expected +func TestMetricsBasic(t *testing.T) { + if testing.Short() { + t.Skip("skipping test in short mode.") + } + + streamName := "TestMetricsBasic_stream" + testStats := &TestStatReceiver{} + + k, d := kinesisAndDynamoInstances(t) + + defer func() { + err := cleanupTestEnvironment(t, k, d, streamName) + require.NoError(t, err, "Problems cleaning up the test environment") + }() + + err := setupTestEnvironment(t, k, d, streamName, 1) + require.NoError(t, err, "Problems setting up the test environment") + + // Configure with our test stat receiver + config := NewConfig().WithBufferSize(100).WithStats(testStats) + config = config.WithShardCheckFrequency(500 * time.Millisecond) + config = config.WithLeaderActionFrequency(500 * time.Millisecond) + + kinsumer, err := NewWithInterfaces(k, d, streamName, *applicationName, "test_client", "", config) + require.NoError(t, err, "NewWithInterfaces() failed") + + err = kinsumer.Run() + require.NoError(t, err, "kinsumer.Run() failed") + defer kinsumer.Stop() + + // Send records through the system + recordCount := 200 + err = spamStream(t, k, int64(recordCount), streamName) + require.NoError(t, err, "Problems sending test data") + + // Wait for metrics to be reported (should happen every 1 second) + require.True(t, testStats.WaitForRecordsInMemoryCalls(2, 5*time.Second), + "Expected RecordsInMemory calls after sending records") + require.True(t, testStats.WaitForRecordsInMemoryBytesCalls(2, 5*time.Second), + "Expected RecordsInMemoryBytes calls after sending records") + + // Verify metrics after sending records + recordsAfterSend := testStats.GetLastRecordsInMemory() + bytesAfterSend := testStats.GetLastRecordsInMemoryBytes() + assert.Equal(t, int64(200), recordsAfterSend, "Should have 200 records in memory after sending") + assert.Equal(t, int64(490), bytesAfterSend, "Should have 490 bytes in memory after sending") // 200 records * ~2.45 bytes avg + + // Consume 100 records to trigger metric decreases (should cause 2 batches of 50 decrements) + var consumedCount int + timeout := time.After(15 * time.Second) + for consumedCount < 100 { + select { + case <-timeout: + t.Fatalf("Timeout consuming records, got %d", consumedCount) + default: + data, err := kinsumer.Next() + if err != nil { + t.Fatalf("Error from Next(): %v", err) + } + if data != nil { + consumedCount++ + } + } + } + + // Wait for additional metrics reporting after consumption + time.Sleep(2 * time.Second) + + // Verify metrics after consuming records + recordsAfterConsume := testStats.GetLastRecordsInMemory() + bytesAfterConsume := testStats.GetLastRecordsInMemoryBytes() + assert.Equal(t, int64(100), recordsAfterConsume, "Should have 100 records in memory after consuming 100") + assert.Equal(t, int64(300), bytesAfterConsume, "Should have 300 bytes in memory after consuming 100") + + // Verify the new metrics are working + recordCalls := testStats.GetRecordsInMemoryCalls() + byteCalls := testStats.GetRecordsInMemoryBytesCalls() + + assert.True(t, len(recordCalls) > 0, "Should have RecordsInMemory calls") + assert.True(t, len(byteCalls) > 0, "Should have RecordsInMemoryBytes calls") + assert.Equal(t, len(recordCalls), len(byteCalls), "Record and byte calls should match") + + // Verify at least one call showed non-zero values (records were in memory) + hasNonZeroRecords := false + hasNonZeroBytes := false + for _, count := range recordCalls { + if count > 0 { + hasNonZeroRecords = true + break + } + } + for _, bytes := range byteCalls { + if bytes > 0 { + hasNonZeroBytes = true + break + } + } + + assert.True(t, hasNonZeroRecords, "Should have non-zero records in memory at some point") + assert.True(t, hasNonZeroBytes, "Should have non-zero bytes in memory at some point") + + t.Logf("✓ Metrics test completed - RecordsInMemory calls: %v, RecordsInMemoryBytes calls: %v", + recordCalls, byteCalls) +} + +// TestMetricsFiltering tests that metrics can be selectively enabled/disabled using the filteredStatReceiver +func TestMetricsFiltering(t *testing.T) { + // Test that only RecordsInMemory metrics are captured when others are disabled + testStats := &TestStatReceiver{} + + // Configure to only enable RecordsInMemory + metricsConfig := MetricsConfig{ + EnableRecordsInMemory: true, + // All others default to false + } + + // Create a filtered stat receiver using the new approach (simulating WithStats behavior) + filteredStats := newFilteredStatReceiver(testStats, metricsConfig) + + // Call all methods on the filtered receiver + filteredStats.Checkpoint() + filteredStats.EventToClient(time.Now(), time.Now()) + filteredStats.EventsFromKinesis(5, "shard-001", time.Second) + filteredStats.RecordsInMemory(100) + filteredStats.RecordsInMemoryBytes(500) + + // Verify only RecordsInMemory was called on the underlying receiver + recordCount, byteCount, checkpointCount, eventToClientCount, eventsFromKinesisCount := testStats.GetCallCounts() + + assert.Equal(t, 1, recordCount, "Should have 1 RecordsInMemory call") + assert.Equal(t, 0, byteCount, "Should have 0 RecordsInMemoryBytes calls (disabled)") + assert.Equal(t, 0, checkpointCount, "Should have 0 Checkpoint calls (disabled)") + assert.Equal(t, 0, eventToClientCount, "Should have 0 EventToClient calls (disabled)") + assert.Equal(t, 0, eventsFromKinesisCount, "Should have 0 EventsFromKinesis calls (disabled)") +} + +// TestMetricsFilteringOrderIndependence tests that the order of WithStats() and WithMetricsConfig() doesn't matter +func TestMetricsFilteringOrderIndependence(t *testing.T) { + testStats1 := &TestStatReceiver{} + testStats2 := &TestStatReceiver{} + + metricsConfig := MetricsConfig{ + EnableRecordsInMemory: true, + // All others default to false + } + + // Test order 1: WithStats first, then WithMetricsConfig + config1 := NewConfig(). + WithStats(testStats1). + WithMetricsConfig(metricsConfig) + + // Test order 2: WithMetricsConfig first, then WithStats + config2 := NewConfig(). + WithMetricsConfig(metricsConfig). + WithStats(testStats2) + + // Simulate what happens in NewWithInterfaces - apply the filtering + filteredStats1 := newFilteredStatReceiver(config1.stats, config1.metricsConfig) + filteredStats2 := newFilteredStatReceiver(config2.stats, config2.metricsConfig) + + // Call all methods on both filtered receivers + filteredStats1.Checkpoint() + filteredStats1.RecordsInMemory(100) + filteredStats1.RecordsInMemoryBytes(500) + + filteredStats2.Checkpoint() + filteredStats2.RecordsInMemory(100) + filteredStats2.RecordsInMemoryBytes(500) + + // Both should have identical behavior regardless of order + recordCount1, byteCount1, checkpointCount1, _, _ := testStats1.GetCallCounts() + recordCount2, byteCount2, checkpointCount2, _, _ := testStats2.GetCallCounts() + + assert.Equal(t, recordCount1, recordCount2, "RecordsInMemory calls should be identical regardless of order") + assert.Equal(t, byteCount1, byteCount2, "RecordsInMemoryBytes calls should be identical regardless of order") + assert.Equal(t, checkpointCount1, checkpointCount2, "Checkpoint calls should be identical regardless of order") + + // Verify the expected filtering (only RecordsInMemory should be called) + assert.Equal(t, 1, recordCount1, "Should have 1 RecordsInMemory call") + assert.Equal(t, 0, byteCount1, "Should have 0 RecordsInMemoryBytes calls (disabled)") + assert.Equal(t, 0, checkpointCount1, "Should have 0 Checkpoint calls (disabled)") +} diff --git a/statsd/statsd.go b/statsd/statsd.go index fc828b0..ec08ba2 100644 --- a/statsd/statsd.go +++ b/statsd/statsd.go @@ -14,11 +14,13 @@ type Statsd struct { client statsd.StatSender } -// New creates a new Statsd statreceiver with a new instance of a cactus statter +// New creates a new Statsd statreceiver with buffering enabled by default func New(addr, prefix string) (*Statsd, error) { sd, err := statsd.NewClientWithConfig(&statsd.ClientConfig{ - Address: addr, - Prefix: prefix, + Address: addr, + Prefix: prefix, + UseBuffered: true, // Enable buffering by default + FlushInterval: 1 * time.Second, // Sensible default flush interval }) if err != nil { @@ -58,3 +60,15 @@ func (s *Statsd) EventsFromKinesis(num int, shardID string, lag time.Duration) { _ = s.client.TimingDuration(fmt.Sprintf("kinsumer.%s.lag", shardID), lag, 1.0) _ = s.client.Inc(fmt.Sprintf("kinsumer.%s.retrieved", shardID), int64(num), 1.0) } + +// RecordsInMemory implementation that writes to statsd a gauge metric about +// the current number of records buffered in memory +func (s *Statsd) RecordsInMemory(count int64) { + _ = s.client.Gauge("kinsumer.records_in_memory", count, 1.0) +} + +// RecordsInMemoryBytes implementation that writes to statsd a gauge metric about +// the current total payload bytes buffered in memory +func (s *Statsd) RecordsInMemoryBytes(bytes int64) { + _ = s.client.Gauge("kinsumer.records_in_memory_bytes", bytes, 1.0) +}