calculate sessions separately for each browser/version

Author: innokenty

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

@@ -8,6 +8,7 @@
 import org.springframework.beans.factory.annotation.Autowired;
 import ru.qatools.gridrouter.json.JsonMessage;
 import ru.qatools.gridrouter.json.JsonMessageFactory;
+import ru.qatools.gridrouter.sessions.SessionStorage;
 
 import javax.servlet.ServletConfig;
 import javax.servlet.ServletException;

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

@@ -20,6 +20,7 @@
 import ru.qatools.gridrouter.json.JsonCapabilities;
 import ru.qatools.gridrouter.json.JsonMessage;
 import ru.qatools.gridrouter.json.JsonMessageFactory;
+import ru.qatools.gridrouter.sessions.SessionStorage;
 
 import javax.servlet.ServletConfig;
 import javax.servlet.ServletException;
@@ -123,7 +124,7 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response)
                         replyWithOk(hubMessage, response);
                         LOGGER.info("[SESSION_CREATED] [{}] [{}] [{}] [{}] [{}] [{}]",
                                 user, remoteHost, browser, route, sessionId, attempt);
-                        sessionStorage.put(hubMessage.getSessionId(), user);
+                        sessionStorage.put(hubMessage.getSessionId(), user, browser, actualVersion.getNumber());
                         return;
                     }
                     LOGGER.warn("[SESSION_FAILED] [{}] [{}] [{}] [{}] - {}",

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

@@ -5,6 +5,7 @@
 import org.springframework.context.annotation.Configuration;
 import org.springframework.scheduling.annotation.EnableScheduling;
 import org.springframework.scheduling.annotation.Scheduled;
+import ru.qatools.gridrouter.sessions.SessionStorage;
 
 import java.time.Duration;
 

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

@@ -2,6 +2,8 @@
 
 import org.apache.commons.io.IOUtils;
 import org.springframework.beans.factory.annotation.Autowired;
+import ru.qatools.gridrouter.json.JsonFormatter;
+import ru.qatools.gridrouter.sessions.SessionStorage;
 
 import javax.servlet.ServletConfig;
 import javax.servlet.ServletException;
@@ -16,7 +18,7 @@
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
-import static org.springframework.http.MediaType.TEXT_PLAIN_VALUE;
+import static org.springframework.http.MediaType.APPLICATION_JSON_VALUE;
 import static org.springframework.web.context.support.SpringBeanAutowiringSupport.processInjectionBasedOnServletContext;
 
 /**
@@ -39,10 +41,11 @@ public void init(ServletConfig config) throws ServletException {
     protected void doGet(HttpServletRequest request, HttpServletResponse response)
             throws ServletException, IOException {
         response.setStatus(SC_OK);
-        response.setContentType(TEXT_PLAIN_VALUE);
+        response.setContentType(APPLICATION_JSON_VALUE);
         try (OutputStream output = response.getOutputStream()) {
-            int count = sessionStorage.getCountFor(request.getRemoteUser());
-            IOUtils.write(String.valueOf(count), output, UTF_8);
+            IOUtils.write(JsonFormatter.toJson(
+                    sessionStorage.getBrowsersCountFor(request.getRemoteUser())
+            ), output, UTF_8);
         }
     }
 }

proxy/src/main/java/ru/qatools/gridrouter/sessions/BrowserVersion.java

@@ -0,0 +1,23 @@
+package ru.qatools.gridrouter.sessions;
+
+/**
+ * @author Innokenty Shuvalov [email protected]
+ */
+public class BrowserVersion {
+
+    private final String browser;
+    private final String version;
+
+    public BrowserVersion(String browser, String version) {
+        this.browser = browser;
+        this.version = version;
+    }
+
+    public String getBrowser() {
+        return browser;
+    }
+
+    public String getVersion() {
+        return version;
+    }
+}

proxy/src/main/java/ru/qatools/gridrouter/sessions/BrowsersCountMap.java

@@ -0,0 +1,42 @@
+package ru.qatools.gridrouter.sessions;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
+/**
+ * @author Innokenty Shuvalov [email protected]
+ */
+public class BrowsersCountMap extends HashMap<String, Map<String, Integer>> {
+
+    public void increment(String browser, String version) {
+        putIfAbsent(browser, new HashMap<>());
+        get(browser).compute(version, (v, count) -> Optional.ofNullable(count).orElse(0) + 1);
+    }
+
+    public void decrement(BrowserVersion browser) {
+        decrement(browser.getBrowser(), browser.getVersion());
+    }
+
+    public void decrement(String browser, String version) {
+        if (!containsKey(browser)) {
+            return;
+        }
+
+        Map<String, Integer> versions = get(browser);
+        if (!versions.containsKey(version)) {
+            return;
+        }
+
+        int count = versions.get(version) - 1;
+        if (count > 0) {
+            versions.put(version, count);
+        } else {
+            versions.remove(version);
+        }
+
+        if (versions.isEmpty()) {
+            remove(browser);
+        }
+    }
+}

proxy/src/main/java/ru/qatools/gridrouter/MemorySessionStorage.java renamed to proxy/src/main/java/ru/qatools/gridrouter/sessions/MemorySessionStorage.java

@@ -1,11 +1,10 @@
-package ru.qatools.gridrouter;
+package ru.qatools.gridrouter.sessions;
 
 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;
@@ -16,14 +15,17 @@
 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<>();
+    private final Map<String, String> session2user = new HashMap<>();
+    private final Map<String, BrowserVersion> session2browserVersion = new HashMap<>();
+    private final Map<String, BrowsersCountMap> user2browserCount = new HashMap<>();
 
     @Override
-    public synchronized void put(String sessionId, String user) {
+    public synchronized void put(String sessionId, String user, String browser, String version) {
         if (session2instant.put(sessionId, now()) == null) {
-            user2count.compute(user, (k, count) -> Optional.ofNullable(count).orElse(0) + 1);
             session2user.put(sessionId, user);
+            session2browserVersion.put(sessionId, new BrowserVersion(browser, version));
+            user2browserCount.putIfAbsent(user, new BrowsersCountMap());
+            user2browserCount.get(user).increment(browser, version);
         }
     }
 
@@ -36,7 +38,8 @@ public void update(String sessionId) {
     public synchronized void remove(String sessionId) {
         if (session2instant.remove(sessionId) != null) {
             String user = session2user.remove(sessionId);
-            user2count.compute(user, (k, count) -> count - 1);
+            BrowserVersion browser = session2browserVersion.remove(sessionId);
+            user2browserCount.get(user).decrement(browser);
         }
     }
 
@@ -51,7 +54,7 @@ public synchronized List<String> expireSessionsOlderThan(Duration duration) {
     }
 
     @Override
-    public synchronized int getCountFor(String user) {
-        return user2count.getOrDefault(user, 0);
+    public synchronized BrowsersCountMap getBrowsersCountFor(String user) {
+        return user2browserCount.getOrDefault(user, new BrowsersCountMap());
     }
 }

proxy/src/main/java/ru/qatools/gridrouter/SessionStorage.java renamed to proxy/src/main/java/ru/qatools/gridrouter/sessions/SessionStorage.java

@@ -1,4 +1,4 @@
-package ru.qatools.gridrouter;
+package ru.qatools.gridrouter.sessions;
 
 import java.time.Duration;
 import java.util.List;
@@ -8,13 +8,13 @@
  */
 public interface SessionStorage {
 
-    void put(String sessionId, String user);
+    void put(String sessionId, String user, String browser, String version);
 
     void update(String sessionId);
 
     void remove(String sessionId);
 
     List<String> expireSessionsOlderThan(Duration duration);
 
-    int getCountFor(String user);
+    BrowsersCountMap getBrowsersCountFor(String user);
 }

proxy/src/main/resources/META-INF/spring/application-context.xml

@@ -19,6 +19,6 @@
 
     <bean id="hostSelectionStrategy" class="ru.qatools.gridrouter.config.RandomHostSelectionStrategy"/>
 
-    <bean class="ru.qatools.gridrouter.MemorySessionStorage"/>
+    <bean class="ru.qatools.gridrouter.sessions.MemorySessionStorage"/>
 
 </beans>