Support consumer priority

References rabbitmq/rabbitmq-server#11705
rabbitmq · Jul 16, 2024 · 0a11871 · 0a11871
1 parent 63a09dd
commit 0a11871
Show file tree

Hide file tree

Showing 4 changed files with 90 additions and 8 deletions.
diff --git a/src/main/java/com/rabbitmq/client/amqp/ConsumerBuilder.java b/src/main/java/com/rabbitmq/client/amqp/ConsumerBuilder.java
@@ -27,6 +27,8 @@ public interface ConsumerBuilder {
 
   ConsumerBuilder initialCredits(int initialCredits);
 
+  ConsumerBuilder priority(int priority);
+
   ConsumerBuilder listeners(Resource.StateListener... listeners);
 
   StreamOptions stream();

diff --git a/src/main/java/com/rabbitmq/client/amqp/impl/AmqpConsumer.java b/src/main/java/com/rabbitmq/client/amqp/impl/AmqpConsumer.java
@@ -22,7 +22,6 @@
 import com.rabbitmq.client.amqp.AmqpException;
 import com.rabbitmq.client.amqp.Consumer;
 import com.rabbitmq.client.amqp.metrics.MetricsCollector;
-import java.util.Collections;
 import java.util.Map;
 import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -55,6 +54,7 @@ final class AmqpConsumer extends ResourceBase implements Consumer {
   private final String address;
   private final String queue;
   private final Map<String, Object> filters;
+  private final Map<String, Object> linkProperties;
   private final AmqpConnection connection;
   private final AtomicReference<PauseStatus> pauseStatus =
       new AtomicReference<>(PauseStatus.UNPAUSED);
@@ -83,11 +83,13 @@ final class AmqpConsumer extends ResourceBase implements Consumer {
     addressBuilder.queue(builder.queue());
     this.address = addressBuilder.address();
     this.queue = builder.queue();
-    this.filters = Collections.unmodifiableMap(builder.filters());
+    this.filters = Map.copyOf(builder.filters());
+    this.linkProperties = Map.copyOf(builder.properties());
     this.connection = builder.connection();
     this.sessionHandler = this.connection.createSessionHandler();
     this.nativeReceiver =
-        this.createNativeReceiver(this.sessionHandler.session(), this.address, this.filters);
+        this.createNativeReceiver(
+            this.sessionHandler.session(), this.address, this.linkProperties, this.filters);
     this.initStateFromNativeReceiver(this.nativeReceiver);
     this.metricsCollector = this.connection.metricsCollector();
     this.startReceivingLoop();
@@ -144,14 +146,18 @@ public void close() {
   // internal API
 
   private ClientReceiver createNativeReceiver(
-      Session nativeSession, String address, Map<String, Object> filters) {
+      Session nativeSession,
+      String address,
+      Map<String, Object> properties,
+      Map<String, Object> filters) {
     try {
       ReceiverOptions receiverOptions =
           new ReceiverOptions()
               .deliveryMode(DeliveryMode.AT_LEAST_ONCE)
               .autoAccept(false)
               .autoSettle(false)
-              .creditWindow(0);
+              .creditWindow(0)
+              .properties(properties);
       if (!filters.isEmpty()) {
         receiverOptions.sourceOptions().filters(filters);
       }
@@ -205,7 +211,8 @@ private void startReceivingLoop() {
 
   void recoverAfterConnectionFailure() {
     this.nativeReceiver =
-        createNativeReceiver(this.sessionHandler.sessionNoCheck(), this.address, this.filters);
+        createNativeReceiver(
+            this.sessionHandler.sessionNoCheck(), this.address, this.linkProperties, this.filters);
     this.initStateFromNativeReceiver(this.nativeReceiver);
     this.pauseStatus.set(PauseStatus.UNPAUSED);
     this.unsettledMessageCount.set(0);

diff --git a/src/main/java/com/rabbitmq/client/amqp/impl/AmqpConsumerBuilder.java b/src/main/java/com/rabbitmq/client/amqp/impl/AmqpConsumerBuilder.java
@@ -32,7 +32,8 @@ class AmqpConsumerBuilder implements ConsumerBuilder {
   private Consumer.MessageHandler messageHandler;
   private int initialCredits = 100;
   private final List<Resource.StateListener> listeners = new ArrayList<>();
-  private final Map<String, Object> filters = new HashMap<>();
+  private final Map<String, Object> filters = new LinkedHashMap<>();
+  private final Map<String, Object> properties = new LinkedHashMap<>();
   private final StreamOptions streamOptions = new DefaultStreamOptions(this, this.filters);
 
   AmqpConsumerBuilder(AmqpConnection connection) {
@@ -57,6 +58,17 @@ public ConsumerBuilder initialCredits(int initialCredits) {
     return this;
   }
 
+  @Override
+  public ConsumerBuilder priority(int priority) {
+    if (priority < 0 || priority > 255) {
+      throw new IllegalArgumentException(
+          "The consumer priority must be between 0 and 255. "
+              + "Recommended values are between 0 and 5.");
+    }
+    this.properties.put("rabbitmq:priority", priority);
+    return this;
+  }
+
   @Override
   public ConsumerBuilder listeners(Resource.StateListener... listeners) {
     if (listeners == null || listeners.length == 0) {
@@ -85,7 +97,11 @@ Consumer.MessageHandler messageHandler() {
   }
 
   int initialCredits() {
-    return initialCredits;
+    return this.initialCredits;
+  }
+
+  Map<String, Object> properties() {
+    return this.properties;
   }
 
   List<Resource.StateListener> listeners() {

diff --git a/src/test/java/com/rabbitmq/client/amqp/impl/AmqpTest.java b/src/test/java/com/rabbitmq/client/amqp/impl/AmqpTest.java
@@ -40,6 +40,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
+import java.util.stream.IntStream;
 import org.assertj.core.api.Assertions;
 import org.junit.jupiter.api.*;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -563,6 +564,62 @@ void consumerUnsettledMessagesGoBackToQueueAfterClosing() {
         .hasMessageCount(messageCount - settledCount);
   }
 
+  @Test
+  void consumerWithHigherPriorityShouldGetMessagesFirst() {
+    int messageCount = 100;
+    connection.management().queue(name).exclusive(true).declare();
+    AtomicInteger lowCount = new AtomicInteger(0);
+    AtomicInteger highCount = new AtomicInteger(0);
+    Sync consumeSync = sync(messageCount);
+    com.rabbitmq.client.amqp.Consumer lowPriorityConsumer =
+        connection
+            .consumerBuilder()
+            .queue(name)
+            .priority(1)
+            .messageHandler(
+                (ctx, msg) -> {
+                  ctx.accept();
+                  lowCount.incrementAndGet();
+                  consumeSync.down();
+                })
+            .build();
+
+    com.rabbitmq.client.amqp.Consumer highPriorityConsumer =
+        connection
+            .consumerBuilder()
+            .queue(name)
+            .priority(5)
+            .messageHandler(
+                (ctx, msg) -> {
+                  ctx.accept();
+                  highCount.incrementAndGet();
+                  consumeSync.down();
+                })
+            .build();
+
+    Publisher publisher = connection.publisherBuilder().queue(name).build();
+    Runnable publish =
+        () ->
+            IntStream.range(0, messageCount)
+                .forEach(ignored -> publisher.publish(publisher.message(), ctx -> {}));
+
+    publish.run();
+
+    assertThat(consumeSync).completes();
+    assertThat(lowCount).hasValue(0);
+    assertThat(highCount).hasValue(messageCount);
+
+    highPriorityConsumer.close();
+
+    consumeSync.reset(messageCount);
+    publish.run();
+    assertThat(consumeSync).completes();
+    assertThat(lowCount).hasValue(messageCount);
+    assertThat(highCount).hasValue(messageCount);
+
+    lowPriorityConsumer.close();
+  }
+
   private static String uuid() {
     return UUID.randomUUID().toString();
   }