github.com/zignig/go-ipfs@v0.0.0-20141111235910-c9e5fdf55a52/exchange/bitswap/strategy/strategy.go (about)

     1  package strategy
     2  
     3  import (
     4  	"errors"
     5  	"sync"
     6  
     7  	bsmsg "github.com/jbenet/go-ipfs/exchange/bitswap/message"
     8  	peer "github.com/jbenet/go-ipfs/peer"
     9  	u "github.com/jbenet/go-ipfs/util"
    10  )
    11  
    12  // TODO niceness should be on a per-peer basis. Use-case: Certain peers are
    13  // "trusted" and/or controlled by a single human user. The user may want for
    14  // these peers to exchange data freely
    15  func New(nice bool) Strategy {
    16  	var stratFunc strategyFunc
    17  	if nice {
    18  		stratFunc = yesManStrategy
    19  	} else {
    20  		stratFunc = standardStrategy
    21  	}
    22  	return &strategist{
    23  		ledgerMap:    ledgerMap{},
    24  		strategyFunc: stratFunc,
    25  	}
    26  }
    27  
    28  type strategist struct {
    29  	lock sync.RWMutex
    30  	ledgerMap
    31  	strategyFunc
    32  }
    33  
    34  // LedgerMap lists Ledgers by their Partner key.
    35  type ledgerMap map[peerKey]*ledger
    36  
    37  // FIXME share this externally
    38  type peerKey u.Key
    39  
    40  // Peers returns a list of peers
    41  func (s *strategist) Peers() []peer.Peer {
    42  	s.lock.RLock()
    43  	defer s.lock.RUnlock()
    44  
    45  	response := make([]peer.Peer, 0)
    46  	for _, ledger := range s.ledgerMap {
    47  		response = append(response, ledger.Partner)
    48  	}
    49  	return response
    50  }
    51  
    52  func (s *strategist) BlockIsWantedByPeer(k u.Key, p peer.Peer) bool {
    53  	s.lock.RLock()
    54  	defer s.lock.RUnlock()
    55  
    56  	ledger := s.ledger(p)
    57  	return ledger.WantListContains(k)
    58  }
    59  
    60  func (s *strategist) ShouldSendBlockToPeer(k u.Key, p peer.Peer) bool {
    61  	s.lock.RLock()
    62  	defer s.lock.RUnlock()
    63  
    64  	ledger := s.ledger(p)
    65  	return ledger.ShouldSend()
    66  }
    67  
    68  func (s *strategist) Seed(int64) {
    69  	s.lock.Lock()
    70  	defer s.lock.Unlock()
    71  
    72  	// TODO
    73  }
    74  
    75  func (s *strategist) MessageReceived(p peer.Peer, m bsmsg.BitSwapMessage) error {
    76  	s.lock.Lock()
    77  	defer s.lock.Unlock()
    78  
    79  	// TODO find a more elegant way to handle this check
    80  	if p == nil {
    81  		return errors.New("Strategy received nil peer")
    82  	}
    83  	if m == nil {
    84  		return errors.New("Strategy received nil message")
    85  	}
    86  	l := s.ledger(p)
    87  	for _, key := range m.Wantlist() {
    88  		l.Wants(key)
    89  	}
    90  	for _, block := range m.Blocks() {
    91  		// FIXME extract blocks.NumBytes(block) or block.NumBytes() method
    92  		l.ReceivedBytes(len(block.Data))
    93  	}
    94  	return errors.New("TODO")
    95  }
    96  
    97  // TODO add contents of m.WantList() to my local wantlist? NB: could introduce
    98  // race conditions where I send a message, but MessageSent gets handled after
    99  // MessageReceived. The information in the local wantlist could become
   100  // inconsistent. Would need to ensure that Sends and acknowledgement of the
   101  // send happen atomically
   102  
   103  func (s *strategist) MessageSent(p peer.Peer, m bsmsg.BitSwapMessage) error {
   104  	s.lock.Lock()
   105  	defer s.lock.Unlock()
   106  
   107  	l := s.ledger(p)
   108  	for _, block := range m.Blocks() {
   109  		l.SentBytes(len(block.Data))
   110  	}
   111  
   112  	// TODO remove these blocks from peer's want list
   113  
   114  	return nil
   115  }
   116  
   117  func (s *strategist) NumBytesSentTo(p peer.Peer) uint64 {
   118  	s.lock.RLock()
   119  	defer s.lock.RUnlock()
   120  
   121  	return s.ledger(p).Accounting.BytesSent
   122  }
   123  
   124  func (s *strategist) NumBytesReceivedFrom(p peer.Peer) uint64 {
   125  	s.lock.RLock()
   126  	defer s.lock.RUnlock()
   127  
   128  	return s.ledger(p).Accounting.BytesRecv
   129  }
   130  
   131  // ledger lazily instantiates a ledger
   132  func (s *strategist) ledger(p peer.Peer) *ledger {
   133  	l, ok := s.ledgerMap[peerKey(p.Key())]
   134  	if !ok {
   135  		l = newLedger(p, s.strategyFunc)
   136  		s.ledgerMap[peerKey(p.Key())] = l
   137  	}
   138  	return l
   139  }