From c7c9ecb3f2614c12c07d4fc5f9bb402987efd978 Mon Sep 17 00:00:00 2001
From: Simon Stewart <simon.m.stewart@gmail.com>
Date: Sat, 10 Nov 2018 22:17:13 +0000
Subject: [PATCH] Move tracing into the http client factory

---
 .../selenium/remote/HttpCommandExecutor.java  |  1 -
 .../selenium/remote/ProtocolHandshake.java    |  6 ++---
 .../remote/internal/OkHttpClient.java         |  3 ++-
 .../selenium/remote/tracing/HttpTracing.java  | 27 +++++++++++++++++++
 4 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/java/client/src/org/openqa/selenium/remote/HttpCommandExecutor.java b/java/client/src/org/openqa/selenium/remote/HttpCommandExecutor.java
index 9b0b55cfce62c..8d372992cff08 100644
--- a/java/client/src/org/openqa/selenium/remote/HttpCommandExecutor.java
+++ b/java/client/src/org/openqa/selenium/remote/HttpCommandExecutor.java
@@ -158,7 +158,6 @@ public Response execute(Command command) throws IOException {
     try (Span span = tracer.createSpan(command.getName(), tracer.getActiveSpan())) {
       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
index d1198aea5d566..cee90c882c2f7 100644
--- a/java/client/src/org/openqa/selenium/remote/ProtocolHandshake.java
+++ b/java/client/src/org/openqa/selenium/remote/ProtocolHandshake.java
@@ -39,7 +39,6 @@
 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;
@@ -100,9 +99,8 @@ private Optional<Result> createSession(HttpClient client, InputStream newSession
 
     HttpResponse response;
     long start = System.currentTimeMillis();
-    try (Span span = DistributedTracer.getInstance().getActiveSpan()) {
-      HttpTracing.inject(span, request);
-
+    DistributedTracer tracer = DistributedTracer.getInstance();
+    try (Span span = tracer.createSpan("NEW_SESSION", tracer.getActiveSpan())) {
       request.setHeader(CONTENT_LENGTH, String.valueOf(size));
       request.setHeader(CONTENT_TYPE, JSON_UTF_8.toString());
       request.setContent(newSessionBlob);
diff --git a/java/client/src/org/openqa/selenium/remote/internal/OkHttpClient.java b/java/client/src/org/openqa/selenium/remote/internal/OkHttpClient.java
index 8f2ea9bf25128..667875fb9f155 100644
--- a/java/client/src/org/openqa/selenium/remote/internal/OkHttpClient.java
+++ b/java/client/src/org/openqa/selenium/remote/internal/OkHttpClient.java
@@ -24,6 +24,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 okhttp3.ConnectionPool;
 import okhttp3.Credentials;
@@ -155,7 +156,7 @@ public HttpClient createClient(URL url) {
                    : response;
           });
 
-          return new OkHttpClient(client.build(), url);
+          return HttpTracing.decorate(new OkHttpClient(client.build(), url));
         }
       };
     }
diff --git a/java/client/src/org/openqa/selenium/remote/tracing/HttpTracing.java b/java/client/src/org/openqa/selenium/remote/tracing/HttpTracing.java
index 003d165239aea..cca7af25a6cec 100644
--- a/java/client/src/org/openqa/selenium/remote/tracing/HttpTracing.java
+++ b/java/client/src/org/openqa/selenium/remote/tracing/HttpTracing.java
@@ -17,7 +17,11 @@
 
 package org.openqa.selenium.remote.tracing;
 
+import org.openqa.selenium.remote.http.HttpClient;
 import org.openqa.selenium.remote.http.HttpRequest;
+import org.openqa.selenium.remote.http.HttpResponse;
+
+import io.opentracing.tag.Tags;
 
 import java.util.Objects;
 
@@ -33,6 +37,8 @@ public static void inject(Span span, HttpRequest request) {
       return;
     }
 
+    span.addTag(Tags.HTTP_METHOD.getKey(), request.getMethod().toString());
+    span.addTag(Tags.HTTP_URL.getKey(), request.getUri());
     span.inject(request);
   }
 
@@ -45,4 +51,25 @@ public static void extract(HttpRequest request, Span intoSpan) {
     intoSpan.extract(request);
   }
 
+  public static HttpClient decorate(HttpClient existing) {
+
+    return request -> {
+      Span span = DistributedTracer.getInstance().getActiveSpan();
+      inject(span, request);
+
+      try {
+        HttpResponse response = existing.execute(request);
+
+        if (span != null) {
+          span.addTag(Tags.HTTP_STATUS.getKey(), response.getStatus());
+        }
+
+        return response;
+      } catch (Throwable throwable) {
+        throw throwable;
+      }
+    };
+
+  }
+
 }