Add counter with WebSocket close codes

This commit is contained in:
Chris Eager 2025-02-07 10:10:05 -06:00 committed by ravi-signal
parent a9975e524b
commit 2dfd17af4a
2 changed files with 28 additions and 13 deletions

View File

@ -1,19 +1,23 @@
package org.whispersystems.textsecuregcm.metrics;
import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
import io.micrometer.core.instrument.Metrics;
import io.micrometer.core.instrument.Tag;
import io.micrometer.core.instrument.Tags;
import io.micrometer.core.instrument.Timer;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import org.whispersystems.textsecuregcm.util.EnumMapUtil;
import org.whispersystems.textsecuregcm.util.ua.ClientPlatform;
import org.whispersystems.textsecuregcm.util.ua.UnrecognizedUserAgentException;
import org.whispersystems.textsecuregcm.util.ua.UserAgentUtil;
import org.whispersystems.websocket.session.WebSocketSessionContext;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
public class OpenWebSocketCounter {
private static final String WEBSOCKET_CLOSED_COUNTER_NAME = name(OpenWebSocketCounter.class, "websocketClosed");
private final Map<ClientPlatform, AtomicInteger> openWebsocketsByClientPlatform;
private final AtomicInteger openWebsocketsFromUnknownPlatforms;
@ -81,6 +85,9 @@ public class OpenWebSocketCounter {
context.addWebsocketClosedListener((context1, statusCode, reason) -> {
sample.stop(durationTimer);
openWebSocketCounter.decrementAndGet();
Metrics.counter(WEBSOCKET_CLOSED_COUNTER_NAME, "status", String.valueOf(statusCode))
.increment();
});
}
}

View File

@ -17,11 +17,12 @@ import jakarta.ws.rs.core.Response;
import java.net.URI;
import java.util.List;
import java.util.Map;
import org.eclipse.jetty.util.component.LifeCycle;
import org.eclipse.jetty.websocket.api.Session;
import org.eclipse.jetty.websocket.api.StatusCode;
import org.eclipse.jetty.websocket.client.WebSocketClient;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.whispersystems.textsecuregcm.metrics.NoopAwsSdkMetricPublisher;
@ -47,27 +48,22 @@ class WhisperServerServiceTest {
Resources.getResource("config/test-secrets-bundle.yml").getPath());
}
private static final WebSocketClient webSocketClient = new WebSocketClient();
private static final DropwizardAppExtension<WhisperServerConfiguration> EXTENSION = new DropwizardAppExtension<>(
WhisperServerService.class, Resources.getResource("config/test.yml").getPath());
private WebSocketClient webSocketClient;
@AfterAll
static void teardown() {
System.clearProperty("secrets.bundle.filename");
}
@BeforeEach
void setUp() throws Exception {
webSocketClient = new WebSocketClient();
@BeforeAll
static void setUp() throws Exception {
webSocketClient.start();
}
@AfterEach
void tearDown() throws Exception {
webSocketClient.stop();
}
@Test
void start() throws Exception {
// make sure the service nominally starts and responds to health checks
@ -96,6 +92,16 @@ class WhisperServerServiceTest {
final TestWebsocketListener testWebsocketListener = new TestWebsocketListener();
EXTENSION.getTestSupport().getEnvironment().getApplicationContext().getServer()
.addEventListener(new LifeCycle.Listener() {
@Override
public void lifeCycleStopped(final LifeCycle event) {
// closed by org.eclipse.jetty.websocket.common.SessionTracker during the container Lifecycle stopping phase
assertEquals(StatusCode.SHUTDOWN, testWebsocketListener.closeFuture().getNow(-1));
}
});
// Session is Closeable, but we intentionally keep it open so that we can confirm the container Lifecycle behavior
final Session session = webSocketClient.connect(testWebsocketListener,
URI.create(String.format("ws://localhost:%d/v1/websocket/", EXTENSION.getLocalPort())))
.join();
@ -112,6 +118,8 @@ class WhisperServerServiceTest {
assertEquals(401, whoami.getStatus());
final long whoamiTimestamp = Long.parseLong(whoami.getHeaders().get(HeaderUtils.TIMESTAMP_HEADER.toLowerCase()));
assertTrue(whoamiTimestamp >= start);
}
@Test