[grid] Allow nodes to be configured for particular browsers

SeleniumHQ · May 19, 2020 · aae2902 · aae2902
1 parent e706b23
commit aae2902
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 23 deletions.
diff --git a/java/server/src/org/openqa/selenium/grid/node/config/BUILD.bazel b/java/server/src/org/openqa/selenium/grid/node/config/BUILD.bazel
@@ -1,3 +1,5 @@
+load("@rules_jvm_external//:defs.bzl", "artifact")
+
 java_library(
     name = "config",
     srcs = glob(["*.java"]),
@@ -11,5 +13,6 @@ java_library(
         "//java/server/src/org/openqa/selenium/grid/data",
         "//java/server/src/org/openqa/selenium/grid/node",
         "//java/server/src/org/openqa/selenium/grid/node/local",
+        artifact("com.google.guava:guava"),
     ],
 )
diff --git a/java/server/src/org/openqa/selenium/grid/node/config/NodeOptions.java b/java/server/src/org/openqa/selenium/grid/node/config/NodeOptions.java
@@ -17,16 +17,20 @@
 
 package org.openqa.selenium.grid.node.config;
 
+import com.google.common.collect.HashMultimap;
 import org.openqa.selenium.Capabilities;
 import org.openqa.selenium.WebDriverInfo;
 import org.openqa.selenium.grid.config.Config;
+import org.openqa.selenium.grid.node.SessionFactory;
 import org.openqa.selenium.grid.node.local.LocalNode;
 import org.openqa.selenium.remote.http.HttpClient;
 import org.openqa.selenium.remote.service.DriverService;
 import org.openqa.selenium.remote.tracing.Tracer;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.ServiceLoader;
 import java.util.logging.Logger;
@@ -43,46 +47,65 @@ public NodeOptions(Config config) {
   }
 
   public void configure(Tracer tracer, HttpClient.Factory httpClientFactory, LocalNode.Builder node) {
+    int maxSessions = Math.min(
+      config.getInt("node", "max-concurrent-sessions").orElse(Runtime.getRuntime().availableProcessors()),
+      Runtime.getRuntime().availableProcessors());
+
+    Map<WebDriverInfo, Collection<SessionFactory>> allDrivers = discoverDrivers(tracer, httpClientFactory, maxSessions);
+
+    // If drivers have been specified, use those.
+    List<String> drivers = config.getAll("node", "drivers").orElse(new ArrayList<>()).stream()
+      .map(String::toLowerCase)
+      .collect(Collectors.toList());
+
+    if (!drivers.isEmpty()) {
+      allDrivers.entrySet().stream()
+        .filter(entry -> drivers.contains(entry.getKey().getDisplayName().toLowerCase()))
+        .peek(entry -> LOG.info(String.format("Adding %s for %s %d times", entry.getKey().getDisplayName(), entry.getKey().getCanonicalCapabilities(), entry.getValue().size())))
+        .forEach(entry -> entry.getValue().forEach(factory -> node.add(entry.getKey().getCanonicalCapabilities(), factory)));
+
+      return;
+    }
+
     if (!config.getBool("node", "detect-drivers").orElse(false)) {
       return;
     }
 
-    addSystemDrivers(tracer, httpClientFactory, node);
+    allDrivers.entrySet().stream()
+      .peek(entry -> LOG.info(String.format("Adding %s for %s %d times", entry.getKey().getDisplayName(), entry.getKey().getCanonicalCapabilities(), entry.getValue().size())))
+      .forEach(entry -> entry.getValue().forEach(factory -> node.add(entry.getKey().getCanonicalCapabilities(), factory)));
   }
 
-
-  private void addSystemDrivers(
-      Tracer tracer,
-      HttpClient.Factory clientFactory,
-      LocalNode.Builder node) {
-
+  private Map<WebDriverInfo, Collection<SessionFactory>> discoverDrivers(
+    Tracer tracer,
+    HttpClient.Factory clientFactory,
+    int maxSessions) {
     // We don't expect duplicates, but they're fine
     List<WebDriverInfo> infos =
-        StreamSupport.stream(ServiceLoader.load(WebDriverInfo.class).spliterator(), false)
-            .filter(WebDriverInfo::isAvailable)
-            .collect(Collectors.toList());
+      StreamSupport.stream(ServiceLoader.load(WebDriverInfo.class).spliterator(), false)
+        .filter(WebDriverInfo::isAvailable)
+        .collect(Collectors.toList());
 
     // Same
     List<DriverService.Builder> builders = new ArrayList<>();
     ServiceLoader.load(DriverService.Builder.class).forEach(builders::add);
 
+    HashMultimap<WebDriverInfo, SessionFactory> toReturn = HashMultimap.create();
     infos.forEach(info -> {
       Capabilities caps = info.getCanonicalCapabilities();
       builders.stream()
-          .filter(builder -> builder.score(caps) > 0)
-          .peek(builder -> LOG.info(String.format("Adding %s %d times", caps, info.getMaximumSimultaneousSessions())))
-          .forEach(builder -> {
-            DriverService.Builder freePortBuilder = builder.usingAnyFreePort();
+        .filter(builder -> builder.score(caps) > 0)
+        .forEach(builder -> {
+          for (int i = 0; i < Math.min(info.getMaximumSimultaneousSessions(), maxSessions); i++) {
 
-            for (int i = 0; i < info.getMaximumSimultaneousSessions(); i++) {
-              node.add(
-                  caps,
-                  new DriverServiceSessionFactory(
-                      tracer,
-                      clientFactory, c -> freePortBuilder.score(c) > 0,
-                      freePortBuilder));
-            }
-          });
+            DriverService.Builder freePortBuilder = builder.usingAnyFreePort();
+            toReturn.put(info, new DriverServiceSessionFactory(
+              tracer,
+              clientFactory, c -> freePortBuilder.score(c) > 0,
+              freePortBuilder));
+          }
+        });
     });
+    return toReturn.asMap();
   }
 }
diff --git a/java/server/src/org/openqa/selenium/grid/node/httpd/NodeFlags.java b/java/server/src/org/openqa/selenium/grid/node/httpd/NodeFlags.java
@@ -25,6 +25,7 @@
 import org.openqa.selenium.grid.config.Role;
 
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.Set;
 
 import static org.openqa.selenium.grid.config.StandardGridRoles.NODE_ROLE;
@@ -38,6 +39,18 @@ public class NodeFlags implements HasRoles {
   @ConfigValue(section = "node", name = "detect-drivers")
   public Boolean autoconfigure = true;
 
+  @Parameter(
+    names = "--max-sessions",
+    description = "Maximum number of concurrent sessions.")
+  @ConfigValue(section = "node", name = "max-concurrent-sessions")
+  public int maxSessions = Runtime.getRuntime().availableProcessors();
+
+  @Parameter(
+    names = {"-I", "--driver-implementation"},
+    description = "Drivers that should be checked. If specified, will skip autoconfiguration. Example: -I \"firefox\" -I \"chrome\"")
+  @ConfigValue(section = "node", name = "drivers")
+  public Set<String> driverNames = new HashSet<>();
+
   @Override
   public Set<Role> getRoles() {
     return Collections.singleton(NODE_ROLE);