github.com/koko1123/flow-go-1@v0.29.6/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/koko1123/flow-go-1/consensus/hotstuff/notifications"
    10  	"github.com/koko1123/flow-go-1/model/flow"
    11  )
    12  
    13  type StopperConsumer struct {
    14  	notifications.NoopConsumer
    15  }
    16  
    17  type Stopper struct {
    18  	sync.Mutex
    19  	running        map[flow.Identifier]struct{}
    20  	nodes          []*Node
    21  	stopping       *atomic.Bool
    22  	finalizedCount uint
    23  	tolerate       int
    24  	stopped        chan struct{}
    25  }
    26  
    27  // How to stop nodes?
    28  // We can stop each node as soon as it enters a certain view. But the problem
    29  // is if some fast nodes reaches a view earlier and gets stopped, it won't
    30  // be available for other nodes to sync, and slow nodes will never be able
    31  // to catch up.
    32  // a better strategy is to wait until all nodes has entered a certain view,
    33  // then stop them all.
    34  func NewStopper(finalizedCount uint, tolerate int) *Stopper {
    35  	return &Stopper{
    36  		running:        make(map[flow.Identifier]struct{}),
    37  		nodes:          make([]*Node, 0),
    38  		stopping:       atomic.NewBool(false),
    39  		finalizedCount: finalizedCount,
    40  		tolerate:       tolerate,
    41  		stopped:        make(chan struct{}),
    42  	}
    43  }
    44  
    45  func (s *Stopper) AddNode(n *Node) *StopperConsumer {
    46  	s.Lock()
    47  	defer s.Unlock()
    48  	s.running[n.id.ID()] = struct{}{}
    49  	s.nodes = append(s.nodes, n)
    50  	stopConsumer := &StopperConsumer{}
    51  	return stopConsumer
    52  }
    53  
    54  func (s *Stopper) onFinalizedTotal(id flow.Identifier, total uint) {
    55  	s.Lock()
    56  	defer s.Unlock()
    57  
    58  	if total < s.finalizedCount {
    59  		return
    60  	}
    61  
    62  	// keep track of remaining running nodes
    63  	delete(s.running, id)
    64  
    65  	// if all the honest nodes have reached the total number of
    66  	// finalized blocks, then stop all nodes
    67  	if len(s.running) <= s.tolerate {
    68  		go s.stopAll()
    69  	}
    70  }
    71  
    72  func (s *Stopper) stopAll() {
    73  	// only allow one process to stop all nodes,
    74  	// and stop them exactly once
    75  	if !s.stopping.CompareAndSwap(false, true) {
    76  		return
    77  	}
    78  
    79  	fmt.Println("stopping all nodes")
    80  
    81  	// wait until all nodes has been shut down
    82  	var wg sync.WaitGroup
    83  	for _, node := range s.nodes {
    84  		wg.Add(1)
    85  		go func(node *Node) {
    86  			node.Shutdown()
    87  			wg.Done()
    88  		}(node)
    89  	}
    90  	wg.Wait()
    91  
    92  	close(s.stopped)
    93  }