github.com/wfusion/gofusion@v1.1.14/lock/mongo.go (about)

     1  package lock
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/pkg/errors"
    10  	"go.mongodb.org/mongo-driver/bson"
    11  	"go.mongodb.org/mongo-driver/mongo/options"
    12  	"go.mongodb.org/mongo-driver/mongo/writeconcern"
    13  
    14  	"github.com/wfusion/gofusion/common/utils"
    15  	"github.com/wfusion/gofusion/config"
    16  	"github.com/wfusion/gofusion/mongo"
    17  
    18  	mgoDrv "go.mongodb.org/mongo-driver/mongo"
    19  )
    20  
    21  var (
    22  	mongoInitLocker sync.Mutex
    23  )
    24  
    25  type mongoLocker struct {
    26  	ctx       context.Context
    27  	appName   string
    28  	mongoName string
    29  	collName  string
    30  }
    31  
    32  func newMongoLocker(ctx context.Context, appName, mongoName, collName string) Lockable {
    33  	mongoInitLocker.Lock()
    34  	defer mongoInitLocker.Unlock()
    35  
    36  	db := mongo.Use(mongoName, mongo.AppName(appName), mongo.WriteConcern(writeconcern.Majority()))
    37  	coll := db.Collection(collName)
    38  	colls, err := db.ListCollectionNames(ctx, bson.M{"name": collName})
    39  	if err != nil {
    40  		panic(errors.Errorf("%s lock component mongo %s parse collection %s failed: %s",
    41  			appName, mongoName, collName, err))
    42  	}
    43  	if len(colls) == 0 {
    44  		if err = db.CreateCollection(ctx, collName); err != nil {
    45  			panic(errors.Errorf("%s lock component mongo %s create collection %s failed: %s",
    46  				appName, mongoName, collName, err))
    47  		}
    48  	}
    49  
    50  	ttlIdxModel := mgoDrv.IndexModel{
    51  		Keys:    bson.M{"expires_at": 1},
    52  		Options: options.Index().SetExpireAfterSeconds(0),
    53  	}
    54  	if _, err = coll.Indexes().CreateOne(ctx, ttlIdxModel); err != nil {
    55  		panic(errors.Errorf("%s lock component mongo %s create ttl index %s failed: %s",
    56  			appName, mongoName, collName, err))
    57  	}
    58  
    59  	indexModel := mgoDrv.IndexModel{
    60  		Keys:    bson.M{"lock_key": 1},
    61  		Options: options.Index().SetUnique(true),
    62  	}
    63  	if _, err = coll.Indexes().CreateOne(ctx, indexModel); err != nil {
    64  		panic(errors.Errorf("%s lock component mongo %s create lock index %s failed: %s",
    65  			appName, mongoName, collName, err))
    66  	}
    67  
    68  	return &mongoLocker{ctx: ctx, appName: appName, mongoName: mongoName, collName: collName}
    69  }
    70  
    71  func (m *mongoLocker) Lock(ctx context.Context, key string, opts ...utils.OptionExtender) (err error) {
    72  	opt := utils.ApplyOptions[lockOption](opts...)
    73  	expired := tolerance
    74  	if opt.expired > 0 {
    75  		expired = opt.expired
    76  	}
    77  	now := time.Now()
    78  	lockKey := m.formatLockKey(key)
    79  	filter := bson.M{
    80  		"lock_key": bson.M{"$eq": lockKey, "$exists": true},
    81  		"$or": []bson.M{
    82  			{"expires_at": bson.M{"$lt": now, "$exists": true}},
    83  			{"holder": bson.M{"$eq": opt.reentrantKey, "$exists": true}},
    84  		},
    85  	}
    86  	update := bson.M{
    87  		"$setOnInsert": bson.M{
    88  			"lock_key": lockKey,
    89  			"holder":   opt.reentrantKey,
    90  		},
    91  		"$inc": bson.M{"count": 1},
    92  		"$max": bson.M{"expires_at": now.Add(expired)},
    93  	}
    94  	mopts := options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After)
    95  
    96  	updatedDoc := new(mongoLockDoc)
    97  	err = mongo.Use(m.mongoName, mongo.AppName(m.appName), mongo.WriteConcern(writeconcern.Majority())).
    98  		Collection(m.collName).
    99  		FindOneAndUpdate(ctx, filter, update, mopts).
   100  		Decode(updatedDoc)
   101  	if err != nil {
   102  		return
   103  	}
   104  
   105  	if updatedDoc.LockKey == lockKey && updatedDoc.Holder == opt.reentrantKey {
   106  		return
   107  	}
   108  
   109  	return ErrTimeout
   110  }
   111  
   112  func (m *mongoLocker) Unlock(ctx context.Context, key string, opts ...utils.OptionExtender) (err error) {
   113  	opt := utils.ApplyOptions[lockOption](opts...)
   114  	filter := bson.M{
   115  		"lock_key": m.formatLockKey(key),
   116  		"holder":   opt.reentrantKey,
   117  		"count":    bson.M{"$gt": 0},
   118  	}
   119  	update := bson.M{
   120  		"$inc": bson.M{"count": -1},
   121  		"$max": bson.M{"expires_at": time.Now().Add(opt.expired / 2)},
   122  	}
   123  	mopts := options.FindOneAndUpdate().SetReturnDocument(options.After)
   124  
   125  	updatedDoc := new(mongoLockDoc)
   126  	coll := mongo.
   127  		Use(m.mongoName, mongo.AppName(m.appName), mongo.WriteConcern(writeconcern.Majority())).
   128  		Collection(m.collName)
   129  	err = coll.FindOneAndUpdate(ctx, filter, update, mopts).Decode(&updatedDoc)
   130  	if err != nil {
   131  		return
   132  	}
   133  	if updatedDoc.Count <= 0 {
   134  		_, err = coll.DeleteOne(ctx, bson.M{
   135  			"lock_key": m.formatLockKey(key),
   136  			"holder":   opt.reentrantKey,
   137  			"count":    bson.M{"$lte": 0},
   138  		})
   139  	}
   140  
   141  	return
   142  }
   143  
   144  func (m *mongoLocker) ReentrantLock(ctx context.Context, key, reentrantKey string,
   145  	opts ...utils.OptionExtender) (err error) {
   146  	opt := utils.ApplyOptions[lockOption](opts...)
   147  	if utils.IsStrBlank(opt.reentrantKey) {
   148  		return ErrReentrantKeyNotFound
   149  	}
   150  	return m.Lock(ctx, key, append(opts, ReentrantKey(reentrantKey))...)
   151  }
   152  
   153  func (m *mongoLocker) formatLockKey(key string) (format string) {
   154  	return fmt.Sprintf("%s_%s", config.Use(m.appName).AppName(), key)
   155  }
   156  
   157  type mongoLockDoc struct {
   158  	LockKey   string    `bson:"lock_key"`
   159  	Holder    string    `bson:"holder"`
   160  	ExpiresAt time.Time `bson:"expires_at"`
   161  	Count     int       `bson:"count"`
   162  }