implement nats-jetstreaming subscription with multiple consumer modes

This commit implements subscription based on the nats-jetstreaming queue, leveraging JetStream to offer multiple consumer modes and supporting various consumer configurations and consumption methods to meet the demands of different business scenarios.

Author: goldensnitch123

example/natsmq/consumer/config.yaml

@@ -0,0 +1,48 @@
+# config.yaml
+
+nats:
+  url: "nats://127.0.0.1:4222"
+
+streams:
+  - name: "user"
+    description: "user stream"
+    subjects: ["user.>"]
+
+consumerQueues:
+  - consumerConfig:
+      name: "taskRegister"
+      durable: "taskRegister"
+      filterSubjects: ["user.register"]
+    delivery:
+      consumptionMethod: "consumer"
+    queueConsumerCount: 3
+    streamName: "user"
+
+  - consumerConfig:
+      name: "activityRegister"
+      durable: "activityRegister"
+      filterSubjects: ["user.register"]
+    delivery:
+      consumptionMethod: "consumer"
+    streamName: "user"
+
+  - consumerConfig:
+      name: "taskRecharger"
+      durable: "taskRecharger"
+      filterSubjects: ["user.recharge"]
+    delivery:
+      consumptionMethod: "fetch"
+    queueConsumerCount: 2
+    streamName: "user"
+
+  - consumerConfig:
+      filterSubjects: ["user.recharge"]
+    delivery:
+      consumptionMethod: "fetchNoWait"
+    streamName: "user"
+
+  - consumerConfig:
+      filterSubjects: ["subject.activity.*"]
+      ordered: true
+    delivery:
+      consumptionMethod: "consumer"

example/natsmq/consumer/consumer.go

@@ -0,0 +1,65 @@
+package main
+
+import (
+	"flag"
+	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
+	"github.com/zeromicro/go-queue/natsmq/common"
+	"github.com/zeromicro/go-queue/natsmq/consumer"
+	"github.com/zeromicro/go-zero/core/conf"
+	"log"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+)
+
+var configFile = flag.String("f", "config.yaml", "Specify the config file")
+
+type ConsumerExampleConfig struct {
+	Streams        []*common.JetStreamConfig      `json:"streams"`
+	Nats           NatsConf                       `json:"nats"`
+	ConsumerQueues []consumer.ConsumerQueueConfig `json:"consumerQueues"`
+}
+
+type NatsConf struct {
+	URL string `json:"url"`
+}
+
+type MyConsumeHandler struct{}
+
+func (h *MyConsumeHandler) Consume(msg jetstream.Msg) error {
+	log.Printf("subject [%s] Received message: %s", msg.Subject(), string(msg.Data()))
+	return nil
+}
+
+func main() {
+	flag.Parse()
+	var c ConsumerExampleConfig
+	conf.MustLoad(*configFile, &c)
+
+	var queueConfigs []*consumer.ConsumerQueueConfig
+	for i := range c.ConsumerQueues {
+		c.ConsumerQueues[i].Handler = &MyConsumeHandler{}
+		queueConfigs = append(queueConfigs, &c.ConsumerQueues[i])
+	}
+
+	natsConf := &common.NatsConfig{
+		URL:     c.Nats.URL,
+		Options: []nats.Option{},
+	}
+
+	cm, err := consumer.NewConsumerManager(natsConf, c.Streams, queueConfigs)
+	if err != nil {
+		log.Fatalf("failed to create consumer manager: %v", err)
+	}
+
+	go cm.Start()
+
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+	sig := <-sigChan
+	log.Printf("Received signal %s, shutting down...", sig)
+	cm.Stop()
+	time.Sleep(time.Second)
+}

example/natsmq/consumer/nats-jetsteam-compose.yaml

@@ -0,0 +1,15 @@
+version: "3.6"
+
+services:
+  nats:
+    image: nats:2.1.8-alpine3.11
+    command:
+      - "--jetstream"
+      - "--debug"
+      - "--port"
+      - "4222"
+      - "--http_port"
+      - "8222"
+    ports:
+      - "4222:4222"
+      - "8222:8222"

example/natsmq/publisher/config.yaml

@@ -0,0 +1,4 @@
+# config.yaml
+
+nats:
+  url: "nats://127.0.0.1:4222"

example/natsmq/publisher/publisher.go

@@ -0,0 +1,66 @@
+package main
+
+import (
+	"context"
+	"flag"
+	"github.com/nats-io/nats.go"
+	"github.com/zeromicro/go-queue/natsmq/common"
+	"github.com/zeromicro/go-queue/natsmq/publisher"
+	"github.com/zeromicro/go-zero/core/conf"
+
+	"log"
+	"time"
+)
+
+var configFile = flag.String("f", "config.yaml", "Specify the config file")
+
+type PublisherExampleConfig struct {
+	Nats NatsConf `json:"nats"`
+}
+
+type NatsConf struct {
+	URL string `json:"url"`
+}
+
+func main() {
+	flag.Parse()
+	var c PublisherExampleConfig
+	conf.MustLoad(*configFile, &c)
+
+	natsConf := &common.NatsConfig{
+		URL:     c.Nats.URL,
+		Options: []nats.Option{},
+	}
+
+	jSPublisher, err := publisher.NewJetStreamPublisher(natsConf)
+
+	if err != nil {
+		log.Fatalf("failed to NewJetStreamPublisher message: %v", err)
+	}
+
+	subjects := []string{
+		"user.register",
+		"user.recharge",
+		"subject.activity.example",
+	}
+	messages := []string{
+		"Test message: user.register message",
+		"Test message: user.recharge message",
+		"Test message: subject.activity message",
+	}
+	ctx := context.Background()
+	for i, subj := range subjects {
+		msg := []byte(messages[i])
+		for j := 0; j < 3; j++ {
+			go func(s string, m []byte) {
+				ack, err := jSPublisher.Publish(ctx, s, m)
+				if err != nil {
+					log.Fatalf("failed to publish message: %v", err)
+				}
+				log.Printf("published message to %s, ack: %+v", s, ack.Stream)
+			}(subj, msg)
+		}
+	}
+
+	time.Sleep(2 * time.Second)
+}

