Merge pull request #251 from blinklabs-io/fix/shutdown-sync-waitgroup

fix: use sync.WaitGroup to avoid random panics on shutdown

Author: agaffney

muxer/muxer.go

@@ -35,6 +35,7 @@ type Muxer struct {
 	sendMutex         sync.Mutex
 	startChan         chan bool
 	doneChan          chan bool
+	waitGroup         sync.WaitGroup
 	stopMutex         sync.Mutex
 	protocolSenders   map[uint16]chan *Segment
 	protocolReceivers map[uint16]chan *Segment
@@ -51,6 +52,7 @@ func New(conn net.Conn) *Muxer {
 		protocolSenders:   make(map[uint16]chan *Segment),
 		protocolReceivers: make(map[uint16]chan *Segment),
 	}
+	m.waitGroup.Add(1)
 	go m.readLoop()
 	return m
 }
@@ -78,6 +80,8 @@ func (m *Muxer) Stop() {
 	}
 	// Close doneChan to signify that we're shutting down
 	close(m.doneChan)
+	// Wait for other goroutines to shutdown
+	m.waitGroup.Wait()
 	// Close protocol receive channels
 	// We rely on the individual mini-protocols to close the sender channel
 	for _, recvChan := range m.protocolReceivers {
@@ -116,7 +120,9 @@ func (m *Muxer) RegisterProtocol(protocolId uint16) (chan *Segment, chan *Segmen
 	m.protocolSenders[protocolId] = senderChan
 	m.protocolReceivers[protocolId] = receiverChan
 	// Start Goroutine to handle outbound messages
+	m.waitGroup.Add(1)
 	go func() {
+		defer m.waitGroup.Done()
 		for {
 			select {
 			case _, ok := <-m.doneChan:
@@ -157,6 +163,7 @@ func (m *Muxer) Send(msg *Segment) error {
 // readLoop waits for incoming data on the connection, parses the segment, and passes it to the appropriate
 // protocol
 func (m *Muxer) readLoop() {
+	defer m.waitGroup.Done()
 	started := false
 	for {
 		// Break out of read loop if we're shutting down

ouroboros.go

@@ -39,6 +39,7 @@ type Ouroboros struct {
 	protoErrorChan        chan error
 	handshakeFinishedChan chan interface{}
 	doneChan              chan interface{}
+	waitGroup             sync.WaitGroup
 	closeMutex            sync.Mutex
 	sendKeepAlives        bool
 	delayMuxerStart       bool
@@ -137,6 +138,8 @@ func (o *Ouroboros) Close() error {
 			return err
 		}
 	}
+	// Wait for other goroutines to finish
+	o.waitGroup.Wait()
 	// Close channels
 	close(o.errorChan)
 	close(o.protoErrorChan)
@@ -196,7 +199,9 @@ func (o *Ouroboros) TxSubmission() *txsubmission.TxSubmission {
 func (o *Ouroboros) setupConnection() error {
 	o.muxer = muxer.New(o.conn)
 	// Start Goroutine to pass along errors from the muxer
+	o.waitGroup.Add(1)
 	go func() {
+		defer o.waitGroup.Done()
 		select {
 		case <-o.doneChan:
 			return
@@ -274,7 +279,9 @@ func (o *Ouroboros) setupConnection() error {
 		protoOptions.Version = protoOptions.Version - protocolVersionNtCFlag
 	}
 	// Start Goroutine to pass along errors from the mini-protocols
+	o.waitGroup.Add(1)
 	go func() {
+		defer o.waitGroup.Done()
 		err, ok := <-o.protoErrorChan
 		// The channel is closed, which means we're already shutting down
 		if !ok {

protocol/protocol.go

@@ -30,6 +30,7 @@ type Protocol struct {
 	recvReadyChan        chan bool
 	sendReadyChan        chan bool
 	doneChan             chan bool
+	waitGroup            sync.WaitGroup
 	stateTransitionTimer *time.Timer
 }
 
@@ -102,11 +103,16 @@ func (p *Protocol) Start() {
 	p.sendReadyChan = make(chan bool, 1)
 	// Start goroutine to cleanup when shutting down
 	go func() {
+		// Wait for doneChan to be closed
 		<-p.doneChan
+		// Wait for all other goroutines to finish
+		p.waitGroup.Wait()
+		// Close channels
 		close(p.sendQueueChan)
 		close(p.sendStateQueueChan)
 		close(p.recvReadyChan)
 		close(p.sendReadyChan)
+		// Cancel any timer
 		if p.stateTransitionTimer != nil {
 			// Stop timer and drain channel
 			if !p.stateTransitionTimer.Stop() {
@@ -118,6 +124,7 @@ func (p *Protocol) Start() {
 	// Set initial state
 	p.setState(p.config.InitialState)
 	// Start our send and receive Goroutines
+	p.waitGroup.Add(2)
 	go p.recvLoop()
 	go p.sendLoop()
 }
@@ -149,6 +156,7 @@ func (p *Protocol) SendError(err error) {
 }
 
 func (p *Protocol) sendLoop() {
+	defer p.waitGroup.Done()
 	var setNewState bool
 	var newState State
 	var err error
@@ -263,6 +271,7 @@ func (p *Protocol) sendLoop() {
 }
 
 func (p *Protocol) recvLoop() {
+	defer p.waitGroup.Done()
 	leftoverData := false
 	isResponse := false
 	for {