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() { }