Calculate running sessions per user

config/src/main/java/ru/qatools/gridrouter/config/WithRoutesMap.java

@@ -1,17 +1,24 @@
 package ru.qatools.gridrouter.config;
 
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 /**
  * @author Dmitry Baev charlie@yandex-team.ru
  * @author Innokenty Shuvalov innokenty@yandex-team.ru
  */
-public interface WithRoutesMap extends WithHosts {
+public interface WithRoutesMap {
+
+    List<Browser> getBrowsers();
 
     default Map<String, String> getRoutesMap() {
-        Map<String, String> routes = new HashMap<>();
-        getHosts().forEach(h -> routes.put(h.getRouteId(), h.getRoute()));
+        HashMap<String, String> routes = new HashMap<>();
+        getBrowsers().stream()
+                .flatMap(b -> b.getVersions().stream())
+                .flatMap(v -> v.getRegions().stream())
+                .flatMap(r -> r.getHosts().stream())
+                .forEach(h -> routes.put(h.getRouteId(), h.getRoute()));
         return routes;
     }
 }

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

@@ -76,8 +76,8 @@ public Map<String, Browsers> getQuotaMap() {
         return userBrowsers;
     }
 
-    public Map<String, String> getRoutes() {
-        return routes;
+    public String getRoute(String routeId) {
+        return routes.get(routeId);
     }
 
     public Version findVersion(String user, JsonCapabilities caps) {

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

@@ -0,0 +1,70 @@
+package ru.qatools.gridrouter;
+
+import org.apache.http.client.utils.URIBuilder;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.servlet.http.HttpServletRequest;
+import java.io.UnsupportedEncodingException;
+import java.net.URISyntaxException;
+import java.net.URLDecoder;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.springframework.http.HttpMethod.DELETE;
+
+/**
+ * @author Alexander Andyashin aandryashin@yandex-team.ru
+ * @author Innokenty Shuvalov innokenty@yandex-team.ru
+ * @author Dmitry Baev charlie@yandex-team.ru
+ * @author Artem Eroshenko eroshenkoam@yandex-team.ru
+ */
+public final class JsonWireUtils {
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(JsonWireUtils.class);
+
+    public static final String WD_HUB_SESSION = "/wd/hub/session/";
+
+    public static final int SESSION_HASH_LENGTH = 32;
+
+    private JsonWireUtils() {
+    }
+
+    public static boolean isUriValid(String uri) {
+        return uri.length() > getUriPrefixLength();
+    }
+
+    public static boolean isSessionDeleteRequest(HttpServletRequest request, String command) {
+        return DELETE.name().equalsIgnoreCase(request.getMethod()) && !command.contains("/");
+    }
+
+    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;
+    }
+
+    public static String redirectionUrl(String host, String command) throws URISyntaxException {
+        return new URIBuilder(host).setPath(WD_HUB_SESSION + command).build().toString();
+    }
+
+    public static String getCommand(String uri) {
+        String encodedCommand = uri.substring(getUriPrefixLength());
+        try {
+            return URLDecoder.decode(encodedCommand, UTF_8.name());
+        } catch (UnsupportedEncodingException e) {
+            LOGGER.error("[UNABLE_TO_DECODE_COMMAND] - could not decode command: {}", encodedCommand, e);
+            return encodedCommand;
+        }
+    }
+}

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

@@ -1,15 +1,14 @@
 package ru.qatools.gridrouter;
 
 import org.apache.commons.io.IOUtils;
-import org.apache.http.client.utils.URIBuilder;
 import org.eclipse.jetty.client.api.Request;
 import org.eclipse.jetty.client.util.StringContentProvider;
 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;
+import ru.qatools.gridrouter.sessions.SessionStorage;
 
 import javax.servlet.ServletConfig;
 import javax.servlet.ServletException;
@@ -18,13 +17,10 @@
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import java.io.IOException;
-import java.io.UnsupportedEncodingException;
-import java.net.URISyntaxException;
-import java.net.URLDecoder;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
-import static org.springframework.http.HttpMethod.DELETE;
 import static org.springframework.web.context.support.SpringBeanAutowiringSupport.processInjectionBasedOnServletContext;
+import static ru.qatools.gridrouter.JsonWireUtils.*;
 import static ru.qatools.gridrouter.RequestUtils.getRemoteHost;
 
 /**
@@ -34,7 +30,7 @@
  * @author Artem Eroshenko eroshenkoam@yandex-team.ru
  */
 @WebServlet(
-        urlPatterns = {ProxyServlet.WD_HUB_SESSION + "*"},
+        urlPatterns = {WD_HUB_SESSION + "*"},
         asyncSupported = true,
         initParams = {
                 @WebInitParam(name = "timeout", value = "300000"),
@@ -45,15 +41,11 @@ public class ProxyServlet extends org.eclipse.jetty.proxy.ProxyServlet {
 
     private static final Logger LOGGER = LoggerFactory.getLogger(ProxyServlet.class);
 
-    public static final String WD_HUB_SESSION = "/wd/hub/session/";
-
-    public static final int SESSION_HASH_LENGTH = 32;
-
     @Autowired
     private ConfigRepository config;
 
     @Autowired
-    private GridStats stats;
+    private SessionStorage sessionStorage;
 
     @Override
     public void init(ServletConfig config) throws ServletException {
@@ -85,7 +77,7 @@ protected String rewriteTarget(HttpServletRequest request) {
             return null;
         }
 
-        String route = getRoute(uri);
+        String route = config.getRoute(getSessionHash(uri));
         String command = getCommand(uri);
 
         if (route == null) {
@@ -95,7 +87,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 {
@@ -132,39 +126,4 @@ private String removeSessionIdSafe(String content, String remoteHost) {
         }
         return content;
     }
-
-    protected String redirectionUrl(String host, String command) throws URISyntaxException {
-        return new URIBuilder(host).setPath(WD_HUB_SESSION + command).build().toString();
-    }
-
-    protected String getRoute(String uri) {
-        return config.getRoutes().get(getSessionHash(uri));
-    }
-
-    protected String getCommand(String uri) {
-        String encodedCommand = uri.substring(getUriPrefixLength());
-        try {
-            return URLDecoder.decode(encodedCommand, UTF_8.name());
-        } catch (UnsupportedEncodingException e) {
-            LOGGER.error("[UNABLE_TO_DECODE_COMMAND] - could not decode command: {}", encodedCommand, e);
-            return encodedCommand;
-        }
-    }
-
-    protected boolean isUriValid(String uri) {
-        return uri.length() > getUriPrefixLength();
-    }
-
-    protected boolean isSessionDeleteRequest(HttpServletRequest request, String command) {
-        return DELETE.name().equalsIgnoreCase(request.getMethod())
-                && !command.contains("/");
-    }
-
-    protected String getSessionHash(String uri) {
-        return uri.substring(WD_HUB_SESSION.length(), getUriPrefixLength());
-    }
-
-    protected int getUriPrefixLength() {
-        return WD_HUB_SESSION.length() + SESSION_HASH_LENGTH;
-    }
 }

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

@@ -17,10 +17,10 @@
 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;
+import ru.qatools.gridrouter.sessions.SessionStorage;
 
 import javax.servlet.ServletConfig;
 import javax.servlet.ServletException;
@@ -64,7 +64,7 @@ public class RouteServlet extends HttpServlet {
     private HostSelectionStrategy hostSelectionStrategy;
 
     @Autowired
-    private GridStats stats;
+    private SessionStorage sessionStorage;
 
     @Autowired
     private CapabilityProcessorFactory capabilityProcessorFactory;
@@ -103,6 +103,7 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response)
         List<Region> unvisitedRegions = new ArrayList<>(allRegions);
 
         int attempt = 0;
+        JsonMessage hubMessage = null;
         try (CloseableHttpClient client = newHttpClient()) {
             while (!allRegions.isEmpty()) {
                 attempt++;
@@ -116,15 +117,15 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response)
 
                     String target = route + request.getRequestURI();
                     HttpResponse hubResponse = client.execute(post(target, message));
-                    JsonMessage hubMessage = JsonMessageFactory.from(hubResponse.getEntity().getContent());
+                    hubMessage = JsonMessageFactory.from(hubResponse.getEntity().getContent());
 
                     if (hubResponse.getStatusLine().getStatusCode() == SC_OK) {
                         String sessionId = hubMessage.getSessionId();
                         hubMessage.setSessionId(host.getRouteId() + sessionId);
                         replyWithOk(hubMessage, response);
                         LOGGER.info("[SESSION_CREATED] [{}] [{}] [{}] [{}] [{}] [{}]",
                                 user, remoteHost, browser, route, sessionId, attempt);
-                        stats.startSession();
+                        sessionStorage.put(hubMessage.getSessionId(), user, browser, actualVersion.getNumber());
                         return;
                     }
                     LOGGER.warn("[SESSION_FAILED] [{}] [{}] [{}] [{}] - {}",
@@ -150,15 +151,23 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response)
         }
 
         LOGGER.error("[SESSION_NOT_CREATED] [{}] [{}] [{}]", user, remoteHost, browser);
-        replyWithError("Cannot create session on any available node", response);
+        if (hubMessage == null) {
+            replyWithError("Cannot create session on any available node", response);
+        } else {
+            replyWithError(hubMessage, response);
+        }
     }
 
     protected void replyWithOk(JsonMessage message, HttpServletResponse response) throws IOException {
         reply(SC_OK, message, response);
     }
 
     protected void replyWithError(String errorMessage, HttpServletResponse response) throws IOException {
-        reply(SC_INTERNAL_SERVER_ERROR, JsonMessageFactory.error(13, errorMessage), response);
+        replyWithError(JsonMessageFactory.error(13, errorMessage), response);
+    }
+
+    protected void replyWithError(JsonMessage message, HttpServletResponse response) throws IOException {
+        reply(SC_INTERNAL_SERVER_ERROR, message, response);
     }
 
     protected void reply(int code, JsonMessage message, HttpServletResponse response) throws IOException {

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

@@ -0,0 +1,29 @@
+package ru.qatools.gridrouter;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
+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;
+
+/**
+ * @author Innokenty Shuvalov innokenty@yandex-team.ru
+ */
+@Configuration
+@EnableScheduling
+public class SessionStorageEvictionScheduler {
+
+    @Value("${grid.router.evict.sessions.timeout.seconds}")
+    private int timeout;
+
+    @Autowired
+    private SessionStorage sessionStorage;
+
+    @Scheduled(cron = "${grid.router.evict.sessions.cron}")
+    public void expireOldSessions() {
+        sessionStorage.expireSessionsOlderThan(Duration.ofSeconds(timeout));
+    }
+}