chore: remove deprecated registration of the span names for zpages (#27)

* replace deprecated registration of the span names for zpages

* remove unused END_SPAN_OPTIONS

* fix test

Author: mayurkale22

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java

@@ -457,7 +457,7 @@ public final void invalidate() {
 
   @Override
   public void close() {
-    span.end();
+    span.end(TraceUtil.END_SPAN_OPTIONS);
     synchronized (lock) {
       isClosed = true;
     }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java

@@ -847,7 +847,7 @@ protected ResumableStreamIterator(int maxBufferSize, String streamName, Span par
     public void close(@Nullable String message) {
       if (stream != null) {
         stream.close(message);
-        span.end();
+        span.end(TraceUtil.END_SPAN_OPTIONS);
       }
     }
 

google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java

@@ -37,10 +37,6 @@ private enum SessionMode {
     READ_WRITE
   }
 
-  static {
-    TraceUtil.exportSpans(READ_WRITE_TRANSACTION, READ_ONLY_TRANSACTION, PARTITION_DML_TRANSACTION);
-  }
-
   @VisibleForTesting final SessionPool pool;
 
   DatabaseClientImpl(SessionPool pool) {
@@ -73,7 +69,7 @@ public Timestamp apply(Session session) {
       TraceUtil.endSpanWithFailure(span, e);
       throw e;
     } finally {
-      span.end();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
     }
   }
 
@@ -93,7 +89,7 @@ public Timestamp apply(Session session) {
       TraceUtil.endSpanWithFailure(span, e);
       throw e;
     } finally {
-      span.end();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
     }
   }
 

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java

@@ -126,8 +126,8 @@ private BatchCreateSessionsRunnable(
     public void run() {
       List<SessionImpl> sessions = null;
       int remainingSessionsToCreate = sessionCount;
-      try (Scope scope =
-          SpannerImpl.tracer.spanBuilder(SpannerImpl.BATCH_CREATE_SESSIONS).startScopedSpan()) {
+      Span span = SpannerImpl.tracer.spanBuilder(SpannerImpl.BATCH_CREATE_SESSIONS).startSpan();
+      try (Scope s = SpannerImpl.tracer.withSpan(span)) {
         SpannerImpl.tracer
             .getCurrentSpan()
             .addAnnotation(String.format("Creating %d sessions", sessionCount));
@@ -144,6 +144,8 @@ public void run() {
           }
           remainingSessionsToCreate -= sessions.size();
         }
+      } finally {
+        span.end(TraceUtil.END_SPAN_OPTIONS);
       }
     }
   }
@@ -205,7 +207,7 @@ SessionImpl createSession() {
           spanner
               .getRpc()
               .createSession(db.getName(), spanner.getOptions().getSessionLabels(), options);
-      span.end();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
       return new SessionImpl(spanner, session.getName(), options);
     } catch (RuntimeException e) {
       TraceUtil.endSpanWithFailure(span, e);
@@ -284,7 +286,7 @@ private List<SessionImpl> internalBatchCreateSessions(
       span.addAnnotation(
           String.format(
               "Request for %d sessions returned %d sessions", sessionCount, sessions.size()));
-      span.end();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
       List<SessionImpl> res = new ArrayList<>(sessionCount);
       for (com.google.spanner.v1.Session session : sessions) {
         res.add(new SessionImpl(spanner, session.getName(), options));

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java

@@ -138,7 +138,7 @@ public Timestamp writeAtLeastOnce(Iterable<Mutation> mutations) throws SpannerEx
     try (Scope s = tracer.withSpan(span)) {
       CommitResponse response = spanner.getRpc().commit(request, options);
       Timestamp t = Timestamp.fromProto(response.getCommitTimestamp());
-      span.end();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
       return t;
     } catch (IllegalArgumentException e) {
       TraceUtil.endSpanWithFailure(span, e);
@@ -201,7 +201,7 @@ public void close() {
     Span span = tracer.spanBuilder(SpannerImpl.DELETE_SESSION).startSpan();
     try (Scope s = tracer.withSpan(span)) {
       spanner.getRpc().deleteSession(name, options);
-      span.end();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
     } catch (RuntimeException e) {
       TraceUtil.endSpanWithFailure(span, e);
       throw e;
@@ -222,7 +222,7 @@ ByteString beginTransaction() {
       if (txn.getId().isEmpty()) {
         throw newSpannerException(ErrorCode.INTERNAL, "Missing id in transaction\n" + getName());
       }
-      span.end();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
       return txn.getId();
     } catch (RuntimeException e) {
       TraceUtil.endSpanWithFailure(span, e);

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java

@@ -72,10 +72,6 @@ final class SessionPool {
   private static final Tracer tracer = Tracing.getTracer();
   static final String WAIT_FOR_SESSION = "SessionPool.WaitForSession";
 
-  static {
-    TraceUtil.exportSpans(WAIT_FOR_SESSION);
-  }
-
   /**
    * Wrapper around current time so that we can fake it in tests. TODO(user): Replace with Java 8
    * Clock.
@@ -854,7 +850,8 @@ private void put(SpannerException e) {
     private PooledSession take() throws SpannerException {
       long currentTimeout = options.getInitialWaitForSessionTimeoutMillis();
       while (true) {
-        try (Scope waitScope = tracer.spanBuilder(WAIT_FOR_SESSION).startScopedSpan()) {
+        Span span = tracer.spanBuilder(WAIT_FOR_SESSION).startSpan();
+        try (Scope waitScope = tracer.withSpan(span)) {
           SessionOrError s = pollUninterruptiblyWithTimeout(currentTimeout);
           if (s == null) {
             // Set the status to DEADLINE_EXCEEDED and retry.
@@ -870,6 +867,8 @@ private PooledSession take() throws SpannerException {
         } catch (Exception e) {
           TraceUtil.endSpanWithFailure(tracer.getCurrentSpan(), e);
           throw e;
+        } finally {
+          span.end(TraceUtil.END_SPAN_OPTIONS);
         }
       }
     }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java

@@ -71,18 +71,6 @@ class SpannerImpl extends BaseService<SpannerOptions> implements Spanner {
   static final String QUERY = "CloudSpannerOperation.ExecuteStreamingQuery";
   static final String READ = "CloudSpannerOperation.ExecuteStreamingRead";
 
-  static {
-    TraceUtil.exportSpans(
-        BATCH_CREATE_SESSIONS,
-        BATCH_CREATE_SESSIONS_REQUEST,
-        CREATE_SESSION,
-        DELETE_SESSION,
-        BEGIN_TRANSACTION,
-        COMMIT,
-        QUERY,
-        READ);
-  }
-
   private final SpannerRpc gapicRpc;
 
   @GuardedBy("this")

google-cloud-spanner/src/main/java/com/google/cloud/spanner/TraceUtil.java

@@ -21,16 +21,17 @@
 import com.google.spanner.v1.Transaction;
 import io.opencensus.contrib.grpc.util.StatusConverter;
 import io.opencensus.trace.AttributeValue;
+import io.opencensus.trace.EndSpanOptions;
 import io.opencensus.trace.Span;
 import io.opencensus.trace.Status;
-import io.opencensus.trace.Tracing;
-import io.opencensus.trace.export.SampledSpanStore;
-import java.util.Arrays;
 import java.util.Map;
 
 /** Utility methods for tracing. */
 class TraceUtil {
 
+  static final EndSpanOptions END_SPAN_OPTIONS =
+      EndSpanOptions.builder().setSampleToLocalSpanStore(true).build();
+
   static Map<String, AttributeValue> getTransactionAnnotations(Transaction t) {
     return ImmutableMap.of(
         "Id",
@@ -58,21 +59,14 @@ static void endSpanWithFailure(Span span, Throwable e) {
       endSpanWithFailure(span, (SpannerException) e);
     } else {
       span.setStatus(Status.INTERNAL.withDescription(e.getMessage()));
-      span.end();
+      span.end(END_SPAN_OPTIONS);
     }
   }
 
   static void endSpanWithFailure(Span span, SpannerException e) {
     span.setStatus(
         StatusConverter.fromGrpcStatus(e.getErrorCode().getGrpcStatus())
             .withDescription(e.getMessage()));
-    span.end();
-  }
-
-  static void exportSpans(String... spans) {
-    SampledSpanStore store = Tracing.getExportComponent().getSampledSpanStore();
-    if (store != null) {
-      store.registerSpanNamesForCollection(Arrays.asList(spans));
-    }
+    span.end(END_SPAN_OPTIONS);
   }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java

@@ -115,7 +115,7 @@ public void close() {
         txnState = TransactionState.ROLLED_BACK;
       }
     } finally {
-      span.end();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
     }
   }
 

google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java

@@ -130,7 +130,7 @@ void commit() {
               ErrorCode.INTERNAL, "Missing commitTimestamp:\n" + session.getName());
         }
         commitTimestamp = Timestamp.fromProto(commitResponse.getCommitTimestamp());
-        opSpan.end();
+        opSpan.end(TraceUtil.END_SPAN_OPTIONS);
       } catch (RuntimeException e) {
         span.addAnnotation("Commit Failed", TraceUtil.getExceptionAnnotations(e));
         TraceUtil.endSpanWithFailure(opSpan, e);