Add --message-priority for STOMP and MQTT

rabbitmq · Dec 14, 2023 · 333a49f · 333a49f
1 parent 4be3e5e
commit 333a49f
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 2 deletions.
diff --git a/cmd/root.go b/cmd/root.go
@@ -6,6 +6,7 @@ import (
 	"math"
 	"os"
 	"os/signal"
+	"strconv"
 	"strings"
 	"sync"
 	"syscall"
@@ -150,6 +151,14 @@ func RootCmd() *cobra.Command {
 			if cfg.Publishers == 0 || cfg.Consumers == 0 {
 				cfg.UseMillis = true
 			}
+
+			if cfg.MessagePriority != "" {
+				_, err := strconv.ParseUint(cfg.MessagePriority, 10, 8)
+				if err != nil {
+					_, _ = fmt.Fprintf(os.Stderr, "ERROR: invalid message priority: %s\n", cfg.MessagePriority)
+					os.Exit(1)
+				}
+			}
 		},
 	}
 	rootCmd.PersistentFlags().StringVarP(&cfg.PublisherUri, "publisher-uri", "", "", "URI for publishing")
@@ -174,6 +183,7 @@ func RootCmd() *cobra.Command {
 	rootCmd.PersistentFlags().StringVar(&cfg.Amqp.Subject, "amqp-subject", "", "AMQP 1.0 message subject")
 	rootCmd.PersistentFlags().
 		BoolVarP(&cfg.MessageDurability, "message-durability", "d", true, "Mark messages as durable (default=true)")
+	rootCmd.PersistentFlags().StringVar(&cfg.MessagePriority, "message-priority", "", "Message priority (0-255, default=unset)")
 	rootCmd.PersistentFlags().StringVar(&cfg.StreamOffset, "stream-offset", "", "Stream consumer offset specification (default=next)")
 	rootCmd.PersistentFlags().StringVar(&cfg.StreamFilterValues, "stream-filter-values", "", "Stream consumer filter")
 	rootCmd.PersistentFlags().StringVar(&cfg.StreamFilterValueSet, "stream-filter-value-set", "", "Stream filter value for publisher")

diff --git a/pkg/amqp10_client/publisher.go b/pkg/amqp10_client/publisher.go
@@ -3,6 +3,7 @@ package amqp10_client
 import (
 	"context"
 	"math/rand"
+	"strconv"
 	"time"
 
 	"github.com/rabbitmq/omq/pkg/config"
@@ -127,8 +128,14 @@ func (p Amqp10Publisher) Send() {
 		msg.Annotations = amqp.Annotations{"x-stream-filter-value": p.Config.StreamFilterValueSet}
 	}
 
-	msg.Header = &amqp.MessageHeader{
-		Durable: p.Config.MessageDurability}
+	msg.Header = &amqp.MessageHeader{}
+	msg.Header.Durable = p.Config.MessageDurability
+	if p.Config.MessagePriority != "" {
+		// already validated in root.go
+		priority, _ := strconv.ParseUint(p.Config.MessagePriority, 10, 8)
+		msg.Header.Priority = uint8(priority)
+	}
+
 	timer := prometheus.NewTimer(metrics.PublishingLatency.With(prometheus.Labels{"protocol": "amqp-1.0"}))
 	err := p.Sender.Send(context.TODO(), msg, nil)
 	timer.ObserveDuration()

diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -45,6 +45,7 @@ type Config struct {
 	UseMillis            bool
 	QueueDurability      AmqpDurabilityMode
 	MessageDurability    bool
+	MessagePriority      string // to allow for "unset" value and STOMP takes strings anyway
 	StreamOffset         string
 	StreamFilterValues   string
 	StreamFilterValueSet string

diff --git a/pkg/stomp_client/publisher.go b/pkg/stomp_client/publisher.go
@@ -129,6 +129,9 @@ func buildHeaders(cfg config.Config) []func(*frame.Frame) error {
 		msgDurability = "false"
 	}
 	headers = append(headers, stomp.SendOpt.Header("persistent", msgDurability))
+	if cfg.MessagePriority != "" {
+		headers = append(headers, stomp.SendOpt.Header("priority", cfg.MessagePriority))
+	}
 
 	if cfg.StreamFilterValueSet != "" {
 		headers = append(headers, stomp.SendOpt.Header("x-stream-filter-value", cfg.StreamFilterValueSet))