chore: refine waitPrefixEvents (#5)

brokers/etcd/helper.go

@@ -2,7 +2,6 @@ package etcd
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"strconv"
 	"time"
@@ -12,13 +11,6 @@ import (
 	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
@@ -38,27 +30,22 @@ func deleteRevKey(ctx context.Context, kv clientv3.KV, key string, rev int64) er
 }
 
 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
-}
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, haveNoTaskErr
 
-func waitEvents(wc clientv3.WatchChan) *clientv3.Event {
-	for wresp := range wc {
-		for _, ev := range wresp.Events {
-			if ev.Type != mvccpb.PUT {
-				continue
+		case wresp := <-wc:
+			for _, ev := range wresp.Events {
+				if ev.Type != mvccpb.PUT {
+					continue
+				}
+				return ev, nil
 			}
-			return ev
 		}
 	}
-	return nil
 }
 
 func getItem(ctx context.Context, c *clientv3.Client, prefix string, opts []clientv3.OpOption, check checkFunc) (*mvccpb.KeyValue, error) {
@@ -78,10 +65,6 @@ func getItem(ctx context.Context, c *clientv3.Client, prefix string, opts []clie
 			return nil, err
 		}
 
-		if ev == nil {
-			return nil, haveNoTaskErr
-		}
-
 		if err := check(ev.Kv); err != nil {
 			return nil, err
 		}