diff --git a/firestore/client.go b/firestore/client.go index 3acfd1a1bf1..8ee1071d283 100644 --- a/firestore/client.go +++ b/firestore/client.go @@ -199,7 +199,10 @@ func (c *Client) GetAll(ctx context.Context, docRefs []*DocumentRef) (_ []*Docum return c.getAll(ctx, docRefs, nil) } -func (c *Client) getAll(ctx context.Context, docRefs []*DocumentRef, tid []byte) ([]*DocumentSnapshot, error) { +func (c *Client) getAll(ctx context.Context, docRefs []*DocumentRef, tid []byte) (_ []*DocumentSnapshot, err error) { + ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.Client.BatchGetDocuments") + defer func() { trace.EndSpan(ctx, err) }() + var docNames []string docIndices := map[string][]int{} // doc name to positions in docRefs for i, dr := range docRefs { @@ -267,6 +270,9 @@ func (c *Client) getAll(ctx context.Context, docRefs []*DocumentRef, tid []byte) // Collections returns an iterator over the top-level collections. func (c *Client) Collections(ctx context.Context) *CollectionIterator { + ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.Client.ListCollectionIds") + defer func() { trace.EndSpan(ctx, nil) }() + it := &CollectionIterator{ client: c, it: c.c.ListCollectionIds( @@ -286,7 +292,10 @@ func (c *Client) Batch() *WriteBatch { } // commit calls the Commit RPC outside of a transaction. -func (c *Client) commit(ctx context.Context, ws []*pb.Write) ([]*WriteResult, error) { +func (c *Client) commit(ctx context.Context, ws []*pb.Write) (_ []*WriteResult, err error) { + ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.Client.commit") + defer func() { trace.EndSpan(ctx, err) }() + req := &pb.CommitRequest{ Database: c.path(), Writes: ws, diff --git a/firestore/docref.go b/firestore/docref.go index e5d4b08439a..fe5b2edd687 100644 --- a/firestore/docref.go +++ b/firestore/docref.go @@ -677,6 +677,9 @@ func (d *DocumentRef) Update(ctx context.Context, updates []Update, preconds ... // Collections returns an iterator over the immediate sub-collections of the document. func (d *DocumentRef) Collections(ctx context.Context) *CollectionIterator { + ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.DocumentRef.ListCollectionIds") + defer func() { trace.EndSpan(ctx, nil) }() + client := d.Parent.c it := &CollectionIterator{ client: client, diff --git a/firestore/list_documents.go b/firestore/list_documents.go index fc6db8626a6..712697f3d8b 100644 --- a/firestore/list_documents.go +++ b/firestore/list_documents.go @@ -18,6 +18,7 @@ import ( "context" 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" ) @@ -33,6 +34,9 @@ type DocumentRefIterator struct { } func newDocumentRefIterator(ctx context.Context, cr *CollectionRef, tid []byte) *DocumentRefIterator { + ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.ListDocuments") + defer func() { trace.EndSpan(ctx, nil) }() + client := cr.c req := &pb.ListDocumentsRequest{ Parent: cr.parentPath, diff --git a/firestore/query.go b/firestore/query.go index a49a4b262e0..8e731d82618 100644 --- a/firestore/query.go +++ b/firestore/query.go @@ -24,6 +24,7 @@ import ( "time" "cloud.google.com/go/internal/btree" + "cloud.google.com/go/internal/trace" "github.com/golang/protobuf/ptypes/wrappers" "google.golang.org/api/iterator" pb "google.golang.org/genproto/googleapis/firestore/v1" @@ -699,7 +700,10 @@ func newQueryDocumentIterator(ctx context.Context, q *Query, tid []byte) *queryD } } -func (it *queryDocumentIterator) next() (*DocumentSnapshot, error) { +func (it *queryDocumentIterator) next() (_ *DocumentSnapshot, err error) { + it.ctx = trace.StartSpan(it.ctx, "cloud.google.com/go/firestore.Query.RunQuery") + defer func() { trace.EndSpan(it.ctx, err) }() + client := it.q.c if it.streamClient == nil { sq, err := it.q.toProto() @@ -719,7 +723,6 @@ func (it *queryDocumentIterator) next() (*DocumentSnapshot, error) { } } var res *pb.RunQueryResponse - var err error for { res, err = it.streamClient.Recv() if err == io.EOF { diff --git a/firestore/transaction.go b/firestore/transaction.go index 392168c718d..00c2f1d1b12 100644 --- a/firestore/transaction.go +++ b/firestore/transaction.go @@ -116,11 +116,13 @@ func (c *Client) RunTransaction(ctx context.Context, f func(context.Context, *Tr // TODO(jba): get backoff time from gRPC trailer metadata? See // extractRetryDelay in https://code.googlesource.com/gocloud/+/master/spanner/retry.go. for i := 0; i < t.maxAttempts; i++ { + t.ctx = trace.StartSpan(t.ctx, "cloud.google.com/go/firestore.Client.BeginTransaction") var res *pb.BeginTransactionResponse res, err = t.c.c.BeginTransaction(t.ctx, &pb.BeginTransactionRequest{ Database: db, Options: txOpts, }) + trace.EndSpan(t.ctx, err) if err != nil { return err } @@ -136,11 +138,14 @@ func (c *Client) RunTransaction(ctx context.Context, f func(context.Context, *Tr // Prefer f's returned error to rollback error. return err } + t.ctx = trace.StartSpan(t.ctx, "cloud.google.com/go/firestore.Client.Commit") _, err = t.c.c.Commit(t.ctx, &pb.CommitRequest{ Database: t.c.path(), Writes: t.writes, Transaction: t.id, }) + trace.EndSpan(t.ctx, err) + // If a read-write transaction returns Aborted, retry. // On success or other failures, return here. if t.readOnly || status.Code(err) != codes.Aborted { diff --git a/firestore/watch.go b/firestore/watch.go index e4ac7fb3ded..bf4528d069d 100644 --- a/firestore/watch.go +++ b/firestore/watch.go @@ -24,6 +24,7 @@ import ( "time" "cloud.google.com/go/internal/btree" + "cloud.google.com/go/internal/trace" "github.com/golang/protobuf/ptypes" gax "github.com/googleapis/gax-go/v2" pb "google.golang.org/genproto/googleapis/firestore/v1" @@ -492,9 +493,12 @@ func (s *watchStream) recv() (*pb.ListenResponse, error) { } } -func (s *watchStream) open() (pb.Firestore_ListenClient, error) { +func (s *watchStream) open() (lc pb.Firestore_ListenClient, err error) { + s.ctx = trace.StartSpan(s.ctx, "cloud.google.com/go/firestore.watchStream.Listen") + defer func() { trace.EndSpan(s.ctx, err) }() + dbPath := s.c.path() - lc, err := s.c.c.Listen(withResourceHeader(s.ctx, dbPath)) + lc, err = s.c.c.Listen(withResourceHeader(s.ctx, dbPath)) if err == nil { err = lc.Send(&pb.ListenRequest{ Database: dbPath,