[tracing] Add basic tracing to the session map

shs96c · shs96c · commit 7c27774b43ad · 2019-11-08T13:29:50.000Z
diff --git a/java/server/src/org/openqa/selenium/grid/commands/Hub.java b/java/server/src/org/openqa/selenium/grid/commands/Hub.java
@@ -106,7 +106,7 @@ public Executable configure(String... args) {
 
       CombinedHandler handler = new CombinedHandler();
 
-      SessionMap sessions = new LocalSessionMap(bus);
+      SessionMap sessions = new LocalSessionMap(tracer, bus);
       handler.addHandler(sessions);
 
       BaseServerOptions serverOptions = new BaseServerOptions(config);
diff --git a/java/server/src/org/openqa/selenium/grid/commands/Standalone.java b/java/server/src/org/openqa/selenium/grid/commands/Standalone.java
@@ -144,7 +144,7 @@ public Executable configure(String... args) {
           combinedHandler,
           HttpClient.Factory.createDefault()));
 
-      SessionMap sessions = new LocalSessionMap(bus);
+      SessionMap sessions = new LocalSessionMap(tracer, bus);
       combinedHandler.addHandler(sessions);
       Distributor distributor = new LocalDistributor(tracer, bus, clientFactory, sessions);
       combinedHandler.addHandler(distributor);
