web_hgcg_server/pub/mongos/manager.go

package mongos

import (
	"context"
	"errors"
	"time"
	"videoweb/pub/common/log"

	"go.mongodb.org/mongo-driver/mongo/readconcern"
	"go.mongodb.org/mongo-driver/mongo/readpref"
	"go.mongodb.org/mongo-driver/mongo/writeconcern"
	"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"

	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// var (
// 	connPools *mongo.Client
// 	dataBase  *mongo.Database
// )

var (
	//业务库
	Shop *Manager
	//日志库
	ShopLog *Manager
)

type Manager struct {
	client *mongo.Client
	//dataBase *mongo.Database
	url string
	db  string
}

func NewManager(url string) *Manager {
	return &Manager{url: url}
}

func (m *Manager) Database() *mongo.Database {
	return m.client.Database(m.db)
}

func (m *Manager) Coll(name string) *Collection {
	coll := m.Database().Collection(name)
	return &Collection{name: name, coll: coll}
}

func (m *Manager) Init() error {
	cs, err := connstring.Parse(m.url)
	if err != nil {
		log.Error("db connect parse err", log.E(err))
		return err
	}
	if cs.Database == "" {
		log.Error("db connect database MUST contains")
		return errors.New("database MUST contains")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	client, err := mongo.Connect(ctx, options.Client().ApplyURI(m.url))
	if err != nil {
		log.Error("[warn] connect mongo db err", log.E(err))
		return err
	}
	if err = client.Ping(ctx, nil); err != nil {
		log.Error("[warn] ping mongo db err", log.E(err))
		return err
	}

	//connPools = client
	m.client = client
	m.db = cs.Database

	log.Info("mongos:Init ok!", log.Any("dbName", m.db))
	return nil
}

func (m *Manager) DisConnect() {
	if m.client != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := m.client.Disconnect(ctx); err != nil {
			log.Error("db disconnect err", log.E(err))
		}

		m.client = nil
	}
}

// Transaction 开启事务处理包裹处理,里面处理的全是利用的事务的ctx 切记里面的context必须要使用同一个，因为里面有重试机制，循环调用该函数
func (m *Manager) Transaction(ctx context.Context, fn func(sctx mongo.SessionContext) error) error {
	opts := options.Session().SetDefaultReadPreference(readpref.Primary())
	handler := func(sctx mongo.SessionContext) error {
		lHandler := func(ctx mongo.SessionContext) (err error) {
			tOpts := options.Transaction()
			tOpts.SetReadConcern(readconcern.Snapshot())
			tOpts.SetWriteConcern(writeconcern.New(writeconcern.WMajority()))
			err = sctx.StartTransaction(tOpts)
			if err != nil {
				return
			}
			defer func() {
				if rErr := recover(); rErr != nil {
					log.Info("caught panic during transaction, aborting...", log.Any("", rErr))
					aErr := sctx.AbortTransaction(sctx)
					if aErr != nil {
						log.Info("中止事务错误", log.E(aErr))
					}
				}
			}()
			err = fn(sctx)
			if err != nil {
				aErr := sctx.AbortTransaction(sctx)
				if aErr != nil {
					log.Info("中止事务错误", log.E(aErr))
				}
				log.Info("事务内错误", log.E(err))
				return
			}
			return commitWithRetry(sctx)
		}
		return runTransactionWithRetry(sctx, lHandler)
	}
	return m.client.UseSessionWithOptions(ctx, opts, handler)
}

func runTransactionWithRetry(sctx mongo.SessionContext, txnFn func(mongo.SessionContext) error) error {
	var err error
	for i := int64(0); i < 100; i++ {
		err = txnFn(sctx) // Performs transaction.
		if err == nil {
			return nil
		}
		// If transient error, retry the whole transaction
		if cmdErr, ok := err.(mongo.CommandError); ok && cmdErr.HasErrorLabel("TransientTransactionError") {
			continue
		}
		return err
	}
	return err
}

// commitWithRetry is an example function demonstrating transaction retry logic.
func commitWithRetry(sctx mongo.SessionContext) error {
	for {
		err := sctx.CommitTransaction(sctx)
		switch e := err.(type) {
		case nil:
			return nil
		case mongo.CommandError:
			// Can retry commit
			if e.HasErrorLabel("UnknownTransactionCommitResult") {
				continue
			}
			return e
		default:
			return e
		}
	}
}

// func Transaction(ctx context.Context, fn func(sctx mongo.SessionContext) error) error {
// 	return connMgr.Transaction(ctx, fn)
// }