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  }