[grid] Add endpoint to return new session request queue contents (#9078)

* [grid] Add endpoint to return session queue contents

* [grid] Update session queue contents tests.

* [grid] Fix imports in LocalNewSessionQueue

* [grid] Remove unused import

* [grid] Remove unused json import from LocalNewSessionQueue

* [grid] Rename the session queue contents endpoint.

* [grid] Remove queue size endpoint

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

Author: pujagani

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

@@ -99,8 +99,4 @@ public int getTotalSlots() {
   public int getUsedSlots() {
     return getSessionCount();
   }
-
-  public int getSessionQueueSize() {
-    return newSessionQueuer.getQueueSize();
-  }
 }

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

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

java/server/src/org/openqa/selenium/grid/sessionqueue/GetNewSessionQueueSize.java renamed to java/server/src/org/openqa/selenium/grid/sessionqueue/GetSessionQueue.java

@@ -17,40 +17,38 @@
 
 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 java.util.Collections;
+
 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 {
+class GetSessionQueue implements HttpHandler {
 
   private final Tracer tracer;
   private final NewSessionQueuer newSessionQueuer;
 
-  GetNewSessionQueueSize(Tracer tracer, NewSessionQueuer newSessionQueuer) {
+  GetSessionQueue(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")) {
+    try (Span span = newSpanAsChildOf(tracer, req, "sessionqueue.contents")) {
       HTTP_REQUEST.accept(span, req);
 
-      int value = newSessionQueuer.getQueueSize();
-
-      span.setAttribute("request.retry", value);
-
       HttpResponse response = new HttpResponse()
-        .setContent(asJson(ImmutableMap.of("value", value)));
+        .setContent(asJson(Collections.singletonMap(
+          "value", newSessionQueuer.getQueueContents())));
 
       HTTP_RESPONSE.accept(span, response);
 

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

@@ -26,6 +26,7 @@
 
 import java.time.Duration;
 import java.time.Instant;
+import java.util.Map;
 import java.util.Optional;
 
 public abstract class NewSessionQueue implements HasReadyState {
@@ -50,6 +51,8 @@ public abstract class NewSessionQueue implements HasReadyState {
 
   public abstract int getQueueSize();
 
+  public abstract Map<String,Object> getQueueContents();
+
   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

@@ -74,8 +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)),
+      get("/se/grid/newsessionqueuer/queue")
+        .to(() -> new GetSessionQueue(tracer, this)),
       delete("/se/grid/newsessionqueuer/queue")
         .to(() -> new ClearSessionQueue(tracer, this))
         .with(requiresSecret));
@@ -125,7 +125,7 @@ public void validateSessionRequest(HttpRequest request) {
 
   public abstract int clearQueue();
 
-  public abstract int getQueueSize();
+  public abstract Map<String, Object> getQueueContents();
 
   @Override
   public boolean matches(HttpRequest req) {

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

@@ -17,6 +17,8 @@
 
 package org.openqa.selenium.grid.sessionqueue.local;
 
+import com.google.common.collect.ImmutableMap;
+import org.openqa.selenium.Capabilities;
 import org.openqa.selenium.events.EventBus;
 import org.openqa.selenium.grid.config.Config;
 import org.openqa.selenium.grid.data.NewSessionErrorResponse;
@@ -28,6 +30,7 @@
 import org.openqa.selenium.grid.sessionqueue.NewSessionQueue;
 import org.openqa.selenium.grid.sessionqueue.config.NewSessionQueueOptions;
 import org.openqa.selenium.internal.Require;
+import org.openqa.selenium.remote.NewSessionPayload;
 import org.openqa.selenium.remote.http.HttpRequest;
 
 import org.openqa.selenium.remote.server.jmx.JMXHelper;
@@ -40,11 +43,16 @@
 import org.openqa.selenium.remote.tracing.Span;
 import org.openqa.selenium.remote.tracing.Tracer;
 
+import java.io.IOException;
+import java.io.Reader;
 import java.time.Duration;
 import java.util.Deque;
 import java.util.HashMap;
+import java.util.List;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Objects;
 import java.util.concurrent.ConcurrentLinkedDeque;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
@@ -54,6 +62,9 @@
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import java.util.stream.Collectors;
+
+import static org.openqa.selenium.remote.http.Contents.reader;
 
 @ManagedService(objectName = "org.seleniumhq.grid:type=SessionQueue,name=LocalSessionQueue",
   description = "New session queue")
@@ -100,6 +111,36 @@ public int getQueueSize() {
     }
   }
 
+  @Override
+  public Map<String, Object> getQueueContents() {
+    Lock readLock = lock.readLock();
+    readLock.lock();
+    try {
+      List<Capabilities> capabilitiesList = sessionRequests.stream()
+        .map(SessionRequest::getHttpRequest)
+        .map(req -> {
+          try (
+            Reader reader = reader(req);
+            NewSessionPayload payload = NewSessionPayload.create(reader)) {
+            return payload.stream().iterator();
+          } catch (IOException e) {
+            LOG.warning("IOException while mapping to capabilities" + e.getMessage());
+          }
+          return null;
+        })
+        .filter(Objects::nonNull)
+        .filter(Iterator::hasNext)
+        .map(Iterator::next)
+        .collect(Collectors.toList());
+
+      return ImmutableMap.of(
+        "request-count", capabilitiesList.size(),
+        "request-payloads", capabilitiesList);
+    } finally {
+      readLock.unlock();
+    }
+  }
+
   @Override
   public boolean offerLast(HttpRequest request, RequestId requestId) {
     Require.nonNull("New Session request", request);

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

@@ -34,6 +34,7 @@
 import org.openqa.selenium.remote.tracing.Tracer;
 
 import java.time.Duration;
+import java.util.Map;
 import java.util.Optional;
 
 public class LocalNewSessionQueuer extends NewSessionQueuer {
@@ -92,8 +93,8 @@ public int clearQueue() {
   }
 
   @Override
-  public int getQueueSize() {
-    return sessionRequests.getQueueSize();
+  public Map<String, Object> getQueueContents() {
+    return sessionRequests.getQueueContents();
   }
 
   @Override

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

@@ -45,6 +45,7 @@
 import java.io.UncheckedIOException;
 import java.net.MalformedURLException;
 import java.net.URI;
+import java.util.Map;
 import java.util.Optional;
 import java.util.logging.Logger;
 
@@ -130,12 +131,11 @@ public int clearQueue() {
   }
 
   @Override
-  public int getQueueSize() {
-    HttpRequest upstream =
-      new HttpRequest(GET, "/se/grid/newsessionqueuer/queue/size");
+  public Map<String, Object> getQueueContents() {
+    HttpRequest upstream = new HttpRequest(GET, "/se/grid/newsessionqueuer/queue");
     HttpTracing.inject(tracer, tracer.getCurrentContext(), upstream);
     HttpResponse response = client.execute(upstream);
-    return Values.get(response, Integer.class);
+    return Values.get(response, Map.class);
   }
 
   @Override

java/server/test/org/openqa/selenium/grid/graphql/GraphqlHandlerTest.java

@@ -141,26 +141,6 @@ public void shouldBeAbleToGetGridUri() {
             "uri", publicUri.toString()))));
   }
 
-  @Test
-  public void shouldBeAbleToGetSessionQueueSize() {
-    HttpRequest request = new HttpRequest(POST, "/session");
-    request.setContent(asJson(
-      ImmutableMap.of(
-        "capabilities", ImmutableMap.of(
-          "alwaysMatch", caps))));
-
-    localNewSessionQueue.offerLast(request, new RequestId(UUID.randomUUID()));
-    GraphqlHandler handler = new GraphqlHandler(tracer, distributor, queuer, publicUri);
-
-    Map<String, Object> topLevel = executeQuery(handler, "{ grid { sessionQueueSize } }");
-
-    assertThat(topLevel).isEqualTo(
-      singletonMap(
-        "data", singletonMap(
-          "grid", singletonMap(
-            "sessionQueueSize", 1L))));
-  }
-
   @Test
   public void shouldReturnAnEmptyListForNodesIfNoneAreRegistered() {
     GraphqlHandler handler = new GraphqlHandler(tracer, distributor, queuer, publicUri);

java/server/test/org/openqa/selenium/grid/sessionqueue/NewSessionQueuerTest.java

@@ -55,6 +55,8 @@
 import java.net.URISyntaxException;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.Callable;
@@ -229,21 +231,31 @@ public void shouldBeClearQueueRemotely() {
   }
 
   @Test
-  public void shouldBeAbleToGetQueueSize() {
+  public void shouldBeAbleToGetQueueContents() {
     RequestId requestId = new RequestId(UUID.randomUUID());
     sessionQueue.offerLast(request, requestId);
 
-    int size = local.getQueueSize();
-    assertEquals(1, size);
+    Map<String, Object> response = local.getQueueContents();
+    assertThat(response).isNotNull();
+
+    assertEquals(1, response.get("request-count"));
+
+    List<Capabilities> capabilitiesList = (List<Capabilities>) response.get("request-payloads");
+    assertEquals(caps, capabilitiesList.get(0));
   }
 
   @Test
-  public void shouldBeAbleToGetQueueSizeRemotely() {
+  public void shouldBeAbleToGetQueueContentsRemotely() {
     RequestId requestId = new RequestId(UUID.randomUUID());
     sessionQueue.offerLast(request, requestId);
 
-    int size = sessionQueue.getQueueSize();
-    assertEquals(1, size);
+    Map<String, Object> response = sessionQueue.getQueueContents();
+    assertThat(response).isNotNull();
+
+    assertEquals(1, response.get("request-count"));
+
+    List<Capabilities> capabilitiesList = (List<Capabilities>) response.get("request-payloads");
+    assertEquals(caps, capabilitiesList.get(0));
   }
 
   @Test

java/server/test/org/openqa/selenium/grid/sessionqueue/local/LocalNewSessionQueueTest.java

@@ -36,6 +36,8 @@
 import java.io.UncheckedIOException;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
@@ -224,6 +226,42 @@ public void shouldBeAbleToGetQueueSize() {
     assertEquals(1, size);
   }
 
+  @Test
+  public void shouldBeAbleToGetQueueContents() {
+    long timestamp = Instant.now().getEpochSecond();
+
+    ImmutableCapabilities chromeCaps = new ImmutableCapabilities(
+      "browserName", "chrome",
+      "platform", "mac",
+      "version", "87");
+    NewSessionPayload chromePayload = NewSessionPayload.create(chromeCaps);
+    HttpRequest chromeRequest = createRequest(chromePayload, POST, "/session");
+    chromeRequest.addHeader(SESSIONREQUEST_TIMESTAMP_HEADER, Long.toString(timestamp));
+    RequestId chromeRequestId = new RequestId(UUID.randomUUID());
+    boolean addedChromeRequest = sessionQueue.offerLast(chromeRequest, chromeRequestId);
+    assertTrue(addedChromeRequest);
+
+    ImmutableCapabilities firefoxCaps = new ImmutableCapabilities(
+      "browserName", "firefox",
+      "platform", "windows",
+      "version", "84");
+    NewSessionPayload firefoxPayload = NewSessionPayload.create(firefoxCaps);
+    HttpRequest firefoxRequest = createRequest(firefoxPayload, POST, "/session");
+    firefoxRequest.addHeader(SESSIONREQUEST_TIMESTAMP_HEADER, Long.toString(timestamp));
+    RequestId firefoxRequestId = new RequestId(UUID.randomUUID());
+    boolean addFirefoxRequest = sessionQueue.offerLast(firefoxRequest, firefoxRequestId);
+    assertTrue(addFirefoxRequest);
+
+    Map<String, Object> response = sessionQueue.getQueueContents();
+    assertThat(response).isNotNull();
+
+    assertEquals(2, response.get("request-count"));
+
+    List<Capabilities> capabilitiesList = (List<Capabilities>) response.get("request-payloads");
+    assertEquals(chromeCaps, capabilitiesList.get(0));
+    assertEquals(firefoxCaps, capabilitiesList.get(1));
+  }
+
   private HttpRequest createRequest(NewSessionPayload payload, HttpMethod httpMethod, String uri) {
     StringBuilder builder = new StringBuilder();
     try {