fix(singleflight): give waiting callers the real error of the run they waited on

singleflight.go

@@ -6,63 +6,80 @@ import (
 	"time"
 )
 
+// flight is the shared state of one fn execution. err is written exactly once,
+// before ch is closed, so any goroutine that observed the close may read err
+// without further synchronization.
+type flight struct {
+	err error
+	ch  chan struct{}
+}
+
 type call struct {
 	ts time.Time
-	ch chan struct{}
+	fl *flight
 	cn int
 	mu sync.Mutex
 }
 
+// Do runs fn, deduping concurrent callers: a caller that arrives while fn is
+// already running does not start a second run — it waits for the running one
+// and returns its error.
+//
+// The waiting caller must get the real error. It used to get nil even when fn
+// failed, and nil looks like success. sentinelClient.refreshRetry retries
+// refresh() until it returns nil, so a refreshRetry that waited on someone
+// else's failed refresh stopped retrying while the master was still
+// unresolved. clusterClient.pick reported the same failure as ErrNoSlot rather
+// than the refresh error behind it.
 func (c *call) Do(ctx context.Context, fn func() error) error {
 	c.mu.Lock()
 	c.cn++
-	ch := c.ch
-	if ch != nil {
+	fl := c.fl
+	if fl != nil {
 		c.mu.Unlock()
 		if ctxCh := ctx.Done(); ctxCh != nil {
 			select {
-			case <-ch:
+			case <-fl.ch:
 			case <-ctxCh:
 				return ctx.Err()
 			}
 		} else {
-			<-ch
+			<-fl.ch
 		}
-		return nil
+		return fl.err
 	}
-	ch = make(chan struct{})
-	c.ch = ch
+	fl = &flight{ch: make(chan struct{})}
+	c.fl = fl
 	c.mu.Unlock()
-	return c.do(ch, fn)
+	return c.do(fl, fn)
 }
 
 // DelayDo sleeps for delay then runs fn, deduping concurrent callers via singleflight.
 func (c *call) DelayDo(delay time.Duration, fn func() error) {
 	c.mu.Lock()
-	ch := c.ch
-	if ch != nil {
+	if c.fl != nil {
 		c.mu.Unlock()
 		return
 	}
-	ch = make(chan struct{})
-	c.ch = ch
+	fl := &flight{ch: make(chan struct{})}
+	c.fl = fl
 	c.cn++
 	c.mu.Unlock()
-	go func(delay time.Duration, ch chan struct{}, fn func() error) {
+	go func(delay time.Duration, fl *flight, fn func() error) {
 		time.Sleep(delay)
-		c.do(ch, fn)
-	}(delay, ch, fn)
+		c.do(fl, fn)
+	}(delay, fl, fn)
 }
 
-func (c *call) do(ch chan struct{}, fn func() error) (err error) {
-	err = fn()
+func (c *call) do(fl *flight, fn func() error) error {
+	fl.err = fn()
 	c.mu.Lock()
-	c.ch = nil
+	c.fl = nil
 	c.cn = 0
 	c.ts = time.Now()
 	c.mu.Unlock()
-	close(ch)
-	return
+	close(fl.ch)
+	return fl.err
 }
 
 func (c *call) suppressing() int {

singleflight_test.go

@@ -44,8 +44,219 @@ func TestSingleFlight(t *testing.T) {
 		t.Fatalf("singleflight should suppress all concurrent calls, got: %v", v)
 	}
 
-	if atomic.LoadInt64(&err) != 1 {
-		t.Fatalf("singleflight should that one call get the return value")
+	// Every caller must see the error: the one that ran fn and everyone who
+	// waited on it. Waiters used to get nil, and nil looks like success to
+	// code that retries until an operation succeeds.
+	if v := atomic.LoadInt64(&err); v != 1000 {
+		t.Fatalf("all callers should get the error of the run they waited on, got: %v", v)
+	}
+}
+
+// TestSingleFlightJoinerReceivesFlightError: a caller that waits for an
+// already-running fn must get the error that fn actually returned.
+//
+// It used to get nil even when fn failed. nil looks like success, so code that
+// retries an operation until it succeeds stopped retrying after a run that
+// failed: sentinelClient.refreshRetry loops until refresh() returns nil, so
+// joining someone else's failed refresh ended the retry loop with the master
+// still unresolved. The test also checks the other direction: the error of a
+// finished run must not leak into the next run.
+func TestSingleFlightJoinerReceivesFlightError(t *testing.T) {
+	defer ShouldNotLeak(SetupLeakDetection())
+	block := make(chan struct{})
+	flightErr := errors.New("flight failed")
+	sg := call{}
+
+	initiatorDone := make(chan error, 1)
+	go func() {
+		initiatorDone <- sg.Do(context.Background(), func() error {
+			<-block
+			return flightErr
+		})
+	}()
+	for sg.suppressing() != 1 {
+		runtime.Gosched()
+	}
+
+	joinerDone := make(chan error, 1)
+	go func() {
+		joinerDone <- sg.Do(context.Background(), func() error {
+			t.Error("joiner fn must not run")
+			return nil
+		})
+	}()
+	for sg.suppressing() != 2 {
+		runtime.Gosched()
+	}
+
+	close(block)
+	if err := <-initiatorDone; err != flightErr {
+		t.Fatalf("initiator: unexpected err %v", err)
+	}
+	if err := <-joinerDone; err != flightErr {
+		t.Fatalf("joiner: unexpected err %v", err)
+	}
+
+	// A caller arriving after the flight completed starts a fresh flight and
+	// gets its own result, not the previous flight's error.
+	if err := sg.Do(context.Background(), func() error { return nil }); err != nil {
+		t.Fatalf("fresh flight: unexpected err %v", err)
+	}
+}
+
+// TestSingleFlightJoinerKeepsItsOwnFlightError: a waiter must get the error of
+// the flight it waited on even when the next flight overlaps with it.
+//
+// The call struct is reused. do() clears the flight before it closes the
+// channel, so the next flight can start while the previous waiters are still
+// waking up. Anything per-flight kept on call itself is overwritten in that
+// window, and the waiters then read the next flight's result instead of their
+// own — nil, which is the failure this fix is about. Keeping err on the flight
+// avoids it, and also removes the need to synchronize the read: the write
+// happens before close(ch), and the read happens after <-ch.
+func TestSingleFlightJoinerKeepsItsOwnFlightError(t *testing.T) {
+	defer ShouldNotLeak(SetupLeakDetection())
+	flightErr := errors.New("flight failed")
+	const joiners = 50
+
+	for range 200 {
+		var (
+			sg    = call{}
+			block = make(chan struct{})
+			errs  = make([]error, joiners)
+			done  int64
+		)
+
+		go func() {
+			sg.Do(context.Background(), func() error {
+				<-block
+				return flightErr
+			})
+		}()
+		for sg.suppressing() != 1 {
+			runtime.Gosched()
+		}
+
+		for j := range joiners {
+			go func(j int) {
+				errs[j] = sg.Do(context.Background(), func() error { return nil })
+				atomic.AddInt64(&done, 1)
+			}(j)
+		}
+		for sg.suppressing() != joiners+1 {
+			runtime.Gosched()
+		}
+
+		// The next flight starts as soon as the current one clears its
+		// counter, which is before the waiters are released.
+		next := make(chan struct{})
+		go func() {
+			defer close(next)
+			for sg.suppressing() != 0 {
+				runtime.Gosched()
+			}
+			sg.Do(context.Background(), func() error { return nil })
+		}()
+
+		close(block)
+		for atomic.LoadInt64(&done) != joiners {
+			runtime.Gosched()
+		}
+		<-next
+
+		for j := range errs {
+			if errs[j] != flightErr {
+				t.Fatalf("joiner %v got %v, want %v", j, errs[j], flightErr)
+			}
+		}
+	}
+}
+
+// TestSingleFlightCancellableJoinerAtCompletion drives callers into the moment
+// a flight completes, half of them able to be cancelled and half not, since the
+// two take different branches of Do.
+//
+// do() clears c.fl before it releases the waiters, so a caller arriving in that
+// window finds no flight and starts its own instead of joining. Whichever side
+// of it a caller lands on, it must come back with the error of the flight it
+// waited on, and none may stay parked.
+func TestSingleFlightCancellableJoinerAtCompletion(t *testing.T) {
+	defer ShouldNotLeak(SetupLeakDetection())
+	flightErr := errors.New("flight failed")
+	const callers = 20
+
+	for range 200 {
+		var (
+			sg      call
+			done    int64
+			results = make([]error, callers)
+		)
+		for j := range callers {
+			go func(j int) {
+				ctx := context.Background()
+				if j%2 == 0 { // half wait through the ctx branch of Do
+					c, cancel := context.WithCancel(ctx)
+					defer cancel()
+					ctx = c
+				}
+				results[j] = sg.Do(ctx, func() error {
+					runtime.Gosched()
+					return flightErr
+				})
+				atomic.AddInt64(&done, 1)
+			}(j)
+		}
+		for atomic.LoadInt64(&done) != callers {
+			runtime.Gosched()
+		}
+		for j := range results {
+			if results[j] != flightErr {
+				t.Fatalf("caller %v got %v, want %v", j, results[j], flightErr)
+			}
+		}
+	}
+}
+
+// TestSingleFlightCancelledJoinerAtCompletion: the same window, with the
+// cancellable waiters cancelled around the time the flight ends. Each must come
+// back with either its flight's error or the cancellation, never nil and never
+// parked, and the runner must still release the waiters that stayed.
+func TestSingleFlightCancelledJoinerAtCompletion(t *testing.T) {
+	defer ShouldNotLeak(SetupLeakDetection())
+	flightErr := errors.New("flight failed")
+	const callers = 20
+
+	for range 200 {
+		var (
+			sg      call
+			done    int64
+			release = make(chan struct{})
+			results = make([]error, callers)
+		)
+		for j := range callers {
+			go func(j int) {
+				ctx, cancel := context.WithCancel(context.Background())
+				if j%2 == 0 {
+					defer cancel()
+				} else {
+					cancel() // already cancelled on arrival
+				}
+				results[j] = sg.Do(ctx, func() error {
+					<-release
+					return flightErr
+				})
+				atomic.AddInt64(&done, 1)
+			}(j)
+		}
+		close(release)
+		for atomic.LoadInt64(&done) != callers {
+			runtime.Gosched()
+		}
+		for j := range results {
+			if results[j] != flightErr && results[j] != context.Canceled {
+				t.Fatalf("caller %v got %v", j, results[j])
+			}
+		}
 	}
 }