Skip to content

Commit

Permalink
[grid] Adding option to detect if VNC is available or not
Browse files Browse the repository at this point in the history
  • Loading branch information
@diemol
diemol committed Jun 11, 2021
1 parent 2cf47f3 commit 9717d40
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 16 deletions.
9 changes: 9 additions & 0 deletions java/server/src/org/openqa/selenium/grid/node/config/NodeFlags.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_REGISTER_CYCLE;
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_REGISTER_PERIOD;
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_SESSION_TIMEOUT;
import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_VNC_ENV_VAR;
import static org.openqa.selenium.grid.node.config.NodeOptions.NODE_SECTION;
import static org.openqa.selenium.grid.node.config.NodeOptions.OVERRIDE_MAX_SESSIONS;

Expand Down Expand Up @@ -155,6 +156,14 @@ public class NodeFlags implements HasRoles {
@ConfigValue(section = NODE_SECTION, name = "heartbeat-period", example = "10")
public int heartbeatPeriod = DEFAULT_HEARTBEAT_PERIOD;

@Parameter(
names = "--vnc-env-var",
hidden = true,
description = "Environment variable to check in order to determine if a vnc stream is " +
"available or not.")
@ConfigValue(section = NODE_SECTION, name = "vnc-env-var", example = "START_XVFB")
public String vncEnvVar = DEFAULT_VNC_ENV_VAR;

@Override
public Set<Role> getRoles() {
return Collections.singleton(NODE_ROLE);
Expand Down
52 changes: 36 additions & 16 deletions java/server/src/org/openqa/selenium/grid/node/config/NodeOptions.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,15 @@

package org.openqa.selenium.grid.node.config;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Multimap;

import org.openqa.selenium.Capabilities;
import org.openqa.selenium.ImmutableCapabilities;
import org.openqa.selenium.PersistentCapabilities;
import org.openqa.selenium.Platform;
import org.openqa.selenium.SessionNotCreatedException;
import org.openqa.selenium.WebDriver;
Expand All @@ -50,6 +51,7 @@
import java.util.Map;
import java.util.Optional;
import java.util.ServiceLoader;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Function;
import java.util.logging.Level;
import java.util.logging.Logger;
Expand All @@ -65,18 +67,20 @@ public class NodeOptions {
static final String NODE_SECTION = "node";
static final boolean DEFAULT_DETECT_DRIVERS = true;
static final boolean OVERRIDE_MAX_SESSIONS = false;
static final String DEFAULT_VNC_ENV_VAR = "START_XVFB";
static final int DEFAULT_REGISTER_CYCLE = 10;
static final int DEFAULT_REGISTER_PERIOD = 120;

private static final Logger LOG = Logger.getLogger(NodeOptions.class.getName());
private static final Json JSON = new Json();
private static final String DEFAULT_IMPL = "org.openqa.selenium.grid.node.local.LocalNodeFactory";
private static final ImmutableCapabilities CURRENT_PLATFORM =
new ImmutableCapabilities("platformName", Platform.getCurrent());
private static final Platform CURRENT_PLATFORM = Platform.getCurrent();
private static final ImmutableSet<String>
SINGLE_SESSION_DRIVERS = ImmutableSet.of("safari", "safari technology preview");

private final Config config;
private final AtomicBoolean vncEnabled = new AtomicBoolean();
private final AtomicBoolean vncEnabledValueSet = new AtomicBoolean();

public NodeOptions(Config config) {
this.config = Require.nonNull("Config", config);
Expand Down Expand Up @@ -179,15 +183,24 @@ public Duration getSessionTimeout() {
return Duration.ofSeconds(seconds);
}

@VisibleForTesting
boolean isVncEnabled() {
String vncEnvVar = config.get(NODE_SECTION, "vnc-env-var").orElse(DEFAULT_VNC_ENV_VAR);
if (!vncEnabledValueSet.getAndSet(true)) {
vncEnabled.set(Boolean.parseBoolean(System.getenv(vncEnvVar)));
}
return vncEnabled.get();
}

private void addDriverFactoriesFromConfig(ImmutableMultimap.Builder<Capabilities,
SessionFactory> sessionFactories) {
config.getAll(NODE_SECTION, "driver-factories").ifPresent(allConfigs -> {
if (allConfigs.size() % 2 != 0) {
throw new ConfigException("Expected each driver class to be mapped to a config");
}

Map<String, String> configMap = IntStream.range(0, allConfigs.size()/2).boxed()
.collect(Collectors.toMap(i -> allConfigs.get(2*i), i -> allConfigs.get(2*i + 1)));
Map<String, String> configMap = IntStream.range(0, allConfigs.size() / 2).boxed()
.collect(Collectors.toMap(i -> allConfigs.get(2 * i), i -> allConfigs.get(2 * i + 1)));

configMap.forEach((clazz, config) -> {
Capabilities stereotype = JSON.toType(config, Capabilities.class);
Expand Down Expand Up @@ -271,7 +284,8 @@ private void addDriverConfigs(
if (!configMap.containsKey("stereotype")) {
throw new ConfigException("Driver config is missing stereotype value. " + configMap);
}
Capabilities stereotype = JSON.toType(configMap.get("stereotype"), Capabilities.class);
Capabilities stereotype =
enhanceStereotype(JSON.toType(configMap.get("stereotype"), Capabilities.class));
String configName = configMap.getOrDefault("name", "Custom Slot Config");
int driverMaxSessions = Integer.parseInt(configMap.getOrDefault("max-sessions", "1"));
Require.positive("Driver max sessions", driverMaxSessions);
Expand Down Expand Up @@ -319,8 +333,7 @@ private void addDetectedDrivers(
.peek(this::report)
.forEach(
entry -> {
Capabilities capabilities = entry.getKey()
.getCanonicalCapabilities().merge(CURRENT_PLATFORM);
Capabilities capabilities = enhanceStereotype(entry.getKey().getCanonicalCapabilities());
sessionFactories.putAll(capabilities, entry.getValue());
});

Expand Down Expand Up @@ -372,8 +385,7 @@ private void addSpecificDrivers(
.peek(this::report)
.forEach(
entry -> {
Capabilities capabilities = entry.getKey()
.getCanonicalCapabilities().merge(CURRENT_PLATFORM);
Capabilities capabilities = enhanceStereotype(entry.getKey().getCanonicalCapabilities());
sessionFactories.putAll(capabilities, entry.getValue());
});
}
Expand All @@ -400,7 +412,7 @@ private Map<WebDriverInfo, Collection<SessionFactory>> discoverDrivers(

Multimap<WebDriverInfo, SessionFactory> toReturn = HashMultimap.create();
infos.forEach(info -> {
Capabilities caps = info.getCanonicalCapabilities().merge(CURRENT_PLATFORM);
Capabilities caps = enhanceStereotype(info.getCanonicalCapabilities());
builders.stream()
.filter(builder -> builder.score(caps) > 0)
.forEach(builder -> {
Expand Down Expand Up @@ -475,15 +487,23 @@ private int getDriverMaxSessions(WebDriverInfo info, int desiredMaxSessions) {
return Math.min(info.getMaximumSimultaneousSessions(), desiredMaxSessions);
}

private Capabilities enhanceStereotype(Capabilities capabilities) {
if (capabilities.getPlatformName() == null) {
capabilities = new PersistentCapabilities(capabilities)
.setCapability("platformName", CURRENT_PLATFORM);
}
if (isVncEnabled()) {
capabilities = new PersistentCapabilities(capabilities)
.setCapability("se:vncEnabled", true);
}
return capabilities;
}

private void report(Map.Entry<WebDriverInfo, Collection<SessionFactory>> entry) {
StringBuilder caps = new StringBuilder();
try (JsonOutput out = JSON.newOutput(caps)) {
out.setPrettyPrint(false);
Capabilities capabilities = entry.getKey().getCanonicalCapabilities();
if (capabilities.getPlatformName() == null) {
capabilities = capabilities.merge(CURRENT_PLATFORM);
}
out.write(capabilities);
out.write(entry.getKey().getCanonicalCapabilities());
}

LOG.info(String.format(
Expand Down

0 comments on commit 9717d40

Please sign in to comment.