Merge pull request #269 from blinklabs-io/fix/prevent-multiple-start-stop-close

fix: protect stop/start/close functions from being called again

Author: agaffney

muxer/muxer.go

@@ -50,10 +50,11 @@ type Muxer struct {
 	startChan         chan bool
 	doneChan          chan bool
 	waitGroup         sync.WaitGroup
-	stopMutex         sync.Mutex
 	protocolSenders   map[uint16]chan *Segment
 	protocolReceivers map[uint16]chan *Segment
 	diffusionMode     DiffusionMode
+	onceStart         sync.Once
+	onceStop          sync.Once
 }
 
 // New creates a new Muxer object and starts the read loop
@@ -77,32 +78,26 @@ func (m *Muxer) ErrorChan() chan error {
 
 // Start unblocks the read loop after the initial handshake to allow it to start processing messages
 func (m *Muxer) Start() {
-	m.startChan <- true
+	m.onceStart.Do(func() {
+		m.startChan <- true
+	})
 }
 
 // Stop shuts down the muxer
 func (m *Muxer) Stop() {
-	// We use a mutex to prevent this function from being called multiple times
-	// concurrently, which would cause a race condition
-	m.stopMutex.Lock()
-	defer m.stopMutex.Unlock()
-	// Immediately return if we're already shutting down
-	select {
-	case <-m.doneChan:
-		return
-	default:
-	}
-	// 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 {
-		close(recvChan)
-	}
-	// Close ErrorChan to signify to consumer that we're shutting down
-	close(m.errorChan)
+	m.onceStop.Do(func() {
+		// 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 {
+			close(recvChan)
+		}
+		// Close ErrorChan to signify to consumer that we're shutting down
+		close(m.errorChan)
+	})
 }
 
 // SetDiffusionMode sets the muxer diffusion mode after the handshake completes
