bigtable, spanner: consolidate dupe OpenCensus tracing code

Packages bigtable and spanner had identical code in their
trace.go files for starting and ending spans. On ending
spans, if we encountered an error usually from a remote RPC,
each traceEndSpan variant function would just set the span's
Status.Message as the error but unfortunately would keep the
Status.Code as 0 which defaults to Status.OK which is not correct,
and can be misdiagnosed during trace examination.

This CL delegates this functionality to internal/trace which already has
the correct code. That package has comprehensive code that correctly
extracts statuses and messages and converts to them be settable on
the created spans.

Added a TODO to later on perhaps keep the result of:
   ctx, span := trace.StartSpan(ctx, spanName)
instead of passing around only the context and having
to incur extra costs from trace.FromContext(ctx) on every call.
However, we'll have to evaluate these costs if they are actually
expensive and a priority.

Also noticed that perhaps in a later CL we should rename internal/trace
to internal/observability, because some of these packages uses stats
as well.

Fixes #1376

Change-Id: I22da12ad23d4885322636c18e3dd5f2897a9f1a5
Reviewed-on: https://code-review.googlesource.com/c/gocloud/+/39410
Reviewed-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Jean de Klerk <deklerk@google.com>
diff --git a/bigtable/bigtable.go b/bigtable/bigtable.go
index 28c0c92..3b06513 100644
--- a/bigtable/bigtable.go
+++ b/bigtable/bigtable.go
@@ -26,6 +26,7 @@
 
 	"cloud.google.com/go/bigtable/internal/gax"
 	btopt "cloud.google.com/go/bigtable/internal/option"
+	"cloud.google.com/go/internal/trace"
 	"github.com/golang/protobuf/proto"
 	"google.golang.org/api/option"
 	gtransport "google.golang.org/api/transport/grpc"
@@ -147,8 +148,8 @@
 
 	var prevRowKey string
 	var err error
