feat: add lock

examples/etcd/main.go

@@ -80,7 +80,7 @@ func startServer() (*machinery.Server, error) {
 	}
 
 	// lock := redislock.New(cnf, []string{"localhost:6379"}, 3, 2)
-	lock, err := etcdlock.New(cnf, "http://127.0.0.1:2379")
+	lock, err := etcdlock.New(cnf, "http://127.0.0.1:2379", 3)
 	if err != nil {
 		return nil, err
 	}
@@ -206,7 +206,7 @@ func send() error {
 		return err
 	}
 
-	time.Sleep(time.Second * 120)
+	time.Sleep(time.Second * 120000)
 
 	return nil
 }

locks/etcd/etcd.go

@@ -1,13 +1,17 @@
 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 (
@@ -19,10 +23,9 @@ type etcdLock struct {
 	conf       clientv3.Config
 	cli        *clientv3.Client
 	retries    int
-	interval   time.Duration
 }
 
-func New(cnf *config.Config, endpoint string) (iface.Lock, error) {
+func New(cnf *config.Config, endpoint string, retries int) (iface.Lock, error) {
 	etcdConf := clientv3.Config{Endpoints: []string{endpoint}}
 	cli, err := clientv3.New(etcdConf)
 	if err != nil {
@@ -33,6 +36,7 @@ func New(cnf *config.Config, endpoint string) (iface.Lock, error) {
 		globalConf: cnf,
 		conf:       etcdConf,
 		cli:        cli,
+		retries:    retries,
 	}
 
 	return &lock, nil
@@ -40,13 +44,46 @@ func New(cnf *config.Config, endpoint string) (iface.Lock, error) {
 }
 
 func (r etcdLock) LockWithRetries(key string, unixTsToExpireNs int64) error {
-	err := errors.New("LockWithRetries not implemented")
-	log.ERROR.Print(err)
-	return err
+	for i := 0; i < r.retries; i++ {
+		err := r.Lock(key, unixTsToExpireNs)
+		if err == nil {
+			// 成功拿到锁，返回
+			return nil
+		}
+
+		log.INFO.Printf("acquired lock=%s failed, retries=%d", key, i)
+		time.Sleep(time.Millisecond * 100)
+	}
+	return ErrRedisLockFailed
+
 }
 
 func (r etcdLock) Lock(key string, unixTsToExpireNs int64) error {
-	err := errors.New("Lock not implemented")
-	log.ERROR.Print(err)
-	return err
+	now := time.Now().UnixNano()
+	ttl := time.Duration(unixTsToExpireNs + 1 - now)
+	lease := clientv3.NewLease(r.cli)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
+	defer cancel()
+
+	resp, err := lease.Grant(ctx, int64(ttl.Seconds()))
+	if err != nil {
+		return err
+	}
+
+	// 创建一个新的session
+	s, err := concurrency.NewSession(r.cli, concurrency.WithLease(resp.ID))
+	if err != nil {
+		return err
+	}
+
+	key = fmt.Sprintf("/machinery/v2/lock/%s", strings.TrimRight(key, "/"))
+	m := concurrency.NewMutex(s, key)
+
+	if err := m.Lock(ctx); err != nil {
+		return err
+	}
+
+	log.INFO.Printf("acquired lock=%s, duration=%s", key, ttl)
+	return nil
 }