github.com/keltia/go-ipfs@v0.3.8-0.20150909044612-210793031c63/exchange/bitswap/wantmanager.go (about) 1 package bitswap 2 3 import ( 4 "sync" 5 "time" 6 7 context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context" 8 key "github.com/ipfs/go-ipfs/blocks/key" 9 engine "github.com/ipfs/go-ipfs/exchange/bitswap/decision" 10 bsmsg "github.com/ipfs/go-ipfs/exchange/bitswap/message" 11 bsnet "github.com/ipfs/go-ipfs/exchange/bitswap/network" 12 wantlist "github.com/ipfs/go-ipfs/exchange/bitswap/wantlist" 13 peer "github.com/ipfs/go-ipfs/p2p/peer" 14 ) 15 16 type WantManager struct { 17 // sync channels for Run loop 18 incoming chan []*bsmsg.Entry 19 connect chan peer.ID // notification channel for new peers connecting 20 disconnect chan peer.ID // notification channel for peers disconnecting 21 22 // synchronized by Run loop, only touch inside there 23 peers map[peer.ID]*msgQueue 24 wl *wantlist.ThreadSafe 25 26 network bsnet.BitSwapNetwork 27 ctx context.Context 28 } 29 30 func NewWantManager(ctx context.Context, network bsnet.BitSwapNetwork) *WantManager { 31 return &WantManager{ 32 incoming: make(chan []*bsmsg.Entry, 10), 33 connect: make(chan peer.ID, 10), 34 disconnect: make(chan peer.ID, 10), 35 peers: make(map[peer.ID]*msgQueue), 36 wl: wantlist.NewThreadSafe(), 37 network: network, 38 ctx: ctx, 39 } 40 } 41 42 type msgPair struct { 43 to peer.ID 44 msg bsmsg.BitSwapMessage 45 } 46 47 type cancellation struct { 48 who peer.ID 49 blk key.Key 50 } 51 52 type msgQueue struct { 53 p peer.ID 54 55 outlk sync.Mutex 56 out bsmsg.BitSwapMessage 57 network bsnet.BitSwapNetwork 58 59 work chan struct{} 60 done chan struct{} 61 } 62 63 func (pm *WantManager) WantBlocks(ks []key.Key) { 64 log.Infof("want blocks: %s", ks) 65 pm.addEntries(ks, false) 66 } 67 68 func (pm *WantManager) CancelWants(ks []key.Key) { 69 pm.addEntries(ks, true) 70 } 71 72 func (pm *WantManager) addEntries(ks []key.Key, cancel bool) { 73 var entries []*bsmsg.Entry 74 for i, k := range ks { 75 entries = append(entries, &bsmsg.Entry{ 76 Cancel: cancel, 77 Entry: wantlist.Entry{ 78 Key: k, 79 Priority: kMaxPriority - i, 80 }, 81 }) 82 } 83 select { 84 case pm.incoming <- entries: 85 case <-pm.ctx.Done(): 86 } 87 } 88 89 func (pm *WantManager) SendBlock(ctx context.Context, env *engine.Envelope) { 90 // Blocks need to be sent synchronously to maintain proper backpressure 91 // throughout the network stack 92 defer env.Sent() 93 94 msg := bsmsg.New(false) 95 msg.AddBlock(env.Block) 96 log.Infof("Sending block %s to %s", env.Peer, env.Block) 97 err := pm.network.SendMessage(ctx, env.Peer, msg) 98 if err != nil { 99 log.Infof("sendblock error: %s", err) 100 } 101 } 102 103 func (pm *WantManager) startPeerHandler(p peer.ID) *msgQueue { 104 _, ok := pm.peers[p] 105 if ok { 106 // TODO: log an error? 107 return nil 108 } 109 110 mq := pm.newMsgQueue(p) 111 112 // new peer, we will want to give them our full wantlist 113 fullwantlist := bsmsg.New(true) 114 for _, e := range pm.wl.Entries() { 115 fullwantlist.AddEntry(e.Key, e.Priority) 116 } 117 mq.out = fullwantlist 118 mq.work <- struct{}{} 119 120 pm.peers[p] = mq 121 go mq.runQueue(pm.ctx) 122 return mq 123 } 124 125 func (pm *WantManager) stopPeerHandler(p peer.ID) { 126 pq, ok := pm.peers[p] 127 if !ok { 128 // TODO: log error? 129 return 130 } 131 132 close(pq.done) 133 delete(pm.peers, p) 134 } 135 136 func (mq *msgQueue) runQueue(ctx context.Context) { 137 for { 138 select { 139 case <-mq.work: // there is work to be done 140 mq.doWork(ctx) 141 case <-mq.done: 142 return 143 case <-ctx.Done(): 144 return 145 } 146 } 147 } 148 149 func (mq *msgQueue) doWork(ctx context.Context) { 150 // allow ten minutes for connections 151 // this includes looking them up in the dht 152 // dialing them, and handshaking 153 conctx, cancel := context.WithTimeout(ctx, time.Minute*10) 154 defer cancel() 155 156 err := mq.network.ConnectTo(conctx, mq.p) 157 if err != nil { 158 log.Infof("cant connect to peer %s: %s", mq.p, err) 159 // TODO: cant connect, what now? 160 return 161 } 162 163 // grab outgoing message 164 mq.outlk.Lock() 165 wlm := mq.out 166 if wlm == nil || wlm.Empty() { 167 mq.outlk.Unlock() 168 return 169 } 170 mq.out = nil 171 mq.outlk.Unlock() 172 173 sendctx, cancel := context.WithTimeout(ctx, time.Minute*5) 174 defer cancel() 175 176 // send wantlist updates 177 err = mq.network.SendMessage(sendctx, mq.p, wlm) 178 if err != nil { 179 log.Infof("bitswap send error: %s", err) 180 // TODO: what do we do if this fails? 181 return 182 } 183 } 184 185 func (pm *WantManager) Connected(p peer.ID) { 186 select { 187 case pm.connect <- p: 188 case <-pm.ctx.Done(): 189 } 190 } 191 192 func (pm *WantManager) Disconnected(p peer.ID) { 193 select { 194 case pm.disconnect <- p: 195 case <-pm.ctx.Done(): 196 } 197 } 198 199 // TODO: use goprocess here once i trust it 200 func (pm *WantManager) Run() { 201 tock := time.NewTicker(rebroadcastDelay.Get()) 202 defer tock.Stop() 203 for { 204 select { 205 case entries := <-pm.incoming: 206 207 // add changes to our wantlist 208 for _, e := range entries { 209 if e.Cancel { 210 pm.wl.Remove(e.Key) 211 } else { 212 pm.wl.Add(e.Key, e.Priority) 213 } 214 } 215 216 // broadcast those wantlist changes 217 for _, p := range pm.peers { 218 p.addMessage(entries) 219 } 220 221 case <-tock.C: 222 // resend entire wantlist every so often (REALLY SHOULDNT BE NECESSARY) 223 var es []*bsmsg.Entry 224 for _, e := range pm.wl.Entries() { 225 es = append(es, &bsmsg.Entry{Entry: e}) 226 } 227 for _, p := range pm.peers { 228 p.outlk.Lock() 229 p.out = bsmsg.New(true) 230 p.outlk.Unlock() 231 232 p.addMessage(es) 233 } 234 case p := <-pm.connect: 235 pm.startPeerHandler(p) 236 case p := <-pm.disconnect: 237 pm.stopPeerHandler(p) 238 case <-pm.ctx.Done(): 239 return 240 } 241 } 242 } 243 244 func (wm *WantManager) newMsgQueue(p peer.ID) *msgQueue { 245 mq := new(msgQueue) 246 mq.done = make(chan struct{}) 247 mq.work = make(chan struct{}, 1) 248 mq.network = wm.network 249 mq.p = p 250 251 return mq 252 } 253 254 func (mq *msgQueue) addMessage(entries []*bsmsg.Entry) { 255 mq.outlk.Lock() 256 defer func() { 257 mq.outlk.Unlock() 258 select { 259 case mq.work <- struct{}{}: 260 default: 261 } 262 }() 263 264 // if we have no message held, or the one we are given is full 265 // overwrite the one we are holding 266 if mq.out == nil { 267 mq.out = bsmsg.New(false) 268 } 269 270 // TODO: add a msg.Combine(...) method 271 // otherwise, combine the one we are holding with the 272 // one passed in 273 for _, e := range entries { 274 if e.Cancel { 275 mq.out.Cancel(e.Key) 276 } else { 277 mq.out.AddEntry(e.Key, e.Priority) 278 } 279 } 280 }