Add connection metrics per username (#75)

Author: antoinebhs

src/main/java/org/gridsuite/study/notification/server/NotificationWebSocketHandler.java

@@ -8,8 +8,9 @@
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import io.micrometer.core.instrument.Gauge;
 import io.micrometer.core.instrument.MeterRegistry;
+import io.micrometer.core.instrument.MultiGauge;
+import io.micrometer.core.instrument.Tags;
 import org.gridsuite.study.notification.server.dto.Filters;
 import org.gridsuite.study.notification.server.dto.FiltersToAdd;
 import org.gridsuite.study.notification.server.dto.FiltersToRemove;
@@ -38,6 +39,8 @@
 import java.util.function.Consumer;
 import java.util.logging.Level;
 
+import static java.util.stream.Collectors.toList;
+
 /**
  * A WebSocketHandler that sends messages from a broker to websockets opened by clients, interleaving with pings to keep connections open.
  * <p>
@@ -73,25 +76,21 @@ public class NotificationWebSocketHandler implements WebSocketHandler {
     static final String HEADER_INSERT_MODE = "insertMode";
     static final String HEADER_REFERENCE_NODE_UUID = "referenceNodeUuid";
     static final String HEADER_INDEXATION_STATUS = "indexation_status";
-
     static final String USERS_METER_NAME = "app.users";
-    static final String CONNECTIONS_METER_NAME = "app.connections";
+    static final String USER_TAG = "user";
 
     private final ObjectMapper jacksonObjectMapper;
 
     private final int heartbeatInterval;
 
     private final Map<String, Integer> userConnections = new ConcurrentHashMap<>();
 
+    private final MeterRegistry meterRegistry;
+
     public NotificationWebSocketHandler(ObjectMapper jacksonObjectMapper, MeterRegistry meterRegistry, @Value("${notification.websocket.heartbeat.interval:30}") int heartbeatInterval) {
         this.jacksonObjectMapper = jacksonObjectMapper;
         this.heartbeatInterval = heartbeatInterval;
-        initMetrics(meterRegistry);
-    }
-
-    private void initMetrics(MeterRegistry meterRegistry) {
-        Gauge.builder(USERS_METER_NAME, userConnections::size).register(meterRegistry);
-        Gauge.builder(CONNECTIONS_METER_NAME, () -> userConnections.values().stream().mapToInt(Integer::intValue).sum()).register(meterRegistry);
+        this.meterRegistry = meterRegistry;
     }
 
     Flux<Message<String>> flux;
@@ -235,11 +234,22 @@ private void updateConnectionMetrics(WebSocketSession webSocketSession) {
         LOGGER.info("New websocket connection id={} for user={} studyUuid={}, updateType={}", webSocketSession.getId(), userId,
                 webSocketSession.getAttributes().get(FILTER_STUDY_UUID), webSocketSession.getAttributes().get(FILTER_UPDATE_TYPE));
         userConnections.compute(userId, (k, v) -> (v == null) ? 1 : v + 1);
+        updateConnectionMetricsRegistry();
     }
 
     private void updateDisconnectionMetrics(WebSocketSession webSocketSession) {
         var userId = webSocketSession.getHandshakeInfo().getHeaders().getFirst(HEADER_USER_ID);
         LOGGER.info("Websocket disconnection id={} for user={}", webSocketSession.getId(), userId);
         userConnections.computeIfPresent(userId, (k, v) -> v > 1 ? v - 1 : null);
+        updateConnectionMetricsRegistry();
+    }
+
+    private void updateConnectionMetricsRegistry() {
+        MultiGauge.builder(USERS_METER_NAME)
+                .description("The current number of connections per user")
+                .register(meterRegistry)
+                .register(userConnections.entrySet().stream()
+                .map(e -> MultiGauge.Row.of(Tags.of(USER_TAG, e.getKey()), e.getValue()))
+                .collect(toList()), true);
     }
 }

src/test/java/org/gridsuite/study/notification/server/NotificationWebSocketIT.java

@@ -6,7 +6,6 @@
  */
 package org.gridsuite.study.notification.server;
 
-import io.micrometer.core.instrument.Gauge;
 import io.micrometer.core.instrument.MeterRegistry;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -22,10 +21,10 @@
 import reactor.core.publisher.Mono;
 
 import java.net.URI;
+import java.util.Map;
 
 import static org.gridsuite.study.notification.server.NotificationWebSocketHandler.*;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
 
 /**
  * @author Jon Harper <jon.harper at rte-france.com>
@@ -50,27 +49,47 @@ public void echo() {
         client.execute(getUrl("/notify"), httpHeaders, ws -> Mono.empty()).block();
     }
 
-    @Test
-    public void metrics() {
-        testMeters(0);
-        WebSocketClient client = new StandardWebSocketClient();
-        HttpHeaders httpHeaders = new HttpHeaders();
-        httpHeaders.add(HEADER_USER_ID, "test");
-        client.execute(getUrl("/notify"), httpHeaders, ws -> Mono.fromRunnable(() -> testMeters(1))).block();
-    }
-
     protected URI getUrl(String path) {
         return URI.create("ws://localhost:" + this.port + path);
     }
 
-    private void testMeters(int val) {
-        testMeter(USERS_METER_NAME, val);
-        testMeter(CONNECTIONS_METER_NAME, val);
+    @Test
+    public void metricsMapOneUserTwoConnections() {
+        WebSocketClient client1 = new StandardWebSocketClient();
+        HttpHeaders httpHeaders1 = new HttpHeaders();
+        String user = "test";
+        httpHeaders1.add(HEADER_USER_ID, user);
+        Map<String, Double> exp = Map.of(user, 2d);
+        Mono<Void> connection1 = client1.execute(getUrl("/notify"), httpHeaders1, ws -> Mono.fromRunnable(() -> testMeterMap(exp)));
+        Mono<Void> connection2 = client1.execute(getUrl("/notify"), httpHeaders1, ws -> Mono.fromRunnable(() -> testMeterMap(exp)));
+
+        Mono.zip(connection1, connection2).block();
+    }
+
+    @Test
+    public void metricsMapTwoUsers() {
+        // First WebSocketClient for connections related to 'test' user
+        WebSocketClient client1 = new StandardWebSocketClient();
+        HttpHeaders httpHeaders1 = new HttpHeaders();
+        String user1 = "test";
+        httpHeaders1.add(HEADER_USER_ID, user1);
+        String user2 = "test1";
+        Map<String, Double> exp = Map.of(user1, 2d, user2, 1d);
+        Mono<Void> connection1 = client1.execute(getUrl("/notify"), httpHeaders1, ws -> Mono.fromRunnable(() -> testMeterMap(exp)));
+        Mono<Void> connection2 = client1.execute(getUrl("/notify"), httpHeaders1, ws -> Mono.fromRunnable(() -> testMeterMap(exp)));
+
+        // Second WebSocketClient for connections related to 'test1' user
+        WebSocketClient client2 = new StandardWebSocketClient();
+        HttpHeaders httpHeaders2 = new HttpHeaders();
+        httpHeaders2.add(HEADER_USER_ID, user2);
+        Mono<Void> connection3 = client2.execute(getUrl("/notify"), httpHeaders2, ws -> Mono.fromRunnable(() -> testMeterMap(exp)));
+
+        Mono.zip(connection1, connection2, connection3).block();
     }
 
-    private void testMeter(String name, int val) {
-        Gauge meter = meterRegistry.get(name).gauge();
-        assertNotNull(meter);
-        assertEquals(val, Double.valueOf(meter.value()).intValue());
+    private void testMeterMap(Map<String, Double> userMap) {
+        for (Map.Entry<String, Double> userEntry : userMap.entrySet()) {
+            assertEquals(userEntry.getValue(), meterRegistry.get(USERS_METER_NAME).tag(USER_TAG, userEntry.getKey()).gauge().value(), 0);
+        }
     }
 }