sonet/pkg/plugins/mq/nats_jet_stream.go

package mq

import (
	"context"
	"github.com/google/uuid"
	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"
	"regexp"
	"sonet/pkg/utils/logger"
	"sync"
	"time"
)

var (
	uuidPattern = regexp.MustCompile("^\\w+(-\\w+){4}$")
)

// NatsJetStreamProducer use for msg ack, msg resend,
type NatsJetStreamProducer struct {
	options nats.Options
	nc      *nats.Conn
	js      jetstream.JetStream
}

func NewNatsJetStreamProducer(options nats.Options) *NatsJetStreamProducer {
	return &NatsJetStreamProducer{
		options: options,
	}
}

func (mq *NatsJetStreamProducer) Init() error {
	nc, err := mq.options.Connect()
	if err != nil {
		return err
	}
	mq.nc = nc

	mq.js, err = jetstream.New(mq.nc)
	return err
}

func (mq *NatsJetStreamProducer) Publish(topic string, msg []byte) error {
	// TODO 是否ack性能430倍差距，通过配置确定publish方式，大协程池等待ack的网络io(自动重试，失败通知[后续优化])
	//err := mq.nc.Publish(topic, msg)
	_, err := mq.js.Publish(context.Background(), topic, msg)
	return err
}

func (mq *NatsJetStreamProducer) MultiPublish(topic string, msgs [][]byte) error {
	for _, msg := range msgs {
		err := mq.Publish(topic, msg)
		if err != nil {
			return err
		}
	}
	return nil
}

func (mq *NatsJetStreamProducer) DelayPublish(topic string, delay time.Duration, body []byte) error {
	panic("Nats JetStream nonsupport delay publish")
}

func (mq *NatsJetStreamProducer) Stop() {
	mq.nc.Close()
}

type NatsJetStreamConsumer struct {
	options        nats.Options
	streamConfig   jetstream.StreamConfig
	consumerConfig jetstream.ConsumerConfig

	nc     *nats.Conn
	js     jetstream.JetStream
	stream jetstream.Stream

	lock      *sync.Mutex
	consumers map[string]map[string]jetstream.Consumer
	consumes  map[string]jetstream.ConsumeContext
}

// NewNatsJetStreamConsumer
// streamConfig name,subjects需配置
// consumerConfig name, filterSubject 无需配置，在 Subscribe 时配置
func NewNatsJetStreamConsumer(
	options nats.Options,
	streamConfig jetstream.StreamConfig,
	consumerConfig jetstream.ConsumerConfig) *NatsJetStreamConsumer {

	return &NatsJetStreamConsumer{
		options:        options,
		streamConfig:   streamConfig,
		consumerConfig: consumerConfig,
		lock:           &sync.Mutex{},
		consumers:      make(map[string]map[string]jetstream.Consumer),
		consumes:       make(map[string]jetstream.ConsumeContext),
	}
}

func (mq *NatsJetStreamConsumer) Init(ctx context.Context) error {
	nc, err := mq.options.Connect()
	if err != nil {
		return err
	}
	mq.nc = nc
	mq.js, err = jetstream.New(nc)

	mq.stream, err = mq.js.CreateStream(ctx, mq.streamConfig)
	if err != nil {
		return err
	}
	return nil
}

func (mq *NatsJetStreamConsumer) Subscribe(filterSubject string, channel string, handler func(*Message) error) error {
	mq.lock.Lock()
	defer mq.lock.Unlock()

	// get consumer
	channelConsumers, ok := mq.consumers[channel]
	if !ok {
		channelConsumers = make(map[string]jetstream.Consumer)
		mq.consumers[channel] = channelConsumers
	}
	consumer, ok := channelConsumers[filterSubject]
	if !ok {
		consumerConfig := mq.consumerConfig
		consumerConfig.Name = channel
		consumerConfig.Durable = channel
		consumerConfig.FilterSubject = filterSubject
		var err error
		consumer, err = mq.stream.CreateOrUpdateConsumer(context.Background(), consumerConfig)
		if err != nil {
			return err
		}
		channelConsumers[filterSubject] = consumer
	}

	// consume message
	consumeKey := channel + "/" + filterSubject
	consumeContext, ok := mq.consumes[consumeKey]
	if ok {
		consumeContext.Stop()
		delete(mq.consumers, consumeKey)
	}

	consumeContext, err := consumer.Consume(func(msg jetstream.Msg) {
		message := &Message{
			Body: msg.Data(),
			Time: time.Now().UnixMilli(),
		}
		err := handler(message)
		if err == nil {
			// ack message
			err := msg.Ack()
			if err != nil {
				logger.Error("ack message error:", err, message)
			}
		}
	})
	if err != nil {
		return err
	}
	mq.consumes[consumeKey] = consumeContext

	return nil
}

func (mq *NatsJetStreamConsumer) SubscribeBroadcast(filterSubject string, handler func(*Message) error) error {
	return mq.Subscribe(filterSubject, uuid.New().String(), handler)
}

func (mq *NatsJetStreamConsumer) Stop() {
	mq.lock.Lock()
	defer mq.lock.Unlock()

	for _, consume := range mq.consumes {
		consume.Stop()
	}
	for channel, _ := range mq.consumers {
		// 删除 uuid 的广播 consumer
		if uuidPattern.MatchString(channel) {
			err := mq.js.DeleteConsumer(context.Background(), mq.streamConfig.Name, channel)
			if err != nil {
				logger.Errorf("delete consumer error, stream=%s, consumer=%s\n", mq.streamConfig.Name, channel)
			}
		}
	}

	mq.nc.Close()
}