finish task

backends/etcd/etcd.go

@@ -1,9 +1,12 @@
 package etcd
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
+	"log"
+	"sync"
 	"time"
 
 	"github.com/RichardKnop/machinery/v2/backends/iface"
@@ -11,6 +14,7 @@ import (
 	"github.com/RichardKnop/machinery/v2/config"
 	"github.com/RichardKnop/machinery/v2/tasks"
 	clientv3 "go.etcd.io/etcd/client/v3"
+	"golang.org/x/sync/errgroup"
 )
 
 type etcdBackend struct {
@@ -86,7 +90,24 @@ func (b *etcdBackend) GroupCompleted(groupUUID string, groupTaskCount int) (bool
 }
 
 func (b *etcdBackend) getGroupMeta(groupUUID string) (*tasks.GroupMeta, error) {
-	return nil, nil
+	key := fmt.Sprintf("/machinery/v2/backend/%s", groupUUID)
+	resp, err := b.client.Get(b.ctx, key)
+	if err != nil {
+		return nil, err
+	}
+	if len(resp.Kvs) == 0 {
+		return nil, fmt.Errorf("task %s not exist", groupUUID)
+	}
+	kv := resp.Kvs[0]
+
+	meta := new(tasks.GroupMeta)
+
+	decoder := json.NewDecoder(bytes.NewReader(kv.Value))
+	decoder.UseNumber()
+	if err := decoder.Decode(meta); err != nil {
+		return nil, err
+	}
+	return meta, nil
 }
 
 func (b *etcdBackend) GroupTaskStates(groupUUID string, groupTaskCount int) ([]*tasks.TaskState, error) {
@@ -99,7 +120,7 @@ func (b *etcdBackend) GroupTaskStates(groupUUID string, groupTaskCount int) ([]*
 }
 
 func (b *etcdBackend) TriggerChord(groupUUID string) (bool, error) {
-	return false, nil
+	return false, fmt.Errorf("not support")
 }
 
 // Setting / getting task state
@@ -145,7 +166,28 @@ func (b *etcdBackend) SetStateFailure(signature *tasks.Signature, err string) er
 }
 
 func (b *etcdBackend) GetState(taskUUID string) (*tasks.TaskState, error) {
-	return nil, nil
+	return b.getState(b.ctx, taskUUID)
+}
+
+func (b *etcdBackend) getState(ctx context.Context, taskUUID string) (*tasks.TaskState, error) {
+	key := fmt.Sprintf("/machinery/v2/backend/%s", taskUUID)
+	resp, err := b.client.Get(b.ctx, key)
+	if err != nil {
+		return nil, err
+	}
+	if len(resp.Kvs) == 0 {
+		return nil, fmt.Errorf("task %s not exist", taskUUID)
+	}
+	kv := resp.Kvs[0]
+
+	state := new(tasks.TaskState)
+
+	decoder := json.NewDecoder(bytes.NewReader(kv.Value))
+	decoder.UseNumber()
+	if err := decoder.Decode(state); err != nil {
+		return nil, err
+	}
+	return state, nil
 }
 
 // Purging stored stored tasks states and group meta data
@@ -162,20 +204,58 @@ func (b *etcdBackend) mergeNewTaskState(newState *tasks.TaskState) {
 }
 
 func (b *etcdBackend) PurgeState(taskUUID string) error {
-	return nil
+	key := fmt.Sprintf("/machinery/v2/backend/%s", taskUUID)
+	_, err := b.client.KV.Delete(b.ctx, key)
+	return err
 }
+
 func (b *etcdBackend) PurgeGroupMeta(groupUUID string) error {
-	return nil
+	key := fmt.Sprintf("/machinery/v2/backend/%s", groupUUID)
+	_, err := b.client.KV.Delete(b.ctx, key)
+	return err
 }
 
 // getStates returns multiple task states
 func (b *etcdBackend) getStates(taskUUIDs ...string) ([]*tasks.TaskState, error) {
-	taskStates := make([]*tasks.TaskState, len(taskUUIDs))
+	eg, ctx := errgroup.WithContext(b.ctx)
+	eg.SetLimit(10)
+	taskStates := make([]*tasks.TaskState, 0, len(taskUUIDs))
+	var mtx sync.Mutex
+	for _, taskUUID := range taskUUIDs {
+		t := taskUUID
+		eg.Go(func() error {
+			state, err := b.getState(ctx, t)
+			if err != nil {
+				return err
+			}
+
+			mtx.Lock()
+			taskStates = append(taskStates, state)
+			mtx.Unlock()
+			return nil
+		})
+	}
+	if err := eg.Wait(); err != nil {
+		return nil, err
+	}
+
 	return taskStates, nil
 }
 
 // updateState saves current task state
 func (b *etcdBackend) updateState(taskState *tasks.TaskState) error {
+	encoded, err := json.Marshal(taskState)
+	if err != nil {
+		return err
+	}
+
+	key := fmt.Sprintf("/machinery/v2/backend/%s", taskState.TaskUUID)
+	_, err = b.client.Put(b.ctx, key, string(encoded))
+	if err != nil {
+		return err
+	}
+
+	log.Default().Printf("update taskstate %s %s, %s", taskState.TaskName, taskState.TaskUUID, encoded)
 	return nil
 }