94 lines
2.0 KiB
Go
94 lines
2.0 KiB
Go
package etcd
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/RichardKnop/machinery/v2/config"
|
|
"github.com/RichardKnop/machinery/v2/locks/iface"
|
|
"github.com/RichardKnop/machinery/v2/log"
|
|
clientv3 "go.etcd.io/etcd/client/v3"
|
|
"go.etcd.io/etcd/client/v3/concurrency"
|
|
)
|
|
|
|
var (
|
|
// ErrLockFailed ..
|
|
ErrLockFailed = errors.New("etcd lock: failed to acquire lock")
|
|
)
|
|
|
|
type etcdLock struct {
|
|
conf *config.Config
|
|
client *clientv3.Client
|
|
retries int
|
|
}
|
|
|
|
// New ..
|
|
func New(ctx context.Context, conf *config.Config, retries int) (iface.Lock, error) {
|
|
etcdConf := clientv3.Config{
|
|
Endpoints: []string{conf.Lock},
|
|
Context: ctx,
|
|
DialTimeout: time.Second * 5,
|
|
}
|
|
|
|
client, err := clientv3.New(etcdConf)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
lock := etcdLock{
|
|
conf: conf,
|
|
client: client,
|
|
retries: retries,
|
|
}
|
|
|
|
return &lock, nil
|
|
}
|
|
|
|
// LockWithRetries ..
|
|
func (r etcdLock) LockWithRetries(key string, unixTsToExpireNs int64) error {
|
|
i := 0
|
|
for ; i < r.retries; i++ {
|
|
err := r.Lock(key, unixTsToExpireNs)
|
|
if err == nil {
|
|
// 成功拿到锁，返回
|
|
return nil
|
|
}
|
|
|
|
log.DEBUG.Printf("acquired lock=%s failed, retries=%d, err=%s", key, i, err)
|
|
time.Sleep(time.Millisecond * 100)
|
|
}
|
|
|
|
log.INFO.Printf("acquired lock=%s failed, retries=%d", key, i)
|
|
return ErrLockFailed
|
|
}
|
|
|
|
// Lock ..
|
|
func (r etcdLock) Lock(key string, unixTsToExpireNs int64) error {
|
|
now := time.Now().UnixNano()
|
|
ttl := time.Duration(unixTsToExpireNs + 1 - now)
|
|
|
|
// 创建一个新的session
|
|
s, err := concurrency.NewSession(r.client, concurrency.WithTTL(int(ttl.Seconds())))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer s.Orphan()
|
|
|
|
lockKey := fmt.Sprintf("/machinery/v2/lock/%s", strings.TrimRight(key, "/"))
|
|
m := concurrency.NewMutex(s, lockKey)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
|
|
defer cancel()
|
|
|
|
if err := m.Lock(ctx); err != nil {
|
|
_ = s.Close()
|
|
return err
|
|
}
|
|
|
|
log.INFO.Printf("acquired lock=%s, duration=%s", key, ttl)
|
|
return nil
|
|
}
|