[grid] Add Grid Model flag

java/server/src/org/openqa/selenium/grid/commands/Hub.java

@@ -159,6 +159,7 @@ protected Handlers createHandlers(Config config) {
       clientFactory,
       sessions,
       queue,
+      distributorOptions.getGridModel(),
       distributorOptions.getSlotSelector(),
       secret,
       distributorOptions.getHealthCheckInterval(),

java/server/src/org/openqa/selenium/grid/commands/Standalone.java

@@ -157,6 +157,7 @@ protected Handlers createHandlers(Config config) {
       clientFactory,
       sessions,
       queue,
+      distributorOptions.getGridModel(),
       distributorOptions.getSlotSelector(),
       registrationSecret,
       distributorOptions.getHealthCheckInterval(),

java/server/src/org/openqa/selenium/grid/distributor/config/DistributorFlags.java

@@ -30,6 +30,7 @@
 
 import static org.openqa.selenium.grid.config.StandardGridRoles.DISTRIBUTOR_ROLE;
 import static org.openqa.selenium.grid.distributor.config.DistributorOptions.DEFAULT_DISTRIBUTOR_IMPLEMENTATION;
+import static org.openqa.selenium.grid.distributor.config.DistributorOptions.DEFAULT_GRID_MODEL_IMPLEMENTATION;
 import static org.openqa.selenium.grid.distributor.config.DistributorOptions.DEFAULT_HEALTHCHECK_INTERVAL;
 import static org.openqa.selenium.grid.distributor.config.DistributorOptions.DEFAULT_REJECT_UNSUPPORTED_CAPS;
 import static org.openqa.selenium.grid.distributor.config.DistributorOptions.DEFAULT_SLOT_MATCHER;
@@ -65,6 +66,12 @@ public class DistributorFlags implements HasRoles {
     example = DEFAULT_DISTRIBUTOR_IMPLEMENTATION)
   private String implementation = DEFAULT_DISTRIBUTOR_IMPLEMENTATION;
 
+  @Parameter(
+    names = {"--grid-model"},
+    description = "Full classname of non-default grid model. This is used to store states of the all the registered Nodes.")
+  @ConfigValue(section = DISTRIBUTOR_SECTION, name = "grid-model", example = DEFAULT_GRID_MODEL_IMPLEMENTATION)
+  private String gridModel = DEFAULT_GRID_MODEL_IMPLEMENTATION;
+
   @Parameter(
     names = {"--slot-matcher"},
     description = "Full classname of non-default slot matcher to use. This is used to determine whether a Node can support a particular session.")

java/server/src/org/openqa/selenium/grid/distributor/config/DistributorOptions.java

@@ -21,6 +21,7 @@
 import org.openqa.selenium.grid.config.ConfigException;
 import org.openqa.selenium.grid.data.SlotMatcher;
 import org.openqa.selenium.grid.distributor.Distributor;
+import org.openqa.selenium.grid.distributor.GridModel;
 import org.openqa.selenium.grid.distributor.selector.SlotSelector;
 
 import java.net.URI;
@@ -37,6 +38,8 @@ public class DistributorOptions {
   static final String DEFAULT_SLOT_MATCHER = "org.openqa.selenium.grid.data.DefaultSlotMatcher";
   static final String DEFAULT_SLOT_SELECTOR_IMPLEMENTATION =
     "org.openqa.selenium.grid.distributor.selector.DefaultSlotSelector";
+  static final String DEFAULT_GRID_MODEL_IMPLEMENTATION =
+    "org.openqa.selenium.grid.distributor.gridmodel.local.LocalGridModel";
   static final boolean DEFAULT_REJECT_UNSUPPORTED_CAPS = false;
   private final Config config;
 
@@ -101,6 +104,14 @@ public Distributor getDistributor() {
       DEFAULT_DISTRIBUTOR_IMPLEMENTATION);
   }
 
+  public GridModel getGridModel() {
+    return config.getClass(
+      DISTRIBUTOR_SECTION,
+      "grid-model",
+      GridModel.class,
+      DEFAULT_GRID_MODEL_IMPLEMENTATION);
+  }
+
   public SlotMatcher getSlotMatcher() {
     return config.getClass(
       DISTRIBUTOR_SECTION,

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

@@ -45,8 +45,8 @@
 import org.openqa.selenium.grid.data.SlotId;
 import org.openqa.selenium.grid.data.TraceSessionRequest;
 import org.openqa.selenium.grid.distributor.Distributor;
+import org.openqa.selenium.grid.distributor.GridModel;
 import org.openqa.selenium.grid.distributor.config.DistributorOptions;
-import org.openqa.selenium.grid.distributor.gridmodel.local.LocalGridModel;
 import org.openqa.selenium.grid.distributor.selector.SlotSelector;
 import org.openqa.selenium.grid.log.LoggingOptions;
 import org.openqa.selenium.grid.node.HealthCheck;
@@ -116,7 +116,7 @@ public class LocalDistributor extends Distributor {
   private final Duration healthcheckInterval;
 
   private final ReadWriteLock lock = new ReentrantReadWriteLock(/* fair */ true);
-  private final LocalGridModel model;
+  private final GridModel model;
   private final Map<NodeId, Node> nodes;
 
   private final NewSessionQueue sessionQueue;
@@ -130,6 +130,7 @@ public LocalDistributor(
     HttpClient.Factory clientFactory,
     SessionMap sessions,
     NewSessionQueue sessionQueue,
+    GridModel model,
     SlotSelector slotSelector,
     Secret registrationSecret,
     Duration healthcheckInterval,
@@ -143,7 +144,7 @@ public LocalDistributor(
     this.slotSelector = Require.nonNull("Slot selector", slotSelector);
     this.registrationSecret = Require.nonNull("Registration secret", registrationSecret);
     this.healthcheckInterval = Require.nonNull("Health check interval", healthcheckInterval);
-    this.model = new LocalGridModel(bus);
+    this.model = model;
     this.nodes = new ConcurrentHashMap<>();
     this.rejectUnsupportedCaps = rejectUnsupportedCaps;
 
@@ -189,6 +190,7 @@ public static Distributor create(Config config) {
       clientFactory,
       sessions,
       sessionQueue,
+      distributorOptions.getGridModel(),
       distributorOptions.getSlotSelector(),
       secretOptions.getRegistrationSecret(),
       distributorOptions.getHealthCheckInterval(),

java/server/test/org/openqa/selenium/grid/distributor/AddingNodesTest.java

@@ -38,6 +38,7 @@
 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.distributor.gridmodel.local.LocalGridModel;
 import org.openqa.selenium.grid.distributor.local.LocalDistributor;
 import org.openqa.selenium.grid.distributor.remote.RemoteDistributor;
 import org.openqa.selenium.grid.distributor.selector.DefaultSlotSelector;
@@ -120,6 +121,7 @@ public void setUpDistributor() throws MalformedURLException {
       clientFactory,
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),

java/server/test/org/openqa/selenium/grid/distributor/BUILD.bazel

@@ -20,6 +20,7 @@ java_test_suite(
         "//java/server/src/org/openqa/selenium/grid/component",
         "//java/server/src/org/openqa/selenium/grid/data",
         "//java/server/src/org/openqa/selenium/grid/distributor",
+        "//java/server/src/org/openqa/selenium/grid/distributor/gridmodel/local",
         "//java/server/src/org/openqa/selenium/grid/distributor/local",
         "//java/server/src/org/openqa/selenium/grid/distributor/remote",
         "//java/server/src/org/openqa/selenium/grid/distributor/selector",

java/server/test/org/openqa/selenium/grid/distributor/DistributorTest.java

@@ -43,6 +43,7 @@
 import org.openqa.selenium.grid.data.Session;
 import org.openqa.selenium.grid.data.SessionRequest;
 import org.openqa.selenium.grid.data.Slot;
+import org.openqa.selenium.grid.distributor.gridmodel.local.LocalGridModel;
 import org.openqa.selenium.grid.distributor.local.LocalDistributor;
 import org.openqa.selenium.grid.distributor.remote.RemoteDistributor;
 import org.openqa.selenium.grid.distributor.selector.DefaultSlotSelector;
@@ -135,6 +136,7 @@ public void setUp() throws URISyntaxException {
       HttpClient.Factory.createDefault(),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -173,6 +175,7 @@ public void shouldStartHeartBeatOnNodeRegistration() {
       new PassthroughHttpClient.Factory(node),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -222,6 +225,7 @@ public void shouldBeAbleToAddANodeAndCreateASession() {
       new PassthroughHttpClient.Factory(node),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -263,6 +267,7 @@ public void creatingASessionAddsItToTheSessionMap() {
       new PassthroughHttpClient.Factory(node),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -305,6 +310,7 @@ public void shouldBeAbleToRemoveANode() throws MalformedURLException {
       new PassthroughHttpClient.Factory(node),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -344,6 +350,7 @@ public void testDrainingNodeDoesNotAcceptNewSessions() {
       new PassthroughHttpClient.Factory(node),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -382,6 +389,7 @@ public void testDrainedNodeShutsDownOnceEmpty() throws InterruptedException {
       new PassthroughHttpClient.Factory(node),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -427,6 +435,7 @@ public void drainedNodeDoesNotShutDownIfNotEmpty() throws InterruptedException {
       new PassthroughHttpClient.Factory(node),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -475,6 +484,7 @@ public void drainedNodeShutsDownAfterSessionsFinish() throws InterruptedExceptio
       new PassthroughHttpClient.Factory(node),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -548,6 +558,7 @@ public void theMostLightlyLoadedNodeIsSelectedFirst() {
       new PassthroughHttpClient.Factory(handler),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -591,6 +602,7 @@ public void shouldUseLastSessionCreatedTimeAsTieBreaker() {
       new PassthroughHttpClient.Factory(handler),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -683,6 +695,7 @@ public void shouldIncludeHostsThatAreUpInHostList() {
       new PassthroughHttpClient.Factory(handler),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofSeconds(1),
@@ -724,6 +737,7 @@ public void shouldNotScheduleAJobIfAllSlotsAreBeingUsed() {
       new PassthroughHttpClient.Factory(node),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -764,6 +778,7 @@ public void shouldReleaseSlotOnceSessionEnds() {
       new PassthroughHttpClient.Factory(node),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -819,6 +834,7 @@ public void shouldNotStartASessionIfTheCapabilitiesAreNotSupported() {
       new PassthroughHttpClient.Factory(handler),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -859,6 +875,7 @@ public void attemptingToStartASessionWhichFailsMarksAsTheSlotAsAvailable() {
       new PassthroughHttpClient.Factory(node),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -904,6 +921,7 @@ public void shouldReturnNodesThatWereDownToPoolOfNodesOnceTheyMarkTheirHealthChe
       new PassthroughHttpClient.Factory(handler),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofSeconds(1),
@@ -969,6 +987,7 @@ public void shouldPrioritizeHostsWithTheMostSlotsAvailableForASessionType() {
       new PassthroughHttpClient.Factory(handler),
       sessions,
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),

java/server/test/org/openqa/selenium/grid/distributor/local/BUILD.bazel

@@ -16,6 +16,7 @@ java_test_suite(
         "//java/server/src/org/openqa/selenium/events/local",
         "//java/server/src/org/openqa/selenium/grid/data",
         "//java/server/src/org/openqa/selenium/grid/distributor",
+        "//java/server/src/org/openqa/selenium/grid/distributor/gridmodel/local",
         "//java/server/src/org/openqa/selenium/grid/distributor/local",
         "//java/server/src/org/openqa/selenium/grid/distributor/selector",
         "//java/server/src/org/openqa/selenium/grid/node",

java/server/test/org/openqa/selenium/grid/distributor/local/LocalDistributorTest.java

@@ -32,6 +32,7 @@
 import org.openqa.selenium.grid.data.RequestId;
 import org.openqa.selenium.grid.data.Session;
 import org.openqa.selenium.grid.distributor.Distributor;
+import org.openqa.selenium.grid.distributor.gridmodel.local.LocalGridModel;
 import org.openqa.selenium.grid.distributor.selector.DefaultSlotSelector;
 import org.openqa.selenium.grid.node.Node;
 import org.openqa.selenium.grid.node.local.LocalNode;
@@ -111,6 +112,7 @@ public void testAddNodeToDistributor() {
       clientFactory,
       new LocalSessionMap(tracer, bus),
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -144,6 +146,7 @@ public void testShouldNotAddNodeWithWrongSecret() {
       clientFactory,
       new LocalSessionMap(tracer, bus),
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       secret,
       Duration.ofMinutes(5),
@@ -171,6 +174,7 @@ public void testRemoveNodeFromDistributor() {
       clientFactory,
       new LocalSessionMap(tracer, bus),
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -204,6 +208,7 @@ public void testAddSameNodeTwice() {
       clientFactory,
       new LocalSessionMap(tracer, bus),
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -232,6 +237,7 @@ public void shouldBeAbleToAddMultipleSessionsConcurrently() throws Exception {
       clientFactory,
       new LocalSessionMap(tracer, bus),
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -314,6 +320,7 @@ public void testDrainNodeFromDistributor() {
       clientFactory,
       new LocalSessionMap(tracer, bus),
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),
@@ -354,6 +361,7 @@ public void testDrainNodeFromNode() {
       clientFactory,
       new LocalSessionMap(tracer, bus),
       queue,
+      new LocalGridModel(bus),
       new DefaultSlotSelector(),
       registrationSecret,
       Duration.ofMinutes(5),

java/server/test/org/openqa/selenium/grid/distributor/local/LocalGridModelTest.java

@@ -22,6 +22,7 @@
 import org.openqa.selenium.events.local.GuavaEventBus;
 import org.openqa.selenium.grid.data.DefaultSlotMatcher;
 import org.openqa.selenium.grid.distributor.Distributor;
+import org.openqa.selenium.grid.distributor.gridmodel.local.LocalGridModel;
 import org.openqa.selenium.grid.distributor.selector.DefaultSlotSelector;
 import org.openqa.selenium.grid.security.Secret;
 import org.openqa.selenium.grid.sessionmap.SessionMap;
@@ -54,6 +55,7 @@ public class LocalGridModelTest {
     clientFactory,
     sessions,
     queue,
+    new LocalGridModel(events),
     new DefaultSlotSelector(),
     secret,
     Duration.ofMinutes(5),

java/server/test/org/openqa/selenium/grid/graphql/BUILD.bazel

@@ -18,6 +18,7 @@ java_test_suite(
         "//java/server/src/org/openqa/selenium/events/local",
         "//java/server/src/org/openqa/selenium/grid/data",
         "//java/server/src/org/openqa/selenium/grid/distributor",
+        "//java/server/src/org/openqa/selenium/grid/distributor/gridmodel/local",
         "//java/server/src/org/openqa/selenium/grid/distributor/local",
         "//java/server/src/org/openqa/selenium/grid/distributor/selector",
         "//java/server/src/org/openqa/selenium/grid/graphql",