diff --git a/platform/consumer/BUILD.bazel b/platform/consumer/BUILD.bazel index cdc035058..5a04ed9fc 100644 --- a/platform/consumer/BUILD.bazel +++ b/platform/consumer/BUILD.bazel @@ -41,6 +41,8 @@ go_test( "@com_github_stretchr_testify//require:go_default_library", "@com_github_uber_go_tally//:go_default_library", "@org_uber_go_mock//gomock:go_default_library", + "@org_uber_go_zap//:go_default_library", "@org_uber_go_zap//zaptest:go_default_library", + "@org_uber_go_zap//zaptest/observer:go_default_library", ], ) diff --git a/platform/consumer/consumer.go b/platform/consumer/consumer.go index d94d20a9d..7aab1b906 100644 --- a/platform/consumer/consumer.go +++ b/platform/consumer/consumer.go @@ -414,13 +414,16 @@ func (m *consumer) processDelivery(ctx context.Context, controller Controller, d } topicKey := controller.TopicKey() + ownerLogFields := deliveryOwnerLogFields(delivery) m.logger.Debugw("processing delivery", - "controller", controller.Name(), - "topic_key", topicKey, - "message_id", msg.ID, - "partition_key", msg.PartitionKey, - "attempt", delivery.Attempt(), + append([]any{ + "controller", controller.Name(), + "topic_key", topicKey, + "message_id", msg.ID, + "partition_key", msg.PartitionKey, + "attempt", delivery.Attempt(), + }, ownerLogFields...)..., ) // Wrap delivery to hide Ack/Nack from controller @@ -453,10 +456,12 @@ func (m *consumer) processDelivery(ctx context.Context, controller Controller, d if wrapped.held { metrics.NamedCounter(controllerScope, opName, "hold_ignored", 1, metrics.TagsFromContext(ctx)...) m.logger.Warnw("hold recorded but controller returned error, failure outcome wins", - "controller", controller.Name(), - "topic_key", topicKey, - "message_id", msg.ID, - "partition_key", msg.PartitionKey, + append([]any{ + "controller", controller.Name(), + "topic_key", topicKey, + "message_id", msg.ID, + "partition_key", msg.PartitionKey, + }, ownerLogFields...)..., ) } @@ -473,13 +478,15 @@ func (m *consumer) processDelivery(ctx context.Context, controller Controller, d // Check if the error is non-retryable (poison pill message) if !errs.IsRetryable(err) { m.logger.Errorw("non-retryable controller error, rejecting message", - "controller", controller.Name(), - "topic_key", controller.TopicKey(), - "message_id", msg.ID, - "partition_key", msg.PartitionKey, - "attempt", delivery.Attempt(), - "error", err, - "elapsed_ms", elapsed.Milliseconds(), + append([]any{ + "controller", controller.Name(), + "topic_key", controller.TopicKey(), + "message_id", msg.ID, + "partition_key", msg.PartitionKey, + "attempt", delivery.Attempt(), + "error", err, + "elapsed_ms", elapsed.Milliseconds(), + }, ownerLogFields...)..., ) // Reject moves to DLQ (or acks if DLQ disabled) @@ -488,10 +495,12 @@ func (m *consumer) processDelivery(ctx context.Context, controller Controller, d rejectOp.Complete(rejectErr) if rejectErr != nil { m.logger.Errorw("failed to reject non-retryable message", - "controller", controller.Name(), - "topic_key", controller.TopicKey(), - "message_id", msg.ID, - "error", rejectErr, + append([]any{ + "controller", controller.Name(), + "topic_key", controller.TopicKey(), + "message_id", msg.ID, + "error", rejectErr, + }, ownerLogFields...)..., ) } return @@ -504,14 +513,16 @@ func (m *consumer) processDelivery(ctx context.Context, controller Controller, d what = "cancel" } m.logger.Errorw("controller error or cancel, nacking message", - "what", what, - "controller", controller.Name(), - "topic_key", topicKey, - "message_id", msg.ID, - "partition_key", msg.PartitionKey, - "attempt", delivery.Attempt(), - "error", err, - "elapsed_ms", elapsed.Milliseconds(), + append([]any{ + "what", what, + "controller", controller.Name(), + "topic_key", topicKey, + "message_id", msg.ID, + "partition_key", msg.PartitionKey, + "attempt", delivery.Attempt(), + "error", err, + "elapsed_ms", elapsed.Milliseconds(), + }, ownerLogFields...)..., ) nackOp := metrics.Begin(controllerScope, "nack", metrics.StorageLatencyBuckets, metrics.TagsFromContext(ctx)...) @@ -519,10 +530,12 @@ func (m *consumer) processDelivery(ctx context.Context, controller Controller, d nackOp.Complete(nackErr) if nackErr != nil { m.logger.Errorw("failed to nack message", - "controller", controller.Name(), - "topic_key", topicKey, - "message_id", msg.ID, - "error", nackErr, + append([]any{ + "controller", controller.Name(), + "topic_key", topicKey, + "message_id", msg.ID, + "error", nackErr, + }, ownerLogFields...)..., ) } return @@ -764,3 +777,18 @@ func (m *consumer) unsubscribeAll(timeoutMs int64) error { m.logger.Debugw("all controllers stopped gracefully") return nil } + +func deliveryOwnerLogFields(delivery extqueue.Delivery) []any { + meta := delivery.Metadata() + if len(meta) == 0 { + return nil + } + var fields []any + if v := meta["leased_by"]; v != "" { + fields = append(fields, "leased_by", v) + } + if v := meta["consumer_group"]; v != "" { + fields = append(fields, "consumer_group", v) + } + return fields +} diff --git a/platform/consumer/consumer_test.go b/platform/consumer/consumer_test.go index ff3d68cfa..61de23d91 100644 --- a/platform/consumer/consumer_test.go +++ b/platform/consumer/consumer_test.go @@ -35,7 +35,9 @@ import ( queuemock "github.com/uber/submitqueue/platform/extension/messagequeue/mock" "github.com/uber/submitqueue/platform/metrics" "go.uber.org/mock/gomock" + "go.uber.org/zap" "go.uber.org/zap/zaptest" + "go.uber.org/zap/zaptest/observer" ) const ( @@ -627,6 +629,123 @@ func TestConsumer_ProcessDelivery_NonRetryableError(t *testing.T) { require.NoError(t, err) } +func TestConsumer_ProcessDelivery_LogsLeasedBy(t *testing.T) { + ctrl := gomock.NewController(t) + core, logs := observer.New(zap.DebugLevel) + logger := zap.New(core).Sugar() + + deliveryChan := make(chan extqueue.Delivery, 1) + mockSub := queuemock.NewMockSubscriber(ctrl) + mockSub.EXPECT().Subscribe(gomock.Any(), gomock.Any(), gomock.Any()).Return(deliveryChan, nil) + + mockQ := queuemock.NewMockQueue(ctrl) + mockQ.EXPECT().Subscriber().Return(mockSub) + + reg := newRegistry(t, mockQ, testTopicKeyStart, "test-group") + c := New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) + + handler := &testController{} + setupController(handler, "test-handler", testTopicKeyStart, "test-group", + func(ctx context.Context, delivery Delivery) error { + return fmt.Errorf("bad payload") + }, + ) + require.NoError(t, c.Register(handler)) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + require.NoError(t, c.Start(ctx)) + + msg := entityqueue.NewMessage("poison-msg", []byte("bad"), "partition1", nil) + msg.Tenant = testTenant + done := make(chan struct{}) + mockDel := queuemock.NewMockDelivery(ctrl) + mockDel.EXPECT().Message().Return(msg).AnyTimes() + mockDel.EXPECT().Attempt().Return(1).AnyTimes() + mockDel.EXPECT().ReceivedAt().Return(time.Now().UnixMilli()).AnyTimes() + mockDel.EXPECT().Metadata().Return(map[string]string{ + "leased_by": "host-1", + "consumer_group": "test-group", + }).AnyTimes() + mockDel.EXPECT().DeliveryID().Return(msg.ID).AnyTimes() + mockDel.EXPECT().Reject(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, _ failure.Failure) error { + close(done) + return nil + }) + + deliveryChan <- mockDel + <-done + require.NoError(t, c.Stop(30000)) + + processLogs := logs.FilterMessage("processing delivery").All() + require.NotEmpty(t, processLogs) + assert.Equal(t, "host-1", processLogs[0].ContextMap()["leased_by"]) + assert.Equal(t, "test-group", processLogs[0].ContextMap()["consumer_group"]) + + rejectLogs := logs.FilterMessage("non-retryable controller error, rejecting message").All() + require.NotEmpty(t, rejectLogs) + assert.Equal(t, "host-1", rejectLogs[0].ContextMap()["leased_by"]) + assert.Equal(t, "partition1", rejectLogs[0].ContextMap()["partition_key"]) +} + +func TestConsumer_ProcessDelivery_HoldIgnoredLogsLeasedBy(t *testing.T) { + ctrl := gomock.NewController(t) + core, logs := observer.New(zap.DebugLevel) + logger := zap.New(core).Sugar() + + deliveryChan := make(chan extqueue.Delivery, 1) + mockSub := queuemock.NewMockSubscriber(ctrl) + mockSub.EXPECT().Subscribe(gomock.Any(), gomock.Any(), gomock.Any()).Return(deliveryChan, nil) + + mockQ := queuemock.NewMockQueue(ctrl) + mockQ.EXPECT().Subscriber().Return(mockSub) + + reg := newRegistry(t, mockQ, testTopicKeyStart, "test-group") + c := New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) + + deliveryMetadata := map[string]string{ + "leased_by": "host-1", + "consumer_group": "test-group", + } + handler := &testController{} + setupController(handler, "test-handler", testTopicKeyStart, "test-group", + func(ctx context.Context, delivery Delivery) error { + delivery.Metadata()["leased_by"] = "mutated-host" + delivery.Metadata()["consumer_group"] = "mutated-group" + delivery.Hold(5000) + return errs.NewRetryableError(fmt.Errorf("processing failed")) + }, + ) + require.NoError(t, c.Register(handler)) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + require.NoError(t, c.Start(ctx)) + + msg := entityqueue.NewMessage("held-msg", []byte("payload"), "partition1", nil) + msg.Tenant = testTenant + done := make(chan struct{}) + mockDel := queuemock.NewMockDelivery(ctrl) + mockDel.EXPECT().Message().Return(msg).AnyTimes() + mockDel.EXPECT().Attempt().Return(1).AnyTimes() + mockDel.EXPECT().ReceivedAt().Return(time.Now().UnixMilli()).AnyTimes() + mockDel.EXPECT().Metadata().Return(deliveryMetadata).AnyTimes() + mockDel.EXPECT().DeliveryID().Return(msg.ID).AnyTimes() + mockDel.EXPECT().Nack(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, _ failure.Failure) error { + close(done) + return nil + }) + + deliveryChan <- mockDel + <-done + require.NoError(t, c.Stop(30000)) + + holdLogs := logs.FilterMessage("hold recorded but controller returned error, failure outcome wins").All() + require.Len(t, holdLogs, 1) + assert.Equal(t, "host-1", holdLogs[0].ContextMap()["leased_by"]) + assert.Equal(t, "test-group", holdLogs[0].ContextMap()["consumer_group"]) +} + // The failure handed to the queue is built from whatever the controller // attributed, and a controller that attributes nothing must still produce // exactly what callers sent before failures carried structure: the error text @@ -1176,6 +1295,7 @@ func TestConsumer_SamePartitionKeyAcrossTenantsProcessesIndependently(t *testing delA := queuemock.NewMockDelivery(ctrl) delA.EXPECT().Message().Return(msgA).AnyTimes() delA.EXPECT().Attempt().Return(1).AnyTimes() + delA.EXPECT().Metadata().Return(nil).AnyTimes() delA.EXPECT().Ack(gomock.Any()).Return(nil).MaxTimes(1) deliveryChan <- delA <-tenantABlocked @@ -1185,6 +1305,7 @@ func TestConsumer_SamePartitionKeyAcrossTenantsProcessesIndependently(t *testing delB := queuemock.NewMockDelivery(ctrl) delB.EXPECT().Message().Return(msgB).AnyTimes() delB.EXPECT().Attempt().Return(1).AnyTimes() + delB.EXPECT().Metadata().Return(nil).AnyTimes() delB.EXPECT().Ack(gomock.Any()).Return(nil).MaxTimes(1) deliveryChan <- delB <-tenantBProcessed diff --git a/platform/extension/messagequeue/mysql/BUILD.bazel b/platform/extension/messagequeue/mysql/BUILD.bazel index ea6bf5c3b..d59b3f4dd 100644 --- a/platform/extension/messagequeue/mysql/BUILD.bazel +++ b/platform/extension/messagequeue/mysql/BUILD.bazel @@ -56,6 +56,7 @@ go_test( "@com_github_uber_go_tally//:go_default_library", "@org_uber_go_mock//gomock:go_default_library", "@org_uber_go_zap//:go_default_library", + "@org_uber_go_zap//zapcore:go_default_library", "@org_uber_go_zap//zaptest:go_default_library", "@org_uber_go_zap//zaptest/observer:go_default_library", ], diff --git a/platform/extension/messagequeue/mysql/constants.go b/platform/extension/messagequeue/mysql/constants.go index ca9da3adc..14c2518fa 100644 --- a/platform/extension/messagequeue/mysql/constants.go +++ b/platform/extension/messagequeue/mysql/constants.go @@ -18,9 +18,15 @@ package mysql const ( // Common log field names (used extensively across all stores) - logTenant = "tenant" - logTopic = "topic" - logPartitionKey = "partition_key" - logMessageID = "message_id" - logError = "error" + logTenant = "tenant" + logTopic = "topic" + logPartitionKey = "partition_key" + logMessageID = "message_id" + logError = "error" + logLeasedBy = "leased_by" + logConsumerGroup = "consumer_group" + logPreviousOwner = "previous_owner" + logReason = "reason" + logOwnedPartitions = "owned_partitions" + logRenewed = "renewed" ) diff --git a/platform/extension/messagequeue/mysql/mock_stores.go b/platform/extension/messagequeue/mysql/mock_stores.go index b9ee70448..57239ee6c 100644 --- a/platform/extension/messagequeue/mysql/mock_stores.go +++ b/platform/extension/messagequeue/mysql/mock_stores.go @@ -296,11 +296,12 @@ func (mr *MockpartitionLeaseStoreMockRecorder) PurgeStaleForTenants(ctx, tenants } // ReleaseLease mocks base method. -func (m *MockpartitionLeaseStore) ReleaseLease(ctx context.Context, tenant, topic, partitionKey, subscriberName, consumerGroup string) error { +func (m *MockpartitionLeaseStore) ReleaseLease(ctx context.Context, tenant, topic, partitionKey, subscriberName, consumerGroup string) (int64, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "ReleaseLease", ctx, tenant, topic, partitionKey, subscriberName, consumerGroup) - ret0, _ := ret[0].(error) - return ret0 + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 } // ReleaseLease indicates an expected call of ReleaseLease. @@ -310,11 +311,12 @@ func (mr *MockpartitionLeaseStoreMockRecorder) ReleaseLease(ctx, tenant, topic, } // ReleaseOwnedLeases mocks base method. -func (m *MockpartitionLeaseStore) ReleaseOwnedLeases(ctx context.Context, tenants []string, topic, subscriberName, consumerGroup string) error { +func (m *MockpartitionLeaseStore) ReleaseOwnedLeases(ctx context.Context, tenants []string, topic, subscriberName, consumerGroup string) (int64, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "ReleaseOwnedLeases", ctx, tenants, topic, subscriberName, consumerGroup) - ret0, _ := ret[0].(error) - return ret0 + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 } // ReleaseOwnedLeases indicates an expected call of ReleaseOwnedLeases. @@ -324,11 +326,12 @@ func (mr *MockpartitionLeaseStoreMockRecorder) ReleaseOwnedLeases(ctx, tenants, } // RenewOwnedLeases mocks base method. -func (m *MockpartitionLeaseStore) RenewOwnedLeases(ctx context.Context, tenants []string, topic, subscriberName, consumerGroup string) error { +func (m *MockpartitionLeaseStore) RenewOwnedLeases(ctx context.Context, tenants []string, topic, subscriberName, consumerGroup string) (int64, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "RenewOwnedLeases", ctx, tenants, topic, subscriberName, consumerGroup) - ret0, _ := ret[0].(error) - return ret0 + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 } // RenewOwnedLeases indicates an expected call of RenewOwnedLeases. diff --git a/platform/extension/messagequeue/mysql/partition_lease_store.go b/platform/extension/messagequeue/mysql/partition_lease_store.go index 1395b1ef2..9ccc53377 100644 --- a/platform/extension/messagequeue/mysql/partition_lease_store.go +++ b/platform/extension/messagequeue/mysql/partition_lease_store.go @@ -92,7 +92,7 @@ func (s *sqlpartitionLeaseStore) TryAcquireLease(ctx context.Context, tenant str } // ReleaseLease releases the lease for a partition owned by this worker -func (s *sqlpartitionLeaseStore) ReleaseLease(ctx context.Context, tenant string, topic string, partitionKey string, subscriberName string, consumerGroup string) (retErr error) { +func (s *sqlpartitionLeaseStore) ReleaseLease(ctx context.Context, tenant string, topic string, partitionKey string, subscriberName string, consumerGroup string) (_ int64, retErr error) { op := metrics.Begin(s.scope, "release_lease", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() @@ -102,7 +102,7 @@ func (s *sqlpartitionLeaseStore) ReleaseLease(ctx context.Context, tenant string `, PartitionLeasesTableName), tenant, consumerGroup, topic, partitionKey, subscriberName) if err != nil { - return fmt.Errorf("release lease tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) + return 0, fmt.Errorf("release lease tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } // RowsAffected error is swallowed because the DELETE query itself succeeded. @@ -116,6 +116,7 @@ func (s *sqlpartitionLeaseStore) ReleaseLease(ctx context.Context, tenant string logPartitionKey, partitionKey, logError, err, ) + return 0, nil } if rows > 0 { s.logger.Debugw("released lease", @@ -125,7 +126,7 @@ func (s *sqlpartitionLeaseStore) ReleaseLease(ctx context.Context, tenant string ) } - return nil + return rows, nil } func (s *sqlpartitionLeaseStore) GetLeasedPartitionsForTenants(ctx context.Context, tenants []string, topic string, subscriberName string, consumerGroup string) (_ map[string][]string, retErr error) { @@ -197,49 +198,65 @@ func (s *sqlpartitionLeaseStore) GetAllLeasesForTenants(ctx context.Context, ten return byTenant, nil } -func (s *sqlpartitionLeaseStore) RenewOwnedLeases(ctx context.Context, tenants []string, topic string, subscriberName string, consumerGroup string) (retErr error) { +func (s *sqlpartitionLeaseStore) RenewOwnedLeases(ctx context.Context, tenants []string, topic string, subscriberName string, consumerGroup string) (_ int64, retErr error) { op := metrics.Begin(s.scope, "renew_owned_leases", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() placeholders, ok := inListPlaceholders(len(tenants)) if !ok { - return nil + return 0, nil } now := currentTimeMillis() args := []any{now} args = appendStrings(args, tenants) args = append(args, consumerGroup, topic, subscriberName) - _, err := s.db.ExecContext(ctx, fmt.Sprintf(` + result, err := s.db.ExecContext(ctx, fmt.Sprintf(` UPDATE %s SET lease_renewed_at = ? WHERE tenant IN (%s) AND consumer_group = ? AND topic = ? AND leased_by = ? `, PartitionLeasesTableName, placeholders), args...) if err != nil { - return fmt.Errorf("renew owned leases topic=%s: %w", topic, err) + return 0, fmt.Errorf("renew owned leases topic=%s: %w", topic, err) } - return nil + rows, err := result.RowsAffected() + if err != nil { + s.logger.Warnw("failed to get rows affected after renew owned leases", + logTopic, topic, + logError, err, + ) + return 0, nil + } + return rows, nil } -func (s *sqlpartitionLeaseStore) ReleaseOwnedLeases(ctx context.Context, tenants []string, topic string, subscriberName string, consumerGroup string) (retErr error) { +func (s *sqlpartitionLeaseStore) ReleaseOwnedLeases(ctx context.Context, tenants []string, topic string, subscriberName string, consumerGroup string) (_ int64, retErr error) { op := metrics.Begin(s.scope, "release_owned_leases", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() placeholders, ok := inListPlaceholders(len(tenants)) if !ok { - return nil + return 0, nil } args := appendStrings(nil, tenants) args = append(args, consumerGroup, topic, subscriberName) - _, err := s.db.ExecContext(ctx, fmt.Sprintf(` + result, err := s.db.ExecContext(ctx, fmt.Sprintf(` DELETE FROM %s WHERE tenant IN (%s) AND consumer_group = ? AND topic = ? AND leased_by = ? `, PartitionLeasesTableName, placeholders), args...) if err != nil { - return fmt.Errorf("release owned leases topic=%s: %w", topic, err) + return 0, fmt.Errorf("release owned leases topic=%s: %w", topic, err) } - return nil + rows, err := result.RowsAffected() + if err != nil { + s.logger.Warnw("failed to get rows affected after release owned leases", + logTopic, topic, + logError, err, + ) + return 0, nil + } + return rows, nil } func (s *sqlpartitionLeaseStore) PurgeStaleForTenants(ctx context.Context, tenants []string, topic string, consumerGroup string, olderThanMs int64) (retErr error) { diff --git a/platform/extension/messagequeue/mysql/partition_lease_store_test.go b/platform/extension/messagequeue/mysql/partition_lease_store_test.go index 7ea852ba1..760e69e9a 100644 --- a/platform/extension/messagequeue/mysql/partition_lease_store_test.go +++ b/platform/extension/messagequeue/mysql/partition_lease_store_test.go @@ -17,9 +17,11 @@ package mysql import ( "context" "database/sql" + "errors" "testing" "github.com/DATA-DOG/go-sqlmock" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/uber-go/tally" "go.uber.org/zap/zaptest" @@ -99,9 +101,10 @@ func TestPartitionLeaseStore_TryAcquireLease(t *testing.T) { func TestPartitionLeaseStore_ReleaseLease(t *testing.T) { tests := []struct { - name string - setup func(mock sqlmock.Sqlmock) - wantErr bool + name string + setup func(mock sqlmock.Sqlmock) + wantRows int64 + wantErr bool }{ { name: "successfully release lease", @@ -110,7 +113,8 @@ func TestPartitionLeaseStore_ReleaseLease(t *testing.T) { WithArgs(testTenant, testConsumerGroup, "test_topic", "part1", testSubscriberName). WillReturnResult(sqlmock.NewResult(0, 1)) }, - wantErr: false, + wantRows: 1, + wantErr: false, }, { name: "idempotent - already released", @@ -119,7 +123,18 @@ func TestPartitionLeaseStore_ReleaseLease(t *testing.T) { WithArgs(testTenant, testConsumerGroup, "test_topic", "part1", testSubscriberName). WillReturnResult(sqlmock.NewResult(0, 0)) }, - wantErr: false, + wantRows: 0, + wantErr: false, + }, + { + name: "rows affected unavailable", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("DELETE FROM queue_partition_leases"). + WithArgs(testTenant, testConsumerGroup, "test_topic", "part1", testSubscriberName). + WillReturnResult(sqlmock.NewErrorResult(errors.New("rows affected unsupported"))) + }, + wantRows: 0, + wantErr: false, }, } @@ -134,12 +149,13 @@ func TestPartitionLeaseStore_ReleaseLease(t *testing.T) { tt.setup(mock) - err := store.ReleaseLease(ctx, testTenant, topic, partitionKey, testSubscriberName, testConsumerGroup) + rows, err := store.ReleaseLease(ctx, testTenant, topic, partitionKey, testSubscriberName, testConsumerGroup) if tt.wantErr { require.Error(t, err) } else { require.NoError(t, err) } + assert.Equal(t, tt.wantRows, rows) require.NoError(t, mock.ExpectationsWereMet()) }) } @@ -231,8 +247,12 @@ func TestPartitionLeaseStore_RenewAndReleaseOwnedLeases(t *testing.T) { WithArgs("alpha", "beta", testConsumerGroup, "test_topic", sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(0, 1)) - require.NoError(t, store.RenewOwnedLeases(context.Background(), tenants, "test_topic", testSubscriberName, testConsumerGroup)) - require.NoError(t, store.ReleaseOwnedLeases(context.Background(), tenants, "test_topic", testSubscriberName, testConsumerGroup)) + renewed, err := store.RenewOwnedLeases(context.Background(), tenants, "test_topic", testSubscriberName, testConsumerGroup) + require.NoError(t, err) + assert.Equal(t, int64(2), renewed) + released, err := store.ReleaseOwnedLeases(context.Background(), tenants, "test_topic", testSubscriberName, testConsumerGroup) + require.NoError(t, err) + assert.Equal(t, int64(2), released) require.NoError(t, store.PurgeStaleForTenants(context.Background(), tenants, "test_topic", testConsumerGroup, 300_000)) require.NoError(t, mock.ExpectationsWereMet()) } diff --git a/platform/extension/messagequeue/mysql/stores.go b/platform/extension/messagequeue/mysql/stores.go index 496dc5af3..b9ff1957f 100644 --- a/platform/extension/messagequeue/mysql/stores.go +++ b/platform/extension/messagequeue/mysql/stores.go @@ -116,8 +116,11 @@ type partitionLeaseStore interface { // TryAcquireLease attempts to acquire or renew a lease for a partition TryAcquireLease(ctx context.Context, tenant string, topic string, partitionKey string, subscriberName string, consumerGroup string, leaseDurationMs int64) (bool, error) - // ReleaseLease releases the lease for a partition owned by this worker - ReleaseLease(ctx context.Context, tenant string, topic string, partitionKey string, subscriberName string, consumerGroup string) error + // ReleaseLease releases the lease for a partition owned by this worker. + // The int64 is rows deleted; 0 means the lease was already gone (stolen or + // purged) or the backend could not report a count, neither of which is an + // error. Callers must treat the partition as released regardless. + ReleaseLease(ctx context.Context, tenant string, topic string, partitionKey string, subscriberName string, consumerGroup string) (int64, error) // GetLeasedPartitionsForTenants returns partitions leased by this worker // across the given tenants, keyed by tenant. Tenants with no leases are absent. @@ -129,10 +132,12 @@ type partitionLeaseStore interface { // RenewOwnedLeases refreshes lease_renewed_at on every row this subscriber // still holds across tenants. Stolen leases are ignored; discovery drops them. - RenewOwnedLeases(ctx context.Context, tenants []string, topic string, subscriberName string, consumerGroup string) error + // The int64 is rows updated. + RenewOwnedLeases(ctx context.Context, tenants []string, topic string, subscriberName string, consumerGroup string) (int64, error) // ReleaseOwnedLeases deletes every lease row this subscriber holds across tenants. - ReleaseOwnedLeases(ctx context.Context, tenants []string, topic string, subscriberName string, consumerGroup string) error + // The int64 is rows deleted. + ReleaseOwnedLeases(ctx context.Context, tenants []string, topic string, subscriberName string, consumerGroup string) (int64, error) // PurgeStaleForTenants deletes lease rows not renewed within olderThanMs // across the given tenants. diff --git a/platform/extension/messagequeue/mysql/subscriber.go b/platform/extension/messagequeue/mysql/subscriber.go index 26657d0b2..fa17776c1 100644 --- a/platform/extension/messagequeue/mysql/subscriber.go +++ b/platform/extension/messagequeue/mysql/subscriber.go @@ -72,6 +72,10 @@ const ( // refresh or remove) a stale lease on a partition with no messages. leasePurgeAfterLeaseDurations = 10 + leaseReasonRebalance = "rebalance" + leaseReasonIdle = "idle" + leaseReasonShutdown = "shutdown" + // maxRetryBackoffMs bounds how long one failed message can pin its // partition's contiguous ack watermark. Callers may choose a lower cap; // this ceiling also applies when MaxBackoffMs is unset. @@ -627,9 +631,9 @@ func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) { cfg := sub.config // Common log fields for all operations in this subscription's lifecycle. logFields := []interface{}{ - "topic", sub.topic, - "consumer_group", cfg.ConsumerGroup, - "subscriber_name", cfg.SubscriberName, + logTopic, sub.topic, + logConsumerGroup, cfg.ConsumerGroup, + logLeasedBy, cfg.SubscriberName, } discoveryTicker := time.NewTicker(time.Duration(cfg.PartitionDiscoveryIntervalMs) * time.Millisecond) @@ -764,6 +768,7 @@ func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subsc if len(acquired) > 0 { leasedByTenant[tenant] = append(append([]string{}, leased...), acquired...) } + s.observeOwnedPartitions(sub, tenant, len(leasedByTenant[tenant])) } } @@ -814,7 +819,8 @@ func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subsc for _, partition := range expired { tenant := partition.Tenant pk := partition.PartitionKey - if err := s.leaseStore.ReleaseLease(ctx, tenant, sub.topic, pk, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { + rows, err := s.leaseStore.ReleaseLease(ctx, tenant, sub.topic, pk, cfg.SubscriberName, cfg.ConsumerGroup) + if err != nil { s.logger.Errorw("release lease for drained partition failed", "tenant", tenant, "topic", sub.topic, @@ -825,6 +831,12 @@ func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subsc } released[partition] = struct{}{} s.stopPartitionWorker(sub, partition) + if rows > 0 { + s.recordLeaseReleased(sub, tenant, leaseReasonIdle, rows) + s.logger.Infow("released idle partition lease", + append(leaseLogFields(sub, tenant, pk), logReason, leaseReasonIdle)..., + ) + } if err := s.offsetStore.DeleteOffset(ctx, tenant, sub.topic, pk, cfg.ConsumerGroup); err != nil { s.logger.Errorw("delete offsets for drained partition failed", @@ -836,14 +848,6 @@ func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subsc continue } delete(sub.drainedSince, partition) - - metrics.NamedCounter(s.scope, "idle_lease", "released", 1, metrics.NewTag("topic", sub.topic)) - s.logger.Infow("released idle partition lease", - "tenant", tenant, - "topic", sub.topic, - "consumer_group", cfg.ConsumerGroup, - "partition_key", pk, - ) } if len(released) > 0 { kept := make([]entityqueue.PartitionIdentity, 0, len(allLeased)) @@ -853,6 +857,13 @@ func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subsc } } allLeased = kept + ownedAfter := make(map[string]int, len(s.tenants)) + for _, partition := range allLeased { + ownedAfter[partition.Tenant]++ + } + for _, tenant := range s.tenants { + s.observeOwnedPartitions(sub, tenant, ownedAfter[tenant]) + } } } @@ -872,6 +883,7 @@ func (s *subscriber) acquireUnownedPartitions(ctx context.Context, sub *subscrip ownedCount := 0 ownedSet := make(map[string]struct{}) heldByOther := make(map[string]struct{}) + previousOwner := make(map[string]string) for _, lease := range leases { switch { case lease.LeasedBy == cfg.SubscriberName: @@ -879,6 +891,8 @@ func (s *subscriber) acquireUnownedPartitions(ctx context.Context, sub *subscrip ownedCount++ case lease.LeaseRenewedAt >= staleThreshold: heldByOther[lease.PartitionKey] = struct{}{} + default: + previousOwner[lease.PartitionKey] = lease.LeasedBy // pre-acquire snapshot, not the SQL winner } } @@ -896,21 +910,95 @@ func (s *subscriber) acquireUnownedPartitions(ctx context.Context, sub *subscrip ok, err := s.leaseStore.TryAcquireLease(ctx, tenant, sub.topic, partitionKey, cfg.SubscriberName, cfg.ConsumerGroup, cfg.LeaseDurationMs) if err != nil { s.logger.Errorw("failed to acquire lease for partition", - "tenant", tenant, - "topic", sub.topic, - "partition_key", partitionKey, - "error", err, + append(leaseLogFields(sub, tenant, partitionKey), logError, err)..., ) continue } if ok { acquired = append(acquired, partitionKey) ownedCount++ + if prior, stolen := previousOwner[partitionKey]; stolen && prior != "" && prior != cfg.SubscriberName { + s.recordLeaseCounter(sub, tenant, "stolen", 1) + s.logger.Infow("lease stolen", + append(leaseLogFields(sub, tenant, partitionKey), logPreviousOwner, prior)..., + ) + } else { + s.recordLeaseCounter(sub, tenant, "acquired", 1) + s.logger.Debugw("lease acquired", leaseLogFields(sub, tenant, partitionKey)...) + } } } return acquired } +func leaseLogFields(sub *subscription, tenant, partitionKey string) []any { + fields := []any{ + logTenant, tenant, + logTopic, sub.topic, + logConsumerGroup, sub.config.ConsumerGroup, + logLeasedBy, sub.config.SubscriberName, + } + if partitionKey != "" { + fields = append(fields, logPartitionKey, partitionKey) + } + return fields +} + +func (s *subscriber) leaseTags(sub *subscription, tenant string) []metrics.Tag { + tags := []metrics.Tag{ + metrics.NewTag(logLeasedBy, sub.config.SubscriberName), + metrics.NewTag("topic", sub.topic), + metrics.NewTag(logConsumerGroup, sub.config.ConsumerGroup), + } + if tenant != "" { + tags = append(tags, metrics.NewTag(logTenant, tenant)) + } + return tags +} + +func (s *subscriber) recordLeaseCounter(sub *subscription, tenant, counter string, n int64) { + metrics.NamedCounter(s.scope, "lease", counter, n, s.leaseTags(sub, tenant)...) +} + +func (s *subscriber) recordLeaseReleased(sub *subscription, tenant, reason string, n int64) { + tags := append(s.leaseTags(sub, tenant), metrics.NewTag(logReason, reason)) + metrics.NamedCounter(s.scope, "lease", "released", n, tags...) +} + +func (s *subscriber) observeOwnedPartitions(sub *subscription, tenant string, owned int) { + metrics.NamedGauge(s.scope, "lease", "partitions_owned", float64(owned), s.leaseTags(sub, tenant)...) +} + +func (s *subscriber) observeLeaseShare(sub *subscription, tenant string, active, cap int) { + tags := s.leaseTags(sub, tenant) + metrics.NamedGauge(s.scope, "lease", "active_subscribers", float64(active), tags...) + metrics.NamedGauge(s.scope, "lease", "fair_share_cap", float64(cap), tags...) +} + +// reportedFairShareCap maps unlimited (shareCap 0) to owned ∪ discovered so partitions_owned > fair_share_cap stays false for a lone replica. +func reportedFairShareCap(shareCap int, owned, discovered []string) int { + if shareCap > 0 { + return shareCap + } + return partitionUniverseSize(owned, discovered) +} + +func partitionUniverseSize(owned, discovered []string) int { + set := make(map[string]struct{}, len(owned)+len(discovered)) + for _, pk := range owned { + set[pk] = struct{}{} + } + for _, pk := range discovered { + set[pk] = struct{}{} + } + return len(set) +} + +func (s *subscriber) observeLeaseOwnership(sub *subscription, tenant string, owned, active, cap int) { + s.observeOwnedPartitions(sub, tenant, owned) + s.observeLeaseShare(sub, tenant, active, cap) +} + // updateDrainedTracking recomputes, for every owned partition absent from // this tick's discovery, when it was first observed drained. A partition is // drained only when zero of its messages remain stored — in-flight, @@ -1236,9 +1324,11 @@ func (w *partitionWorker) pollAndDeliver(ctx context.Context) (retErr error) { // Create delivery metadata deliveryMetadata := map[string]string{ - "topic": sub.topic, - "partition_key": partitionKey, - "offset": deliveryID, + "topic": sub.topic, + "partition_key": partitionKey, + "offset": deliveryID, + logLeasedBy: cfg.SubscriberName, + logConsumerGroup: cfg.ConsumerGroup, } // Add DLQ-specific metadata if this is a DLQ message @@ -1364,24 +1454,68 @@ func (s *subscriber) runLeaseTick(ctx context.Context, sub *subscription, timeou tickCtx, cancel := context.WithTimeout(ctx, timeout) defer cancel() + ownedPartitions := 0 + activeSubscribers := 0 + fairShareCap := 0 + shareObserved := false + var remainingByTenant map[string]int leasedByTenant, err := s.leaseStore.GetLeasedPartitionsForTenants(tickCtx, s.tenants, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) if err != nil { s.logger.Errorw("get leased partitions failed", append(logFields, "error", err)...) } else { + remainingByTenant = make(map[string]int, len(s.tenants)) + for _, tenant := range s.tenants { + n := len(leasedByTenant[tenant]) + ownedPartitions += n + remainingByTenant[tenant] = n + } activeByTenant, err := s.heartbeatStore.ActiveSubscribersForTenants(tickCtx, s.tenants, sub.topic, cfg.ConsumerGroup, cfg.LeaseDurationMs) if err != nil { s.logger.Errorw("active subscribers failed", append(logFields, "error", err)...) } else { + shareObserved = true + sub.workersMu.Lock() + discovered := append([]entityqueue.PartitionIdentity(nil), sub.lastDiscoveredPartitions...) + sub.workersMu.Unlock() for _, tenant := range s.tenants { - if _, err := s.rebalance(tickCtx, sub, tenant, leasedByTenant[tenant], activeByTenant[tenant]); err != nil { + owned := leasedByTenant[tenant] + active := activeByTenant[tenant] + discoveredKeys := partitionKeysForTenant(discovered, tenant) + shareCap := s.fairShareCap(sub, owned, discoveredKeys, active) + cap := reportedFairShareCap(shareCap, owned, discoveredKeys) + activeSubscribers += len(active) + fairShareCap += cap + s.observeLeaseOwnership(sub, tenant, len(owned), len(active), cap) + released, err := s.rebalance(tickCtx, sub, tenant, owned, active) + if err != nil { s.logger.Errorw("rebalance failed", append(logFields, "tenant", tenant, "error", err)...) + continue + } + if n := len(released); n > 0 { + ownedPartitions -= n + remainingByTenant[tenant] = len(owned) - n + s.observeLeaseOwnership(sub, tenant, remainingByTenant[tenant], len(active), cap) } } } } - if err := s.leaseStore.RenewOwnedLeases(tickCtx, s.tenants, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { + renewed, err := s.leaseStore.RenewOwnedLeases(tickCtx, s.tenants, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) + if err != nil { s.logger.Errorw("lease renewal failed", append(logFields, "error", err)...) + } else { + s.recordLeaseRenewed(sub, remainingByTenant, renewed) + } + tickFields := append(append([]any{}, logFields...), + logOwnedPartitions, ownedPartitions, + logRenewed, renewed, + ) + if len(s.tenants) == 1 && shareObserved { + tickFields = append(tickFields, + "active_subscribers", activeSubscribers, + "fair_share_cap", fairShareCap, + ) } + s.logger.Infow("lease tick", tickFields...) if err := s.sendHeartbeats(tickCtx, sub); err != nil { s.logger.Errorw("periodic heartbeat failed", append(logFields, "error", err)...) } @@ -1393,9 +1527,58 @@ func (s *subscriber) runLeaseTick(ctx context.Context, sub *subscription, timeou } } +func (s *subscriber) recordLeaseRenewed(sub *subscription, remainingByTenant map[string]int, renewed int64) { + if renewed <= 0 { + return + } + if remainingByTenant != nil { + var sum int64 + for _, tenant := range s.tenants { + sum += int64(remainingByTenant[tenant]) + } + if sum == renewed { + for _, tenant := range s.tenants { + if n := remainingByTenant[tenant]; n > 0 { + s.recordLeaseCounter(sub, tenant, "renewed", int64(n)) + } + } + return + } + } + metrics.NamedCounter(s.scope, "lease", "renewed", renewed, + metrics.NewTag(logLeasedBy, sub.config.SubscriberName), + metrics.NewTag("topic", sub.topic), + metrics.NewTag(logConsumerGroup, sub.config.ConsumerGroup), + ) +} + func (s *subscriber) releaseAllLeases(ctx context.Context, sub *subscription) error { cfg := sub.config - return s.leaseStore.ReleaseOwnedLeases(ctx, s.tenants, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) + var released int64 + var firstErr error + for _, tenant := range s.tenants { + n, err := s.leaseStore.ReleaseOwnedLeases(ctx, []string{tenant}, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) + if err != nil { + if firstErr == nil { + firstErr = fmt.Errorf("release owned leases tenant=%s: %w", tenant, err) + } + continue + } + if n > 0 { + released += n + s.recordLeaseReleased(sub, tenant, leaseReasonShutdown, n) + } + } + if released > 0 { + s.logger.Infow("released leases on shutdown", + "topic", sub.topic, + logConsumerGroup, cfg.ConsumerGroup, + logLeasedBy, cfg.SubscriberName, + logReason, leaseReasonShutdown, + "released", released, + ) + } + return firstErr } func (s *subscriber) sendHeartbeats(ctx context.Context, sub *subscription) error { @@ -1437,21 +1620,23 @@ func (s *subscriber) rebalance(ctx context.Context, sub *subscription, tenant st sort.Strings(sortedOwned) for _, pk := range sortedOwned[maxPart:] { - if err := s.leaseStore.ReleaseLease(ctx, tenant, sub.topic, pk, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { + rows, err := s.leaseStore.ReleaseLease(ctx, tenant, sub.topic, pk, cfg.SubscriberName, cfg.ConsumerGroup) + if err != nil { return released, fmt.Errorf("release partition %s during rebalance: %w", pk, err) } released = append(released, pk) - s.stopPartitionWorker(sub, entityqueue.PartitionIdentity{Tenant: tenant, PartitionKey: pk}) - s.logger.Infow("released partition for rebalance", - "tenant", tenant, - "topic", sub.topic, - "consumer_group", cfg.ConsumerGroup, - "partition_key", pk, - "owned", len(owned), - "max_partitions", maxPart, - ) + if rows > 0 { + s.recordLeaseReleased(sub, tenant, leaseReasonRebalance, rows) + s.logger.Infow("released partition for rebalance", + append(leaseLogFields(sub, tenant, pk), + logReason, leaseReasonRebalance, + "owned", len(owned), + "max_partitions", maxPart, + )..., + ) + } } return released, nil } diff --git a/platform/extension/messagequeue/mysql/subscriber_test.go b/platform/extension/messagequeue/mysql/subscriber_test.go index 3750d1cf6..cff0337bb 100644 --- a/platform/extension/messagequeue/mysql/subscriber_test.go +++ b/platform/extension/messagequeue/mysql/subscriber_test.go @@ -28,6 +28,7 @@ import ( "github.com/uber-go/tally" "go.uber.org/mock/gomock" "go.uber.org/zap" + "go.uber.org/zap/zapcore" "go.uber.org/zap/zaptest" "go.uber.org/zap/zaptest/observer" @@ -66,8 +67,8 @@ func allowSupervisorLeaseCalls(mockLeaseStore *MockpartitionLeaseStore) { mockLeaseStore.EXPECT().GetLeasedPartitionsForTenants(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(map[string][]string{}, nil).AnyTimes() mockLeaseStore.EXPECT().DiscoverPartitions(gomock.Any(), gomock.Any(), gomock.Any()).Return(map[string][]string{}, nil).AnyTimes() mockLeaseStore.EXPECT().GetAllLeasesForTenants(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(map[string][]leaseInfo{}, nil).AnyTimes() - mockLeaseStore.EXPECT().RenewOwnedLeases(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() - mockLeaseStore.EXPECT().ReleaseOwnedLeases(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + mockLeaseStore.EXPECT().RenewOwnedLeases(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() + mockLeaseStore.EXPECT().ReleaseOwnedLeases(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() mockLeaseStore.EXPECT().PurgeStaleForTenants(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() mockLeaseStore.EXPECT().TryAcquireLease(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(false, nil).AnyTimes() } @@ -1260,6 +1261,17 @@ func TestSubscriber_AcquireUnownedPartitions(t *testing.T) { }, want: nil, }, + { + name: "stale lease held by other is stolen", + discovered: []string{"part1"}, + leases: []leaseInfo{ + {PartitionKey: "part1", LeasedBy: "other-worker", LeaseRenewedAt: staleMs}, + }, + setup: func(mockLeaseStore *MockpartitionLeaseStore) { + expectAcquire(mockLeaseStore, "part1", true) + }, + want: []string{"part1"}, + }, } for _, tt := range tests { @@ -1301,7 +1313,7 @@ func TestSubscriber_LeaseTickRenewsWhenActiveSubscribersFails(t *testing.T) { Return(nil, activeErr) mockLeaseStore.EXPECT(). RenewOwnedLeases(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). - Return(nil) + Return(int64(0), nil) mockHeartbeatStore.EXPECT(). HeartbeatForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). Return(nil) @@ -1344,7 +1356,7 @@ func TestSubscriber_LeaseTickRenewsWhenGetLeasedPartitionsFails(t *testing.T) { Return(nil, leasedErr) mockLeaseStore.EXPECT(). RenewOwnedLeases(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). - Return(nil) + Return(int64(0), nil) mockHeartbeatStore.EXPECT(). HeartbeatForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). Return(nil) @@ -1553,7 +1565,7 @@ func TestSubscriber_DrainedPartitionKeepsOffsetWhenLeaseReleaseFails(t *testing. Return(map[string][]leaseInfo{}, nil) mockLeaseStore.EXPECT(). ReleaseLease(gomock.Any(), testTenant, "test-topic", partition.PartitionKey, cfg.SubscriberName, cfg.ConsumerGroup). - Return(errors.New("release failed")) + Return(int64(0), errors.New("release failed")) s := NewSubscriber( zaptest.NewLogger(t).Sugar(), @@ -1592,8 +1604,14 @@ func TestSubscriber_ReleaseAllLeasesContinuesAfterErrors(t *testing.T) { tenants := []string{"tenant-1", "tenant-2", "tenant-3"} mockLeaseStore.EXPECT(). - ReleaseOwnedLeases(gomock.Any(), tenants, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). - Return(releaseErr) + ReleaseOwnedLeases(gomock.Any(), []string{"tenant-1"}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(int64(0), releaseErr) + mockLeaseStore.EXPECT(). + ReleaseOwnedLeases(gomock.Any(), []string{"tenant-2"}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(int64(0), nil) + mockLeaseStore.EXPECT(). + ReleaseOwnedLeases(gomock.Any(), []string{"tenant-3"}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(int64(0), nil) s := NewSubscriber( zaptest.NewLogger(t).Sugar(), tally.NoopScope, @@ -1701,6 +1719,8 @@ func TestSubscriber_PartitionWorkerPollAndDeliver(t *testing.T) { select { case del := <-deliveryCh: assert.Equal(t, "msg-1", del.Message().ID) + assert.Equal(t, cfg.SubscriberName, del.Metadata()[logLeasedBy]) + assert.Equal(t, cfg.ConsumerGroup, del.Metadata()[logConsumerGroup]) default: t.Fatal("expected delivery but channel was empty") } @@ -2204,6 +2224,14 @@ func TestSubscriber_FairShareCap(t *testing.T) { }) } +func TestReportedFairShareCap(t *testing.T) { + assert.Equal(t, 2, reportedFairShareCap(2, []string{"a"}, []string{"b", "c"})) + assert.Equal(t, 2, reportedFairShareCap(0, []string{"a"}, []string{"a", "b"})) + owned, cap := 1, reportedFairShareCap(0, []string{"part-1"}, []string{"part-1"}) + assert.Equal(t, 1, cap) + assert.LessOrEqual(t, owned, cap) +} + // partitionKeysN generates n distinct partition keys. func partitionKeysN(n int) []string { keys := make([]string, n) @@ -2228,10 +2256,10 @@ func TestSubscriber_RebalanceReleasesExcess(t *testing.T) { mockLease := NewMockpartitionLeaseStore(ctrl) mockLease.EXPECT(). ReleaseLease(gomock.Any(), testTenant, "test-topic", "pk-c", "s1", "test-cg"). - Return(nil) + Return(int64(1), nil) mockLease.EXPECT(). ReleaseLease(gomock.Any(), testTenant, "test-topic", "pk-d", "s1", "test-cg"). - Return(nil) + Return(int64(1), nil) s := NewSubscriber( zaptest.NewLogger(t).Sugar(), tally.NoopScope, @@ -2342,3 +2370,598 @@ func TestUpdateDrainedTracking(t *testing.T) { }) } } + +func TestSubscriber_AcquireUnownedPartitions_LeaseEvents(t *testing.T) { + cfg := testSubscriptionConfig() + staleMs := time.Now().UnixMilli() - cfg.LeaseDurationMs - 60_000 + + tests := []struct { + name string + leases []leaseInfo + wantMsg string + wantLevel zapcore.Level + wantStolen bool + wantPrior string + }{ + { + name: "first claim", + wantMsg: "lease acquired", + wantLevel: zapcore.DebugLevel, + }, + { + name: "steal stale owner", + leases: []leaseInfo{ + {PartitionKey: "part1", LeasedBy: "other-worker", LeaseRenewedAt: staleMs}, + }, + wantMsg: "lease stolen", + wantLevel: zapcore.InfoLevel, + wantStolen: true, + wantPrior: "other-worker", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + mockLeaseStore.EXPECT(). + TryAcquireLease(gomock.Any(), testTenant, "test-topic", "part1", cfg.SubscriberName, cfg.ConsumerGroup, cfg.LeaseDurationMs). + Return(true, nil) + + core, logs := observer.New(zap.DebugLevel) + scope := tally.NewTestScope("test", nil) + s := NewSubscriber( + zap.New(core).Sugar(), + scope, + NewMockmessageStore(ctrl), + NewMockoffsetStore(ctrl), + mockLeaseStore, + NewMocksubscriberHeartbeatStore(ctrl), + NewMockdeliveryStateStore(ctrl), + []string{testTenant}, + ) + sub := &subscription{topic: "test-topic", config: cfg} + got := s.acquireUnownedPartitions(context.Background(), sub, testTenant, []string{"part1"}, tt.leases, 0) + assert.Equal(t, []string{"part1"}, got) + + entries := logs.FilterMessage(tt.wantMsg).All() + require.Len(t, entries, 1) + assert.Equal(t, tt.wantLevel, entries[0].Level) + assert.Equal(t, testTenant, entries[0].ContextMap()[logTenant]) + assert.Equal(t, "test-topic", entries[0].ContextMap()[logTopic]) + assert.Equal(t, cfg.ConsumerGroup, entries[0].ContextMap()[logConsumerGroup]) + assert.Equal(t, cfg.SubscriberName, entries[0].ContextMap()[logLeasedBy]) + assert.Equal(t, "part1", entries[0].ContextMap()[logPartitionKey]) + if tt.wantStolen { + assert.Equal(t, tt.wantPrior, entries[0].ContextMap()[logPreviousOwner]) + } + + snapshot := scope.Snapshot() + counter := "test.subscriber.lease.acquired" + if tt.wantStolen { + counter = "test.subscriber.lease.stolen" + } + assert.Equal(t, int64(1), testCounterValue(t, snapshot, counter, map[string]string{ + logLeasedBy: cfg.SubscriberName, + "topic": "test-topic", + logTenant: testTenant, + logConsumerGroup: cfg.ConsumerGroup, + })) + for _, c := range snapshot.Counters() { + assert.NotContains(t, c.Tags(), "partition_key") + } + }) + } +} + +func TestSubscriber_LeaseTick_OwnershipSummary(t *testing.T) { + ctrl := gomock.NewController(t) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + mockHeartbeatStore := NewMocksubscriberHeartbeatStore(ctrl) + cfg := testSubscriptionConfig() + + mockLeaseStore.EXPECT(). + GetLeasedPartitionsForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(map[string][]string{testTenant: {"part-1"}}, nil) + mockHeartbeatStore.EXPECT(). + ActiveSubscribersForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup, cfg.LeaseDurationMs). + Return(map[string][]string{testTenant: {cfg.SubscriberName, "other-subscriber"}}, nil) + mockLeaseStore.EXPECT(). + RenewOwnedLeases(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(int64(1), nil) + mockHeartbeatStore.EXPECT(). + HeartbeatForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(nil) + mockHeartbeatStore.EXPECT(). + PurgeStaleForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup, heartbeatPurgeAfterLeaseDurations*cfg.LeaseDurationMs). + Return(nil) + mockLeaseStore.EXPECT(). + PurgeStaleForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup, leasePurgeAfterLeaseDurations*cfg.LeaseDurationMs). + Return(nil) + + core, logs := observer.New(zap.InfoLevel) + scope := tally.NewTestScope("test", nil) + s := NewSubscriber( + zap.New(core).Sugar(), + scope, + NewMockmessageStore(ctrl), + NewMockoffsetStore(ctrl), + mockLeaseStore, + mockHeartbeatStore, + newTestDeliveryStateStore(ctrl), + []string{testTenant}, + ) + sub := &subscription{ + topic: "test-topic", + config: cfg, + deliveryCh: make(chan extqueue.Delivery, 1), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), + lastDiscoveredPartitions: tenantPartitionKeys(testTenant, []string{"part-1"}), + } + + s.runLeaseTick(context.Background(), sub, time.Second, []any{logTopic, sub.topic, logLeasedBy, cfg.SubscriberName}) + + entries := logs.FilterMessage("lease tick").All() + require.Len(t, entries, 1) + fields := entries[0].ContextMap() + assert.Equal(t, int64(1), fields[logOwnedPartitions]) + assert.Equal(t, int64(1), fields[logRenewed]) + assert.Equal(t, int64(2), fields["active_subscribers"]) + assert.Equal(t, int64(1), fields["fair_share_cap"]) + assert.Equal(t, cfg.SubscriberName, fields[logLeasedBy]) + + snapshot := scope.Snapshot() + assert.Equal(t, int64(1), testCounterValue(t, snapshot, "test.subscriber.lease.renewed", map[string]string{ + logLeasedBy: cfg.SubscriberName, + "topic": "test-topic", + logTenant: testTenant, + logConsumerGroup: cfg.ConsumerGroup, + })) + assert.Equal(t, float64(1), testGaugeValue(t, snapshot, "test.subscriber.lease.partitions_owned", map[string]string{ + logLeasedBy: cfg.SubscriberName, + "topic": "test-topic", + logTenant: testTenant, + logConsumerGroup: cfg.ConsumerGroup, + })) + assert.Equal(t, float64(2), testGaugeValue(t, snapshot, "test.subscriber.lease.active_subscribers", map[string]string{ + logLeasedBy: cfg.SubscriberName, + "topic": "test-topic", + logTenant: testTenant, + logConsumerGroup: cfg.ConsumerGroup, + })) + assert.Equal(t, float64(1), testGaugeValue(t, snapshot, "test.subscriber.lease.fair_share_cap", map[string]string{ + logLeasedBy: cfg.SubscriberName, + "topic": "test-topic", + logTenant: testTenant, + logConsumerGroup: cfg.ConsumerGroup, + })) +} + +func TestSubscriber_RebalanceReleasesExcess_RecordsReason(t *testing.T) { + ctrl := gomock.NewController(t) + mockLease := NewMockpartitionLeaseStore(ctrl) + mockLease.EXPECT(). + ReleaseLease(gomock.Any(), testTenant, "test-topic", "pk-c", "s1", "test-cg"). + Return(int64(1), nil) + mockLease.EXPECT(). + ReleaseLease(gomock.Any(), testTenant, "test-topic", "pk-d", "s1", "test-cg"). + Return(int64(1), nil) + + core, logs := observer.New(zap.InfoLevel) + scope := tally.NewTestScope("test", nil) + s := NewSubscriber( + zap.New(core).Sugar(), scope, + NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), + mockLease, NewMocksubscriberHeartbeatStore(ctrl), NewMockdeliveryStateStore(ctrl), + []string{testTenant}, + ) + sub := &subscription{ + topic: "test-topic", + config: extqueue.DefaultSubscriptionConfig("s1", "test-cg"), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), + } + + released, err := s.rebalance(context.Background(), sub, testTenant, []string{"pk-d", "pk-a", "pk-c", "pk-b"}, []string{"s1", "s2"}) + require.NoError(t, err) + assert.Equal(t, []string{"pk-c", "pk-d"}, released) + + entries := logs.FilterMessage("released partition for rebalance").All() + require.Len(t, entries, 2) + for _, entry := range entries { + assert.Equal(t, leaseReasonRebalance, entry.ContextMap()[logReason]) + assert.Equal(t, "s1", entry.ContextMap()[logLeasedBy]) + assert.Contains(t, []string{"pk-c", "pk-d"}, entry.ContextMap()[logPartitionKey]) + } + assert.Equal(t, int64(2), testCounterValue(t, scope.Snapshot(), "test.subscriber.lease.released", map[string]string{ + logLeasedBy: "s1", + "topic": "test-topic", + logTenant: testTenant, + logConsumerGroup: "test-cg", + logReason: leaseReasonRebalance, + })) +} + +func TestSubscriber_UncappedDiscoverDoesNotZeroShareGauges(t *testing.T) { + ctrl := gomock.NewController(t) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + cfg := testSubscriptionConfig() + cfg.PollIntervalMs = int64(time.Hour / time.Millisecond) + + mockLeaseStore.EXPECT(). + DiscoverPartitions(gomock.Any(), []string{testTenant}, "test-topic"). + Return(map[string][]string{testTenant: {"part-1"}}, nil) + mockLeaseStore.EXPECT(). + GetLeasedPartitionsForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(map[string][]string{testTenant: {"part-1"}}, nil) + mockLeaseStore.EXPECT(). + GetAllLeasesForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup). + Return(map[string][]leaseInfo{ + testTenant: {{PartitionKey: "part-1", LeasedBy: cfg.SubscriberName, LeaseRenewedAt: time.Now().UnixMilli()}}, + }, nil) + + scope := tally.NewTestScope("test", nil) + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), + scope, + NewMockmessageStore(ctrl), + NewMockoffsetStore(ctrl), + mockLeaseStore, + newTestHeartbeatStore(ctrl), + newTestDeliveryStateStore(ctrl), + []string{testTenant}, + ) + sub := &subscription{ + topic: "test-topic", + config: cfg, + deliveryCh: make(chan extqueue.Delivery, 1), + workers: map[entityqueue.PartitionIdentity]*partitionWorker{ + {Tenant: testTenant, PartitionKey: "part-1"}: { + cancelFunc: func() {}, + done: make(chan struct{}), + }, + }, + } + s.observeLeaseShare(sub, testTenant, 3, 2) + + require.NoError(t, s.discoverAndReconcileWorkers(context.Background(), sub, true)) + s.stopAllWorkers(sub) + + snapshot := scope.Snapshot() + assert.Equal(t, float64(1), testGaugeValue(t, snapshot, "test.subscriber.lease.partitions_owned", leaseMetricTags(cfg, testTenant))) + assert.Equal(t, float64(3), testGaugeValue(t, snapshot, "test.subscriber.lease.active_subscribers", leaseMetricTags(cfg, testTenant))) + assert.Equal(t, float64(2), testGaugeValue(t, snapshot, "test.subscriber.lease.fair_share_cap", leaseMetricTags(cfg, testTenant))) +} + +func TestSubscriber_LeaseTick_SingleSubscriberWritesUniverseCap(t *testing.T) { + ctrl := gomock.NewController(t) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + mockHeartbeatStore := NewMocksubscriberHeartbeatStore(ctrl) + cfg := testSubscriptionConfig() + + mockLeaseStore.EXPECT(). + GetLeasedPartitionsForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(map[string][]string{testTenant: {"part-1"}}, nil) + mockHeartbeatStore.EXPECT(). + ActiveSubscribersForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup, cfg.LeaseDurationMs). + Return(map[string][]string{testTenant: {cfg.SubscriberName}}, nil) + mockLeaseStore.EXPECT(). + RenewOwnedLeases(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(int64(1), nil) + mockHeartbeatStore.EXPECT(). + HeartbeatForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(nil) + mockHeartbeatStore.EXPECT(). + PurgeStaleForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup, heartbeatPurgeAfterLeaseDurations*cfg.LeaseDurationMs). + Return(nil) + mockLeaseStore.EXPECT(). + PurgeStaleForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup, leasePurgeAfterLeaseDurations*cfg.LeaseDurationMs). + Return(nil) + + scope := tally.NewTestScope("test", nil) + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), + scope, + NewMockmessageStore(ctrl), + NewMockoffsetStore(ctrl), + mockLeaseStore, + mockHeartbeatStore, + newTestDeliveryStateStore(ctrl), + []string{testTenant}, + ) + sub := &subscription{ + topic: "test-topic", + config: cfg, + deliveryCh: make(chan extqueue.Delivery, 1), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), + lastDiscoveredPartitions: tenantPartitionKeys(testTenant, []string{"part-1"}), + } + s.observeLeaseShare(sub, testTenant, 2, 4) + + s.runLeaseTick(context.Background(), sub, time.Second, []any{logTopic, sub.topic, logLeasedBy, cfg.SubscriberName}) + + snapshot := scope.Snapshot() + owned := testGaugeValue(t, snapshot, "test.subscriber.lease.partitions_owned", leaseMetricTags(cfg, testTenant)) + cap := testGaugeValue(t, snapshot, "test.subscriber.lease.fair_share_cap", leaseMetricTags(cfg, testTenant)) + assert.Equal(t, float64(1), cap) + assert.LessOrEqual(t, owned, cap) + assert.Equal(t, float64(1), testGaugeValue(t, snapshot, "test.subscriber.lease.active_subscribers", leaseMetricTags(cfg, testTenant))) +} + +func TestSubscriber_IdleReleaseRecordsTenant(t *testing.T) { + tests := []struct { + name string + releasedRows int64 + offsetErr error + wantReleased int64 + }{ + {name: "offset deleted", releasedRows: 1, wantReleased: 1}, + {name: "offset delete fails after release", releasedRows: 1, offsetErr: errors.New("delete offset failed"), wantReleased: 1}, + {name: "lease already gone", releasedRows: 0, wantReleased: 0}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + mockOffsetStore := NewMockoffsetStore(ctrl) + cfg := testSubscriptionConfig() + partition := entityqueue.PartitionIdentity{Tenant: testTenant, PartitionKey: "drained"} + + mockLeaseStore.EXPECT(). + DiscoverPartitions(gomock.Any(), []string{testTenant}, "test-topic"). + Return(map[string][]string{}, nil) + mockLeaseStore.EXPECT(). + GetLeasedPartitionsForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(map[string][]string{testTenant: {partition.PartitionKey}}, nil) + mockLeaseStore.EXPECT(). + GetAllLeasesForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup). + Return(map[string][]leaseInfo{}, nil) + mockLeaseStore.EXPECT(). + ReleaseLease(gomock.Any(), testTenant, "test-topic", partition.PartitionKey, cfg.SubscriberName, cfg.ConsumerGroup). + Return(tt.releasedRows, nil) + mockOffsetStore.EXPECT(). + DeleteOffset(gomock.Any(), testTenant, "test-topic", partition.PartitionKey, cfg.ConsumerGroup). + Return(tt.offsetErr) + + scope := tally.NewTestScope("test", nil) + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), + scope, + NewMockmessageStore(ctrl), + mockOffsetStore, + mockLeaseStore, + newTestHeartbeatStore(ctrl), + newTestDeliveryStateStore(ctrl), + []string{testTenant}, + ) + workerDone := make(chan struct{}) + close(workerDone) + sub := &subscription{ + topic: "test-topic", + config: cfg, + deliveryCh: make(chan extqueue.Delivery), + workers: map[entityqueue.PartitionIdentity]*partitionWorker{ + partition: {cancelFunc: func() {}, done: workerDone}, + }, + drainedSince: map[entityqueue.PartitionIdentity]time.Time{partition: time.Now().Add(-time.Hour)}, + } + + require.NoError(t, s.discoverAndReconcileWorkers(context.Background(), sub, true)) + snapshot := scope.Snapshot() + released, emitted := lookupCounter(snapshot, "test.subscriber.lease.released", map[string]string{ + logLeasedBy: cfg.SubscriberName, + "topic": "test-topic", + logTenant: testTenant, + logConsumerGroup: cfg.ConsumerGroup, + logReason: leaseReasonIdle, + }) + assert.Equal(t, tt.wantReleased > 0, emitted) + assert.Equal(t, tt.wantReleased, released) + assert.Equal(t, float64(0), testGaugeValue(t, snapshot, "test.subscriber.lease.partitions_owned", leaseMetricTags(cfg, testTenant))) + _, tracked := sub.drainedSince[partition] + assert.Equal(t, tt.offsetErr != nil, tracked) + }) + } +} + +func TestSubscriber_ReleaseAllLeases_PerTenant(t *testing.T) { + ctrl := gomock.NewController(t) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + cfg := testSubscriptionConfig() + tenants := []string{"tenant-a", "tenant-b"} + + mockLeaseStore.EXPECT(). + ReleaseOwnedLeases(gomock.Any(), []string{"tenant-a"}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(int64(2), nil) + mockLeaseStore.EXPECT(). + ReleaseOwnedLeases(gomock.Any(), []string{"tenant-b"}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(int64(1), nil) + + scope := tally.NewTestScope("test", nil) + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), scope, + NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), + mockLeaseStore, NewMocksubscriberHeartbeatStore(ctrl), + NewMockdeliveryStateStore(ctrl), + tenants, + ) + sub := &subscription{topic: "test-topic", config: cfg} + + require.NoError(t, s.releaseAllLeases(context.Background(), sub)) + snapshot := scope.Snapshot() + assert.Equal(t, int64(2), testCounterValue(t, snapshot, "test.subscriber.lease.released", map[string]string{ + logLeasedBy: cfg.SubscriberName, + "topic": "test-topic", + logTenant: "tenant-a", + logConsumerGroup: cfg.ConsumerGroup, + logReason: leaseReasonShutdown, + })) + assert.Equal(t, int64(1), testCounterValue(t, snapshot, "test.subscriber.lease.released", map[string]string{ + logLeasedBy: cfg.SubscriberName, + "topic": "test-topic", + logTenant: "tenant-b", + logConsumerGroup: cfg.ConsumerGroup, + logReason: leaseReasonShutdown, + })) +} + +func TestSubscriber_LeaseTick_LogsOnRenewError(t *testing.T) { + ctrl := gomock.NewController(t) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + mockHeartbeatStore := NewMocksubscriberHeartbeatStore(ctrl) + cfg := testSubscriptionConfig() + renewErr := errors.New("renew failed") + + mockLeaseStore.EXPECT(). + GetLeasedPartitionsForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(map[string][]string{testTenant: {"part-1"}}, nil) + mockHeartbeatStore.EXPECT(). + ActiveSubscribersForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup, cfg.LeaseDurationMs). + Return(map[string][]string{testTenant: {cfg.SubscriberName, "other"}}, nil) + mockLeaseStore.EXPECT(). + RenewOwnedLeases(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(int64(0), renewErr) + mockHeartbeatStore.EXPECT(). + HeartbeatForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(nil) + mockHeartbeatStore.EXPECT(). + PurgeStaleForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup, heartbeatPurgeAfterLeaseDurations*cfg.LeaseDurationMs). + Return(nil) + mockLeaseStore.EXPECT(). + PurgeStaleForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup, leasePurgeAfterLeaseDurations*cfg.LeaseDurationMs). + Return(nil) + + core, logs := observer.New(zap.InfoLevel) + s := NewSubscriber( + zap.New(core).Sugar(), + tally.NoopScope, + NewMockmessageStore(ctrl), + NewMockoffsetStore(ctrl), + mockLeaseStore, + mockHeartbeatStore, + newTestDeliveryStateStore(ctrl), + []string{testTenant}, + ) + sub := &subscription{ + topic: "test-topic", + config: cfg, + deliveryCh: make(chan extqueue.Delivery, 1), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), + lastDiscoveredPartitions: tenantPartitionKeys(testTenant, []string{"part-1"}), + } + + s.runLeaseTick(context.Background(), sub, time.Second, []any{logTopic, sub.topic, logLeasedBy, cfg.SubscriberName}) + + entries := logs.FilterMessage("lease tick").All() + require.Len(t, entries, 1) + assert.Equal(t, int64(1), entries[0].ContextMap()[logOwnedPartitions]) + assert.Equal(t, int64(2), entries[0].ContextMap()["active_subscribers"]) +} + +func TestSubscriber_LeaseTick_RenewMismatchOmitsTenant(t *testing.T) { + ctrl := gomock.NewController(t) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + mockHeartbeatStore := NewMocksubscriberHeartbeatStore(ctrl) + cfg := testSubscriptionConfig() + + mockLeaseStore.EXPECT(). + GetLeasedPartitionsForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(map[string][]string{testTenant: {"part-1"}}, nil) + mockHeartbeatStore.EXPECT(). + ActiveSubscribersForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup, cfg.LeaseDurationMs). + Return(map[string][]string{testTenant: {cfg.SubscriberName, "other"}}, nil) + mockLeaseStore.EXPECT(). + RenewOwnedLeases(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(int64(2), nil) + mockHeartbeatStore.EXPECT(). + HeartbeatForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(nil) + mockHeartbeatStore.EXPECT(). + PurgeStaleForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup, heartbeatPurgeAfterLeaseDurations*cfg.LeaseDurationMs). + Return(nil) + mockLeaseStore.EXPECT(). + PurgeStaleForTenants(gomock.Any(), []string{testTenant}, "test-topic", cfg.ConsumerGroup, leasePurgeAfterLeaseDurations*cfg.LeaseDurationMs). + Return(nil) + + scope := tally.NewTestScope("test", nil) + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), + scope, + NewMockmessageStore(ctrl), + NewMockoffsetStore(ctrl), + mockLeaseStore, + mockHeartbeatStore, + newTestDeliveryStateStore(ctrl), + []string{testTenant}, + ) + sub := &subscription{ + topic: "test-topic", + config: cfg, + deliveryCh: make(chan extqueue.Delivery, 1), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), + lastDiscoveredPartitions: tenantPartitionKeys(testTenant, []string{"part-1"}), + } + + s.runLeaseTick(context.Background(), sub, time.Second, []any{logTopic, sub.topic, logLeasedBy, cfg.SubscriberName}) + + snapshot := scope.Snapshot() + assert.Equal(t, int64(2), testCounterValue(t, snapshot, "test.subscriber.lease.renewed", map[string]string{ + logLeasedBy: cfg.SubscriberName, + "topic": "test-topic", + logConsumerGroup: cfg.ConsumerGroup, + })) + for _, c := range snapshot.Counters() { + if c.Name() == "test.subscriber.lease.renewed" { + assert.Empty(t, c.Tags()[logTenant]) + } + } +} + +func leaseMetricTags(cfg extqueue.SubscriptionConfig, tenant string) map[string]string { + return map[string]string{ + logLeasedBy: cfg.SubscriberName, + "topic": "test-topic", + logTenant: tenant, + logConsumerGroup: cfg.ConsumerGroup, + } +} + +func testCounterValue(t *testing.T, snapshot tally.Snapshot, name string, tags map[string]string) int64 { + t.Helper() + value, found := lookupCounter(snapshot, name, tags) + if !found { + t.Fatalf("counter %s %+v not found", name, tags) + } + return value +} + +func lookupCounter(snapshot tally.Snapshot, name string, tags map[string]string) (int64, bool) { + for _, c := range snapshot.Counters() { + if c.Name() != name || !tagsMatch(c.Tags(), tags) { + continue + } + return c.Value(), true + } + return 0, false +} + +func testGaugeValue(t *testing.T, snapshot tally.Snapshot, name string, tags map[string]string) float64 { + t.Helper() + for _, g := range snapshot.Gauges() { + if g.Name() != name || !tagsMatch(g.Tags(), tags) { + continue + } + return g.Value() + } + t.Fatalf("gauge %s %+v not found", name, tags) + return 0 +} + +func tagsMatch(got, want map[string]string) bool { + for k, v := range want { + if got[k] != v { + return false + } + } + return true +} diff --git a/service/runway/server/main_test.go b/service/runway/server/main_test.go index 2379662d1..e5804c9dd 100644 --- a/service/runway/server/main_test.go +++ b/service/runway/server/main_test.go @@ -130,6 +130,7 @@ func TestPrimaryConsumer_GitFailureDisposition(t *testing.T) { delivery := queuemock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(message).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() + delivery.EXPECT().Metadata().Return(nil).AnyTimes() done := make(chan struct{}) if tt.wantOutcome == "nack" { delivery.EXPECT().Nack(gomock.Any(), gomock.Any()).DoAndReturn(func(context.Context, failure.Failure) error {