[tracing] Better handling of tracing through distributor

With these changes in place, it's now possible to follow the 'New
Session' command through all the services it touches. *phew*

Author: shs96c

java/client/src/org/openqa/selenium/remote/tracing/HttpTracing.java

@@ -29,9 +29,12 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.logging.Logger;
 
 public class HttpTracing {
 
+  private static final Logger LOG = Logger.getLogger(HttpTracing.class.getName());
+
   private HttpTracing() {
     // Utility classes
   }
@@ -52,6 +55,9 @@ public static void inject(Tracer tracer, Span span, HttpRequest request) {
     Objects.requireNonNull(tracer, "Tracer to use must be set.");
     Objects.requireNonNull(request, "Request must be set.");
 
+    StackTraceElement caller = Thread.currentThread().getStackTrace()[2];
+    LOG.info(String.format("Injecting %s into %s at %s:%d", request, span, caller.getClassName(), caller.getLineNumber()));
+
     span.setTag(Tags.HTTP_METHOD.getKey(), request.getMethod().toString());
     span.setTag(Tags.HTTP_URL.getKey(), request.getUri());
 

java/server/src/org/openqa/selenium/grid/distributor/local/LocalDistributor.java

@@ -106,6 +106,7 @@ public CreateSessionResponse newSession(HttpRequest request)
     Span previous = tracer.scopeManager().activeSpan();
     SpanContext parent = HttpTracing.extract(tracer, request);
     Span span = tracer.buildSpan("distributor.new_session").asChildOf(parent).start();
+    tracer.scopeManager().activate(span);
 
     try (Reader reader = reader(request);
     NewSessionPayload payload = NewSessionPayload.create(reader)) {

java/server/src/org/openqa/selenium/grid/distributor/remote/RemoteDistributor.java

@@ -90,8 +90,7 @@ public RemoteDistributor add(Node node) {
   public void remove(UUID nodeId) {
     Objects.requireNonNull(nodeId, "Node ID must be set");
     HttpRequest request = new HttpRequest(DELETE, "/se/grid/distributor/node/" + nodeId);
-    Span span = tracer.scopeManager().activeSpan();
-    HttpTracing.inject(tracer, span, request);
+    HttpTracing.inject(tracer, tracer.scopeManager().activeSpan(), request);
 
     HttpResponse response = client.execute(request);
 

java/server/src/org/openqa/selenium/grid/node/Node.java

@@ -111,14 +111,14 @@ protected Node(Tracer tracer, UUID id, URI uri) {
         // "getSessionId" is aggressive about finding session ids, so this needs to be the last
         // route the is checked.
         matching(req -> getSessionId(req.getUri()).map(SessionId::new).map(this::isSessionOwner).orElse(false))
-            .to(() -> new ForwardWebDriverCommand(this)),
+            .to(() -> new ForwardWebDriverCommand(this)).with(new SpanDecorator(tracer, req -> "node.forward_command")),
         get("/se/grid/node/owner/{sessionId}")
             .to((params) -> new IsSessionOwner(this, json, new SessionId(params.get("sessionId")))).with(new SpanDecorator(tracer, req -> "node.is_session_owner")),
         delete("/se/grid/node/session/{sessionId}")
-            .to((params) -> new StopNodeSession(this, new SessionId(params.get("sessionId")))),
+            .to((params) -> new StopNodeSession(this, new SessionId(params.get("sessionId")))).with(new SpanDecorator(tracer, req -> "node.stop_session")),
         get("/se/grid/node/session/{sessionId}")
-            .to((params) -> new GetNodeSession(this, json, new SessionId(params.get("sessionId")))),
-        post("/se/grid/node/session").to(() -> new NewNodeSession(this, json)),
+            .to((params) -> new GetNodeSession(this, json, new SessionId(params.get("sessionId")))).with(new SpanDecorator(tracer, req -> "node.get_session")),
+        post("/se/grid/node/session").to(() -> new NewNodeSession(this, json)).with(new SpanDecorator(tracer, req -> "node.new_session")),
         get("/se/grid/node/status")
             .to(() -> req -> new HttpResponse().setContent(utf8String(json.toJson(getStatus())))).with(new SpanDecorator(tracer, req -> "node.node_status")),
         get("/status").to(() -> new StatusHandler(this, json)).with(new SpanDecorator(tracer, req -> "node.status")));

java/server/src/org/openqa/selenium/grid/node/local/LocalNode.java

@@ -26,6 +26,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import io.opentracing.Span;
 import io.opentracing.Tracer;
 import org.openqa.selenium.Capabilities;
 import org.openqa.selenium.NoSuchSessionException;
@@ -53,6 +54,7 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.logging.Logger;
 import java.util.stream.Collectors;
 
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
@@ -130,8 +132,12 @@ public boolean isSupporting(Capabilities capabilities) {
 
   @Override
   public Optional<CreateSessionResponse> newSession(CreateSessionRequest sessionRequest) {
+    Span span = tracer.scopeManager().activeSpan();
+    Logger.getLogger(LocalNode.class.getName()).info("Creating new session using span: " + span);
     Objects.requireNonNull(sessionRequest, "Session request has not been set.");
 
+    if (span != null) {span.setTag("session_count", getCurrentSessionCount());}
+
     if (getCurrentSessionCount() >= maxSessionCount) {
       return Optional.empty();
     }
@@ -157,6 +163,14 @@ public Optional<CreateSessionResponse> newSession(CreateSessionRequest sessionRe
     ActiveSession session = possibleSession.get();
     currentSessions.put(session.getId(), slot);
 
+    if (span != null) {
+      span.setTag("session.id", session.getId().toString());
+      span.setTag("session.capabilities", session.getCapabilities().toString());
+      span.setTag("session.downstream.dialect", session.getDownstreamDialect().toString());
+      span.setTag("session.upstream.dialect", session.getUpstreamDialect().toString());
+      span.setTag("session.uri", session.getUri().toString());
+    }
+
     // The session we return has to look like it came from the node, since we might be dealing
     // with a webdriver implementation that only accepts connections from localhost
     Session externalSession = createExternalSession(session, externalUri);