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 }