strategy.go 3.59 KB
Newer Older
1 2 3 4
package strategy

import (
	"errors"
5
	"sync"
Jeromy's avatar
Jeromy committed
6
	"time"
7 8

	bsmsg "github.com/jbenet/go-ipfs/exchange/bitswap/message"
Brian Tiger Chow's avatar
Brian Tiger Chow committed
9
	peer "github.com/jbenet/go-ipfs/peer"
10 11 12
	u "github.com/jbenet/go-ipfs/util"
)

13 14
var log = u.Logger("strategy")

15 16 17 18 19 20 21 22 23 24
// TODO niceness should be on a per-peer basis. Use-case: Certain peers are
// "trusted" and/or controlled by a single human user. The user may want for
// these peers to exchange data freely
func New(nice bool) Strategy {
	var stratFunc strategyFunc
	if nice {
		stratFunc = yesManStrategy
	} else {
		stratFunc = standardStrategy
	}
25 26
	return &strategist{
		ledgerMap:    ledgerMap{},
27
		strategyFunc: stratFunc,
28 29 30 31
	}
}

type strategist struct {
32
	lock sync.RWMutex
33 34 35 36 37 38 39 40 41 42 43
	ledgerMap
	strategyFunc
}

// LedgerMap lists Ledgers by their Partner key.
type ledgerMap map[peerKey]*ledger

// FIXME share this externally
type peerKey u.Key

// Peers returns a list of peers
44
func (s *strategist) Peers() []peer.Peer {
45 46 47
	s.lock.RLock()
	defer s.lock.RUnlock()

48
	response := make([]peer.Peer, 0)
49 50 51 52 53 54
	for _, ledger := range s.ledgerMap {
		response = append(response, ledger.Partner)
	}
	return response
}

55
func (s *strategist) BlockIsWantedByPeer(k u.Key, p peer.Peer) bool {
56 57 58
	s.lock.RLock()
	defer s.lock.RUnlock()

59 60 61 62
	ledger := s.ledger(p)
	return ledger.WantListContains(k)
}

63
func (s *strategist) ShouldSendBlockToPeer(k u.Key, p peer.Peer) bool {
64 65 66
	s.lock.RLock()
	defer s.lock.RUnlock()

67
	ledger := s.ledger(p)
68 69 70 71 72 73

	// Dont resend blocks
	if _, ok := ledger.sentToPeer[k]; ok {
		return false
	}

74 75 76
	return ledger.ShouldSend()
}

77 78 79 80 81 82 83 84
func (s *strategist) BlockSentToPeer(k u.Key, p peer.Peer) {
	s.lock.Lock()
	defer s.lock.Unlock()

	ledger := s.ledger(p)
	ledger.sentToPeer[k] = struct{}{}
}

85
func (s *strategist) Seed(int64) {
86 87 88
	s.lock.Lock()
	defer s.lock.Unlock()

89 90 91
	// TODO
}

92 93
// MessageReceived performs book-keeping. Returns error if passed invalid
// arguments.
94
func (s *strategist) MessageReceived(p peer.Peer, m bsmsg.BitSwapMessage) error {
95 96 97
	s.lock.Lock()
	defer s.lock.Unlock()

Brian Tiger Chow's avatar
Brian Tiger Chow committed
98 99 100 101 102 103 104
	// TODO find a more elegant way to handle this check
	if p == nil {
		return errors.New("Strategy received nil peer")
	}
	if m == nil {
		return errors.New("Strategy received nil message")
	}
105 106 107 108 109 110 111 112
	l := s.ledger(p)
	for _, key := range m.Wantlist() {
		l.Wants(key)
	}
	for _, block := range m.Blocks() {
		// FIXME extract blocks.NumBytes(block) or block.NumBytes() method
		l.ReceivedBytes(len(block.Data))
	}
113
	return nil
114 115 116 117 118 119 120 121
}

// TODO add contents of m.WantList() to my local wantlist? NB: could introduce
// race conditions where I send a message, but MessageSent gets handled after
// MessageReceived. The information in the local wantlist could become
// inconsistent. Would need to ensure that Sends and acknowledgement of the
// send happen atomically

122
func (s *strategist) MessageSent(p peer.Peer, m bsmsg.BitSwapMessage) error {
123 124 125
	s.lock.Lock()
	defer s.lock.Unlock()

126 127 128 129
	l := s.ledger(p)
	for _, block := range m.Blocks() {
		l.SentBytes(len(block.Data))
	}
130 131 132

	// TODO remove these blocks from peer's want list

133 134 135
	return nil
}

136
func (s *strategist) NumBytesSentTo(p peer.Peer) uint64 {
137 138 139
	s.lock.RLock()
	defer s.lock.RUnlock()

140 141 142
	return s.ledger(p).Accounting.BytesSent
}

143
func (s *strategist) NumBytesReceivedFrom(p peer.Peer) uint64 {
144 145 146
	s.lock.RLock()
	defer s.lock.RUnlock()

147 148 149
	return s.ledger(p).Accounting.BytesRecv
}

150
// ledger lazily instantiates a ledger
151
func (s *strategist) ledger(p peer.Peer) *ledger {
152 153 154 155 156 157 158
	l, ok := s.ledgerMap[peerKey(p.Key())]
	if !ok {
		l = newLedger(p, s.strategyFunc)
		s.ledgerMap[peerKey(p.Key())] = l
	}
	return l
}
Jeromy's avatar
Jeromy committed
159 160 161 162 163 164

func (s *strategist) GetBatchSize() int {
	return 10
}

func (s *strategist) GetRebroadcastDelay() time.Duration {
Jeromy's avatar
Jeromy committed
165
	return time.Second * 5
Jeromy's avatar
Jeromy committed
166
}