Add a test to check for trace propagation

It appears to work. Somewhat distressingly, we had to write our own OpenTracing implementation, since all the other ones out there come with a ton of dependencies we don't have, and they're not really designed to surface information for testing (though ours doesn't surface information well either).
SeleniumHQ · Nov 12, 2018 · 48c3eb6 · 48c3eb6
1 parent 46bdda0
commit 48c3eb6
Show file tree

Hide file tree

Showing 17 changed files with 577 additions and 47 deletions.
diff --git a/java/client/src/org/openqa/selenium/remote/BUCK b/java/client/src/org/openqa/selenium/remote/BUCK
@@ -134,12 +134,13 @@ java_library(name = 'remote-lib',
     '//java/client/src/org/openqa/selenium/json:json',
     '//java/client/src/org/openqa/selenium/remote/session:session',
     '//third_party/java/okhttp3:okhttp',
+    '//third_party/java/opencensus:opencensus-api',
+    '//third_party/java/opentracing:opentracing-api',
   ],
   deps = [
     ':http-session-id',
     '//java/client/src/org/openqa/selenium:selenium',
     '//third_party/java/contrib:opentracing-tracerresolver',
-    '//third_party/java/opencensus:opencensus-api',
     '//third_party/java/opentracing:opentracing-noop',
     '//third_party/java/guava:guava',
   ],

diff --git a/java/client/src/org/openqa/selenium/remote/tracing/CompoundSpan.java b/java/client/src/org/openqa/selenium/remote/tracing/CompoundSpan.java
@@ -65,6 +65,16 @@ public Span addTraceTag(String key, String value) {
     return this;
   }
 
+  @Override
+  public String getTraceTag(String key) {
+    // We assume that all the values are the same. This may not actually be true.
+    return allSpans.stream()
+        .map(span -> span.getTraceTag(key))
+        .filter(Objects::nonNull)
+        .findFirst()
+        .orElse(null);
+  }
+
   @Override
   public Span addTag(String key, String value) {
     Objects.requireNonNull(key, "Key must be set");
@@ -97,10 +107,4 @@ 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/DistributedTracer.java b/java/client/src/org/openqa/selenium/remote/tracing/DistributedTracer.java
@@ -17,15 +17,24 @@
 
 package org.openqa.selenium.remote.tracing;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 
 import org.openqa.selenium.BuildInfo;
 
+import io.opentracing.SpanContext;
 import io.opentracing.contrib.tracerresolver.TracerResolver;
 import io.opentracing.noop.NoopTracerFactory;
+import io.opentracing.propagation.Format;
+import io.opentracing.propagation.TextMap;
+import io.opentracing.propagation.TextMapExtractAdapter;
 
+import java.util.Iterator;
 import java.util.LinkedList;
+import java.util.Map;
 import java.util.Objects;
+import java.util.function.BiConsumer;
+import java.util.function.Supplier;
 
 /**
  * Represents an entry point for accessing all aspects of distributed tracing.
@@ -116,6 +125,59 @@ void remove(Span span) {
     ACTIVE_SPANS.get().removeIf(span::equals);
   }
 
+  public Span extract(String operationName, Iterator<Map.Entry<String, String>> traceTags) {
+    ImmutableMap.Builder<String, String> repeatableTags = ImmutableMap.builder();
+    traceTags.forEachRemaining(item -> repeatableTags.put(item.getKey(), item.getValue()));
+
+    ImmutableSet.Builder<Span> spans = ImmutableSet.builder();
+
+    TextMapExtractAdapter adapter = new TextMapExtractAdapter(repeatableTags.build());
+    for (io.opentracing.Tracer tracer : otTracers) {
+      SpanContext context = tracer.extract(Format.Builtin.TEXT_MAP, adapter);
+
+      spans.add(new OpenTracingSpan(this, tracer, context, operationName));
+    }
+
+    for (io.opencensus.trace.Tracer ocTracer : ocTracers) {
+      // TODO: implement for OpenCensus
+    }
+
+    return new CompoundSpan(this, spans.build()).activate();
+  }
+
+  private static class TextMapAdapter implements TextMap {
+
+    private final BiConsumer<String, String> setter;
+    private final Supplier<Iterator<Map.Entry<String, String>>> getter;
+
+    private TextMapAdapter(
+        BiConsumer<String, String> setter,
+        Supplier<Iterator<Map.Entry<String, String>>> getter) {
+
+      this.setter = setter;
+      this.getter = getter;
+    }
+
+    @Override
+    public Iterator<Map.Entry<String, String>> iterator() {
+      if (getter == null) {
+        throw new UnsupportedOperationException("iterator");
+      }
+      return getter.get();
+    }
+
+    @Override
+    public void put(String key, String value) {
+      if (setter == null) {
+        throw new UnsupportedOperationException("put");
+      }
+      if (key == null || value == null) {
+        return;
+      }
+      setter.accept(key, value);
+    }
+  }
+
   public static class Builder {
 
     private ImmutableSet.Builder<io.opencensus.trace.Tracer> ocTracers = ImmutableSet.builder();

diff --git a/java/client/src/org/openqa/selenium/remote/tracing/HttpTracing.java b/java/client/src/org/openqa/selenium/remote/tracing/HttpTracing.java
@@ -23,7 +23,11 @@
 
 import io.opentracing.tag.Tags;
 
+import java.util.AbstractMap;
+import java.util.Iterator;
+import java.util.Map;
 import java.util.Objects;
+import java.util.stream.StreamSupport;
 
 public class HttpTracing {
 
@@ -42,13 +46,18 @@ public static void inject(Span span, HttpRequest request) {
     span.inject(request);
   }
 
-  public static void extract(HttpRequest request, Span intoSpan) {
+  public static Span extract(DistributedTracer tracer, String operationName, HttpRequest request) {
     Objects.requireNonNull(request, "Request must be set.");
-    if (intoSpan == null) {
-      return;
-    }
 
-    intoSpan.extract(request);
+    Iterator<Map.Entry<String, String>> iterator =
+        StreamSupport.stream(request.getHeaderNames().spliterator(), false)
+            .filter(name -> request.getHeader(name) != null)
+            .map(name -> (Map.Entry<String, String>) new AbstractMap.SimpleImmutableEntry<>(
+                name,
+                request.getHeader(name)))
+            .iterator();
+
+    return tracer.extract(operationName, iterator);
   }
 
   public static HttpClient decorate(HttpClient existing) {
@@ -69,7 +78,6 @@ public static HttpClient decorate(HttpClient existing) {
         throw throwable;
       }
     };
-
   }
 
 }
diff --git a/java/client/src/org/openqa/selenium/remote/tracing/OpenCensusSpan.java b/java/client/src/org/openqa/selenium/remote/tracing/OpenCensusSpan.java
@@ -63,6 +63,11 @@ public Span addTraceTag(String key, String value) {
     return this;
   }
 
+  @Override
+  public String getTraceTag(String key) {
+    return span.getContext().getTracestate().get(key);
+  }
+
   @Override
   public Span addTag(String key, String value) {
     span.putAttribute(Objects.requireNonNull(key), AttributeValue.stringAttributeValue(value));
@@ -92,11 +97,6 @@ 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
@@ -21,7 +21,6 @@
 
 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;
@@ -43,7 +42,7 @@ class OpenTracingSpan extends Span {
   OpenTracingSpan(
       DistributedTracer distributedTracer,
       Tracer tracer,
-      io.opentracing.Span parent,
+      io.opentracing.SpanContext parent,
       String operation) {
     this.distributedTracer = Objects.requireNonNull(distributedTracer);
     this.tracer = Objects.requireNonNull(tracer);
@@ -72,6 +71,11 @@ public Span addTraceTag(String key, String value) {
     return this;
   }
 
+  @Override
+  public String getTraceTag(String key) {
+    return span.getBaggageItem(key);
+  }
+
   @Override
   public Span addTag(String key, String value) {
     span.setTag(key, value);
@@ -92,7 +96,7 @@ public Span addTag(String key, long value) {
 
   @Override
   public Span createChild(String operation) {
-    Span child = new OpenTracingSpan(distributedTracer, tracer, span, operation);
+    Span child = new OpenTracingSpan(distributedTracer, tracer, span.context(), operation);
     return child.activate();
   }
 
@@ -101,17 +105,6 @@ 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));
-    if (context == null) {
-      return;
-    }
-    for (Map.Entry<String, String> item : context.baggageItems()) {
-      addTraceTag(item.getKey(), item.getValue());
-    }
-  }
-
   @Override
   public void close() {
     span.finish();

diff --git a/java/client/src/org/openqa/selenium/remote/tracing/Span.java b/java/client/src/org/openqa/selenium/remote/tracing/Span.java
@@ -39,6 +39,8 @@ public abstract class Span implements Closeable {
    */
   public abstract Span addTraceTag(String key, String value);
 
+  public abstract String getTraceTag(String key);
+
   /**
    * 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.
@@ -53,6 +55,4 @@ public abstract class Span implements Closeable {
   public abstract void close();
 
   abstract void inject(HttpRequest request);
-
-  abstract void extract(HttpRequest request);
 }
diff --git a/java/client/test/org/openqa/selenium/environment/webserver/JreAppServer.java b/java/client/test/org/openqa/selenium/environment/webserver/JreAppServer.java
@@ -20,6 +20,7 @@
 import static com.google.common.net.HttpHeaders.CONTENT_TYPE;
 import static com.google.common.net.MediaType.JSON_UTF_8;
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.util.concurrent.TimeUnit.SECONDS;
 import static java.util.stream.Collectors.mapping;
 import static java.util.stream.Collectors.toList;
 import static org.openqa.selenium.remote.http.HttpMethod.GET;
@@ -48,11 +49,14 @@
 import java.net.URL;
 import java.nio.file.Paths;
 import java.util.AbstractMap;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 import java.util.function.BiConsumer;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -73,10 +77,13 @@ public JreAppServer() {
             HttpRequest req = new SunHttpRequest(httpExchange);
             HttpResponse resp = new SunHttpResponse(httpExchange);
 
-            mappings.entrySet().stream()
-                .filter(entry -> entry.getKey().test(req))
+            List<Predicate<HttpRequest>> reversedKeys = new ArrayList<>(mappings.keySet());
+            Collections.reverse(reversedKeys);
+
+            reversedKeys.stream()
+                .filter(pred -> pred.test(req))
                 .findFirst()
-                .map(Map.Entry::getValue)
+                .map(mappings::get)
                 .orElseGet(() -> (in, out) -> {
                   out.setStatus(404);
                   out.setContent("".getBytes(UTF_8));
@@ -91,21 +98,22 @@ public JreAppServer() {
   }
 
   protected JreAppServer emulateJettyAppServer() {
-    addHandler(GET, "/encoding", new EncodingHandler());
-    addHandler(GET, "/page", new PageHandler());
-    addHandler(GET, "/redirect", new RedirectHandler(whereIs("/")));
-    addHandler(GET, "/sleep", new SleepingHandler());
-    addHandler(POST, "/upload", new UploadHandler());
-
     String javascript = locate("javascript").toAbsolutePath().toString();
     String common = locate("common/src/web").toAbsolutePath().toString();
+    // Listed first, so considered last
     addHandler(
         GET,
         "/",
         new StaticContent(
             path -> Paths.get(common + path),
             path -> Paths.get(javascript + path.substring("/javascript".length()))));
 
+    addHandler(GET, "/encoding", new EncodingHandler());
+    addHandler(GET, "/page", new PageHandler());
+    addHandler(GET, "/redirect", new RedirectHandler(whereIs("/")));
+    addHandler(GET, "/sleep", new SleepingHandler());
+    addHandler(POST, "/upload", new UploadHandler());
+
     return this;
   }
 
@@ -119,6 +127,7 @@ public JreAppServer addHandler(
 
   public void start() {
     server.start();
+    PortProber.waitForPortUp(server.getAddress().getPort(), 5, SECONDS);
   }
 
   public void stop() {

diff --git a/java/client/test/org/openqa/selenium/remote/tracing/BUCK b/java/client/test/org/openqa/selenium/remote/tracing/BUCK
@@ -7,6 +7,8 @@ java_test(
     srcs = glob(["*.java"]),
     deps = [
         "//java/client/src/org/openqa/selenium/remote:remote",
+        "//java/client/test/org/openqa/selenium/environment:environment",
+        "//java/client/test/org/openqa/selenium/remote/tracing/simple:simple",
         "//third_party/java/assertj:assertj",
         "//third_party/java/junit:junit",
     ],