firestore: add tracing for exported RPC-making methods

Adds tracing with OpenCensus for exported methods
that make RPCs.

Since the methods have different receivers,
their spans have varying basename prefixes e.g.

* "cloud.google.com/go/firestore.DocumentRef.Get"
* "cloud.google.com/go/firestore.Transaction.Rollback"
* "cloud.google.com/go/firestore.WriteBatch.Commit"

Fixes #1375

Change-Id: I7734fb529dd5d2c5559c25ba6634baa26fc4428f
Reviewed-on: https://code-review.googlesource.com/c/gocloud/+/39590
Reviewed-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Jean de Klerk <deklerk@google.com>

firestore/client.go

@@ -24,6 +24,7 @@ import (
 	"time"
 
 	vkit "cloud.google.com/go/firestore/apiv1"
+	"cloud.google.com/go/internal/trace"
 	"cloud.google.com/go/internal/version"
 	"github.com/golang/protobuf/ptypes"
 	gax "github.com/googleapis/gax-go/v2"
@@ -137,7 +138,10 @@ func (c *Client) idsToRef(IDs []string, dbPath string) (*CollectionRef, *Documen
 // returned in the order of the given DocumentRefs.
 //
 // If a document is not present, the corresponding DocumentSnapshot's Exists method will return false.
-func (c *Client) GetAll(ctx context.Context, docRefs []*DocumentRef) ([]*DocumentSnapshot, error) {
+func (c *Client) GetAll(ctx context.Context, docRefs []*DocumentRef) (_ []*DocumentSnapshot, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.GetAll")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	return c.getAll(ctx, docRefs, nil)
 }
 

firestore/docref.go

@@ -23,6 +23,7 @@ import (
 	"sort"
 
 	vkit "cloud.google.com/go/firestore/apiv1"
+	"cloud.google.com/go/internal/trace"
 	"google.golang.org/api/iterator"
 	pb "google.golang.org/genproto/googleapis/firestore/v1"
 	"google.golang.org/grpc/codes"
@@ -67,7 +68,10 @@ func (d *DocumentRef) Collection(id string) *CollectionRef {
 //    grpc.Code(err) == codes.NotFound
 // In that case, Get returns a non-nil DocumentSnapshot whose Exists method return false and whose
 // ReadTime is the time of the failed read operation.
-func (d *DocumentRef) Get(ctx context.Context) (*DocumentSnapshot, error) {
+func (d *DocumentRef) Get(ctx context.Context) (_ *DocumentSnapshot, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.DocumentRef.Get")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	if d == nil {
 		return nil, errNilDocRef
 	}
@@ -121,7 +125,10 @@ func (d *DocumentRef) Get(ctx context.Context) (*DocumentSnapshot, error) {
 //   - serverTimestamp: The field must be of type time.Time. When writing, if
 //     the field has the zero value, the server will populate the stored document with
 //     the time that the request is processed.
-func (d *DocumentRef) Create(ctx context.Context, data interface{}) (*WriteResult, error) {
+func (d *DocumentRef) Create(ctx context.Context, data interface{}) (_ *WriteResult, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.DocumentRef.Create")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	ws, err := d.newCreateWrites(data)
 	if err != nil {
 		return nil, err
@@ -150,7 +157,10 @@ func (d *DocumentRef) newCreateWrites(data interface{}) ([]*pb.Write, error) {
 // completely. Specify one of the Merge options to preserve an existing document's
 // fields. To delete some fields, use a Merge option with firestore.Delete as the
 // field value.
-func (d *DocumentRef) Set(ctx context.Context, data interface{}, opts ...SetOption) (*WriteResult, error) {
+func (d *DocumentRef) Set(ctx context.Context, data interface{}, opts ...SetOption) (_ *WriteResult, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.DocumentRef.Set")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	ws, err := d.newSetWrites(data, opts)
 	if err != nil {
 		return nil, err
@@ -227,7 +237,10 @@ func fpvsFromData(v reflect.Value, prefix FieldPath, fpvs *[]fpv) {
 
 // Delete deletes the document. If the document doesn't exist, it does nothing
 // and returns no error.
-func (d *DocumentRef) Delete(ctx context.Context, preconds ...Precondition) (*WriteResult, error) {
+func (d *DocumentRef) Delete(ctx context.Context, preconds ...Precondition) (_ *WriteResult, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.DocumentRef.Delete")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	ws, err := d.newDeleteWrites(preconds)
 	if err != nil {
 		return nil, err
@@ -520,15 +533,18 @@ func (u *Update) process() (fpv, error) {
 
 // Update updates the document. The values at the given
 // field paths are replaced, but other fields of the stored document are untouched.
-func (d *DocumentRef) Update(ctx context.Context, updates []Update, preconds ...Precondition) (*WriteResult, error) {
+func (d *DocumentRef) Update(ctx context.Context, updates []Update, preconds ...Precondition) (_ *WriteResult, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.DocumentRef.Update")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	ws, err := d.newUpdatePathWrites(updates, preconds)
 	if err != nil {
 		return nil, err
 	}
 	return d.Parent.c.commit1(ctx, ws)
 }
 
-// Collections returns an interator over the immediate sub-collections of the document.
+// Collections returns an iterator over the immediate sub-collections of the document.
 func (d *DocumentRef) Collections(ctx context.Context) *CollectionIterator {
 	client := d.Parent.c
 	it := &CollectionIterator{

firestore/transaction.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"errors"
 
+	"cloud.google.com/go/internal/trace"
 	gax "github.com/googleapis/gax-go/v2"
 	pb "google.golang.org/genproto/googleapis/firestore/v1"
 	"google.golang.org/grpc"
@@ -89,7 +90,10 @@ type transactionInProgressKey struct{}
 //
 // Since f may be called more than once, f should usually be idempotent – that is, it
 // should have the same result when called multiple times.
-func (c *Client) RunTransaction(ctx context.Context, f func(context.Context, *Transaction) error, opts ...TransactionOption) error {
+func (c *Client) RunTransaction(ctx context.Context, f func(context.Context, *Transaction) error, opts ...TransactionOption) (err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.Client.RunTransaction")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	if ctx.Value(transactionInProgressKey{}) != nil {
 		return errNestedTransaction
 	}
@@ -112,7 +116,6 @@ func (c *Client) RunTransaction(ctx context.Context, f func(context.Context, *Tr
 	// TODO(jba): use other than the standard backoff parameters?
 	// TODO(jba): get backoff time from gRPC trailer metadata? See
 	// extractRetryDelay in https://code.googlesource.com/gocloud/+/master/spanner/retry.go.
-	var err error
 	for i := 0; i < t.maxAttempts; i++ {
 		var res *pb.BeginTransactionResponse
 		res, err = t.c.c.BeginTransaction(t.ctx, &pb.BeginTransactionRequest{

firestore/writebatch.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"errors"
 
+	"cloud.google.com/go/internal/trace"
 	pb "google.golang.org/genproto/googleapis/firestore/v1"
 )
 
@@ -70,7 +71,10 @@ func (b *WriteBatch) Update(dr *DocumentRef, data []Update, opts ...Precondition
 // Commit applies all the writes in the batch to the database atomically. Commit
 // returns an error if there are no writes in the batch, if any errors occurred in
 // constructing the writes, or if the Commmit operation fails.
-func (b *WriteBatch) Commit(ctx context.Context) ([]*WriteResult, error) {
+func (b *WriteBatch) Commit(ctx context.Context) (_ []*WriteResult, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.WriteBatch.Commit")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	if b.err != nil {
 		return nil, b.err
 	}