package worker
import (
"context"
"fmt"
"go.etcd.io/etcd/clientv3"
"time"
)
type JobLock struct {
Client *clientv3.Client
Kv clientv3.KV
Lease clientv3.Lease
JobName string
CancelFunc context.CancelFunc
LeaseId clientv3.LeaseID
IsLocked bool
}
var(
G_jobLock *JobLock
)
func InitJobLock() (err error) {
var (
config clientv3.Config
client *clientv3.Client
kv clientv3.KV
lease clientv3.Lease
)
config = clientv3.Config{
Endpoints:[]string{"127.0.0.1:2379"},
DialTimeout: 5000 * time.Millisecond,
}
if client, err = clientv3.New(config); err != nil {
fmt.Println(err)
}
kv = clientv3.NewKV(client)
lease = clientv3.NewLease(client)
G_jobLock = &JobLock{
Client:client,
Kv:kv,
Lease:lease,
}
return
}
func (jobLock *JobLock) TryLock(jobName string) (err error){
var(
ctx context.Context
cancelFunc context.CancelFunc
leaseGrantRes *clientv3.LeaseGrantResponse
leaseId clientv3.LeaseID
KeepResChan <- chan *clientv3.LeaseKeepAliveResponse
txn clientv3.Txn
lockKey string
txnRes *clientv3.TxnResponse
)
ctx, cancelFunc = context.WithCancel(context.TODO())
leaseGrantRes, err = jobLock.Lease.Grant(ctx, 5)
leaseId = leaseGrantRes.ID
if KeepResChan, err = jobLock.Lease.KeepAlive(ctx, leaseId); err != nil{
fmt.Println(err)
goto FAIL
}
go func() {
var keepRes *clientv3.LeaseKeepAliveResponse
for{
select {
case keepRes = <- KeepResChan:
if keepRes == nil{
goto END
}
}
time.Sleep(1*time.Second)
}
END:
}()
txn = jobLock.Kv.Txn(context.TODO())
lockKey = common.JOB_LOCK_DIR + jobName
txn.If(clientv3.Compare(clientv3.CreateRevision(lockKey), "=", 0)).
Then(clientv3.OpPut(lockKey, "", clientv3.WithLease(leaseId))).
Else(clientv3.OpGet(lockKey))
if txnRes, err = txn.Commit(); err != nil{
fmt.Println("提交创建锁失败", err)
goto FAIL
}
if !txnRes.Succeeded{
err = "锁被占用"
goto FAIL
}
jobLock.CancelFunc = cancelFunc
jobLock.LeaseId = leaseId
jobLock.IsLocked = true
return
FAIL:
cancelFunc()
jobLock.Lease.Revoke(ctx, leaseId)
return
}
func (jobLock *JobLock)UnLock() {
if jobLock.IsLocked{
jobLock.CancelFunc()
jobLock.Lease.Revoke(context.TODO(), jobLock.LeaseId)
}
}