Merge remote-tracking branch 'origin/trunk' into router-cache

Author: VietND96

java/src/org/openqa/selenium/firefox/GeckoDriverService.java

@@ -147,6 +147,8 @@ public static class Builder
     private @Nullable FirefoxDriverLogLevel logLevel;
     private @Nullable Boolean logTruncate;
     private @Nullable File profileRoot;
+    private @Nullable Integer marionettePort;
+    private @Nullable Integer websocketPort;
 
     @Override
     public int score(Capabilities capabilities) {
@@ -204,6 +206,31 @@ public GeckoDriverService.Builder withProfileRoot(@Nullable File root) {
       return this;
     }
 
+    /**
+     * Configures geckodriver to connect to an existing Firefox instance via the specified
+     * Marionette port.
+     *
+     * @param marionettePort The port where Marionette is listening on the existing Firefox
+     *     instance.
+     * @return A self reference.
+     */
+    public GeckoDriverService.Builder connectToExisting(int marionettePort) {
+      this.marionettePort = marionettePort;
+      return this;
+    }
+
+    /**
+     * Configures the WebSocket port for BiDi. A value of 0 will automatically allocate a free port.
+     *
+     * @param websocketPort The port to use for WebSocket communication, or 0 for automatic
+     *     allocation.
+     * @return A self reference.
+     */
+    public GeckoDriverService.Builder withWebSocketPort(@Nullable Integer websocketPort) {
+      this.websocketPort = websocketPort;
+      return this;
+    }
+
     @Override
     protected void loadSystemProperties() {
       parseLogOutput(GECKO_DRIVER_LOG_PROPERTY);
@@ -229,13 +256,27 @@ protected List<String> createArgs() {
       List<String> args = new ArrayList<>();
       args.add(String.format(Locale.ROOT, "--port=%d", getPort()));
 
-      int wsPort = PortProber.findFreePort();
-      args.add(String.format("--websocket-port=%d", wsPort));
-
-      args.add("--allow-origins");
-      args.add(String.format("http://127.0.0.1:%d", wsPort));
-      args.add(String.format("http://localhost:%d", wsPort));
-      args.add(String.format("http://[::1]:%d", wsPort));
+      // Check if marionette port is specified via connectToExisting method
+      if (marionettePort != null) {
+        args.add("--connect-existing");
+        args.add("--marionette-port");
+        args.add(String.valueOf(marionettePort));
+      } else {
+        // Configure websocket port for BiDi communication
+        if (websocketPort != null) {
+          args.add("--websocket-port");
+          args.add(String.valueOf(websocketPort));
+
+          args.add("--allow-origins");
+          args.add(String.format("http://127.0.0.1:%d", websocketPort));
+          args.add(String.format("http://localhost:%d", websocketPort));
+          args.add(String.format("http://[::1]:%d", websocketPort));
+        } else {
+          // Use 0 to auto-allocate a free port
+          args.add("--websocket-port");
+          args.add("0");
+        }
+      }
 
       if (logLevel != null) {
         args.add("--log");

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

@@ -539,7 +539,7 @@ public void run() {
         // up starving a session request.
         Map<Capabilities, Long> stereotypes =
             getAvailableNodes().stream()
-                .filter(node -> node.hasCapacity())
+                .filter(NodeStatus::hasCapacity)
                 .flatMap(node -> node.getSlots().stream().map(Slot::getStereotype))
                 .collect(
                     Collectors.groupingBy(ImmutableCapabilities::copyOf, Collectors.counting()));

java/src/org/openqa/selenium/grid/distributor/local/LocalNodeRegistry.java

@@ -361,9 +361,8 @@ public Set<NodeStatus> getAvailableNodes() {
     readLock.lock();
     try {
       return model.getSnapshot().stream()
-          .filter(
-              node ->
-                  !DOWN.equals(node.getAvailability()) && !DRAINING.equals(node.getAvailability()))
+          // Filter nodes are UP and have capacity (available slots)
+          .filter(node -> UP.equals(node.getAvailability()) && node.hasCapacity())
           .collect(ImmutableSet.toImmutableSet());
     } finally {
       readLock.unlock();

java/src/org/openqa/selenium/grid/sessionqueue/local/LocalNewSessionQueue.java

@@ -55,6 +55,7 @@
 import org.openqa.selenium.grid.data.TraceSessionRequest;
 import org.openqa.selenium.grid.distributor.config.DistributorOptions;
 import org.openqa.selenium.grid.jmx.JMXHelper;
+import org.openqa.selenium.grid.jmx.MBean;
 import org.openqa.selenium.grid.jmx.ManagedAttribute;
 import org.openqa.selenium.grid.jmx.ManagedService;
 import org.openqa.selenium.grid.log.LoggingOptions;
@@ -110,6 +111,7 @@ public class LocalNewSessionQueue extends NewSessionQueue implements Closeable {
             thread.setName(NAME);
             return thread;
           });
+  private final MBean jmxBean;
 
   public LocalNewSessionQueue(
       Tracer tracer,
@@ -139,7 +141,8 @@ public LocalNewSessionQueue(
         requestTimeoutCheck.toMillis(),
         MILLISECONDS);
 
-    new JMXHelper().register(this);
+    // Manage JMX and unregister on close()
+    this.jmxBean = new JMXHelper().register(this);
   }
 
   public static NewSessionQueue create(Config config) {
@@ -502,6 +505,10 @@ public boolean isReady() {
   @Override
   public void close() {
     shutdownGracefully(NAME, service);
+
+    if (jmxBean != null) {
+      new JMXHelper().unregister(jmxBean.getObjectName());
+    }
   }
 
   private void failDueToTimeout(RequestId reqId) {

java/test/org/openqa/selenium/firefox/FirefoxDriverConcurrentTest.java

@@ -28,6 +28,7 @@
 import org.openqa.selenium.ParallelTestRunner.Worker;
 import org.openqa.selenium.WebDriver;
 import org.openqa.selenium.WebElement;
+import org.openqa.selenium.bidi.BiDi;
 import org.openqa.selenium.testing.JupiterTestBase;
 import org.openqa.selenium.testing.drivers.WebDriverBuilder;
 
@@ -164,4 +165,69 @@ void shouldBeAbleToUseTheSameProfileMoreThanOnce() {
       if (two != null) two.quit();
     }
   }
+
+  @Test
+  void multipleFirefoxInstancesWithBiDiEnabledCanRunSimultaneously() {
+    // Create two Firefox instances with BiDi enabled, should use different ports
+    FirefoxOptions options1 = new FirefoxOptions().enableBiDi();
+    FirefoxOptions options2 = new FirefoxOptions().enableBiDi();
+
+    WebDriver driver1 = null;
+    WebDriver driver2 = null;
+
+    try {
+      driver1 = new WebDriverBuilder().get(options1);
+      BiDi biDi1 = ((FirefoxDriver) driver1).getBiDi();
+      assertThat(biDi1).isNotNull();
+
+      // Extract the BiDi websocket URL and port for the first instance
+      String webSocketUrl1 =
+          (String) ((FirefoxDriver) driver1).getCapabilities().getCapability("webSocketUrl");
+      String port1 = webSocketUrl1.replaceAll("^ws://[^:]+:(\\d+)/.*$", "$1");
+
+      driver2 = new WebDriverBuilder().get(options2);
+      BiDi biDi2 = ((FirefoxDriver) driver2).getBiDi();
+      assertThat(biDi2).isNotNull();
+
+      // Extract the BiDi websocket URL and port for the second instance
+      String webSocketUrl2 =
+          (String) ((FirefoxDriver) driver2).getCapabilities().getCapability("webSocketUrl");
+      String port2 = webSocketUrl2.replaceAll("^ws://[^:]+:(\\d+)/.*$", "$1");
+
+      // Verify that the ports are different
+      assertThat(port1).isNotEqualTo(port2);
+    } finally {
+      // Clean up
+      if (driver1 != null) {
+        driver1.quit();
+      }
+      if (driver2 != null) {
+        driver2.quit();
+      }
+    }
+  }
+
+  @Test
+  void geckoDriverServiceConnectToExistingFirefox() {
+    GeckoDriverService.Builder builder = new GeckoDriverService.Builder();
+
+    // Test connectToExisting method
+    builder.connectToExisting(2829);
+    GeckoDriverService service = builder.build();
+
+    assertThat(service).isNotNull();
+    service.stop();
+  }
+
+  @Test
+  void geckoDriverServiceCustomWebSocketPort() {
+    GeckoDriverService.Builder builder = new GeckoDriverService.Builder();
+
+    // Test withWebSocketPort method
+    builder.withWebSocketPort(9225);
+    GeckoDriverService service = builder.build();
+
+    assertThat(service).isNotNull();
+    service.stop();
+  }
 }