Add alarm test

Publishers are blocked when an alarm is triggered, but
consumers are not.

Author: acogoluegnes

src/test/java/com/rabbitmq/model/amqp/AlarmsTest.java

@@ -0,0 +1,97 @@
+// 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.model.amqp;
+
+import static com.rabbitmq.model.amqp.TestUtils.*;
+import static java.util.stream.IntStream.range;
+
+import com.rabbitmq.model.Connection;
+import com.rabbitmq.model.Environment;
+import com.rabbitmq.model.Publisher;
+import org.junit.jupiter.api.*;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class AlarmsTest {
+
+  static Environment environment;
+  Connection connection;
+
+  @BeforeAll
+  static void initAll() {
+    environment = environmentBuilder().build();
+  }
+
+  @BeforeEach
+  void init() {
+    this.connection = environment.connectionBuilder().build();
+  }
+
+  @AfterEach
+  void tearDown() {
+    this.connection.close();
+  }
+
+  @AfterAll
+  static void tearDownAll() {
+    environment.close();
+  }
+
+  @ParameterizedTest
+  @ValueSource(strings = {"disk", "memory"})
+  void alarmShouldBlockPublisher(String alarmType) throws Exception {
+    String q = connection.management().queue().exclusive(true).declare().name();
+    Publisher publisher = connection.publisherBuilder().queue(q).build();
+    int messageCount = 100;
+    Sync publishSync = sync(messageCount);
+    range(0, messageCount)
+        .forEach(ignored -> publisher.publish(publisher.message(), ctx -> publishSync.down()));
+    assertThat(publishSync).completes();
+    publishSync.reset(messageCount + 1);
+    Sync consumeSync = sync(messageCount);
+    try (AutoCloseable ignored = alarm(alarmType)) {
+      new Thread(() -> publisher.publish(publisher.message(), ctx -> publishSync.down())).start();
+      connection
+          .consumerBuilder()
+          .queue(q)
+          .messageHandler(
+              (ctx, msg) -> {
+                ctx.accept();
+                consumeSync.down();
+              })
+          .build();
+      assertThat(consumeSync).completes();
+      consumeSync.reset(messageCount + 1);
+    }
+
+    range(0, messageCount)
+        .forEach(ignored -> publisher.publish(publisher.message(), ctx -> publishSync.down()));
+    assertThat(publishSync).completes();
+    assertThat(consumeSync).completes();
+  }
+
+  private static AutoCloseable alarm(String type) throws Exception {
+    if ("disk".equals(type)) {
+      return Cli.diskAlarm();
+    } else if ("memory".equals(type)) {
+      return Cli.memoryAlarm();
+    } else {
+      throw new IllegalArgumentException();
+    }
+  }
+}

src/test/java/com/rabbitmq/model/amqp/Cli.java

@@ -24,6 +24,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.Callable;
 import java.util.function.Predicate;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -199,6 +200,54 @@ static QueueInfo queueInfo(String q) {
     return listQueues().stream().filter(info -> q.equals(info.name())).findFirst().get();
   }
 
+  public static AutoCloseable diskAlarm() throws Exception {
+    return new CallableAutoCloseable(
+        () -> {
+          setDiskAlarm();
+          return null;
+        },
+        () -> {
+          clearDiskAlarm();
+          return null;
+        });
+  }
+
+  public static AutoCloseable memoryAlarm() throws Exception {
+    return new CallableAutoCloseable(
+        () -> {
+          setMemoryAlarm();
+          return null;
+        },
+        () -> {
+          clearMemoryAlarm();
+          return null;
+        });
+  }
+
+  private static void setDiskAlarm() {
+    setResourceAlarm("disk");
+  }
+
+  private static void clearDiskAlarm() {
+    clearResourceAlarm("disk");
+  }
+
+  private static void setMemoryAlarm() {
+    setResourceAlarm("memory");
+  }
+
+  static void clearMemoryAlarm() {
+    clearResourceAlarm("memory");
+  }
+
+  private static void setResourceAlarm(String source) {
+    rabbitmqctl("eval 'rabbit_alarm:set_alarm({{resource_limit, " + source + ", node()}, []}).'");
+  }
+
+  private static void clearResourceAlarm(String source) {
+    rabbitmqctl("eval 'rabbit_alarm:clear_alarm({resource_limit, " + source + ", node()}).'");
+  }
+
   private static class ConnectionInfo {
     private final String pid;
     private final int peerPort;
@@ -325,4 +374,19 @@ private static String extractConnectionName(String clientProperties) {
   static String hostname() {
     return executeCommand("hostname").output();
   }
+
+  private static final class CallableAutoCloseable implements AutoCloseable {
+
+    private final Callable<Void> end;
+
+    private CallableAutoCloseable(Callable<Void> start, Callable<Void> end) throws Exception {
+      this.end = end;
+      start.call();
+    }
+
+    @Override
+    public void close() throws Exception {
+      this.end.call();
+    }
+  }
 }

src/test/java/com/rabbitmq/model/amqp/TestUtils.java

@@ -49,13 +49,9 @@
 import org.junit.jupiter.api.extension.ExecutionCondition;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.api.extension.ExtensionContext;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 public abstract class TestUtils {
 
-  private static final Logger LOGGER = LoggerFactory.getLogger(TestUtils.class);
-
   private static final Duration DEFAULT_CONDITION_TIMEOUT = Duration.ofSeconds(10);
 
   private TestUtils() {}
@@ -140,6 +136,34 @@ static QueueInfoAssert assertThat(Management.QueueInfo queueInfo) {
     return new QueueInfoAssert(queueInfo);
   }
 
+  static SyncAssert assertThat(Sync sync) {
+    return new SyncAssert(sync);
+  }
+
+  static class SyncAssert extends AbstractObjectAssert<SyncAssert, Sync> {
+
+    private SyncAssert(Sync sync) {
+      super(sync, SyncAssert.class);
+    }
+
+    SyncAssert completes() {
+      return this.completes(DEFAULT_CONDITION_TIMEOUT);
+    }
+
+    SyncAssert completes(Duration timeout) {
+      try {
+        boolean completed = actual.await(timeout);
+        if (!completed) {
+          fail("Sync timed out after %d ms", timeout.toMillis());
+        }
+      } catch (InterruptedException e) {
+        Thread.interrupted();
+        throw new RuntimeException(e);
+      }
+      return this;
+    }
+  }
+
   static CountDownLatchAssert assertThat(AtomicReference<CountDownLatch> reference) {
     return new CountDownLatchAssert(reference);
   }
@@ -510,6 +534,14 @@ private QueueInfoAssert flag(String label, boolean expected, boolean actual) {
     }
   }
 
+  static Sync sync() {
+    return sync(1);
+  }
+
+  static Sync sync(int count) {
+    return new Sync(count);
+  }
+
   private static class CloseableResourceWrapper<T>
       implements ExtensionContext.Store.CloseableResource {
 
@@ -530,4 +562,25 @@ public void close() {
       this.closing.accept(this.resource);
     }
   }
+
+  static class Sync {
+
+    private final AtomicReference<CountDownLatch> latch = new AtomicReference<>();
+
+    private Sync(int count) {
+      this.latch.set(new CountDownLatch(count));
+    }
+
+    void down() {
+      this.latch.get().countDown();
+    }
+
+    private boolean await(Duration timeout) throws InterruptedException {
+      return this.latch.get().await(timeout.toMillis(), TimeUnit.MILLISECONDS);
+    }
+
+    void reset(int count) {
+      this.latch.set(new CountDownLatch(count));
+    }
+  }
 }