Wait for flows to finish before shutting down (#1605)

* Wait for flows to finish before shutting down

* Use CompareAndSwap

* Add comment

* Fix error message

Co-authored-by: stasatdaglabs <39559713+stasatdaglabs@users.noreply.github.com>

app/component_manager.go

@@ -72,6 +72,8 @@ func (a *ComponentManager) Stop() {
 		log.Errorf("Error stopping the net adapter: %+v", err)
 	}
 
+	a.protocolManager.Close()
+
 	return
 }
 

app/protocol/flowcontext/flow_context.go

@@ -61,6 +61,8 @@ type FlowContext struct {
 
 	orphans      map[externalapi.DomainHash]*externalapi.DomainBlock
 	orphansMutex sync.RWMutex
+
+	shutdownChan chan struct{}
 }
 
 // New returns a new instance of FlowContext.
@@ -79,9 +81,21 @@ func New(cfg *config.Config, domain domain.Domain, addressManager *addressmanage
 		transactionsToRebroadcast:   make(map[externalapi.DomainTransactionID]*externalapi.DomainTransaction),
 		orphans:                     make(map[externalapi.DomainHash]*externalapi.DomainBlock),
 		timeStarted:                 mstime.Now().UnixMilliseconds(),
+		shutdownChan:                make(chan struct{}),
 	}
 }
 
+// Close signals to all flows the the protocol manager is closed.
+func (f *FlowContext) Close() {
+	close(f.shutdownChan)
+}
+
+// ShutdownChan is a chan where flows can subscribe to shutdown
+// event.
+func (f *FlowContext) ShutdownChan() <-chan struct{} {
+	return f.shutdownChan
+}
+
 // SetOnBlockAddedToDAGHandler sets the onBlockAddedToDAG handler
 func (f *FlowContext) SetOnBlockAddedToDAGHandler(onBlockAddedToDAGHandler OnBlockAddedToDAGHandler) {
 	f.onBlockAddedToDAGHandler = onBlockAddedToDAGHandler

app/protocol/flows/ping/send.go

@@ -13,6 +13,7 @@ import (
 
 // SendPingsContext is the interface for the context needed for the SendPings flow.
 type SendPingsContext interface {
+	ShutdownChan() <-chan struct{}
 }
 
 type sendPingsFlow struct {
@@ -39,7 +40,13 @@ func (flow *sendPingsFlow) start() error {
 	ticker := time.NewTicker(pingInterval)
 	defer ticker.Stop()
 
-	for range ticker.C {
+	for {
+		select {
+		case <-flow.ShutdownChan():
+			return nil
+		case <-ticker.C:
+		}
+
 		nonce, err := random.Uint64()
 		if err != nil {
 			return err
@@ -62,5 +69,4 @@ func (flow *sendPingsFlow) start() error {
 		}
 		flow.peer.SetPingIdle()
 	}
-	return nil
 }

app/protocol/manager.go

@@ -2,6 +2,9 @@ package protocol
 
 import (
 	"fmt"
+	"github.com/pkg/errors"
+	"sync"
+	"sync/atomic"
 
 	"github.com/kaspanet/kaspad/domain"
 
@@ -17,7 +20,9 @@ import (
 
 // Manager manages the p2p protocol
 type Manager struct {
-	context *flowcontext.FlowContext
+	context          *flowcontext.FlowContext
+	routersWaitGroup sync.WaitGroup
+	isClosed         uint32
 }
 
 // NewManager creates a new instance of the p2p protocol manager
@@ -32,6 +37,18 @@ func NewManager(cfg *config.Config, domain domain.Domain, netAdapter *netadapter
 	return &manager, nil
 }
 
+// Close closes the protocol manager and waits until all p2p flows
+// finish.
+func (m *Manager) Close() {
+	if !atomic.CompareAndSwapUint32(&m.isClosed, 0, 1) {
+		panic(errors.New("The protocol manager was already closed"))
+	}
+
+	atomic.StoreUint32(&m.isClosed, 1)
+	m.context.Close()
+	m.routersWaitGroup.Wait()
+}
+
 // Peers returns the currently active peers
 func (m *Manager) Peers() []*peerpkg.Peer {
 	return m.context.Peers()
@@ -53,11 +70,13 @@ func (m *Manager) AddBlock(block *externalapi.DomainBlock) error {
 	return m.context.AddBlock(block)
 }
 
-func (m *Manager) runFlows(flows []*flow, peer *peerpkg.Peer, errChan <-chan error) error {
+func (m *Manager) runFlows(flows []*flow, peer *peerpkg.Peer, errChan <-chan error, flowsWaitGroup *sync.WaitGroup) error {
+	flowsWaitGroup.Add(len(flows))
 	for _, flow := range flows {
 		executeFunc := flow.executeFunc // extract to new variable so that it's not overwritten
 		spawn(fmt.Sprintf("flow-%s", flow.name), func() {
 			executeFunc(peer)
+			flowsWaitGroup.Done()
 		})
 	}
 

app/protocol/protocol.go

@@ -1,10 +1,10 @@
 package protocol
 
 import (
-	"sync/atomic"
-
 	"github.com/kaspanet/kaspad/app/protocol/flows/rejects"
 	"github.com/kaspanet/kaspad/infrastructure/network/connmanager"
+	"sync"
+	"sync/atomic"
 
 	"github.com/kaspanet/kaspad/app/appmessage"
 	"github.com/kaspanet/kaspad/app/protocol/flows/addressexchange"
@@ -41,6 +41,13 @@ func (m *Manager) routerInitializer(router *routerpkg.Router, netConnection *net
 	// After flows were registered - spawn a new thread that will wait for connection to finish initializing
 	// and start receiving messages
 	spawn("routerInitializer-runFlows", func() {
+		m.routersWaitGroup.Add(1)
+		defer m.routersWaitGroup.Done()
+
+		if atomic.LoadUint32(&m.isClosed) == 1 {
+			panic(errors.Errorf("tried to initialize router when the protocol manager is closed"))
+		}
+
 		isBanned, err := m.context.ConnectionManager().IsBanned(netConnection)
 		if err != nil && !errors.Is(err, addressmanager.ErrAddressNotFound) {
 			panic(err)
@@ -79,11 +86,17 @@ func (m *Manager) routerInitializer(router *routerpkg.Router, netConnection *net
 
 		removeHandshakeRoutes(router)
 
-		err = m.runFlows(flows, peer, errChan)
+		flowsWaitGroup := &sync.WaitGroup{}
+		err = m.runFlows(flows, peer, errChan, flowsWaitGroup)
 		if err != nil {
 			m.handleError(err, netConnection, router.OutgoingRoute())
+			// We call `flowsWaitGroup.Wait()` in two places instead of deferring, because
+			// we already defer `m.routersWaitGroup.Done()`, so we try to avoid error prone
+			// and confusing use of multiple dependent defers.
+			flowsWaitGroup.Wait()
 			return
 		}
+		flowsWaitGroup.Wait()
 	})
 }