DefaultTrace impl with basic test

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/DefaultTrace.java

@@ -0,0 +1,147 @@
+package org.springframework.cloud.sleuth.trace;
+
+import static org.springframework.cloud.sleuth.trace.Utils.error;
+
+import org.springframework.cloud.sleuth.trace.sampler.IsTracingSampler;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.UUID;
+
+/**
+ * @author Spencer Gibb
+ */
+public class DefaultTrace implements Trace {
+
+	//TODO: no default?, let autoconfig
+	private Sampler<?> defaultSampler = new IsTracingSampler(this);
+
+	private IdGenerator idGenerator = new IdGenerator() {
+		@Override
+		public String create() {
+			return UUID.randomUUID().toString();
+		}
+	};
+
+	private Collection<SpanReceiver> spanReceivers;
+
+	@Override
+	public TraceScope startSpan(String description) {
+		return this.startSpan(description, defaultSampler);
+	}
+
+	@Override
+	public TraceScope startSpan(String description, TraceInfo tinfo) {
+		if (tinfo == null) return continueSpan(null);
+		MilliSpan span = MilliSpan.builder()
+				.begin(System.currentTimeMillis())
+				.description(description)
+				.traceId(tinfo.getTraceId())
+				.spanId(idGenerator.create())
+				.parents(Collections.singletonList(tinfo.getSpanId()))
+				//TODO: when lombok plugin supports @Singular parent(tinfo.getSpanId()).
+				.build();
+		return continueSpan(span);
+	}
+
+	@Override
+	public TraceScope startSpan(String description, Span parent) {
+		if (parent == null) {
+			return startSpan(description);
+		}
+		Span currentSpan = getCurrentSpan();
+		if ((currentSpan != null) && (currentSpan != parent)) {
+			error("HTrace client error: thread " +
+					Thread.currentThread().getName() + " tried to start a new Span " +
+					"with parent " + parent.toString() + ", but there is already a " +
+					"currentSpan " + currentSpan);
+		}
+		return continueSpan(createChild(parent, description));
+	}
+
+	@Override
+	public <T> TraceScope startSpan(String description, Sampler<T> s) {
+		return startSpan(description, s, null);
+	}
+
+	@Override
+	public <T> TraceScope startSpan(String description, Sampler<T> s, T info) {
+		Span span = null;
+		if (isTracing() || s.next(info)) {
+			span = createNew(description);
+		}
+		return continueSpan(span);
+	}
+
+	protected Span createNew(String description) {
+		Span parent = getCurrentSpan();
+		if (parent == null) {
+			return MilliSpan.builder()
+					.begin(System.currentTimeMillis())
+					.description(description)
+					.traceId(idGenerator.create())
+					.spanId(idGenerator.create())
+					.build();
+		} else {
+			return createChild(parent, description);
+		}
+	}
+
+	protected Span createChild(Span parent, String childDescription) {
+		return MilliSpan.builder().
+				begin(System.currentTimeMillis()).
+				description(childDescription).
+				traceId(parent.getTraceId()).
+				parents(Collections.singletonList(parent.getSpanId())).
+				//TODO: when lombok plugin supports @Singular parent(parent.getSpanId()).
+				spanId(idGenerator.create()).
+				processId(parent.getProcessId()).
+				build();
+	}
+
+	@Override
+	public TraceScope continueSpan(Span span) {
+		// Return an empty TraceScope that does nothing on close
+		if (span == null) return NullScope.INSTANCE;
+		Span oldSpan = getCurrentSpan();
+		SpanHolder.setCurrentSpan(span);
+		return new TraceScope(this, span, oldSpan);
+	}
+
+	protected Span getCurrentSpan() {
+		return SpanHolder.getCurrentSpan();
+	}
+
+	@Override
+	public void addKVAnnotation(String key, String value) {
+		Span s = getCurrentSpan();
+		if (s != null) {
+			s.addKVAnnotation(key, value);
+		}
+	}
+
+	@Override
+	public boolean isTracing() {
+		return getCurrentSpan() != null;
+	}
+
+	@Override
+	public void deliver(Span span) {
+		for (SpanReceiver receiver : spanReceivers) {
+			receiver.receiveSpan(span);
+		}
+	}
+
+	public void setDefaultSampler(Sampler<?> defaultSampler) {
+		this.defaultSampler = defaultSampler;
+	}
+
+	public void setIdGenerator(IdGenerator idGenerator) {
+		this.idGenerator = idGenerator;
+	}
+
+	@Override
+	public void setSpanReceivers(Collection<SpanReceiver> spanReceivers) {
+		this.spanReceivers = spanReceivers;
+	}
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/MilliSpan.java

@@ -1,29 +1,31 @@
 package org.springframework.cloud.sleuth.trace;
 
-import lombok.Builder;
-import lombok.Data;
-import lombok.Singular;
-
 import java.util.List;
 import java.util.Map;
 
+import lombok.Builder;
+import lombok.Singular;
+import lombok.Value;
+import lombok.experimental.NonFinal;
+
 /**
  * @author Spencer Gibb
  */
-@Data
+@Value
 @Builder
 public class MilliSpan implements Span {
-	private final long begin;
-	private long end;
-	private final String description;
-	private final String traceId;
+	private long begin;
+	@NonFinal
+	private long end = 0;
+	private String description;
+	private String traceId;
 	@Singular
-	private final List<String> parents;
-	private final String spanId;
-	private final Map<String, String> kVAnnotations;
-	private final String processId;
+	private List<String> parents;
+	private String spanId;
+	private Map<String, String> kVAnnotations;
+	private String processId;
 	@Singular
-	private final List<TimelineAnnotation> timelineAnnotations;
+	private List<TimelineAnnotation> timelineAnnotations;
 
 	@Override
 	public synchronized void stop() {
@@ -32,7 +34,6 @@ public class MilliSpan implements Span {
 				throw new IllegalStateException("Span for " + description
 						+ " has not been started");
 			end = System.currentTimeMillis();
-			//TODO figure out how to Trace.deliver(this)
 		}
 	}
 
@@ -50,11 +51,6 @@ public class MilliSpan implements Span {
 		return begin != 0 && end == 0;
 	}
 
-	@Override
-	public Span child(String description) {
-		return null;
-	}
-
 	@Override
 	public void addKVAnnotation(String key, String value) {
 		kVAnnotations.put(key, value);

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/NullScope.java

@@ -0,0 +1,31 @@
+package org.springframework.cloud.sleuth.trace;
+
+/**
+ * @author Spencer Gibb
+ */
+/**
+ * Singleton instance representing an empty {@link TraceScope}.
+ */
+public final class NullScope extends TraceScope {
+
+	public static final TraceScope INSTANCE = new NullScope();
+
+	private NullScope() {
+		super(null, null, null);
+	}
+
+	@Override
+	public Span detach() {
+		return null;
+	}
+
+	@Override
+	public void close() {
+		return;
+	}
+
+	@Override
+	public String toString() {
+		return "NullScope";
+	}
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/Span.java

@@ -59,11 +59,6 @@ public interface Span {
 	 */
 	String getTraceId();
 
-	/**
-	 * Create a child span of this span with the given description
-	 */
-	Span child(String description);
-
 	/**
 	 * Returns the parent IDs of the span.<p/>
 	 * <p/>

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/SpanHolder.java

@@ -1,14 +1,13 @@
 package org.springframework.cloud.sleuth.trace;
 
-import org.apache.commons.logging.Log;
-import org.apache.commons.logging.LogFactory;
+import lombok.extern.apachecommons.CommonsLog;
 
 /**
  * @author Spencer Gibb
  */
+@CommonsLog
 public class SpanHolder {
 	private static final ThreadLocal<Span> currentSpan = new ThreadLocal<>();
-	private static final Log log = LogFactory.getLog(SpanHolder.class);
 
 	public static Span getCurrentSpan() {
 		return currentSpan.get();

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/Trace.java

@@ -81,12 +81,7 @@ public interface Trace {
 	 */
 	boolean isTracing();
 
-	/**
-	 * If we are tracing, return the current span, else null
-	 *
-	 * @return Span representing the current trace, or null if not tracing.
-	 */
-	Span currentSpan();
+	public void deliver(Span span);
 
 	void setSpanReceivers(Collection<SpanReceiver> spanReceivers);
 }

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/TraceScope.java

@@ -1,17 +1,18 @@
 package org.springframework.cloud.sleuth.trace;
 
-import lombok.Data;
-import lombok.extern.apachecommons.CommonsLog;
-
 import java.io.Closeable;
-import java.io.IOException;
+
+import lombok.Data;
+import lombok.SneakyThrows;
 
 /**
  * @author Spencer Gibb
  */
 @Data
-@CommonsLog
 public class TraceScope implements Closeable {
+
+	private final Trace trace;
+
 	/**
 	 * the span for this scope
 	 */
@@ -24,6 +25,12 @@ public class TraceScope implements Closeable {
 
 	private boolean detached = false;
 
+	public TraceScope(Trace trace, Span span, Span savedSpan) {
+		this.trace = trace;
+		this.span = span;
+		this.savedSpan = savedSpan;
+	}
+
 	/**
 	 * Remove this span as the current thread, but don't stop it yet or
 	 * send it for collection. This is useful if the span object is then
@@ -33,14 +40,14 @@ public class TraceScope implements Closeable {
 	 */
 	public Span detach() {
 		if (detached) {
-			error("Tried to detach trace span " + span + " but " +
+			Utils.error("Tried to detach trace span " + span + " but " +
 					"it has already been detached.");
 		}
 		detached = true;
 
 		Span cur = SpanHolder.getCurrentSpan();
 		if (cur != span) {
-			error("Tried to detach trace span " + span + " but " +
+			Utils.error("Tried to detach trace span " + span + " but " +
 					"it is not the current span for the " +
 					Thread.currentThread().getName() + " thread.  You have " +
 					"probably forgotten to close or detach " + cur);
@@ -51,25 +58,23 @@ public class TraceScope implements Closeable {
 	}
 
 	@Override
-	public void close() throws IOException {
+	@SneakyThrows
+	public void close() {
 		if (detached) {
 			return;
 		}
 		detached = true;
 		Span cur = SpanHolder.getCurrentSpan();
 		if (cur != span) {
-			error("Tried to close trace span " + span + " but " +
+			Utils.error("Tried to close trace span " + span + " but " +
 					"it is not the current span for the " +
 					Thread.currentThread().getName() + " thread.  You have " +
 					"probably forgotten to close or detach " + cur);
 		} else {
 			span.stop();
+			trace.deliver(span);
 			SpanHolder.setCurrentSpan(savedSpan);
 		}
 	}
 
-	private void error(String msg) {
-		log.error(msg);
-		throw new RuntimeException(msg);
-	}
 }

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/Utils.java

@@ -0,0 +1,14 @@
+package org.springframework.cloud.sleuth.trace;
+
+import lombok.extern.apachecommons.CommonsLog;
+
+/**
+ * @author Spencer Gibb
+ */
+@CommonsLog
+public abstract class Utils {
+	public static void error(String msg) {
+		error(msg);
+		throw new RuntimeException(msg);
+	}
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/receiver/ArrayListSpanReceiver.java

@@ -0,0 +1,29 @@
+package org.springframework.cloud.sleuth.trace.receiver;
+
+import org.springframework.cloud.sleuth.trace.Span;
+import org.springframework.cloud.sleuth.trace.SpanReceiver;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * @author Spencer Gibb
+ */
+public class ArrayListSpanReceiver implements SpanReceiver {
+	private final ArrayList<Span> spans = new ArrayList<>();
+
+	@Override
+	public void receiveSpan(Span span) {
+		spans.add(span);
+	}
+
+	public List<Span> getSpans() {
+		return spans;
+	}
+
+	@Override
+	public void close() throws IOException {
+
+	}
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/sampler/AlwaysSampler.java

@@ -0,0 +1,13 @@
+package org.springframework.cloud.sleuth.trace.sampler;
+
+import org.springframework.cloud.sleuth.trace.Sampler;
+
+/**
+ * @author Spencer Gibb
+ */
+public class AlwaysSampler implements Sampler<Object> {
+	@Override
+	public boolean next(Object info) {
+		return true;
+	}
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/trace/sampler/IsTracingSampler.java

@@ -0,0 +1,20 @@
+package org.springframework.cloud.sleuth.trace.sampler;
+
+import org.springframework.cloud.sleuth.trace.Sampler;
+import org.springframework.cloud.sleuth.trace.Trace;
+
+/**
+ * @author Spencer Gibb
+ */
+public class IsTracingSampler implements Sampler<Object> {
+	private final Trace trace;
+
+	public IsTracingSampler(Trace trace) {
+		this.trace = trace;
+	}
+
+	@Override
+	public boolean next(Object info) {
+		return this.trace.isTracing();
+	}
+}

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/trace/DefaultTraceTests.java

@@ -0,0 +1,100 @@
+package org.springframework.cloud.sleuth.trace;
+
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
+import static org.hamcrest.Matchers.notNullValue;
+import static org.junit.Assert.assertThat;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+import org.junit.Test;
+import org.springframework.cloud.sleuth.trace.receiver.ArrayListSpanReceiver;
+import org.springframework.cloud.sleuth.trace.sampler.AlwaysSampler;
+
+/**
+ * @author Spencer Gibb
+ */
+public class DefaultTraceTests {
+
+	public static final String CREATE_SIMPLE_TRACE = "createSimpleTrace";
+	public static final String IMPORTANT_WORK_1 = "important work 1";
+	public static final String IMPORTANT_WORK_2 = "important work 2";
+
+	@Test
+	public void tracingWorks() {
+		DefaultTrace trace = new DefaultTrace();
+		ArrayListSpanReceiver spanReceiver = new ArrayListSpanReceiver();
+		trace.setSpanReceivers(Collections.<SpanReceiver> singletonList(spanReceiver));
+
+		TraceScope scope = trace.startSpan(CREATE_SIMPLE_TRACE, new AlwaysSampler());
+		try {
+			importantWork1(trace);
+		}
+		finally {
+			scope.close();
+		}
+
+		List<Span> spans = spanReceiver.getSpans();
+		assertThat("spans was null", spans, is(notNullValue()));
+		assertThat("spans was empty", spans.isEmpty(), not(true));
+		assertThat("spans was wrong size", spans.size(), is(3));
+
+		Span root = assertSpan(spans, null, CREATE_SIMPLE_TRACE);
+		Span child = assertSpan(spans, root.getSpanId(), IMPORTANT_WORK_1);
+		Span grandChild = assertSpan(spans, child.getSpanId(), IMPORTANT_WORK_2);
+
+		List<Span> gen4 = findSpans(spans, grandChild.getSpanId());
+		assertThat("gen4 was non-empty", gen4.isEmpty(), is(true));
+	}
+
+	private Span assertSpan(List<Span> spans, String parentId, String desc) {
+		List<Span> found = findSpans(spans, parentId);
+		assertThat("more than one span with parentId " + parentId, found.size(), is(1));
+		Span span = found.get(0);
+		assertThat("description is wrong for span with parentId " + parentId,
+				span.getDescription(), is(desc));
+		return span;
+	}
+
+	private List<Span> findSpans(List<Span> spans, String parentId) {
+		List<Span> found = new ArrayList<>();
+		for (Span span : spans) {
+			if (parentId == null && span.getParents().isEmpty()) {
+				found.add(span);
+			}
+			else if (span.getParents().contains(parentId)) {
+				found.add(span);
+			}
+		}
+		return found;
+	}
+
+	private void importantWork1(Trace trace) {
+		TraceScope cur = trace.startSpan(IMPORTANT_WORK_1);
+		try {
+			Thread.sleep((long) (50 * Math.random()));
+			importantWork2(trace);
+		}
+		catch (InterruptedException e) {
+			Thread.currentThread().interrupt();
+		}
+		finally {
+			cur.close();
+		}
+	}
+
+	private void importantWork2(Trace trace) {
+		TraceScope cur = trace.startSpan(IMPORTANT_WORK_2);
+		try {
+			Thread.sleep((long) (50 * Math.random()));
+		}
+		catch (InterruptedException e) {
+			Thread.currentThread().interrupt();
+		}
+		finally {
+			cur.close();
+		}
+	}
+}