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