diff --git a/platform/extension/messagequeue/mysql/mock_stores.go b/platform/extension/messagequeue/mysql/mock_stores.go index f8d7d105..28c446bf 100644 --- a/platform/extension/messagequeue/mysql/mock_stores.go +++ b/platform/extension/messagequeue/mysql/mock_stores.go @@ -251,6 +251,21 @@ func (mr *MockpartitionLeaseStoreMockRecorder) DiscoverAndAcquirePartitions(ctx, return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscoverAndAcquirePartitions", reflect.TypeOf((*MockpartitionLeaseStore)(nil).DiscoverAndAcquirePartitions), ctx, topic, subscriberName, consumerGroup, leaseDurationMs, maxPartitions) } +// GetAllLeases mocks base method. +func (m *MockpartitionLeaseStore) GetAllLeases(ctx context.Context, topic, consumerGroup string) ([]leaseInfo, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAllLeases", ctx, topic, consumerGroup) + ret0, _ := ret[0].([]leaseInfo) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAllLeases indicates an expected call of GetAllLeases. +func (mr *MockpartitionLeaseStoreMockRecorder) GetAllLeases(ctx, topic, consumerGroup any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAllLeases", reflect.TypeOf((*MockpartitionLeaseStore)(nil).GetAllLeases), ctx, topic, consumerGroup) +} + // GetLeasedPartitions mocks base method. func (m *MockpartitionLeaseStore) GetLeasedPartitions(ctx context.Context, topic, subscriberName, consumerGroup string) ([]string, error) { m.ctrl.T.Helper() diff --git a/platform/extension/messagequeue/mysql/partition_lease_store.go b/platform/extension/messagequeue/mysql/partition_lease_store.go index 6e2fb227..8d8287c6 100644 --- a/platform/extension/messagequeue/mysql/partition_lease_store.go +++ b/platform/extension/messagequeue/mysql/partition_lease_store.go @@ -196,9 +196,51 @@ func (s *sqlpartitionLeaseStore) GetLeasedPartitions(ctx context.Context, topic return partitions, nil } +// GetAllLeases returns the lease row for every partition currently leased +// under (topic, consumerGroup) by any subscriber. +func (s *sqlpartitionLeaseStore) GetAllLeases(ctx context.Context, topic string, consumerGroup string) (_ []leaseInfo, retErr error) { + op := metrics.Begin(s.scope, "get_all_leases", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) + defer func() { op.Complete(retErr) }() + + rows, err := s.db.QueryContext(ctx, fmt.Sprintf(` + SELECT partition_key, leased_by, lease_renewed_at FROM %s + WHERE consumer_group = ? AND topic = ? + `, PartitionLeasesTableName), consumerGroup, topic) + + if err != nil { + return nil, fmt.Errorf("get all leases topic=%s: %w", topic, err) + } + defer rows.Close() + + var leases []leaseInfo + for rows.Next() { + var lease leaseInfo + if err := rows.Scan(&lease.PartitionKey, &lease.LeasedBy, &lease.LeaseRenewedAt); err != nil { + return nil, fmt.Errorf("scan lease topic=%s: %w", topic, err) + } + leases = append(leases, lease) + } + + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("row iteration topic=%s: %w", topic, err) + } + + return leases, nil +} + // DiscoverAndAcquirePartitions discovers partitions from messages table and tries to acquire leases. // Returns the number of new leases acquired and the full list of discovered partitions. // maxPartitions limits how many total partitions this subscriber can own (0 = unlimited) +// +// Acquisition is lease-aware: one GetAllLeases read classifies every +// discovered partition, and TryAcquireLease is attempted only for partitions +// that are unleased or whose lease is stale (stealable). Partitions already +// owned by this subscriber are counted against the cap but not re-probed +// (renewal is the lease tick's job), and partitions validly held by another +// subscriber are skipped entirely — probing them is a guaranteed-futile +// write on a contended lease row. The classification is advisory (a lease +// can expire or renew between the read and the attempt); TryAcquireLease +// remains the atomic arbiter. func (s *sqlpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Context, topic string, subscriberName string, consumerGroup string, leaseDurationMs int64, maxPartitions int) (_ int, _ []string, retErr error) { op := metrics.Begin(s.scope, "discover_and_acquire", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() @@ -234,28 +276,44 @@ func (s *sqlpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Contex "count", len(partitions), ) - // Query owned partitions once before the loop to avoid N+1 queries. - // Build a set of already-owned partition keys so we can distinguish - // re-acquiring an already-owned partition from acquiring a new one. + // One read of every lease row classifies the discovered partitions: + // self-owned (count toward the cap, no re-probe), validly held by + // another subscriber (skip), or unleased/stale (acquisition candidates). + allLeases, err := s.GetAllLeases(ctx, topic, consumerGroup) + if err != nil { + return 0, nil, fmt.Errorf("get all leases for acquisition topic=%s: %w", topic, err) + } + staleThreshold := currentTimeMillis() - leaseDurationMs ownedCount := 0 ownedSet := make(map[string]struct{}) - if maxPartitions > 0 { - owned, err := s.GetLeasedPartitions(ctx, topic, subscriberName, consumerGroup) - if err != nil { - return 0, nil, fmt.Errorf("get owned partitions for cap check topic=%s: %w", topic, err) - } - ownedCount = len(owned) - for _, pk := range owned { - ownedSet[pk] = struct{}{} + heldByOther := make(map[string]struct{}) + for _, lease := range allLeases { + switch { + case lease.LeasedBy == subscriberName: + // Self-owned, fresh or stale: a stale self-lease means our own + // renewals are lagging, not that ownership moved. + ownedSet[lease.PartitionKey] = struct{}{} + ownedCount++ + case lease.LeaseRenewedAt >= staleThreshold: + heldByOther[lease.PartitionKey] = struct{}{} } } // Sort partitions deterministically sort.Strings(partitions) - // Try to acquire leases for discovered partitions + // Try to acquire leases for unleased or stale discovered partitions acquiredCount := 0 + skippedCount := 0 for _, partitionKey := range partitions { + if _, owned := ownedSet[partitionKey]; owned { + continue + } + if _, held := heldByOther[partitionKey]; held { + skippedCount++ + continue + } + // Enforce maxPartitions cap using local count if maxPartitions > 0 && ownedCount >= maxPartitions { s.logger.Debugw("reached max partitions cap, stopping acquisition", @@ -279,22 +337,19 @@ func (s *sqlpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Contex continue } if acquired { - // Only count as newly acquired if not already owned. - // TryAcquireLease returns true for already-owned partitions (renew), - // so we must not double-count them against the maxPartitions cap. - if _, alreadyOwned := ownedSet[partitionKey]; !alreadyOwned { - acquiredCount++ - ownedCount++ - } + acquiredCount++ + ownedCount++ } } metrics.NamedCounter(s.scope, "discover_and_acquire", "partitions_discovered", int64(len(partitions)), metrics.NewTag("topic", topic)) metrics.NamedCounter(s.scope, "discover_and_acquire", "partitions_acquired", int64(acquiredCount), metrics.NewTag("topic", topic)) + metrics.NamedCounter(s.scope, "discover_and_acquire", "lease_aware_skipped", int64(skippedCount), metrics.NewTag("topic", topic)) s.logger.Debugw("completed partition discovery and acquisition", logTopic, topic, "discovered_count", len(partitions), "acquired_count", acquiredCount, + "skipped_held_by_other", skippedCount, ) return acquiredCount, partitions, nil diff --git a/platform/extension/messagequeue/mysql/partition_lease_store_test.go b/platform/extension/messagequeue/mysql/partition_lease_store_test.go index 6d405833..386fc094 100644 --- a/platform/extension/messagequeue/mysql/partition_lease_store_test.go +++ b/platform/extension/messagequeue/mysql/partition_lease_store_test.go @@ -18,6 +18,7 @@ import ( "context" "database/sql" "testing" + "time" "github.com/DATA-DOG/go-sqlmock" "github.com/stretchr/testify/require" @@ -216,37 +217,124 @@ func TestPartitionLeaseStore_GetLeasedPartitions(t *testing.T) { require.NoError(t, mock.ExpectationsWereMet()) } +func TestPartitionLeaseStore_GetAllLeases(t *testing.T) { + tests := []struct { + name string + setup func(mock sqlmock.Sqlmock) + want []leaseInfo + }{ + { + name: "returns leases held by any subscriber", + setup: func(mock sqlmock.Sqlmock) { + rows := sqlmock.NewRows([]string{"partition_key", "leased_by", "lease_renewed_at"}). + AddRow("part1", testSubscriberName, int64(1000)). + AddRow("part2", "other-worker", int64(2000)) + mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic"). + WillReturnRows(rows) + }, + want: []leaseInfo{ + {PartitionKey: "part1", LeasedBy: testSubscriberName, LeaseRenewedAt: 1000}, + {PartitionKey: "part2", LeasedBy: "other-worker", LeaseRenewedAt: 2000}, + }, + }, + { + name: "no leases returns empty", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic"). + WillReturnRows(sqlmock.NewRows([]string{"partition_key", "leased_by", "lease_renewed_at"})) + }, + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db, mock, store := setuppartitionLeaseStoreTest(t) + defer db.Close() + + tt.setup(mock) + + leases, err := store.GetAllLeases(context.Background(), "test_topic", testConsumerGroup) + require.NoError(t, err) + require.Equal(t, tt.want, leases) + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { + leaseColumns := []string{"partition_key", "leased_by", "lease_renewed_at"} + freshMs := time.Now().UnixMilli() + staleMs := freshMs - testLeaseDurationMs - 60_000 + + // expectDiscover mocks the DISTINCT partition scan. + expectDiscover := func(mock sqlmock.Sqlmock, partitions ...string) { + rows := sqlmock.NewRows([]string{"partition_key"}) + for _, pk := range partitions { + rows.AddRow(pk) + } + mock.ExpectQuery("SELECT DISTINCT partition_key FROM queue_messages"). + WithArgs("test_topic"). + WillReturnRows(rows) + } + + // expectAcquire mocks one TryAcquireLease attempt whose ownership check + // reports the given owner. + expectAcquire := func(mock sqlmock.Sqlmock, owner string) { + mock.ExpectExec("INSERT INTO queue_partition_leases"). + WillReturnResult(sqlmock.NewResult(1, 1)) + mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). + WillReturnRows(sqlmock.NewRows([]string{"leased_by"}).AddRow(owner)) + } + tests := []struct { name string maxPartitions int setup func(mock sqlmock.Sqlmock) wantAcquired int - wantErr bool }{ { - name: "unlimited - acquires all available", + name: "acquires unleased, skips fresh lease held by other", maxPartitions: 0, setup: func(mock sqlmock.Sqlmock) { - // Discover partitions - rows := sqlmock.NewRows([]string{"partition_key"}). - AddRow("part1"). - AddRow("part2") - mock.ExpectQuery("SELECT DISTINCT partition_key FROM queue_messages"). - WithArgs("test_topic"). - WillReturnRows(rows) - - // Acquire part1 - success - mock.ExpectExec("INSERT INTO queue_partition_leases"). - WillReturnResult(sqlmock.NewResult(1, 1)) - mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). - WillReturnRows(sqlmock.NewRows([]string{"leased_by"}).AddRow(testSubscriberName)) - - // Acquire part2 - taken by other worker - mock.ExpectExec("INSERT INTO queue_partition_leases"). - WillReturnResult(sqlmock.NewResult(1, 1)) - mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). - WillReturnRows(sqlmock.NewRows([]string{"leased_by"}).AddRow("other-worker")) + expectDiscover(mock, "part1", "part2") + mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic"). + WillReturnRows(sqlmock.NewRows(leaseColumns). + AddRow("part2", "other-worker", freshMs)) + // Only unleased part1 is attempted; part2's fresh lease is + // never write-probed. + expectAcquire(mock, testSubscriberName) + }, + wantAcquired: 1, + }, + { + name: "stale lease held by other is stealable", + maxPartitions: 0, + setup: func(mock sqlmock.Sqlmock) { + expectDiscover(mock, "part1") + mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic"). + WillReturnRows(sqlmock.NewRows(leaseColumns). + AddRow("part1", "other-worker", staleMs)) + expectAcquire(mock, testSubscriberName) + }, + wantAcquired: 1, + }, + { + name: "self-owned partitions are not re-probed", + maxPartitions: 0, + setup: func(mock sqlmock.Sqlmock) { + expectDiscover(mock, "part1", "part2") + mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic"). + WillReturnRows(sqlmock.NewRows(leaseColumns). + AddRow("part1", testSubscriberName, freshMs)) + // Only part2 is attempted; renewal of part1 is the lease + // tick's job. + expectAcquire(mock, testSubscriberName) }, wantAcquired: 1, }, @@ -254,33 +342,13 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { name: "stops acquiring when cap reached", maxPartitions: 2, setup: func(mock sqlmock.Sqlmock) { - // Discover 3 partitions - rows := sqlmock.NewRows([]string{"partition_key"}). - AddRow("part1"). - AddRow("part2"). - AddRow("part3") - mock.ExpectQuery("SELECT DISTINCT partition_key FROM queue_messages"). - WithArgs("test_topic"). - WillReturnRows(rows) - - // Pre-loop GetLeasedPartitions: owns 0 partitions - mock.ExpectQuery("SELECT partition_key FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName). - WillReturnRows(sqlmock.NewRows([]string{"partition_key"})) - - // Acquire part1 - success - mock.ExpectExec("INSERT INTO queue_partition_leases"). - WillReturnResult(sqlmock.NewResult(1, 1)) - mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). - WillReturnRows(sqlmock.NewRows([]string{"leased_by"}).AddRow(testSubscriberName)) - - // Acquire part2 - success (now at cap of 2, stops) - mock.ExpectExec("INSERT INTO queue_partition_leases"). - WillReturnResult(sqlmock.NewResult(1, 1)) - mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). - WillReturnRows(sqlmock.NewRows([]string{"leased_by"}).AddRow(testSubscriberName)) - - // part3 is never attempted because ownedCount (2) >= maxPartitions (2) + expectDiscover(mock, "part1", "part2", "part3") + mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic"). + WillReturnRows(sqlmock.NewRows(leaseColumns)) + // part1 and part2 acquired; part3 never attempted at the cap. + expectAcquire(mock, testSubscriberName) + expectAcquire(mock, testSubscriberName) }, wantAcquired: 2, }, @@ -288,52 +356,42 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { name: "pre-owned partitions count toward cap", maxPartitions: 3, setup: func(mock sqlmock.Sqlmock) { - // Discover 3 partitions - rows := sqlmock.NewRows([]string{"partition_key"}). - AddRow("part1"). - AddRow("part2"). - AddRow("part3") - mock.ExpectQuery("SELECT DISTINCT partition_key FROM queue_messages"). - WithArgs("test_topic"). - WillReturnRows(rows) - - // Pre-loop GetLeasedPartitions: already owns 2 partitions - mock.ExpectQuery("SELECT partition_key FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName). - WillReturnRows(sqlmock.NewRows([]string{"partition_key"}). - AddRow("existing1"). - AddRow("existing2")) - - // Acquire part1 - success (now at 3, cap reached) - mock.ExpectExec("INSERT INTO queue_partition_leases"). - WillReturnResult(sqlmock.NewResult(1, 1)) - mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). - WillReturnRows(sqlmock.NewRows([]string{"leased_by"}).AddRow(testSubscriberName)) - - // part2, part3 never attempted because ownedCount (3) >= maxPartitions (3) + expectDiscover(mock, "part1", "part2", "part3") + mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic"). + WillReturnRows(sqlmock.NewRows(leaseColumns). + AddRow("existing1", testSubscriberName, freshMs). + AddRow("existing2", testSubscriberName, freshMs)) + // One acquisition reaches the cap of 3; part2/part3 skipped. + expectAcquire(mock, testSubscriberName) }, wantAcquired: 1, }, { - name: "already at cap - acquires nothing", + name: "already at cap acquires nothing", maxPartitions: 2, setup: func(mock sqlmock.Sqlmock) { - // Discover 2 partitions - rows := sqlmock.NewRows([]string{"partition_key"}). - AddRow("part1"). - AddRow("part2") - mock.ExpectQuery("SELECT DISTINCT partition_key FROM queue_messages"). - WithArgs("test_topic"). - WillReturnRows(rows) - - // Pre-loop GetLeasedPartitions: already owns 2 partitions (at cap) - mock.ExpectQuery("SELECT partition_key FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName). - WillReturnRows(sqlmock.NewRows([]string{"partition_key"}). - AddRow("existing1"). - AddRow("existing2")) - - // No acquire attempts - immediately breaks + expectDiscover(mock, "part1", "part2") + mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic"). + WillReturnRows(sqlmock.NewRows(leaseColumns). + AddRow("existing1", testSubscriberName, freshMs). + AddRow("existing2", testSubscriberName, freshMs)) + // No acquire attempts. + }, + wantAcquired: 0, + }, + { + name: "lost race counts nothing", + maxPartitions: 0, + setup: func(mock sqlmock.Sqlmock) { + expectDiscover(mock, "part1") + mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic"). + WillReturnRows(sqlmock.NewRows(leaseColumns)) + // Attempted while unleased, but another subscriber won the + // atomic acquire between the read and the write. + expectAcquire(mock, "other-worker") }, wantAcquired: 0, }, @@ -344,19 +402,12 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { db, mock, store := setuppartitionLeaseStoreTest(t) defer db.Close() - ctx := context.Background() - topic := "test_topic" - tt.setup(mock) - acquired, discoveredPartitions, err := store.DiscoverAndAcquirePartitions(ctx, topic, testSubscriberName, testConsumerGroup, testLeaseDurationMs, tt.maxPartitions) - if tt.wantErr { - require.Error(t, err) - } else { - require.NoError(t, err) - require.Equal(t, tt.wantAcquired, acquired) - require.NotNil(t, discoveredPartitions) - } + acquired, discoveredPartitions, err := store.DiscoverAndAcquirePartitions(context.Background(), "test_topic", testSubscriberName, testConsumerGroup, testLeaseDurationMs, tt.maxPartitions) + require.NoError(t, err) + require.Equal(t, tt.wantAcquired, acquired) + require.NotNil(t, discoveredPartitions) require.NoError(t, mock.ExpectationsWereMet()) }) } diff --git a/platform/extension/messagequeue/mysql/stores.go b/platform/extension/messagequeue/mysql/stores.go index 8c39070f..9b28c26d 100644 --- a/platform/extension/messagequeue/mysql/stores.go +++ b/platform/extension/messagequeue/mysql/stores.go @@ -103,6 +103,17 @@ type offsetStore interface { GetMinAckedOffset(ctx context.Context, topic string, partitionKey string) (offset int64, found bool, err error) } +// leaseInfo describes one partition's current lease row (internal use only) +type leaseInfo struct { + // PartitionKey is the partition this lease covers + PartitionKey string + // LeasedBy is the subscriber name currently holding the lease + LeasedBy string + // LeaseRenewedAt is the epoch milliseconds of the last renewal; a lease + // is stale (stealable) once this is older than the lease duration + LeaseRenewedAt int64 +} + // partitionLeaseStore handles partition lease operations (internal use only) type partitionLeaseStore interface { // TryAcquireLease attempts to acquire or renew a lease for a partition @@ -120,6 +131,12 @@ type partitionLeaseStore interface { // GetLeasedPartitions returns all partitions currently leased by this worker GetLeasedPartitions(ctx context.Context, topic string, subscriberName string, consumerGroup string) ([]string, error) + // GetAllLeases returns the lease row for every partition currently leased + // under (topic, consumerGroup) by any subscriber. One PK-prefix read that + // lets acquisition skip partitions validly held by other subscribers + // instead of write-probing every lease row each discovery tick. + GetAllLeases(ctx context.Context, topic string, consumerGroup string) ([]leaseInfo, error) + // DiscoverAndAcquirePartitions discovers partitions from messages table and tries to acquire leases. // Returns the number of new leases acquired and the full list of discovered partitions. // leaseDurationMs is how long the lease is valid (in milliseconds)