Merge pull request #11 from keybase/joshblum/bugfix

bug fixes

Author: zoom-ua

pipeliner.go

@@ -7,12 +7,19 @@ import (
 
 // Pipeliner coordinates a flow of parallel requests, rate-limiting so that
 // only a fixed number are outstanding at any one given time.
+//
+// Once an error has been recorded (via CompleteOne, or via context
+// cancellation), it is sticky: every subsequent call to WaitForRoom or
+// Flush on the same instance returns that error. A Pipeliner is therefore
+// single-use per batch — construct a new one for each independent batch of
+// work rather than reusing an instance after an error or cancellation.
 type Pipeliner struct {
-	sync.RWMutex
-	window int
-	numOut int
-	ch     chan struct{}
-	err    error
+	mu      sync.RWMutex
+	window  int
+	numOut  int
+	ch      chan struct{}
+	err     error
+	pending sync.WaitGroup // counts reservations awaiting a CompleteOne call
 }
 
 // NewPipeliner makes a pipeliner with window size `w`.
@@ -24,21 +31,24 @@ func NewPipeliner(w int) *Pipeliner {
 }
 
 func (p *Pipeliner) getError() error {
-	p.RLock()
-	defer p.RUnlock()
+	p.mu.RLock()
+	defer p.mu.RUnlock()
 	return p.err
 }
 
-func (p *Pipeliner) hasRoom() bool {
-	p.RLock()
-	defer p.RUnlock()
-	return p.numOut < p.window
-}
-
-func (p *Pipeliner) launchOne() {
-	p.Lock()
-	defer p.Unlock()
-	p.numOut++
+// tryReserve atomically checks for room in the window and, if available,
+// reserves a slot. Checking and reserving under a single lock acquisition
+// prevents concurrent callers from both observing room and over-committing
+// past the window.
+func (p *Pipeliner) tryReserve() bool {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.numOut < p.window {
+		p.numOut++
+		p.pending.Add(1)
+		return true
+	}
+	return false
 }
 
 // WaitForRoom will block until there is room in the window to fire
@@ -50,42 +60,42 @@ func (p *Pipeliner) WaitForRoom(ctx context.Context) error {
 	for {
 		p.checkContextDone(ctx)
 		if err := p.getError(); err != nil {
+			p.drain()
 			return err
 		}
-		if p.hasRoom() {
-			break
+		if p.tryReserve() {
+			return nil
 		}
 		p.wait(ctx)
 	}
-	p.launchOne()
-	return nil
 }
 
 // CompleteOne should be called when a request is completed, to make
 // room for subsequent requests. Call it with an error if you want the
 // rest of the pipeline to be short-circuited. This is the error that
 // is returned from WaitForRoom.
 func (p *Pipeliner) CompleteOne(e error) {
+	defer p.pending.Done()
 	p.setError(e)
 	p.landOne()
 	p.ch <- struct{}{}
 }
 
 func (p *Pipeliner) landOne() {
-	p.Lock()
-	defer p.Unlock()
+	p.mu.Lock()
+	defer p.mu.Unlock()
 	p.numOut--
 }
 
 func (p *Pipeliner) hasOutstanding() bool {
-	p.RLock()
-	defer p.RUnlock()
+	p.mu.RLock()
+	defer p.mu.RUnlock()
 	return p.numOut > 0
 }
 
 func (p *Pipeliner) setError(e error) {
-	p.Lock()
-	defer p.Unlock()
+	p.mu.Lock()
+	defer p.mu.Unlock()
 	if e != nil && p.err == nil {
 		p.err = e
 	}
@@ -107,13 +117,35 @@ func (p *Pipeliner) wait(ctx context.Context) {
 	}
 }
 
+// drain blocks until every reserved-but-not-yet-completed request has called
+// CompleteOne, so WaitForRoom/Flush can never return an error while leaving
+// a goroutine blocked sending on ch. A "receive while numOut > 0" loop won't
+// do: numOut can hit zero before the last CompleteOne's send actually lands.
+// So instead we race a reader against pending.Wait, which only unblocks once
+// every reservation's CompleteOne call has truly returned.
+func (p *Pipeliner) drain() {
+	done := make(chan struct{})
+	go func() {
+		p.pending.Wait()
+		close(done)
+	}()
+	for {
+		select {
+		case <-p.ch:
+		case <-done:
+			return
+		}
+	}
+}
+
 // Flush any outstanding requests, blocking until the last completes.
 // Returns an error set by CompleteOne, or a context-based error
 // if any request was canceled mid-flight.
 func (p *Pipeliner) Flush(ctx context.Context) error {
 	for {
 		p.checkContextDone(ctx)
 		if err := p.getError(); err != nil {
+			p.drain()
 			return err
 		}
 		if !p.hasOutstanding() {

pipeliner_test.go

@@ -4,16 +4,20 @@ import (
 	"context"
 	"errors"
 	"math/rand"
+	"reflect"
+	"runtime"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
 )
 
 func TestPipeliner(t *testing.T) {
-	v, err := testPipeliner(false, false)
+	v, launched, err := testPipeliner(false, false)
 	require.NoError(t, err)
+	require.Equal(t, len(v), launched)
 	for i, e := range v {
 		require.Equal(t, i, e)
 	}
@@ -24,21 +28,25 @@ func TestPipelinerError(t *testing.T) {
 }
 
 func testPipelinerError(t *testing.T) {
-	v, err := testPipeliner(true, false)
+	v, launched, err := testPipeliner(true, false)
 	require.Error(t, err)
 	if err != nil {
 		require.Equal(t, err.Error(), "errored out")
 	}
-	for _, e := range v[28:] {
+	// Indices below `launched` were handed to a goroutine and may have
+	// written a value; how many of those race ahead of the error being
+	// noticed varies run to run. Only indices at or past `launched` were
+	// never launched at all, so they must be untouched.
+	for _, e := range v[launched:] {
 		require.Equal(t, 0, e)
 	}
 }
 
 func TestPipelinerCancel(t *testing.T) {
-	v, err := testPipeliner(false, true)
+	v, launched, err := testPipeliner(false, true)
 	require.Error(t, err)
 	require.Equal(t, err.Error(), "context canceled")
-	for _, e := range v[28:] {
+	for _, e := range v[launched:] {
 		require.Equal(t, 0, e)
 	}
 }
@@ -49,8 +57,121 @@ func TestPipelinerErrorStress(t *testing.T) {
 	}
 }
 
-func testPipeliner(doError bool, doCancel bool) ([]int, error) {
-	v := make([]int, 100)
+// TestPipelinerNoGoroutineLeakOnError is a regression test: callers that
+// follow the documented pattern of returning immediately when WaitForRoom
+// errors out must not leak the goroutines that were still in flight.
+func TestPipelinerNoGoroutineLeakOnError(t *testing.T) {
+	runtime.GC()
+	before := runtime.NumGoroutine()
+
+	for range 200 {
+		testPipelinerError(t)
+	}
+
+	// Give any leaked goroutines a chance to be scheduled and block, so
+	// they show up in the goroutine count rather than racing this check.
+	time.Sleep(50 * time.Millisecond)
+	runtime.GC()
+	after := runtime.NumGoroutine()
+
+	require.LessOrEqual(t, after, before+2,
+		"goroutine count grew from %d to %d, suggesting a leak", before, after)
+}
+
+// TestTryReserveConcurrentNeverExceedsWindow hammers tryReserve from many
+// goroutines to guard the check-and-reserve race. It calls tryReserve
+// directly rather than through WaitForRoom: concurrent WaitForRoom callers
+// can race over CompleteOne's wakeups, a separate hazard this library
+// doesn't promise to handle.
+func TestTryReserveConcurrentNeverExceedsWindow(t *testing.T) {
+	const window = 4
+	const attempts = 5000
+
+	p := NewPipeliner(window)
+
+	var wg sync.WaitGroup
+	var active atomic.Int32
+	var maxActive atomic.Int32
+
+	for range attempts {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if !p.tryReserve() {
+				return
+			}
+			defer p.pending.Done()
+
+			n := active.Add(1)
+			for {
+				old := maxActive.Load()
+				if n <= old || maxActive.CompareAndSwap(old, n) {
+					break
+				}
+			}
+			active.Add(-1)
+			p.landOne()
+		}()
+	}
+
+	wg.Wait()
+	require.LessOrEqual(t, int(maxActive.Load()), window,
+		"observed %d reservations active at once, exceeding window %d", maxActive.Load(), window)
+}
+
+// TestPublicAPIDoesNotExposeLockMethods guards against re-embedding
+// sync.RWMutex (rather than holding it as an unexported field), which would
+// silently promote Lock/Unlock/RLock/RUnlock onto the public API again.
+func TestPublicAPIDoesNotExposeLockMethods(t *testing.T) {
+	typ := reflect.TypeFor[*Pipeliner]()
+	for _, name := range []string{"Lock", "Unlock", "RLock", "RUnlock"} {
+		_, ok := typ.MethodByName(name)
+		require.False(t, ok, "Pipeliner must not export %s", name)
+	}
+}
+
+// TestSetErrorFirstWins pins down the sticky-error invariant directly: once
+// an error is recorded, later errors must never overwrite it.
+func TestSetErrorFirstWins(t *testing.T) {
+	p := NewPipeliner(4)
+	first := errors.New("first")
+	second := errors.New("second")
+
+	p.setError(first)
+	p.setError(second)
+
+	require.Equal(t, first, p.getError())
+}
+
+// TestStickyErrorReturnedByWaitForRoomAndFlush verifies the sticky error
+// surfaces consistently through the public API: once set, WaitForRoom and
+// Flush keep returning it, even across further calls on the same instance.
+func TestStickyErrorReturnedByWaitForRoomAndFlush(t *testing.T) {
+	p := NewPipeliner(1)
+	ctx := context.Background()
+
+	require.NoError(t, p.WaitForRoom(ctx))
+	first := errors.New("first")
+	go p.CompleteOne(first)
+
+	err := p.Flush(ctx)
+	require.Equal(t, first, err)
+
+	// Reusing the instance after an error keeps returning the same error
+	// rather than doing further work or surfacing a different one.
+	err = p.WaitForRoom(ctx)
+	require.Equal(t, first, err)
+
+	err = p.Flush(ctx)
+	require.Equal(t, first, err)
+}
+
+// testPipeliner runs a batch of 100 simulated requests through a Pipeliner.
+// It returns the results, the number of requests actually handed to a
+// goroutine before WaitForRoom stopped returning nil, and any resulting
+// error.
+func testPipeliner(doError bool, doCancel bool) (v []int, launched int, err error) {
+	v = make([]int, 100)
 	var vlock sync.Mutex
 	pipeliner := NewPipeliner(4)
 	ctx := context.Background()
@@ -74,17 +195,16 @@ func testPipeliner(doError bool, doCancel bool) ([]int, error) {
 	}
 
 	for i := range v {
-		err := pipeliner.WaitForRoom(ctx)
-		if err != nil {
-			// Ensure any in-flight goroutines are complete before returning. Note that Flush does not actually wait if there is an error.
-			for pipeliner.hasOutstanding() {
-				<-pipeliner.ch
-			}
-			return v, err
+		werr := pipeliner.WaitForRoom(ctx)
+		if werr != nil {
+			// WaitForRoom drains any in-flight goroutines before returning
+			// an error, so it's safe to return immediately here.
+			return v, launched, werr
 		}
+		launched++
 		go func(i int) { pipeliner.CompleteOne(f(ctx, i)) }(i)
 	}
 
-	err := pipeliner.Flush(ctx)
-	return v, err
+	err = pipeliner.Flush(ctx)
+	return v, launched, err
 }