finish delivery

brokers/etcd/delivery.go

@@ -55,17 +55,21 @@ func (d *deliver) assign(key string, node string) error {
 		return err
 	}
 
-	cmp := clientv3.Compare(clientv3.CreateRevision(assignKey), "=", 0)
+	keyExist := clientv3.Compare(clientv3.CreateRevision(key), ">", 0)
+	assignNotExist := clientv3.Compare(clientv3.CreateRevision(assignKey), "=", 0)
 	putReq := clientv3.OpPut(assignKey, node, clientv3.WithLease(grantResp.ID))
 	getReq := clientv3.OpGet(key)
-	resp, err := d.client.Txn(ctx).If(cmp).Then(putReq, getReq).Commit()
+	resp, err := d.client.Txn(ctx).If(keyExist, assignNotExist).Then(putReq, getReq).Commit()
 	if err != nil {
 		return err
 	}
 	if !resp.Succeeded {
-		return fmt.Errorf("key %s already assign", key)
+		return fmt.Errorf("key %s not exist or already assign ", key)
 	}
 
+	if len(resp.Responses) < 2 {
+		return fmt.Errorf("have no resp %s", key)
+	}
 	getResp := resp.Responses[1].GetResponseRange()
 	if len(getResp.Kvs) == 0 {
 		return fmt.Errorf("have no task %s", key)

brokers/etcd/etcd.go

@@ -30,7 +30,7 @@ type etcdBroker struct {
 	ctx       context.Context
 	client    *clientv3.Client
 	wg        sync.WaitGroup
-	keyMap map[string]struct{}
+	assignMap map[string]bool
 	mtx       sync.Mutex
 }
 
@@ -52,7 +52,7 @@ func New(ctx context.Context, conf *config.Config) (iface.Broker, error) {
 		Broker:    common.NewBroker(conf),
 		ctx:       ctx,
 		client:    client,
-		keyMap: map[string]struct{}{},
+		assignMap: map[string]bool{},
 	}
 
 	return &broker, nil
@@ -296,7 +296,10 @@ func (b *etcdBroker) GetDelayedTasks() ([]*tasks.Signature, error) {
 }
 
 func (b *etcdBroker) nextTask(queue string, consumerTag string) (Delivery, error) {
-	for k := range b.keyMap {
+	for k, assigned := range b.assignMap {
+		if assigned {
+			continue
+		}
 		if !strings.Contains(k, queue) {
 			continue
 		}
@@ -313,6 +316,22 @@ func (b *etcdBroker) nextTask(queue string, consumerTag string) (Delivery, error
 	return b.nextTask(queue, consumerTag)
 }
 
+func (b *etcdBroker) setAssign(key string, assign bool) bool {
+	if !strings.Contains(key, "/assign") {
+		return false
+	}
+	k := strings.TrimSuffix(key, "/assign")
+
+	b.mtx.Lock()
+	defer b.mtx.Unlock()
+
+	if _, ok := b.assignMap[k]; ok {
+		b.assignMap[k] = assign
+	}
+
+	return true
+}
+
 func (b *etcdBroker) listWatchTasks(ctx context.Context, queue string) error {
 	keyPrefix := fmt.Sprintf("/machinery/v2/broker/pending_tasks/%s", queue)
 
@@ -325,11 +344,11 @@ func (b *etcdBroker) listWatchTasks(ctx context.Context, queue string) error {
 	}
 
 	for _, kv := range resp.Kvs {
-		if bytes.Contains(kv.Key, []byte("/assign")) {
-			delete(b.keyMap, string(kv.Key)[:23])
+		key := string(kv.Key)
+		if b.setAssign(key, true) {
 			continue
 		}
-		b.keyMap[string(kv.Key)] = struct{}{}
+		b.assignMap[key] = false
 	}
 
 	// Watch
@@ -338,20 +357,21 @@ func (b *etcdBroker) listWatchTasks(ctx context.Context, queue string) error {
 	wc := b.client.Watch(watchCtx, keyPrefix, clientv3.WithPrefix(), clientv3.WithKeysOnly(), clientv3.WithRev(resp.Header.Revision))
 	for wresp := range wc {
 		for _, ev := range wresp.Events {
+			key := string(ev.Kv.Key)
 			if ev.Type == clientv3.EventTypeDelete {
-				if bytes.Contains(ev.Kv.Key, []byte("/assign")) {
-					b.keyMap[string(ev.Kv.Key)] = struct{}{}
-				}
+				if b.setAssign(key, false) {
 					continue
 				}
 
+				delete(b.assignMap, key)
+			}
+
 			if ev.Type == clientv3.EventTypePut {
-				if bytes.Contains(ev.Kv.Key, []byte("/assign")) {
-					delete(b.keyMap, string(ev.Kv.Key)[:23])
+				if b.setAssign(key, true) {
 					continue
 				}
 
-				b.keyMap[string(ev.Kv.Key)] = struct{}{}
+				b.assignMap[key] = false
 			}
 		}
 	}