Allow us to inject tracing information into headers

SeleniumHQ · Nov 10, 2018 · 173e41a · 173e41a
1 parent 6aa186b
commit 173e41a
Show file tree

Hide file tree

Showing 7 changed files with 191 additions and 36 deletions.
diff --git a/java/client/src/org/openqa/selenium/remote/HttpCommandExecutor.java b/java/client/src/org/openqa/selenium/remote/HttpCommandExecutor.java
@@ -24,7 +24,6 @@
 
 import com.google.common.collect.ImmutableMap;
 
-import org.openqa.selenium.BuildInfo;
 import org.openqa.selenium.NoSuchSessionException;
 import org.openqa.selenium.SessionNotCreatedException;
 import org.openqa.selenium.UnsupportedCommandException;
@@ -38,6 +37,7 @@
 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;
@@ -136,19 +136,17 @@ public Response execute(Command command) throws IOException {
       if (commandCodec != null) {
         throw new SessionNotCreatedException("Session already exists");
       }
-      try (Span span = tracer.createSpan("new-session", tracer.getActiveSpan())) {
-        ProtocolHandshake handshake = new ProtocolHandshake();
-        log(LogType.PROFILER, new HttpProfilerLogEntry(command.getName(), true));
-        ProtocolHandshake.Result result = handshake.createSession(client, command);
-        Dialect dialect = result.getDialect();
-        commandCodec = dialect.getCommandCodec();
-        for (Map.Entry<String, CommandInfo> entry : additionalCommands.entrySet()) {
-          defineCommand(entry.getKey(), entry.getValue());
-        }
-        responseCodec = dialect.getResponseCodec();
-        log(LogType.PROFILER, new HttpProfilerLogEntry(command.getName(), false));
-        return result.createResponse();
+      ProtocolHandshake handshake = new ProtocolHandshake();
+      log(LogType.PROFILER, new HttpProfilerLogEntry(command.getName(), true));
+      ProtocolHandshake.Result result = handshake.createSession(client, command);
+      Dialect dialect = result.getDialect();
+      commandCodec = dialect.getCommandCodec();
+      for (Map.Entry<String, CommandInfo> entry : additionalCommands.entrySet()) {
+        defineCommand(entry.getKey(), entry.getValue());
       }
+      responseCodec = dialect.getResponseCodec();
+      log(LogType.PROFILER, new HttpProfilerLogEntry(command.getName(), false));
+      return result.createResponse();
     }
 
     if (commandCodec == null || responseCodec == null) {
@@ -158,8 +156,9 @@ public Response execute(Command command) throws IOException {
 
     HttpRequest httpRequest = commandCodec.encode(command);
     try (Span span = tracer.createSpan(command.getName(), tracer.getActiveSpan())) {
-      span.addTag("selenium-sessionid", String.valueOf(command.getSessionId()));
       log(LogType.PROFILER, new HttpProfilerLogEntry(command.getName(), true));
+      span.addTag("selenium-sessionid", String.valueOf(command.getSessionId()));
+      HttpTracing.inject(span, httpRequest);
       HttpResponse httpResponse = client.execute(httpRequest);
       log(LogType.PROFILER, new HttpProfilerLogEntry(command.getName(), false));
 

diff --git a/java/client/src/org/openqa/selenium/remote/ProtocolHandshake.java b/java/client/src/org/openqa/selenium/remote/ProtocolHandshake.java
@@ -38,6 +38,9 @@
 import org.openqa.selenium.remote.http.HttpMethod;
 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.BufferedInputStream;
 import java.io.IOException;
@@ -56,7 +59,7 @@ public class ProtocolHandshake {
   private final static Logger LOG = Logger.getLogger(ProtocolHandshake.class.getName());
 
   public Result createSession(HttpClient client, Command command)
-    throws IOException {
+      throws IOException {
     Capabilities desired = (Capabilities) command.getParameters().get("desiredCapabilities");
     desired = desired == null ? new ImmutableCapabilities() : desired;
 
@@ -84,22 +87,28 @@ public Result createSession(HttpClient client, Command command)
     }
 
     throw new SessionNotCreatedException(
-      String.format(
-        "Unable to create new remote session. " +
-        "desired capabilities = %s",
-        desired));
+        String.format(
+            "Unable to create new remote session. " +
+            "desired capabilities = %s",
+            desired));
   }
 
   private Optional<Result> createSession(HttpClient client, InputStream newSessionBlob, long size)
-    throws IOException {
+      throws IOException {
     // Create the http request and send it
     HttpRequest request = new HttpRequest(HttpMethod.POST, "/session");
 
-    request.setHeader(CONTENT_LENGTH, String.valueOf(size));
-    request.setHeader(CONTENT_TYPE, JSON_UTF_8.toString());
-    request.setContent(newSessionBlob);
+    HttpResponse response;
     long start = System.currentTimeMillis();
-    HttpResponse response = client.execute(request);
+    try (Span span = DistributedTracer.getInstance().getActiveSpan()) {
+      HttpTracing.inject(span, request);
+
+      request.setHeader(CONTENT_LENGTH, String.valueOf(size));
+      request.setHeader(CONTENT_TYPE, JSON_UTF_8.toString());
+      request.setContent(newSessionBlob);
+
+      response = client.execute(request);
+    }
     long time = System.currentTimeMillis() - start;
 
     // Ignore the content type. It may not have been set. Strictly speaking we're not following the
@@ -126,6 +135,7 @@ private Optional<Result> createSession(HttpClient client, InputStream newSession
   }
 
   public static class Result {
+
     private static Function<Object, Proxy> massageProxy = obj -> {
       if (obj instanceof Proxy) {
         return (Proxy) obj;
@@ -158,7 +168,8 @@ public static class Result {
 
       if (capabilities.containsKey(PROXY)) {
         //noinspection unchecked
-        ((Map<String, Object>)capabilities).put(PROXY, massageProxy.apply(capabilities.get(PROXY)));
+        ((Map<String, Object>) capabilities)
+            .put(PROXY, massageProxy.apply(capabilities.get(PROXY)));
       }
     }
 

diff --git a/java/client/src/org/openqa/selenium/remote/tracing/CompoundSpan.java b/java/client/src/org/openqa/selenium/remote/tracing/CompoundSpan.java
@@ -19,9 +19,11 @@
 
 import com.google.common.collect.ImmutableSet;
 
+import org.openqa.selenium.remote.http.HttpRequest;
+
 import java.util.Objects;
 
-class CompoundSpan implements Span {
+class CompoundSpan extends Span {
 
   private final DistributedTracer tracer;
   private final ImmutableSet<Span> allSpans;
@@ -83,4 +85,16 @@ public void close() {
     allSpans.forEach(Span::close);
     tracer.remove(this);
   }
+
+  @Override
+  void inject(HttpRequest request) {
+    Objects.requireNonNull(request);
+    allSpans.forEach(span -> span.inject(request));
+  }
+
+  @Override
+  void extract(HttpRequest request) {
+    Objects.requireNonNull(request);
+    allSpans.forEach(span -> span.extract(request));
+  }
 }
diff --git a/java/client/src/org/openqa/selenium/remote/tracing/HttpTracing.java b/java/client/src/org/openqa/selenium/remote/tracing/HttpTracing.java
@@ -0,0 +1,48 @@
+// Licensed to the Software Freedom Conservancy (SFC) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The SFC licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.openqa.selenium.remote.tracing;
+
+import org.openqa.selenium.remote.http.HttpRequest;
+
+import java.util.Objects;
+
+public class HttpTracing {
+
+  private HttpTracing() {
+    // Utility classes
+  }
+
+  public static void inject(Span span, HttpRequest request) {
+    Objects.requireNonNull(request, "Request must be set.");
+    if (span == null) {
+      return;
+    }
+
+    span.inject(request);
+  }
+
+  public static void extract(HttpRequest request, Span intoSpan) {
+    Objects.requireNonNull(request, "Request must be set.");
+    if (intoSpan == null) {
+      return;
+    }
+
+    intoSpan.extract(request);
+  }
+
+}
diff --git a/java/client/src/org/openqa/selenium/remote/tracing/OpenCensusSpan.java b/java/client/src/org/openqa/selenium/remote/tracing/OpenCensusSpan.java
@@ -17,12 +17,14 @@
 
 package org.openqa.selenium.remote.tracing;
 
+import org.openqa.selenium.remote.http.HttpRequest;
+
 import io.opencensus.trace.AttributeValue;
 import io.opencensus.trace.Tracer;
 
 import java.util.Objects;
 
-class OpenCensusSpan implements Span {
+class OpenCensusSpan extends Span {
 
   private final io.opencensus.trace.Span span;
   private final DistributedTracer distributedTracer;
@@ -76,6 +78,16 @@ public Span createChild(String operation) {
     return child.activate();
   }
 
+  @Override
+  void inject(HttpRequest request) {
+    throw new UnsupportedOperationException("inject");
+  }
+
+  @Override
+  void extract(HttpRequest request) {
+    throw new UnsupportedOperationException("extract");
+  }
+
   @Override
   public void close() {
     span.end();

diff --git a/java/client/src/org/openqa/selenium/remote/tracing/OpenTracingSpan.java b/java/client/src/org/openqa/selenium/remote/tracing/OpenTracingSpan.java
@@ -17,11 +17,24 @@
 
 package org.openqa.selenium.remote.tracing;
 
+import com.google.common.collect.ImmutableSet;
+
+import org.openqa.selenium.remote.http.HttpRequest;
+
+import io.opentracing.SpanContext;
 import io.opentracing.Tracer;
+import io.opentracing.propagation.Format;
+import io.opentracing.propagation.TextMap;
 
+import java.util.AbstractMap;
+import java.util.Iterator;
+import java.util.Locale;
+import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
+import java.util.stream.StreamSupport;
 
-class OpenTracingSpan implements Span {
+class OpenTracingSpan extends Span {
 
   private final DistributedTracer distributedTracer;
   private final Tracer tracer;
@@ -76,9 +89,62 @@ public Span createChild(String operation) {
     return child.activate();
   }
 
+  @Override
+  void inject(HttpRequest request) {
+    tracer.inject(span.context(), Format.Builtin.HTTP_HEADERS, new HttpRequestInjector(request));
+  }
+
+  @Override
+  void extract(HttpRequest request) {
+    SpanContext context = tracer.extract(Format.Builtin.HTTP_HEADERS, new HttpRequestInjector(request));
+    for (Map.Entry<String, String> item : context.baggageItems()) {
+      addTraceTag(item.getKey(), item.getValue());
+    }
+  }
+
   @Override
   public void close() {
     span.finish();
     distributedTracer.remove(this);
   }
+
+  private class HttpRequestInjector implements TextMap {
+
+    private final Set<String> names = ImmutableSet.<String>builder()
+        .add("cache-control")
+        .add("connection")
+        .add("content-length")
+        .add("content-type")
+        .add("date")
+        .add("keep-alive")
+        .add("proxy-authorization")
+        .add("proxy-authenticate")
+        .add("proxy-connection")
+        .add("referer")
+        .add("te")
+        .add("trailer")
+        .add("transfer-encoding")
+        .add("upgrade")
+        .add("user-agent")
+        .build();
+    private final HttpRequest request;
+
+    HttpRequestInjector(HttpRequest request) {
+      this.request = request;
+    }
+
+    @Override
+    public Iterator<Map.Entry<String, String>> iterator() {
+      return StreamSupport.stream(request.getHeaderNames().spliterator(), false)
+          .filter(name -> names.contains(name.toLowerCase(Locale.US)))
+          .map(name -> (Map.Entry<String, String>) new AbstractMap.SimpleImmutableEntry<>(
+              name, request.getHeader(name)))
+          .iterator();
+    }
+
+    @Override
+    public void put(String key, String value) {
+      request.setHeader(key, value);
+    }
+  }
 }
diff --git a/java/client/src/org/openqa/selenium/remote/tracing/Span.java b/java/client/src/org/openqa/selenium/remote/tracing/Span.java
@@ -17,35 +17,40 @@
 
 package org.openqa.selenium.remote.tracing;
 
+import org.openqa.selenium.remote.http.HttpRequest;
+
 import java.io.Closeable;
 
-public interface Span extends Closeable {
+public abstract class Span implements Closeable {
 
-  Span createChild(String operation);
+  public abstract Span createChild(String operation);
 
   /**
    * Allows subclasses to indicate that this is the currently active span
    */
-  Span activate();
+  public abstract Span activate();
 
   /**
    * Add a tag that will be transmitted across the wire to allow remote traces
    * to also have the value. This is equivalent to OpenTracing's concept of
    * &quot;baggage&quot;.
    */
-  Span addTraceTag(String key, String value);
+  public abstract Span addTraceTag(String key, String value);
 
   /**
    * Add a piece of metadata to the span, which allows high cardinality data to
    * be added to the span. This data will not be propogated to other spans.
    */
-  Span addTag(String key, String value);
+  public abstract Span addTag(String key, String value);
 
-  Span addTag(String key, boolean value);
+  public abstract Span addTag(String key, boolean value);
 
-  Span addTag(String key, long value);
+  public abstract Span addTag(String key, long value);
 
   @Override
-  void close();
+  public abstract void close();
+
+  abstract void inject(HttpRequest request);
 
+  abstract void extract(HttpRequest request);
 }