diff --git a/bitswap/ledger.go b/bitswap/ledger.go index b8b58e5a665..6ddc0a71107 100644 --- a/bitswap/ledger.go +++ b/bitswap/ledger.go @@ -1,6 +1,7 @@ package bitswap import ( + "sync" "time" peer "github.com/jbenet/go-ipfs/peer" @@ -9,6 +10,7 @@ import ( // Ledger stores the data exchange relationship between two peers. type Ledger struct { + lock sync.RWMutex // Partner is the remote Peer. Partner *peer.Peer @@ -35,16 +37,25 @@ type Ledger struct { type LedgerMap map[u.Key]*Ledger func (l *Ledger) ShouldSend() bool { + l.lock.Lock() + defer l.lock.Unlock() + return l.Strategy(l) } func (l *Ledger) SentBytes(n int) { + l.lock.Lock() + defer l.lock.Unlock() + l.exchangeCount++ l.lastExchange = time.Now() l.Accounting.BytesSent += uint64(n) } func (l *Ledger) ReceivedBytes(n int) { + l.lock.Lock() + defer l.lock.Unlock() + l.exchangeCount++ l.lastExchange = time.Now() l.Accounting.BytesRecv += uint64(n) @@ -52,10 +63,22 @@ func (l *Ledger) ReceivedBytes(n int) { // TODO: this needs to be different. We need timeouts. func (l *Ledger) Wants(k u.Key) { + l.lock.Lock() + defer l.lock.Unlock() + l.wantList[k] = struct{}{} } func (l *Ledger) WantListContains(k u.Key) bool { + l.lock.RLock() + defer l.lock.RUnlock() + _, ok := l.wantList[k] return ok } + +func (l *Ledger) ExchangeCount() uint64 { + l.lock.RLock() + defer l.lock.RUnlock() + return l.exchangeCount +} diff --git a/bitswap/ledger_test.go b/bitswap/ledger_test.go new file mode 100644 index 00000000000..d651d485ff7 --- /dev/null +++ b/bitswap/ledger_test.go @@ -0,0 +1,23 @@ +package bitswap + +import ( + "sync" + "testing" +) + +func TestRaceConditions(t *testing.T) { + const numberOfExpectedExchanges = 10000 + l := new(Ledger) + var wg sync.WaitGroup + for i := 0; i < numberOfExpectedExchanges; i++ { + wg.Add(1) + go func() { + defer wg.Done() + l.ReceivedBytes(1) + }() + } + wg.Wait() + if l.ExchangeCount() != numberOfExpectedExchanges { + t.Fail() + } +}