diff --git a/exchange/bitswap/strategy/interface.go b/exchange/bitswap/strategy/interface.go index 8608c52ce..a95ea8bd2 100644 --- a/exchange/bitswap/strategy/interface.go +++ b/exchange/bitswap/strategy/interface.go @@ -24,6 +24,10 @@ type Strategy interface { // MessageSent records sending of message for accounting purposes MessageSent(*peer.Peer, bsmsg.BitSwapMessage) error + + NumBytesSentTo(*peer.Peer) uint64 + + NumBytesReceivedFrom(*peer.Peer) uint64 } type WantList interface { diff --git a/exchange/bitswap/strategy/strategy.go b/exchange/bitswap/strategy/strategy.go index 208811561..406508d6e 100644 --- a/exchange/bitswap/strategy/strategy.go +++ b/exchange/bitswap/strategy/strategy.go @@ -76,6 +76,14 @@ func (s *strategist) MessageSent(p *peer.Peer, m bsmsg.BitSwapMessage) error { return nil } +func (s *strategist) NumBytesSentTo(p *peer.Peer) uint64 { + return s.ledger(p).Accounting.BytesSent +} + +func (s *strategist) NumBytesReceivedFrom(p *peer.Peer) uint64 { + return s.ledger(p).Accounting.BytesRecv +} + // ledger lazily instantiates a ledger func (s *strategist) ledger(p *peer.Peer) *ledger { l, ok := s.ledgerMap[peerKey(p.Key())] diff --git a/exchange/bitswap/strategy/strategy_test.go b/exchange/bitswap/strategy/strategy_test.go index dfa216849..e90bcd4ec 100644 --- a/exchange/bitswap/strategy/strategy_test.go +++ b/exchange/bitswap/strategy/strategy_test.go @@ -1,11 +1,12 @@ package strategy import ( + "strings" "testing" message "github.com/jbenet/go-ipfs/exchange/bitswap/message" - "github.com/jbenet/go-ipfs/peer" - "github.com/jbenet/go-ipfs/util/testutil" + peer "github.com/jbenet/go-ipfs/peer" + testutil "github.com/jbenet/go-ipfs/util/testutil" ) type peerAndStrategist struct { @@ -20,6 +21,38 @@ func newPeerAndStrategist(idStr string) peerAndStrategist { } } +func TestConsistentAccounting(t *testing.T) { + sender := newPeerAndStrategist("Ernie") + receiver := newPeerAndStrategist("Bert") + + // Send messages from Ernie to Bert + for i := 0; i < 1000; i++ { + + m := message.New() + content := []string{"this", "is", "message", "i"} + m.AppendBlock(testutil.NewBlockOrFail(t, strings.Join(content, " "))) + + sender.MessageSent(receiver.Peer, m) + receiver.MessageReceived(sender.Peer, m) + } + + // Ensure sender records the change + if sender.NumBytesSentTo(receiver.Peer) == 0 { + t.Fatal("Sent bytes were not recorded") + } + + // Ensure sender and receiver have the same values + if sender.NumBytesSentTo(receiver.Peer) != receiver.NumBytesReceivedFrom(sender.Peer) { + t.Fatal("Inconsistent book-keeping. Strategies don't agree") + } + + // Ensure sender didn't record receving anything. And that the receiver + // didn't record sending anything + if receiver.NumBytesSentTo(sender.Peer) != 0 || sender.NumBytesReceivedFrom(receiver.Peer) != 0 { + t.Fatal("Bert didn't send bytes to Ernie") + } +} + func TestBlockRecordedAsWantedAfterMessageReceived(t *testing.T) { beggar := newPeerAndStrategist("can't be chooser") chooser := newPeerAndStrategist("chooses JIF")