feat: Nats Jetstream consumer influxdata#11046

b3rtram · Jun 23, 2022 · 0700bf8 · 0700bf8
1 parent 44a4df8
commit 0700bf8
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 4 deletions.
diff --git a/plugins/inputs/nats_consumer/README.md b/plugins/inputs/nats_consumer/README.md
@@ -17,6 +17,9 @@ instances of telegraf can read from a NATS cluster in parallel.
   ## subject(s) to consume
   subjects = ["telegraf"]
 
+  ## jetstream subjects
+  js_subjects = ["js_telegraf"]
+
   ## name a queue group
   queue_group = "telegraf_consumers"
 

diff --git a/plugins/inputs/nats_consumer/nats_consumer.go b/plugins/inputs/nats_consumer/nats_consumer.go
@@ -46,6 +46,7 @@ type natsConsumer struct {
 	Username    string   `toml:"username"`
 	Password    string   `toml:"password"`
 	Credentials string   `toml:"credentials"`
+	JsSubjects  []string `toml:"js_subjects"`
 
 	tls.ClientConfig
 
@@ -58,8 +59,10 @@ type natsConsumer struct {
 	MaxUndeliveredMessages int `toml:"max_undelivered_messages"`
 	MetricBuffer           int `toml:"metric_buffer" deprecated:"0.10.3;2.0.0;option is ignored"`
 
-	conn *nats.Conn
-	subs []*nats.Subscription
+	conn   *nats.Conn
+	jsConn nats.JetStreamContext
+	subs   []*nats.Subscription
+	jsSubs []*nats.Subscription
 
 	parser parsers.Parser
 	// channel for all incoming NATS messages
@@ -142,6 +145,32 @@ func (n *natsConsumer) Start(acc telegraf.Accumulator) error {
 
 			n.subs = append(n.subs, sub)
 		}
+
+		if len(n.JsSubjects) > 0 {
+			var connErr error
+			n.jsConn, connErr = n.conn.JetStream(nats.PublishAsyncMaxPending(256))
+			if connErr != nil {
+				return connErr
+			}
+
+			for _, jsSub := range n.JsSubjects {
+				sub, err := n.jsConn.QueueSubscribe(jsSub, n.QueueGroup, func(m *nats.Msg) {
+					n.in <- m
+				})
+				if err != nil {
+					return err
+				}
+
+				// set the subscription pending limits
+				err = sub.SetPendingLimits(n.PendingMessageLimit, n.PendingBytesLimit)
+				if err != nil {
+					return err
+				}
+
+				n.jsSubs = append(n.jsSubs, sub)
+			}
+		}
+
 	}
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -154,8 +183,8 @@ func (n *natsConsumer) Start(acc telegraf.Accumulator) error {
 		go n.receiver(ctx)
 	}()
 
-	n.Log.Infof("Started the NATS consumer service, nats: %v, subjects: %v, queue: %v",
-		n.conn.ConnectedUrl(), n.Subjects, n.QueueGroup)
+	n.Log.Infof("Started the NATS consumer service, nats: %v, subjects: %v, jssubjects: %v, queue: %v",
+		n.conn.ConnectedUrl(), n.Subjects, n.JsSubjects, n.QueueGroup)
 
 	return nil
 }
@@ -205,6 +234,13 @@ func (n *natsConsumer) clean() {
 		}
 	}
 
+	for _, sub := range n.jsSubs {
+		if err := sub.Unsubscribe(); err != nil {
+			n.Log.Errorf("Error unsubscribing from subject %s in queue %s: %s",
+				sub.Subject, sub.Queue, err.Error())
+		}
+	}
+
 	if n.conn != nil && !n.conn.IsClosed() {
 		n.conn.Close()
 	}

diff --git a/plugins/inputs/nats_consumer/sample.conf b/plugins/inputs/nats_consumer/sample.conf
@@ -6,6 +6,9 @@
   ## subject(s) to consume
   subjects = ["telegraf"]
 
+  ## jetstream subjects
+  js_subjects = ["js_telegraf"]
+
   ## name a queue group
   queue_group = "telegraf_consumers"