[grid] Make distributor pick up session requests from the queue only if the grid is not full (#8940)

* [grid] Make distributor pick up session requests from the queue only if the grid is not full

* [grid] Move the handleSession method into the runnable

* [grid] Move private methods to the inner class NewSessionRunnable.

* Code formatting and fixing test

* Adding a tear down for NewSessionCreationTest

* Moving `exerciseDriver` back to EndToEndTest

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

Author: pujagani

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

@@ -79,7 +79,12 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Queue;
 import java.util.Set;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
@@ -98,6 +103,9 @@ public class LocalDistributor extends Distributor {
   private final Secret registrationSecret;
   private final Regularly hostChecker = new Regularly("distributor host checker");
   private final Map<NodeId, Runnable> allChecks = new HashMap<>();
+  private final Queue<RequestId> requestIds = new ConcurrentLinkedQueue<>();
+  private final ScheduledExecutorService executorService =
+    Executors.newSingleThreadScheduledExecutor();
 
   private final ReadWriteLock lock = new ReentrantReadWriteLock(/* fair */ true);
   private final GridModel model;
@@ -125,64 +133,92 @@ public LocalDistributor(
     bus.addListener(NodeStatusEvent.listener(this::register));
     bus.addListener(NodeStatusEvent.listener(model::refresh));
     bus.addListener(NodeDrainComplete.listener(this::remove));
+    bus.addListener(NewSessionRequestEvent.listener(requestIds::offer));
 
-    bus.addListener(NewSessionRequestEvent.listener(reqId -> {
-      Optional<HttpRequest> sessionRequest = this.sessionRequests.remove(reqId);
-      // Check if polling the queue did not return null
-      if (sessionRequest.isPresent()) {
-        handleNewSessionRequest(sessionRequest.get(), reqId);
-      } else {
-        fireSessionRejectedEvent(
-          "Unable to poll request from the new session request queue.",
-          reqId);
-      }
-    }));
+    Thread shutdownHook = new Thread(this::callExecutorShutdown);
+    Runtime.getRuntime().addShutdownHook(shutdownHook);
+    NewSessionRunnable runnable = new NewSessionRunnable();
+    executorService.scheduleAtFixedRate(runnable, 0, 1000, TimeUnit.MILLISECONDS);
   }
 
-  private void handleNewSessionRequest(HttpRequest sessionRequest, RequestId reqId) {
-
-    try (Span span = newSpanAsChildOf(tracer, sessionRequest, "distributor.poll_queue")) {
-      Map<String, EventAttributeValue> attributeMap = new HashMap<>();
-      attributeMap.put(
-          AttributeKey.LOGGER_CLASS.getKey(), EventAttribute.setValue(getClass().getName()));
-      span.setAttribute(AttributeKey.REQUEST_ID.getKey(), reqId.toString());
-      attributeMap.put(AttributeKey.REQUEST_ID.getKey(), EventAttribute.setValue(reqId.toString()));
+  public class NewSessionRunnable implements Runnable {
+    @Override
+    public void run() {
+      Lock writeLock = lock.writeLock();
+      writeLock.lock();
+      try {
+        if (!requestIds.isEmpty()) {
+          boolean hasCapacity = nodes
+            .keySet()
+            .stream()
+            .anyMatch(key -> nodes.get(key).getStatus().hasCapacity());
+          if (hasCapacity) {
+            RequestId reqId = requestIds.poll();
+            if (reqId != null) {
+              Optional<HttpRequest> optionalHttpRequest = sessionRequests.remove(reqId);
+              // Check if polling the queue did not return null
+              if (optionalHttpRequest.isPresent()) {
+                handleNewSessionRequest(optionalHttpRequest.get(), reqId);
+              } else {
+                fireSessionRejectedEvent(
+                  "Unable to poll request from the new session request queue.",
+                  reqId);
+              }
+            }
+          }
+        }
+      } finally {
+        writeLock.unlock();
+      }
+    }
 
-      attributeMap.put("request", EventAttribute.setValue(sessionRequest.toString()));
-      Either<SessionNotCreatedException, CreateSessionResponse> response =
+    private void handleNewSessionRequest(HttpRequest sessionRequest, RequestId reqId) {
+      try (Span span = newSpanAsChildOf(tracer, sessionRequest, "distributor.poll_queue")) {
+        Map<String, EventAttributeValue> attributeMap = new HashMap<>();
+        attributeMap.put(
+          AttributeKey.LOGGER_CLASS.getKey(),
+          EventAttribute.setValue(getClass().getName()));
+        span.setAttribute(AttributeKey.REQUEST_ID.getKey(), reqId.toString());
+        attributeMap.put(
+          AttributeKey.REQUEST_ID.getKey(),
+          EventAttribute.setValue(reqId.toString()));
+
+        attributeMap.put("request", EventAttribute.setValue(sessionRequest.toString()));
+        Either<SessionNotCreatedException, CreateSessionResponse> response =
           newSession(sessionRequest);
-      if (response.isRight()) {
-        CreateSessionResponse sessionResponse = response.right();
-        NewSessionResponse newSessionResponse =
+        if (response.isRight()) {
+          CreateSessionResponse sessionResponse = response.right();
+          NewSessionResponse newSessionResponse =
             new NewSessionResponse(
-                reqId,
-                sessionResponse.getSession(),
-                sessionResponse.getDownstreamEncodedResponse());
+              reqId,
+              sessionResponse.getSession(),
+              sessionResponse.getDownstreamEncodedResponse());
 
-        bus.fire(new NewSessionResponseEvent(newSessionResponse));
-      } else {
-        SessionNotCreatedException exception = response.left();
+          bus.fire(new NewSessionResponseEvent(newSessionResponse));
+        } else {
+          SessionNotCreatedException exception = response.left();
 
-        if (exception instanceof RetrySessionRequestException) {
-          boolean retried = this.sessionRequests.retryAddToQueue(sessionRequest, reqId);
+          if (exception instanceof RetrySessionRequestException) {
+            boolean retried = sessionRequests.retryAddToQueue(sessionRequest, reqId);
 
-          attributeMap.put("request.retry_add", EventAttribute.setValue(retried));
-          span.addEvent("Retry adding to front of queue. All slots are busy.", attributeMap);
+            attributeMap.put("request.retry_add", EventAttribute.setValue(retried));
+            span.addEvent("Retry adding to front of queue. No slot available.", attributeMap);
 
-          if (!retried) {
-            span.addEvent("Retry adding to front of queue failed.", attributeMap);
+            if (!retried) {
+              span.addEvent("Retry adding to front of queue failed.", attributeMap);
+              fireSessionRejectedEvent(exception.getMessage(), reqId);
+            }
+          } else {
             fireSessionRejectedEvent(exception.getMessage(), reqId);
           }
-        } else {
-          fireSessionRejectedEvent(exception.getMessage(), reqId);
         }
       }
     }
-  }
 
-  private void fireSessionRejectedEvent(String message, RequestId reqId) {
-    bus.fire(
+    private void fireSessionRejectedEvent(String message, RequestId reqId) {
+      bus.fire(
         new NewSessionRejectedEvent(new NewSessionErrorResponse(reqId, message)));
+    }
   }
 
   public static Distributor create(Config config) {
@@ -393,4 +429,9 @@ protected Supplier<CreateSessionResponse> reserve(SlotId slotId, CreateSessionRe
       writeLock.unlock();
     }
   }
+
+  public void callExecutorShutdown() {
+    LOG.info("Shutting down Distributor executor service");
+    executorService.shutdownNow();
+  }
 }

java/server/test/org/openqa/selenium/grid/router/EndToEndTest.java

@@ -141,7 +141,7 @@ private static TestData createStandalone() {
       "port = " + PortProber.findFreePort(),
       "registration-secret = \"provolone\"",
       "[sessionqueue]",
-      "session-request-timeout = 2",
+      "session-request-timeout = 10",
       "session-retry-interval = 1"
     };
     Config config = new MemoizedConfig(
@@ -181,7 +181,7 @@ private static TestData createHubAndNode() {
       "[server]",
       "registration-secret = \"feta\"",
       "[sessionqueue]",
-      "session-request-timeout = 2",
+      "session-request-timeout = 10",
       "session-retry-interval = 1"
     };
 
@@ -228,7 +228,7 @@ private static TestData createFullyDistributed() {
       "[server]",
       "registration-secret = \"colby\"",
       "[sessionqueue]",
-      "session-request-timeout = 2",
+      "session-request-timeout = 10",
       "session-retry-interval = 1"
     };
 
@@ -310,26 +310,27 @@ private static TestData createFullyDistributed() {
       eventServer::stop);
   }
 
-  private static void waitUntilReady(Server<?> server, Duration duration) {
-    waitUntilReady(server, duration, false);
-  }
-
-  private static void waitUntilReady(Server<?> server, Duration duration, boolean printOutput) {
+  private static void waitUntilReady(Server<?> server, Duration duration, boolean logOutput) {
     HttpClient client = HttpClient.Factory.createDefault().createClient(server.getUrl());
 
     new FluentWait<>(client)
       .withTimeout(duration)
       .pollingEvery(Duration.ofSeconds(1))
       .until(c -> {
         HttpResponse response = c.execute(new HttpRequest(GET, "/status"));
-        if (printOutput) {
+        if (logOutput) {
           System.out.println(Contents.string(response));
         }
         Map<String, Object> status = Values.get(response, MAP_TYPE);
-        return Boolean.TRUE.equals(status.get("ready"));
+        return Boolean.TRUE.equals(status != null &&
+                                   Boolean.parseBoolean(status.get("ready").toString()));
       });
   }
 
+  private static void waitUntilReady(Server<?> server, Duration duration) {
+    waitUntilReady(server, duration, false);
+  }
+
   private static Config setRandomPort(Config config) {
     return new MemoizedConfig(
       new CompoundConfig(
@@ -397,7 +398,7 @@ public void exerciseDriver() {
     // Kill the session, and wait until the grid says it's ready
     driver.quit();
 
-    waitUntilReady(server, Duration.ofSeconds(200), true);
+    waitUntilReady(server, Duration.ofSeconds(100), true);
 
     // And now we're good to go.
     driver = new RemoteWebDriver(server.getUrl(), caps);
@@ -433,10 +434,6 @@ public void shouldAllowPassthroughForW3CMode() {
 
   @Test
   public void shouldRejectSessionRequestIfCapsNotSupported() {
-    Capabilities caps = new ImmutableCapabilities("browserName", "cheese", "type", "cheddar");
-    WebDriver driver = new RemoteWebDriver(server.getUrl(), caps);
-    driver.get("http://www.google.com");
-
     try {
       Capabilities unsupportedCaps = new ImmutableCapabilities("browserName", "brie");
       WebDriver disposable = new RemoteWebDriver(server.getUrl(), unsupportedCaps);
@@ -445,8 +442,6 @@ public void shouldRejectSessionRequestIfCapsNotSupported() {
     } catch (SessionNotCreatedException expected) {
       // Fall through
     }
-
-    driver.quit();
   }
 
   @Test

java/server/test/org/openqa/selenium/grid/router/NewSessionCreationTest.java

@@ -20,12 +20,13 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.openqa.selenium.Capabilities;
 import org.openqa.selenium.ImmutableCapabilities;
 import org.openqa.selenium.SessionNotCreatedException;
-import org.openqa.selenium.WebDriverInfo;
 import org.openqa.selenium.chrome.ChromeDriverInfo;
 import org.openqa.selenium.events.EventBus;
 import org.openqa.selenium.events.local.GuavaEventBus;
@@ -35,7 +36,6 @@
 import org.openqa.selenium.grid.distributor.Distributor;
 import org.openqa.selenium.grid.distributor.local.LocalDistributor;
 import org.openqa.selenium.grid.node.Node;
-import org.openqa.selenium.grid.node.config.DriverServiceSessionFactory;
 import org.openqa.selenium.grid.node.local.LocalNode;
 import org.openqa.selenium.grid.security.Secret;
 import org.openqa.selenium.grid.server.BaseServerOptions;
@@ -55,12 +55,10 @@
 import org.openqa.selenium.remote.http.HttpRequest;
 import org.openqa.selenium.remote.http.HttpResponse;
 import org.openqa.selenium.remote.http.Routable;
-import org.openqa.selenium.remote.service.DriverService;
 import org.openqa.selenium.remote.tracing.DefaultTestTracer;
 import org.openqa.selenium.remote.tracing.Tracer;
 import org.openqa.selenium.testing.drivers.Browser;
 
-import java.net.MalformedURLException;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.time.Duration;
@@ -81,6 +79,7 @@ public class NewSessionCreationTest {
   private EventBus events;
   private HttpClient.Factory clientFactory;
   private Secret registrationSecret;
+  private Server<?> server;
 
   @Before
   public void setup() {
@@ -90,6 +89,11 @@ public void setup() {
     registrationSecret = new Secret("hereford hop");
   }
 
+  @After
+  public void stopServer() {
+    server.stop();
+  }
+
   @Test
   public void ensureJsCannotCreateANewSession() throws URISyntaxException {
     ChromeDriverInfo chromeDriverInfo = new ChromeDriverInfo();
@@ -116,7 +120,7 @@ public void ensureJsCannotCreateANewSession() throws URISyntaxException {
     Routable router = new Router(tracer, clientFactory, sessions, queuer, distributor)
       .with(new EnsureSpecCompliantHeaders(ImmutableList.of(), ImmutableSet.of()));
 
-    Server<?> server = new NettyServer(
+    server = new NettyServer(
       new BaseServerOptions(new MapConfig(ImmutableMap.of())),
       router,
       new ProxyCdpIntoGrid(clientFactory, sessions))
@@ -167,8 +171,7 @@ public void ensureJsCannotCreateANewSession() throws URISyntaxException {
   }
 
   @Test
-  public void shouldRetryNewSessionRequestOnUnexpectedError() throws
-    URISyntaxException, MalformedURLException {
+  public void shouldRetryNewSessionRequestOnUnexpectedError() throws URISyntaxException {
     Capabilities capabilities = new ImmutableCapabilities("browserName", "cheese");
     URI nodeUri = new URI("http://localhost:4444");
     CombinedHandler handler = new CombinedHandler();
@@ -218,7 +221,7 @@ public void shouldRetryNewSessionRequestOnUnexpectedError() throws
     Router router = new Router(tracer, clientFactory, sessions, queuer, distributor);
     handler.addHandler(router);
 
-    Server<?> server = new NettyServer(
+    server = new NettyServer(
       new BaseServerOptions(
         new MapConfig(ImmutableMap.of())),
       handler);
@@ -236,17 +239,4 @@ public void shouldRetryNewSessionRequestOnUnexpectedError() throws
     assertThat(httpResponse.getStatus()).isEqualTo(HTTP_OK);
   }
 
-  private LocalNode.Builder addDriverFactory(
-    LocalNode.Builder builder,
-    WebDriverInfo info,
-    DriverService.Builder<?, ?> driverService) {
-    return builder.add(
-      info.getCanonicalCapabilities(),
-      new DriverServiceSessionFactory(
-        tracer,
-        clientFactory,
-        info.getCanonicalCapabilities(),
-        info::isSupporting,
-        driverService));
-  }
 }