From 46fc208fc1bcdad7aebdb762192b2ec44af4bddd Mon Sep 17 00:00:00 2001
From: Puja Jagani <puja.jagani93@gmail.com>
Date: Tue, 21 Sep 2021 01:28:28 +0530
Subject: [PATCH] [grid] Consider max-session value while selecting the slot
 and identifying Node capacity (#9838)

---
 .../openqa/selenium/grid/data/NodeStatus.java |  5 +--
 .../selector/DefaultSlotSelectorTest.java     | 33 ++++++++++++++++++-
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/java/src/org/openqa/selenium/grid/data/NodeStatus.java b/java/src/org/openqa/selenium/grid/data/NodeStatus.java
index da487d0016dfc..c38b5c742ff17 100644
--- a/java/src/org/openqa/selenium/grid/data/NodeStatus.java
+++ b/java/src/org/openqa/selenium/grid/data/NodeStatus.java
@@ -139,9 +139,10 @@ public boolean hasCapacity() {
     return slots.stream().anyMatch(slot -> slot.getSession() == null);
   }
 
+  // Check if the Node's max session limit is not exceeded and has a free slot that supports the capability.
   public boolean hasCapacity(Capabilities caps) {
-    return slots.stream()
-      .anyMatch(slot -> slot.getSession() == null && slot.isSupporting(caps));
+    return slots.stream().filter(slot -> slot.getSession() != null).count() < maxSessionCount &&
+           slots.stream().anyMatch(slot -> slot.getSession() == null && slot.isSupporting(caps));
   }
 
   public int getMaxSessionCount() {
diff --git a/java/test/org/openqa/selenium/grid/distributor/selector/DefaultSlotSelectorTest.java b/java/test/org/openqa/selenium/grid/distributor/selector/DefaultSlotSelectorTest.java
index 1a7ccb98a4459..111db3f7c237a 100644
--- a/java/test/org/openqa/selenium/grid/distributor/selector/DefaultSlotSelectorTest.java
+++ b/java/test/org/openqa/selenium/grid/distributor/selector/DefaultSlotSelectorTest.java
@@ -131,6 +131,37 @@ public void theMostLightlyLoadedNodeIsSelectedFirst() {
     assertThat(lightest.getSlots().stream()).anyMatch(slot -> expected.equals(slot.getId()));
   }
 
+  @Test
+  public void theNodeWhichHasExceededMaxSessionsIsNotSelected() {
+    Capabilities chrome = new ImmutableCapabilities("browserName", "chrome");
+
+    NodeStatus lightLoad =
+      createNode(ImmutableList.of(chrome), 12, 2);
+    NodeStatus mediumLoad =
+      createNode(ImmutableList.of(chrome), 12, 5);
+    NodeStatus maximumLoad =
+      createNode(ImmutableList.of(chrome), 12, 12);
+
+    Set<SlotId> ids = selector.selectSlot(chrome, ImmutableSet.of(maximumLoad, mediumLoad, lightLoad));
+    SlotId expected = ids.iterator().next();
+
+    // The slot should belong to the Node with light load
+    assertThat(lightLoad.getSlots().stream())
+      .anyMatch(slot -> expected.equals(slot.getId()));
+
+    // The node whose current number of sessions is greater than or equal to the max sessions is not included
+    // Hence, the node with the maximum load is skipped
+    ImmutableSet<NodeId> nodeIds = ids.stream()
+      .map(SlotId::getOwningNodeId)
+      .distinct()
+      .collect(toImmutableSet());
+    assertThat(nodeIds).doesNotContain(maximumLoad.getNodeId());
+    assertThat(nodeIds)
+      .containsSequence(
+        lightLoad.getNodeId(),
+        mediumLoad.getNodeId());
+  }
+
   @Test
   public void nodesAreOrderedByNumberOfSupportedBrowsersAndLoad() {
     Capabilities chrome = new ImmutableCapabilities("browserName", "chrome");
@@ -161,6 +192,7 @@ public void nodesAreOrderedByNumberOfSupportedBrowsersAndLoad() {
       .anyMatch(slot -> expected.equals(slot.getId()));
 
     // Nodes are ordered by the diversity of supported browsers, then by load
+    // The node whose current number of sessions is greater than or equal to the max sessions is not included
     ImmutableSet<NodeId> nodeIds = ids.stream()
       .map(SlotId::getOwningNodeId)
       .distinct()
@@ -169,7 +201,6 @@ public void nodesAreOrderedByNumberOfSupportedBrowsersAndLoad() {
       .containsSequence(
         highLoadAndOneBrowser.getNodeId(),
         mediumLoadAndTwoBrowsers.getNodeId(),
-        mediumLoadAndOtherTwoBrowsers.getNodeId(),
         lightLoadAndThreeBrowsers.getNodeId());
   }