github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/consensus/integration/stopper_test.go (about)

     1  package integration_test
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	"go.uber.org/atomic"
     8  
     9  	"github.com/onflow/flow-go/model/flow"
    10  )
    11  
    12  // Stopper is responsible for detecting a stopping condition, and stopping all nodes.
    13  //
    14  // Design motivation:
    15  //   - We can stop each node as soon as it enters a certain view. But the problem
    16  //     is if some fast node reaches a view earlier and gets stopped, it won't
    17  //     be available for other nodes to sync, and slow nodes will never be able
    18  //     to catch up.
    19  //   - A better strategy is to wait until all nodes have entered a certain view,
    20  //     then stop them all - this is what the Stopper does.
    21  type Stopper struct {
    22  	sync.Mutex
    23  	running  map[flow.Identifier]struct{}
    24  	stopping *atomic.Bool
    25  	// finalizedCount is the number of blocks which must be finalized (by each node)
    26  	// before the stopFunc is called
    27  	finalizedCount uint
    28  	// tolerate is the number of nodes which we will tolerate NOT finalizing the
    29  	// expected number of blocks
    30  	tolerate int
    31  	stopFunc func()
    32  	stopped  chan struct{}
    33  }
    34  
    35  func NewStopper(finalizedCount uint, tolerate int) *Stopper {
    36  	return &Stopper{
    37  		running:        make(map[flow.Identifier]struct{}),
    38  		stopping:       atomic.NewBool(false),
    39  		finalizedCount: finalizedCount,
    40  		tolerate:       tolerate,
    41  		stopped:        make(chan struct{}),
    42  	}
    43  }
    44  
    45  // AddNode registers a node with the Stopper, so that the stopping condition is
    46  // adjusted to account for this node (ie. we will now also wait for the added
    47  // node to finalize the desired number of blocks).
    48  func (s *Stopper) AddNode(n *Node) {
    49  	s.Lock()
    50  	defer s.Unlock()
    51  	s.running[n.id.ID()] = struct{}{}
    52  }
    53  
    54  // WithStopFunc adds a function to use to stop all nodes (typically the cancel function of the context used to start them).
    55  // Caution: not safe for concurrent use by multiple goroutines.
    56  func (s *Stopper) WithStopFunc(stop func()) {
    57  	s.stopFunc = stop
    58  }
    59  
    60  // onFinalizedTotal is called via CounterConsumer each time a node finalizes a block.
    61  // When called, the node with ID `id` has finalized `total` blocks.
    62  func (s *Stopper) onFinalizedTotal(id flow.Identifier, total uint) {
    63  	s.Lock()
    64  	defer s.Unlock()
    65  
    66  	if total < s.finalizedCount {
    67  		return
    68  	}
    69  
    70  	// keep track of remaining running nodes
    71  	delete(s.running, id)
    72  
    73  	// if all the honest nodes have reached the total number of
    74  	// finalized blocks, then stop all nodes
    75  	if len(s.running) <= s.tolerate {
    76  		go s.stopAll()
    77  	}
    78  }
    79  
    80  // stopAll stops all registered nodes, using the provided stopFunc.
    81  func (s *Stopper) stopAll() {
    82  	// only allow one process to stop all nodes, and stop them exactly once
    83  	if !s.stopping.CompareAndSwap(false, true) {
    84  		return
    85  	}
    86  	if s.stopFunc == nil {
    87  		panic("Stopper used without a stopFunc - use WithStopFunc to specify how to stop nodes once stop condition is reached")
    88  	}
    89  	fmt.Println("stopping all nodes...")
    90  	s.stopFunc()
    91  	close(s.stopped)
    92  }