calculate open sessions separately for each user

Author: innokenty

proxy/src/main/java/ru/qatools/gridrouter/JsonWireUtils.java

@@ -41,6 +41,15 @@ public static String getSessionHash(String uri) {
         return uri.substring(WD_HUB_SESSION.length(), getUriPrefixLength());
     }
 
+    public static String getFullSessionId(String uri) {
+        String tail = uri.substring(WD_HUB_SESSION.length());
+        int end = tail.indexOf('/');
+        if (end < 0) {
+            return tail;
+        }
+        return tail.substring(0, end);
+    }
+
     public static int getUriPrefixLength() {
         return WD_HUB_SESSION.length() + SESSION_HASH_LENGTH;
     }

proxy/src/main/java/ru/qatools/gridrouter/MemorySessionStorage.java

@@ -0,0 +1,57 @@
+package ru.qatools.gridrouter;
+
+import java.time.Duration;
+import java.time.temporal.Temporal;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static java.time.ZonedDateTime.now;
+import static java.util.stream.Collectors.toList;
+
+/**
+ * @author Innokenty Shuvalov [email protected]
+ */
+public class MemorySessionStorage implements SessionStorage {
+
+    private final Map<String, Temporal> session2instant = new HashMap<>();
+    private final Map<String, Integer> user2count      = new HashMap<>();
+    private final Map<String, String>  session2user    = new HashMap<>();
+
+    @Override
+    public synchronized void put(String sessionId, String user) {
+        if (session2instant.put(sessionId, now()) == null) {
+            user2count.compute(user, (k, count) -> Optional.ofNullable(count).orElse(0) + 1);
+            session2user.put(sessionId, user);
+        }
+    }
+
+    @Override
+    public void update(String sessionId) {
+        session2instant.replace(sessionId, now());
+    }
+
+    @Override
+    public synchronized void remove(String sessionId) {
+        if (session2instant.remove(sessionId) != null) {
+            String user = session2user.remove(sessionId);
+            user2count.compute(user, (k, count) -> count - 1);
+        }
+    }
+
+    @Override
+    public synchronized List<String> expireSessionsOlderThan(Duration duration) {
+        List<String> sessions2delete = session2instant.entrySet().stream()
+                .filter(e -> duration.compareTo(Duration.between(e.getValue(), now())) < 0)
+                .map(Map.Entry::getKey)
+                .collect(toList());
+        sessions2delete.stream().forEach(this::remove);
+        return sessions2delete;
+    }
+
+    @Override
+    public synchronized int getCountFor(String user) {
+        return user2count.getOrDefault(user, 0);
+    }
+}

proxy/src/main/java/ru/qatools/gridrouter/ProxyServlet.java

@@ -6,7 +6,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
-import ru.qatools.gridrouter.json.GridStats;
 import ru.qatools.gridrouter.json.JsonMessage;
 import ru.qatools.gridrouter.json.JsonMessageFactory;
 
