feat(queue/sql): add subscriber with partition leasing

## Why?
Implement the Subscriber interface with distributed partition leasing to enable reliable message consumption across multiple workers.

## What?
- Subscriber with automatic partition discovery and lease acquisition
- Per-partition polling with ack/nack support
- DLQ handling for messages exceeding retry limits
- Lease renewal and graceful shutdown
- Fast test execution (~4s) with comprehensive coverage

## Test Plan
- Message delivery and ordering verified
- Ack/nack operations work correctly
- DLQ handling tested (max retries exceeded)
- Multi-partition and multi-consumer scenarios validated
- Lease management (acquire, renew, release) tested

Author: behinddwalls

extensions/queue/sql/BUILD.bazel

@@ -9,6 +9,7 @@ go_library(
         "offset_store.go",
         "partition_lease_store.go",
         "publisher.go",
+        "subscriber.go",
         "test_helpers.go",
     ],
     importpath = "github.com/uber/submitqueue/extensions/queue/sql",
@@ -36,6 +37,7 @@ go_test(
         "offset_store_test.go",
         "partition_lease_store_test.go",
         "publisher_test.go",
+        "subscriber_test.go",
     ],
     data = ["//schema/queue/mysql:schema.sql"],
     embed = [":sql"],

extensions/queue/sql/subscriber.go

