machinery-plugins/brokers/etcd/helper.go

package etcd

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"time"

	"github.com/RichardKnop/machinery/v2/log"
	"go.etcd.io/etcd/api/v3/mvccpb"
	clientv3 "go.etcd.io/etcd/client/v3"
)

var (
	ErrKeyExists      = errors.New("key already exists")
	ErrWaitMismatch   = errors.New("unexpected wait result")
	ErrTooManyClients = errors.New("too many clients")
	ErrNoWatcher      = errors.New("no watcher channel")
)

type checkFunc func(*mvccpb.KeyValue) error

// deleteRevKey deletes a key by revision, returning false if key is missing
func deleteRevKey(ctx context.Context, kv clientv3.KV, key string, rev int64) error {
	cmp := clientv3.Compare(clientv3.ModRevision(key), "=", rev)
	req := clientv3.OpDelete(key)
	txnresp, err := kv.Txn(ctx).If(cmp).Then(req).Commit()
	if err != nil {
		return err
	}

	if !txnresp.Succeeded {
		return fmt.Errorf("delete key failed")
	}

	return nil
}

func waitPrefixEvents(ctx context.Context, c *clientv3.Client, prefix string, rev int64) (*clientv3.Event, error) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	wc := c.Watch(ctx, prefix, clientv3.WithPrefix(), clientv3.WithRev(rev))
	if wc == nil {
		return nil, ErrNoWatcher
	}

	return waitEvents(wc), nil
}

func waitEvents(wc clientv3.WatchChan) *clientv3.Event {
	for wresp := range wc {
		for _, ev := range wresp.Events {
			if ev.Type != mvccpb.PUT {
				continue
			}
			return ev
		}
	}
	return nil
}

func getItem(ctx context.Context, c *clientv3.Client, prefix string, opts []clientv3.OpOption, check checkFunc) (*mvccpb.KeyValue, error) {
	resp, err := c.Get(ctx, prefix, opts...)
	if err != nil {
		return nil, err
	}

	// nothing yet; wait on elements
	if len(resp.Kvs) == 0 {
		ev, err := waitPrefixEvents(
			ctx,
			c,
			prefix,
			resp.Header.Revision)
		if err != nil {
			return nil, err
		}

		if ev == nil {
			return nil, haveNoTaskErr
		}

		if err := check(ev.Kv); err != nil {
			return nil, err
		}

		if err := deleteRevKey(ctx, c, string(ev.Kv.Key), ev.Kv.ModRevision); err != nil {
			return nil, err
		}

		return ev.Kv, nil
	}

	kv := resp.Kvs[0]
	if err := check(kv); err != nil {
		return nil, err
	}

	if err := deleteRevKey(ctx, c, string(kv.Key), kv.ModRevision); err != nil {
		return nil, err
	}

	return kv, nil
}

func getFirstItem(ctx context.Context, c *clientv3.Client, prefix string) ([]byte, error) {
	// 按创建时间排序, 只返回单个数据
	opts := clientv3.WithFirstRev()

	check := func(*mvccpb.KeyValue) error {
		return nil
	}

	kv, err := getItem(ctx, c, prefix, opts, check)
	if err != nil {
		return nil, err
	}

	return kv.Value, nil
}

func getFirstETAItem(ctx context.Context, c *clientv3.Client, prefix string) ([]byte, error) {
	// 按创建key排序, 只返回单个数据
	opts := clientv3.WithFirstKey()

	check := func(kv *mvccpb.KeyValue) error {
		k := string(kv.Key)
		if len(k) < 55 {
			log.ERROR.Print("not valid eta key: %s", k)
			return fmt.Errorf("not valid eta key: %s", k)
		}

		nsStr := k[36:55]
		ns, err := strconv.ParseInt(nsStr, 10, 64)
		if err != nil {
			return fmt.Errorf("parse key: %w", err)
		}

		t := time.Unix(0, ns)
		if t.After(time.Now()) {
			time.Sleep(time.Millisecond * 100)
			return haveNoTaskErr
		}

		return nil
	}

	kv, err := getItem(ctx, c, prefix, opts, check)
	if err != nil {
		return nil, err
	}

	return kv.Value, nil
}

func putItem() {

}