Expose current task statuses count on /metrics endpoint

Author: mcornaton

CHANGELOG.md

@@ -12,7 +12,7 @@ All notable changes to this project will be documented in this file.
 - Add prometheus endpoint with custom metrics. (#632)
 - Expose version through prometheus endpoint. (#637, #639)
 - Stop fetching completed tasks count from DB. (#638)
-- Expose current task statuses count to Prometheus. (#640)
+- Expose current task statuses count to Prometheus and `/metrics` endpoint. (#640, #654)
 - Add `tasks` endpoints to `iexec-core-library`. (#645)
 
 ### Quality

iexec-core-library/src/main/java/com/iexec/core/metric/PlatformMetric.java

@@ -18,10 +18,13 @@
 
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 import com.fasterxml.jackson.databind.annotation.JsonPOJOBuilder;
+import com.iexec.core.task.TaskStatus;
 import lombok.Builder;
 import lombok.Value;
 
 import java.math.BigInteger;
+import java.util.LinkedHashMap;
+import java.util.concurrent.atomic.AtomicLong;
 
 @Value
 @Builder
@@ -32,7 +35,7 @@ public class PlatformMetric {
     int aliveAvailableCpu;
     int aliveTotalGpu;
     int aliveAvailableGpu;
-    long completedTasks;
+    LinkedHashMap<TaskStatus, AtomicLong> currentTaskStatusCounts;
     long dealEventsCount;
     long dealsCount;
     long replayDealsCount;

src/main/java/com/iexec/core/metric/MetricService.java

@@ -17,22 +17,27 @@
 package com.iexec.core.metric;
 
 import com.iexec.core.chain.DealWatcherService;
-import com.iexec.core.task.TaskService;
+import com.iexec.core.task.TaskStatus;
+import com.iexec.core.task.event.TaskStatusesCountUpdatedEvent;
 import com.iexec.core.worker.WorkerService;
+import org.springframework.context.event.EventListener;
 import org.springframework.stereotype.Service;
 
+import java.util.LinkedHashMap;
+import java.util.concurrent.atomic.AtomicLong;
+
 @Service
 public class MetricService {
     private final DealWatcherService dealWatcherService;
     private final WorkerService workerService;
-    private final TaskService taskService;
+    private LinkedHashMap<TaskStatus, AtomicLong> currentTaskStatusesCount;
 
     public MetricService(DealWatcherService dealWatcherService,
-                         WorkerService workerService,
-                         TaskService taskService) {
+                         WorkerService workerService) {
         this.dealWatcherService = dealWatcherService;
         this.workerService = workerService;
-        this.taskService = taskService;
+
+        this.currentTaskStatusesCount = new LinkedHashMap<>();
     }
 
     public PlatformMetric getPlatformMetrics() {
@@ -42,11 +47,16 @@ public PlatformMetric getPlatformMetrics() {
                 .aliveAvailableCpu(workerService.getAliveAvailableCpu())
                 .aliveTotalGpu(workerService.getAliveTotalGpu())
                 .aliveAvailableGpu(workerService.getAliveAvailableGpu())
-                .completedTasks(taskService.getCompletedTasksCount())
+                .currentTaskStatusCounts(currentTaskStatusesCount)
                 .dealEventsCount(dealWatcherService.getDealEventsCount())
                 .dealsCount(dealWatcherService.getDealsCount())
                 .replayDealsCount(dealWatcherService.getReplayDealsCount())
                 .latestBlockNumberWithDeal(dealWatcherService.getLatestBlockNumberWithDeal())
                 .build();
     }
+
+    @EventListener
+    void onTaskStatusesCountUpdateEvent(TaskStatusesCountUpdatedEvent event) {
+        this.currentTaskStatusesCount = event.getCurrentTaskStatusesCount();
+    }
 }

src/main/java/com/iexec/core/task/event/TaskStatusesCountUpdatedEvent.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2024-2024 IEXEC BLOCKCHAIN TECH
+ *
+ * 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.
+ */
+
+package com.iexec.core.task.event;
+
+import com.iexec.core.task.TaskStatus;
+import lombok.Value;
+
+import java.util.LinkedHashMap;
+import java.util.concurrent.atomic.AtomicLong;
+
+@Value
+public class TaskStatusesCountUpdatedEvent {
+    LinkedHashMap<TaskStatus, AtomicLong> currentTaskStatusesCount;
+}

src/main/java/com/iexec/core/task/update/TaskUpdateManager.java

@@ -65,7 +65,7 @@ class TaskUpdateManager {
     private final BlockchainAdapterService blockchainAdapterService;
     private final SmsService smsService;
 
-    private final Map<TaskStatus, AtomicLong> currentTaskStatusesCount;
+    private final LinkedHashMap<TaskStatus, AtomicLong> currentTaskStatusesCount;
     private final ExecutorService taskStatusesCountExecutor;
 
     public TaskUpdateManager(TaskService taskService,
@@ -104,18 +104,25 @@ public TaskUpdateManager(TaskService taskService,
     @PostConstruct
     Future<Void> init() {
         return taskStatusesCountExecutor.submit(
-                // The following could take a bit of time, depending on how many tasks are in DB.
-                // It is expected to take ~1.7s for 1,000,000 tasks and to be linear (so, ~17s for 10,000,000 tasks).
-                // As we use AtomicLongs, the final count should be accurate - no race conditions to expect,
-                // even though new deals are detected during the count.
-                () -> currentTaskStatusesCount
-                        .entrySet()
-                        .parallelStream()
-                        .forEach(entry -> entry.getValue().addAndGet(taskService.countByCurrentStatus(entry.getKey()))),
+                this::initializeCurrentTaskStatusesCount,
                 null    // Trick to get a `Future<Void>` instead of a `Future<?>`
         );
     }
 
+    /**
+     * The following could take a bit of time, depending on how many tasks are in DB.
+     * It is expected to take ~1.7s for 1,000,000 tasks and to be linear (so, ~17s for 10,000,000 tasks).
+     * As we use AtomicLongs, the final count should be accurate - no race conditions to expect,
+     * even though new deals are detected during the count.
+     */
+    private void initializeCurrentTaskStatusesCount() {
+        currentTaskStatusesCount
+                .entrySet()
+                .parallelStream()
+                .forEach(entry -> entry.getValue().addAndGet(taskService.countByCurrentStatus(entry.getKey())));
+        publishTaskStatusesCountUpdate();
+    }
+
     void updateTask(String chainTaskId) {
         Optional<Task> optional = taskService.getTaskByChainTaskId(chainTaskId);
         if (optional.isEmpty()) {
@@ -710,10 +717,17 @@ void toFailed(Task task, TaskStatus reason) {
     void updateMetricsAfterStatusUpdate(TaskStatus previousStatus, TaskStatus newStatus) {
         currentTaskStatusesCount.get(previousStatus).decrementAndGet();
         currentTaskStatusesCount.get(newStatus).incrementAndGet();
+        publishTaskStatusesCountUpdate();
     }
 
     @EventListener(TaskCreatedEvent.class)
     void onTaskCreatedEvent() {
         currentTaskStatusesCount.get(RECEIVED).incrementAndGet();
+        publishTaskStatusesCountUpdate();
+    }
+
+    private void publishTaskStatusesCountUpdate() {
+        final TaskStatusesCountUpdatedEvent event = new TaskStatusesCountUpdatedEvent(currentTaskStatusesCount);
+        applicationEventPublisher.publishEvent(event);
     }
 }

src/test/java/com/iexec/core/metric/MetricServiceTests.java

@@ -17,17 +17,21 @@
 package com.iexec.core.metric;
 
 import com.iexec.core.chain.DealWatcherService;
-import com.iexec.core.task.TaskService;
+import com.iexec.core.task.TaskStatus;
+import com.iexec.core.task.event.TaskStatusesCountUpdatedEvent;
 import com.iexec.core.worker.Worker;
 import com.iexec.core.worker.WorkerService;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.InjectMocks;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 import java.math.BigInteger;
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicLong;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.when;
@@ -38,8 +42,6 @@ class MetricServiceTests {
     private DealWatcherService dealWatcherService;
     @Mock
     private WorkerService workerService;
-    @Mock
-    private TaskService taskService;
 
     @InjectMocks
     private MetricService metricService;
@@ -51,30 +53,60 @@ void init() {
 
     @Test
     void shouldGetPlatformMetrics() {
+        final LinkedHashMap<TaskStatus, AtomicLong> expectedCurrentTaskStatusesCount = createExpectedCurrentTaskStatusesCount();
+
         List<Worker> aliveWorkers = List.of(new Worker());
         when(workerService.getAliveWorkers()).thenReturn(aliveWorkers);
         when(workerService.getAliveTotalCpu()).thenReturn(1);
         when(workerService.getAliveAvailableCpu()).thenReturn(1);
         when(workerService.getAliveTotalGpu()).thenReturn(1);
         when(workerService.getAliveAvailableGpu()).thenReturn(1);
-        when(taskService.getCompletedTasksCount())
-                .thenReturn(10L);
         when(dealWatcherService.getDealEventsCount()).thenReturn(10L);
         when(dealWatcherService.getDealsCount()).thenReturn(8L);
         when(dealWatcherService.getReplayDealsCount()).thenReturn(2L);
         when(dealWatcherService.getLatestBlockNumberWithDeal()).thenReturn(BigInteger.valueOf(255L));
 
         PlatformMetric metric = metricService.getPlatformMetrics();
-        assertThat(metric.getAliveWorkers()).isEqualTo(aliveWorkers.size());
-        assertThat(metric.getAliveTotalCpu()).isEqualTo(1);
-        assertThat(metric.getAliveAvailableCpu()).isEqualTo(1);
-        assertThat(metric.getAliveTotalGpu()).isEqualTo(1);
-        assertThat(metric.getAliveAvailableGpu()).isEqualTo(1);
-        assertThat(metric.getCompletedTasks()).isEqualTo(10L);
-        assertThat(metric.getDealEventsCount()).isEqualTo(10);
-        assertThat(metric.getDealsCount()).isEqualTo(8);
-        assertThat(metric.getReplayDealsCount()).isEqualTo(2);
-        assertThat(metric.getLatestBlockNumberWithDeal()).isEqualTo(255);
+        Assertions.assertAll(
+                () -> assertThat(metric.getAliveWorkers()).isEqualTo(aliveWorkers.size()),
+                () -> assertThat(metric.getAliveTotalCpu()).isEqualTo(1),
+                () -> assertThat(metric.getAliveAvailableCpu()).isEqualTo(1),
+                () -> assertThat(metric.getAliveTotalGpu()).isEqualTo(1),
+                () -> assertThat(metric.getAliveAvailableGpu()).isEqualTo(1),
+                () -> assertThat(metric.getCurrentTaskStatusCounts()).isEqualTo(expectedCurrentTaskStatusesCount),
+                () -> assertThat(metric.getDealEventsCount()).isEqualTo(10),
+                () -> assertThat(metric.getDealsCount()).isEqualTo(8),
+                () -> assertThat(metric.getReplayDealsCount()).isEqualTo(2),
+                () -> assertThat(metric.getLatestBlockNumberWithDeal()).isEqualTo(255)
+        );
+    }
+
+    private LinkedHashMap<TaskStatus, AtomicLong> createExpectedCurrentTaskStatusesCount() {
+        final LinkedHashMap<TaskStatus, AtomicLong> expectedCurrentTaskStatusesCount = new LinkedHashMap<>(TaskStatus.values().length);
+        expectedCurrentTaskStatusesCount.put(TaskStatus.RECEIVED, new AtomicLong(1));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.INITIALIZING, new AtomicLong(2));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.INITIALIZED, new AtomicLong(3));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.INITIALIZE_FAILED, new AtomicLong(4));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.RUNNING, new AtomicLong(5));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.RUNNING_FAILED, new AtomicLong(6));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.CONTRIBUTION_TIMEOUT, new AtomicLong(7));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.CONSENSUS_REACHED, new AtomicLong(8));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.REOPENING, new AtomicLong(9));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.REOPENED, new AtomicLong(10));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.REOPEN_FAILED, new AtomicLong(11));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.AT_LEAST_ONE_REVEALED, new AtomicLong(12));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.RESULT_UPLOADING, new AtomicLong(13));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.RESULT_UPLOADED, new AtomicLong(14));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.RESULT_UPLOAD_TIMEOUT, new AtomicLong(15));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.FINALIZING, new AtomicLong(16));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.FINALIZED, new AtomicLong(17));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.FINALIZE_FAILED, new AtomicLong(18));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.FINAL_DEADLINE_REACHED, new AtomicLong(19));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.COMPLETED, new AtomicLong(20));
+        expectedCurrentTaskStatusesCount.put(TaskStatus.FAILED, new AtomicLong(21));
+        metricService.onTaskStatusesCountUpdateEvent(new TaskStatusesCountUpdatedEvent(expectedCurrentTaskStatusesCount));
+
+        return expectedCurrentTaskStatusesCount;
     }
 
 }

