[grid] Update for node shutdown gracefully

Signed-off-by: Viet Nguyen Duc <[email protected]>

Author: VietND96

java/src/org/openqa/selenium/grid/node/Node.java

@@ -152,7 +152,10 @@ protected Node(
                     req ->
                         getSessionId(req.getUri())
                             .map(SessionId::new)
-                            .map(this::isSessionOwner)
+                            .map(
+                                sessionId ->
+                                    this.getSession(sessionId) != null
+                                        && this.isSessionOwner(sessionId))
                             .orElse(false))
                 .to(() -> new ForwardWebDriverCommand(this))
                 .with(spanDecorator("node.forward_command")),

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

@@ -297,7 +297,13 @@ protected LocalNode(
         heartbeatPeriod.getSeconds(),
         TimeUnit.SECONDS);
 
-    Runtime.getRuntime().addShutdownHook(new Thread(this::stopAllSessions));
+    Runtime.getRuntime()
+        .addShutdownHook(
+            new Thread(
+                () -> {
+                  stopAllSessions();
+                  drain();
+                }));
     new JMXHelper().register(this);
   }
 
@@ -316,7 +322,6 @@ private void stopTimedOutSession(RemovalNotification<SessionId, SessionSlot> not
       }
       // Attempt to stop the session
       slot.stop();
-      this.sessionToDownloadsDir.invalidate(id);
       // Decrement pending sessions if Node is draining
       if (this.isDraining()) {
         int done = pendingSessions.decrementAndGet();
@@ -473,8 +478,6 @@ public Either<WebDriverException, CreateSessionResponse> newSession(
         sessionToDownloadsDir.put(session.getId(), uuidForSessionDownloads);
         currentSessions.put(session.getId(), slotToUse);
 
-        checkSessionCount();
-
         SessionId sessionId = session.getId();
         Capabilities caps = session.getCapabilities();
         SESSION_ID.accept(span, sessionId);
@@ -513,6 +516,8 @@ public Either<WebDriverException, CreateSessionResponse> newSession(
         span.addEvent("Unable to create session with the driver", attributeMap);
         return Either.left(possibleSession.left());
       }
+    } finally {
+      checkSessionCount();
     }
   }
 
@@ -765,6 +770,10 @@ public HttpResponse uploadFile(HttpRequest req, SessionId id) {
   public void stop(SessionId id) throws NoSuchSessionException {
     Require.nonNull("Session ID", id);
 
+    if (sessionToDownloadsDir.getIfPresent(id) != null) {
+      sessionToDownloadsDir.invalidate(id);
+    }
+
     SessionSlot slot = currentSessions.getIfPresent(id);
     if (slot == null) {
       throw new NoSuchSessionException("Cannot find session with id: " + id);

java/test/org/openqa/selenium/grid/e2e/BUILD.bazel

@@ -0,0 +1,20 @@
+load("@rules_jvm_external//:defs.bzl", "artifact")
+load("//java:defs.bzl", "JUNIT5_DEPS", "java_selenium_test_suite")
+
+java_selenium_test_suite(
+    name = "large-tests",
+    size = "large",
+    srcs = glob(["*.java"]),
+    browsers = [
+        "chrome",
+    ],
+    deps = [
+        "//java/src/org/openqa/selenium:core",
+        "//java/src/org/openqa/selenium/chrome",
+        "//java/src/org/openqa/selenium/grid",
+        "//java/src/org/openqa/selenium/remote/http",
+        artifact("com.google.guava:guava"),
+        artifact("org.junit.jupiter:junit-jupiter-api"),
+        artifact("org.assertj:assertj-core"),
+    ] + JUNIT5_DEPS,
+)

java/test/org/openqa/selenium/grid/e2e/SessionTimeoutTest.java

@@ -0,0 +1,49 @@
+// Licensed to the Software Freedom Conservancy (SFC) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The SFC licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.openqa.selenium.grid.e2e;
+
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.net.URI;
+import org.junit.jupiter.api.Test;
+import org.openqa.selenium.NoSuchSessionException;
+import org.openqa.selenium.WebDriver;
+import org.openqa.selenium.chrome.ChromeOptions;
+import org.openqa.selenium.grid.Bootstrap;
+import org.openqa.selenium.remote.RemoteWebDriver;
+
+public class SessionTimeoutTest {
+  @Test
+  void testSessionTimeout() throws Exception {
+    Bootstrap.main(("hub --host 127.0.0.1 --port 4444").split(" "));
+    Bootstrap.main(
+        ("node --host 127.0.0.1 --port 5555 --session-timeout 15 --selenium-manager true")
+            .split(" "));
+
+    var options = new ChromeOptions();
+    options.addArguments("--disable-search-engine-choice-screen");
+    options.addArguments("--headless=new");
+
+    WebDriver driver = new RemoteWebDriver(URI.create("http://localhost:4444").toURL(), options);
+    driver.get("http://localhost:4444");
+    Thread.sleep(15000);
+    NoSuchSessionException exception = assertThrows(NoSuchSessionException.class, driver::getTitle);
+    assertTrue(exception.getMessage().startsWith("Cannot find session with id:"));
+  }
+}