add list queue

brokers/etcd/etcd.go

@@ -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,11 +378,11 @@ func (b *etcdBroker) listWatchTasks(ctx context.Context, queue string) error {
 					continue
 				}
 
-				b.mtx.Lock()
 				b.assignMap[key] = false
-				b.mtx.Unlock()
 			}
 		}
+
+		b.mtx.Unlock()
 	}
 
 	return nil