refactor(pkg/finality-grandpa): make closing globalIn the voter's shutdown signal (#4852)

Author: timwu20

internal/client/consensus/grandpa/grandpa.go

@@ -436,7 +436,7 @@ type voterWork[
 	E runtime.Extrinsic,
 ] struct {
 	voter            *grandpa.Voter[H, N, primitives.AuthoritySignature, primitives.AuthorityID]
-	voterErrChan     <-chan error
+	voterDone        <-chan error
 	sharedVoterState *SharedVoterState[primitives.AuthorityID]
 	env              *environment[H, N, Hasher, Header, E]
 	voterCommandsRx  <-chan voterCommand
@@ -554,14 +554,9 @@ func (vw *voterWork[H, N, Hasher, Header, E]) rebuildVoter() {
 		// Repoint shared_voter_state so that the RPC endpoint can query the state
 		vw.sharedVoterState.reset(voter.VoterState())
 
+		// NewVoter runs the voter; Done yields why it stopped, once it has.
 		vw.voter = voter
-		errChan := make(chan error)
-		go func() {
-			err := voter.Start()
-			errChan <- err
-			close(errChan)
-		}()
-		vw.voterErrChan = errChan
+		vw.voterDone = voter.Done()
 	case voterSetStatePaused[H, N]:
 	default:
 		panic("unreachable")
@@ -651,9 +646,9 @@ func (vw *voterWork[H, N, Hasher, Header, E]) handleVoterCommand(command voterCo
 
 func (vw *voterWork[H, N, Hasher, Header, E]) poll() error {
 	select {
-	case err := <-vw.voterErrChan:
+	case err := <-vw.voterDone:
 		if err == nil {
-			// voters don't conclude naturally
+			// nothing here closes globalIn, so the voter has no orderly way to stop
 			return fmt.Errorf("consensus-grandpa inner voter has concluded: %w", ErrSafety)
 		}
 		vc, isVoterCommand := err.(voterCommand)

pkg/finality-grandpa/bridge_state.go

@@ -17,13 +17,16 @@ func newWaker() *waker {
 }
 
 func (w *waker) wake() {
+	// Read under the lock and hand the value to the goroutine, which outlives the
+	// lock and would otherwise race register's write.
 	w.RLock()
-	defer w.RUnlock()
-	if w.wakeCh == nil {
+	ch := w.wakeCh
+	w.RUnlock()
+	if ch == nil {
 		return
 	}
 	go func() {
-		w.wakeCh <- struct{}{}
+		ch <- struct{}{}
 	}()
 }
 

pkg/finality-grandpa/environment_test.go

@@ -26,7 +26,11 @@ type environment struct {
 	network                  *Network
 	listeners                []chan listenerItem
 	lastCompleteAndConcluded [2]uint64
-	mtx                      sync.Mutex
+	// roundIn holds the inbound channels handed to the voter, per round, so they
+	// can be closed once the round concludes. RoundData is called more than once
+	// for a round number, hence a slice.
+	roundIn map[uint64][]chan SignedMessageError[string, uint32, Signature, ID]
+	mtx     sync.Mutex
 
 	concludedCalled chan struct{}
 }
@@ -36,6 +40,7 @@ func newEnvironment(network *Network, localID ID) environment {
 		chain:           newDummyChain(),
 		localID:         localID,
 		network:         network,
+		roundIn:         make(map[uint64][]chan SignedMessageError[string, uint32, Signature, ID]),
 		concludedCalled: make(chan struct{}),
 	}
 }
@@ -84,6 +89,12 @@ func (e *environment) RoundData(
 	outgoing := make(Output[string, uint32])
 	incoming := e.network.MakeRoundComms(round, e.localID, outgoing)
 
+	// Remember it so Concluded can close it: the voter reads this channel through
+	// a forwarding goroutine that ends only when the channel does.
+	e.mtx.Lock()
+	e.roundIn[round] = append(e.roundIn[round], incoming)
+	e.mtx.Unlock()
+
 	var outgoingFunc = func(m Message[string, uint32]) error {
 		outgoing <- m
 		return nil
@@ -123,8 +134,16 @@ func (e *environment) Concluded(
 	_ HistoricalVotes[string, uint32, Signature, ID],
 ) error {
 	e.mtx.Lock()
-	defer e.mtx.Unlock()
 	e.lastCompleteAndConcluded[1] = round
+	incoming := e.roundIn[round]
+	delete(e.roundIn, round)
+	e.mtx.Unlock()
+
+	// The round is over, so release the inbound channels handed out for it.
+	for _, in := range incoming {
+		e.network.StopRoundComms(round, in)
+	}
+
 	go func() {
 		e.concludedCalled <- struct{}{}
 	}()
@@ -259,13 +278,29 @@ func (bm *BroadcastNetwork[M, N]) AddNode(f func(N) M, out chan N) (in chan M) {
 func (bm *BroadcastNetwork[M, N]) route() {
 	defer bm.routeWG.Done()
 	for msg := range bm.receiver {
+		// Under the lock: RemoveNode closes a node's channel, and closing one a
+		// producer is about to send on panics. Senders are buffered, so holding it
+		// across the delivery does not block.
 		bm.mu.Lock()
 		bm.history = append(bm.history, msg)
-		senders := append([]chan M(nil), bm.senders...)
-		bm.mu.Unlock()
-		for _, sender := range senders {
+		for _, sender := range bm.senders {
 			sender <- msg
 		}
+		bm.mu.Unlock()
+	}
+}
+
+// RemoveNode deregisters a node's inbound channel and closes it, shutting down
+// the voter reading it. Held under bm.mu so it cannot race a delivery in route.
+func (bm *BroadcastNetwork[M, N]) RemoveNode(in chan M) {
+	bm.mu.Lock()
+	defer bm.mu.Unlock()
+	for i, sender := range bm.senders {
+		if sender == in {
+			bm.senders = append(bm.senders[:i], bm.senders[i+1:]...)
+			close(in)
+			return
+		}
 	}
 }
 
@@ -402,6 +437,31 @@ func (n *Network) MakeGlobalComms(
 		}, out)
 }
 
+// StopRoundComms closes one inbound channel handed out by MakeRoundComms. Only
+// that node's channel: the round network is shared, and other voters may still
+// be in this round.
+func (n *Network) StopRoundComms(
+	roundNumber uint64,
+	in chan SignedMessageError[string, uint32, Signature, ID],
+) {
+	n.mtx.Lock()
+	round, ok := n.rounds[roundNumber]
+	n.mtx.Unlock()
+
+	if ok {
+		round.RemoveNode(in)
+	}
+}
+
+// StopGlobalComms closes the inbound channel handed to a voter by
+// MakeGlobalComms, which is how that voter is shut down.
+func (n *Network) StopGlobalComms(in chan GlobalInItem[string, uint32, Signature, ID]) {
+	n.mtx.Lock()
+	defer n.mtx.Unlock()
+
+	n.globalMessages.RemoveNode(in)
+}
+
 func (n *Network) SendMessage(message CommunicationIn[string, uint32, Signature, ID]) {
 	n.globalMessages.SendMessage(GlobalInItem[string, uint32, Signature, ID]{message, nil})
 }

pkg/finality-grandpa/timer.go

@@ -9,39 +9,46 @@ import (
 	"time"
 )
 
+// timer reports whether a deadline has passed and wakes whoever is polling it
+// when that changes. Rounds create several and discard them as they advance, so
+// Close releases one whose round finished before it fired.
 type timer struct {
-	wakerChan *wakerChan[error]
+	waker     atomic.Pointer[waker]
+	stop      chan struct{}
 	closeOnce sync.Once
 	expired   atomic.Bool
 }
 
 func newTimer(in <-chan time.Time) *timer {
-	inErr := make(chan error)
-	wc := newWakerChan(inErr)
-	t := timer{wakerChan: wc}
+	t := timer{stop: make(chan struct{})}
 	go t.poll(in)
 	return &t
 }
 
 func (t *timer) poll(in <-chan time.Time) {
-	<-in
-	t.closeOnce.Do(func() {
-		t.wakerChan.in <- nil
-		close(t.wakerChan.in)
-	})
+	select {
+	case <-in:
+	case <-t.stop:
+		return
+	}
+	// Ordered: waking before expired is set would send the poller back to sleep
+	// having seen the timer as still pending.
 	t.expired.Store(true)
+	if w := t.waker.Load(); w != nil {
+		w.wake()
+	}
 }
 
 func (t *timer) SetWaker(waker *waker) {
-	t.wakerChan.setWaker(waker)
+	t.waker.Store(waker)
 }
 
 func (t *timer) Elapsed() (bool, error) {
 	return t.expired.Load(), nil
 }
 
+// Close releases a timer that has not fired. Idempotent, and a no-op once the
+// timer has elapsed.
 func (t *timer) Close() {
-	t.closeOnce.Do(func() {
-		close(t.wakerChan.in)
-	})
+	t.closeOnce.Do(func() { close(t.stop) })
 }