feat: add trace in redis & mon & sql (#1799)

* feat: add sub spanId with redis Signed-off-by: chenquan <chenquan.dev@gmail.com> * add tests Signed-off-by: chenquan <chenquan.dev@gmail.com> * fix a bug Signed-off-by: chenquan <chenquan.dev@gmail.com> * feat: add sub spanId in sql Signed-off-by: chenquan <chenquan.dev@gmail.com> * feat: add sub spanId in mon Signed-off-by: chenquan <chenquan.dev@gmail.com> * chore: optimize code Signed-off-by: chenquan <chenquan.dev@gmail.com> * feat: add breaker in warpSession Signed-off-by: chenquan <chenquan.dev@gmail.com> * chore: optimize code Signed-off-by: chenquan <chenquan.dev@gmail.com> * test: add tests Signed-off-by: chenquan <chenquan.dev@gmail.com> * chore: reformat code Signed-off-by: chenquan <chenquan.dev@gmail.com> * fix: fix typo Signed-off-by: chenquan <chenquan.dev@gmail.com> * fix a bug Signed-off-by: chenquan <chenquan.dev@gmail.com>
2022-04-21 20:04:44 -05:00
parent 94ddb3380e
commit 162e9cef86
7 changed files with 268 additions and 15 deletions
--- a/core/stores/mon/collection.go
+++ b/core/stores/mon/collection.go
@@ -8,8 +8,12 @@ import (
 	"github.com/zeromicro/go-zero/core/breaker"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/timex"
+	"github.com/zeromicro/go-zero/core/trace"
 	"go.mongodb.org/mongo-driver/mongo"
 	mopt "go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/x/mongo/driver/session"
+	"go.opentelemetry.io/otel"
+	tracesdk "go.opentelemetry.io/otel/trace"
 )

 const defaultSlowThreshold = time.Millisecond * 500
@@ -112,6 +116,9 @@ func newCollection(collection *mongo.Collection, brk breaker.Breaker) Collection

 func (c *decoratedCollection) Aggregate(ctx context.Context, pipeline interface{},
 	opts ...*mopt.AggregateOptions) (cur *mongo.Cursor, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		starTime := timex.Now()
 		defer func() {
@@ -126,6 +133,9 @@ func (c *decoratedCollection) Aggregate(ctx context.Context, pipeline interface{

 func (c *decoratedCollection) BulkWrite(ctx context.Context, models []mongo.WriteModel,
 	opts ...*mopt.BulkWriteOptions) (res *mongo.BulkWriteResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -140,6 +150,9 @@ func (c *decoratedCollection) BulkWrite(ctx context.Context, models []mongo.Writ

 func (c *decoratedCollection) CountDocuments(ctx context.Context, filter interface{},
 	opts ...*mopt.CountOptions) (count int64, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -154,6 +167,9 @@ func (c *decoratedCollection) CountDocuments(ctx context.Context, filter interfa

 func (c *decoratedCollection) DeleteMany(ctx context.Context, filter interface{},
 	opts ...*mopt.DeleteOptions) (res *mongo.DeleteResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -168,6 +184,9 @@ func (c *decoratedCollection) DeleteMany(ctx context.Context, filter interface{}

 func (c *decoratedCollection) DeleteOne(ctx context.Context, filter interface{},
 	opts ...*mopt.DeleteOptions) (res *mongo.DeleteResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -182,6 +201,9 @@ func (c *decoratedCollection) DeleteOne(ctx context.Context, filter interface{},

 func (c *decoratedCollection) Distinct(ctx context.Context, fieldName string, filter interface{},
 	opts ...*mopt.DistinctOptions) (val []interface{}, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -196,6 +218,9 @@ func (c *decoratedCollection) Distinct(ctx context.Context, fieldName string, fi

 func (c *decoratedCollection) EstimatedDocumentCount(ctx context.Context,
 	opts ...*mopt.EstimatedDocumentCountOptions) (val int64, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -210,6 +235,9 @@ func (c *decoratedCollection) EstimatedDocumentCount(ctx context.Context,

 func (c *decoratedCollection) Find(ctx context.Context, filter interface{},
 	opts ...*mopt.FindOptions) (cur *mongo.Cursor, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -224,6 +252,9 @@ func (c *decoratedCollection) Find(ctx context.Context, filter interface{},

 func (c *decoratedCollection) FindOne(ctx context.Context, filter interface{},
 	opts ...*mopt.FindOneOptions) (res *mongo.SingleResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -239,6 +270,9 @@ func (c *decoratedCollection) FindOne(ctx context.Context, filter interface{},

 func (c *decoratedCollection) FindOneAndDelete(ctx context.Context, filter interface{},
 	opts ...*mopt.FindOneAndDeleteOptions) (res *mongo.SingleResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -255,6 +289,9 @@ func (c *decoratedCollection) FindOneAndDelete(ctx context.Context, filter inter
 func (c *decoratedCollection) FindOneAndReplace(ctx context.Context, filter interface{},
 	replacement interface{}, opts ...*mopt.FindOneAndReplaceOptions) (
 	res *mongo.SingleResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -270,6 +307,9 @@ func (c *decoratedCollection) FindOneAndReplace(ctx context.Context, filter inte

 func (c *decoratedCollection) FindOneAndUpdate(ctx context.Context, filter interface{}, update interface{},
 	opts ...*mopt.FindOneAndUpdateOptions) (res *mongo.SingleResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -285,6 +325,9 @@ func (c *decoratedCollection) FindOneAndUpdate(ctx context.Context, filter inter

 func (c *decoratedCollection) InsertMany(ctx context.Context, documents []interface{},
 	opts ...*mopt.InsertManyOptions) (res *mongo.InsertManyResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -299,6 +342,9 @@ func (c *decoratedCollection) InsertMany(ctx context.Context, documents []interf

 func (c *decoratedCollection) InsertOne(ctx context.Context, document interface{},
 	opts ...*mopt.InsertOneOptions) (res *mongo.InsertOneResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -313,6 +359,9 @@ func (c *decoratedCollection) InsertOne(ctx context.Context, document interface{

 func (c *decoratedCollection) ReplaceOne(ctx context.Context, filter interface{}, replacement interface{},
 	opts ...*mopt.ReplaceOptions) (res *mongo.UpdateResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -327,6 +376,9 @@ func (c *decoratedCollection) ReplaceOne(ctx context.Context, filter interface{}

 func (c *decoratedCollection) UpdateByID(ctx context.Context, id interface{}, update interface{},
 	opts ...*mopt.UpdateOptions) (res *mongo.UpdateResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -341,6 +393,9 @@ func (c *decoratedCollection) UpdateByID(ctx context.Context, id interface{}, up

 func (c *decoratedCollection) UpdateMany(ctx context.Context, filter interface{}, update interface{},
 	opts ...*mopt.UpdateOptions) (res *mongo.UpdateResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -355,6 +410,9 @@ func (c *decoratedCollection) UpdateMany(ctx context.Context, filter interface{}

 func (c *decoratedCollection) UpdateOne(ctx context.Context, filter interface{}, update interface{},
 	opts ...*mopt.UpdateOptions) (res *mongo.UpdateResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = c.brk.DoWithAcceptable(func() error {
 		startTime := timex.Now()
 		defer func() {
@@ -414,5 +472,14 @@ func (p keepablePromise) keep(err error) error {

 func acceptable(err error) bool {
 	return err == nil || err == mongo.ErrNoDocuments || err == mongo.ErrNilValue ||
-		err == mongo.ErrNilDocument || err == mongo.ErrNilCursor || err == mongo.ErrEmptySlice
+		err == mongo.ErrNilDocument || err == mongo.ErrNilCursor || err == mongo.ErrEmptySlice ||
+		// session err
+		err == session.ErrSessionEnded || err == session.ErrNoTransactStarted || err == session.ErrTransactInProgress ||
+		err == session.ErrAbortAfterCommit || err == session.ErrAbortTwice || err == session.ErrCommitAfterAbort ||
+		err == session.ErrUnackWCUnsupported || err == session.ErrSnapshotTransaction
+}
+
+func startSpan(ctx context.Context) (context.Context, tracesdk.Span) {
+	tracer := otel.GetTracerProvider().Tracer(trace.TraceName)
+	return tracer.Start(ctx, "mongo")
 }
--- a/core/stores/mon/model.go
+++ b/core/stores/mon/model.go
@@ -11,14 +11,21 @@ import (
 	mopt "go.mongodb.org/mongo-driver/mongo/options"
 )

-// Model is a mongodb store model that represents a collection.
-type Model struct {
-	Collection
-	name string
-	cli  *mongo.Client
-	brk  breaker.Breaker
-	opts []Option
-}
+type (
+	// Model is a mongodb store model that represents a collection.
+	Model struct {
+		Collection
+		name string
+		cli  *mongo.Client
+		brk  breaker.Breaker
+		opts []Option
+	}
+
+	wrapSession struct {
+		mongo.Session
+		brk breaker.Breaker
+	}
+)

 // MustNewModel returns a Model, exits on errors.
 func MustNewModel(uri, db, collection string, opts ...Option) *Model {
@@ -62,8 +69,14 @@ func (m *Model) StartSession(opts ...*mopt.SessionOptions) (sess mongo.Session,
 			logDuration(m.name, "StartSession", starTime, err)
 		}()

-		sess, err = m.cli.StartSession(opts...)
-		return err
+		session, sessionErr := m.cli.StartSession(opts...)
+		if sessionErr != nil {
+			return sessionErr
+		}
+
+		sess = &wrapSession{Session: session, brk: m.brk}
+
+		return nil
 	}, acceptable)
 	return
 }
@@ -152,3 +165,43 @@ func (m *Model) FindOneAndUpdate(ctx context.Context, v, filter interface{}, upd

 	return res.Decode(v)
 }
+
+func (w *wrapSession) AbortTransaction(ctx context.Context) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	return w.brk.DoWithAcceptable(func() error {
+		return w.Session.AbortTransaction(ctx)
+	}, acceptable)
+}
+
+func (w *wrapSession) CommitTransaction(ctx context.Context) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	return w.brk.DoWithAcceptable(func() error {
+		return w.Session.CommitTransaction(ctx)
+	}, acceptable)
+}
+
+func (w *wrapSession) WithTransaction(ctx context.Context, fn func(sessCtx mongo.SessionContext) (interface{}, error), opts ...*mopt.TransactionOptions) (res interface{}, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = w.brk.DoWithAcceptable(func() error {
+		res, err = w.Session.WithTransaction(ctx, fn, opts...)
+		return err
+	}, acceptable)
+
+	return
+}
+
+func (w *wrapSession) EndSession(ctx context.Context) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	_ = w.brk.DoWithAcceptable(func() error {
+		w.Session.EndSession(ctx)
+		return nil
+	}, acceptable)
+}
--- a/core/stores/mon/model_test.go
+++ b/core/stores/mon/model_test.go
@@ -18,6 +18,17 @@ func TestModel_StartSession(t *testing.T) {
 		m := createModel(mt)
 		sess, err := m.StartSession()
 		assert.Nil(t, err)
+
+		_, err = sess.WithTransaction(context.Background(), func(sessCtx mongo.SessionContext) (interface{}, error) {
+			_ = sessCtx.StartTransaction()
+			sessCtx.Client().Database("1")
+			sessCtx.EndSession(context.Background())
+			return nil, nil
+		})
+		assert.Nil(t, err)
+
+		assert.NoError(t, sess.CommitTransaction(context.Background()))
+		assert.Error(t, sess.AbortTransaction(context.Background()))
 		sess.EndSession(context.Background())
 	})
 }