[grid] Add new session request queue size endpoint and GraphQL support (#9030)

* [grid] Add new session request queue size endpoint and GraphQL support

* [grid] Remove unused imports from GraphQL classes. Fix assert for NewSessionQueuer tests.

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

Author: pujagani

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

@@ -165,7 +165,11 @@ protected Handlers createHandlers(Config config) {
     handler.addHandler(distributor);
 
     Router router = new Router(tracer, clientFactory, sessions, queuer, distributor);
-    GraphqlHandler graphqlHandler = new GraphqlHandler(tracer, distributor, serverOptions.getExternalUri());
+    GraphqlHandler graphqlHandler = new GraphqlHandler(
+      tracer,
+      distributor,
+      queuer,
+      serverOptions.getExternalUri());
     HttpHandler readinessCheck = req -> {
       boolean ready = router.isReady() && bus.isReady();
       return new HttpResponse()

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

@@ -175,7 +175,11 @@ protected Handlers createHandlers(Config config) {
         .setContent(Contents.utf8String("Standalone is " + ready));
     };
 
-    GraphqlHandler graphqlHandler = new GraphqlHandler(tracer, distributor, serverOptions.getExternalUri());
+    GraphqlHandler graphqlHandler = new GraphqlHandler(
+      tracer,
+      distributor,
+      queuer,
+      serverOptions.getExternalUri());
 
     Routable ui;
     URL uiRoot = getClass().getResource("/javascript/grid-ui/build");

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

@@ -18,6 +18,7 @@ java_library(
         "//java/client/src/org/openqa/selenium/remote",
         "//java/server/src/org/openqa/selenium/grid/data",
         "//java/server/src/org/openqa/selenium/grid/distributor",
+        "//java/server/src/org/openqa/selenium/grid/sessionqueue",
         artifact("com.google.guava:guava"),
         artifact("com.graphql-java:graphql-java"),
     ],

java/server/src/org/openqa/selenium/grid/graphql/GraphqlHandler.java

@@ -29,6 +29,7 @@
 import graphql.schema.idl.SchemaParser;
 import graphql.schema.idl.TypeDefinitionRegistry;
 import org.openqa.selenium.grid.distributor.Distributor;
+import org.openqa.selenium.grid.sessionqueue.NewSessionQueuer;
 import org.openqa.selenium.internal.Require;
 import org.openqa.selenium.json.Json;
 import org.openqa.selenium.remote.http.Contents;
@@ -68,12 +69,15 @@ public class GraphqlHandler implements HttpHandler {
   public static final Json JSON = new Json();
   private final Tracer tracer;
   private final Distributor distributor;
+  private final NewSessionQueuer newSessionQueuer;
   private final URI publicUri;
   private final GraphQL graphQl;
 
 
-  public GraphqlHandler(Tracer tracer, Distributor distributor, URI publicUri) {
+  public GraphqlHandler(Tracer tracer, Distributor distributor, NewSessionQueuer newSessionQueuer,
+                        URI publicUri) {
     this.distributor = Require.nonNull("Distributor", distributor);
+    this.newSessionQueuer = Require.nonNull("NewSessionQueuer", newSessionQueuer);
     this.publicUri = Require.nonNull("Uri", publicUri);
     this.tracer = Require.nonNull("Tracer", tracer);
 
@@ -169,7 +173,7 @@ private RuntimeWiring buildRuntimeWiring() {
       .scalar(Types.Uri)
       .scalar(Types.Url)
       .type("GridQuery", typeWiring -> typeWiring
-          .dataFetcher("grid", new GridData(distributor, publicUri))
+          .dataFetcher("grid", new GridData(distributor, newSessionQueuer, publicUri))
           .dataFetcher("session", new SessionData(distributor)))
       .build();
   }

java/server/src/org/openqa/selenium/grid/graphql/Grid.java

@@ -24,6 +24,7 @@
 import org.openqa.selenium.grid.data.NodeStatus;
 import org.openqa.selenium.grid.data.Slot;
 import org.openqa.selenium.grid.distributor.Distributor;
+import org.openqa.selenium.grid.sessionqueue.NewSessionQueuer;
 import org.openqa.selenium.internal.Require;
 
 import java.net.URI;
@@ -37,10 +38,12 @@ public class Grid {
 
   private final URI uri;
   private final Supplier<DistributorStatus> distributorStatus;
+  private final NewSessionQueuer newSessionQueuer;
 
-  public Grid(Distributor distributor, URI uri) {
+  public Grid(Distributor distributor, NewSessionQueuer newSessionQueuer, URI uri) {
     Require.nonNull("Distributor", distributor);
     this.uri = Require.nonNull("Grid's public URI", uri);
+    this.newSessionQueuer = Require.nonNull("NewSessionQueuer", newSessionQueuer);
     this.distributorStatus = Suppliers.memoize(distributor::getStatus);
   }
 
@@ -96,4 +99,8 @@ public int getTotalSlots() {
   public int getUsedSlots() {
     return getSessionCount();
   }
+
+  public int getSessionQueueSize() {
+    return newSessionQueuer.getQueueSize();
+  }
 }

java/server/src/org/openqa/selenium/grid/graphql/GridData.java

@@ -20,21 +20,24 @@
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import org.openqa.selenium.grid.distributor.Distributor;
+import org.openqa.selenium.grid.sessionqueue.NewSessionQueuer;
 import org.openqa.selenium.internal.Require;
 
 import java.net.URI;
 
 public class GridData implements DataFetcher {
   private final Distributor distributor;
+  private final NewSessionQueuer newSessionQueuer;
   private final URI publicUri;
 
-  public GridData(Distributor distributor, URI publicUri) {
+  public GridData(Distributor distributor, NewSessionQueuer newSessionQueuer, URI publicUri) {
     this.distributor = Require.nonNull("Distributor", distributor);
     this.publicUri = Require.nonNull("Grid's public URI", publicUri);
+    this.newSessionQueuer = Require.nonNull("NewSessionQueuer", newSessionQueuer);
   }
 
   @Override
   public Object get(DataFetchingEnvironment environment) {
-    return new Grid(distributor, publicUri);
+    return new Grid(distributor, newSessionQueuer, publicUri);
   }
 }

java/server/src/org/openqa/selenium/grid/graphql/selenium-grid-schema.graphqls

@@ -47,4 +47,5 @@ type Grid {
   totalSlots: Int
   usedSlots: Int
   sessionCount: Int!
+  sessionQueueSize: Int!
 }

java/server/src/org/openqa/selenium/grid/router/httpd/RouterServer.java

@@ -133,7 +133,11 @@ protected Handlers createHandlers(Config config) {
       distributorUrl,
       secret);
 
-    GraphqlHandler graphqlHandler = new GraphqlHandler(tracer, distributor, serverOptions.getExternalUri());
+    GraphqlHandler graphqlHandler = new GraphqlHandler(
+      tracer,
+      distributor,
+      queuer,
+      serverOptions.getExternalUri());
 
     Route handler = Route.combine(
       new Router(tracer, clientFactory, sessions, queuer, distributor).with(networkOptions.getSpecComplianceChecks()),

java/server/src/org/openqa/selenium/grid/sessionqueue/GetNewSessionQueueSize.java

@@ -0,0 +1,61 @@
+// 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.sessionqueue;
+
+import com.google.common.collect.ImmutableMap;
+import org.openqa.selenium.internal.Require;
+import org.openqa.selenium.remote.http.HttpHandler;
+import org.openqa.selenium.remote.http.HttpRequest;
+import org.openqa.selenium.remote.http.HttpResponse;
+import org.openqa.selenium.remote.tracing.Span;
+import org.openqa.selenium.remote.tracing.Tracer;
+
+import static org.openqa.selenium.remote.http.Contents.asJson;
+import static org.openqa.selenium.remote.tracing.HttpTracing.newSpanAsChildOf;
+import static org.openqa.selenium.remote.tracing.Tags.HTTP_REQUEST;
+import static org.openqa.selenium.remote.tracing.Tags.HTTP_RESPONSE;
+
+class GetNewSessionQueueSize implements HttpHandler {
+
+  private final Tracer tracer;
+  private final NewSessionQueuer newSessionQueuer;
+
+  GetNewSessionQueueSize(Tracer tracer, NewSessionQueuer newSessionQueuer) {
+    this.tracer = Require.nonNull("Tracer", tracer);
+    this.newSessionQueuer = Require.nonNull("New Session Queuer", newSessionQueuer);
+  }
+
+  @Override
+  public HttpResponse execute(HttpRequest req) {
+    try (Span span = newSpanAsChildOf(tracer, req, "sessionqueue.size")) {
+      HTTP_REQUEST.accept(span, req);
+
+      int value = newSessionQueuer.getQueueSize();
+
+      span.setAttribute("request.retry", value);
+
+      HttpResponse response = new HttpResponse()
+        .setContent(asJson(ImmutableMap.of("value", value)));
+
+      HTTP_RESPONSE.accept(span, response);
+
+      return response;
+    }
+  }
+}
+

java/server/src/org/openqa/selenium/grid/sessionqueue/NewSessionQueue.java

@@ -48,6 +48,8 @@ public abstract class NewSessionQueue implements HasReadyState {
 
   public abstract int clear();
 
+  public abstract int getQueueSize();
+
   public void addRequestHeaders(HttpRequest request, RequestId reqId) {
     long timestamp = Instant.now().getEpochSecond();
     request.addHeader(SESSIONREQUEST_TIMESTAMP_HEADER, Long.toString(timestamp));

java/server/src/org/openqa/selenium/grid/sessionqueue/NewSessionQueuer.java

@@ -63,6 +63,7 @@ protected NewSessionQueuer(Tracer tracer, Secret registrationSecret) {
     RequiresSecretFilter requiresSecret = new RequiresSecretFilter(registrationSecret);
 
     routes = combine(
+
       post("/session")
         .to(() -> this::addToQueue),
       post("/se/grid/newsessionqueuer/session")
@@ -73,6 +74,8 @@ protected NewSessionQueuer(Tracer tracer, Secret registrationSecret) {
       get("/se/grid/newsessionqueuer/session/{requestId}")
         .to(params -> new RemoveFromSessionQueue(tracer, this, requestIdFrom(params)))
         .with(requiresSecret),
+      get("/se/grid/newsessionqueuer/queue/size")
+        .to(() -> new GetNewSessionQueueSize(tracer, this)),
       delete("/se/grid/newsessionqueuer/queue")
         .to(() -> new ClearSessionQueue(tracer, this))
         .with(requiresSecret));
@@ -122,6 +125,8 @@ public void validateSessionRequest(HttpRequest request) {
 
   public abstract int clearQueue();
 
+  public abstract int getQueueSize();
+
   @Override
   public boolean matches(HttpRequest req) {
     return routes.matches(req);

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

@@ -88,6 +88,7 @@ public boolean isReady() {
     return bus.isReady();
   }
 
+  @Override
   @ManagedAttribute(name = "NewSessionQueueSize")
   public int getQueueSize() {
     Lock readLock = lock.readLock();

java/server/src/org/openqa/selenium/grid/sessionqueue/local/LocalNewSessionQueuer.java

@@ -91,6 +91,11 @@ public int clearQueue() {
     return sessionRequests.clear();
   }
 
+  @Override
+  public int getQueueSize() {
+    return sessionRequests.getQueueSize();
+  }
+
   @Override
   public boolean isReady() {
     return bus.isReady();

java/server/src/org/openqa/selenium/grid/sessionqueue/remote/RemoteNewSessionQueuer.java

@@ -129,6 +129,15 @@ public int clearQueue() {
     return Values.get(response, Integer.class);
   }
 
+  @Override
+  public int getQueueSize() {
+    HttpRequest upstream =
+      new HttpRequest(GET, "/se/grid/newsessionqueuer/queue/size");
+    HttpTracing.inject(tracer, tracer.getCurrentContext(), upstream);
+    HttpResponse response = client.execute(upstream);
+    return Values.get(response, Integer.class);
+  }
+
   @Override
   public boolean isReady() {
     try {