101 lines
3.1 KiB
Go
101 lines
3.1 KiB
Go
package nats
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
|
|
"github.com/nats-io/nats.go"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
// ConsumerManager handles JetStream consumer lifecycle management
|
|
type ConsumerManager struct {
|
|
js nats.JetStreamContext
|
|
streamName string
|
|
consumerName string
|
|
subject string
|
|
logger *zap.Logger
|
|
}
|
|
|
|
// NewConsumerManager creates a new consumer manager
|
|
func NewConsumerManager(js nats.JetStreamContext, streamName, consumerName, subject string, logger *zap.Logger) *ConsumerManager {
|
|
if logger == nil {
|
|
logger = zap.NewNop()
|
|
}
|
|
return &ConsumerManager{
|
|
js: js,
|
|
streamName: streamName,
|
|
consumerName: consumerName,
|
|
subject: subject,
|
|
logger: logger,
|
|
}
|
|
}
|
|
|
|
// EnsureConsumer creates the consumer if it doesn't exist; if it already exists, it is reused.
|
|
func (cm *ConsumerManager) EnsureConsumer(config *nats.ConsumerConfig) error {
|
|
// First check if the consumer already exists to make this initialization idempotent.
|
|
info, err := cm.js.ConsumerInfo(cm.streamName, cm.consumerName)
|
|
if err == nil && info != nil {
|
|
cm.logger.Info("Using existing JetStream consumer",
|
|
zap.String("consumer", cm.consumerName),
|
|
zap.String("stream", cm.streamName),
|
|
zap.String("subject", cm.subject),
|
|
)
|
|
return nil
|
|
}
|
|
if err != nil && !errors.Is(err, nats.ErrConsumerNotFound) {
|
|
return fmt.Errorf("failed to fetch consumer info: %w", err)
|
|
}
|
|
|
|
// Consumer does not exist; create it.
|
|
if _, err := cm.js.AddConsumer(cm.streamName, config); err != nil {
|
|
return fmt.Errorf("failed to create consumer: %w", err)
|
|
}
|
|
|
|
cm.logger.Info("Created JetStream consumer",
|
|
zap.String("consumer", cm.consumerName),
|
|
zap.String("stream", cm.streamName),
|
|
zap.String("subject", cm.subject),
|
|
zap.Duration("ack_wait", config.AckWait),
|
|
)
|
|
|
|
return nil
|
|
}
|
|
|
|
// CreatePullSubscription creates a pull subscription with recovery logic
|
|
func (cm *ConsumerManager) CreatePullSubscription() (*nats.Subscription, error) {
|
|
return cm.js.PullSubscribe(cm.subject, cm.consumerName, nats.Bind(cm.streamName, cm.consumerName))
|
|
}
|
|
|
|
// CreatePullSubscriptionWithRecovery creates a pull subscription
|
|
func (cm *ConsumerManager) CreatePullSubscriptionWithRecovery(streamManager *StreamManager, consumerConfig *nats.ConsumerConfig) (*nats.Subscription, error) {
|
|
sub, err := cm.CreatePullSubscription()
|
|
if err == nil {
|
|
return sub, nil
|
|
}
|
|
|
|
// Attempt recovery when the consumer or stream is missing.
|
|
if !errors.Is(err, nats.ErrConsumerNotFound) && !errors.Is(err, nats.ErrStreamNotFound) {
|
|
return nil, fmt.Errorf("create pull subscription: %w", err)
|
|
}
|
|
|
|
if streamManager != nil {
|
|
if streamErr := streamManager.EnsureStream(nil); streamErr != nil {
|
|
return nil, fmt.Errorf("recover stream %s: %w", cm.streamName, streamErr)
|
|
}
|
|
}
|
|
|
|
if consumerConfig == nil {
|
|
return nil, fmt.Errorf("consumer config is required for recovery")
|
|
}
|
|
if err := cm.EnsureConsumer(consumerConfig); err != nil {
|
|
return nil, fmt.Errorf("recover consumer %s: %w", cm.consumerName, err)
|
|
}
|
|
|
|
sub, err = cm.CreatePullSubscription()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create pull subscription after recovery: %w", err)
|
|
}
|
|
return sub, nil
|
|
}
|