github.com/yinchengtsinghua/golang-Eos-dpos-Ethereum@v0.0.0-20190121132951-92cc4225ed8e/swarm/network/simulation/node.go (about)

     1  
     2  //此源码被清华学神尹成大魔王专业翻译分析并修改
     3  //尹成QQ77025077
     4  //尹成微信18510341407
     5  //尹成所在QQ群721929980
     6  //尹成邮箱 yinc13@mails.tsinghua.edu.cn
     7  //尹成毕业于清华大学,微软区块链领域全球最有价值专家
     8  //https://mvp.microsoft.com/zh-cn/PublicProfile/4033620
     9  //
    10  //
    11  //
    12  //
    13  //
    14  //
    15  //
    16  //
    17  //
    18  //
    19  //
    20  //
    21  //
    22  //
    23  //
    24  
    25  package simulation
    26  
    27  import (
    28  	"encoding/json"
    29  	"errors"
    30  	"io/ioutil"
    31  	"math/rand"
    32  	"os"
    33  	"time"
    34  
    35  	"github.com/ethereum/go-ethereum/log"
    36  	"github.com/ethereum/go-ethereum/p2p/discover"
    37  	"github.com/ethereum/go-ethereum/p2p/simulations"
    38  	"github.com/ethereum/go-ethereum/p2p/simulations/adapters"
    39  )
    40  
    41  //
    42  func (s *Simulation) NodeIDs() (ids []discover.NodeID) {
    43  	nodes := s.Net.GetNodes()
    44  	ids = make([]discover.NodeID, len(nodes))
    45  	for i, node := range nodes {
    46  		ids[i] = node.ID()
    47  	}
    48  	return ids
    49  }
    50  
    51  //
    52  func (s *Simulation) UpNodeIDs() (ids []discover.NodeID) {
    53  	nodes := s.Net.GetNodes()
    54  	for _, node := range nodes {
    55  		if node.Up {
    56  			ids = append(ids, node.ID())
    57  		}
    58  	}
    59  	return ids
    60  }
    61  
    62  //
    63  func (s *Simulation) DownNodeIDs() (ids []discover.NodeID) {
    64  	nodes := s.Net.GetNodes()
    65  	for _, node := range nodes {
    66  		if !node.Up {
    67  			ids = append(ids, node.ID())
    68  		}
    69  	}
    70  	return ids
    71  }
    72  
    73  //
    74  //
    75  type AddNodeOption func(*adapters.NodeConfig)
    76  
    77  //
    78  //
    79  func AddNodeWithMsgEvents(enable bool) AddNodeOption {
    80  	return func(o *adapters.NodeConfig) {
    81  		o.EnableMsgEvents = enable
    82  	}
    83  }
    84  
    85  //
    86  //
    87  //
    88  //
    89  func AddNodeWithService(serviceName string) AddNodeOption {
    90  	return func(o *adapters.NodeConfig) {
    91  		o.Services = append(o.Services, serviceName)
    92  	}
    93  }
    94  
    95  //
    96  //
    97  //
    98  //
    99  func (s *Simulation) AddNode(opts ...AddNodeOption) (id discover.NodeID, err error) {
   100  	conf := adapters.RandomNodeConfig()
   101  	for _, o := range opts {
   102  		o(conf)
   103  	}
   104  	if len(conf.Services) == 0 {
   105  		conf.Services = s.serviceNames
   106  	}
   107  	node, err := s.Net.NewNodeWithConfig(conf)
   108  	if err != nil {
   109  		return id, err
   110  	}
   111  	return node.ID(), s.Net.Start(node.ID())
   112  }
   113  
   114  //
   115  //
   116  func (s *Simulation) AddNodes(count int, opts ...AddNodeOption) (ids []discover.NodeID, err error) {
   117  	ids = make([]discover.NodeID, 0, count)
   118  	for i := 0; i < count; i++ {
   119  		id, err := s.AddNode(opts...)
   120  		if err != nil {
   121  			return nil, err
   122  		}
   123  		ids = append(ids, id)
   124  	}
   125  	return ids, nil
   126  }
   127  
   128  //
   129  //
   130  func (s *Simulation) AddNodesAndConnectFull(count int, opts ...AddNodeOption) (ids []discover.NodeID, err error) {
   131  	if count < 2 {
   132  		return nil, errors.New("count of nodes must be at least 2")
   133  	}
   134  	ids, err = s.AddNodes(count, opts...)
   135  	if err != nil {
   136  		return nil, err
   137  	}
   138  	err = s.ConnectNodesFull(ids)
   139  	if err != nil {
   140  		return nil, err
   141  	}
   142  	return ids, nil
   143  }
   144  
   145  //
   146  //
   147  //
   148  func (s *Simulation) AddNodesAndConnectChain(count int, opts ...AddNodeOption) (ids []discover.NodeID, err error) {
   149  	if count < 2 {
   150  		return nil, errors.New("count of nodes must be at least 2")
   151  	}
   152  	id, err := s.AddNode(opts...)
   153  	if err != nil {
   154  		return nil, err
   155  	}
   156  	err = s.ConnectToLastNode(id)
   157  	if err != nil {
   158  		return nil, err
   159  	}
   160  	ids, err = s.AddNodes(count-1, opts...)
   161  	if err != nil {
   162  		return nil, err
   163  	}
   164  	ids = append([]discover.NodeID{id}, ids...)
   165  	err = s.ConnectNodesChain(ids)
   166  	if err != nil {
   167  		return nil, err
   168  	}
   169  	return ids, nil
   170  }
   171  
   172  //
   173  //
   174  func (s *Simulation) AddNodesAndConnectRing(count int, opts ...AddNodeOption) (ids []discover.NodeID, err error) {
   175  	if count < 2 {
   176  		return nil, errors.New("count of nodes must be at least 2")
   177  	}
   178  	ids, err = s.AddNodes(count, opts...)
   179  	if err != nil {
   180  		return nil, err
   181  	}
   182  	err = s.ConnectNodesRing(ids)
   183  	if err != nil {
   184  		return nil, err
   185  	}
   186  	return ids, nil
   187  }
   188  
   189  //
   190  //
   191  func (s *Simulation) AddNodesAndConnectStar(count int, opts ...AddNodeOption) (ids []discover.NodeID, err error) {
   192  	if count < 2 {
   193  		return nil, errors.New("count of nodes must be at least 2")
   194  	}
   195  	ids, err = s.AddNodes(count, opts...)
   196  	if err != nil {
   197  		return nil, err
   198  	}
   199  	err = s.ConnectNodesStar(ids[0], ids[1:])
   200  	if err != nil {
   201  		return nil, err
   202  	}
   203  	return ids, nil
   204  }
   205  
   206  //
   207  //
   208  //
   209  func (s *Simulation) UploadSnapshot(snapshotFile string, opts ...AddNodeOption) error {
   210  	f, err := os.Open(snapshotFile)
   211  	if err != nil {
   212  		return err
   213  	}
   214  	defer func() {
   215  		err := f.Close()
   216  		if err != nil {
   217  			log.Error("Error closing snapshot file", "err", err)
   218  		}
   219  	}()
   220  	jsonbyte, err := ioutil.ReadAll(f)
   221  	if err != nil {
   222  		return err
   223  	}
   224  	var snap simulations.Snapshot
   225  	err = json.Unmarshal(jsonbyte, &snap)
   226  	if err != nil {
   227  		return err
   228  	}
   229  
   230  //
   231  //
   232  //
   233  	for _, n := range snap.Nodes {
   234  		n.Node.Config.EnableMsgEvents = true
   235  		n.Node.Config.Services = s.serviceNames
   236  		for _, o := range opts {
   237  			o(n.Node.Config)
   238  		}
   239  	}
   240  
   241  	log.Info("Waiting for p2p connections to be established...")
   242  
   243  //
   244  	err = s.Net.Load(&snap)
   245  	if err != nil {
   246  		return err
   247  	}
   248  	log.Info("Snapshot loaded")
   249  	return nil
   250  }
   251  
   252  //
   253  //
   254  //
   255  //
   256  //
   257  func (s *Simulation) SetPivotNode(id discover.NodeID) {
   258  	s.mu.Lock()
   259  	defer s.mu.Unlock()
   260  	s.pivotNodeID = &id
   261  }
   262  
   263  //
   264  //
   265  func (s *Simulation) PivotNodeID() (id *discover.NodeID) {
   266  	s.mu.Lock()
   267  	defer s.mu.Unlock()
   268  	return s.pivotNodeID
   269  }
   270  
   271  //
   272  func (s *Simulation) StartNode(id discover.NodeID) (err error) {
   273  	return s.Net.Start(id)
   274  }
   275  
   276  //
   277  func (s *Simulation) StartRandomNode() (id discover.NodeID, err error) {
   278  	n := s.randomDownNode()
   279  	if n == nil {
   280  		return id, ErrNodeNotFound
   281  	}
   282  	return n.ID, s.Net.Start(n.ID)
   283  }
   284  
   285  //
   286  func (s *Simulation) StartRandomNodes(count int) (ids []discover.NodeID, err error) {
   287  	ids = make([]discover.NodeID, 0, count)
   288  	downIDs := s.DownNodeIDs()
   289  	for i := 0; i < count; i++ {
   290  		n := s.randomNode(downIDs, ids...)
   291  		if n == nil {
   292  			return nil, ErrNodeNotFound
   293  		}
   294  		err = s.Net.Start(n.ID)
   295  		if err != nil {
   296  			return nil, err
   297  		}
   298  		ids = append(ids, n.ID)
   299  	}
   300  	return ids, nil
   301  }
   302  
   303  //
   304  func (s *Simulation) StopNode(id discover.NodeID) (err error) {
   305  	return s.Net.Stop(id)
   306  }
   307  
   308  //
   309  func (s *Simulation) StopRandomNode() (id discover.NodeID, err error) {
   310  	n := s.RandomUpNode()
   311  	if n == nil {
   312  		return id, ErrNodeNotFound
   313  	}
   314  	return n.ID, s.Net.Stop(n.ID)
   315  }
   316  
   317  //
   318  func (s *Simulation) StopRandomNodes(count int) (ids []discover.NodeID, err error) {
   319  	ids = make([]discover.NodeID, 0, count)
   320  	upIDs := s.UpNodeIDs()
   321  	for i := 0; i < count; i++ {
   322  		n := s.randomNode(upIDs, ids...)
   323  		if n == nil {
   324  			return nil, ErrNodeNotFound
   325  		}
   326  		err = s.Net.Stop(n.ID)
   327  		if err != nil {
   328  			return nil, err
   329  		}
   330  		ids = append(ids, n.ID)
   331  	}
   332  	return ids, nil
   333  }
   334  
   335  //
   336  func init() {
   337  	rand.Seed(time.Now().UnixNano())
   338  }
   339  
   340  //
   341  //
   342  func (s *Simulation) RandomUpNode(exclude ...discover.NodeID) *adapters.SimNode {
   343  	return s.randomNode(s.UpNodeIDs(), exclude...)
   344  }
   345  
   346  //
   347  func (s *Simulation) randomDownNode(exclude ...discover.NodeID) *adapters.SimNode {
   348  	return s.randomNode(s.DownNodeIDs(), exclude...)
   349  }
   350  
   351  //
   352  func (s *Simulation) randomNode(ids []discover.NodeID, exclude ...discover.NodeID) *adapters.SimNode {
   353  	for _, e := range exclude {
   354  		var i int
   355  		for _, id := range ids {
   356  			if id == e {
   357  				ids = append(ids[:i], ids[i+1:]...)
   358  			} else {
   359  				i++
   360  			}
   361  		}
   362  	}
   363  	l := len(ids)
   364  	if l == 0 {
   365  		return nil
   366  	}
   367  	n := s.Net.GetNode(ids[rand.Intn(l)])
   368  	node, _ := n.Node.(*adapters.SimNode)
   369  	return node
   370  }