refactor: keep wg.Add and wg.Donecloser

This commit is contained in:
Halimao 2024-05-22 15:59:44 +08:00
parent 7fd7c1f7dd
commit cda2483f39
14 changed files with 99 additions and 54 deletions

View file

@ -47,7 +47,10 @@ func startEngineClient(config *lightClientConfig, rpc *rpc.Client, headCh <-chan
cancelRoot: cancel, cancelRoot: cancel,
} }
ec.wg.Add(1) ec.wg.Add(1)
go ec.updateLoop(headCh) go func() {
defer ec.wg.Done()
ec.updateLoop(headCh)
}()
return ec return ec
} }
@ -57,8 +60,6 @@ func (ec *engineClient) stop() {
} }
func (ec *engineClient) updateLoop(headCh <-chan types.ChainHeadEvent) { func (ec *engineClient) updateLoop(headCh <-chan types.ChainHeadEvent) {
defer ec.wg.Done()
for { for {
select { select {
case <-ec.rootCtx.Done(): case <-ec.rootCtx.Done():

View file

@ -120,7 +120,10 @@ func New(config Config) (*Console, error) {
} }
console.wg.Add(1) console.wg.Add(1)
go console.interruptHandler() go func() {
defer console.wg.Done()
console.interruptHandler()
}()
return console, nil return console, nil
} }
@ -363,8 +366,6 @@ func (c *Console) Evaluate(statement string) {
// interruptHandler runs in its own goroutine and waits for signals. // interruptHandler runs in its own goroutine and waits for signals.
// When a signal is received, it interrupts the JS interpreter. // When a signal is received, it interrupts the JS interpreter.
func (c *Console) interruptHandler() { func (c *Console) interruptHandler() {
defer c.wg.Done()
// During Interactive, liner inhibits the signal while it is prompting for // During Interactive, liner inhibits the signal while it is prompting for
// input. However, the signal will be received while evaluating JS. // input. However, the signal will be received while evaluating JS.
// //

View file

@ -63,8 +63,14 @@ func (s *scheduler) run(sections chan uint64, dist chan *request, done chan []by
// Start the pipeline schedulers to forward between user -> distributor -> user // Start the pipeline schedulers to forward between user -> distributor -> user
wg.Add(2) wg.Add(2)
go s.scheduleRequests(sections, dist, pend, quit, wg) go func() {
go s.scheduleDeliveries(pend, done, quit, wg) defer wg.Done()
s.scheduleRequests(sections, dist, pend, quit)
}()
go func() {
defer wg.Done()
s.scheduleDeliveries(pend, done, quit)
}()
} }
// reset cleans up any leftovers from previous runs. This is required before a // reset cleans up any leftovers from previous runs. This is required before a
@ -84,9 +90,8 @@ func (s *scheduler) reset() {
// scheduleRequests reads section retrieval requests from the input channel, // scheduleRequests reads section retrieval requests from the input channel,
// deduplicates the stream and pushes unique retrieval tasks into the distribution // deduplicates the stream and pushes unique retrieval tasks into the distribution
// channel for a database or network layer to honour. // channel for a database or network layer to honour.
func (s *scheduler) scheduleRequests(reqs chan uint64, dist chan *request, pend chan uint64, quit chan struct{}, wg *sync.WaitGroup) { func (s *scheduler) scheduleRequests(reqs chan uint64, dist chan *request, pend chan uint64, quit chan struct{}) {
// Clean up the goroutine and pipeline when done // Clean up the goroutine and pipeline when done
defer wg.Done()
defer close(pend) defer close(pend)
// Keep reading and scheduling section requests // Keep reading and scheduling section requests
@ -131,9 +136,8 @@ func (s *scheduler) scheduleRequests(reqs chan uint64, dist chan *request, pend
// scheduleDeliveries reads section acceptance notifications and waits for them // scheduleDeliveries reads section acceptance notifications and waits for them
// to be delivered, pushing them into the output data buffer. // to be delivered, pushing them into the output data buffer.
func (s *scheduler) scheduleDeliveries(pend chan uint64, done chan []byte, quit chan struct{}, wg *sync.WaitGroup) { func (s *scheduler) scheduleDeliveries(pend chan uint64, done chan []byte, quit chan struct{}) {
// Clean up the goroutine and pipeline when done // Clean up the goroutine and pipeline when done
defer wg.Done()
defer close(done) defer close(done)
// Keep reading notifications and scheduling deliveries // Keep reading notifications and scheduling deliveries

View file

@ -197,16 +197,24 @@ func TestCopy(t *testing.T) {
} }
// Finalise the changes on all concurrently // Finalise the changes on all concurrently
finalise := func(wg *sync.WaitGroup, db *StateDB) { finalise := func(db *StateDB) {
defer wg.Done()
db.Finalise(true) db.Finalise(true)
} }
var wg sync.WaitGroup var wg sync.WaitGroup
wg.Add(3) wg.Add(3)
go finalise(&wg, orig) go func() {
go finalise(&wg, copy) defer wg.Done()
go finalise(&wg, ccopy) finalise(orig)
}()
go func() {
defer wg.Done()
finalise(copy)
}()
go func() {
defer wg.Done()
finalise(ccopy)
}()
wg.Wait() wg.Wait()
// Verify that the three states have been updated independently // Verify that the three states have been updated independently

View file

@ -198,7 +198,6 @@ func newHandler(config *handlerConfig) (*handler, error) {
// protoTracker tracks the number of active protocol handlers. // protoTracker tracks the number of active protocol handlers.
func (h *handler) protoTracker() { func (h *handler) protoTracker() {
defer h.wg.Done()
var active int var active int
for { for {
select { select {
@ -426,14 +425,20 @@ func (h *handler) Start(maxPeers int) {
h.wg.Add(1) h.wg.Add(1)
h.txsCh = make(chan core.NewTxsEvent, txChanSize) h.txsCh = make(chan core.NewTxsEvent, txChanSize)
h.txsSub = h.txpool.SubscribeTransactions(h.txsCh, false) h.txsSub = h.txpool.SubscribeTransactions(h.txsCh, false)
go h.txBroadcastLoop() go func() {
defer h.wg.Done()
h.txBroadcastLoop()
}()
// start sync handlers // start sync handlers
h.txFetcher.Start() h.txFetcher.Start()
// start peer handler tracker // start peer handler tracker
h.wg.Add(1) h.wg.Add(1)
go h.protoTracker() go func() {
defer h.wg.Done()
h.protoTracker()
}()
} }
func (h *handler) Stop() { func (h *handler) Stop() {
@ -535,7 +540,6 @@ func (h *handler) BroadcastTransactions(txs types.Transactions) {
// txBroadcastLoop announces new transactions to connected peers. // txBroadcastLoop announces new transactions to connected peers.
func (h *handler) txBroadcastLoop() { func (h *handler) txBroadcastLoop() {
defer h.wg.Done()
for { for {
select { select {
case event := <-h.txsCh: case event := <-h.txsCh:

View file

@ -178,8 +178,14 @@ func newDialScheduler(config dialConfig, it enode.Iterator, setupFunc dialSetupF
d.lastStatsLog = d.clock.Now() d.lastStatsLog = d.clock.Now()
d.ctx, d.cancel = context.WithCancel(context.Background()) d.ctx, d.cancel = context.WithCancel(context.Background())
d.wg.Add(2) d.wg.Add(2)
go d.readNodes(it) go func() {
go d.loop(it) defer d.wg.Done()
d.readNodes(it)
}()
go func() {
defer d.wg.Done()
d.loop(it)
}()
return d return d
} }
@ -311,14 +317,11 @@ loop:
for range d.dialing { for range d.dialing {
<-d.doneCh <-d.doneCh
} }
d.wg.Done()
} }
// readNodes runs in its own goroutine and delivers nodes from // readNodes runs in its own goroutine and delivers nodes from
// the input iterator to the nodesIn channel. // the input iterator to the nodesIn channel.
func (d *dialScheduler) readNodes(it enode.Iterator) { func (d *dialScheduler) readNodes(it enode.Iterator) {
defer d.wg.Done()
for it.Next() { for it.Next() {
select { select {
case d.nodesIn <- it.Node(): case d.nodesIn <- it.Node():

View file

@ -150,8 +150,14 @@ func ListenV4(c UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv4, error) {
go tab.loop() go tab.loop()
t.wg.Add(2) t.wg.Add(2)
go t.loop() go func() {
go t.readLoop(cfg.Unhandled) defer t.wg.Done()
t.loop()
}()
go func() {
defer t.wg.Done()
t.readLoop(cfg.Unhandled)
}()
return t, nil return t, nil
} }

View file

@ -138,8 +138,14 @@ func ListenV5(conn UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv5, error) {
} }
go t.tab.loop() go t.tab.loop()
t.wg.Add(2) t.wg.Add(2)
go t.readLoop() go func() {
go t.dispatch() defer t.wg.Done()
t.readLoop()
}()
go func() {
defer t.wg.Done()
t.dispatch()
}()
return t, nil return t, nil
} }
@ -513,8 +519,6 @@ func (t *UDPv5) callDone(c *callV5) {
// When that happens the call is simply re-sent to complete the handshake. We allow one // When that happens the call is simply re-sent to complete the handshake. We allow one
// handshake attempt per call. // handshake attempt per call.
func (t *UDPv5) dispatch() { func (t *UDPv5) dispatch() {
defer t.wg.Done()
// Arm first read. // Arm first read.
t.readNextCh <- struct{}{} t.readNextCh <- struct{}{}
@ -651,8 +655,6 @@ func (t *UDPv5) send(toID enode.ID, toAddr *net.UDPAddr, packet v5wire.Packet, c
// readLoop runs in its own goroutine and reads packets from the network. // readLoop runs in its own goroutine and reads packets from the network.
func (t *UDPv5) readLoop() { func (t *UDPv5) readLoop() {
defer t.wg.Done()
buf := make([]byte, maxPacketSize) buf := make([]byte, maxPacketSize)
for range t.readNextCh { for range t.readNextCh {
nbytes, from, err := t.conn.ReadFromUDP(buf) nbytes, from, err := t.conn.ReadFromUDP(buf)

View file

@ -174,10 +174,13 @@ func (m *FairMix) AddSource(it Iterator) {
if m.closed == nil { if m.closed == nil {
return return
} }
m.wg.Add(1)
source := &mixSource{it, make(chan *Node), m.timeout} source := &mixSource{it, make(chan *Node), m.timeout}
m.sources = append(m.sources, source) m.sources = append(m.sources, source)
go m.runSource(m.closed, source) m.wg.Add(1)
go func() {
defer m.wg.Done()
m.runSource(m.closed, source)
}()
} }
// Close shuts down the mixer and all current sources. // Close shuts down the mixer and all current sources.
@ -281,7 +284,6 @@ func (m *FairMix) deleteSource(s *mixSource) {
// runSource reads a single source in a loop. // runSource reads a single source in a loop.
func (m *FairMix) runSource(closed chan struct{}, s *mixSource) { func (m *FairMix) runSource(closed chan struct{}, s *mixSource) {
defer m.wg.Done()
defer close(s.next) defer close(s.next)
for s.it.Next() { for s.it.Next() {
n := s.it.Node() n := s.it.Node()

View file

@ -252,8 +252,14 @@ func (p *Peer) run() (remoteRequested bool, err error) {
reason DiscReason // sent to the peer reason DiscReason // sent to the peer
) )
p.wg.Add(2) p.wg.Add(2)
go p.readLoop(readErr) go func() {
go p.pingLoop() defer p.wg.Done()
p.readLoop(readErr)
}()
go func() {
defer p.wg.Done()
p.pingLoop()
}()
// Start all protocol handlers. // Start all protocol handlers.
writeStart <- struct{}{} writeStart <- struct{}{}
@ -295,8 +301,6 @@ loop:
} }
func (p *Peer) pingLoop() { func (p *Peer) pingLoop() {
defer p.wg.Done()
ping := time.NewTimer(pingInterval) ping := time.NewTimer(pingInterval)
defer ping.Stop() defer ping.Stop()
@ -319,7 +323,6 @@ func (p *Peer) pingLoop() {
} }
func (p *Peer) readLoop(errc chan<- error) { func (p *Peer) readLoop(errc chan<- error) {
defer p.wg.Done()
for { for {
msg, err := p.rw.ReadMsg() msg, err := p.rw.ReadMsg()
if err != nil { if err != nil {

View file

@ -656,7 +656,12 @@ func (srv *Server) setupListening() error {
} }
srv.loopWG.Add(1) srv.loopWG.Add(1)
go srv.listenLoop() go func() {
// Wait for slots to be returned on exit. This ensures all connection goroutines
// are down before listenLoop returns.
defer srv.loopWG.Done()
srv.listenLoop()
}()
return nil return nil
} }
@ -859,9 +864,6 @@ func (srv *Server) listenLoop() {
slots <- struct{}{} slots <- struct{}{}
} }
// Wait for slots to be returned on exit. This ensures all connection goroutines
// are down before listenLoop returns.
defer srv.loopWG.Done()
defer func() { defer func() {
for i := 0; i < cap(slots); i++ { for i := 0; i < cap(slots); i++ {
<-slots <-slots

View file

@ -334,15 +334,20 @@ func (n *ExecNode) ServeRPC(clientConn *websocket.Conn) error {
} }
var wg sync.WaitGroup var wg sync.WaitGroup
wg.Add(2) wg.Add(2)
go wsCopy(&wg, conn, clientConn) go func() {
go wsCopy(&wg, clientConn, conn) defer wg.Done()
wsCopy(conn, clientConn)
}()
go func() {
defer wg.Done()
wsCopy(clientConn, conn)
}()
wg.Wait() wg.Wait()
conn.Close() conn.Close()
return nil return nil
} }
func wsCopy(wg *sync.WaitGroup, src, dst *websocket.Conn) { func wsCopy(src, dst *websocket.Conn) {
defer wg.Done()
for { for {
msgType, r, err := src.NextReader() msgType, r, err := src.NextReader()
if err != nil { if err != nil {

View file

@ -352,7 +352,6 @@ func testClientCancel(transport string, t *testing.T) {
ncallers = 10 ncallers = 10
) )
caller := func(index int) { caller := func(index int) {
defer wg.Done()
for i := 0; i < nreqs; i++ { for i := 0; i < nreqs; i++ {
var ( var (
ctx context.Context ctx context.Context
@ -386,7 +385,10 @@ func testClientCancel(transport string, t *testing.T) {
} }
wg.Add(ncallers) wg.Add(ncallers)
for i := 0; i < ncallers; i++ { for i := 0; i < ncallers; i++ {
go caller(i) go func(idx int) {
defer wg.Done()
caller(idx)
}(i)
} }
wg.Wait() wg.Wait()
} }

View file

@ -319,7 +319,10 @@ func newWebsocketCodec(conn *websocket.Conn, host string, req http.Header, readL
return nil return nil
}) })
wc.wg.Add(1) wc.wg.Add(1)
go wc.pingLoop() go func() {
defer wc.wg.Done()
wc.pingLoop()
}()
return wc return wc
} }
@ -347,7 +350,6 @@ func (wc *websocketCodec) writeJSON(ctx context.Context, v interface{}, isError
// pingLoop sends periodic ping frames when the connection is idle. // pingLoop sends periodic ping frames when the connection is idle.
func (wc *websocketCodec) pingLoop() { func (wc *websocketCodec) pingLoop() {
var pingTimer = time.NewTimer(wsPingInterval) var pingTimer = time.NewTimer(wsPingInterval)
defer wc.wg.Done()
defer pingTimer.Stop() defer pingTimer.Stop()
for { for {