src/test/java/com/iexec/core/task/update/TaskUpdateManagerTest.java

@@ -39,6 +39,7 @@
 import com.iexec.core.task.TaskService;
 import com.iexec.core.task.TaskStatus;
 import com.iexec.core.task.event.PleaseUploadEvent;
+import com.iexec.core.task.event.TaskStatusesCountUpdatedEvent;
 import com.iexec.core.worker.Worker;
 import com.iexec.core.worker.WorkerService;
 import io.micrometer.core.instrument.Gauge;
@@ -53,15 +54,14 @@
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 import org.springframework.context.ApplicationEventPublisher;
+import org.springframework.test.util.ReflectionTestUtils;
 
 import java.math.BigInteger;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
-import java.util.ArrayList;
-import java.util.Date;
-import java.util.List;
-import java.util.Optional;
+import java.util.*;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 
 import static com.iexec.core.task.TaskStatus.*;
@@ -128,7 +128,7 @@ void afterEach() {
 
     // region init
     @Test
-    void shouldBuildGauges() throws ExecutionException, InterruptedException {
+    void shouldBuildGaugesAndFireEvent() throws ExecutionException, InterruptedException {
         for (final TaskStatus status : TaskStatus.values()) {
             // Give a unique initial count for each status
             when(taskService.countByCurrentStatus(status)).thenReturn((long) status.ordinal());
@@ -143,6 +143,8 @@ void shouldBuildGauges() throws ExecutionException, InterruptedException {
                     .extracting(Gauge::value)
                     .isEqualTo(((double) status.ordinal()));
         }
+
+        verify(applicationEventPublisher, times(1)).publishEvent(any(TaskStatusesCountUpdatedEvent.class));
     }
     // endregion
 
@@ -2034,6 +2036,23 @@ void shouldUpdateMetricsAfterStatusUpdate() throws ExecutionException, Interrupt
 
         assertThat(currentReceivedTasks.value()).isZero();
         assertThat(currentInitializingTasks.value()).isOne();
+        // Called a first time during init, then a second time during update
+        verify(applicationEventPublisher, times(2)).publishEvent(any(TaskStatusesCountUpdatedEvent.class));
+    }
+    // endregion
+
+    // region onTaskCreatedEvent
+    @Test
+    void shouldUpdateCurrentReceivedCountAndFireEvent() {
+        final AtomicLong receivedCount =
+                (AtomicLong) ((LinkedHashMap<?, ?>) ReflectionTestUtils.getField(taskUpdateManager, "currentTaskStatusesCount"))
+                        .get(RECEIVED);
+        final long initialCount = receivedCount.get();
+
+        taskUpdateManager.onTaskCreatedEvent();
+
+        assertThat(receivedCount.get() - initialCount).isOne();
+        verify(applicationEventPublisher, times(1)).publishEvent(any(TaskStatusesCountUpdatedEvent.class));
     }
     // endregion