@@ -45,7 +44,7 @@ public class ProxyServlet extends org.eclipse.jetty.proxy.ProxyServlet {
     private ConfigRepository config;
 
     @Autowired
-    private GridStats stats;
+    private SessionStorage sessionStorage;
 
     @Override
     public void init(ServletConfig config) throws ServletException {
@@ -87,7 +86,9 @@ protected String rewriteTarget(HttpServletRequest request) {
 
         if (isSessionDeleteRequest(request, command)) {
             LOGGER.info("[SESSION_DELETED] [{}] [{}] [{}]", remoteHost, route, command);
-            stats.stopSession();
+            sessionStorage.remove(getFullSessionId(uri));
+        } else {
+            sessionStorage.update(getFullSessionId(uri));
         }
 
         try {

proxy/src/main/java/ru/qatools/gridrouter/RouteServlet.java

@@ -17,7 +17,6 @@
 import ru.qatools.gridrouter.config.HostSelectionStrategy;
 import ru.qatools.gridrouter.config.Region;
 import ru.qatools.gridrouter.config.Version;
-import ru.qatools.gridrouter.json.GridStats;
 import ru.qatools.gridrouter.json.JsonCapabilities;
 import ru.qatools.gridrouter.json.JsonMessage;
 import ru.qatools.gridrouter.json.JsonMessageFactory;
@@ -64,7 +63,7 @@ public class RouteServlet extends HttpServlet {
     private HostSelectionStrategy hostSelectionStrategy;
 
     @Autowired
-    private GridStats stats;
+    private SessionStorage sessionStorage;
 
     @Autowired
     private CapabilityProcessorFactory capabilityProcessorFactory;
@@ -124,7 +123,7 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response)
                         replyWithOk(hubMessage, response);
                         LOGGER.info("[SESSION_CREATED] [{}] [{}] [{}] [{}] [{}] [{}]",
                                 user, remoteHost, browser, route, sessionId, attempt);
-                        stats.startSession();
+                        sessionStorage.put(hubMessage.getSessionId(), user);
                         return;
                     }
                     LOGGER.warn("[SESSION_FAILED] [{}] [{}] [{}] [{}] - {}",

proxy/src/main/java/ru/qatools/gridrouter/SessionStorage.java

@@ -0,0 +1,20 @@
+package ru.qatools.gridrouter;
+
+import java.time.Duration;
+import java.util.List;
+
+/**
+ * @author Innokenty Shuvalov [email protected]
+ */
+public interface SessionStorage {
+
+    void put(String sessionId, String user);
+
+    void update(String sessionId);
+
+    void remove(String sessionId);
+
+    List<String> expireSessionsOlderThan(Duration duration);
+
+    int getCountFor(String user);
+}

proxy/src/main/java/ru/qatools/gridrouter/StatsServlet.java

@@ -2,10 +2,11 @@
 
 import org.apache.commons.io.IOUtils;
 import org.springframework.beans.factory.annotation.Autowired;
-import ru.qatools.gridrouter.json.GridStats;
 
 import javax.servlet.ServletConfig;
 import javax.servlet.ServletException;
+import javax.servlet.annotation.HttpConstraint;
+import javax.servlet.annotation.ServletSecurity;
 import javax.servlet.annotation.WebServlet;
 import javax.servlet.http.HttpServlet;
 import javax.servlet.http.HttpServletRequest;
@@ -15,17 +16,18 @@
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
-import static org.springframework.http.MediaType.APPLICATION_JSON_VALUE;
+import static org.springframework.http.MediaType.TEXT_PLAIN_VALUE;
 import static org.springframework.web.context.support.SpringBeanAutowiringSupport.processInjectionBasedOnServletContext;
 
 /**
  * @author Dmitry Baev [email protected]
  */
 @WebServlet(urlPatterns = {"/stats"}, asyncSupported = true)
+@ServletSecurity(value = @HttpConstraint(rolesAllowed = {"user"}))
 public class StatsServlet extends HttpServlet {
 
     @Autowired
-    private GridStats stats;
+    private SessionStorage sessionStorage;
 
     @Override
     public void init(ServletConfig config) throws ServletException {
@@ -34,11 +36,13 @@ public void init(ServletConfig config) throws ServletException {
     }
 
     @Override
-    protected void doGet(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException {
-        resp.setStatus(SC_OK);
-        resp.setContentType(APPLICATION_JSON_VALUE);
-        try (OutputStream output = resp.getOutputStream()) {
-            IOUtils.write(stats.toJson(), output, UTF_8);
+    protected void doGet(HttpServletRequest request, HttpServletResponse response)
+            throws ServletException, IOException {
+        response.setStatus(SC_OK);
+        response.setContentType(TEXT_PLAIN_VALUE);
+        try (OutputStream output = response.getOutputStream()) {
+            int count = sessionStorage.getCountFor(request.getRemoteUser());
+            IOUtils.write(String.valueOf(count), output, UTF_8);
         }
     }
 }

proxy/src/main/java/ru/qatools/gridrouter/json/WithOpenSessions.java

@@ -3,10 +3,8 @@
        xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
        xmlns:context="http://www.springframework.org/schema/context"
        xsi:schemaLocation="
-      http://www.springframework.org/schema/beans
-      http://www.springframework.org/schema/beans/spring-beans-4.0.xsd
-      http://www.springframework.org/schema/context
-      http://www.springframework.org/schema/context/spring-context-4.0.xsd">
+            http://www.springframework.org/schema/beans http://www.springframework.org/schema/beans/spring-beans-4.0.xsd
+            http://www.springframework.org/schema/context http://www.springframework.org/schema/context/spring-context-4.0.xsd">
 
     <context:component-scan base-package="ru.qatools.gridrouter"/>
 
@@ -21,6 +19,6 @@
 
     <bean id="hostSelectionStrategy" class="ru.qatools.gridrouter.config.RandomHostSelectionStrategy"/>
 
-    <bean id="stats" class="ru.qatools.gridrouter.json.GridStats"/>
+    <bean class="ru.qatools.gridrouter.MemorySessionStorage"/>
 
 </beans>