github.com/MetalBlockchain/metalgo@v1.11.9/snow/consensus/snowman/network_test.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package snowman
     5  
     6  import (
     7  	"context"
     8  	"testing"
     9  
    10  	"github.com/MetalBlockchain/metalgo/ids"
    11  	"github.com/MetalBlockchain/metalgo/snow/choices"
    12  	"github.com/MetalBlockchain/metalgo/snow/consensus/snowball"
    13  	"github.com/MetalBlockchain/metalgo/snow/consensus/snowman/snowmantest"
    14  	"github.com/MetalBlockchain/metalgo/snow/snowtest"
    15  	"github.com/MetalBlockchain/metalgo/utils"
    16  	"github.com/MetalBlockchain/metalgo/utils/bag"
    17  	"github.com/MetalBlockchain/metalgo/utils/sampler"
    18  )
    19  
    20  type Network struct {
    21  	params         snowball.Parameters
    22  	colors         []*snowmantest.Block
    23  	rngSource      sampler.Source
    24  	nodes, running []Consensus
    25  }
    26  
    27  func NewNetwork(params snowball.Parameters, numColors int, rngSource sampler.Source) *Network {
    28  	n := &Network{
    29  		params: params,
    30  		colors: []*snowmantest.Block{{
    31  			TestDecidable: choices.TestDecidable{
    32  				IDV:     ids.Empty.Prefix(rngSource.Uint64()),
    33  				StatusV: choices.Processing,
    34  			},
    35  			ParentV: snowmantest.GenesisID,
    36  			HeightV: snowmantest.GenesisHeight + 1,
    37  		}},
    38  		rngSource: rngSource,
    39  	}
    40  
    41  	s := sampler.NewDeterministicUniform(n.rngSource)
    42  	for i := 1; i < numColors; i++ {
    43  		s.Initialize(uint64(len(n.colors)))
    44  		dependencyInd, _ := s.Next()
    45  		dependency := n.colors[dependencyInd]
    46  		n.colors = append(n.colors, &snowmantest.Block{
    47  			TestDecidable: choices.TestDecidable{
    48  				IDV:     ids.Empty.Prefix(rngSource.Uint64()),
    49  				StatusV: choices.Processing,
    50  			},
    51  			ParentV: dependency.IDV,
    52  			HeightV: dependency.HeightV + 1,
    53  		})
    54  	}
    55  	return n
    56  }
    57  
    58  func (n *Network) shuffleColors() {
    59  	s := sampler.NewDeterministicUniform(n.rngSource)
    60  	s.Initialize(uint64(len(n.colors)))
    61  	indices, _ := s.Sample(len(n.colors))
    62  	colors := []*snowmantest.Block(nil)
    63  	for _, index := range indices {
    64  		colors = append(colors, n.colors[int(index)])
    65  	}
    66  	n.colors = colors
    67  	utils.Sort(n.colors)
    68  }
    69  
    70  func (n *Network) AddNode(t testing.TB, sm Consensus) error {
    71  	snowCtx := snowtest.Context(t, snowtest.CChainID)
    72  	ctx := snowtest.ConsensusContext(snowCtx)
    73  	if err := sm.Initialize(ctx, n.params, snowmantest.GenesisID, snowmantest.GenesisHeight, snowmantest.GenesisTimestamp); err != nil {
    74  		return err
    75  	}
    76  
    77  	n.shuffleColors()
    78  	deps := map[ids.ID]ids.ID{}
    79  	for _, blk := range n.colors {
    80  		myDep, found := deps[blk.ParentV]
    81  		if !found {
    82  			myDep = blk.Parent()
    83  		}
    84  		myBlock := &snowmantest.Block{
    85  			TestDecidable: choices.TestDecidable{
    86  				IDV:     blk.ID(),
    87  				StatusV: blk.Status(),
    88  			},
    89  			ParentV: myDep,
    90  			HeightV: blk.Height(),
    91  			VerifyV: blk.Verify(context.Background()),
    92  			BytesV:  blk.Bytes(),
    93  		}
    94  		if err := sm.Add(myBlock); err != nil {
    95  			return err
    96  		}
    97  		deps[myBlock.ID()] = myDep
    98  	}
    99  	n.nodes = append(n.nodes, sm)
   100  	n.running = append(n.running, sm)
   101  	return nil
   102  }
   103  
   104  func (n *Network) Finalized() bool {
   105  	return len(n.running) == 0
   106  }
   107  
   108  func (n *Network) Round() error {
   109  	if len(n.running) == 0 {
   110  		return nil
   111  	}
   112  
   113  	s := sampler.NewDeterministicUniform(n.rngSource)
   114  	s.Initialize(uint64(len(n.running)))
   115  
   116  	runningInd, _ := s.Next()
   117  	running := n.running[runningInd]
   118  
   119  	s.Initialize(uint64(len(n.nodes)))
   120  	indices, _ := s.Sample(n.params.K)
   121  	sampledColors := bag.Bag[ids.ID]{}
   122  	for _, index := range indices {
   123  		peer := n.nodes[int(index)]
   124  		sampledColors.Add(peer.Preference())
   125  	}
   126  
   127  	if err := running.RecordPoll(context.Background(), sampledColors); err != nil {
   128  		return err
   129  	}
   130  
   131  	// If this node has been finalized, remove it from the poller
   132  	if running.NumProcessing() == 0 {
   133  		newSize := len(n.running) - 1
   134  		n.running[runningInd] = n.running[newSize]
   135  		n.running = n.running[:newSize]
   136  	}
   137  
   138  	return nil
   139  }
   140  
   141  func (n *Network) Agreement() bool {
   142  	if len(n.nodes) == 0 {
   143  		return true
   144  	}
   145  	pref := n.nodes[0].Preference()
   146  	for _, node := range n.nodes {
   147  		if pref != node.Preference() {
   148  			return false
   149  		}
   150  	}
   151  	return true
   152  }