Expose session request metadata from the new session payload

java/client/src/org/openqa/selenium/remote/NewSessionPayload.java

@@ -322,6 +322,37 @@ public Set<Dialect> getDownstreamDialects() {
     return dialects.isEmpty() ? ImmutableSet.of(DEFAULT_DIALECT) : dialects;
   }
 
+  public Map<String, Object> getMetadata() {
+    Set<String> ignoredMetadataKeys = ImmutableSet.of("capabilities", "desiredCapabilities");
+
+    CharSource charSource = backingStore.asByteSource().asCharSource(UTF_8);
+    try (Reader reader = charSource.openBufferedStream();
+         JsonInput input = json.newInput(reader)) {
+      ImmutableMap.Builder<String, Object> toReturn = ImmutableMap.builder();
+
+      input.beginObject();
+      while (input.hasNext()) {
+        String name = input.nextName();
+        if (ignoredMetadataKeys.contains(name)) {
+          input.skipValue();
+          continue;
+        }
+
+        Object value = input.read(Object.class);
+        if (value == null) {
+          continue;
+        }
+
+        toReturn.put(name, value);
+      }
+      input.endObject();
+
+      return toReturn.build();
+    } catch (IOException e) {
+      throw new UncheckedIOException(e);
+    }
+  }
+
   @Override
   public void close() {
     try {

java/client/test/org/openqa/selenium/remote/NewSessionPayloadTest.java

@@ -42,6 +42,7 @@
 
 import java.io.IOException;
 import java.io.StringReader;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -271,6 +272,43 @@ public void shouldPreserveMetadata() throws IOException {
     }
   }
 
+  @Test
+  public void shouldExposeMetaData() {
+    Map<String, Object> raw = ImmutableMap.of(
+      "capabilities", singletonMap("alwaysMatch", singletonMap("browserName", "cheese")),
+      "se:meta", "cheese is good");
+
+    try (NewSessionPayload payload = NewSessionPayload.create(raw)) {
+      Map<String, Object> seen = payload.getMetadata();
+      assertThat(seen).isEqualTo(Map.of("se:meta", "cheese is good"));
+    }
+  }
+
+  @Test
+  public void nullValuesInMetaDataAreIgnored() {
+    Map<String, Object> raw = new HashMap<>();
+    raw.put("capabilities", singletonMap("alwaysMatch", singletonMap("browserName", "cheese")));
+    raw.put("se:bad", null);
+    raw.put("se:good", "cheese");
+
+    try (NewSessionPayload payload = NewSessionPayload.create(raw)) {
+      Map<String, Object> seen = payload.getMetadata();
+      assertThat(seen).isEqualTo(Map.of("se:good", "cheese"));
+    }
+  }
+
+  @Test
+  public void keysUsedForStoringCapabilitiesAreIgnoredFromMetadata() {
+    Map<String, Object> raw = ImmutableMap.of(
+      "capabilities", singletonMap("alwaysMatch", singletonMap("browserName", "cheese")),
+      "desiredCapabilities", emptyMap());
+
+    try (NewSessionPayload payload = NewSessionPayload.create(raw)) {
+      Map<String, Object> seen = payload.getMetadata();
+      assertThat(seen).isEqualTo(emptyMap());
+    }
+  }
+
   private List<Capabilities> create(Map<String, ?> source) {
     List<Capabilities> presumablyFromMemory;
     List<Capabilities> fromDisk;

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

@@ -70,7 +70,8 @@ protected NewSessionQueue(Tracer tracer, Secret registrationSecret) {
               new RequestId(UUID.randomUUID()),
               Instant.now(),
               payload.getDownstreamDialects(),
-              payload.stream().collect(Collectors.toSet()));
+              payload.stream().collect(Collectors.toSet()),
+              payload.getMetadata());
             return addToQueue(sessionRequest);
           } catch (IOException e) {
             throw new UncheckedIOException(e);

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

@@ -18,6 +18,7 @@
 package org.openqa.selenium.grid.sessionqueue;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import org.openqa.selenium.Capabilities;
 import org.openqa.selenium.grid.data.RequestId;
 import org.openqa.selenium.internal.Require;
@@ -34,6 +35,9 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static java.util.Collections.emptyMap;
+import static org.openqa.selenium.json.Json.MAP_TYPE;
+
 public class SessionRequest {
 
   private static final Type SET_OF_CAPABILITIES = new TypeToken<Set<Capabilities>>() {}.getType();
@@ -42,6 +46,7 @@ public class SessionRequest {
   private final Instant enqueued;
   private final Set<Capabilities> desiredCapabilities;
   private final Set<Dialect> downstreamDialects;
+  private final Map<String, Object> metadata;
 
   public SessionRequest(RequestId requestId, HttpRequest request, Instant enqueued) {
     this.requestId = Require.nonNull("Request ID", requestId);
@@ -51,18 +56,21 @@ public SessionRequest(RequestId requestId, HttpRequest request, Instant enqueued
     try (NewSessionPayload payload = NewSessionPayload.create(Contents.reader(request))) {
       desiredCapabilities = payload.stream().collect(Collectors.toSet());
       downstreamDialects = payload.getDownstreamDialects();
+      metadata = payload.getMetadata();
     }
   }
 
   public SessionRequest(
     RequestId requestId,
     Instant enqueued,
     Set<Dialect> downstreamDialects,
-    Set<Capabilities> desiredCapabilities) {
+    Set<Capabilities> desiredCapabilities,
+    Map<String, Object> metadata) {
     this.requestId = Require.nonNull("Request ID", requestId);
     this.enqueued = Require.nonNull("Enqueud time", enqueued);
-    this.downstreamDialects = Require.nonNull("Downstream dialects", downstreamDialects);
-    this.desiredCapabilities = Require.nonNull("Capabilities", desiredCapabilities);
+    this.downstreamDialects = ImmutableSet.copyOf(Require.nonNull("Downstream dialects", downstreamDialects));
+    this.desiredCapabilities = ImmutableSet.copyOf(Require.nonNull("Capabilities", desiredCapabilities));
+    this.metadata = ImmutableMap.copyOf(Require.nonNull("Metadata", metadata));
   }
 
   public RequestId getRequestId() {
@@ -86,14 +94,16 @@ private Map<String, Object> toJson() {
       "requestId", requestId,
       "enqueued", enqueued,
       "dialects", downstreamDialects,
-      "capabilities", desiredCapabilities);
+      "capabilities", desiredCapabilities,
+      "metadata", metadata);
   }
 
   private static SessionRequest fromJson(JsonInput input) {
     RequestId id = null;
     Instant enqueued = null;
     Set<Capabilities> capabilities = null;
     Set<Dialect> dialects = null;
+    Map<String, Object> metadata = emptyMap();
 
     input.beginObject();
     while (input.hasNext()) {
@@ -110,6 +120,10 @@ private static SessionRequest fromJson(JsonInput input) {
           enqueued = input.read(Instant.class);
           break;
 
+        case "metadata":
+          metadata = input.read(MAP_TYPE);
+          break;
+
         case "requestId":
           id = input.read(RequestId.class);
           break;
@@ -121,6 +135,6 @@ private static SessionRequest fromJson(JsonInput input) {
     }
     input.endObject();
 
-    return new SessionRequest(id, enqueued, dialects, capabilities);
+    return new SessionRequest(id, enqueued, dialects, capabilities, metadata);
   }
 }

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

@@ -59,7 +59,6 @@
 import org.openqa.selenium.internal.Either;
 import org.openqa.selenium.net.PortProber;
 import org.openqa.selenium.remote.Dialect;
-import org.openqa.selenium.remote.NewSessionPayload;
 import org.openqa.selenium.remote.SessionId;
 import org.openqa.selenium.remote.http.HttpClient;
 import org.openqa.selenium.remote.http.HttpHandler;
@@ -70,7 +69,6 @@
 import org.openqa.selenium.support.ui.FluentWait;
 import org.openqa.selenium.support.ui.Wait;
 
-import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.net.MalformedURLException;
 import java.net.URI;
@@ -97,8 +95,6 @@
 import static org.openqa.selenium.grid.data.Availability.DOWN;
 import static org.openqa.selenium.grid.data.Availability.UP;
 import static org.openqa.selenium.remote.Dialect.W3C;
-import static org.openqa.selenium.remote.http.Contents.utf8String;
-import static org.openqa.selenium.remote.http.HttpMethod.POST;
 
 public class DistributorTest {
 
@@ -1043,21 +1039,8 @@ private SessionRequest createRequest(Capabilities... allCaps) {
       new RequestId(UUID.randomUUID()),
       Instant.now(),
       Set.of(W3C),
-      Set.of(allCaps));
-  }
-
-  private HttpRequest createRequest(NewSessionPayload payload) {
-    StringBuilder builder = new StringBuilder();
-    try {
-      payload.writeTo(builder);
-    } catch (IOException e) {
-      throw new UncheckedIOException(e);
-    }
-
-    HttpRequest request = new HttpRequest(POST, "/se/grid/distributor/session");
-    request.setContent(utf8String(builder.toString()));
-
-    return request;
+      Set.of(allCaps),
+      Map.of());
   }
 
   private URI createUri() {

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

@@ -56,6 +56,7 @@
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
@@ -244,11 +245,13 @@ public HttpResponse execute(HttpRequest req) {
       .build();
     distributor.add(node);
 
-    SessionRequest sessionRequest = new SessionRequest(
-      new RequestId(UUID.randomUUID()),
-      Instant.now(),
-      Set.of(W3C),
-      Set.of(new ImmutableCapabilities("browserName", "cheese")));
+    SessionRequest sessionRequest =
+        new SessionRequest(
+            new RequestId(UUID.randomUUID()),
+            Instant.now(),
+            Set.of(W3C),
+            Set.of(new ImmutableCapabilities("browserName", "cheese")),
+            Map.of());
 
     List<Callable<SessionId>> callables = new ArrayList<>();
     for (int i = 0; i < 3; i++) {

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

@@ -110,7 +110,8 @@ public void setupGrid() {
       new RequestId(UUID.randomUUID()),
       Instant.now(),
       Set.of(OSS, W3C),
-      Set.of(caps));
+      Set.of(caps),
+      Map.of());
 
     queue = new LocalNewSessionQueue(
       tracer,
@@ -176,7 +177,8 @@ public void shouldBeAbleToGetSessionQueueSize() {
       new RequestId(UUID.randomUUID()),
       Instant.now(),
       Set.of(W3C),
-      Set.of(caps));
+      Set.of(caps),
+      Map.of());
 
     continueOnceAddedToQueue(request);
     GraphqlHandler handler = new GraphqlHandler(tracer, distributor, queue, publicUri, version);
@@ -196,7 +198,8 @@ public void shouldBeAbleToGetSessionQueueRequests() {
       new RequestId(UUID.randomUUID()),
       Instant.now(),
       Set.of(W3C),
-      Set.of(caps));
+      Set.of(caps),
+      Map.of());
 
     continueOnceAddedToQueue(request);
 

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

@@ -56,6 +56,7 @@
 import java.util.Collection;
 import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
@@ -122,7 +123,8 @@ public LocalNewSessionQueueTest(Supplier<TestData> supplier) {
       new RequestId(UUID.randomUUID()),
       Instant.now(),
       Set.of(W3C),
-      Set.of(CAPS));
+      Set.of(CAPS),
+      Map.of());
   }
 
   @Parameterized.Parameters
@@ -377,7 +379,8 @@ public void shouldBeAbleToHandleMultipleSessionRequestsAtTheSameTime() {
         new RequestId(UUID.randomUUID()),
         Instant.now(),
         Set.of(W3C),
-        Set.of(CAPS));
+        Set.of(CAPS),
+        Map.of());
 
       return queue.addToQueue(sessionRequest);
     };
@@ -409,7 +412,8 @@ public void shouldBeAbleToTimeoutARequestOnRetry() {
       new RequestId(UUID.randomUUID()),
       LONG_AGO,
       Set.of(W3C),
-      Set.of(CAPS));
+      Set.of(CAPS),
+      Map.of());
 
     AtomicInteger count = new AtomicInteger();
 
@@ -438,7 +442,8 @@ public void shouldBeAbleToTimeoutARequestOnRemove() throws InterruptedException
       new RequestId(UUID.randomUUID()),
       LONG_AGO,
       Set.of(W3C),
-      Set.of(CAPS));
+      Set.of(CAPS),
+      Map.of());
     localQueue.injectIntoQueue(sessionRequest);
 
     queue.remove(sessionRequest.getRequestId());
@@ -456,7 +461,8 @@ public void shouldBeAbleToClearQueueAndRejectMultipleRequests() {
         new RequestId(UUID.randomUUID()),
         Instant.now(),
         Set.of(W3C),
-        Set.of(CAPS));
+        Set.of(CAPS),
+        Map.of());
       return queue.addToQueue(sessionRequest);
     };