fix(finality-grandpa): stop the voter without racing Start or leaking its reader

pkg/finality-grandpa/voter.go

@@ -16,13 +16,16 @@ import (
 type wakerChan[Item any] struct {
 	in    chan Item
 	out   chan Item
+	stop  chan struct{}
+	once  sync.Once
 	waker atomic.Pointer[waker]
 }
 
 func newWakerChan[Item any](in chan Item) *wakerChan[Item] {
 	wc := &wakerChan[Item]{
-		in:  in,
-		out: make(chan Item),
+		in:   in,
+		out:  make(chan Item),
+		stop: make(chan struct{}),
 	}
 	go wc.start()
 	return wc
@@ -33,14 +36,39 @@ func (wc *wakerChan[Item]) start() {
 	if wc.in == nil {
 		return
 	}
-	for item := range wc.in {
+	for {
+		var (
+			item Item
+			ok   bool
+		)
+		select {
+		case <-wc.stop:
+			return
+		case item, ok = <-wc.in:
+			if !ok {
+				return
+			}
+		}
 		if w := wc.waker.Load(); w != nil {
 			w.wake()
 		}
-		wc.out <- item
+		// Also selects on stop: out is unbuffered, so a forwarder parked here with
+		// no consumer left would otherwise never return.
+		select {
+		case wc.out <- item:
+		case <-wc.stop:
+			return
+		}
 	}
 }
 
+// close releases the forwarding goroutine. The input channel belongs to whoever
+// constructed the wakerChan, so when that channel is not closed this is the only
+// way to stop reading from it. Idempotent.
+func (wc *wakerChan[Item]) close() {
+	wc.once.Do(func() { close(wc.stop) })
+}
+
 func (wc *wakerChan[Item]) setWaker(waker *waker) {
 	wc.waker.Store(waker)
 }
@@ -538,8 +566,25 @@ type Voter[Hash constraints.Ordered, Number constraints.Unsigned, Signature comp
 	stopTimeout time.Duration
 	stopChan    chan struct{}
 	wg          sync.WaitGroup
+	// stopOnce keeps the teardown to a single run. Stop closes channels, so a
+	// second call used to panic; an owner that both supervises the voter and
+	// shuts down can reach it twice.
+	stopOnce sync.Once
+	stopErr  error
+	// runState claims the single wg token NewVoter takes out; whichever of Start
+	// and Stop reaches it first owns releasing it. Taking the token in Start
+	// instead raced Stop's Wait, and a Stop that won the race would tear the voter
+	// down — closing channels it still polls — while Start was coming up.
+	runState atomic.Int32
 }
 
+// Voter lifecycle states for Voter.runState.
+const (
+	voterFresh int32 = iota
+	voterRunning
+	voterStopped
+)
+
 // NewVoter creates a new `Voter` tracker with given round number and base block.
 //
 // Provide data about the last completed round. If there is no
@@ -597,7 +642,7 @@ func NewVoter[Hash constraints.Ordered, Number constraints.Unsigned, Signature c
 		bestRound:  bestRound,
 		pastRounds: *pastRounds,
 	}
-	return &Voter[Hash, Number, Signature, ID]{
+	v := &Voter[Hash, Number, Signature, ID]{
 		env:                    env,
 		voters:                 voters,
 		inner:                  inner,
@@ -609,6 +654,9 @@ func NewVoter[Hash constraints.Ordered, Number constraints.Unsigned, Signature c
 		stopChan:               make(chan struct{}),
 		stopTimeout:            30 * time.Second,
 	}
+	// Held until Start returns, or until Stop claims it because Start never ran.
+	v.wg.Add(1)
+	return v
 }
 
 func (v *Voter[Hash, Number, Signature, ID]) pruneBackgroundRounds(waker *waker) error {
@@ -895,7 +943,9 @@ func (v *Voter[Hash, Number, Signature, ID]) setLastFinalizedNumber(finalizedNum
 }
 
 func (v *Voter[Hash, Number, Signature, ID]) Start() error { //skipcq: RVV-B0001
-	v.wg.Add(1)
+	if !v.runState.CompareAndSwap(voterFresh, voterRunning) {
+		return fmt.Errorf("voter was stopped before it started")
+	}
 	defer v.wg.Done()
 	waker := newWaker()
 	for {
@@ -914,9 +964,21 @@ func (v *Voter[Hash, Number, Signature, ID]) Start() error { //skipcq: RVV-B0001
 	}
 }
 
+// Stop tears the voter down and waits for its loop to return. It is idempotent:
+// later calls block until the first has finished and return the same result.
 func (v *Voter[Hash, Number, Signature, ID]) Stop() error {
+	v.stopOnce.Do(func() { v.stopErr = v.stop() })
+	return v.stopErr
+}
+
+func (v *Voter[Hash, Number, Signature, ID]) stop() error {
 	close(v.stopChan)
 	v.globalOut.Close()
+	// Start never ran and now never will, so release its token or Wait would block
+	// for the whole stop timeout.
+	if v.runState.CompareAndSwap(voterFresh, voterStopped) {
+		v.wg.Done()
+	}
 	timeout := time.NewTimer(v.stopTimeout)
 	wgDone := make(chan struct{})
 	go func() {
@@ -930,6 +992,10 @@ func (v *Voter[Hash, Number, Signature, ID]) Stop() error {
 	}
 
 	close(v.finalizedNotifications.in)
+	// globalIn belongs to the caller and is not ours to close, so the forwarder
+	// reading it has to be released explicitly or it outlives the voter — still
+	// taking items off a channel the caller may reuse for the next voter.
+	v.globalIn.close()
 	switch state := v.inner.bestRound.state.(type) {
 	case statePrecommitted:
 	case statePrevoted[timerI]:

pkg/finality-grandpa/voter_test.go

@@ -728,3 +728,132 @@ func TestBuffered(t *testing.T) {
 	run.Store(false)
 	wg.Wait()
 }
+
+// newLifecycleVoter builds a voter over its own inbound channel, for tests about
+// starting and stopping rather than about voting.
+func newLifecycleVoter(
+	t *testing.T,
+	network *Network,
+	globalIn chan GlobalInItem[string, uint32, Signature, ID],
+) *Voter[string, uint32, Signature, ID] {
+	t.Helper()
+	var localID ID = 5
+	voters := NewVoterSet([]IDWeight[ID]{{localID, 100}})
+
+	env := newEnvironment(network, localID)
+	var lastFinalized HashNumber[string, uint32]
+	env.WithChain(func(chain *dummyChain) {
+		chain.PushBlocks(GenesisHash, []string{"A", "B", "C"})
+		lastFinalized.Hash, lastFinalized.Number = chain.LastFinalized()
+	})
+
+	return NewVoter[string, uint32, Signature, ID](
+		&env,
+		*voters,
+		globalIn,
+		func(CommunicationOut[string, uint32, Signature, ID]) error { return nil },
+		0,
+		nil,
+		lastFinalized,
+		lastFinalized,
+	)
+}
+
+// Start runs in its own goroutine — it blocks — so an owner that shuts down right
+// after starting up can reach Stop first. The WaitGroup token is taken in NewVoter
+// so that ordering is safe: taking it in Start raced Stop's Wait, and a Stop that
+// won could close channels the voter was about to poll.
+func TestVoter_StartAndStopRace(t *testing.T) {
+	network := NewNetwork()
+	defer network.Stop()
+
+	for i := 0; i < 50; i++ {
+		voter := newLifecycleVoter(t, network, make(chan GlobalInItem[string, uint32, Signature, ID]))
+		voter.stopTimeout = 5 * time.Second
+
+		var wg sync.WaitGroup
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			_ = voter.Start()
+		}()
+		assert.NoError(t, voter.Stop())
+		wg.Wait()
+	}
+}
+
+// A voter that was stopped before it ever started must decline to start, and Stop
+// must not sit out its whole timeout waiting for a Start that will never come.
+func TestVoter_StopBeforeStart(t *testing.T) {
+	network := NewNetwork()
+	defer network.Stop()
+
+	voter := newLifecycleVoter(t, network, make(chan GlobalInItem[string, uint32, Signature, ID]))
+	voter.stopTimeout = 10 * time.Second
+
+	start := time.Now()
+	assert.NoError(t, voter.Stop())
+	assert.Less(t, time.Since(start), time.Second, "Stop waited on a Start that never ran")
+	assert.Error(t, voter.Start(), "a stopped voter must not start")
+}
+
+// globalIn belongs to the caller, so Stop cannot close it — but the forwarder
+// reading it has to go, or it outlives the voter and keeps taking items off a
+// channel the caller may well hand to the next voter.
+func TestVoter_StopReleasesGlobalIn(t *testing.T) {
+	network := NewNetwork()
+	defer network.Stop()
+
+	globalIn := make(chan GlobalInItem[string, uint32, Signature, ID], 2)
+	voter := newLifecycleVoter(t, network, globalIn)
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		_ = voter.Start()
+	}()
+	time.Sleep(50 * time.Millisecond)
+	assert.NoError(t, voter.Stop())
+	<-done
+
+	globalIn <- GlobalInItem[string, uint32, Signature, ID]{}
+	globalIn <- GlobalInItem[string, uint32, Signature, ID]{}
+	time.Sleep(200 * time.Millisecond)
+	assert.Len(t, globalIn, 2, "something is still reading globalIn after Stop")
+}
+
+// Stop closes channels, so it has to be single-shot: an owner that both supervises
+// the voter and shuts the node down can reach it twice, and the second call used to
+// panic on a closed channel. Concurrent callers must all get the same answer.
+func TestVoter_StopIsIdempotent(t *testing.T) {
+	network := NewNetwork()
+	defer network.Stop()
+
+	voter := newLifecycleVoter(t, network, make(chan GlobalInItem[string, uint32, Signature, ID]))
+	voter.stopTimeout = 5 * time.Second
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		_ = voter.Start()
+	}()
+	time.Sleep(50 * time.Millisecond)
+
+	const callers = 4
+	errs := make([]error, callers)
+	var wg sync.WaitGroup
+	for i := 0; i < callers; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			errs[i] = voter.Stop()
+		}(i)
+	}
+	wg.Wait()
+	<-done
+
+	for i, err := range errs {
+		assert.Equal(t, errs[0], err, "Stop caller %d saw a different result", i)
+	}
+	assert.NoError(t, voter.Stop())
+}