feat(firestore): add opencensus tracing support (#2942)
* feat(firestore): add opencensus tracing support
* add BeginTransaction span, rollback span delete
* use lowercase name for commit
Co-authored-by: Christopher Wilcox <crwilcox@google.com>
Co-authored-by: Chris Cotter <cjcotter@google.com>
diff --git a/firestore/client.go b/firestore/client.go
index 3acfd1a..8ee1071 100644
--- a/firestore/client.go
+++ b/firestore/client.go
@@ -199,7 +199,10 @@
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 @@
// 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 @@
}
// 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 e5d4b08..fe5b2ed 100644
--- a/firestore/docref.go
+++ b/firestore/docref.go
@@ -677,6 +677,9 @@
// 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 fc6db86..712697f 100644
--- a/firestore/list_documents.go
+++ b/firestore/list_documents.go
@@ -18,6 +18,7 @@
"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 @@
}
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 a49a4b2..8e731d8 100644
--- a/firestore/query.go
+++ b/firestore/query.go
@@ -24,6 +24,7 @@
"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 (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 @@
}
}
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 392168c..00c2f1d 100644
--- a/firestore/transaction.go
+++ b/firestore/transaction.go
@@ -116,11 +116,13 @@
// 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 @@
// 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 e4ac7fb..bf4528d 100644
--- a/firestore/watch.go
+++ b/firestore/watch.go
@@ -24,6 +24,7 @@
"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) 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,
