|
| 1 | +package natsq |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "errors" |
| 6 | + "log" |
| 7 | + "sync" |
| 8 | + |
| 9 | + "github.com/nats-io/nats.go" |
| 10 | + "github.com/nats-io/nats.go/jetstream" |
| 11 | + "github.com/zeromicro/go-zero/core/logx" |
| 12 | + "github.com/zeromicro/go-zero/core/queue" |
| 13 | +) |
| 14 | + |
| 15 | +const ( |
| 16 | + NatDefaultMode = iota |
| 17 | + NatJetMode |
| 18 | +) |
| 19 | + |
| 20 | +type ( |
| 21 | + Msg struct { |
| 22 | + Subject string |
| 23 | + Data []byte |
| 24 | + } |
| 25 | + |
| 26 | + ConsumeHandle func(m *Msg) error |
| 27 | + |
| 28 | + // ConsumeHandler Consumer interface, used to define the methods required by the consumer |
| 29 | + ConsumeHandler interface { |
| 30 | + HandleMessage(m *Msg) error |
| 31 | + } |
| 32 | + |
| 33 | + // ConsumerQueue Consumer queue, used to maintain the relationship between a consumer queue |
| 34 | + ConsumerQueue struct { |
| 35 | + StreamName string // stream name |
| 36 | + QueueName string // queue name |
| 37 | + Subjects []string // Subscribe subject |
| 38 | + Consumer ConsumeHandler // consumer object |
| 39 | + JetOption []jetstream.PullConsumeOpt // Jetstream configuration |
| 40 | + } |
| 41 | + |
| 42 | + // ConsumerManager Consumer manager for managing multiple consumer queues |
| 43 | + ConsumerManager struct { |
| 44 | + mutex sync.RWMutex // read-write lock |
| 45 | + conn *nats.Conn // nats connect |
| 46 | + mode uint // nats mode |
| 47 | + queues []ConsumerQueue // consumer queue list |
| 48 | + options []nats.Option // Connection configuration items |
| 49 | + doneChan chan struct{} // close channel |
| 50 | + } |
| 51 | +) |
| 52 | + |
| 53 | +// MustNewConsumerManager creates a new ConsumerManager instance. |
| 54 | +// It connects to NATS server, registers the provided consumer queues, and returns the ConsumerManager. |
| 55 | +// If any error occurs during the process, it logs the error and continues. |
| 56 | +func MustNewConsumerManager(cfg *NatsConfig, cq []*ConsumerQueue, mode uint) queue.MessageQueue { |
| 57 | + sc, err := nats.Connect(cfg.ServerUri, cfg.Options...) |
| 58 | + if err != nil { |
| 59 | + logx.Errorf("failed to connect nats, error: %v", err) |
| 60 | + } |
| 61 | + cm := &ConsumerManager{ |
| 62 | + conn: sc, |
| 63 | + options: cfg.Options, |
| 64 | + mode: mode, |
| 65 | + doneChan: make(chan struct{}), |
| 66 | + } |
| 67 | + if len(cq) == 0 { |
| 68 | + logx.Errorf("failed consumerQueue register to nats, error: cq len is 0") |
| 69 | + } |
| 70 | + for _, item := range cq { |
| 71 | + err = cm.registerQueue(item) |
| 72 | + if err != nil { |
| 73 | + logx.Errorf("failed to register nats, error: %v", err) |
| 74 | + } |
| 75 | + } |
| 76 | + |
| 77 | + return cm |
| 78 | +} |
| 79 | + |
| 80 | +// Start starts consuming messages from all the registered consumer queues. |
| 81 | +// It launches a goroutine for each consumer queue to subscribe and process messages. |
| 82 | +// The method blocks until the doneChan is closed. |
| 83 | +func (cm *ConsumerManager) Start() { |
| 84 | + cm.mutex.RLock() |
| 85 | + defer cm.mutex.RUnlock() |
| 86 | + |
| 87 | + if len(cm.queues) == 0 { |
| 88 | + logx.Errorf("no consumer queues found") |
| 89 | + } |
| 90 | + for _, consumerQueue := range cm.queues { |
| 91 | + go cm.subscribe(consumerQueue) |
| 92 | + } |
| 93 | + <-cm.doneChan |
| 94 | +} |
| 95 | + |
| 96 | +// Stop closes the NATS connection and stops the ConsumerManager. |
| 97 | +func (cm *ConsumerManager) Stop() { |
| 98 | + if cm.conn != nil { |
| 99 | + cm.conn.Close() |
| 100 | + } |
| 101 | +} |
| 102 | + |
| 103 | +// registerQueue registers a new consumer queue with the ConsumerManager. |
| 104 | +// It validates the required fields of the ConsumerQueue and adds it to the list of queues. |
| 105 | +// If any required field is missing, it returns an error. |
| 106 | +func (cm *ConsumerManager) registerQueue(queue *ConsumerQueue) error { |
| 107 | + cm.mutex.Lock() |
| 108 | + defer cm.mutex.Unlock() |
| 109 | + |
| 110 | + if cm.mode == NatJetMode && queue.StreamName == "" { |
| 111 | + return errors.New("stream name is required") |
| 112 | + } |
| 113 | + |
| 114 | + if queue.QueueName == "" { |
| 115 | + return errors.New("queue name is required") |
| 116 | + } |
| 117 | + if len(queue.Subjects) == 0 { |
| 118 | + return errors.New("subject is required") |
| 119 | + } |
| 120 | + if queue.Consumer == nil { |
| 121 | + return errors.New("consumer is required") |
| 122 | + } |
| 123 | + |
| 124 | + cm.queues = append(cm.queues, *queue) |
| 125 | + return nil |
| 126 | +} |
| 127 | + |
| 128 | +// subscribe subscribes to the specified consumer queue and starts processing messages. |
| 129 | +// If the NATS mode is NatJetMode, it creates a JetStream consumer and consumes messages using the provided options. |
| 130 | +// If the NATS mode is NatDefaultMode, it subscribes to the specified subjects using the queue name. |
| 131 | +// The method blocks until the doneChan is closed. |
| 132 | +func (cm *ConsumerManager) subscribe(queue ConsumerQueue) { |
| 133 | + ctx := context.Background() |
| 134 | + if cm.mode == NatJetMode { |
| 135 | + js, _ := jetstream.New(cm.conn) |
| 136 | + stream, err := js.Stream(ctx, "ccc") |
| 137 | + if err != nil { |
| 138 | + log.Fatalf("Error creating stream: %v", err) |
| 139 | + return |
| 140 | + } |
| 141 | + consumer, _ := stream.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{ |
| 142 | + Name: queue.QueueName, |
| 143 | + AckPolicy: jetstream.AckExplicitPolicy, |
| 144 | + FilterSubjects: queue.Subjects, |
| 145 | + }) |
| 146 | + consContext, subErr := consumer.Consume(func(msg jetstream.Msg) { |
| 147 | + err := queue.Consumer.HandleMessage(&Msg{Subject: msg.Subject(), Data: msg.Data()}) |
| 148 | + if err != nil { |
| 149 | + logx.Errorf("error handling message: %v", err.Error()) |
| 150 | + } else { |
| 151 | + msg.Ack() |
| 152 | + } |
| 153 | + }, queue.JetOption...) |
| 154 | + if subErr != nil { |
| 155 | + logx.Errorf("error subscribing to queue %s: %v", queue.QueueName, subErr.Error()) |
| 156 | + return |
| 157 | + } |
| 158 | + defer consContext.Stop() |
| 159 | + } |
| 160 | + if cm.mode == NatDefaultMode { |
| 161 | + for _, subject := range queue.Subjects { |
| 162 | + cm.conn.QueueSubscribe(subject, queue.QueueName, func(m *nats.Msg) { |
| 163 | + err := queue.Consumer.HandleMessage(&Msg{Subject: m.Subject, Data: m.Data}) |
| 164 | + if err != nil { |
| 165 | + logx.Errorf("error handling message: %v", err.Error()) |
| 166 | + } else { |
| 167 | + m.Ack() |
| 168 | + } |
| 169 | + }) |
| 170 | + } |
| 171 | + } |
| 172 | + |
| 173 | + <-cm.doneChan |
| 174 | +} |
0 commit comments