Change prometheus selenium active sessions gauge to custom collector

src/main/java/de/zalando/ep/zalenium/prometheus/TestSessionCollectorExports.java

@@ -0,0 +1,43 @@
+package de.zalando.ep.zalenium.prometheus;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+
+import org.openqa.grid.internal.ProxySet;
+import org.openqa.grid.internal.RemoteProxy;
+
+import io.prometheus.client.Collector;
+import io.prometheus.client.GaugeMetricFamily;
+
+public class TestSessionCollectorExports extends Collector {
+
+    private ProxySet proxySet;
+
+    public TestSessionCollectorExports(ProxySet proxySet) {
+        super();
+        this.proxySet = proxySet;
+    }
+
+    @Override
+    public List<MetricFamilySamples> collect() {
+        GaugeMetricFamily testSessionMetric = new GaugeMetricFamily("selenium_test_sessions_running",
+                "The number of Selenium test sessions that are running by proxy type",
+                Collections.singletonList("proxy"));
+
+        Iterable<RemoteProxy> iterable = () -> proxySet.iterator();
+        Map<String, Integer> countByProxies = StreamSupport.stream(iterable.spliterator(), false).collect(
+                Collectors.groupingBy(p -> p.getClass().getSimpleName(), Collectors.summingInt(p -> p.getTotalUsed())));
+
+        countByProxies.entrySet().stream()
+                .forEach(e -> testSessionMetric.addMetric(Collections.singletonList(e.getKey()), e.getValue()));
+
+        List<MetricFamilySamples> mfs = new ArrayList<MetricFamilySamples>();
+        mfs.add(testSessionMetric);
+        return mfs;
+    }
+
+}

src/main/java/de/zalando/ep/zalenium/proxy/DockerSeleniumRemoteProxy.java

@@ -1,20 +1,17 @@
 package de.zalando.ep.zalenium.proxy;
 
-import com.google.common.annotations.VisibleForTesting;
-import com.google.gson.JsonElement;
-import com.google.gson.JsonObject;
-import com.google.gson.JsonParser;
-import de.zalando.ep.zalenium.container.ContainerClient;
-import de.zalando.ep.zalenium.container.ContainerClientRegistration;
-import de.zalando.ep.zalenium.container.ContainerFactory;
-import de.zalando.ep.zalenium.dashboard.DashboardCollection;
-import de.zalando.ep.zalenium.dashboard.TestInformation;
-import de.zalando.ep.zalenium.matcher.DockerSeleniumCapabilityMatcher;
-import de.zalando.ep.zalenium.matcher.ZaleniumCapabilityType;
-import de.zalando.ep.zalenium.util.CommonProxyUtilities;
-import de.zalando.ep.zalenium.util.Environment;
-import de.zalando.ep.zalenium.util.GoogleAnalyticsApi;
-import io.prometheus.client.Gauge;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
 
 import org.apache.commons.compress.archivers.tar.TarArchiveEntry;
 import org.apache.commons.compress.archivers.tar.TarArchiveInputStream;
@@ -35,21 +32,25 @@
 import org.openqa.selenium.Platform;
 import org.openqa.selenium.remote.CapabilityType;
 import org.openqa.selenium.remote.server.jmx.ManagedService;
-
-import javax.servlet.http.HttpServletRequest;
-import javax.servlet.http.HttpServletResponse;
-import java.io.File;
-import java.io.FileOutputStream;
-import java.io.IOException;
-import java.io.OutputStream;
-import java.util.Map;
-import java.util.Optional;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.atomic.AtomicBoolean;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.annotations.VisibleForTesting;
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonParser;
+
+import de.zalando.ep.zalenium.container.ContainerClient;
+import de.zalando.ep.zalenium.container.ContainerClientRegistration;
+import de.zalando.ep.zalenium.container.ContainerFactory;
+import de.zalando.ep.zalenium.dashboard.DashboardCollection;
+import de.zalando.ep.zalenium.dashboard.TestInformation;
+import de.zalando.ep.zalenium.matcher.DockerSeleniumCapabilityMatcher;
+import de.zalando.ep.zalenium.matcher.ZaleniumCapabilityType;
+import de.zalando.ep.zalenium.util.CommonProxyUtilities;
+import de.zalando.ep.zalenium.util.Environment;
+import de.zalando.ep.zalenium.util.GoogleAnalyticsApi;
+
 /*
     The implementation of this class was inspired on https://gist.github.com/krmahadevan/4649607
  */
@@ -95,9 +96,6 @@ public class DockerSeleniumRemoteProxy extends DefaultRemoteProxy {
 
     private static final ExecutorService EXECUTOR_SERVICE = Executors.newFixedThreadPool(5);
 
-    static final Gauge seleniumTestSessionsRunning = Gauge.build()
-            .name("selenium_test_sessions_running").help("The number of Selenium test sessions that are running in a container").register();
-
     public DockerSeleniumRemoteProxy(RegistrationRequest request, GridRegistry registry) {
         super(request, registry);
         try {
@@ -203,9 +201,6 @@ public synchronized TestSession getNewSession(Map<String, Object> requestedCapab
 
         if (!this.isBusy() && increaseCounter()) {
             TestSession testSession = createNewSession(requestedCapability);
-            if (testSession != null) {
-                seleniumTestSessionsRunning.inc();
-            }
             return testSession;
         }
         LOGGER.debug("{} No more sessions allowed", getContainerId());
@@ -371,7 +366,6 @@ public void afterSession(TestSession session) {
         } catch (Exception e) {
             LOGGER.warn(getContainerId() + " " + e.toString(), e);
         } finally {
-            seleniumTestSessionsRunning.dec();
             super.afterSession(session);
         }
     }

src/main/java/de/zalando/ep/zalenium/registry/ZaleniumRegistry.java

@@ -1,6 +1,7 @@
 package de.zalando.ep.zalenium.registry;
 
 import de.zalando.ep.zalenium.dashboard.Dashboard;
+import de.zalando.ep.zalenium.prometheus.TestSessionCollectorExports;
 import net.jcip.annotations.ThreadSafe;
 
 import org.openqa.grid.internal.ActiveTestSessions;
@@ -99,6 +100,8 @@ public ZaleniumRegistry(Hub hub) {
 
         proxies = new AutoStartProxySet(false, minContainers, maxContainers, timeToWaitToStart, waitForAvailableNodes, starter, Clock.systemDefaultZone());
         this.matcherThread.setUncaughtExceptionHandler(new UncaughtExceptionHandler());
+
+        new TestSessionCollectorExports(proxies).register();
     }
 
     /**