Add distributed tracing to the router and sessionmap

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

@@ -93,9 +93,9 @@ public Executable configure(String... args) {
       DistributedTracer tracer = new LoggingOptions(config).getTracer();
       GlobalDistributedTracer.setInstance(tracer);
 
-      SessionMap sessions = new LocalSessionMap();
+      SessionMap sessions = new LocalSessionMap(tracer);
       Distributor distributor = new LocalDistributor(tracer);
-      Router router = new Router(sessions, distributor);
+      Router router = new Router(tracer, sessions, distributor);
 
       Server<?> server = new BaseServer<>(
           new BaseServerOptions(config));

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

@@ -99,9 +99,9 @@ public Executable configure(String... args) {
       DistributedTracer tracer = new LoggingOptions(config).getTracer();
       GlobalDistributedTracer.setInstance(tracer);
 
-      SessionMap sessions = new LocalSessionMap();
+      SessionMap sessions = new LocalSessionMap(tracer);
       Distributor distributor = new LocalDistributor(tracer);
-      Router router = new Router(sessions, distributor);
+      Router router = new Router(tracer, sessions, distributor);
 
       String hostName;
       try {

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

@@ -6,6 +6,7 @@ java_library(
         "//java/server/test/org/openqa/selenium/grid:__subpackages__",
     ],
     deps = [
+        "//java/client/src/org/openqa/selenium/remote/tracing",
         "//java/server/src/org/openqa/selenium/grid/distributor",
         "//java/server/src/org/openqa/selenium/grid/sessionmap",
         "//java/server/src/org/openqa/selenium/injector",

java/server/src/org/openqa/selenium/grid/router/HandleSession.java

@@ -28,6 +28,9 @@
 import org.openqa.selenium.remote.SessionId;
 import org.openqa.selenium.remote.http.HttpRequest;
 import org.openqa.selenium.remote.http.HttpResponse;
+import org.openqa.selenium.remote.tracing.DistributedTracer;
+import org.openqa.selenium.remote.tracing.HttpTracing;
+import org.openqa.selenium.remote.tracing.Span;
 
 import java.io.IOException;
 import java.time.Duration;
@@ -37,8 +40,10 @@
 class HandleSession implements CommandHandler {
 
   private final LoadingCache<SessionId, CommandHandler> knownSessions;
+  private final DistributedTracer tracer;
 
-  public HandleSession(SessionMap sessions) {
+  public HandleSession(DistributedTracer tracer, SessionMap sessions) {
+    this.tracer = Objects.requireNonNull(tracer);
     Objects.requireNonNull(sessions);
 
     this.knownSessions = CacheBuilder.newBuilder()
@@ -57,17 +62,27 @@ public CommandHandler load(SessionId id) throws Exception {
 
   @Override
   public void execute(HttpRequest req, HttpResponse resp) throws IOException {
-    String[] split = req.getUri().split("/", 4);
-    SessionId id = new SessionId(split[2]);
+    try (Span span = tracer.createSpan("router.webdriver-command", tracer.getActiveSpan())) {
+      span.addTag("http.method", req.getMethod());
+      span.addTag("http.url", req.getUri());
 
-    try {
-      knownSessions.get(id).execute(req, resp);
-    } catch (ExecutionException e) {
-      Throwable cause = e.getCause();
-      if (cause instanceof RuntimeException) {
-        throw (RuntimeException) cause;
+      String[] split = req.getUri().split("/", 4);
+      SessionId id = new SessionId(split[2]);
+
+      span.addTag("session.id", id);
+
+      try {
+        knownSessions.get(id).execute(req, resp);
+        span.addTag("http.status", resp.getStatus());
+      } catch (ExecutionException e) {
+        span.addTag("exception", e.getMessage());
+
+        Throwable cause = e.getCause();
+        if (cause instanceof RuntimeException) {
+          throw (RuntimeException) cause;
+        }
+        throw new RuntimeException(cause);
       }
-      throw new RuntimeException(cause);
     }
   }
 }

java/server/src/org/openqa/selenium/grid/router/Router.java

@@ -29,6 +29,7 @@
 import org.openqa.selenium.json.Json;
 import org.openqa.selenium.remote.http.HttpRequest;
 import org.openqa.selenium.remote.http.HttpResponse;
+import org.openqa.selenium.remote.tracing.DistributedTracer;
 
 import java.io.IOException;
 import java.util.Optional;
@@ -42,8 +43,9 @@ public class Router implements Predicate<HttpRequest>, CommandHandler {
   private final Injector injector;
   private final Routes routes;
 
-  public Router(SessionMap sessions, Distributor distributor) {
+  public Router(DistributedTracer tracer, SessionMap sessions, Distributor distributor) {
     injector = Injector.builder()
+        .register(tracer)
         .register(sessions)
         .register(distributor)
         .register(new Json())

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

@@ -113,7 +113,7 @@ public Executable configure(String... args) {
           tracer,
           HttpClient.Factory.createDefault().createClient(distributorUrl));
 
-      Router router = new Router(sessions, distributor);
+      Router router = new Router(tracer, sessions, distributor);
 
       Server<?> server = new BaseServer<>(serverOptions);
       server.addRoute(Routes.matching(router).using(router).decorateWith(W3CCommandHandler.class));

java/server/src/org/openqa/selenium/grid/sessionmap/httpd/SessionMapServer.java

@@ -89,7 +89,7 @@ public Executable configure(String... args) {
       DistributedTracer tracer = new LoggingOptions(config).getTracer();
       GlobalDistributedTracer.setInstance(tracer);
 
-      SessionMap sessions = new LocalSessionMap();
+      SessionMap sessions = new LocalSessionMap(tracer);
 
       BaseServerOptions serverOptions = new BaseServerOptions(config);
 

java/server/src/org/openqa/selenium/grid/sessionmap/local/BUCK

@@ -4,6 +4,7 @@ java_library(
     srcs = glob(["*.java"]),
     deps = [
         "//java/client/src/org/openqa/selenium/remote:remote",
+        "//java/client/src/org/openqa/selenium/remote/tracing:tracing",
         "//java/server/src/org/openqa/selenium/grid/sessionmap:sessionmap",
         "//java/server/src/org/openqa/selenium/grid/data:data",
         "//java/server/src/org/openqa/selenium/grid/web:web",

java/server/src/org/openqa/selenium/grid/sessionmap/local/BUILD.bazel

@@ -6,6 +6,7 @@ java_library(
         "//java/server/test/org/openqa/selenium/grid:__subpackages__",
     ],
     deps = [
+        "//java/client/src/org/openqa/selenium/remote/tracing",
         "//java/server/src/org/openqa/selenium/grid/sessionmap",
     ],
 )

java/server/src/org/openqa/selenium/grid/sessionmap/local/LocalSessionMap.java

@@ -21,6 +21,8 @@
 import org.openqa.selenium.grid.data.Session;
 import org.openqa.selenium.grid.sessionmap.SessionMap;
 import org.openqa.selenium.remote.SessionId;
+import org.openqa.selenium.remote.tracing.DistributedTracer;
+import org.openqa.selenium.remote.tracing.Span;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -31,51 +33,74 @@
 
 public class LocalSessionMap extends SessionMap {
 
-  private Map<SessionId, Session> knownSessions = new HashMap<>();
-  private ReadWriteLock lock = new ReentrantReadWriteLock(/* be fair */ true);
+  private final DistributedTracer tracer;
+  private final Map<SessionId, Session> knownSessions = new HashMap<>();
+  private final ReadWriteLock lock = new ReentrantReadWriteLock(/* be fair */ true);
+
+  public LocalSessionMap(DistributedTracer tracer) {
+    this.tracer = Objects.requireNonNull(tracer);
+  }
 
   @Override
   public boolean add(Session session) {
     Objects.requireNonNull(session, "Session has not been set");
 
-    Lock writeLock = lock.writeLock();
-    writeLock.lock();
-    try {
-      knownSessions.put(session.getId(), session);
-    } finally {
-      writeLock.unlock();
-    }
+    try (Span span = tracer.createSpan("sessionmap.add", tracer.getActiveSpan())) {
+      span.addTag("session.id", session.getId());
+      span.addTag("session.capabilities", session.getCapabilities());
+      span.addTag("session.uri", session.getUri());
 
-    return true;
+      Lock writeLock = lock.writeLock();
+      writeLock.lock();
+      try {
+        knownSessions.put(session.getId(), session);
+      } finally {
+        writeLock.unlock();
+      }
+
+      return true;
+    }
   }
 
   @Override
   public Session get(SessionId id) {
     Objects.requireNonNull(id, "Session ID has not been set");
 
-    Lock readLock = lock.readLock();
-    readLock.lock();
-    try {
-      Session session = knownSessions.get(id);
-      if (session == null) {
-        throw new NoSuchSessionException("Unable to find session with ID: " + id);
+    try (Span span = tracer.createSpan("sessionmap.get", tracer.getActiveSpan())) {
+      span.addTag("session.id", id);
+
+      Lock readLock = lock.readLock();
+      readLock.lock();
+      try {
+        Session session = knownSessions.get(id);
+        if (session == null) {
+          throw new NoSuchSessionException("Unable to find session with ID: " + id);
+        }
+
+        span.addTag("session.capabilities", session.getCapabilities());
+        span.addTag("session.uri", session.getUri());
+
+        return session;
+      } finally {
+        readLock.unlock();
       }
-      return session;
-    } finally {
-      readLock.unlock();
     }
   }
 
   @Override
   public void remove(SessionId id) {
     Objects.requireNonNull(id, "Session ID has not been set");
 
-    Lock writeLock = lock.writeLock();
-    writeLock.lock();
-    try {
-      knownSessions.remove(id);
-    } finally {
-      writeLock.unlock();
+    try (Span span = tracer.createSpan("sessionmap.remove", tracer.getActiveSpan())) {
+      span.addTag("session.id", id);
+
+      Lock writeLock = lock.writeLock();
+      writeLock.lock();
+      try {
+        knownSessions.remove(id);
+      } finally {
+        writeLock.unlock();
+      }
     }
   }
 }

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

@@ -68,7 +68,7 @@ public void shouldBeAbleToAddANodeAndCreateASession() throws URISyntaxException
     URI nodeUri = new URI("http://example:5678");
     URI routableUri = new URI("http://localhost:1234");
 
-    LocalSessionMap sessions = new LocalSessionMap();
+    LocalSessionMap sessions = new LocalSessionMap(tracer);
     LocalNode node = LocalNode.builder(tracer, routableUri, sessions)
         .add(caps, c -> new Session(new SessionId(UUID.randomUUID()), nodeUri, c))
         .build();
@@ -91,7 +91,7 @@ public void shouldBeAbleToRemoveANode() throws URISyntaxException {
     URI nodeUri = new URI("http://example:5678");
     URI routableUri = new URI("http://localhost:1234");
 
-    LocalSessionMap sessions = new LocalSessionMap();
+    LocalSessionMap sessions = new LocalSessionMap(tracer);
     LocalNode node = LocalNode.builder(tracer, routableUri, sessions)
         .add(caps, c -> new Session(new SessionId(UUID.randomUUID()), nodeUri, c))
         .build();
@@ -113,7 +113,7 @@ public void registeringTheSameNodeMultipleTimesOnlyCountsTheFirstTime()
     URI nodeUri = new URI("http://example:5678");
     URI routableUri = new URI("http://localhost:1234");
 
-    LocalSessionMap sessions = new LocalSessionMap();
+    LocalSessionMap sessions = new LocalSessionMap(tracer);
     LocalNode node = LocalNode.builder(tracer, routableUri, sessions)
         .add(caps, c -> new Session(new SessionId(UUID.randomUUID()), nodeUri, c))
         .build();

java/server/test/org/openqa/selenium/grid/node/NodeTest.java

@@ -71,7 +71,7 @@ public void setUp() throws URISyntaxException {
 
     uri = new URI("http://localhost:1234");
 
-    sessions = new LocalSessionMap();
+    sessions = new LocalSessionMap(tracer);
 
     class Handler extends Session implements CommandHandler {
 

java/server/test/org/openqa/selenium/grid/router/EndToEndTest.java

@@ -58,15 +58,15 @@ public class EndToEndTest {
   @Test
   public void inMemory() throws URISyntaxException {
 
-    SessionMap sessions = new LocalSessionMap();
+    SessionMap sessions = new LocalSessionMap(tracer);
     Distributor distributor = new LocalDistributor(tracer);
     URI nodeUri = new URI("http://localhost:4444");
     LocalNode node = LocalNode.builder(tracer, nodeUri, sessions)
         .add(driverCaps, createFactory(nodeUri))
         .build();
     distributor.add(node);
 
-    Router router = new Router(sessions, distributor);
+    Router router = new Router(tracer, sessions, distributor);
 
     Server<?> server = createServer();
     server.addRoute(Routes.matching(router).using(router));
@@ -86,7 +86,7 @@ private void exerciseDriver(Server<?> server) {
   @Test
   public void withServers() throws URISyntaxException {
 
-    LocalSessionMap localSessions = new LocalSessionMap();
+    LocalSessionMap localSessions = new LocalSessionMap(tracer);
     Server<?> sessionServer = createServer();
     sessionServer.addRoute(Routes.matching(localSessions).using(localSessions));
     sessionServer.start();
@@ -113,7 +113,7 @@ public void withServers() throws URISyntaxException {
 
     distributor.add(localNode);
 
-    Router router = new Router(sessions, distributor);
+    Router router = new Router(tracer, sessions, distributor);
     Server<?> routerServer = createServer();
     routerServer.addRoute(Routes.matching(router).using(router));
     routerServer.start();

java/server/test/org/openqa/selenium/grid/sessionmap/BUCK

@@ -7,6 +7,7 @@ java_test(
     ],
     deps = [
         "//java/client/src/org/openqa/selenium/remote:remote",
+        "//java/client/src/org/openqa/selenium/remote/tracing:tracing",
         "//java/server/src/org/openqa/selenium/grid/data:data",
         "//java/server/src/org/openqa/selenium/grid/sessionmap:sessionmap",
         "//java/server/src/org/openqa/selenium/grid/sessionmap/local:local",

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

@@ -4,6 +4,7 @@ gen_java_tests(
     size = "small",
     srcs = glob(["*.java"]),
     deps = [
+        "//java/client/src/org/openqa/selenium/remote/tracing",
         "//java/server/src/org/openqa/selenium/grid/sessionmap",
         "//java/server/src/org/openqa/selenium/grid/sessionmap/local",
         "//java/server/src/org/openqa/selenium/grid/sessionmap/remote",

java/server/test/org/openqa/selenium/grid/sessionmap/SessionMapTest.java

@@ -18,7 +18,6 @@
 package org.openqa.selenium.grid.sessionmap;
 
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
-import static org.assertj.core.api.Assertions.catchThrowableOfType;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
@@ -32,6 +31,7 @@
 import org.openqa.selenium.grid.web.PassthroughHttpClient;
 import org.openqa.selenium.remote.SessionId;
 import org.openqa.selenium.remote.http.HttpClient;
+import org.openqa.selenium.remote.tracing.DistributedTracer;
 
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -57,7 +57,7 @@ public void setUp() throws URISyntaxException {
         new URI("http://localhost:1234"),
         new ImmutableCapabilities());
 
-    local = new LocalSessionMap();
+    local = new LocalSessionMap(DistributedTracer.builder().build());
     client = new PassthroughHttpClient<>(local);
     remote = new RemoteSessionMap(client);
   }