@@ -158,6 +153,12 @@ func (m *Muxer) RegisterProtocol(protocolId uint16) (chan *Segment, chan *Segmen
 // Send takes a populated Segment and writes it to the connection. A mutex is used to prevent more than
 // one protocol from sending at once
 func (m *Muxer) Send(msg *Segment) error {
+	// Immediately return if we're already shutting down
+	select {
+	case <-m.doneChan:
+		return fmt.Errorf("shutting down")
+	default:
+	}
 	// We use a mutex to make sure only one protocol can send at a time
 	m.sendMutex.Lock()
 	defer m.sendMutex.Unlock()

ouroboros.go

@@ -54,7 +54,7 @@ type Ouroboros struct {
 	handshakeFinishedChan chan interface{}
 	doneChan              chan interface{}
 	waitGroup             sync.WaitGroup
-	closeMutex            sync.Mutex
+	onceClose             sync.Once
 	sendKeepAlives        bool
 	delayMuxerStart       bool
 	fullDuplex            bool
@@ -130,47 +130,40 @@ func (o *Ouroboros) Dial(proto string, address string) error {
 
 // Close will shutdown the Ouroboros connection
 func (o *Ouroboros) Close() error {
-	// We use a mutex to prevent this function from being called multiple times
-	// concurrently, which would cause a race condition
-	o.closeMutex.Lock()
-	defer o.closeMutex.Unlock()
-	// Immediately return if we're already shutting down
-	select {
-	case <-o.doneChan:
-		return nil
-	default:
-	}
-	// Close doneChan to signify that we're shutting down
-	close(o.doneChan)
-	// Gracefully stop the muxer
-	if o.muxer != nil {
-		o.muxer.Stop()
-	}
-	// Close the underlying connection
-	if o.conn != nil {
-		if err := o.conn.Close(); err != nil {
-			return err
+	var err error
+	o.onceClose.Do(func() {
+		// Close doneChan to signify that we're shutting down
+		close(o.doneChan)
+		// Gracefully stop the muxer
+		if o.muxer != nil {
+			o.muxer.Stop()
 		}
-	}
-	// Wait for other goroutines to finish
-	o.waitGroup.Wait()
-	// Close channels
-	close(o.errorChan)
-	close(o.protoErrorChan)
-	// We can only close a channel once, so we have to jump through a few hoops
-	select {
-	// The channel is either closed or has an item pending
-	case _, ok := <-o.handshakeFinishedChan:
-		// We successfully retrieved an item
-		// This will probably never happen, but it doesn't hurt to cover this case
-		if ok {
+		// Close the underlying connection
+		if o.conn != nil {
+			if err = o.conn.Close(); err != nil {
+				return
+			}
+		}
+		// Wait for other goroutines to finish
+		o.waitGroup.Wait()
+		// Close channels
+		close(o.errorChan)
+		close(o.protoErrorChan)
+		// We can only close a channel once, so we have to jump through a few hoops
+		select {
+		// The channel is either closed or has an item pending
+		case _, ok := <-o.handshakeFinishedChan:
+			// We successfully retrieved an item
+			// This will probably never happen, but it doesn't hurt to cover this case
+			if ok {
+				close(o.handshakeFinishedChan)
+			}
+		// The channel is open and has no pending items
+		default:
 			close(o.handshakeFinishedChan)
 		}
-	// The channel is open and has no pending items
-	default:
-		close(o.handshakeFinishedChan)
-	}
-	return nil
+	})
+	return err
 }
 
 // ChainSync returns the chain-sync protocol handler

protocol/blockfetch/client.go

@@ -32,6 +32,7 @@ type Client struct {
 	startBatchResultChan chan error
 	busyMutex            sync.Mutex
 	blockUseCallback     bool
+	onceStop             sync.Once
 }
 
 func NewClient(protoOptions protocol.ProtocolOptions, cfg *Config) *Client {
@@ -77,8 +78,12 @@ func NewClient(protoOptions protocol.ProtocolOptions, cfg *Config) *Client {
 }
 
 func (c *Client) Stop() error {
-	msg := NewMsgClientDone()
-	return c.SendMessage(msg)
+	var err error
+	c.onceStop.Do(func() {
+		msg := NewMsgClientDone()
+		err = c.SendMessage(msg)
+	})
+	return err
 }
 
 // GetBlockRange starts an async process to fetch all blocks in the specified range (inclusive)

protocol/chainsync/client.go

@@ -35,6 +35,7 @@ type Client struct {
 	currentTipChan        chan Tip
 	wantFirstBlock        bool
 	firstBlockChan        chan common.Point
+	onceStop              sync.Once
 }
 
 // NewClient returns a new ChainSync client object
@@ -116,13 +117,16 @@ func (c *Client) messageHandler(msg protocol.Message, isResponse bool) error {
 
 // Stop transitions the protocol to the Done state. No more protocol operations will be possible afterward
 func (c *Client) Stop() error {
-	c.busyMutex.Lock()
-	defer c.busyMutex.Unlock()
-	msg := NewMsgDone()
-	if err := c.SendMessage(msg); err != nil {
-		return err
-	}
-	return nil
+	var err error
+	c.onceStop.Do(func() {
+		c.busyMutex.Lock()
+		defer c.busyMutex.Unlock()
+		msg := NewMsgDone()
+		if err = c.SendMessage(msg); err != nil {
+			return
+		}
+	})
+	return err
 }
 
 // GetCurrentTip returns the current chain tip

protocol/handshake/client.go

@@ -16,14 +16,16 @@ package handshake
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/blinklabs-io/gouroboros/protocol"
 )
 
 // Client implements the Handshake client
 type Client struct {
 	*protocol.Protocol
-	config *Config
+	config    *Config
+	onceStart sync.Once
 }
 
 // NewClient returns a new Handshake client object
@@ -60,22 +62,24 @@ func NewClient(protoOptions protocol.ProtocolOptions, cfg *Config) *Client {
 
 // Start begins the handshake process
 func (c *Client) Start() {
-	c.Protocol.Start()
-	// Send our ProposeVersions message
-	versionMap := make(map[uint16]interface{})
-	diffusionMode := DiffusionModeInitiatorOnly
-	if c.config.ClientFullDuplex {
-		diffusionMode = DiffusionModeInitiatorAndResponder
-	}
-	for _, version := range c.config.ProtocolVersions {
-		if c.Mode() == protocol.ProtocolModeNodeToNode {
-			versionMap[version] = []interface{}{c.config.NetworkMagic, diffusionMode}
-		} else {
-			versionMap[version] = c.config.NetworkMagic
+	c.onceStart.Do(func() {
+		c.Protocol.Start()
+		// Send our ProposeVersions message
+		versionMap := make(map[uint16]interface{})
+		diffusionMode := DiffusionModeInitiatorOnly
+		if c.config.ClientFullDuplex {
+			diffusionMode = DiffusionModeInitiatorAndResponder
 		}
-	}
-	msg := NewMsgProposeVersions(versionMap)
-	_ = c.SendMessage(msg)
+		for _, version := range c.config.ProtocolVersions {
+			if c.Mode() == protocol.ProtocolModeNodeToNode {
+				versionMap[version] = []interface{}{c.config.NetworkMagic, diffusionMode}
+			} else {
+				versionMap[version] = c.config.NetworkMagic
+			}
+		}
+		msg := NewMsgProposeVersions(versionMap)
+		_ = c.SendMessage(msg)
+	})
 }
 
 func (c *Client) handleMessage(msg protocol.Message, isResponse bool) error {

protocol/keepalive/client.go

@@ -16,14 +16,17 @@ package keepalive
 
 import (
 	"fmt"
-	"github.com/blinklabs-io/gouroboros/protocol"
+	"sync"
 	"time"
+
+	"github.com/blinklabs-io/gouroboros/protocol"
 )
 
 type Client struct {
 	*protocol.Protocol
-	config *Config
-	timer  *time.Timer
+	config    *Config
+	timer     *time.Timer
+	onceStart sync.Once
 }
 
 func NewClient(protoOptions protocol.ProtocolOptions, cfg *Config) *Client {
@@ -68,8 +71,10 @@ func NewClient(protoOptions protocol.ProtocolOptions, cfg *Config) *Client {
 }
 
 func (c *Client) Start() {
-	c.Protocol.Start()
-	c.startTimer()
+	c.onceStart.Do(func() {
+		c.Protocol.Start()
+		c.startTimer()
+	})
 }
 
 func (c *Client) startTimer() {

protocol/localtxmonitor/client.go

@@ -16,8 +16,9 @@ package localtxmonitor
 
 import (
 	"fmt"
-	"github.com/blinklabs-io/gouroboros/protocol"
 	"sync"
+
+	"github.com/blinklabs-io/gouroboros/protocol"
 )
 
 // Client implements the LocalTxMonitor client
@@ -31,6 +32,7 @@ type Client struct {
 	hasTxResultChan    chan bool
 	nextTxResultChan   chan []byte
 	getSizesResultChan chan MsgReplyGetSizesResult
+	onceStop           sync.Once
 }
 
 // NewClient returns a new LocalTxMonitor client object
@@ -133,13 +135,16 @@ func (c *Client) Release() error {
 
 // Stop transitions the protocol to the Done state. No more operations will be possible
 func (c *Client) Stop() error {
-	c.busyMutex.Lock()
-	defer c.busyMutex.Unlock()
-	msg := NewMsgDone()
-	if err := c.SendMessage(msg); err != nil {
-		return err
-	}
-	return nil
+	var err error
+	c.onceStop.Do(func() {
+		c.busyMutex.Lock()
+		defer c.busyMutex.Unlock()
+		msg := NewMsgDone()
+		if err = c.SendMessage(msg); err != nil {
+			return
+		}
+	})
+	return err
 }
 
 // HasTx returns whether or not the specified transaction ID exists in the mempool snapshot

protocol/localtxsubmission/client.go

@@ -16,9 +16,10 @@ package localtxsubmission
 
 import (
 	"fmt"
+	"sync"
+
 	"github.com/blinklabs-io/gouroboros/ledger"
 	"github.com/blinklabs-io/gouroboros/protocol"
-	"sync"
 )
 
 // Client implements the LocalTxSubmission client
@@ -27,6 +28,7 @@ type Client struct {
 	config           *Config
 	busyMutex        sync.Mutex
 	submitResultChan chan error
+	onceStop         sync.Once
 }
 
 // NewClient returns a new LocalTxSubmission client object
@@ -94,13 +96,16 @@ func (c *Client) SubmitTx(eraId uint16, tx []byte) error {
 
 // Stop transitions the protocol to the Done state. No more operations will be possible
 func (c *Client) Stop() error {
-	c.busyMutex.Lock()
-	defer c.busyMutex.Unlock()
-	msg := NewMsgDone()
-	if err := c.SendMessage(msg); err != nil {
-		return err
-	}
-	return nil
+	var err error
+	c.onceStop.Do(func() {
+		c.busyMutex.Lock()
+		defer c.busyMutex.Unlock()
+		msg := NewMsgDone()
+		if err = c.SendMessage(msg); err != nil {
+			return
+		}
+	})
+	return err
 }
 
 func (c *Client) handleAcceptTx() error {