Support discard and requeue with annotations

Maps to modified outcome.

Author: acogoluegnes

src/main/java/com/rabbitmq/client/amqp/Consumer.java

@@ -17,6 +17,8 @@
 // [email protected].
 package com.rabbitmq.client.amqp;
 
+import java.util.Map;
+
 /**
  * API to consume messages from a RabbitMQ queue.
  *
@@ -75,12 +77,48 @@ interface Context {
      */
     void discard();
 
+    /**
+     * Discard the message with annotations to combine with the existing message annotations.
+     *
+     * <p>This means the message cannot be processed because it is invalid, the broker can drop it
+     * or dead-letter it if it is configured.
+     *
+     * <p>Annotation keys must start with the <code>x-opt-</code> prefix.
+     *
+     * <p>This maps to the AMQP 1.0 <code>
+     * modified{delivery-failed = true, undeliverable-here = true}</code> outcome.
+     *
+     * @param annotations message annotations to combine with existing ones
+     * @see <a
+     *     href="https://docs.oasis-open.org/amqp/core/v1.0/os/amqp-core-messaging-v1.0-os.html#type-modified">AMQP
+     *     1.0 <code>modified</code> outcome</a>
+     */
+    void discard(Map<String, Object> annotations);
+
     /**
      * Requeue the message (AMQP 1.0 <code>released</code> outcome).
      *
      * <p>This means the message has not been processed and the broker can requeue it and deliver it
      * to the same or a different consumer.
      */
     void requeue();
+
+    /**
+     * Requeue the message with annotations to combine with the existing message annotations.
+     *
+     * <p>This means the message has not been processed and the broker can requeue it and deliver it
+     * to the same or a different consumer.
+     *
+     * <p>Annotation keys must start with the <code>x-opt-</code> prefix.
+     *
+     * <p>This maps to the AMQP 1.0 <code>
+     * modified{delivery-failed = false, undeliverable-here = false}</code> outcome.
+     *
+     * @param annotations message annotations to combine with existing ones
+     * @see <a
+     *     href="https://docs.oasis-open.org/amqp/core/v1.0/os/amqp-core-messaging-v1.0-os.html#type-modified">AMQP
+     *     1.0 <code>modified</code> outcome</a>
+     */
+    void requeue(Map<String, Object> annotations);
   }
 }

src/main/java/com/rabbitmq/client/amqp/impl/AmqpConsumer.java

@@ -24,6 +24,7 @@
 import com.rabbitmq.client.amqp.BackOffDelayPolicy;
 import com.rabbitmq.client.amqp.Consumer;
 import com.rabbitmq.client.amqp.metrics.MetricsCollector;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.*;
@@ -387,6 +388,24 @@ public void discard() {
       }
     }
 
+    @Override
+    public void discard(Map<String, Object> annotations) {
+      if (settled.compareAndSet(false, true)) {
+        try {
+          annotations = annotations == null ? Collections.emptyMap() : annotations;
+          checkAnnotations(annotations);
+          protonExecutor.execute(replenishCreditOperation);
+          delivery.disposition(DeliveryState.modified(true, true, annotations), true);
+          unsettledMessageCount.decrementAndGet();
+          metricsCollector.consumeDisposition(MetricsCollector.ConsumeDisposition.DISCARDED);
+        } catch (ClientIllegalStateException | RejectedExecutionException | ClientIOException e) {
+          LOGGER.debug("message discard (modified) failed: {}", e.getMessage());
+        } catch (ClientException e) {
+          throw ExceptionUtils.convert(e);
+        }
+      }
+    }
+
     @Override
     public void requeue() {
       if (settled.compareAndSet(false, true)) {
@@ -402,5 +421,33 @@ public void requeue() {
         }
       }
     }
+
+    @Override
+    public void requeue(Map<String, Object> annotations) {
+      if (settled.compareAndSet(false, true)) {
+        try {
+          annotations = annotations == null ? Collections.emptyMap() : annotations;
+          checkAnnotations(annotations);
+          protonExecutor.execute(replenishCreditOperation);
+          delivery.disposition(DeliveryState.modified(false, false, annotations), true);
+          unsettledMessageCount.decrementAndGet();
+          metricsCollector.consumeDisposition(MetricsCollector.ConsumeDisposition.REQUEUED);
+        } catch (ClientIllegalStateException | RejectedExecutionException | ClientIOException e) {
+          LOGGER.debug("message requeue (modified) failed: {}", e.getMessage());
+        } catch (ClientException e) {
+          throw ExceptionUtils.convert(e);
+        }
+      }
+    }
+  }
+
+  static void checkAnnotations(Map<String, Object> annotations) {
+    annotations.forEach(
+        (k, v) -> {
+          if (!k.startsWith("x-opt-")) {
+            throw new IllegalArgumentException(
+                "Message annotation keys must start with 'x-opt-': " + k);
+          }
+        });
   }
 }

