machinery-plugins/locks/etcd/etcd.go

91 lines
1.9 KiB
Go
Raw Normal View History

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