@@ -0,0 +1,369 @@
+package sql
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/uber-go/tally/v4"
+	"go.uber.org/zap"
+
+	"github.com/uber/submitqueue/entities/queue"
+)
+
+type subscriber struct {
+	config      Config
+	logger      *zap.SugaredLogger
+	metrics     tally.Scope
+	messageRepo MessageStore
+	offsetRepo  OffsetStore
+	leaseRepo   PartitionLeaseStore
+	mu          sync.RWMutex
+	closed      bool
+
+	// Active subscriptions
+	subscriptions map[string]*subscription
+	subMu         sync.Mutex
+}
+
+type subscription struct {
+	topic      string
+	deliveryCh chan *queue.Delivery
+	cancelFunc context.CancelFunc
+	wg         sync.WaitGroup
+}
+
+func newSubscriber(config Config, logger *zap.SugaredLogger, metrics tally.Scope, messageRepo MessageStore, offsetRepo OffsetStore, leaseRepo PartitionLeaseStore) *subscriber {
+	logger.Infow("created subscriber",
+		"consumer_group", config.ConsumerGroup,
+		"worker_id", config.WorkerID,
+		"poll_interval", config.PollInterval,
+		"batch_size", config.BatchSize,
+		"max_retry_attempts", config.Retry.MaxAttempts,
+		"lease_renewal_interval", config.LeaseRenewalInterval,
+	)
+
+	return &subscriber{
+		config:        config,
+		logger:        logger,
+		metrics:       metrics,
+		messageRepo:   messageRepo,
+		offsetRepo:    offsetRepo,
+		leaseRepo:     leaseRepo,
+		subscriptions: make(map[string]*subscription),
+	}
+}
+
+// Subscribe starts consuming messages from the specified topic
+func (s *subscriber) Subscribe(ctx context.Context, topic string) (<-chan *queue.Delivery, error) {
+	s.mu.RLock()
+	if s.closed {
+		s.mu.RUnlock()
+		s.logger.Errorw("subscribe failed: subscriber is closed", "topic", topic)
+		return nil, fmt.Errorf("subscriber is closed")
+	}
+	s.mu.RUnlock()
+
+	// Validate topic name
+	if err := validateTopicName(topic); err != nil {
+		s.logger.Errorw("subscribe failed: invalid topic name", "topic", topic, "error", err)
+		return nil, err
+	}
+
+	s.subMu.Lock()
+	defer s.subMu.Unlock()
+
+	// Check if already subscribed
+	if sub, exists := s.subscriptions[topic]; exists {
+		s.logger.Debugw("reusing existing subscription", "topic", topic)
+		return sub.deliveryCh, nil
+	}
+
+	s.logger.Infow("creating new subscription", "topic", topic)
+
+	// Create new subscription
+	subCtx, cancel := context.WithCancel(ctx)
+	sub := &subscription{
+		topic:      topic,
+		deliveryCh: make(chan *queue.Delivery),
+		cancelFunc: cancel,
+	}
+
+	s.subscriptions[topic] = sub
+
+	// Track active subscription
+	s.metrics.Tagged(map[string]string{"topic": topic}).Gauge("active_subscriptions").Update(1)
+
+	// Start partition leasing and polling goroutine
+	sub.wg.Add(1)
+	go s.managePartitions(subCtx, sub)
+
+	s.logger.Infow("subscription created", "topic", topic)
+	return sub.deliveryCh, nil
+}
+
+// managePartitions discovers partitions, acquires leases, and polls messages
+func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) {
+	defer sub.wg.Done()
+	defer close(sub.deliveryCh)
+
+	pollTicker := time.NewTicker(s.config.PollInterval)
+	defer pollTicker.Stop()
+
+	leaseTicker := time.NewTicker(s.config.LeaseRenewalInterval)
+	defer leaseTicker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			// Release all leases on shutdown
+			s.releaseAllLeases(context.Background(), sub.topic)
+			return
+
+		case <-leaseTicker.C:
+			// Renew existing leases
+			s.renewLeases(ctx, sub.topic)
+
+		case <-pollTicker.C:
+			// Fetch and deliver messages from leased partitions
+			s.pollLeasedPartitions(ctx, sub)
+		}
+	}
+}
+
+// renewLeases renews leases for all partitions owned by this worker
+func (s *subscriber) renewLeases(ctx context.Context, topic string) {
+	leasedPartitions, err := s.leaseRepo.GetLeasedPartitions(ctx, topic)
+	if err != nil {
+		return
+	}
+
+	for _, partitionKey := range leasedPartitions {
+		s.leaseRepo.RenewLease(ctx, topic, partitionKey)
+	}
+}
+
+// releaseAllLeases releases all leases for a topic
+func (s *subscriber) releaseAllLeases(ctx context.Context, topic string) {
+	leasedPartitions, err := s.leaseRepo.GetLeasedPartitions(ctx, topic)
+	if err != nil {
+		return
+	}
+
+	for _, partitionKey := range leasedPartitions {
+		s.leaseRepo.ReleaseLease(ctx, topic, partitionKey)
+	}
+}
+
+// pollLeasedPartitions fetches and delivers messages from all leased partitions
+func (s *subscriber) pollLeasedPartitions(ctx context.Context, sub *subscription) {
+	// Discover and try to acquire leases for new partitions
+	acquiredCount, err := s.leaseRepo.DiscoverAndAcquirePartitions(ctx, sub.topic)
+	if err == nil && acquiredCount > 0 {
+		s.metrics.Tagged(map[string]string{"topic": sub.topic}).Counter("leases_acquired").Inc(int64(acquiredCount))
+	}
+
+	// Get currently leased partitions
+	leasedPartitions, err := s.leaseRepo.GetLeasedPartitions(ctx, sub.topic)
+	if err != nil {
+		s.logger.Errorw("failed to get leased partitions", "topic", sub.topic, "error", err)
+		return
+	}
+
+	// Poll each leased partition
+	for _, partitionKey := range leasedPartitions {
+		s.fetchAndDeliverPartition(ctx, sub, partitionKey)
+	}
+}
+
+// fetchAndDeliverPartition fetches messages from a specific partition and delivers them
+func (s *subscriber) fetchAndDeliverPartition(ctx context.Context, sub *subscription, partitionKey string) {
+	start := time.Now()
+
+	// Initialize offset for this partition if needed
+	if err := s.offsetRepo.Initialize(ctx, sub.topic, partitionKey); err != nil {
+		s.logger.Errorw("failed to initialize offset", "topic", sub.topic, "partition_key", partitionKey, "error", err)
+		return
+	}
+
+	// Get current offset for this partition
+	currentOffset, err := s.offsetRepo.GetAckedOffset(ctx, sub.topic, partitionKey)
+	if err != nil {
+		s.logger.Errorw("failed to get current offset", "topic", sub.topic, "partition_key", partitionKey, "error", err)
+		return
+	}
+
+	// Fetch messages for this partition
+	rows, err := s.messageRepo.FetchByOffset(ctx, sub.topic, partitionKey, currentOffset, s.config.BatchSize)
+	if err != nil {
+		return
+	}
+
+	messageCount := 0
+	for _, row := range rows {
+		// Check if message has exceeded retry limit (persistent retry_count from DB)
+		if row.RetryCount >= s.config.Retry.MaxAttempts {
+			s.logger.Warnw("message exceeded retry limit",
+				"topic", sub.topic,
+				"partition_key", partitionKey,
+				"message_id", row.ID,
+				"retry_count", row.RetryCount,
+			)
+
+			// Move to DLQ if enabled
+			if s.config.DLQ.Enabled {
+				if err := s.messageRepo.MoveToDLQ(ctx, sub.topic, row.ID, row.RetryCount, "exceeded retry limit"); err != nil {
+					s.logger.Errorw("failed to move message to DLQ",
+						"topic", sub.topic,
+						"message_id", row.ID,
+						"error", err,
+					)
+				} else {
+					s.logger.Infow("moved message to DLQ",
+						"topic", sub.topic,
+						"message_id", row.ID,
+						"retry_count", row.RetryCount,
+					)
+					s.metrics.Tagged(map[string]string{"topic": sub.topic}).Counter("messages_moved_to_dlq").Inc(1)
+
+					// Update offset since message is now processed (moved to DLQ)
+					if err := s.offsetRepo.UpdateAckedOffset(ctx, sub.topic, partitionKey, row.Offset); err != nil {
+						s.logger.Errorw("failed to update offset after DLQ move",
+							"topic", sub.topic,
+							"partition_key", partitionKey,
+							"offset", row.Offset,
+							"error", err,
+						)
+					}
+				}
+			}
+			continue
+		}
+
+		// Create message
+		msg := queue.NewMessage(row.ID, row.Payload)
+		msg.Metadata = row.Metadata
+		msg.PartitionKey = row.PartitionKey
+		msg.PublishedAt = row.PublishedAt
+
+		// Calculate message age for metrics
+		messageAge := time.Since(row.PublishedAt)
+		s.metrics.Tagged(map[string]string{"topic": sub.topic}).Timer("message_age").Record(messageAge)
+
+		// Create delivery with ack/nack handlers
+		delivery := queue.NewDelivery(
+			msg,
+			s.createAckFunc(sub.topic, partitionKey, row.Offset, msg.ID),
+			s.createNackFunc(sub.topic, partitionKey, msg.ID, row.Offset),
+			s.createExtendVisibilityTimeoutFunc(sub.topic, msg.ID),
+		)
+
+		// Deliver message
+		select {
+		case sub.deliveryCh <- delivery:
+			messageCount++
+		case <-ctx.Done():
+			return
+		}
+	}
+
+	// Record metrics
+	if messageCount > 0 {
+		elapsed := time.Since(start)
+		tags := map[string]string{"topic": sub.topic, "partition_key": partitionKey}
+		s.metrics.Tagged(tags).Counter("messages_received").Inc(int64(messageCount))
+		s.metrics.Tagged(map[string]string{"topic": sub.topic}).Timer("poll_latency").Record(elapsed)
+
+		s.logger.Debugw("delivered messages",
+			"topic", sub.topic,
+			"partition_key", partitionKey,
+			"count", messageCount,
+			"duration_ms", elapsed.Milliseconds(),
+		)
+	}
+}
+
+// createAckFunc creates an ack function for a message
+func (s *subscriber) createAckFunc(topic string, partitionKey string, offset int64, messageID string) func(ctx context.Context) error {
+	return func(ctx context.Context) error {
+		if err := s.offsetRepo.AckMessage(ctx, topic, partitionKey, messageID, offset, s.messageRepo); err != nil {
+			return err
+		}
+
+		s.metrics.Tagged(map[string]string{"topic": topic}).Counter("messages_acked").Inc(1)
+		return nil
+	}
+}
+
+// createNackFunc creates a nack function for a message with DLQ support
+// Note: Retry tracking is now persistent via retry_count in database
+func (s *subscriber) createNackFunc(topic string, partitionKey string, messageID string, offset int64) func(ctx context.Context, requeueAfter time.Duration) error {
+	return func(ctx context.Context, requeueAfter time.Duration) error {
+		// Set visibility timeout to make message visible after requeueAfter duration
+		if err := s.messageRepo.SetVisibilityTimeout(ctx, topic, messageID, requeueAfter); err != nil {
+			s.logger.Errorw("failed to set visibility timeout for nack",
+				"topic", topic,
+				"partition_key", partitionKey,
+				"message_id", messageID,
+				"error", err,
+			)
+			return err
+		}
+
+		s.metrics.Tagged(map[string]string{"topic": topic}).Counter("messages_nacked").Inc(1)
+
+		s.logger.Infow("message nacked",
+			"topic", topic,
+			"partition_key", partitionKey,
+			"message_id", messageID,
+			"requeue_after", requeueAfter,
+		)
+
+		return nil
+	}
+}
+
+// createExtendVisibilityTimeoutFunc creates a function to extend visibility timeout for a message
+func (s *subscriber) createExtendVisibilityTimeoutFunc(topic string, messageID string) func(ctx context.Context, duration time.Duration) error {
+	return func(ctx context.Context, duration time.Duration) error {
+		if err := s.messageRepo.SetVisibilityTimeout(ctx, topic, messageID, duration); err != nil {
+			return err
+		}
+
+		s.metrics.Tagged(map[string]string{"topic": topic}).Counter("visibility_extended").Inc(1)
+		return nil
+	}
+}
+
+// Close gracefully shuts down the subscriber
+func (s *subscriber) Close() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if s.closed {
+		return nil
+	}
+
+	s.logger.Info("closing subscriber")
+
+	s.subMu.Lock()
+	defer s.subMu.Unlock()
+
+	// Cancel all subscriptions
+	for topic, sub := range s.subscriptions {
+		s.logger.Debugw("closing subscription", "topic", topic)
+		sub.cancelFunc()
+		sub.wg.Wait()
+
+		// Update metrics
+		s.metrics.Tagged(map[string]string{"topic": topic}).Gauge("active_subscriptions").Update(0)
+	}
+
+	s.subscriptions = make(map[string]*subscription)
+
+	s.closed = true
+
+	s.logger.Info("subscriber closed")
+	return nil
+}