Add support for child span count. (#1695)

* Add support for child span count.

* Fix checker and cleanup implementation

* Fix after rebase.

Author: Bogdan Drutu

exporters/trace/stackdriver/src/test/java/io/opencensus/exporter/trace/stackdriver/StackdriverV2ExporterHandlerProtoTest.java

@@ -389,6 +389,7 @@ public void mapHttpAttributes() {
             parentSpanId,
             /* hasRemoteParent= */ true,
             SPAN_NAME,
+            null,
             startTimestamp,
             httpAttributes,
             annotations,
@@ -413,6 +414,32 @@ public void mapHttpAttributes() {
         .containsEntry("/http/status_code", AttributeValue.newBuilder().setIntValue(200L).build());
   }
 
+  @Test
+  public void exportChildSpanCount() {
+    SpanData spanData =
+        SpanData.create(
+            spanContext,
+            parentSpanId,
+            /* hasRemoteParent= */ true,
+            SPAN_NAME,
+            Kind.SERVER,
+            startTimestamp,
+            attributes,
+            annotations,
+            messageEvents,
+            links,
+            CHILD_SPAN_COUNT,
+            status,
+            endTimestamp);
+    assertThat(
+            handler
+                .generateSpan(
+                    spanData, EMPTY_RESOURCE_LABELS, Collections.<String, AttributeValue>emptyMap())
+                .getChildSpanCount()
+                .getValue())
+        .isEqualTo(CHILD_SPAN_COUNT);
+  }
+
   @Test
   public void generateSpanName_ForServer() {
     SpanData spanData =

impl_core/src/main/java/io/opencensus/implcore/trace/RecordEventsSpanImpl.java

@@ -93,6 +93,9 @@ public final class RecordEventsSpanImpl extends Span implements Element<RecordEv
   @GuardedBy("this")
   @Nullable
   private TraceEvents<Link> links;
+  // The number of children.
+  @GuardedBy("this")
+  private int numberOfChildren;
   // The status of the span.
   @GuardedBy("this")
   @Nullable
@@ -263,7 +266,7 @@ public SpanData toSpanData() {
           annotationsSpanData,
           messageEventsSpanData,
           linksSpanData,
-          null, // Not supported yet.
+          numberOfChildren,
           hasBeenEnded ? getStatusWithDefault() : null,
           hasBeenEnded ? timestampConverter.convertNanoTime(endNanoTime) : null);
     }
@@ -381,6 +384,16 @@ public void end(EndSpanOptions options) {
     startEndHandler.onEnd(this);
   }
 
+  void addChild() {
+    synchronized (this) {
+      if (hasBeenEnded) {
+        logger.log(Level.FINE, "Calling end() on an ended Span.");
+        return;
+      }
+      numberOfChildren++;
+    }
+  }
+
   @GuardedBy("this")
   private AttributesWithCapacity getInitializedAttributes() {
     if (attributes == null) {
@@ -568,6 +581,7 @@ private RecordEventsSpanImpl(
     this.clock = clock;
     this.hasBeenEnded = false;
     this.sampleToLocalSpanStore = false;
+    this.numberOfChildren = 0;
     this.timestampConverter =
         timestampConverter != null ? timestampConverter : TimestampConverter.now(clock);
     startNanoTime = clock.nowNanos();

impl_core/src/main/java/io/opencensus/implcore/trace/SpanBuilderImpl.java

@@ -58,35 +58,35 @@ final class SpanBuilderImpl extends SpanBuilder {
   @Nullable private Kind kind;
 
   private Span startSpanInternal(
-      @Nullable SpanContext parent,
+      @Nullable SpanContext parentContext,
       @Nullable Boolean hasRemoteParent,
       String name,
       @Nullable Sampler sampler,
       List<Span> parentLinks,
       @Nullable Boolean recordEvents,
       @Nullable Kind kind,
-      @Nullable TimestampConverter timestampConverter) {
+      @Nullable Span parentSpan) {
     TraceParams activeTraceParams = options.traceConfig.getActiveTraceParams();
     Random random = options.randomHandler.current();
     TraceId traceId;
     SpanId spanId = SpanId.generateRandomId(random);
     SpanId parentSpanId = null;
     // TODO(bdrutu): Handle tracestate correctly not just propagate.
     Tracestate tracestate = TRACESTATE_DEFAULT;
-    if (parent == null || !parent.isValid()) {
+    if (parentContext == null || !parentContext.isValid()) {
       // New root span.
       traceId = TraceId.generateRandomId(random);
       // This is a root span so no remote or local parent.
       hasRemoteParent = null;
     } else {
       // New child span.
-      traceId = parent.getTraceId();
-      parentSpanId = parent.getSpanId();
-      tracestate = parent.getTracestate();
+      traceId = parentContext.getTraceId();
+      parentSpanId = parentContext.getSpanId();
+      tracestate = parentContext.getTracestate();
     }
     TraceOptions traceOptions =
         makeSamplingDecision(
-                parent,
+                parentContext,
                 hasRemoteParent,
                 name,
                 sampler,
@@ -96,22 +96,33 @@ private Span startSpanInternal(
                 activeTraceParams)
             ? SAMPLED_TRACE_OPTIONS
             : NOT_SAMPLED_TRACE_OPTIONS;
-    Span span =
-        (traceOptions.isSampled() || Boolean.TRUE.equals(recordEvents))
-            ? RecordEventsSpanImpl.startSpan(
-                SpanContext.create(traceId, spanId, traceOptions, tracestate),
-                name,
-                kind,
-                parentSpanId,
-                hasRemoteParent,
-                activeTraceParams,
-                options.startEndHandler,
-                timestampConverter,
-                options.clock)
-            : NoRecordEventsSpanImpl.create(
-                SpanContext.create(traceId, spanId, traceOptions, tracestate));
-    linkSpans(span, parentLinks);
-    return span;
+
+    if (traceOptions.isSampled() || Boolean.TRUE.equals(recordEvents)) {
+      // Pass the timestamp converter from the parent to ensure that the recorded events are in
+      // the right order. Implementation uses System.nanoTime() which is monotonically increasing.
+      TimestampConverter timestampConverter = null;
+      if (parentSpan instanceof RecordEventsSpanImpl) {
+        RecordEventsSpanImpl parentRecordEventsSpan = (RecordEventsSpanImpl) parentSpan;
+        timestampConverter = parentRecordEventsSpan.getTimestampConverter();
+        parentRecordEventsSpan.addChild();
+      }
+      Span span =
+          RecordEventsSpanImpl.startSpan(
+              SpanContext.create(traceId, spanId, traceOptions, tracestate),
+              name,
+              kind,
+              parentSpanId,
+              hasRemoteParent,
+              activeTraceParams,
+              options.startEndHandler,
+              timestampConverter,
+              options.clock);
+      linkSpans(span, parentLinks);
+      return span;
+    } else {
+      return NoRecordEventsSpanImpl.create(
+          SpanContext.create(traceId, spanId, traceOptions, tracestate));
+    }
   }
 
   private static boolean makeSamplingDecision(
@@ -179,34 +190,28 @@ private SpanBuilderImpl(
 
   @Override
   public Span startSpan() {
-    SpanContext parentContext = remoteParentSpanContext;
-    Boolean hasRemoteParent = Boolean.TRUE;
-    TimestampConverter timestampConverter = null;
-    if (remoteParentSpanContext == null) {
+    if (remoteParentSpanContext != null) {
+      return startSpanInternal(
+          remoteParentSpanContext,
+          Boolean.TRUE,
+          name,
+          sampler,
+          parentLinks,
+          recordEvents,
+          kind,
+          null);
+    } else {
       // This is not a child of a remote Span. Get the parent SpanContext from the parent Span if
       // any.
-      Span parent = this.parent;
-      hasRemoteParent = Boolean.FALSE;
+      SpanContext parentContext = null;
+      Boolean hasRemoteParent = null;
       if (parent != null) {
         parentContext = parent.getContext();
-        // Pass the timestamp converter from the parent to ensure that the recorded events are in
-        // the right order. Implementation uses System.nanoTime() which is monotonically increasing.
-        if (parent instanceof RecordEventsSpanImpl) {
-          timestampConverter = ((RecordEventsSpanImpl) parent).getTimestampConverter();
-        }
-      } else {
-        hasRemoteParent = null;
+        hasRemoteParent = Boolean.FALSE;
       }
+      return startSpanInternal(
+          parentContext, hasRemoteParent, name, sampler, parentLinks, recordEvents, kind, parent);
     }
-    return startSpanInternal(
-        parentContext,
-        hasRemoteParent,
-        name,
-        sampler,
-        parentLinks,
-        recordEvents,
-        kind,
-        timestampConverter);
   }
 
   static final class Options {

impl_core/src/test/java/io/opencensus/implcore/trace/SpanBuilderImplTest.java

@@ -138,6 +138,33 @@ public void startSpanNullParentWithRecordEvents() {
     assertThat(spanData.getHasRemoteParent()).isNull();
   }
 
+  @Test
+  public void startSpanIncreaseNumberOfChildren() {
+    RecordEventsSpanImpl parent =
+        (RecordEventsSpanImpl)
+            SpanBuilderImpl.createWithParent(SPAN_NAME, null, spanBuilderOptions)
+                .setSampler(Samplers.alwaysSample())
+                .startSpan();
+    assertThat(parent.getContext().getTraceOptions().isSampled()).isTrue();
+    assertThat(parent.toSpanData().getChildSpanCount()).isEqualTo(0);
+    RecordEventsSpanImpl span =
+        (RecordEventsSpanImpl)
+            SpanBuilderImpl.createWithParent(SPAN_NAME, parent, spanBuilderOptions)
+                .setSampler(Samplers.alwaysSample())
+                .startSpan();
+    assertThat(span.getContext().getTraceOptions().isSampled()).isTrue();
+    assertThat(span.toSpanData().getChildSpanCount()).isEqualTo(0);
+    assertThat(parent.toSpanData().getChildSpanCount()).isEqualTo(1);
+    span =
+        (RecordEventsSpanImpl)
+            SpanBuilderImpl.createWithParent(SPAN_NAME, parent, spanBuilderOptions)
+                .setSampler(Samplers.alwaysSample())
+                .startSpan();
+    assertThat(span.getContext().getTraceOptions().isSampled()).isTrue();
+    assertThat(span.toSpanData().getChildSpanCount()).isEqualTo(0);
+    assertThat(parent.toSpanData().getChildSpanCount()).isEqualTo(2);
+  }
+
   @Test
   public void startSpanNullParentNoRecordOptions() {
     Span span =