refactor(queue/mysql): immutable log, delivery state, heartbeat store, fair leasing

- Make message log immutable (no per-message mutation during delivery)
- Add delivery_state_store with explicit acked boolean (per-consumer-group)
- Add subscriber_heartbeat_store for fair partition distribution
- Refactor partition_lease_store with consumer group scoping and discovery
- Simplify offset_store (remove AckMessage, use metrics.Begin pattern)
- Rewrite subscriber with per-partition workers, watermark advancement,
  non-blocking nack, and comprehensive observability (Warnw + metrics
  for all swallowed errors)
- Each store queries only its own table (no cross-table JOINs)
- Logger hierarchy: sql.go passes parent logger, each component names itself

Author: behinddwalls

extension/queue/mysql/BUILD.bazel

@@ -17,6 +17,7 @@ go_library(
     name = "mysql",
     srcs = [
         "constants.go",
+        "delivery_state_store.go",
         "errors.go",
         "message_store.go",
         ":mock_stores_src",
@@ -26,10 +27,12 @@ go_library(
         "sql.go",
         "stores.go",
         "subscriber.go",
+        "subscriber_heartbeat_store.go",
     ],
     importpath = "github.com/uber/submitqueue/extension/queue/mysql",
     visibility = ["//visibility:public"],
     deps = [
+        "//core/metrics",
         "//entity/queue",
         "//extension/queue",
         "@com_github_uber_go_tally_v4//:tally",
@@ -41,11 +44,13 @@ go_library(
 go_test(
     name = "mysql_test",
     srcs = [
+        "delivery_state_store_test.go",
         "message_store_test.go",
         "offset_store_test.go",
         "partition_lease_store_test.go",
         "publisher_test.go",
         "sql_test.go",
+        "subscriber_heartbeat_store_test.go",
         "subscriber_test.go",
     ],
     embed = [":mysql"],

extension/queue/mysql/constants.go

@@ -17,16 +17,9 @@ package mysql
 // Common constants for frequently repeated strings across stores
 
 const (
-	// Tag key (used in every Tagged() call)
-	tagErrorType = "error_type"
-
 	// Common log field names (used extensively across all stores)
 	logTopic        = "topic"
 	logPartitionKey = "partition_key"
 	logMessageID    = "message_id"
 	logError        = "error"
-
-	// Error types used across multiple methods/stores
-	errorBeginTx = "begin_transaction"
-	errorCommit  = "commit"
 )

extension/queue/mysql/delivery_state_store.go

@@ -0,0 +1,295 @@
+// Copyright (c) 2025 Uber Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package mysql
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"time"
+
+	"github.com/uber-go/tally/v4"
+	"github.com/uber/submitqueue/core/metrics"
+	"go.uber.org/zap"
+)
+
+// sqldeliveryStateStore is the SQL implementation of deliveryStateStore
+type sqldeliveryStateStore struct {
+	db     *sql.DB
+	logger *zap.SugaredLogger
+	scope  tally.Scope
+}
+
+// newDeliveryStateStore creates a new SQL delivery state store
+func newDeliveryStateStore(db *sql.DB, logger *zap.SugaredLogger, scope tally.Scope) deliveryStateStore {
+	return &sqldeliveryStateStore{
+		db:     db,
+		logger: logger.Named("delivery_state_store"),
+		scope:  scope.SubScope("delivery_state_store"),
+	}
+}
+
+// MarkDelivered inserts a row marking message as in-flight for this consumer group.
+func (s *sqldeliveryStateStore) MarkDelivered(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64, visibilityTimeoutMs int64) (retErr error) {
+	op := metrics.Begin(s.scope, "mark_delivered",
+		metrics.NewTag("topic", topic),
+		metrics.NewTag("consumer_group", consumerGroup),
+		metrics.NewTag("partition_key", partitionKey))
+	defer func() { op.Complete(retErr) }()
+
+	now := time.Now().UnixMilli()
+	invisibleUntil := now + visibilityTimeoutMs
+
+	_, err := s.db.ExecContext(ctx, fmt.Sprintf(`
+		INSERT INTO %s (consumer_group, topic, partition_key, message_offset, acked, invisible_until, retry_count)
+		VALUES (?, ?, ?, ?, FALSE, ?, 0)
+		ON DUPLICATE KEY UPDATE
+			invisible_until = IF(acked = FALSE, VALUES(invisible_until), invisible_until),
+			retry_count = IF(acked = FALSE, retry_count + 1, retry_count)
+	`, DeliveryStateTableName),
+		consumerGroup, topic, partitionKey, offset, invisibleUntil)
+
+	if err != nil {
+		return fmt.Errorf("mark delivered topic=%s partition=%s offset=%d: %w", topic, partitionKey, offset, err)
+	}
+
+	return nil
+}
+
+// ExtendVisibility extends the visibility timeout for an in-flight message
+// without incrementing retry_count. Used by ExtendVisibilityTimeout.
+func (s *sqldeliveryStateStore) ExtendVisibility(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64, visibilityTimeoutMs int64) (retErr error) {
+	op := metrics.Begin(s.scope, "extend_visibility",
+		metrics.NewTag("topic", topic),
+		metrics.NewTag("consumer_group", consumerGroup),
+		metrics.NewTag("partition_key", partitionKey))
+	defer func() { op.Complete(retErr) }()
+
+	now := time.Now().UnixMilli()
+	invisibleUntil := now + visibilityTimeoutMs
+
+	_, err := s.db.ExecContext(ctx, fmt.Sprintf(`
+		UPDATE %s
+		SET invisible_until = ?
+		WHERE consumer_group = ? AND topic = ? AND partition_key = ? AND message_offset = ? AND acked = FALSE
+	`, DeliveryStateTableName),
+		invisibleUntil, consumerGroup, topic, partitionKey, offset)
+
+	if err != nil {
+		return fmt.Errorf("extend visibility topic=%s partition=%s offset=%d: %w", topic, partitionKey, offset, err)
+	}
+
+	return nil
+}
+
+// MarkAcked sets acked = TRUE to indicate this group has processed the message.
+func (s *sqldeliveryStateStore) MarkAcked(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64) (retErr error) {
+	op := metrics.Begin(s.scope, "mark_acked",
+		metrics.NewTag("topic", topic),
+		metrics.NewTag("consumer_group", consumerGroup),
+		metrics.NewTag("partition_key", partitionKey))
+	defer func() { op.Complete(retErr) }()
+
+	_, err := s.db.ExecContext(ctx, fmt.Sprintf(`
+		INSERT INTO %s (consumer_group, topic, partition_key, message_offset, acked, invisible_until, retry_count)
+		VALUES (?, ?, ?, ?, TRUE, 0, 0)
+		ON DUPLICATE KEY UPDATE acked = TRUE
+	`, DeliveryStateTableName),
+		consumerGroup, topic, partitionKey, offset)
+
+	if err != nil {
+		return fmt.Errorf("mark acked topic=%s partition=%s offset=%d: %w", topic, partitionKey, offset, err)
+	}
+
+	return nil
+}
+
+// MarkNacked sets invisible_until = now + delay to schedule redelivery.
+// retry_count is NOT incremented here — it is incremented by MarkDelivered on redelivery.
+func (s *sqldeliveryStateStore) MarkNacked(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64, delayMs int64) (retErr error) {
+	op := metrics.Begin(s.scope, "mark_nacked",
+		metrics.NewTag("topic", topic),
+		metrics.NewTag("consumer_group", consumerGroup),
+		metrics.NewTag("partition_key", partitionKey))
+	defer func() { op.Complete(retErr) }()
+
+	now := time.Now().UnixMilli()
+	invisibleUntil := now + delayMs
+
+	_, err := s.db.ExecContext(ctx, fmt.Sprintf(`
+		INSERT INTO %s (consumer_group, topic, partition_key, message_offset, acked, invisible_until, retry_count)
+		VALUES (?, ?, ?, ?, FALSE, ?, 0)
+		ON DUPLICATE KEY UPDATE
+			invisible_until = VALUES(invisible_until)
+	`, DeliveryStateTableName),
+		consumerGroup, topic, partitionKey, offset, invisibleUntil)
+
+	if err != nil {
+		return fmt.Errorf("mark nacked topic=%s partition=%s offset=%d: %w", topic, partitionKey, offset, err)
+	}
+
+	return nil
+}
+
+// GetRetryCount returns the retry count for a specific message and consumer group.
+func (s *sqldeliveryStateStore) GetRetryCount(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64) (_ int, retErr error) {
+	op := metrics.Begin(s.scope, "get_retry_count",
+		metrics.NewTag("topic", topic),
+		metrics.NewTag("consumer_group", consumerGroup),
+		metrics.NewTag("partition_key", partitionKey))
+	defer func() { op.Complete(retErr) }()
+
+	var retryCount int
+	err := s.db.QueryRowContext(ctx, fmt.Sprintf(`
+		SELECT retry_count FROM %s
+		WHERE consumer_group = ? AND topic = ? AND partition_key = ? AND message_offset = ?
+	`, DeliveryStateTableName), consumerGroup, topic, partitionKey, offset).Scan(&retryCount)
+
+	if err == sql.ErrNoRows {
+		return 0, nil
+	}
+	if err != nil {
+		return 0, fmt.Errorf("get retry count topic=%s partition=%s offset=%d: %w", topic, partitionKey, offset, err)
+	}
+
+	return retryCount, nil
+}
+
+// IsDeliverable checks if a message offset is deliverable for this consumer group.
+//
+// Deliverability is determined by the acked flag and invisible_until timestamp:
+//   - No row (never delivered):                      deliverable
+//   - acked = TRUE:                                  not deliverable (already processed)
+//   - acked = FALSE, invisible_until <= now (expired): deliverable (ready for retry)
+//   - acked = FALSE, invisible_until > now (pending):  not deliverable (in-flight or nack delay)
+func (s *sqldeliveryStateStore) IsDeliverable(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64) (_ bool, retErr error) {
+	op := metrics.Begin(s.scope, "is_deliverable",
+		metrics.NewTag("topic", topic),
+		metrics.NewTag("consumer_group", consumerGroup),
+		metrics.NewTag("partition_key", partitionKey))
+	defer func() { op.Complete(retErr) }()
+
+	now := time.Now().UnixMilli()
+
+	var acked bool
+	var invisibleUntil uint64
+	err := s.db.QueryRowContext(ctx, fmt.Sprintf(`
+		SELECT acked, invisible_until FROM %s
+		WHERE consumer_group = ? AND topic = ? AND partition_key = ? AND message_offset = ?
+	`, DeliveryStateTableName), consumerGroup, topic, partitionKey, offset).Scan(&acked, &invisibleUntil)
+
+	if err == sql.ErrNoRows {
+		// No delivery state row -> never delivered -> deliverable
+		return true, nil
+	}
+	if err != nil {
+		return false, fmt.Errorf("check deliverability topic=%s partition=%s offset=%d: %w", topic, partitionKey, offset, err)
+	}
+
+	// Already processed by this consumer group — never redeliver
+	if acked {
+		return false, nil
+	}
+
+	// Message is deliverable if invisible_until has passed
+	return invisibleUntil <= uint64(now), nil
+}
+
+// AdvanceWatermark computes the new contiguous acked watermark and cleans up
+// delivery state rows that are behind it.
+// offsets are the actual message offsets above the current watermark (from messageStore).
+// Returns the new watermark (highest contiguous acked offset from currentWatermark).
+func (s *sqldeliveryStateStore) AdvanceWatermark(ctx context.Context, consumerGroup, topic, partitionKey string, currentWatermark int64, offsets []int64) (_ int64, retErr error) {
+	op := metrics.Begin(s.scope, "advance_watermark",
+		metrics.NewTag("topic", topic),
+		metrics.NewTag("consumer_group", consumerGroup),
+		metrics.NewTag("partition_key", partitionKey))
+	defer func() { op.Complete(retErr) }()
+
+	if len(offsets) == 0 {
+		return currentWatermark, nil
+	}
+
+	// Batch-fetch delivery state for the provided offsets.
+	placeholders := make([]byte, 0, len(offsets)*2-1)
+	args := make([]interface{}, 0, 3+len(offsets))
+	args = append(args, consumerGroup, topic, partitionKey)
+	for i, offset := range offsets {
+		if i > 0 {
+			placeholders = append(placeholders, ',')
+		}
+		placeholders = append(placeholders, '?')
+		args = append(args, offset)
+	}
+
+	rows, err := s.db.QueryContext(ctx, fmt.Sprintf(`
+		SELECT message_offset, acked FROM %s
+		WHERE consumer_group = ? AND topic = ? AND partition_key = ?
+		AND message_offset IN (%s)
+	`, DeliveryStateTableName, string(placeholders)), args...)
+	if err != nil {
+		return currentWatermark, fmt.Errorf("query delivery state for watermark topic=%s partition=%s: %w", topic, partitionKey, err)
+	}
+	defer rows.Close()
+
+	// Build lookup map: offset -> acked
+	ackedMap := make(map[int64]bool, len(offsets))
+	for rows.Next() {
+		var offset int64
+		var acked bool
+		if err := rows.Scan(&offset, &acked); err != nil {
+			return currentWatermark, fmt.Errorf("scan delivery state topic=%s partition=%s: %w", topic, partitionKey, err)
+		}
+		ackedMap[offset] = acked
+	}
+	if err := rows.Err(); err != nil {
+		return currentWatermark, fmt.Errorf("delivery state iteration topic=%s partition=%s: %w", topic, partitionKey, err)
+	}
+
+	// Walk message offsets in order. Advance while contiguous acked.
+	// Stop at first offset that is not acked (in-flight, nacked, or undelivered).
+	newWatermark := currentWatermark
+	for _, offset := range offsets {
+		acked, exists := ackedMap[offset]
+		if !exists || !acked {
+			// No delivery state (undelivered) or not acked — stop
+			break
+		}
+		newWatermark = offset
+	}
+
+	// Cleanup error is swallowed because the watermark was already computed and
+	// will be returned to the caller. The stale delivery state rows behind the
+	// watermark are harmless — they are never read again (all queries use
+	// offset > watermark). Cleanup is retried on the next AdvanceWatermark call.
+	if newWatermark > currentWatermark {
+		_, err := s.db.ExecContext(ctx, fmt.Sprintf(`
+			DELETE FROM %s
+			WHERE consumer_group = ? AND topic = ? AND partition_key = ? AND message_offset <= ?
+		`, DeliveryStateTableName), consumerGroup, topic, partitionKey, newWatermark)
+		if err != nil {
+			metrics.NamedCounter(s.scope, "advance_watermark", "cleanup_errors", 1,
+				metrics.NewTag("topic", topic))
+			s.logger.Warnw("failed to clean up delivery state behind watermark",
+				logTopic, topic,
+				logPartitionKey, partitionKey,
+				"watermark", newWatermark,
+				logError, err,
+			)
+		}
+	}
+
+	return newWatermark, nil
+}