Merge pull request #78 from cloudstruct/feature/close-function

feat: method to shutdown connection

Author: agaffney

muxer/muxer.go

@@ -21,6 +21,7 @@ type Muxer struct {
 	conn              net.Conn
 	sendMutex         sync.Mutex
 	startChan         chan bool
+	doneChan          chan bool
 	ErrorChan         chan error
 	protocolSenders   map[uint16]chan *Segment
 	protocolReceivers map[uint16]chan *Segment
@@ -30,6 +31,7 @@ func New(conn net.Conn) *Muxer {
 	m := &Muxer{
 		conn:              conn,
 		startChan:         make(chan bool, 1),
+		doneChan:          make(chan bool),
 		ErrorChan:         make(chan error, 10),
 		protocolSenders:   make(map[uint16]chan *Segment),
 		protocolReceivers: make(map[uint16]chan *Segment),
@@ -42,6 +44,37 @@ func (m *Muxer) Start() {
 	m.startChan <- true
 }
 
+func (m *Muxer) Stop() {
+	// Immediately return if we're already shutting down
+	select {
+	case <-m.doneChan:
+		return
+	default:
+	}
+	// 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)
+	// Close doneChan to signify that we're shutting down
+	close(m.doneChan)
+}
+
+func (m *Muxer) sendError(err error) {
+	// Immediately return if we're already shutting down
+	select {
+	case <-m.doneChan:
+		return
+	default:
+	}
+	// Send error to consumer
+	m.ErrorChan <- err
+	// Stop the muxer on any error
+	m.Stop()
+}
+
 func (m *Muxer) RegisterProtocol(protocolId uint16) (chan *Segment, chan *Segment) {
 	// Generate channels
 	senderChan := make(chan *Segment, 10)
@@ -52,9 +85,17 @@ func (m *Muxer) RegisterProtocol(protocolId uint16) (chan *Segment, chan *Segmen
 	// Start Goroutine to handle outbound messages
 	go func() {
 		for {
-			msg := <-senderChan
-			if err := m.Send(msg); err != nil {
-				m.ErrorChan <- err
+			select {
+			case _, ok := <-m.doneChan:
+				// doneChan has been closed, which means we're shutting down
+				if !ok {
+					return
+				}
+			case msg := <-senderChan:
+				if err := m.Send(msg); err != nil {
+					m.sendError(err)
+					return
+				}
 			}
 		}
 	}()
@@ -81,9 +122,16 @@ func (m *Muxer) Send(msg *Segment) error {
 func (m *Muxer) readLoop() {
 	started := false
 	for {
+		// Break out of read loop if we're shutting down
+		select {
+		case <-m.doneChan:
+			return
+		default:
+		}
 		header := SegmentHeader{}
 		if err := binary.Read(m.conn, binary.BigEndian, &header); err != nil {
-			m.ErrorChan <- err
+			m.sendError(err)
+			return
 		}
 		msg := &Segment{
 			SegmentHeader: header,
@@ -92,24 +140,32 @@ func (m *Muxer) readLoop() {
 		// We use ReadFull because it guarantees to read the expected number of bytes or
 		// return an error
 		if _, err := io.ReadFull(m.conn, msg.Payload); err != nil {
-			m.ErrorChan <- err
-		}
-		// Wait until the muxer is started to process anything other than handshake messages
-		if !started && msg.GetProtocolId() != PROTOCOL_HANDSHAKE {
-			<-m.startChan
-			started = true
+			m.sendError(err)
+			return
 		}
 		// Send message payload to proper receiver
 		recvChan := m.protocolReceivers[msg.GetProtocolId()]
 		if recvChan == nil {
 			// Try the "unknown protocol" receiver if we didn't find an explicit one
 			recvChan = m.protocolReceivers[PROTOCOL_UNKNOWN]
 			if recvChan == nil {
-				m.ErrorChan <- fmt.Errorf("received message for unknown protocol ID %d", msg.GetProtocolId())
+				m.sendError(fmt.Errorf("received message for unknown protocol ID %d", msg.GetProtocolId()))
+				return
 			}
 		}
 		if recvChan != nil {
 			recvChan <- msg
 		}
+		// Wait until the muxer is started to continue
+		// We don't want to read more than one segment until the handshake is complete
+		if !started {
+			select {
+			case <-m.doneChan:
+				// Break out of read loop if we're shutting down
+				return
+			case <-m.startChan:
+				started = true
+			}
+		}
 	}
 }

ouroboros.go

@@ -1,6 +1,7 @@
 package ouroboros
 
 import (
+	"fmt"
 	"github.com/cloudstruct/go-ouroboros-network/muxer"
 	"github.com/cloudstruct/go-ouroboros-network/protocol"
 	"github.com/cloudstruct/go-ouroboros-network/protocol/blockfetch"
@@ -21,6 +22,7 @@ type Ouroboros struct {
 	handshakeComplete  bool
 	muxer              *muxer.Muxer
 	ErrorChan          chan error
+	protoErrorChan     chan error
 	sendKeepAlives     bool
 	delayMuxerStart    bool
 	// Mini-protocols
@@ -68,6 +70,7 @@ func New(options *OuroborosOptions) (*Ouroboros, error) {
 		ErrorChan:                       options.ErrorChan,
 		sendKeepAlives:                  options.SendKeepAlives,
 		delayMuxerStart:                 options.DelayMuxerStart,
+		protoErrorChan:                  make(chan error, 10),
 	}
 	if o.ErrorChan == nil {
 		o.ErrorChan = make(chan error, 10)
@@ -98,16 +101,32 @@ func (o *Ouroboros) Dial(proto string, address string) error {
 	return nil
 }
 
+func (o *Ouroboros) Close() error {
+	// Gracefully stop the muxer
+	o.muxer.Stop()
+	// Close the underlying connection
+	if err := o.conn.Close(); err != nil {
+		return err
+	}
+	return nil
+}
+
 func (o *Ouroboros) setupConnection() error {
 	o.muxer = muxer.New(o.conn)
 	// Start Goroutine to pass along errors from the muxer
 	go func() {
-		err := <-o.muxer.ErrorChan
-		o.ErrorChan <- err
+		err, ok := <-o.muxer.ErrorChan
+		// Break out of goroutine if muxer's error channel is closed
+		if !ok {
+			return
+		}
+		o.ErrorChan <- fmt.Errorf("muxer error: %s", err)
+		// Close connection on muxer errors
+		o.Close()
 	}()
 	protoOptions := protocol.ProtocolOptions{
 		Muxer:     o.muxer,
-		ErrorChan: o.ErrorChan,
+		ErrorChan: o.protoErrorChan,
 	}
 	var protoVersions []uint16
 	if o.useNodeToNodeProto {
@@ -131,13 +150,26 @@ func (o *Ouroboros) setupConnection() error {
 			return err
 		}
 	}
-	o.handshakeComplete = <-o.Handshake.Finished
+	// Wait for handshake completion or error
+	select {
+	case err := <-o.protoErrorChan:
+		return err
+	case finished := <-o.Handshake.Finished:
+		o.handshakeComplete = finished
+	}
 	// Provide the negotiated protocol version to the various mini-protocols
 	protoOptions.Version = o.Handshake.Version
 	// Drop bit used to signify NtC protocol versions
 	if protoOptions.Version > PROTOCOL_VERSION_NTC_FLAG {
 		protoOptions.Version = protoOptions.Version - PROTOCOL_VERSION_NTC_FLAG
 	}
+	// Start Goroutine to pass along errors from the mini-protocols
+	go func() {
+		err := <-o.protoErrorChan
+		o.ErrorChan <- fmt.Errorf("protocol error: %s", err)
+		// Close connection on mini-protocol errors
+		o.Close()
+	}()
 	// Configure the relevant mini-protocols
 	if o.useNodeToNodeProto {
 		versionNtN := GetProtocolVersionNtN(o.Handshake.Version)
@@ -160,6 +192,7 @@ func (o *Ouroboros) setupConnection() error {
 			o.LocalStateQuery = localstatequery.New(protoOptions, o.localStateQueryCallbackConfig)
 		}
 	}
+	// Start muxer
 	if !o.delayMuxerStart {
 		o.muxer.Start()
 	}

protocol/protocol.go

@@ -25,6 +25,7 @@ type Protocol struct {
 	sendStateQueueChan chan Message
 	recvReadyChan      chan bool
 	sendReadyChan      chan bool
+	doneChan           chan bool
 }
 
 type ProtocolConfig struct {
@@ -79,6 +80,7 @@ func New(config ProtocolConfig) *Protocol {
 		sendStateQueueChan: make(chan Message, 50),
 		recvReadyChan:      make(chan bool, 1),
 		sendReadyChan:      make(chan bool, 1),
+		doneChan:           make(chan bool),
 	}
 	// Set initial state
 	p.setState(config.InitialState)
@@ -110,8 +112,16 @@ func (p *Protocol) sendLoop() {
 	var newState State
 	var err error
 	for {
-		// Wait until ready to send based on state map
-		<-p.sendReadyChan
+		select {
+		case <-p.sendReadyChan:
+			// We are ready to send based on state map
+		case <-p.doneChan:
+			// We are responsible for closing this channel as the sender, even through it
+			// was created by the muxer
+			close(p.muxerSendChan)
+			// Break out of send loop if we're shutting down
+			return
+		}
 		// Lock the state to prevent collisions
 		p.stateMutex.Lock()
 		// Check for queued state changes from previous pipelined sends
@@ -212,7 +222,12 @@ func (p *Protocol) recvLoop() {
 		// Don't grab the next segment from the muxer if we still have data in the buffer
 		if !leftoverData {
 			// Wait for segment
-			segment := <-p.muxerRecvChan
+			segment, ok := <-p.muxerRecvChan
+			// Break out of receive loop if channel is closed
+			if !ok {
+				close(p.doneChan)
+				return
+			}
 			// Add segment payload to buffer
 			p.recvBuffer.Write(segment.Payload)
 			// Save whether it's a response