github.com/oskarth/go-ethereum@v1.6.8-0.20191013093314-dac24a9d3494/swarm/network/simulation/node.go (about) 1 // Copyright 2018 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 simulation 18 19 import ( 20 "encoding/json" 21 "errors" 22 "io/ioutil" 23 "math/rand" 24 "os" 25 "time" 26 27 "github.com/ethereum/go-ethereum/log" 28 "github.com/ethereum/go-ethereum/p2p/enode" 29 "github.com/ethereum/go-ethereum/p2p/simulations" 30 "github.com/ethereum/go-ethereum/p2p/simulations/adapters" 31 ) 32 33 // NodeIDs returns NodeIDs for all nodes in the network. 34 func (s *Simulation) NodeIDs() (ids []enode.ID) { 35 nodes := s.Net.GetNodes() 36 ids = make([]enode.ID, len(nodes)) 37 for i, node := range nodes { 38 ids[i] = node.ID() 39 } 40 return ids 41 } 42 43 // UpNodeIDs returns NodeIDs for nodes that are up in the network. 44 func (s *Simulation) UpNodeIDs() (ids []enode.ID) { 45 nodes := s.Net.GetNodes() 46 for _, node := range nodes { 47 if node.Up { 48 ids = append(ids, node.ID()) 49 } 50 } 51 return ids 52 } 53 54 // DownNodeIDs returns NodeIDs for nodes that are stopped in the network. 55 func (s *Simulation) DownNodeIDs() (ids []enode.ID) { 56 nodes := s.Net.GetNodes() 57 for _, node := range nodes { 58 if !node.Up { 59 ids = append(ids, node.ID()) 60 } 61 } 62 return ids 63 } 64 65 // AddNodeOption defines the option that can be passed 66 // to Simulation.AddNode method. 67 type AddNodeOption func(*adapters.NodeConfig) 68 69 // AddNodeWithMsgEvents sets the EnableMsgEvents option 70 // to NodeConfig. 71 func AddNodeWithMsgEvents(enable bool) AddNodeOption { 72 return func(o *adapters.NodeConfig) { 73 o.EnableMsgEvents = enable 74 } 75 } 76 77 // AddNodeWithService specifies a service that should be 78 // started on a node. This option can be repeated as variadic 79 // argument toe AddNode and other add node related methods. 80 // If AddNodeWithService is not specified, all services will be started. 81 func AddNodeWithService(serviceName string) AddNodeOption { 82 return func(o *adapters.NodeConfig) { 83 o.Services = append(o.Services, serviceName) 84 } 85 } 86 87 // AddNode creates a new node with random configuration, 88 // applies provided options to the config and adds the node to network. 89 // By default all services will be started on a node. If one or more 90 // AddNodeWithService option are provided, only specified services will be started. 91 func (s *Simulation) AddNode(opts ...AddNodeOption) (id enode.ID, err error) { 92 conf := adapters.RandomNodeConfig() 93 for _, o := range opts { 94 o(conf) 95 } 96 if len(conf.Services) == 0 { 97 conf.Services = s.serviceNames 98 } 99 node, err := s.Net.NewNodeWithConfig(conf) 100 if err != nil { 101 return id, err 102 } 103 return node.ID(), s.Net.Start(node.ID()) 104 } 105 106 // AddNodes creates new nodes with random configurations, 107 // applies provided options to the config and adds nodes to network. 108 func (s *Simulation) AddNodes(count int, opts ...AddNodeOption) (ids []enode.ID, err error) { 109 ids = make([]enode.ID, 0, count) 110 for i := 0; i < count; i++ { 111 id, err := s.AddNode(opts...) 112 if err != nil { 113 return nil, err 114 } 115 ids = append(ids, id) 116 } 117 return ids, nil 118 } 119 120 // AddNodesAndConnectFull is a helpper method that combines 121 // AddNodes and ConnectNodesFull. Only new nodes will be connected. 122 func (s *Simulation) AddNodesAndConnectFull(count int, opts ...AddNodeOption) (ids []enode.ID, err error) { 123 if count < 2 { 124 return nil, errors.New("count of nodes must be at least 2") 125 } 126 ids, err = s.AddNodes(count, opts...) 127 if err != nil { 128 return nil, err 129 } 130 err = s.ConnectNodesFull(ids) 131 if err != nil { 132 return nil, err 133 } 134 return ids, nil 135 } 136 137 // AddNodesAndConnectChain is a helpper method that combines 138 // AddNodes and ConnectNodesChain. The chain will be continued from the last 139 // added node, if there is one in simulation using ConnectToLastNode method. 140 func (s *Simulation) AddNodesAndConnectChain(count int, opts ...AddNodeOption) (ids []enode.ID, err error) { 141 if count < 2 { 142 return nil, errors.New("count of nodes must be at least 2") 143 } 144 id, err := s.AddNode(opts...) 145 if err != nil { 146 return nil, err 147 } 148 err = s.ConnectToLastNode(id) 149 if err != nil { 150 return nil, err 151 } 152 ids, err = s.AddNodes(count-1, opts...) 153 if err != nil { 154 return nil, err 155 } 156 ids = append([]enode.ID{id}, ids...) 157 err = s.ConnectNodesChain(ids) 158 if err != nil { 159 return nil, err 160 } 161 return ids, nil 162 } 163 164 // AddNodesAndConnectRing is a helpper method that combines 165 // AddNodes and ConnectNodesRing. 166 func (s *Simulation) AddNodesAndConnectRing(count int, opts ...AddNodeOption) (ids []enode.ID, err error) { 167 if count < 2 { 168 return nil, errors.New("count of nodes must be at least 2") 169 } 170 ids, err = s.AddNodes(count, opts...) 171 if err != nil { 172 return nil, err 173 } 174 err = s.ConnectNodesRing(ids) 175 if err != nil { 176 return nil, err 177 } 178 return ids, nil 179 } 180 181 // AddNodesAndConnectStar is a helpper method that combines 182 // AddNodes and ConnectNodesStar. 183 func (s *Simulation) AddNodesAndConnectStar(count int, opts ...AddNodeOption) (ids []enode.ID, err error) { 184 if count < 2 { 185 return nil, errors.New("count of nodes must be at least 2") 186 } 187 ids, err = s.AddNodes(count, opts...) 188 if err != nil { 189 return nil, err 190 } 191 err = s.ConnectNodesStar(ids[0], ids[1:]) 192 if err != nil { 193 return nil, err 194 } 195 return ids, nil 196 } 197 198 //UploadSnapshot uploads a snapshot to the simulation 199 //This method tries to open the json file provided, applies the config to all nodes 200 //and then loads the snapshot into the Simulation network 201 func (s *Simulation) UploadSnapshot(snapshotFile string, opts ...AddNodeOption) error { 202 f, err := os.Open(snapshotFile) 203 if err != nil { 204 return err 205 } 206 defer func() { 207 err := f.Close() 208 if err != nil { 209 log.Error("Error closing snapshot file", "err", err) 210 } 211 }() 212 jsonbyte, err := ioutil.ReadAll(f) 213 if err != nil { 214 return err 215 } 216 var snap simulations.Snapshot 217 err = json.Unmarshal(jsonbyte, &snap) 218 if err != nil { 219 return err 220 } 221 222 //the snapshot probably has the property EnableMsgEvents not set 223 //just in case, set it to true! 224 //(we need this to wait for messages before uploading) 225 for _, n := range snap.Nodes { 226 n.Node.Config.EnableMsgEvents = true 227 n.Node.Config.Services = s.serviceNames 228 for _, o := range opts { 229 o(n.Node.Config) 230 } 231 } 232 233 log.Info("Waiting for p2p connections to be established...") 234 235 //now we can load the snapshot 236 err = s.Net.Load(&snap) 237 if err != nil { 238 return err 239 } 240 log.Info("Snapshot loaded") 241 return nil 242 } 243 244 // SetPivotNode sets the NodeID of the network's pivot node. 245 // Pivot node is just a specific node that should be treated 246 // differently then other nodes in test. SetPivotNode and 247 // PivotNodeID are just a convenient functions to set and 248 // retrieve it. 249 func (s *Simulation) SetPivotNode(id enode.ID) { 250 s.mu.Lock() 251 defer s.mu.Unlock() 252 s.pivotNodeID = &id 253 } 254 255 // PivotNodeID returns NodeID of the pivot node set by 256 // Simulation.SetPivotNode method. 257 func (s *Simulation) PivotNodeID() (id *enode.ID) { 258 s.mu.Lock() 259 defer s.mu.Unlock() 260 return s.pivotNodeID 261 } 262 263 // StartNode starts a node by NodeID. 264 func (s *Simulation) StartNode(id enode.ID) (err error) { 265 return s.Net.Start(id) 266 } 267 268 // StartRandomNode starts a random node. 269 func (s *Simulation) StartRandomNode() (id enode.ID, err error) { 270 n := s.randomDownNode() 271 if n == nil { 272 return id, ErrNodeNotFound 273 } 274 return n.ID, s.Net.Start(n.ID) 275 } 276 277 // StartRandomNodes starts random nodes. 278 func (s *Simulation) StartRandomNodes(count int) (ids []enode.ID, err error) { 279 ids = make([]enode.ID, 0, count) 280 downIDs := s.DownNodeIDs() 281 for i := 0; i < count; i++ { 282 n := s.randomNode(downIDs, ids...) 283 if n == nil { 284 return nil, ErrNodeNotFound 285 } 286 err = s.Net.Start(n.ID) 287 if err != nil { 288 return nil, err 289 } 290 ids = append(ids, n.ID) 291 } 292 return ids, nil 293 } 294 295 // StopNode stops a node by NodeID. 296 func (s *Simulation) StopNode(id enode.ID) (err error) { 297 return s.Net.Stop(id) 298 } 299 300 // StopRandomNode stops a random node. 301 func (s *Simulation) StopRandomNode() (id enode.ID, err error) { 302 n := s.RandomUpNode() 303 if n == nil { 304 return id, ErrNodeNotFound 305 } 306 return n.ID, s.Net.Stop(n.ID) 307 } 308 309 // StopRandomNodes stops random nodes. 310 func (s *Simulation) StopRandomNodes(count int) (ids []enode.ID, err error) { 311 ids = make([]enode.ID, 0, count) 312 upIDs := s.UpNodeIDs() 313 for i := 0; i < count; i++ { 314 n := s.randomNode(upIDs, ids...) 315 if n == nil { 316 return nil, ErrNodeNotFound 317 } 318 err = s.Net.Stop(n.ID) 319 if err != nil { 320 return nil, err 321 } 322 ids = append(ids, n.ID) 323 } 324 return ids, nil 325 } 326 327 // seed the random generator for Simulation.randomNode. 328 func init() { 329 rand.Seed(time.Now().UnixNano()) 330 } 331 332 // RandomUpNode returns a random SimNode that is up. 333 // Arguments are NodeIDs for nodes that should not be returned. 334 func (s *Simulation) RandomUpNode(exclude ...enode.ID) *adapters.SimNode { 335 return s.randomNode(s.UpNodeIDs(), exclude...) 336 } 337 338 // randomDownNode returns a random SimNode that is not up. 339 func (s *Simulation) randomDownNode(exclude ...enode.ID) *adapters.SimNode { 340 return s.randomNode(s.DownNodeIDs(), exclude...) 341 } 342 343 // randomNode returns a random SimNode from the slice of NodeIDs. 344 func (s *Simulation) randomNode(ids []enode.ID, exclude ...enode.ID) *adapters.SimNode { 345 for _, e := range exclude { 346 var i int 347 for _, id := range ids { 348 if id == e { 349 ids = append(ids[:i], ids[i+1:]...) 350 } else { 351 i++ 352 } 353 } 354 } 355 l := len(ids) 356 if l == 0 { 357 return nil 358 } 359 n := s.Net.GetNode(ids[rand.Intn(l)]) 360 node, _ := n.Node.(*adapters.SimNode) 361 return node 362 }