common
/
machinery-plugins
Private
1
Fork 0
Code Pull Requests Releases Activity

add list queue

Browse Source
Operations
Operations
Create branch
Create tag
Revert
Cherry-pick
main
git 2024-06-10 21:11:02 +08:00
parent 72a4a9731e
commit 3ea43cba66
Signed by: git
GPG Key ID: 3F65EFFA44207ADD
1 changed files with 9 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
brokers/etcd/etcd.go
View File

@ -5,7 +5,6 @@ import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"runtime"
"strconv"
@ -25,8 +24,6 @@ import (
"go.etcd.io/etcd/client/v3/concurrency"
)
var haveNoTaskErr = errors.New("have no task")
type etcdBroker struct {
common.Broker
ctx context.Context
@ -297,7 +294,7 @@ func (b *etcdBroker) GetDelayedTasks() ([]*tasks.Signature, error) {
func (b *etcdBroker) nextTask(queue string, consumerTag string) Delivery {
b.mtx.Lock()
assignMap := make(map[string]bool, len(b.assignMap))
for k, v := range assignMap {
for k, v := range b.assignMap {
assignMap[k] = v
}
b.mtx.Unlock()
@ -327,9 +324,6 @@ func (b *etcdBroker) setAssign(key string, assign bool) bool {
}
k := strings.TrimSuffix(key, "/assign")
b.mtx.Lock()
defer b.mtx.Unlock()
if _, ok := b.assignMap[k]; ok {
b.assignMap[k] = assign
}
@ -341,25 +335,26 @@ func (b *etcdBroker) listWatchTasks(ctx context.Context, queue string) error {
keyPrefix := fmt.Sprintf("/machinery/v2/broker/pending_tasks/%s", queue)
// List
listCtx, listCancel := context.WithTimeout(ctx, time.Second*5)
listCtx, listCancel := context.WithTimeout(ctx, time.Second*10)
defer listCancel()
resp, err := b.client.Get(listCtx, keyPrefix, clientv3.WithPrefix(), clientv3.WithKeysOnly())
if err != nil {
return err
}
b.mtx.Lock()
b.assignMap = map[string]bool{}
for _, kv := range resp.Kvs {
key := string(kv.Key)
if b.setAssign(key, true) {
continue
}
b.mtx.Lock()
b.assignMap[key] = false
b.mtx.Unlock()
}
b.mtx.Unlock()
// Watch
watchCtx, watchCancel := context.WithTimeout(ctx, time.Minute*60)
watchCtx, watchCancel := context.WithTimeout(ctx, time.Minute*10)
defer watchCancel()
wc := b.client.Watch(watchCtx, keyPrefix, clientv3.WithPrefix(), clientv3.WithKeysOnly(), clientv3.WithRev(resp.Header.Revision))
for wresp := range wc {
@ -367,6 +362,7 @@ func (b *etcdBroker) listWatchTasks(ctx context.Context, queue string) error {
return watchCtx.Err()
}
b.mtx.Lock()
for _, ev := range wresp.Events {
key := string(ev.Kv.Key)
if ev.Type == clientv3.EventTypeDelete {
@ -374,9 +370,7 @@ func (b *etcdBroker) listWatchTasks(ctx context.Context, queue string) error {
continue
}
b.mtx.Lock()
delete(b.assignMap, key)
b.mtx.Unlock()
}
if ev.Type == clientv3.EventTypePut {
@ -384,12 +378,12 @@ func (b *etcdBroker) listWatchTasks(ctx context.Context, queue string) error {
continue
}
b.mtx.Lock()
b.assignMap[key] = false
}
}
b.mtx.Unlock()
}
}
}
return nil
}