feat: adds getHealth function in TaskManagerAmqp #1638

ICIJ · Jan 27, 2025 · f858417 · f858417
1 parent 0ff4b35
commit f858417
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 0 deletions.
diff --git a/datashare-tasks/src/main/java/org/icij/datashare/asynctasks/TaskManagerAmqp.java b/datashare-tasks/src/main/java/org/icij/datashare/asynctasks/TaskManagerAmqp.java
@@ -111,4 +111,9 @@ public void close() throws IOException {
     public void clear() {
         tasks.clear();
     }
+
+    @Override
+    public boolean getHealth() {
+        return amqp.getHealth();
+    }
 }
diff --git a/datashare-tasks/src/main/java/org/icij/datashare/asynctasks/bus/amqp/AmqpInterlocutor.java b/datashare-tasks/src/main/java/org/icij/datashare/asynctasks/bus/amqp/AmqpInterlocutor.java
@@ -133,6 +133,18 @@ public void deleteQueues(AmqpQueue... amqpQueues) throws IOException {
         }
     }
 
+    public boolean getHealth() {
+        if (!connection.isOpen()) {
+            logger.error("TaskManager AMQP Health error : Connection is closed");
+            return false;
+        }
+        if (connection.getHeartbeat() <= 0) {
+            logger.error("TaskManager AMQP Health error : heartbeat equals 0");
+            return false;
+        }
+        return true;
+    }
+
     private static class UnknownChannelException extends RuntimeException {
         public UnknownChannelException(AmqpQueue queue) {
             super("Unknown channel for queue " + queue);

diff --git a/datashare-tasks/src/test/java/org/icij/datashare/asynctasks/TaskManagerAmqpTest.java b/datashare-tasks/src/test/java/org/icij/datashare/asynctasks/TaskManagerAmqpTest.java
@@ -168,6 +168,27 @@ public void test_task_canceled() throws Exception {
         assertThat(taskManager.getTask(task.id).getState()).isEqualTo(Task.State.CANCELLED);
     }
 
+    @Test
+    public void test_health_ok() {
+        assertThat(taskManager.getHealth()).isTrue();
+    }
+
+    @Test
+    public void test_health_ko() throws Exception {
+        AmqpInterlocutor amqpKo = new AmqpInterlocutor(new PropertiesProvider(new HashMap<>() {{
+            put("messageBusAddress", "amqp://admin:admin@localhost?rabbitMq=false");
+        }}));
+        amqpKo.createAmqpChannelForPublish(AmqpQueue.TASK);
+        amqpKo.createAmqpChannelForPublish(AmqpQueue.MANAGER_EVENT);
+        RedissonClient redissonClientKo = new RedissonClientFactory().withOptions(
+                Options.from(new PropertiesProvider(Map.of("redisAddress", "redis://redis:6379")).getProperties())).create();
+        TaskManagerAmqp taskManagerAmqpKo = new TaskManagerAmqp(amqpKo, new TaskRepositoryRedis(redissonClientKo, "tasks:queue:test"), RoutingStrategy.UNIQUE, () -> nextMessage.countDown());
+
+        amqpKo.close();
+
+        assertThat(taskManagerAmqpKo.getHealth()).isFalse();
+    }
+
     @BeforeClass
     public static void beforeClass() throws Exception {
         AMQP = new AmqpInterlocutor(new PropertiesProvider(new HashMap<>() {{