...34import org.openqa.selenium.grid.data.NodeStatus;35import org.openqa.selenium.grid.data.Session;36import org.openqa.selenium.grid.data.SessionClosedEvent;37import org.openqa.selenium.grid.data.Slot;38import org.openqa.selenium.grid.data.SlotId;39import org.openqa.selenium.grid.log.LoggingOptions;40import org.openqa.selenium.grid.node.HealthCheck;41import org.openqa.selenium.grid.node.Node;42import org.openqa.selenium.grid.node.config.NodeOptions;43import org.openqa.selenium.grid.security.Secret;44import org.openqa.selenium.grid.server.BaseServerOptions;45import org.openqa.selenium.grid.server.EventBusOptions;46import org.openqa.selenium.internal.Require;47import org.openqa.selenium.json.Json;48import org.openqa.selenium.remote.CommandExecutor;49import org.openqa.selenium.remote.RemoteWebDriver;50import org.openqa.selenium.remote.SessionId;51import org.openqa.selenium.remote.http.HttpClient;52import org.openqa.selenium.remote.http.HttpRequest;53import org.openqa.selenium.remote.http.HttpResponse;54import org.openqa.selenium.remote.tracing.Tracer;55import java.lang.reflect.Field;56import java.net.URI;57import java.net.URISyntaxException;58import java.time.Instant;59import java.util.HashMap;60import java.util.Map;61import java.util.Optional;62import java.util.ServiceLoader;63import java.util.TreeMap;64import java.util.UUID;65import java.util.logging.Logger;66import java.util.stream.StreamSupport;67import static java.nio.charset.StandardCharsets.UTF_8;68import static org.openqa.selenium.grid.data.Availability.DRAINING;69import static org.openqa.selenium.grid.data.Availability.UP;70import static org.openqa.selenium.json.Json.MAP_TYPE;71import static org.openqa.selenium.remote.http.HttpMethod.DELETE;72/**73 * An implementation of {@link Node} that marks itself as draining immediately74 * after starting, and which then shuts down after usage. This will allow an75 * appropriately configured Kubernetes cluster to start a new node once the76 * session is finished.77 */78public class OneShotNode extends Node {79 private static final Logger LOG = Logger.getLogger(OneShotNode.class.getName());80 private static final Json JSON = new Json();81 private final EventBus events;82 private final WebDriverInfo driverInfo;83 private final Capabilities stereotype;84 private final Secret registrationSecret;85 private final URI gridUri;86 private final UUID slotId = UUID.randomUUID();87 private RemoteWebDriver driver;88 private SessionId sessionId;89 private HttpClient client;90 private Capabilities capabilities;91 private Instant sessionStart = Instant.EPOCH;92 private OneShotNode(93 Tracer tracer,94 EventBus events,95 Secret registrationSecret,96 NodeId id,97 URI uri,98 URI gridUri,99 Capabilities stereotype,100 WebDriverInfo driverInfo) {101 super(tracer, id, uri, registrationSecret);102 this.registrationSecret = registrationSecret;103 this.events = Require.nonNull("Event bus", events);104 this.gridUri = Require.nonNull("Public Grid URI", gridUri);105 this.stereotype = ImmutableCapabilities.copyOf(Require.nonNull("Stereotype", stereotype));106 this.driverInfo = Require.nonNull("Driver info", driverInfo);107 }108 public static Node create(Config config) {109 LoggingOptions loggingOptions = new LoggingOptions(config);110 EventBusOptions eventOptions = new EventBusOptions(config);111 BaseServerOptions serverOptions = new BaseServerOptions(config);112 NodeOptions nodeOptions = new NodeOptions(config);113 Map<String, Object> raw = new Json().toType(114 config.get("k8s", "stereotype")115 .orElseThrow(() -> new ConfigException("Unable to find node stereotype")),116 MAP_TYPE);117 Capabilities stereotype = new ImmutableCapabilities(raw);118 Optional<String> driverName = config.get("k8s", "driver_name").map(String::toLowerCase);119 // Find the webdriver info corresponding to the driver name120 WebDriverInfo driverInfo = StreamSupport.stream(ServiceLoader.load(WebDriverInfo.class).spliterator(), false)121 .filter(info -> info.isSupporting(stereotype))122 .filter(info -> driverName.map(name -> name.equals(info.getDisplayName().toLowerCase())).orElse(true))123 .findFirst()124 .orElseThrow(() -> new ConfigException(125 "Unable to find matching driver for %s and %s", stereotype, driverName.orElse("any driver")));126 LOG.info(String.format("Creating one-shot node for %s with stereotype %s", driverInfo, stereotype));127 LOG.info("Grid URI is: " + nodeOptions.getPublicGridUri());128 return new OneShotNode(129 loggingOptions.getTracer(),130 eventOptions.getEventBus(),131 serverOptions.getRegistrationSecret(),132 new NodeId(UUID.randomUUID()),133 serverOptions.getExternalUri(),134 nodeOptions.getPublicGridUri().orElseThrow(() -> new ConfigException("Unable to determine public grid address")),135 stereotype,136 driverInfo);137 }138 @Override139 public Optional<CreateSessionResponse> newSession(CreateSessionRequest sessionRequest) {140 if (driver != null) {141 throw new IllegalStateException("Only expected one session at a time");142 }143 Optional<WebDriver> driver = driverInfo.createDriver(sessionRequest.getCapabilities());144 if (!driver.isPresent()) {145 return Optional.empty();146 }147 if (!(driver.get() instanceof RemoteWebDriver)) {148 driver.get().quit();149 return Optional.empty();150 }151 this.driver = (RemoteWebDriver) driver.get();152 this.sessionId = this.driver.getSessionId();153 this.client = extractHttpClient(this.driver);154 this.capabilities = rewriteCapabilities(this.driver);155 this.sessionStart = Instant.now();156 LOG.info("Encoded response: " + JSON.toJson(ImmutableMap.of(157 "value", ImmutableMap.of(158 "sessionId", sessionId,159 "capabilities", capabilities))));160 events.fire(new NodeDrainStarted(getId()));161 return Optional.of(162 new CreateSessionResponse(163 getSession(sessionId),164 JSON.toJson(ImmutableMap.of(165 "value", ImmutableMap.of(166 "sessionId", sessionId,167 "capabilities", capabilities))).getBytes(UTF_8)));168 }169 private HttpClient extractHttpClient(RemoteWebDriver driver) {170 CommandExecutor executor = driver.getCommandExecutor();171 try {172 Field client = null;173 Class<?> current = executor.getClass();174 while (client == null && (current != null || Object.class.equals(current))) {175 client = findClientField(current);176 current = current.getSuperclass();177 }178 if (client == null) {179 throw new IllegalStateException("Unable to find client field in " + executor.getClass());180 }181 if (!HttpClient.class.isAssignableFrom(client.getType())) {182 throw new IllegalStateException("Client field is not assignable to http client");183 }184 client.setAccessible(true);185 return (HttpClient) client.get(executor);186 } catch (ReflectiveOperationException e) {187 throw new IllegalStateException(e);188 }189 }190 private Field findClientField(Class<?> clazz) {191 try {192 return clazz.getDeclaredField("client");193 } catch (NoSuchFieldException e) {194 return null;195 }196 }197 private Capabilities rewriteCapabilities(RemoteWebDriver driver) {198 // Rewrite the se:options if necessary199 Object rawSeleniumOptions = driver.getCapabilities().getCapability("se:options");200 if (rawSeleniumOptions == null || rawSeleniumOptions instanceof Map) {201 @SuppressWarnings("unchecked") Map<String, Object> original = (Map<String, Object>) rawSeleniumOptions;202 Map<String, Object> updated = new TreeMap<>(original == null ? new HashMap<>() : original);203 String cdpPath = String.format("/session/%s/se/cdp", driver.getSessionId());204 updated.put("cdp", rewrite(cdpPath));205 return new PersistentCapabilities(driver.getCapabilities()).setCapability("se:options", updated);206 }207 return ImmutableCapabilities.copyOf(driver.getCapabilities());208 }209 private URI rewrite(String path) {210 try {211 return new URI(212 gridUri.getScheme(),213 gridUri.getUserInfo(),214 gridUri.getHost(),215 gridUri.getPort(),216 path,217 null,218 null);219 } catch (URISyntaxException e) {220 throw new RuntimeException(e);221 }222 }223 @Override224 public HttpResponse executeWebDriverCommand(HttpRequest req) {225 LOG.info("Executing " + req);226 HttpResponse res = client.execute(req);227 if (DELETE.equals(req.getMethod()) && req.getUri().equals("/session/" + sessionId)) {228 // Ensure the response is sent before we viciously kill the node229 new Thread(230 () -> {231 try {232 Thread.sleep(500);233 } catch (InterruptedException e) {234 Thread.currentThread().interrupt();235 throw new RuntimeException(e);236 }237 LOG.info("Stopping session: " + sessionId);238 stop(sessionId);239 },240 "Node clean up: " + getId())241 .start();242 }243 return res;244 }245 @Override246 public Session getSession(SessionId id) throws NoSuchSessionException {247 if (!isSessionOwner(id)) {248 throw new NoSuchSessionException("Unable to find session with id: " + id);249 }250 return new Session(251 sessionId,252 getUri(),253 stereotype,254 capabilities,255 sessionStart); }256 @Override257 public HttpResponse uploadFile(HttpRequest req, SessionId id) {258 return null;259 }260 @Override261 public void stop(SessionId id) throws NoSuchSessionException {262 LOG.info("Stop has been called: " + id);263 Require.nonNull("Session ID", id);264 if (!isSessionOwner(id)) {265 throw new NoSuchSessionException("Unable to find session " + id);266 }267 LOG.info("Quitting session " + id);268 try {269 driver.quit();270 } catch (Exception e) {271 // It's possible that the driver has already quit.272 }273 events.fire(new SessionClosedEvent(id));274 LOG.info("Firing node drain complete message");275 events.fire(new NodeDrainComplete(getId()));276 }277 @Override278 public boolean isSessionOwner(SessionId id) {279 return driver != null && sessionId.equals(id);280 }281 @Override282 public boolean isSupporting(Capabilities capabilities) {283 return driverInfo.isSupporting(capabilities);284 }285 @Override286 public NodeStatus getStatus() {287 return new NodeStatus(288 getId(),289 getUri(),290 1,291 ImmutableSet.of(292 new Slot(293 new SlotId(getId(), slotId),294 stereotype,295 Instant.EPOCH,296 driver == null ?297 Optional.empty() :298 Optional.of(new Session(sessionId, getUri(), stereotype, capabilities, Instant.now())))),299 isDraining() ? DRAINING : UP,300 registrationSecret);301 }302 @Override303 public void drain() {304 events.fire(new NodeDrainStarted(getId()));305 draining = true;306 }307 @Override...