github.com/teng231/glock@v1.1.11/distributed_lock.go (about)

     1  package glock
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"time"
     7  
     8  	"github.com/go-redis/redis/v8"
     9  	"github.com/go-redsync/redsync/v4"
    10  	"github.com/go-redsync/redsync/v4/redis/goredis/v8"
    11  )
    12  
    13  type IDistributedLock interface {
    14  	// Lock is close request in. Lock all request come to the gate
    15  	Lock(key string) (*LockContext, error)
    16  	// Unlock release the gate
    17  	Unlock(lc *LockContext) error
    18  }
    19  
    20  type DistributedLock struct {
    21  	sync     *redsync.Redsync
    22  	timelock time.Duration
    23  	prefix   string
    24  }
    25  
    26  type LockContext struct {
    27  	mutex *redsync.Mutex
    28  	ctx   context.Context
    29  }
    30  
    31  func StartDistributedLock(cf *ConnectConfig) (*DistributedLock, error) {
    32  	if cf.MaxRetries == 0 {
    33  		cf.MaxRetries = 10
    34  	}
    35  	if cf.MinRetryBackoff == 0 {
    36  		cf.MinRetryBackoff = 50 * time.Millisecond
    37  	}
    38  	if cf.MaxRetryBackoff == 0 {
    39  		cf.MaxRetryBackoff = 2 * time.Second
    40  	}
    41  	if cf.PoolSize == 0 {
    42  		cf.PoolSize = 1000
    43  	}
    44  	client := redis.NewClient(&redis.Options{
    45  		Password:        cf.RedisPw,
    46  		Addr:            cf.RedisAddr,
    47  		Username:        cf.RedisUsername,
    48  		MaxRetries:      cf.MaxRetries,
    49  		MinRetryBackoff: cf.MinRetryBackoff,
    50  		MaxRetryBackoff: cf.MaxRetryBackoff,
    51  		DialTimeout:     10 * time.Second,
    52  		PoolSize:        cf.PoolSize,
    53  		DB:              cf.RedisDb, // use default DB
    54  	})
    55  	if cf.Timelock < 0 {
    56  		return nil, errors.New("timelock is required")
    57  	}
    58  	ctx, cancel := context.WithTimeout(context.Background(), cf.Timelock)
    59  	defer cancel()
    60  	if err := client.Ping(ctx).Err(); err != nil {
    61  		return nil, err
    62  	}
    63  	pool := goredis.NewPool(client)
    64  
    65  	rs := redsync.New(pool)
    66  	return &DistributedLock{
    67  		sync:     rs,
    68  		timelock: cf.Timelock,
    69  		prefix:   cf.Prefix,
    70  	}, nil
    71  }
    72  
    73  // CreateDistributedLock deprecated
    74  func CreateDistributedLock(addr, pw, prefix string, timelock time.Duration) (*DistributedLock, error) {
    75  	client := redis.NewClient(&redis.Options{
    76  		Password:        pw,
    77  		Addr:            addr,
    78  		MaxRetries:      10,
    79  		MinRetryBackoff: 15 * time.Millisecond,
    80  		MaxRetryBackoff: 1000 * time.Millisecond,
    81  		DialTimeout:     10 * time.Second,
    82  		PoolSize:        1000,
    83  		DB:              1, // use default DB
    84  	})
    85  	if timelock < 0 {
    86  		return nil, errors.New("timelock is required")
    87  	}
    88  	ctx, cancel := context.WithTimeout(context.Background(), timelock)
    89  	defer cancel()
    90  	if err := client.Ping(ctx).Err(); err != nil {
    91  		return nil, err
    92  	}
    93  	pool := goredis.NewPool(client)
    94  
    95  	rs := redsync.New(pool)
    96  	return &DistributedLock{
    97  		sync:     rs,
    98  		timelock: timelock,
    99  		prefix:   prefix,
   100  	}, nil
   101  }
   102  
   103  func (d *DistributedLock) Lock(key string) (*LockContext, error) {
   104  	mutex := d.sync.NewMutex(d.prefix + key)
   105  	ctx, cancel := context.WithTimeout(context.Background(), d.timelock)
   106  	go func() {
   107  		<-ctx.Done()
   108  		cancel()
   109  	}()
   110  	if err := mutex.LockContext(ctx); err != nil {
   111  		cancel()
   112  		return nil, err
   113  	}
   114  	return &LockContext{mutex, ctx}, nil
   115  }
   116  
   117  func (d *DistributedLock) Unlock(lc *LockContext) error {
   118  	if _, err := lc.mutex.UnlockContext(lc.ctx); err != nil {
   119  		return err
   120  	}
   121  	return nil
   122  }