github.com/neatlab/neatio@v1.7.3-0.20220425043230-d903e92fcc75/network/p2p/simulations/simulation.go (about)

     1  package simulations
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/neatlab/neatio/network/p2p/discover"
     8  )
     9  
    10  type Simulation struct {
    11  	network *Network
    12  }
    13  
    14  func NewSimulation(network *Network) *Simulation {
    15  	return &Simulation{
    16  		network: network,
    17  	}
    18  }
    19  
    20  func (s *Simulation) Run(ctx context.Context, step *Step) (result *StepResult) {
    21  	result = newStepResult()
    22  
    23  	result.StartedAt = time.Now()
    24  	defer func() { result.FinishedAt = time.Now() }()
    25  
    26  	stop := s.watchNetwork(result)
    27  	defer stop()
    28  
    29  	if err := step.Action(ctx); err != nil {
    30  		result.Error = err
    31  		return
    32  	}
    33  
    34  	nodes := make(map[discover.NodeID]struct{}, len(step.Expect.Nodes))
    35  	for _, id := range step.Expect.Nodes {
    36  		nodes[id] = struct{}{}
    37  	}
    38  	for len(result.Passes) < len(nodes) {
    39  		select {
    40  		case id := <-step.Trigger:
    41  
    42  			if _, ok := nodes[id]; !ok {
    43  				continue
    44  			}
    45  
    46  			if _, ok := result.Passes[id]; ok {
    47  				continue
    48  			}
    49  
    50  			pass, err := step.Expect.Check(ctx, id)
    51  			if err != nil {
    52  				result.Error = err
    53  				return
    54  			}
    55  			if pass {
    56  				result.Passes[id] = time.Now()
    57  			}
    58  		case <-ctx.Done():
    59  			result.Error = ctx.Err()
    60  			return
    61  		}
    62  	}
    63  
    64  	return
    65  }
    66  
    67  func (s *Simulation) watchNetwork(result *StepResult) func() {
    68  	stop := make(chan struct{})
    69  	done := make(chan struct{})
    70  	events := make(chan *Event)
    71  	sub := s.network.Events().Subscribe(events)
    72  	go func() {
    73  		defer close(done)
    74  		defer sub.Unsubscribe()
    75  		for {
    76  			select {
    77  			case event := <-events:
    78  				result.NetworkEvents = append(result.NetworkEvents, event)
    79  			case <-stop:
    80  				return
    81  			}
    82  		}
    83  	}()
    84  	return func() {
    85  		close(stop)
    86  		<-done
    87  	}
    88  }
    89  
    90  type Step struct {
    91  	Action func(context.Context) error
    92  
    93  	Trigger chan discover.NodeID
    94  
    95  	Expect *Expectation
    96  }
    97  
    98  type Expectation struct {
    99  	Nodes []discover.NodeID
   100  
   101  	Check func(context.Context, discover.NodeID) (bool, error)
   102  }
   103  
   104  func newStepResult() *StepResult {
   105  	return &StepResult{
   106  		Passes: make(map[discover.NodeID]time.Time),
   107  	}
   108  }
   109  
   110  type StepResult struct {
   111  	Error error
   112  
   113  	StartedAt time.Time
   114  
   115  	FinishedAt time.Time
   116  
   117  	Passes map[discover.NodeID]time.Time
   118  
   119  	NetworkEvents []*Event
   120  }