src/test/java/com/rabbitmq/client/amqp/impl/AmqpConsumerTest.java

@@ -0,0 +1,42 @@
+// Copyright (c) 2024 Broadcom. All Rights Reserved.
+// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// If you have any questions regarding licensing, please contact us at
+// [email protected].
+package com.rabbitmq.client.amqp.impl;
+
+import static com.rabbitmq.client.amqp.impl.AmqpConsumer.checkAnnotations;
+import static java.util.Map.of;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import org.junit.jupiter.api.Test;
+
+public class AmqpConsumerTest {
+
+  @Test
+  void checkAnnotationsOK() {
+    checkAnnotations(of());
+    checkAnnotations(of("x-opt-foo", "bar"));
+    checkAnnotations(of("x-opt-foo-1", "bar1", "x-opt-foo-2", "bar2"));
+  }
+
+  @Test
+  void checkAnnotationsKO() {
+    assertThatThrownBy(() -> checkAnnotations(of("foo", "bar")))
+        .isInstanceOf(IllegalArgumentException.class);
+    assertThatThrownBy(() -> checkAnnotations(of("x-opt-foo", "bar1", "foo", "bar2")))
+        .isInstanceOf(IllegalArgumentException.class);
+  }
+}

src/test/java/com/rabbitmq/client/amqp/impl/Assertions.java

@@ -20,6 +20,7 @@
 import static org.assertj.core.api.Assertions.fail;
 
 import com.rabbitmq.client.amqp.Management;
+import com.rabbitmq.client.amqp.Message;
 import java.time.Duration;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
@@ -44,6 +45,10 @@ static SyncAssert assertThat(TestUtils.Sync sync) {
     return new SyncAssert(sync);
   }
 
+  static MessageAssert assertThat(Message message) {
+    return new MessageAssert(message);
+  }
+
   static ConnectionAssert assertThat(AmqpConnection connection) {
     return new ConnectionAssert(connection);
   }
@@ -208,6 +213,51 @@ private QueueInfoAssert flag(String label, boolean expected, boolean actual) {
     }
   }
 
+  static class MessageAssert extends AbstractObjectAssert<MessageAssert, Message> {
+
+    private MessageAssert(Message message) {
+      super(message, MessageAssert.class);
+    }
+
+    MessageAssert hasId(Object id) {
+      isNotNull();
+      if (!actual.messageId().equals(id)) {
+        fail("Message ID should be '%s' but is '%s'", id, actual.messageId());
+      }
+      return this;
+    }
+
+    MessageAssert hasAnnotation(String key) {
+      isNotNull();
+      if (!actual.hasAnnotation(key)) {
+        fail("Message should have annotation '%s' but does not", key);
+      }
+      return this;
+    }
+
+    MessageAssert hasAnnotation(String key, Object value) {
+      if (key == null || value == null) {
+        throw new IllegalArgumentException();
+      }
+      isNotNull();
+      hasAnnotation(key);
+      if (!value.equals(this.actual.annotation(key))) {
+        fail(
+            "Message should have annotation '%s = %s' but has '%s = %s'",
+            key, value, key, this.actual.annotation(key));
+      }
+      return this;
+    }
+
+    MessageAssert doesNotHaveAnnotation(String key) {
+      isNotNull();
+      if (actual.hasAnnotation(key)) {
+        fail("Message should not have annotation '%s' but has it", key);
+      }
+      return this;
+    }
+  }
+
   static class ConnectionAssert extends AbstractObjectAssert<ConnectionAssert, AmqpConnection> {
 
     private ConnectionAssert(AmqpConnection connection) {
@@ -219,7 +269,7 @@ ConnectionAssert hasNodename(String nodename) {
       isNotNull();
       if (!actual.connectionNodename().equals(nodename)) {
         fail(
-            "Connection should on node '%s' but is on node '%s'",
+            "Connection should be on node '%s' but is on node '%s'",
             nodename, actual.connectionNodename());
       }
       return this;