diff --git a/java/server/src/org/openqa/selenium/grid/distributor/httpd/DistributorServer.java b/java/server/src/org/openqa/selenium/grid/distributor/httpd/DistributorServer.java
@@ -111,7 +111,7 @@ public Executable configure(String... args) {
 
       HttpClient.Factory clientFactory = new TracedHttpClient.Factory(tracer, HttpClient.Factory.createDefault());
 
-      SessionMap sessions = new SessionMapOptions(config).getSessionMap(clientFactory);
+      SessionMap sessions = new SessionMapOptions(config).getSessionMap(tracer, clientFactory);
 
       Distributor distributor = new LocalDistributor(
           tracer,
diff --git a/java/server/src/org/openqa/selenium/grid/router/httpd/RouterServer.java b/java/server/src/org/openqa/selenium/grid/router/httpd/RouterServer.java
@@ -105,7 +105,7 @@ public Executable configure(String... args) {
       HttpClient.Factory clientFactory = new TracedHttpClient.Factory(tracer, HttpClient.Factory.createDefault());
 
       SessionMapOptions sessionsOptions = new SessionMapOptions(config);
-      SessionMap sessions = sessionsOptions.getSessionMap(clientFactory);
+      SessionMap sessions = sessionsOptions.getSessionMap(tracer, clientFactory);
 
       BaseServerOptions serverOptions = new BaseServerOptions(config);
 
diff --git a/java/server/src/org/openqa/selenium/grid/sessionmap/AddToSessionMap.java b/java/server/src/org/openqa/selenium/grid/sessionmap/AddToSessionMap.java
@@ -18,34 +18,59 @@
 package org.openqa.selenium.grid.sessionmap;
 
 import com.google.common.collect.ImmutableMap;
+import io.opentracing.Span;
+import io.opentracing.SpanContext;
+import io.opentracing.Tracer;
+import io.opentracing.tag.Tags;
 import org.openqa.selenium.grid.data.Session;
 import org.openqa.selenium.json.Json;
 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.HttpTracing;
 
 import java.util.Objects;
 
+import static io.opentracing.tag.Tags.HTTP_METHOD;
+import static io.opentracing.tag.Tags.HTTP_URL;
 import static org.openqa.selenium.remote.http.Contents.string;
 import static org.openqa.selenium.remote.http.Contents.utf8String;
 
 class AddToSessionMap implements HttpHandler {
 
+  private final Tracer tracer;
   private final Json json;
   private final SessionMap sessions;
 
-  AddToSessionMap(Json json, SessionMap sessions) {
+  AddToSessionMap(Tracer tracer, Json json, SessionMap sessions) {
+    this.tracer = tracer;
     this.json = Objects.requireNonNull(json);
     this.sessions = Objects.requireNonNull(sessions);
   }
 
   @Override
   public HttpResponse execute(HttpRequest req) {
-    Session session = json.toType(string(req), Session.class);
-    Objects.requireNonNull(session, "Session to add must be set");
+    SpanContext parent = HttpTracing.extract(tracer, req);
+    Span current = tracer.scopeManager().activeSpan();
+    Span span = tracer.buildSpan("sessions.add_session").asChildOf(parent).start();
+    tracer.scopeManager().activate(span);
+    try {
+      HTTP_METHOD.set(span, req.getMethod().toString());
+      HTTP_URL.set(span, req.getUri());
 
-    sessions.add(session);
+      Session session = json.toType(string(req), Session.class);
+      Objects.requireNonNull(session, "Session to add must be set");
 
-    return new HttpResponse().setContent(utf8String(json.toJson(ImmutableMap.of("value", true))));
+      span.setTag("session.id", session.getId().toString());
+      span.setTag("session.capabilities", session.getCapabilities().toString());
+      span.setTag("session.uri", session.getUri().toString());
+
+      sessions.add(session);
+
+      return new HttpResponse().setContent(utf8String(json.toJson(ImmutableMap.of("value", true))));
+    } finally {
+      span.finish();
+      tracer.scopeManager().activate(current);
+    }
   }
 }
diff --git a/java/server/src/org/openqa/selenium/grid/sessionmap/GetFromSessionMap.java b/java/server/src/org/openqa/selenium/grid/sessionmap/GetFromSessionMap.java
@@ -18,33 +18,58 @@
 package org.openqa.selenium.grid.sessionmap;
 
 import com.google.common.collect.ImmutableMap;
+import io.opentracing.Span;
+import io.opentracing.SpanContext;
+import io.opentracing.Tracer;
 import org.openqa.selenium.grid.data.Session;
 import org.openqa.selenium.json.Json;
 import org.openqa.selenium.remote.SessionId;
 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.HttpTracing;
 
 import java.util.Objects;
 
+import static io.opentracing.tag.Tags.HTTP_METHOD;
+import static io.opentracing.tag.Tags.HTTP_URL;
 import static org.openqa.selenium.remote.http.Contents.utf8String;
 
 class GetFromSessionMap implements HttpHandler {
 
+  private final Tracer tracer;
   private final Json json;
   private final SessionMap sessions;
-  private SessionId id;
+  private final SessionId id;
 
-  public GetFromSessionMap(Json json, SessionMap sessions, SessionId id) {
+  public GetFromSessionMap(Tracer tracer, Json json, SessionMap sessions, SessionId id) {
+    this.tracer = Objects.requireNonNull(tracer);
     this.json = Objects.requireNonNull(json);
     this.sessions = Objects.requireNonNull(sessions);
     this.id = Objects.requireNonNull(id);
   }
 
   @Override
   public HttpResponse execute(HttpRequest req) {
-    Session session = sessions.get(id);
+    SpanContext parent = HttpTracing.extract(tracer, req);
+    Span current = tracer.scopeManager().activeSpan();
+    Span span = tracer.buildSpan("sessions.get_session").asChildOf(parent).start();
+    tracer.scopeManager().activate(span);
 
-    return new HttpResponse().setContent(utf8String(json.toJson(ImmutableMap.of("value", session))));
+    try {
+      HTTP_METHOD.set(span, req.getMethod().toString());
+      HTTP_URL.set(span, req.getUri());
+
+      Session session = sessions.get(id);
+
+      span.setTag("session.id", session.getId().toString());
+      span.setTag("session.capabilities", session.getCapabilities().toString());
+      span.setTag("session.uri", session.getUri().toString());
+
+      return new HttpResponse().setContent(utf8String(json.toJson(ImmutableMap.of("value", session))));
+    } finally {
+      span.finish();
+      tracer.scopeManager().activate(current);
+    }
   }
 }
diff --git a/java/server/src/org/openqa/selenium/grid/sessionmap/RemoveFromSession.java b/java/server/src/org/openqa/selenium/grid/sessionmap/RemoveFromSession.java
@@ -17,26 +17,50 @@
 
 package org.openqa.selenium.grid.sessionmap;
 
+import io.opentracing.Span;
+import io.opentracing.SpanContext;
+import io.opentracing.Tracer;
 import org.openqa.selenium.remote.SessionId;
 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.HttpTracing;
 
 import java.util.Objects;
 
+import static io.opentracing.tag.Tags.HTTP_METHOD;
+import static io.opentracing.tag.Tags.HTTP_URL;
+
+
 class RemoveFromSession implements HttpHandler {
 
+  private final Tracer tracer;
   private final SessionMap sessions;
-  private SessionId id;
+  private final SessionId id;
 
-  public RemoveFromSession(SessionMap sessions, SessionId id) {
+  public RemoveFromSession(Tracer tracer, SessionMap sessions, SessionId id) {
+    this.tracer = Objects.requireNonNull(tracer);
     this.sessions = Objects.requireNonNull(sessions);
     this.id = Objects.requireNonNull(id);
   }
 
   @Override
   public HttpResponse execute(HttpRequest req) {
-    sessions.remove(id);
-    return new HttpResponse();
+    SpanContext parent = HttpTracing.extract(tracer, req);
+    Span current = tracer.scopeManager().activeSpan();
+    Span span = tracer.buildSpan("sessions.remove_session").asChildOf(parent).start();
+    tracer.scopeManager().activate(span);
+
+    try {
+      HTTP_METHOD.set(span, req.getMethod().toString());
+      HTTP_URL.set(span, req.getUri());
+      span.setTag("session.id", String.valueOf(id));
+
+      sessions.remove(id);
+      return new HttpResponse();
+    } finally {
+      span.finish();
+      tracer.scopeManager().activate(current);
+    }
   }
 }
diff --git a/java/server/src/org/openqa/selenium/grid/sessionmap/SessionMap.java b/java/server/src/org/openqa/selenium/grid/sessionmap/SessionMap.java
@@ -17,6 +17,7 @@
 
 package org.openqa.selenium.grid.sessionmap;
 
+import io.opentracing.Tracer;
 import org.openqa.selenium.NoSuchSessionException;
 import org.openqa.selenium.grid.data.Session;
 import org.openqa.selenium.json.Json;
@@ -28,6 +29,7 @@
 import org.openqa.selenium.remote.http.Route;
 
 import java.net.URI;
+import java.util.Objects;
 
 import static org.openqa.selenium.remote.http.Route.combine;
 import static org.openqa.selenium.remote.http.Route.delete;
@@ -67,6 +69,8 @@
  */
 public abstract class SessionMap implements Routable, HttpHandler {
 
+  protected final Tracer tracer;
+
   private final Route routes;
 
   public abstract boolean add(Session session);
@@ -75,14 +79,16 @@ public abstract class SessionMap implements Routable, HttpHandler {
 
   public abstract void remove(SessionId id);
 
-  public SessionMap() {
+  public SessionMap(Tracer tracer) {
+    this.tracer = Objects.requireNonNull(tracer);
+
     Json json = new Json();
     routes = combine(
-        post("/se/grid/session").to(() -> new AddToSessionMap(json, this)),
+        post("/se/grid/session").to(() -> new AddToSessionMap(tracer, json, this)),
         Route.get("/se/grid/session/{sessionId}")
-            .to((params) -> new GetFromSessionMap(json, this, new SessionId(params.get("sessionId")))),
+            .to((params) -> new GetFromSessionMap(tracer, json, this, new SessionId(params.get("sessionId")))),
         delete("/se/grid/session/{sessionId}")
-            .to((params) -> new RemoveFromSession(this, new SessionId(params.get("sessionId")))));
+            .to((params) -> new RemoveFromSession(tracer, this, new SessionId(params.get("sessionId")))));
   }
 
   @Override
diff --git a/java/server/src/org/openqa/selenium/grid/sessionmap/config/BUILD.bazel b/java/server/src/org/openqa/selenium/grid/sessionmap/config/BUILD.bazel
@@ -10,5 +10,6 @@ java_library(
         "//java/server/src/org/openqa/selenium/grid/sessionmap",
         "//java/server/src/org/openqa/selenium/grid/sessionmap/remote",
         "//third_party/java/beust:jcommander",
+        "//third_party/java/opentracing:opentracing-api",
     ],
 )
diff --git a/java/server/src/org/openqa/selenium/grid/sessionmap/config/SessionMapOptions.java b/java/server/src/org/openqa/selenium/grid/sessionmap/config/SessionMapOptions.java
@@ -17,6 +17,7 @@
 
 package org.openqa.selenium.grid.sessionmap.config;
 
+import io.opentracing.Tracer;
 import org.openqa.selenium.grid.config.Config;
 import org.openqa.selenium.grid.config.ConfigException;
 import org.openqa.selenium.grid.sessionmap.SessionMap;
@@ -35,9 +36,9 @@ public SessionMapOptions(Config config) {
     this.config = config;
   }
 
-  public SessionMap getSessionMap(HttpClient.Factory clientFactory) {
+  public SessionMap getSessionMap(Tracer tracer, HttpClient.Factory clientFactory) {
     HttpClient client = clientFactory.createClient(getSessionMapUrl());
-    return new RemoteSessionMap(client);
+    return new RemoteSessionMap(tracer, client);
   }
 
   private URL getSessionMapUrl() {
diff --git a/java/server/src/org/openqa/selenium/grid/sessionmap/httpd/SessionMapServer.java b/java/server/src/org/openqa/selenium/grid/sessionmap/httpd/SessionMapServer.java
@@ -20,6 +20,7 @@
 import com.beust.jcommander.JCommander;
 import com.beust.jcommander.ParameterException;
 import com.google.auto.service.AutoService;
+import io.opentracing.Tracer;
 import org.openqa.selenium.BuildInfo;
 import org.openqa.selenium.cli.CliCommand;
 import org.openqa.selenium.events.EventBus;
@@ -93,11 +94,12 @@ public Executable configure(String... args) {
 
       LoggingOptions loggingOptions = new LoggingOptions(config);
       loggingOptions.configureLogging();
+      Tracer tracer = loggingOptions.getTracer();
 
       EventBusConfig events = new EventBusConfig(config);
       EventBus bus = events.getEventBus();
 
-      SessionMap sessions = new LocalSessionMap(bus);
+      SessionMap sessions = new LocalSessionMap(tracer, bus);
 
       BaseServerOptions serverOptions = new BaseServerOptions(config);
 
diff --git a/java/server/src/org/openqa/selenium/grid/sessionmap/local/LocalSessionMap.java b/java/server/src/org/openqa/selenium/grid/sessionmap/local/LocalSessionMap.java
@@ -17,6 +17,7 @@
 
 package org.openqa.selenium.grid.sessionmap.local;
 
+import io.opentracing.Tracer;
 import org.openqa.selenium.NoSuchSessionException;
 import org.openqa.selenium.events.EventBus;
 import org.openqa.selenium.grid.data.Session;
@@ -38,7 +39,9 @@ public class LocalSessionMap extends SessionMap {
   private final Map<SessionId, Session> knownSessions = new HashMap<>();
   private final ReadWriteLock lock = new ReentrantReadWriteLock(/* be fair */ true);
 
-  public LocalSessionMap(EventBus bus) {
+  public LocalSessionMap(Tracer tracer, EventBus bus) {
+    super(tracer);
+
     this.bus = Objects.requireNonNull(bus);
 
     bus.addListener(SESSION_CLOSED, event -> {
diff --git a/java/server/src/org/openqa/selenium/grid/sessionmap/remote/RemoteSessionMap.java b/java/server/src/org/openqa/selenium/grid/sessionmap/remote/RemoteSessionMap.java
@@ -17,6 +17,8 @@
 
 package org.openqa.selenium.grid.sessionmap.remote;
 
+import io.opentracing.Span;
+import io.opentracing.Tracer;
 import org.openqa.selenium.NoSuchSessionException;
 import org.openqa.selenium.grid.data.Session;
 import org.openqa.selenium.grid.sessionmap.SessionMap;
@@ -26,6 +28,7 @@
 import org.openqa.selenium.remote.http.HttpClient;
 import org.openqa.selenium.remote.http.HttpRequest;
 import org.openqa.selenium.remote.http.HttpResponse;
+import org.openqa.selenium.remote.tracing.HttpTracing;
 
 import java.lang.reflect.Type;
 import java.util.Objects;
@@ -40,7 +43,9 @@ public class RemoteSessionMap extends SessionMap {
   public static final Json JSON = new Json();
   private final HttpClient client;
 
-  public RemoteSessionMap(HttpClient client) {
+  public RemoteSessionMap(Tracer tracer, HttpClient client) {
+    super(tracer);
+
     this.client = Objects.requireNonNull(client);
   }
 
@@ -73,6 +78,9 @@ public void remove(SessionId id) {
   }
 
   private <T> T makeRequest(HttpRequest request, Type typeOfT) {
+    Span activeSpan = tracer.scopeManager().activeSpan();
+    HttpTracing.inject(tracer, activeSpan, request);
+
     HttpResponse response = client.execute(request);
     return Values.get(response, typeOfT);
   }
diff --git a/java/server/test/org/openqa/selenium/grid/distributor/AddingNodesTest.java b/java/server/test/org/openqa/selenium/grid/distributor/AddingNodesTest.java
@@ -92,7 +92,7 @@ public void setUpDistributor() throws MalformedURLException {
         handler,
         HttpClient.Factory.createDefault());
 
-    LocalSessionMap sessions = new LocalSessionMap(bus);
+    LocalSessionMap sessions = new LocalSessionMap(tracer, bus);
     Distributor local = new LocalDistributor(tracer, bus, clientFactory, sessions);
     handler.addHandler(local);
     distributor = new RemoteDistributor(tracer, clientFactory, externalUrl);
diff --git a/java/server/test/org/openqa/selenium/grid/distributor/DistributorTest.java b/java/server/test/org/openqa/selenium/grid/distributor/DistributorTest.java
diff --git a/java/server/test/org/openqa/selenium/grid/distributor/local/LocalDistributorTest.java b/java/server/test/org/openqa/selenium/grid/distributor/local/LocalDistributorTest.java
diff --git a/java/server/test/org/openqa/selenium/grid/router/EndToEndTest.java b/java/server/test/org/openqa/selenium/grid/router/EndToEndTest.java
diff --git a/java/server/test/org/openqa/selenium/grid/router/RouterTest.java b/java/server/test/org/openqa/selenium/grid/router/RouterTest.java
diff --git a/java/server/test/org/openqa/selenium/grid/sessionmap/BUILD.bazel b/java/server/test/org/openqa/selenium/grid/sessionmap/BUILD.bazel
diff --git a/java/server/test/org/openqa/selenium/grid/sessionmap/SessionMapTest.java b/java/server/test/org/openqa/selenium/grid/sessionmap/SessionMapTest.java

Original file line number	Diff line number	Diff line change
`@@ -10,5 +10,6 @@ java_library(`
`10`	`10`	`"//java/server/src/org/openqa/selenium/grid/sessionmap",`
`11`	`11`	`"//java/server/src/org/openqa/selenium/grid/sessionmap/remote",`
`12`	`12`	`"//third_party/java/beust:jcommander",`
	`13`	`+ "//third_party/java/opentracing:opentracing-api",`
`13`	`14`	`],`
`14`	`15`	`)`