package etcd import ( "bytes" "context" "encoding/json" "fmt" "time" "github.com/RichardKnop/machinery/v2/brokers/errs" "github.com/RichardKnop/machinery/v2/log" "github.com/RichardKnop/machinery/v2/tasks" clientv3 "go.etcd.io/etcd/client/v3" ) type Delivery interface { Ack() Nack() Body() []byte Signature() *tasks.Signature } type deliver struct { ctx context.Context client *clientv3.Client signature *tasks.Signature value []byte key string node string aliveCancel func() } func NewDelivery(ctx context.Context, client *clientv3.Client, key string, node string) (Delivery, error) { d := &deliver{ ctx: ctx, client: client, key: key, node: node, } if err := d.assign(key, node); err != nil { return nil, err } return d, nil } func (d *deliver) assign(key string, node string) error { assignKey := fmt.Sprintf("%s/assign", key) ctx, cancel := context.WithTimeout(d.ctx, time.Second*5) defer cancel() grantResp, err := d.client.Grant(ctx, 30) if err != nil { return err } cmp := 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() if err != nil { return err } if !resp.Succeeded { return fmt.Errorf("key %s already assign", key) } getResp := resp.Responses[1].GetResponseRange() if len(getResp.Kvs) == 0 { return fmt.Errorf("have no task %s", key) } kv := getResp.Kvs[0] signature := new(tasks.Signature) decoder := json.NewDecoder(bytes.NewReader(kv.Value)) decoder.UseNumber() if err = decoder.Decode(signature); err != nil { return errs.NewErrCouldNotUnmarshalTaskSignature(kv.Value, err) } aliveCtx, aliveCancel := context.WithCancel(d.ctx) if _, err = d.client.KeepAlive(aliveCtx, grantResp.ID); err != nil { aliveCancel() return err } d.aliveCancel = aliveCancel d.signature = signature return nil } func (d *deliver) Ack() { defer d.aliveCancel() ctx, cancel := context.WithTimeout(d.ctx, time.Second*2) defer cancel() _, err := d.client.Delete(ctx, d.key, clientv3.WithPrefix()) if err != nil { log.ERROR.Printf("ack task %s err: %s", d.value, err) } } func (d *deliver) Nack() { defer d.aliveCancel() assignKey := fmt.Sprintf("%s/assign", d.key) ctx, cancel := context.WithTimeout(d.ctx, time.Second*2) defer cancel() _, err := d.client.Delete(ctx, assignKey) if err != nil { log.ERROR.Printf("nack task %s err: %s", d.value, err) } } func (d *deliver) Signature() *tasks.Signature { return d.signature } func (d *deliver) Body() []byte { return d.value }