[grid] Extract Node session creation to use Either to help retry sessions accurately (#9168)

* [grid] Make LocalNode communicate whether to retry session request or not.

* [grid] Use WebDriverException to encapsulate errors from Node. Extract session creation in One Shot node to use Either.

* [grid] Small text changes

Co-authored-by: Diego Molina <[email protected]>
Co-authored-by: Diego Molina <[email protected]>

Author: pujagani

java/server/src/org/openqa/selenium/grid/distributor/RetrySessionRequestException.java renamed to java/client/src/org/openqa/selenium/RetrySessionRequestException.java

@@ -15,9 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-package org.openqa.selenium.grid.distributor;
-
-import org.openqa.selenium.SessionNotCreatedException;
+package org.openqa.selenium;
 
 public class RetrySessionRequestException extends SessionNotCreatedException {
   public RetrySessionRequestException(String msg) {

java/server/src/org/openqa/selenium/grid/distributor/Distributor.java

@@ -20,6 +20,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import org.openqa.selenium.Capabilities;
+import org.openqa.selenium.RetrySessionRequestException;
 import org.openqa.selenium.SessionNotCreatedException;
 import org.openqa.selenium.grid.data.CreateSessionRequest;
 import org.openqa.selenium.grid.data.CreateSessionResponse;

java/server/src/org/openqa/selenium/grid/distributor/local/LocalDistributor.java

@@ -26,6 +26,7 @@
 import org.openqa.selenium.Beta;
 import org.openqa.selenium.Capabilities;
 import org.openqa.selenium.ImmutableCapabilities;
+import org.openqa.selenium.RetrySessionRequestException;
 import org.openqa.selenium.SessionNotCreatedException;
 import org.openqa.selenium.concurrent.Regularly;
 import org.openqa.selenium.events.EventBus;
@@ -48,7 +49,6 @@
 import org.openqa.selenium.grid.data.Slot;
 import org.openqa.selenium.grid.data.SlotId;
 import org.openqa.selenium.grid.distributor.Distributor;
-import org.openqa.selenium.grid.distributor.RetrySessionRequestException;
 import org.openqa.selenium.grid.distributor.selector.DefaultSlotSelector;
 import org.openqa.selenium.grid.log.LoggingOptions;
 import org.openqa.selenium.grid.node.HealthCheck;

java/server/src/org/openqa/selenium/grid/node/k8s/OneShotNode.java

@@ -24,6 +24,7 @@
 import org.openqa.selenium.NoSuchSessionException;
 import org.openqa.selenium.PersistentCapabilities;
 import org.openqa.selenium.WebDriver;
+import org.openqa.selenium.WebDriverException;
 import org.openqa.selenium.WebDriverInfo;
 import org.openqa.selenium.events.EventBus;
 import org.openqa.selenium.grid.config.Config;
@@ -46,6 +47,7 @@
 import org.openqa.selenium.grid.security.SecretOptions;
 import org.openqa.selenium.grid.server.BaseServerOptions;
 import org.openqa.selenium.grid.server.EventBusOptions;
+import org.openqa.selenium.internal.Either;
 import org.openqa.selenium.internal.Require;
 import org.openqa.selenium.json.Json;
 import org.openqa.selenium.remote.CommandExecutor;
@@ -154,18 +156,28 @@ public static Node create(Config config) {
 
   @Override
   public Optional<CreateSessionResponse> newSession(CreateSessionRequest sessionRequest) {
+    Either<WebDriverException, CreateSessionResponse> result = createNewSession(sessionRequest);
+
+    if (result.isRight()) {
+      return Optional.of(result.right());
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  public Either<WebDriverException, CreateSessionResponse> createNewSession(CreateSessionRequest sessionRequest) {
     if (driver != null) {
       throw new IllegalStateException("Only expected one session at a time");
     }
 
     Optional<WebDriver> driver = driverInfo.createDriver(sessionRequest.getCapabilities());
     if (!driver.isPresent()) {
-      return Optional.empty();
+      return Either.left(new WebDriverException("Unable to create a driver instance"));
     }
 
     if (!(driver.get() instanceof RemoteWebDriver)) {
       driver.get().quit();
-      return Optional.empty();
+      return Either.left(new WebDriverException("Driver is not a RemoteWebDriver instance"));
     }
 
     this.driver = (RemoteWebDriver) driver.get();
@@ -181,7 +193,7 @@ public Optional<CreateSessionResponse> newSession(CreateSessionRequest sessionRe
 
     events.fire(new NodeDrainStarted(getId()));
 
-    return Optional.of(
+    return Either.right(
       new CreateSessionResponse(
         getSession(sessionId),
         JSON.toJson(ImmutableMap.of(

java/server/src/org/openqa/selenium/grid/node/local/LocalNode.java

@@ -28,6 +28,8 @@
 import org.openqa.selenium.ImmutableCapabilities;
 import org.openqa.selenium.NoSuchSessionException;
 import org.openqa.selenium.PersistentCapabilities;
+import org.openqa.selenium.RetrySessionRequestException;
+import org.openqa.selenium.SessionNotCreatedException;
 import org.openqa.selenium.WebDriverException;
 import org.openqa.selenium.concurrent.Regularly;
 import org.openqa.selenium.events.EventBus;
@@ -42,21 +44,22 @@
 import org.openqa.selenium.grid.data.SessionClosedEvent;
 import org.openqa.selenium.grid.data.Slot;
 import org.openqa.selenium.grid.data.SlotId;
+import org.openqa.selenium.grid.jmx.JMXHelper;
+import org.openqa.selenium.grid.jmx.ManagedAttribute;
+import org.openqa.selenium.grid.jmx.ManagedService;
 import org.openqa.selenium.grid.node.ActiveSession;
 import org.openqa.selenium.grid.node.HealthCheck;
 import org.openqa.selenium.grid.node.Node;
 import org.openqa.selenium.grid.node.SessionFactory;
 import org.openqa.selenium.grid.security.Secret;
+import org.openqa.selenium.internal.Either;
 import org.openqa.selenium.internal.Require;
 import org.openqa.selenium.io.TemporaryFilesystem;
 import org.openqa.selenium.io.Zip;
 import org.openqa.selenium.json.Json;
 import org.openqa.selenium.remote.SessionId;
 import org.openqa.selenium.remote.http.HttpRequest;
 import org.openqa.selenium.remote.http.HttpResponse;
-import org.openqa.selenium.grid.jmx.JMXHelper;
-import org.openqa.selenium.grid.jmx.ManagedAttribute;
-import org.openqa.selenium.grid.jmx.ManagedService;
 import org.openqa.selenium.remote.tracing.AttributeKey;
 import org.openqa.selenium.remote.tracing.EventAttribute;
 import org.openqa.selenium.remote.tracing.EventAttributeValue;
@@ -252,17 +255,27 @@ public boolean isSupporting(Capabilities capabilities) {
 
   @Override
   public Optional<CreateSessionResponse> newSession(CreateSessionRequest sessionRequest) {
+    Either<WebDriverException, CreateSessionResponse> result = createNewSession(sessionRequest);
+
+    if (result.isRight()) {
+      return Optional.of(result.right());
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  public Either<WebDriverException, CreateSessionResponse> createNewSession(CreateSessionRequest sessionRequest) {
     Require.nonNull("Session request", sessionRequest);
 
     try (Span span = tracer.getCurrentContext().createSpan("node.new_session")) {
       Map<String, EventAttributeValue> attributeMap = new HashMap<>();
       attributeMap
-          .put(AttributeKey.LOGGER_CLASS.getKey(), EventAttribute.setValue(getClass().getName()));
-      LOG.fine("Creating new session using span: " + span);
+        .put(AttributeKey.LOGGER_CLASS.getKey(), EventAttribute.setValue(getClass().getName()));
       attributeMap.put("session.request.capabilities",
-                       EventAttribute.setValue(sessionRequest.getCapabilities().toString()));
+        EventAttribute.setValue(sessionRequest.getCapabilities().toString()));
       attributeMap.put("session.request.downstreamdialect",
-                       EventAttribute.setValue(sessionRequest.getDownstreamDialects().toString()));
+        EventAttribute.setValue(sessionRequest.getDownstreamDialects().toString()));
+
       int currentSessionCount = getCurrentSessionCount();
       span.setAttribute("current.session.count", currentSessionCount);
       attributeMap.put("current.session.count", EventAttribute.setValue(currentSessionCount));
@@ -272,11 +285,12 @@ public Optional<CreateSessionResponse> newSession(CreateSessionRequest sessionRe
         span.setStatus(Status.RESOURCE_EXHAUSTED);
         attributeMap.put("max.session.count", EventAttribute.setValue(maxSessionCount));
         span.addEvent("Max session count reached", attributeMap);
-        return Optional.empty();
+        return Either.left(new RetrySessionRequestException("Max session count reached."));
       }
       if (isDraining()) {
         span.setStatus(Status.UNAVAILABLE.withDescription("The node is draining. Cannot accept new sessions."));
-        return Optional.empty();
+        return Either.left(
+          new RetrySessionRequestException("The node is draining. Cannot accept new sessions."));
       }
 
       // Identify possible slots to use as quickly as possible to enable concurrent session starting
@@ -296,8 +310,9 @@ public Optional<CreateSessionResponse> newSession(CreateSessionRequest sessionRe
       if (slotToUse == null) {
         span.setAttribute("error", true);
         span.setStatus(Status.NOT_FOUND);
-        span.addEvent("No slot matched capabilities ", attributeMap);
-        return Optional.empty();
+        span.addEvent("No slot matched the requested capabilities. All slots are busy.", attributeMap);
+        return Either.left(
+          new RetrySessionRequestException("No slot matched the requested capabilities. All slots are busy."));
       }
 
       Optional<ActiveSession> possibleSession = slotToUse.apply(sessionRequest);
@@ -306,8 +321,8 @@ public Optional<CreateSessionResponse> newSession(CreateSessionRequest sessionRe
         slotToUse.release();
         span.setAttribute("error", true);
         span.setStatus(Status.NOT_FOUND);
-        span.addEvent("No slots available for capabilities ", attributeMap);
-        return Optional.empty();
+        span.addEvent("Unable to create session with the driver", attributeMap);
+        return Either.left(new SessionNotCreatedException("Unable to create session with the driver"));
       }
 
       ActiveSession session = possibleSession.get();
@@ -327,9 +342,9 @@ public Optional<CreateSessionResponse> newSession(CreateSessionRequest sessionRe
       // The session we return has to look like it came from the node, since we might be dealing
       // with a webdriver implementation that only accepts connections from localhost
       Session externalSession = createExternalSession(session, externalUri);
-      return Optional.of(new CreateSessionResponse(
-          externalSession,
-          getEncoder(session.getDownstreamDialect()).apply(externalSession)));
+      return Either.right(new CreateSessionResponse(
+        externalSession,
+        getEncoder(session.getDownstreamDialect()).apply(externalSession)));
     }
   }