154 lines
3.4 KiB
Go
154 lines
3.4 KiB
Go
/*
|
|
* Tencent is pleased to support the open source community by making Blueking Container Service available.
|
|
* Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
|
|
* Licensed under the MIT License (the "License"); you may not use this file except
|
|
* in compliance with the License. You may obtain a copy of the License at
|
|
* http://opensource.org/licenses/MIT
|
|
* Unless required by applicable law or agreed to in writing, software distributed under
|
|
* the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
|
|
* either express or implied. See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
package etcd
|
|
|
|
import (
|
|
"context"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"github.com/RichardKnop/machinery/v2/log"
|
|
"go.etcd.io/etcd/api/v3/mvccpb"
|
|
clientv3 "go.etcd.io/etcd/client/v3"
|
|
)
|
|
|
|
const (
|
|
revokePrefix = "/machinery/v2/revoker/tasks"
|
|
)
|
|
|
|
type revokeSign struct {
|
|
taskID string
|
|
registerTime time.Time
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
}
|
|
|
|
// Revoke etcd revoker send sign
|
|
func (b *etcdRevoker) Revoke(ctx context.Context, taskID string) error {
|
|
key := revokePrefix + "/" + taskID
|
|
|
|
// 2分钟自动过期
|
|
lease, err := b.client.Grant(ctx, 120)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
_, err = b.client.Put(ctx, key, time.Now().Format(time.RFC3339), clientv3.WithLease(lease.ID))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// RevokeCtx etcd revoker ctx
|
|
func (b *etcdRevoker) RevokeCtx(taskID string) context.Context {
|
|
b.mtx.Lock()
|
|
defer b.mtx.Unlock()
|
|
|
|
sign, ok := b.revokeSignMap[taskID]
|
|
if ok {
|
|
sign.registerTime = time.Now()
|
|
return sign.ctx
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
sign = &revokeSign{
|
|
taskID: taskID,
|
|
registerTime: time.Now(),
|
|
ctx: ctx,
|
|
cancel: cancel,
|
|
}
|
|
b.revokeSignMap[taskID] = sign
|
|
|
|
return sign.ctx
|
|
}
|
|
|
|
func (b *etcdRevoker) tryRevoke(kv *mvccpb.KeyValue) {
|
|
key := string(kv.Key)
|
|
taskID := filepath.Base(key)
|
|
|
|
b.mtx.Lock()
|
|
sign, ok := b.revokeSignMap[taskID]
|
|
if ok {
|
|
sign.cancel()
|
|
delete(b.revokeSignMap, taskID)
|
|
}
|
|
b.mtx.Unlock()
|
|
|
|
if !ok {
|
|
return
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
|
|
defer cancel()
|
|
|
|
_, err := b.client.Delete(ctx, key)
|
|
if err != nil {
|
|
log.ERROR.Printf("revoke %s failed: %s", key, err)
|
|
}
|
|
}
|
|
|
|
func (b *etcdRevoker) listWatchRevoke(ctx context.Context) error {
|
|
// List
|
|
listCtx, listCancel := context.WithTimeout(ctx, time.Second*10)
|
|
defer listCancel()
|
|
|
|
resp, err := b.client.Get(listCtx, revokePrefix, clientv3.WithPrefix(), clientv3.WithKeysOnly())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
for _, kv := range resp.Kvs {
|
|
b.tryRevoke(kv)
|
|
}
|
|
|
|
// Watch
|
|
watchCtx, watchCancel := context.WithTimeout(ctx, time.Minute*60)
|
|
defer watchCancel()
|
|
|
|
watchOpts := []clientv3.OpOption{
|
|
clientv3.WithPrefix(),
|
|
clientv3.WithKeysOnly(),
|
|
clientv3.WithRev(resp.Header.Revision),
|
|
}
|
|
wc := b.client.Watch(watchCtx, revokePrefix, watchOpts...)
|
|
for wresp := range wc {
|
|
if wresp.Err() != nil {
|
|
return wresp.Err()
|
|
}
|
|
|
|
for _, ev := range wresp.Events {
|
|
if ev.Type != clientv3.EventTypePut {
|
|
continue
|
|
}
|
|
|
|
b.tryRevoke(ev.Kv)
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (b *etcdRevoker) cleanupRevokeSign() {
|
|
b.mtx.Lock()
|
|
defer b.mtx.Unlock()
|
|
|
|
for taskID, sign := range b.revokeSignMap {
|
|
if time.Since(sign.registerTime) > time.Hour*24*3 {
|
|
sign.cancel()
|
|
delete(b.revokeSignMap, taskID)
|
|
}
|
|
}
|
|
}
|