Add additional metrics for user/client/session count (#31)

src/main/java/io/kokuwa/keycloak/metrics/stats/MetricsStatsFactory.java

@@ -0,0 +1,10 @@
+package io.kokuwa.keycloak.metrics.stats;
+
+import org.keycloak.provider.ProviderFactory;
+
+/**
+ * Factory for Keycloak metrics.
+ *
+ * @author Stephan Schnabel
+ */
+public interface MetricsStatsFactory extends ProviderFactory<MetricsStatsTask> {}

src/main/java/io/kokuwa/keycloak/metrics/stats/MetricsStatsFactoryImpl.java

@@ -0,0 +1,69 @@
+package io.kokuwa.keycloak.metrics.stats;
+
+import java.time.Duration;
+import java.util.Optional;
+
+import org.jboss.logging.Logger;
+import org.keycloak.Config.Scope;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.models.utils.KeycloakModelUtils;
+import org.keycloak.timer.TimerProvider;
+
+import io.micrometer.core.instrument.Metrics;
+
+/**
+ * Implementation of {@link MetricsStatsFactory}.
+ *
+ * @author Stephan Schnabel
+ */
+public class MetricsStatsFactoryImpl implements MetricsStatsFactory {
+
+	private static final Logger log = Logger.getLogger(MetricsStatsFactory.class);
+
+	@Override
+	public String getId() {
+		return "default";
+	}
+
+	@Override
+	public void init(Scope config) {}
+
+	@Override
+	public void postInit(KeycloakSessionFactory factory) {
+
+		if (!"true".equals(System.getenv().get("KC_METRICS_STATS_ENABLED"))) {
+			log.infov("Keycloak stats not enabled.");
+			return;
+		}
+
+		var intervalDuration = Optional
+				.ofNullable(System.getenv("KC_METRICS_STATS_INTERVAL"))
+				.map(Duration::parse)
+				.orElse(Duration.ofSeconds(60));
+		var infoThreshold = Optional
+				.ofNullable(System.getenv("KC_METRICS_STATS_INFO_THRESHOLD"))
+				.map(Duration::parse)
+				.orElse(Duration.ofMillis(Double.valueOf(intervalDuration.toMillis() * 0.5).longValue()));
+		var warnThreshold = Optional
+				.ofNullable(System.getenv("KC_METRICS_STATS_WARN_THRESHOLD"))
+				.map(Duration::parse)
+				.orElse(Duration.ofMillis(Double.valueOf(intervalDuration.toMillis() * 0.75).longValue()));
+		log.infov("Keycloak stats enabled with interval of {0} and info/warn after {1}/{2}.",
+				intervalDuration, infoThreshold, warnThreshold);
+
+		var interval = intervalDuration.toMillis();
+		var task = new MetricsStatsTask(Metrics.globalRegistry, intervalDuration, infoThreshold, warnThreshold);
+		KeycloakModelUtils.runJobInTransaction(factory, session -> session
+				.getProvider(TimerProvider.class)
+				.schedule(() -> KeycloakModelUtils.runJobInTransaction(factory, task), interval, "metrics"));
+	}
+
+	@Override
+	public MetricsStatsTask create(KeycloakSession session) {
+		return null;
+	}
+
+	@Override
+	public void close() {}
+}

src/main/java/io/kokuwa/keycloak/metrics/stats/MetricsStatsSpi.java

@@ -0,0 +1,34 @@
+package io.kokuwa.keycloak.metrics.stats;
+
+import org.keycloak.provider.Provider;
+import org.keycloak.provider.ProviderFactory;
+import org.keycloak.provider.Spi;
+
+/**
+ * SPI for Keycloak metrics.
+ *
+ * @author Stephan Schnabel
+ */
+public class MetricsStatsSpi implements Spi {
+
+	@Override
+	public boolean isInternal() {
+		return false;
+	}
+
+	@Override
+	public String getName() {
+		return "metrics";
+	}
+
+	@Override
+	public Class<? extends Provider> getProviderClass() {
+		return MetricsStatsTask.class;
+	}
+
+	@Override
+	public Class<? extends ProviderFactory<? extends Provider>> getProviderFactoryClass() {
+		// this must be an interface, otherwise spi will be silenty ignored
+		return MetricsStatsFactory.class;
+	}
+}

src/main/java/io/kokuwa/keycloak/metrics/stats/MetricsStatsTask.java

@@ -0,0 +1,89 @@
+package io.kokuwa.keycloak.metrics.stats;
+
+import java.time.Duration;
+import java.time.Instant;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+
+import org.jboss.logging.Logger;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakSessionTask;
+import org.keycloak.provider.Provider;
+
+import io.micrometer.core.instrument.MeterRegistry;
+import io.micrometer.core.instrument.Tag;
+
+/**
+ * Keycloak metrics.
+ *
+ * @author Stephan Schnabel
+ */
+public class MetricsStatsTask implements Provider, KeycloakSessionTask {
+
+	private static final Logger log = Logger.getLogger(MetricsStatsTask.class);
+	private final Map<String, AtomicLong> values = new HashMap<>();
+	private final MeterRegistry registry;
+	private final Duration interval;
+	private final Duration infoThreshold;
+	private final Duration warnThreshold;
+
+	MetricsStatsTask(MeterRegistry registry, Duration interval, Duration infoThreshold, Duration warnThreshold) {
+		this.registry = registry;
+		this.interval = interval;
+		this.infoThreshold = infoThreshold;
+		this.warnThreshold = warnThreshold;
+	}
+
+	@Override
+	public void run(KeycloakSession session) {
+		log.tracev("Triggered metrics stats task.");
+		var start = Instant.now();
+
+		try {
+			scrape(session);
+		} catch (Exception e) {
+			if (e instanceof org.hibernate.exception.SQLGrammarException) {
+				log.infov("Metrics status task skipped, database not ready");
+			} else {
+				log.errorv(e, "Failed to scrape stats.");
+			}
+			return;
+		}
+
+		var duration = Duration.between(start, Instant.now());
+		if (duration.compareTo(interval) > 0) {
+			log.errorv("Finished scrapping keycloak stats in {0}, consider to increase interval", duration);
+		} else if (duration.compareTo(warnThreshold) > 0) {
+			log.warnv("Finished scrapping keycloak stats in {0}, consider to increase interval", duration);
+		} else if (duration.compareTo(infoThreshold) > 0) {
+			log.infov("Finished scrapping keycloak stats in {0}", duration);
+		} else {
+			log.debugv("Finished scrapping keycloak stats in {0}", duration);
+		}
+	}
+
+	@Override
+	public void close() {}
+
+	private void scrape(KeycloakSession session) {
+		session.realms().getRealmsStream().forEach(realm -> {
+			var tagRealm = Tag.of("realm", realm.getName());
+			gauge("keycloak_users", Set.of(tagRealm), session.users().getUsersCount(realm));
+			gauge("keycloak_clients", Set.of(tagRealm), session.clients().getClientsCount(realm));
+			var sessions = session.sessions();
+			var activeSessions = sessions.getActiveClientSessionStats(realm, false);
+			realm.getClientsStream().forEach(client -> {
+				var tags = Set.of(tagRealm, Tag.of("client", client.getClientId()));
+				gauge("keycloak_offline_sessions", tags, sessions.getOfflineSessionsCount(realm, client));
+				gauge("keycloak_active_user_sessions", tags, sessions.getActiveUserSessions(realm, client));
+				gauge("keycloak_active_client_sessions", tags, activeSessions.getOrDefault(client.getId(), 0L));
+			});
+		});
+	}
+
+	private void gauge(String name, Set<Tag> tags, long value) {
+		values.computeIfAbsent(name + tags, s -> registry.gauge(name, tags, new AtomicLong())).set(value);
+	}
+}

src/main/resources/META-INF/services/io.kokuwa.keycloak.metrics.stats.MetricsStatsFactory

@@ -0,0 +1 @@
+io.kokuwa.keycloak.metrics.stats.MetricsStatsFactoryImpl

src/main/resources/META-INF/services/org.keycloak.provider.Spi

@@ -0,0 +1 @@
+io.kokuwa.keycloak.metrics.stats.MetricsStatsSpi