From 798b7745d3616194e5c6b85aac90f17536fc7c70 Mon Sep 17 00:00:00 2001
From: Bogdan Drutu <bogdandrutu@gmail.com>
Date: Sat, 15 Sep 2018 09:26:21 -0700
Subject: [PATCH 1/2] Fix parent/child spans relationship in Spanner.

---
 .../com/google/cloud/spanner/SpannerImpl.java    | 16 ++++++++++------
 .../spanner/ResumableStreamIteratorTest.java     |  2 +-
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/google-cloud-clients/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-clients/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java
index 579da6797ced..1173f9fd0ecd 100644
--- a/google-cloud-clients/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java
+++ b/google-cloud-clients/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java
@@ -1063,7 +1063,7 @@ ResultSet executeQueryInternalWithOptions(
       final int prefetchChunks =
           readOptions.hasPrefetchChunks() ? readOptions.prefetchChunks() : defaultPrefetchChunks;
       ResumableStreamIterator stream =
-          new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, QUERY) {
+          new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, QUERY, span) {
             @Override
             CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
               GrpcStreamIterator stream = new GrpcStreamIterator(prefetchChunks);
@@ -1176,7 +1176,7 @@ ResultSet readInternalWithOptions(
       final int prefetchChunks =
           readOptions.hasPrefetchChunks() ? readOptions.prefetchChunks() : defaultPrefetchChunks;
       ResumableStreamIterator stream =
-          new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, READ) {
+          new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, READ, span) {
             @Override
             CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
               GrpcStreamIterator stream = new GrpcStreamIterator(prefetchChunks);
@@ -2452,20 +2452,20 @@ abstract static class ResumableStreamIterator extends AbstractIterator<PartialRe
      */
     private boolean safeToRetry = true;
 
-    protected ResumableStreamIterator(int maxBufferSize, String streamName) {
+    protected ResumableStreamIterator(int maxBufferSize, String streamName, Span parent) {
       checkArgument(maxBufferSize >= 0);
       this.maxBufferSize = maxBufferSize;
-      this.span = tracer.spanBuilder(streamName).startSpan();
+      this.span = tracer.spanBuilderWithExplicitParent(streamName, parent).startSpan();
     }
 
     abstract CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken);
 
     @Override
     public void close(@Nullable String message) {
-      span.end();
       if (stream != null) {
         stream.close(message);
       }
+      span.end();
     }
 
     @Override
@@ -2478,7 +2478,11 @@ protected PartialResultSet computeNext() {
               ImmutableMap.of("ResumeToken",
                   AttributeValue.stringAttributeValue(
                       resumeToken == null ? "null" : resumeToken.toStringUtf8())));
-          stream = checkNotNull(startStream(resumeToken));
+          try (Scope s = tracer.withSpan(span)) {
+            // When start a new stream set the Span as current to make the gRPC Span a child of
+            // this Span.
+            stream = checkNotNull(startStream(resumeToken));
+          }
         }
         // Buffer contains items up to a resume token or has reached capacity: flush.
         if (!buffer.isEmpty()
diff --git a/google-cloud-clients/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java b/google-cloud-clients/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java
index 79ecf761650f..112ed1103cf9 100644
--- a/google-cloud-clients/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java
+++ b/google-cloud-clients/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java
@@ -98,7 +98,7 @@ public void setUp() {
 
   private void initWithLimit(int maxBufferSize) {
     iterator =
-        new SpannerImpl.ResumableStreamIterator(maxBufferSize, "") {
+        new SpannerImpl.ResumableStreamIterator(maxBufferSize, "", null) {
           @Override
           SpannerImpl.CloseableIterator<PartialResultSet> startStream(
               @Nullable ByteString resumeToken) {

From 0bd72c6811aa8fb0c23b544ef9ff4a9cc776825f Mon Sep 17 00:00:00 2001
From: Bogdan Drutu <bogdandrutu@gmail.com>
Date: Sat, 15 Sep 2018 11:26:56 -0700
Subject: [PATCH 2/2] Make the commit Span a child of the ReadContext Span.

---
 .../src/main/java/com/google/cloud/spanner/SpannerImpl.java     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/google-cloud-clients/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-clients/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java
index 1173f9fd0ecd..b997103dbaf2 100644
--- a/google-cloud-clients/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java
+++ b/google-cloud-clients/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java
@@ -1426,7 +1426,7 @@ void commit() {
         mutations = null;
       }
       final CommitRequest commitRequest = builder.build();
-      Span opSpan = tracer.spanBuilder(COMMIT).startSpan();
+      Span opSpan = tracer.spanBuilderWithExplicitParent(COMMIT, span).startSpan();
       try (Scope s = tracer.withSpan(opSpan)) {
         CommitResponse commitResponse =
             runWithRetries(