github.com/etsc3259/etsc@v0.0.0-20190109113336-a9c2c10f9c95/swarm/network/simulation/events.go (about) 1 // Copyright 2018 The go-etsc Authors 2 // This file is part of the go-etsc library. 3 // 4 // The go-etsc 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-etsc 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-etsc library. If not, see <http://www.gnu.org/licenses/>. 16 17 package simulation 18 19 import ( 20 "context" 21 "sync" 22 23 "github.com/ETSC3259/etsc/p2p" 24 "github.com/ETSC3259/etsc/p2p/enode" 25 ) 26 27 // PeerEvent is the type of the channel returned by Simulation.PeerEvents. 28 type PeerEvent struct { 29 // NodeID is the ID of node that the event is caught on. 30 NodeID enode.ID 31 // Event is the event that is caught. 32 Event *p2p.PeerEvent 33 // Error is the error that may have happened during event watching. 34 Error error 35 } 36 37 // PeerEventsFilter defines a filter on PeerEvents to exclude messages with 38 // defined properties. Use PeerEventsFilter methods to set required options. 39 type PeerEventsFilter struct { 40 t *p2p.PeerEventType 41 protocol *string 42 msgCode *uint64 43 } 44 45 // NewPeerEventsFilter returns a new PeerEventsFilter instance. 46 func NewPeerEventsFilter() *PeerEventsFilter { 47 return &PeerEventsFilter{} 48 } 49 50 // Type sets the filter to only one peer event type. 51 func (f *PeerEventsFilter) Type(t p2p.PeerEventType) *PeerEventsFilter { 52 f.t = &t 53 return f 54 } 55 56 // Protocol sets the filter to only one message protocol. 57 func (f *PeerEventsFilter) Protocol(p string) *PeerEventsFilter { 58 f.protocol = &p 59 return f 60 } 61 62 // MsgCode sets the filter to only one msg code. 63 func (f *PeerEventsFilter) MsgCode(c uint64) *PeerEventsFilter { 64 f.msgCode = &c 65 return f 66 } 67 68 // PeerEvents returns a channel of events that are captured by admin peerEvents 69 // subscription nodes with provided NodeIDs. Additional filters can be set to ignore 70 // events that are not relevant. 71 func (s *Simulation) PeerEvents(ctx context.Context, ids []enode.ID, filters ...*PeerEventsFilter) <-chan PeerEvent { 72 eventC := make(chan PeerEvent) 73 74 // wait group to make sure all subscriptions to admin peerEvents are established 75 // before this function returns. 76 var subsWG sync.WaitGroup 77 for _, id := range ids { 78 s.shutdownWG.Add(1) 79 subsWG.Add(1) 80 go func(id enode.ID) { 81 defer s.shutdownWG.Done() 82 83 client, err := s.Net.GetNode(id).Client() 84 if err != nil { 85 subsWG.Done() 86 eventC <- PeerEvent{NodeID: id, Error: err} 87 return 88 } 89 events := make(chan *p2p.PeerEvent) 90 sub, err := client.Subscribe(ctx, "admin", events, "peerEvents") 91 if err != nil { 92 subsWG.Done() 93 eventC <- PeerEvent{NodeID: id, Error: err} 94 return 95 } 96 defer sub.Unsubscribe() 97 98 subsWG.Done() 99 100 for { 101 select { 102 case <-ctx.Done(): 103 if err := ctx.Err(); err != nil { 104 select { 105 case eventC <- PeerEvent{NodeID: id, Error: err}: 106 case <-s.Done(): 107 } 108 } 109 return 110 case <-s.Done(): 111 return 112 case e := <-events: 113 match := len(filters) == 0 // if there are no filters match all events 114 for _, f := range filters { 115 if f.t != nil && *f.t != e.Type { 116 continue 117 } 118 if f.protocol != nil && *f.protocol != e.Protocol { 119 continue 120 } 121 if f.msgCode != nil && e.MsgCode != nil && *f.msgCode != *e.MsgCode { 122 continue 123 } 124 // all filter parameters matched, break the loop 125 match = true 126 break 127 } 128 if match { 129 select { 130 case eventC <- PeerEvent{NodeID: id, Event: e}: 131 case <-ctx.Done(): 132 if err := ctx.Err(); err != nil { 133 select { 134 case eventC <- PeerEvent{NodeID: id, Error: err}: 135 case <-s.Done(): 136 } 137 } 138 return 139 case <-s.Done(): 140 return 141 } 142 } 143 case err := <-sub.Err(): 144 if err != nil { 145 select { 146 case eventC <- PeerEvent{NodeID: id, Error: err}: 147 case <-ctx.Done(): 148 if err := ctx.Err(); err != nil { 149 select { 150 case eventC <- PeerEvent{NodeID: id, Error: err}: 151 case <-s.Done(): 152 } 153 } 154 return 155 case <-s.Done(): 156 return 157 } 158 } 159 } 160 } 161 }(id) 162 } 163 164 // wait all subscriptions 165 subsWG.Wait() 166 return eventC 167 }