-	ctx = traceStartSpan(ctx, "cloud.google.com/go/bigtable.ReadRows")
-	defer func() { traceEndSpan(ctx, err) }()
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigtable.ReadRows")
+	defer func() { trace.EndSpan(ctx, err) }()
 	attrMap := make(map[string]interface{})
 	err = gax.Invoke(ctx, func(ctx context.Context) error {
 		if !arg.valid() {
@@ -185,12 +186,12 @@
 				attrMap["rowKey"] = prevRowKey
 				attrMap["error"] = err.Error()
 				attrMap["time_secs"] = time.Since(startTime).Seconds()
-				tracePrintf(ctx, attrMap, "Retry details in ReadRows")
+				trace.TracePrintf(ctx, attrMap, "Retry details in ReadRows")
 				return err
 			}
 			attrMap["time_secs"] = time.Since(startTime).Seconds()
 			attrMap["rowCount"] = len(res.Chunks)
-			tracePrintf(ctx, attrMap, "Details in ReadRows")
+			trace.TracePrintf(ctx, attrMap, "Details in ReadRows")
 
 			for _, cc := range res.Chunks {
 				row, err := cr.Process(cc)
@@ -477,8 +478,8 @@
 	}
 
 	var err error
-	ctx = traceStartSpan(ctx, "cloud.google.com/go/bigtable/Apply")
-	defer func() { traceEndSpan(ctx, err) }()
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigtable/Apply")
+	defer func() { trace.EndSpan(ctx, err) }()
 	var callOptions []gax.CallOption
 	if m.cond == nil {
 		req := &btpb.MutateRowRequest{
@@ -657,14 +658,14 @@
 	}
 
 	var err error
-	ctx = traceStartSpan(ctx, "cloud.google.com/go/bigtable/ApplyBulk")
-	defer func() { traceEndSpan(ctx, err) }()
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigtable/ApplyBulk")
+	defer func() { trace.EndSpan(ctx, err) }()
 
 	for _, group := range groupEntries(origEntries, maxMutations) {
 		attrMap := make(map[string]interface{})
 		err = gax.Invoke(ctx, func(ctx context.Context) error {
 			attrMap["rowCount"] = len(group)
-			tracePrintf(ctx, attrMap, "Row count in ApplyBulk")
+			trace.TracePrintf(ctx, attrMap, "Row count in ApplyBulk")
 			err := t.doApplyBulk(ctx, group, opts...)
 			if err != nil {
 				// We want to retry the entire request with the current group
diff --git a/bigtable/trace.go b/bigtable/trace.go
deleted file mode 100644
index 6c579b1..0000000
--- a/bigtable/trace.go
+++ /dev/null
@@ -1,57 +0,0 @@
-// Copyright 2018 Google LLC
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package bigtable
-
-import (
-	"context"
-	"fmt"
-
-	"go.opencensus.io/trace"
-)
-
-func traceStartSpan(ctx context.Context, name string) context.Context {
-	ctx, _ = trace.StartSpan(ctx, name)
-	return ctx
-}
-
-func traceEndSpan(ctx context.Context, err error) {
-	span := trace.FromContext(ctx)
-	if err != nil {
-		span.SetStatus(trace.Status{Message: err.Error()})
-	}
-
-	span.End()
-}
-
-func tracePrintf(ctx context.Context, attrMap map[string]interface{}, format string, args ...interface{}) {
-	var attrs []trace.Attribute
-	for k, v := range attrMap {
-		var a trace.Attribute
-		switch v := v.(type) {
-		case string:
-			a = trace.StringAttribute(k, v)
-		case bool:
-			a = trace.BoolAttribute(k, v)
-		case int:
-			a = trace.Int64Attribute(k, int64(v))
-		case int64:
-			a = trace.Int64Attribute(k, v)
-		default:
-			a = trace.StringAttribute(k, fmt.Sprintf("%#v", v))
-		}
-		attrs = append(attrs, a)
-	}
-	trace.FromContext(ctx).Annotatef(attrs, format, args...)
-}
diff --git a/internal/trace/trace.go b/internal/trace/trace.go
index 95c7821..e2156f7 100644
--- a/internal/trace/trace.go
+++ b/internal/trace/trace.go
@@ -16,6 +16,7 @@
 
 import (
 	"context"
+	"fmt"
 
 	"go.opencensus.io/trace"
 	"google.golang.org/api/googleapi"
@@ -38,7 +39,7 @@
 	span.End()
 }
 
-// ToStatus interrogates an error and converts it to an appropriate
+// toStatus interrogates an error and converts it to an appropriate
 // OpenCensus status.
 func toStatus(err error) trace.Status {
 	if err2, ok := err.(*googleapi.Error); ok {
@@ -82,3 +83,27 @@
 		return int32(code.Code_UNKNOWN)
 	}
 }
+
+// TODO: (odeke-em): perhaps just pass around spans due to the cost
+// incurred from using trace.FromContext(ctx) yet we could avoid
+// throwing away the work done by ctx, span := trace.StartSpan.
+func TracePrintf(ctx context.Context, attrMap map[string]interface{}, format string, args ...interface{}) {
+	var attrs []trace.Attribute
+	for k, v := range attrMap {
+		var a trace.Attribute
+		switch v := v.(type) {
+		case string:
+			a = trace.StringAttribute(k, v)
+		case bool:
+			a = trace.BoolAttribute(k, v)
+		case int:
+			a = trace.Int64Attribute(k, int64(v))
+		case int64:
+			a = trace.Int64Attribute(k, v)
+		default:
+			a = trace.StringAttribute(k, fmt.Sprintf("%#v", v))
+		}
+		attrs = append(attrs, a)
+	}
+	trace.FromContext(ctx).Annotatef(attrs, format, args...)
+}
diff --git a/spanner/client.go b/spanner/client.go
index 53913be..1ecc0c6 100644
--- a/spanner/client.go
+++ b/spanner/client.go
@@ -23,6 +23,7 @@
 	"sync/atomic"
 	"time"
 
+	"cloud.google.com/go/internal/trace"
 	"cloud.google.com/go/internal/version"
 	"google.golang.org/api/option"
 	gtransport "google.golang.org/api/transport/grpc"
@@ -119,8 +120,8 @@
 // NewClientWithConfig creates a client to a database. A valid database name has the
 // form projects/PROJECT_ID/instances/INSTANCE_ID/databases/DATABASE_ID.
 func NewClientWithConfig(ctx context.Context, database string, config ClientConfig, opts ...option.ClientOption) (c *Client, err error) {
-	ctx = startSpan(ctx, "cloud.google.com/go/spanner.NewClient")
-	defer func() { endSpan(ctx, err) }()
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/spanner.NewClient")
+	defer func() { trace.EndSpan(ctx, err) }()
 
 	// Validate database path.
 	if err := validDatabaseName(database); err != nil {
@@ -354,8 +355,8 @@
 // See https://godoc.org/cloud.google.com/go/spanner#ReadWriteTransaction for
 // more details.
 func (c *Client) ReadWriteTransaction(ctx context.Context, f func(context.Context, *ReadWriteTransaction) error) (commitTimestamp time.Time, err error) {
-	ctx = startSpan(ctx, "cloud.google.com/go/spanner.ReadWriteTransaction")
-	defer func() { endSpan(ctx, err) }()
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/spanner.ReadWriteTransaction")
+	defer func() { trace.EndSpan(ctx, err) }()
 	if err := checkNestedTxn(ctx); err != nil {
 		return time.Time{}, err
 	}
@@ -385,7 +386,7 @@
 			}
 		}
 		t.txReadOnly.txReadEnv = t
-		statsPrintf(ctx, map[string]interface{}{"transactionID": string(sh.getTransactionID())},
+		trace.TracePrintf(ctx, map[string]interface{}{"transactionID": string(sh.getTransactionID())},
 			"Starting transaction attempt")
 		if err = t.begin(ctx); err != nil {
 			// Mask error from begin operation as retryable error.
@@ -438,8 +439,8 @@
 		})
 	}
 
-	ctx = startSpan(ctx, "cloud.google.com/go/spanner.Apply")
-	defer func() { endSpan(ctx, err) }()
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/spanner.Apply")
+	defer func() { trace.EndSpan(ctx, err) }()
 	t := &writeOnlyTransaction{c.idleSessions}
 	return t.applyAtLeastOnce(ctx, ms...)
 }
diff --git a/spanner/pdml.go b/spanner/pdml.go
index 4ccf791..090c12f 100644
--- a/spanner/pdml.go
+++ b/spanner/pdml.go
@@ -18,6 +18,7 @@
 	"context"
 	"time"
 
+	"cloud.google.com/go/internal/trace"
 	"google.golang.org/api/iterator"
 	sppb "google.golang.org/genproto/googleapis/spanner/v1"
 	"google.golang.org/grpc/codes"
@@ -32,8 +33,8 @@
 // PartitionedUpdate returns an estimated count of the number of rows affected. The actual
 // number of affected rows may be greater than the estimate.
 func (c *Client) PartitionedUpdate(ctx context.Context, statement Statement) (count int64, err error) {
-	ctx = startSpan(ctx, "cloud.google.com/go/spanner.PartitionedUpdate")
-	defer func() { endSpan(ctx, err) }()
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/spanner.PartitionedUpdate")
+	defer func() { trace.EndSpan(ctx, err) }()
 	if err := checkNestedTxn(ctx); err != nil {
 		return 0, err
 	}
diff --git a/spanner/read.go b/spanner/read.go
index 6faccdb..6c2a301 100644
--- a/spanner/read.go
+++ b/spanner/read.go
@@ -25,6 +25,7 @@
 	"time"
 
 	"cloud.google.com/go/internal/protostruct"
+	"cloud.google.com/go/internal/trace"
 	"cloud.google.com/go/spanner/internal/backoff"
 	proto "github.com/golang/protobuf/proto"
 	proto3 "github.com/golang/protobuf/ptypes/struct"
@@ -48,7 +49,7 @@
 // Cloud Spanner.
 func stream(ctx context.Context, rpc func(ct context.Context, resumeToken []byte) (streamingReceiver, error), setTimestamp func(time.Time), release func(error)) *RowIterator {
 	ctx, cancel := context.WithCancel(ctx)
-	ctx = startSpan(ctx, "cloud.google.com/go/spanner.RowIterator")
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/spanner.RowIterator")
 	return &RowIterator{
 		streamd:      newResumableStreamDecoder(ctx, rpc),
 		rowd:         &partialResultSetDecoder{},
@@ -168,7 +169,7 @@
 // Stop terminates the iteration. It should be called after you finish using the iterator.
 func (r *RowIterator) Stop() {
 	if r.streamd != nil {
-		defer endSpan(r.streamd.ctx, r.err)
+		defer trace.EndSpan(r.streamd.ctx, r.err)
 	}
 	if r.cancel != nil {
 		r.cancel()
@@ -532,7 +533,7 @@
 // doBackoff does an exponential backoff sleep.
 func (d *resumableStreamDecoder) doBackOff() {
 	delay := d.backoff.Delay(d.retryCount)
-	statsPrintf(d.ctx, nil, "Backing off stream read for %s", delay)
+	trace.TracePrintf(d.ctx, nil, "Backing off stream read for %s", delay)
 	ticker := time.NewTicker(delay)
 	defer ticker.Stop()
 	d.retryCount++
diff --git a/spanner/retry.go b/spanner/retry.go
index a172493..b2937d0 100644
--- a/spanner/retry.go
+++ b/spanner/retry.go
@@ -22,6 +22,7 @@
 	"strings"
 	"time"
 
+	"cloud.google.com/go/internal/trace"
 	"cloud.google.com/go/spanner/internal/backoff"
 	"github.com/golang/protobuf/proto"
 	"github.com/golang/protobuf/ptypes"
@@ -185,7 +186,7 @@
 			if !ok {
 				b = backoff.DefaultBackoff.Delay(retryCount)
 			}
-			statsPrintf(ctx, nil, "Backing off for %s, then retrying", b)
+			trace.TracePrintf(ctx, nil, "Backing off for %s, then retrying", b)
 			select {
 			case <-ctx.Done():
 				return errContextCanceled(ctx, funcErr)
diff --git a/spanner/session.go b/spanner/session.go
index d41b862..bf2b7ad 100644
--- a/spanner/session.go
+++ b/spanner/session.go
@@ -27,6 +27,7 @@
 	"sync"
 	"time"
 
+	"cloud.google.com/go/internal/trace"
 	sppb "google.golang.org/genproto/googleapis/spanner/v1"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -450,7 +451,7 @@
 }
 
 func (p *sessionPool) createSession(ctx context.Context) (*session, error) {
-	statsPrintf(ctx, nil, "Creating a new session")
+	trace.TracePrintf(ctx, nil, "Creating a new session")
 	doneCreate := func(done bool) {
 		p.mu.Lock()
 		if !done {
@@ -517,7 +518,7 @@
 // take returns a cached session if there are available ones; if there isn't any, it tries to allocate a new one.
 // Session returned by take should be used for read operations.
 func (p *sessionPool) take(ctx context.Context) (*sessionHandle, error) {
-	statsPrintf(ctx, nil, "Acquiring a read-only session")
+	trace.TracePrintf(ctx, nil, "Acquiring a read-only session")
 	ctx = contextWithOutgoingMetadata(ctx, p.md)
 	for {
 		var (
@@ -533,11 +534,11 @@
 		if p.idleList.Len() > 0 {
 			// Idle sessions are available, get one from the top of the idle list.
 			s = p.idleList.Remove(p.idleList.Front()).(*session)
-			statsPrintf(ctx, map[string]interface{}{"sessionID": s.getID()},
+			trace.TracePrintf(ctx, map[string]interface{}{"sessionID": s.getID()},
 				"Acquired read-only session")
 		} else if p.idleWriteList.Len() > 0 {
 			s = p.idleWriteList.Remove(p.idleWriteList.Front()).(*session)
-			statsPrintf(ctx, map[string]interface{}{"sessionID": s.getID()},
+			trace.TracePrintf(ctx, map[string]interface{}{"sessionID": s.getID()},
 				"Acquired read-write session")
 		}
 		if s != nil {
@@ -555,10 +556,10 @@
 		if (p.MaxOpened > 0 && p.numOpened >= p.MaxOpened) || (p.MaxBurst > 0 && p.createReqs >= p.MaxBurst) {
 			mayGetSession := p.mayGetSession
 			p.mu.Unlock()
-			statsPrintf(ctx, nil, "Waiting for read-only session to become available")
+			trace.TracePrintf(ctx, nil, "Waiting for read-only session to become available")
 			select {
 			case <-ctx.Done():
-				statsPrintf(ctx, nil, "Context done waiting for session")
+				trace.TracePrintf(ctx, nil, "Context done waiting for session")
 				return nil, errGetSessionTimeout()
 			case <-mayGetSession:
 			}
@@ -570,10 +571,10 @@
 		p.createReqs++
 		p.mu.Unlock()
 		if s, err = p.createSession(ctx); err != nil {
-			statsPrintf(ctx, nil, "Error creating session: %v", err)
+			trace.TracePrintf(ctx, nil, "Error creating session: %v", err)
 			return nil, toSpannerError(err)
 		}
-		statsPrintf(ctx, map[string]interface{}{"sessionID": s.getID()},
+		trace.TracePrintf(ctx, map[string]interface{}{"sessionID": s.getID()},
 			"Created session")
 		return &sessionHandle{session: s}, nil
 	}
@@ -582,7 +583,7 @@
 // takeWriteSession returns a write prepared cached session if there are available ones; if there isn't any, it tries to allocate a new one.
 // Session returned should be used for read write transactions.
 func (p *sessionPool) takeWriteSession(ctx context.Context) (*sessionHandle, error) {
-	statsPrintf(ctx, nil, "Acquiring a read-write session")
+	trace.TracePrintf(ctx, nil, "Acquiring a read-write session")
 	ctx = contextWithOutgoingMetadata(ctx, p.md)
 	for {
 		var (
@@ -598,10 +599,10 @@
 		if p.idleWriteList.Len() > 0 {
 			// Idle sessions are available, get one from the top of the idle list.
 			s = p.idleWriteList.Remove(p.idleWriteList.Front()).(*session)
-			statsPrintf(ctx, map[string]interface{}{"sessionID": s.getID()}, "Acquired read-write session")
+			trace.TracePrintf(ctx, map[string]interface{}{"sessionID": s.getID()}, "Acquired read-write session")
 		} else if p.idleList.Len() > 0 {
 			s = p.idleList.Remove(p.idleList.Front()).(*session)
-			statsPrintf(ctx, map[string]interface{}{"sessionID": s.getID()}, "Acquired read-only session")
+			trace.TracePrintf(ctx, map[string]interface{}{"sessionID": s.getID()}, "Acquired read-only session")
 		}
 		if s != nil {
 			s.setIdleList(nil)
@@ -617,10 +618,10 @@
 			if (p.MaxOpened > 0 && p.numOpened >= p.MaxOpened) || (p.MaxBurst > 0 && p.createReqs >= p.MaxBurst) {
 				mayGetSession := p.mayGetSession
 				p.mu.Unlock()
-				statsPrintf(ctx, nil, "Waiting for read-write session to become available")
+				trace.TracePrintf(ctx, nil, "Waiting for read-write session to become available")
 				select {
 				case <-ctx.Done():
-					statsPrintf(ctx, nil, "Context done waiting for session")
+					trace.TracePrintf(ctx, nil, "Context done waiting for session")
 					return nil, errGetSessionTimeout()
 				case <-mayGetSession:
 				}
@@ -633,16 +634,16 @@
 			p.createReqs++
 			p.mu.Unlock()
 			if s, err = p.createSession(ctx); err != nil {
-				statsPrintf(ctx, nil, "Error creating session: %v", err)
+				trace.TracePrintf(ctx, nil, "Error creating session: %v", err)
 				return nil, toSpannerError(err)
 			}
-			statsPrintf(ctx, map[string]interface{}{"sessionID": s.getID()},
+			trace.TracePrintf(ctx, map[string]interface{}{"sessionID": s.getID()},
 				"Created session")
 		}
 		if !s.isWritePrepared() {
 			if err = s.prepareForWrite(ctx); err != nil {
 				s.recycle()
-				statsPrintf(ctx, map[string]interface{}{"sessionID": s.getID()},
+				trace.TracePrintf(ctx, map[string]interface{}{"sessionID": s.getID()},
 					"Error preparing session for write")
 				return nil, toSpannerError(err)
 			}
diff --git a/spanner/trace.go b/spanner/stats.go
similarity index 61%
rename from spanner/trace.go
rename to spanner/stats.go
index 663ef19..6151b0c 100644
--- a/spanner/trace.go
+++ b/spanner/stats.go
@@ -16,48 +16,11 @@
 
 import (
 	"context"
-	"fmt"
 
 	"go.opencensus.io/stats"
 	"go.opencensus.io/stats/view"
-	"go.opencensus.io/trace"
 )
 
-func startSpan(ctx context.Context, name string) context.Context {
-	ctx, _ = trace.StartSpan(ctx, name)
-	return ctx
-}
-
-func endSpan(ctx context.Context, err error) {
-	span := trace.FromContext(ctx)
-	if err != nil {
-		// TODO(jba): Add error code to the status.
-		span.SetStatus(trace.Status{Message: err.Error()})
-	}
-	span.End()
-}
-
-func statsPrintf(ctx context.Context, attrMap map[string]interface{}, format string, args ...interface{}) {
-	var attrs []trace.Attribute
-	for k, v := range attrMap {
-		var a trace.Attribute
-		switch v := v.(type) {
-		case string:
-			a = trace.StringAttribute(k, v)
-		case bool:
-			a = trace.BoolAttribute(k, v)
-		case int:
-			a = trace.Int64Attribute(k, int64(v))
-		case int64:
-			a = trace.Int64Attribute(k, v)
-		default:
-			a = trace.StringAttribute(k, fmt.Sprintf("%#v", v))
-		}
-		attrs = append(attrs, a)
-	}
-	trace.FromContext(ctx).Annotatef(attrs, format, args...)
-}
-
 const statsPrefix = "cloud.google.com/go/spanner/"
 
 func recordStat(ctx context.Context, m *stats.Int64Measure, n int64) {
diff --git a/spanner/transaction.go b/spanner/transaction.go
index be75c30..b8bec16 100644
--- a/spanner/transaction.go
+++ b/spanner/transaction.go
@@ -22,6 +22,7 @@
 	"sync/atomic"
 	"time"
 
+	"cloud.google.com/go/internal/trace"
 	"google.golang.org/api/iterator"
 	sppb "google.golang.org/genproto/googleapis/spanner/v1"
 	"google.golang.org/grpc"
@@ -80,8 +81,8 @@
 // ReadWithOptions returns a RowIterator for reading multiple rows from the database.
 // Pass a ReadOptions to modify the read operation.
 func (t *txReadOnly) ReadWithOptions(ctx context.Context, table string, keys KeySet, columns []string, opts *ReadOptions) (ri *RowIterator) {
-	ctx = startSpan(ctx, "cloud.google.com/go/spanner.Read")
-	defer func() { endSpan(ctx, ri.err) }()
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/spanner.Read")
+	defer func() { trace.EndSpan(ctx, ri.err) }()
 	var (
 		sh  *sessionHandle
 		ts  *sppb.TransactionSelector
@@ -188,8 +189,8 @@
 }
 
 func (t *txReadOnly) query(ctx context.Context, statement Statement, mode sppb.ExecuteSqlRequest_QueryMode) (ri *RowIterator) {
-	ctx = startSpan(ctx, "cloud.google.com/go/spanner.Query")
-	defer func() { endSpan(ctx, ri.err) }()
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/spanner.Query")
+	defer func() { trace.EndSpan(ctx, ri.err) }()
 	req, sh, err := t.prepareExecuteSQL(ctx, statement, mode)
 	if err != nil {
 		return &RowIterator{err: err}
@@ -659,8 +660,8 @@
 // Update returns an error if the statement is a query. However, the
 // query is executed, and any data read will be validated upon commit.
 func (t *ReadWriteTransaction) Update(ctx context.Context, stmt Statement) (rowCount int64, err error) {
-	ctx = startSpan(ctx, "cloud.google.com/go/spanner.Update")
-	defer func() { endSpan(ctx, err) }()
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/spanner.Update")
+	defer func() { trace.EndSpan(ctx, err) }()
 	req, sh, err := t.prepareExecuteSQL(ctx, stmt, sppb.ExecuteSqlRequest_NORMAL)
 	if err != nil {
 		return 0, err
@@ -682,8 +683,8 @@
 // affected rows for the given query at the same index. If an error occurs,
 // counts will be returned up to the query that encountered the error.
 func (t *ReadWriteTransaction) BatchUpdate(ctx context.Context, stmts []Statement) (_ []int64, err error) {
-	ctx = startSpan(ctx, "cloud.google.com/go/spanner.BatchUpdate")
-	defer func() { endSpan(ctx, err) }()
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/spanner.BatchUpdate")
+	defer func() { trace.EndSpan(ctx, err) }()
 
 	sh, ts, err := t.acquire(ctx)
 	if err != nil {