Merge pull request #192 from manimaul/wk/TACO-35-trace-span-prop

TACO-35 - adding span to Request object
xjdr · Mar 29, 2018 · b6f2294 · b6f2294
2 parents 4137eea + 7e67556
commit b6f2294
Show file tree

Hide file tree

Showing 13 changed files with 155 additions and 23 deletions.
diff --git a/xio-core/src/main/java/com/xjeffrose/xio/client/http/HttpClient.java b/xio-core/src/main/java/com/xjeffrose/xio/client/http/HttpClient.java
@@ -21,7 +21,7 @@ public HttpClient(XioClient client, String hostHeader) {
 
   public void write(HttpRequest request, TraceContext context) {
     request.headers().set("Host", hostHeader);
-    client.write(new XioRequest(request, context));
+    client.write(new XioRequest<>(request, context));
   }
 
   public void write(HttpRequest request) {

diff --git a/xio-core/src/main/java/com/xjeffrose/xio/http/DefaultFullRequest.java b/xio-core/src/main/java/com/xjeffrose/xio/http/DefaultFullRequest.java
@@ -1,12 +1,13 @@
 package com.xjeffrose.xio.http;
 
+import brave.Span;
 import com.google.auto.value.AutoValue;
 import com.xjeffrose.xio.core.internal.UnstableApi;
 import io.netty.buffer.ByteBuf;
 import io.netty.handler.codec.http.HttpHeaderNames;
 import io.netty.handler.codec.http.HttpMethod;
 import java.util.Optional;
-import lombok.Builder;
+import javax.annotation.Nullable;
 import lombok.ToString;
 
 /** Value class for representing an outgoing HTTP1/2 Request, for use in a client. */
@@ -30,6 +31,9 @@ public boolean startOfStream() {
 
   public abstract int streamId();
 
+  @Nullable
+  public abstract Span traceSpan();
+
   /** Not intended to be called. */
   @Override
   public String version() {
@@ -53,6 +57,8 @@ public abstract static class Builder {
 
     public abstract Builder streamId(int streamId);
 
+    public abstract Builder traceSpan(Span span);
+
     public abstract DefaultFullRequest build();
 
     abstract Optional<Headers> headers();

diff --git a/xio-core/src/main/java/com/xjeffrose/xio/http/DefaultStreamingRequest.java b/xio-core/src/main/java/com/xjeffrose/xio/http/DefaultStreamingRequest.java
@@ -1,18 +1,19 @@
 package com.xjeffrose.xio.http;
 
+import brave.Span;
 import com.google.auto.value.AutoValue;
 import com.xjeffrose.xio.core.internal.UnstableApi;
 import io.netty.handler.codec.http.HttpHeaderNames;
 import io.netty.handler.codec.http.HttpMethod;
 import java.util.Optional;
-import lombok.Builder;
+import javax.annotation.Nullable;
 import lombok.ToString;
 
 /** Value class for representing a streaming outgoing HTTP1/2 Request, for use in a client. */
 @UnstableApi
 @AutoValue
 @ToString
-public abstract class DefaultStreamingRequest implements StreamingRequest {
+public abstract class DefaultStreamingRequest implements StreamingRequest, Traceable {
 
   @Override
   public boolean startOfStream() {
@@ -27,6 +28,9 @@ public boolean startOfStream() {
 
   public abstract int streamId();
 
+  @Nullable
+  public abstract Span traceSpan();
+
   /** Not intended to be called. */
   @Override
   public String version() {
@@ -48,6 +52,8 @@ public abstract static class Builder {
 
     public abstract Builder streamId(int streamId);
 
+    public abstract Builder traceSpan(Span span);
+
     public abstract DefaultStreamingRequest build();
 
     abstract Optional<Headers> headers();

diff --git a/xio-core/src/main/java/com/xjeffrose/xio/http/Request.java b/xio-core/src/main/java/com/xjeffrose/xio/http/Request.java
@@ -7,7 +7,7 @@
 import io.netty.handler.codec.http.HttpMethod;
 
 @UnstableApi
-public interface Request {
+public interface Request extends Traceable {
 
   boolean startOfStream();
 

diff --git a/xio-core/src/main/java/com/xjeffrose/xio/http/StreamingRequestData.java b/xio-core/src/main/java/com/xjeffrose/xio/http/StreamingRequestData.java
@@ -1,8 +1,10 @@
 package com.xjeffrose.xio.http;
 
+import brave.Span;
 import com.xjeffrose.xio.core.internal.UnstableApi;
 import io.netty.buffer.ByteBuf;
 import io.netty.handler.codec.http.HttpMethod;
+import javax.annotation.Nullable;
 import lombok.ToString;
 
 @UnstableApi
@@ -11,12 +13,20 @@ public class StreamingRequestData implements Request, StreamingData {
 
   private final Request request;
   private final StreamingData data;
+  private final Span span;
 
-  public StreamingRequestData(Request request, StreamingData data) {
+  public StreamingRequestData(Request request, StreamingData data, Span span) {
     this.request = request;
     this.data = data;
+    this.span = span;
+  }
+
+  public StreamingRequestData(Request request, StreamingData data) {
+    this(request, data, null);
   }
 
+  // region Request
+
   @Override
   public boolean startOfStream() {
     return false;
@@ -62,6 +72,20 @@ public ByteBuf body() {
     return request.body();
   }
 
+  // endregion
+
+  // region Traceable
+
+  @Nullable
+  @Override
+  public Span traceSpan() {
+    return span;
+  }
+
+  // endregion
+
+  // region StreamingData
+
   @Override
   public ByteBuf content() {
     return data.content();
@@ -76,4 +100,7 @@ public boolean endOfStream() {
   public Headers trailingHeaders() {
     return data.trailingHeaders();
   }
+
+  // endregion
+
 }
diff --git a/xio-core/src/main/java/com/xjeffrose/xio/http/Traceable.java b/xio-core/src/main/java/com/xjeffrose/xio/http/Traceable.java
@@ -0,0 +1,10 @@
+package com.xjeffrose.xio.http;
+
+import brave.Span;
+import javax.annotation.Nullable;
+
+public interface Traceable {
+
+  @Nullable
+  Span traceSpan();
+}
diff --git a/xio-core/src/main/java/com/xjeffrose/xio/http/internal/FullHttp1Request.java b/xio-core/src/main/java/com/xjeffrose/xio/http/internal/FullHttp1Request.java
@@ -1,20 +1,28 @@
 package com.xjeffrose.xio.http.internal;
 
+import brave.Span;
 import com.xjeffrose.xio.http.FullRequest;
 import com.xjeffrose.xio.http.Headers;
 import io.netty.buffer.ByteBuf;
 import io.netty.handler.codec.http.FullHttpRequest;
 import io.netty.handler.codec.http.HttpMethod;
 import io.netty.handler.codec.http.HttpUtil;
+import javax.annotation.Nullable;
 
 public class FullHttp1Request implements FullRequest {
 
   private final FullHttpRequest delegate;
   private final Http1Headers headers;
+  private final Span span;
 
-  public FullHttp1Request(FullHttpRequest delegate) {
+  public FullHttp1Request(FullHttpRequest delegate, @Nullable Span span) {
     this.delegate = delegate;
-    headers = new Http1Headers(delegate.headers());
+    this.headers = new Http1Headers(delegate.headers());
+    this.span = span;
+  }
+
+  public FullHttp1Request(FullHttpRequest delegate) {
+    this(delegate, null);
   }
 
   @Override
@@ -61,4 +69,16 @@ public boolean hasBody() {
   public ByteBuf body() {
     return delegate.content();
   }
+
+  // endregion
+
+  // region Traceable
+
+  @Nullable
+  @Override
+  public Span traceSpan() {
+    return null;
+  }
+
+  // endregion
 }
diff --git a/xio-core/src/main/java/com/xjeffrose/xio/http/internal/FullHttp2Request.java b/xio-core/src/main/java/com/xjeffrose/xio/http/internal/FullHttp2Request.java
@@ -1,22 +1,32 @@
 package com.xjeffrose.xio.http.internal;
 
+import brave.Span;
 import com.xjeffrose.xio.http.FullRequest;
 import com.xjeffrose.xio.http.Headers;
 import io.netty.handler.codec.http.HttpMethod;
 import io.netty.handler.codec.http2.Http2Headers;
+import javax.annotation.Nullable;
 
 public class FullHttp2Request implements FullRequest {
 
   private final Http2Headers delegate;
   private final Http2HeadersWrapper headers;
   private final int streamId;
+  private final Span span;
 
-  public FullHttp2Request(Http2Headers delegate, int streamId) {
+  public FullHttp2Request(Http2Headers delegate, int streamId, @Nullable Span span) {
     this.delegate = delegate;
-    headers = new Http2HeadersWrapper(delegate);
+    this.headers = new Http2HeadersWrapper(delegate);
     this.streamId = streamId;
+    this.span = span;
+  }
+
+  public FullHttp2Request(Http2Headers delegate, int streamId) {
+    this(delegate, streamId, null);
   }
 
+  // region Request
+
   @Override
   public boolean startOfStream() {
     return true;
@@ -56,4 +66,17 @@ public Headers headers() {
   public boolean keepAlive() {
     return true;
   }
+
+  // endregion
+
+  // region Traceable
+
+  @Nullable
+  @Override
+  public Span traceSpan() {
+    return span;
+  }
+
+  // endregion
+
 }
diff --git a/xio-core/src/main/java/com/xjeffrose/xio/http/internal/Http1Request.java b/xio-core/src/main/java/com/xjeffrose/xio/http/internal/Http1Request.java
@@ -1,12 +1,14 @@
 package com.xjeffrose.xio.http.internal;
 
+import brave.Span;
 import com.xjeffrose.xio.http.Headers;
 import com.xjeffrose.xio.http.StreamingRequest;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.handler.codec.http.HttpMethod;
 import io.netty.handler.codec.http.HttpRequest;
 import io.netty.handler.codec.http.HttpUtil;
+import javax.annotation.Nullable;
 import lombok.ToString;
 
 // TODO(CK): Rename this to StreamingHttp1Request
@@ -17,12 +19,20 @@ public class Http1Request implements StreamingRequest {
 
   protected final HttpRequest delegate;
   private final Http1Headers headers;
+  private final Span span;
 
-  public Http1Request(HttpRequest delegate) {
+  public Http1Request(HttpRequest delegate, @Nullable Span span) {
     this.delegate = delegate;
     headers = new Http1Headers(delegate.headers());
+    this.span = span;
+  }
+
+  public Http1Request(HttpRequest delegate) {
+    this(delegate, null);
   }
 
+  // region Request
+
   @Override
   public boolean startOfStream() {
     return true;
@@ -62,4 +72,17 @@ public boolean keepAlive() {
   public ByteBuf body() {
     return Unpooled.EMPTY_BUFFER;
   }
+
+  // endregion
+
+  // region Traceable
+
+  @Nullable
+  @Override
+  public Span traceSpan() {
+    return span;
+  }
+
+  // endregion
+
 }
diff --git a/xio-core/src/main/java/com/xjeffrose/xio/http/internal/StreamingHttp2Request.java b/xio-core/src/main/java/com/xjeffrose/xio/http/internal/StreamingHttp2Request.java
@@ -1,11 +1,13 @@
 package com.xjeffrose.xio.http.internal;
 
+import brave.Span;
 import com.xjeffrose.xio.http.Headers;
 import com.xjeffrose.xio.http.StreamingRequest;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.handler.codec.http.HttpMethod;
 import io.netty.handler.codec.http2.Http2Headers;
+import javax.annotation.Nullable;
 import lombok.ToString;
 
 /** Wrap an incoming Http2 Request, for use in a server. */
@@ -15,13 +17,21 @@ public class StreamingHttp2Request implements StreamingRequest {
   private final Http2Headers delegate;
   private final Http2HeadersWrapper headers;
   private final int streamId;
+  private final Span span;
 
-  public StreamingHttp2Request(Http2Headers delegate, int streamId) {
+  public StreamingHttp2Request(Http2Headers delegate, int streamId, @Nullable Span span) {
     this.delegate = delegate;
-    headers = new Http2HeadersWrapper(delegate);
+    this.headers = new Http2HeadersWrapper(delegate);
     this.streamId = streamId;
+    this.span = span;
+  }
+
+  public StreamingHttp2Request(Http2Headers delegate, int streamId) {
+    this(delegate, streamId, null);
   }
 
+  // region Request
+
   @Override
   public boolean startOfStream() {
     return true;
@@ -66,4 +76,17 @@ public boolean keepAlive() {
   public ByteBuf body() {
     return Unpooled.EMPTY_BUFFER;
   }
+
+  // endregion
+
+  // region Traceable
+
+  @Nullable
+  @Override
+  public Span traceSpan() {
+    return span;
+  }
+
+  // endregion
+
 }
diff --git a/xio-core/src/main/java/com/xjeffrose/xio/tracing/HttpClientRequestTracingHandler.java b/xio-core/src/main/java/com/xjeffrose/xio/tracing/HttpClientRequestTracingHandler.java
@@ -1,11 +1,6 @@
 package com.xjeffrose.xio.tracing;
 
-import brave.Span;
-import io.netty.channel.ChannelFuture;
-import io.netty.channel.ChannelFutureListener;
-import io.netty.channel.ChannelHandlerContext;
-import io.netty.channel.ChannelOutboundHandlerAdapter;
-import io.netty.channel.ChannelPromise;
+import io.netty.channel.*;
 import io.netty.handler.codec.http.HttpRequest;
 
 class HttpClientRequestTracingHandler extends ChannelOutboundHandlerAdapter {
@@ -26,7 +21,7 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise)
     }
 
     HttpRequest request = (HttpRequest) msg;
-    Span span = state.onRequest(ctx, request);
+    state.onRequest(ctx, request);
 
     ctx.write(msg, promise)
         .addListener(