github.com/bcnmy/go-ethereum@v1.10.27/p2p/simulations/simulation.go (about)

     1  // Copyright 2017 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package simulations
    18  
    19  import (
    20  	"context"
    21  	"time"
    22  
    23  	"github.com/ethereum/go-ethereum/p2p/enode"
    24  )
    25  
    26  // Simulation provides a framework for running actions in a simulated network
    27  // and then waiting for expectations to be met
    28  type Simulation struct {
    29  	network *Network
    30  }
    31  
    32  // NewSimulation returns a new simulation which runs in the given network
    33  func NewSimulation(network *Network) *Simulation {
    34  	return &Simulation{
    35  		network: network,
    36  	}
    37  }
    38  
    39  // Run performs a step of the simulation by performing the step's action and
    40  // then waiting for the step's expectation to be met
    41  func (s *Simulation) Run(ctx context.Context, step *Step) (result *StepResult) {
    42  	result = newStepResult()
    43  
    44  	result.StartedAt = time.Now()
    45  	defer func() { result.FinishedAt = time.Now() }()
    46  
    47  	// watch network events for the duration of the step
    48  	stop := s.watchNetwork(result)
    49  	defer stop()
    50  
    51  	// perform the action
    52  	if err := step.Action(ctx); err != nil {
    53  		result.Error = err
    54  		return
    55  	}
    56  
    57  	// wait for all node expectations to either pass, error or timeout
    58  	nodes := make(map[enode.ID]struct{}, len(step.Expect.Nodes))
    59  	for _, id := range step.Expect.Nodes {
    60  		nodes[id] = struct{}{}
    61  	}
    62  	for len(result.Passes) < len(nodes) {
    63  		select {
    64  		case id := <-step.Trigger:
    65  			// skip if we aren't checking the node
    66  			if _, ok := nodes[id]; !ok {
    67  				continue
    68  			}
    69  
    70  			// skip if the node has already passed
    71  			if _, ok := result.Passes[id]; ok {
    72  				continue
    73  			}
    74  
    75  			// run the node expectation check
    76  			pass, err := step.Expect.Check(ctx, id)
    77  			if err != nil {
    78  				result.Error = err
    79  				return
    80  			}
    81  			if pass {
    82  				result.Passes[id] = time.Now()
    83  			}
    84  		case <-ctx.Done():
    85  			result.Error = ctx.Err()
    86  			return
    87  		}
    88  	}
    89  
    90  	return
    91  }
    92  
    93  func (s *Simulation) watchNetwork(result *StepResult) func() {
    94  	stop := make(chan struct{})
    95  	done := make(chan struct{})
    96  	events := make(chan *Event)
    97  	sub := s.network.Events().Subscribe(events)
    98  	go func() {
    99  		defer close(done)
   100  		defer sub.Unsubscribe()
   101  		for {
   102  			select {
   103  			case event := <-events:
   104  				result.NetworkEvents = append(result.NetworkEvents, event)
   105  			case <-stop:
   106  				return
   107  			}
   108  		}
   109  	}()
   110  	return func() {
   111  		close(stop)
   112  		<-done
   113  	}
   114  }
   115  
   116  type Step struct {
   117  	// Action is the action to perform for this step
   118  	Action func(context.Context) error
   119  
   120  	// Trigger is a channel which receives node ids and triggers an
   121  	// expectation check for that node
   122  	Trigger chan enode.ID
   123  
   124  	// Expect is the expectation to wait for when performing this step
   125  	Expect *Expectation
   126  }
   127  
   128  type Expectation struct {
   129  	// Nodes is a list of nodes to check
   130  	Nodes []enode.ID
   131  
   132  	// Check checks whether a given node meets the expectation
   133  	Check func(context.Context, enode.ID) (bool, error)
   134  }
   135  
   136  func newStepResult() *StepResult {
   137  	return &StepResult{
   138  		Passes: make(map[enode.ID]time.Time),
   139  	}
   140  }
   141  
   142  type StepResult struct {
   143  	// Error is the error encountered whilst running the step
   144  	Error error
   145  
   146  	// StartedAt is the time the step started
   147  	StartedAt time.Time
   148  
   149  	// FinishedAt is the time the step finished
   150  	FinishedAt time.Time
   151  
   152  	// Passes are the timestamps of the successful node expectations
   153  	Passes map[enode.ID]time.Time
   154  
   155  	// NetworkEvents are the network events which occurred during the step
   156  	NetworkEvents []*Event
   157  }