refine etcd

brokers/etcd/etcd.go

@@ -166,8 +166,7 @@ func (b *etcdBroker) StartConsuming(consumerTag string, concurrency int, taskPro
 // consume takes delivered messages from the channel and manages a worker pool
 // to process tasks concurrently
 func (b *etcdBroker) consume(deliveries <-chan Delivery, concurrency int, taskProcessor iface.TaskProcessor) error {
-	eg, ctx := errgroup.WithContext(context.Background())
-	eg.SetLimit(concurrency)
+	eg, ctx := errgroup.WithContext(b.ctx)
 
 	for i := 0; i < concurrency; i++ {
 		eg.Go(func() error {
@@ -194,24 +193,22 @@ func (b *etcdBroker) consume(deliveries <-chan Delivery, concurrency int, taskPr
 }
 
 // consumeOne processes a single message using TaskProcessor
-func (b *etcdBroker) consumeOne(d Delivery, taskProcessor iface.TaskProcessor) error {
+func (b *etcdBroker) consumeOne(delivery Delivery, taskProcessor iface.TaskProcessor) error {
 	// If the task is not registered, we requeue it,
 	// there might be different workers for processing specific tasks
-	if !b.IsTaskRegistered(d.Signature().Name) {
-		if d.Signature().IgnoreWhenTaskNotRegistered {
-			return nil
+	if !b.IsTaskRegistered(delivery.Signature().Name) {
+		log.INFO.Printf("Task not registered with this worker. Requeuing message: %s", delivery.Body())
+
+		if !delivery.Signature().IgnoreWhenTaskNotRegistered {
+			delivery.Nack()
 		}
-		log.INFO.Printf("Task not registered with this worker. Requeuing message: %s", d.Body())
-		d.Nack()
 		return nil
 	}
 
-	log.DEBUG.Printf("Received new message: %s", d.Body())
+	log.DEBUG.Printf("Received new message: %s", delivery.Body())
+	defer delivery.Ack()
 
-	err := taskProcessor.Process(d.Signature())
-	d.Ack()
-
-	return err
+	return taskProcessor.Process(delivery.Signature())
 }
 
 // StopConsuming 停止
@@ -254,12 +251,16 @@ func (b *etcdBroker) getTasks(ctx context.Context, key string) ([]*tasks.Signatu
 	}
 
 	result := make([]*tasks.Signature, 0, len(resp.Kvs))
-	for _, kvs := range resp.Kvs {
+	for _, kv := range resp.Kvs {
+		if strings.Contains(string(kv.Key), "/assign") {
+			continue
+		}
+
 		signature := new(tasks.Signature)
-		decoder := json.NewDecoder(bytes.NewReader(kvs.Value))
+		decoder := json.NewDecoder(bytes.NewReader(kv.Value))
 		decoder.UseNumber()
 		if err := decoder.Decode(signature); err != nil {
-			return nil, errs.NewErrCouldNotUnmarshalTaskSignature(kvs.Value, err)
+			return nil, errs.NewErrCouldNotUnmarshalTaskSignature(kv.Value, err)
 		}
 
 		result = append(result, signature)
@@ -418,7 +419,7 @@ func (b *etcdBroker) handleDelayedTask(ctx context.Context) error {
 		putReq := clientv3.OpPut(pendingKey, string(kv.Value))
 		c, err := b.client.Txn(b.ctx).If(cmp).Then(deleteReq, putReq).Commit()
 		if err != nil {
-			return err
+			return fmt.Errorf("handle delay task %s: %w", key, err)
 		}
 		if !c.Succeeded {
 			log.WARNING.Printf("handle delay task %s not success", key)