Optimization: refactor picker

common/picker/picker.go

@@ -1,22 +1,53 @@
 package picker
 
-import "context"
+import (
+	"context"
+	"sync"
+)
+
+// Picker provides synchronization, and Context cancelation
+// for groups of goroutines working on subtasks of a common task.
+// Inspired by errGroup
+type Picker struct {
+	cancel func()
+
+	wg sync.WaitGroup
+
+	once   sync.Once
+	result interface{}
+}
+
+// WithContext returns a new Picker and an associated Context derived from ctx.
+func WithContext(ctx context.Context) (*Picker, context.Context) {
+	ctx, cancel := context.WithCancel(ctx)
+	return &Picker{cancel: cancel}, ctx
+}
+
+// Wait blocks until all function calls from the Go method have returned,
+// then returns the first nil error result (if any) from them.
+func (p *Picker) Wait() interface{} {
+	p.wg.Wait()
+	if p.cancel != nil {
+		p.cancel()
+	}
+	return p.result
+}
+
+// Go calls the given function in a new goroutine.
+// The first call to return a nil error cancels the group; its result will be returned by Wait.
+func (p *Picker) Go(f func() (interface{}, error)) {
+	p.wg.Add(1)
 
-func SelectFast(ctx context.Context, in <-chan interface{}) <-chan interface{} {
-	out := make(chan interface{})
 	go func() {
-		select {
-		case p, open := <-in:
-			if open {
-				out <- p
-			}
-		case <-ctx.Done():
-		}
+		defer p.wg.Done()
 
-		close(out)
-		for range in {
+		if ret, err := f(); err == nil {
+			p.once.Do(func() {
+				p.result = ret
+				if p.cancel != nil {
+					p.cancel()
+				}
+			})
 		}
 	}()
-
-	return out
 }

common/picker/picker_test.go

@@ -6,39 +6,37 @@ import (
 	"time"
 )
 
-func sleepAndSend(delay int, in chan<- interface{}, input interface{}) {
-	time.Sleep(time.Millisecond * time.Duration(delay))
-	in <- input
-}
-
-func sleepAndClose(delay int, in chan interface{}) {
-	time.Sleep(time.Millisecond * time.Duration(delay))
-	close(in)
+func sleepAndSend(ctx context.Context, delay int, input interface{}) func() (interface{}, error) {
+	return func() (interface{}, error) {
+		timer := time.NewTimer(time.Millisecond * time.Duration(delay))
+		select {
+		case <-timer.C:
+			return input, nil
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
+	}
 }
 
 func TestPicker_Basic(t *testing.T) {
-	in := make(chan interface{})
-	fast := SelectFast(context.Background(), in)
-	go sleepAndSend(20, in, 1)
-	go sleepAndSend(30, in, 2)
-	go sleepAndClose(40, in)
+	picker, ctx := WithContext(context.Background())
+	picker.Go(sleepAndSend(ctx, 30, 2))
+	picker.Go(sleepAndSend(ctx, 20, 1))
 
-	number, exist := <-fast
-	if !exist || number != 1 {
-		t.Error("should recv 1", exist, number)
+	number := picker.Wait()
+	if number != nil && number.(int) != 1 {
+		t.Error("should recv 1", number)
 	}
 }
 
 func TestPicker_Timeout(t *testing.T) {
-	in := make(chan interface{})
 	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*5)
 	defer cancel()
-	fast := SelectFast(ctx, in)
-	go sleepAndSend(20, in, 1)
-	go sleepAndClose(30, in)
+	picker, ctx := WithContext(ctx)
+	picker.Go(sleepAndSend(ctx, 20, 1))
 
-	_, exist := <-fast
-	if exist {
-		t.Error("should recv false")
+	number := picker.Wait()
+	if number != nil {
+		t.Error("should recv nil")
 	}
 }