natsmq/common/common.go

@@ -0,0 +1,155 @@
+package common
+
+import (
+	"context"
+	"fmt"
+	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
+	"log"
+	"sync"
+	"time"
+)
+
+const (
+	DefaultStream = "defaultStream"
+)
+
+var (
+	streamRegistry  = make(map[string]*JetStreamManager)
+	registryLock    sync.RWMutex
+	streamInstances = make(map[string]jetstream.Stream)
+	streamInstLock  sync.RWMutex
+)
+
+// RegisterManager registers a JetStreamManager with the specified streamID.
+func RegisterManager(streamID string, mgr *JetStreamManager) {
+	registryLock.Lock()
+	defer registryLock.Unlock()
+	streamRegistry[streamID] = mgr
+}
+
+// GetManager retrieves the JetStreamManager for the given streamID.
+// Returns the manager and true if found; otherwise returns nil and false.
+func GetManager(streamID string) (*JetStreamManager, bool) {
+	registryLock.RLock()
+	defer registryLock.RUnlock()
+	mgr, ok := streamRegistry[streamID]
+	return mgr, ok
+}
+
+// RegisterStreamInstances initializes the JetStream contexts (if needed),
+// creates or updates streams based on the provided JetStream configurations,
+// and stores the stream instances in a global map for later usage.
+// Parameters:
+//
+//	nc   - pointer to the NATS connection
+//	cfgs - list of JetStreamConfig configurations to register
+func RegisterStreamInstances(nc *nats.Conn, cfgs []*JetStreamConfig) {
+	// Register managers for each provided configuration if not already registered.
+	if len(cfgs) > 0 {
+		for _, cfg := range cfgs {
+			if _, ok := GetManager(cfg.Name); !ok {
+				mgr := NewJetStream(cfg)
+				RegisterManager(cfg.Name, mgr)
+			} else {
+				log.Printf("manager for stream %q already registered", cfg.Name)
+			}
+		}
+	}
+
+	// Iterate through all registered stream managers to initialize JetStream and create stream instances.
+	for streamName, streamMgr := range streamRegistry {
+		streamInstLock.RLock()
+		_, exists := streamInstances[streamName]
+		streamInstLock.RUnlock()
+		if exists {
+			log.Printf("streamInstance %q already created", streamName)
+			continue
+		}
+		// Initialize JetStream context
+		if err := streamMgr.InitJetStream(nc); err != nil {
+			log.Printf("failed to initialize jetstream for stream %q: %v", streamName, err)
+			continue
+		}
+		ctx := context.Background()
+		stream, err := streamMgr.CreateStream(ctx)
+		if err != nil {
+			log.Printf("failed to create stream %q: %v", streamName, err)
+			continue
+		}
+		streamInstLock.Lock()
+		streamInstances[streamName] = stream
+		streamInstLock.Unlock()
+		log.Printf("streamInstance %q created", streamName)
+	}
+}
+
+// GetStream retrieves a JetStream stream instance by streamID.
+// Returns the stream instance and true if found.
+func GetStream(streamID string) (jetstream.Stream, bool) {
+	streamInstLock.RLock()
+	defer streamInstLock.RUnlock()
+	stream, ok := streamInstances[streamID]
+	return stream, ok
+}
+
+func init() {
+	// Registers a default stream manager if one hasn't been registered.
+	if _, ok := GetManager(DefaultStream); !ok {
+		defaultCfg := &JetStreamConfig{
+			Name:                 DefaultStream,
+			Subjects:             []string{"subject.*.*"},
+			Description:          DefaultStream,
+			Retention:            0,
+			MaxConsumers:         30,
+			MaxMsgs:              -1,
+			MaxBytes:             -1,
+			Discard:              0,
+			DiscardNewPerSubject: false,
+			MaxAge:               0,
+			MaxMsgsPerSubject:    10000,
+			MaxMsgSize:           10000,
+			NoAck:                false,
+		}
+		defaultManager := NewJetStream(defaultCfg)
+		RegisterManager(DefaultStream, defaultManager)
+		log.Printf("default stream %q registered", DefaultStream)
+	}
+}
+
+// InitJetStream initializes the JetStream context for the manager using the given NATS connection.
+// Parameters:
+//
+//	nc - pointer to the NATS connection
+//
+// Returns:
+//
+//	error - non-nil error if JetStream context creation fails
+func (jsm *JetStreamManager) InitJetStream(nc *nats.Conn) error {
+	js, err := jetstream.New(nc)
+	if err != nil {
+		nc.Close()
+		return fmt.Errorf("failed to create JetStream context: %w", err)
+	}
+	jsm.JS = js
+	return nil
+}
+
+// CreateStream creates or updates a JetStream stream using the manager's configuration.
+// Parameters:
+//
+//	ctx - context to control request timeout
+//
+// Returns:
+//
+//	jetstream.Stream - the created/updated stream instance
+//	error            - non-nil error if stream creation fails
+func (jsm *JetStreamManager) CreateStream(ctx context.Context) (jetstream.Stream, error) {
+	streamCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	defer cancel()
+	stream, err := jsm.JS.CreateOrUpdateStream(streamCtx, jsm.streamConf)
+	if err != nil {
+		return nil, err
+	}
+	return stream, nil
+}