Repository: ipfs/go-bitswap
Pull request #495 (2021-06-03 19:56:25)

Fix flaky DontHaveTimeoutManger tests
# Goals

Fix flaky tests for the DontHaveTimeoutManger

# Implementation

Replace clock with a mock, so that we are able to mock out time in tests, modify tests to use mock time and be more predictable.

Also avoid AfterFunc cause of potential problems with channel consumption.

hannahhoward: Part of #493 

Patch:
diff --git a/bitswap.go b/bitswap.go
index b7f763d..ac89043 100644
--- a/bitswap.go
+++ b/bitswap.go
@@ -118,6 +118,12 @@ func WithScoreLedger(scoreLedger deciface.ScoreLedger) Option {
 	}
 }
 
+func SetSimulateDontHavesOnTimeout(send bool) Option {
+	return func(bs *Bitswap) {
+		bs.simulateDontHavesOnTimeout = send
+	}
+}
+
 // New initializes a BitSwap instance that communicates over the provided
 // BitSwapNetwork. This function registers the returned instance as the network
 // delegate. Runs until context is cancelled or bitswap.Close is called.
@@ -149,9 +155,12 @@ func New(parent context.Context, network bsnet.BitSwapNetwork,
 	// has an old version of Bitswap that doesn't support DONT_HAVE messages,
 	// or when no response is received within a timeout.
 	var sm *bssm.SessionManager
+	var bs *Bitswap
 	onDontHaveTimeout := func(p peer.ID, dontHaves []cid.Cid) {
 		// Simulate a message arriving with DONT_HAVEs
-		sm.ReceiveFrom(ctx, p, nil, nil, dontHaves)
+		if bs.simulateDontHavesOnTimeout {
+			sm.ReceiveFrom(ctx, p, nil, nil, dontHaves)
+		}
 	}
 	peerQueueFactory := func(ctx context.Context, p peer.ID) bspm.PeerQueue {
 		return bsmq.New(ctx, p, network, onDontHaveTimeout)
@@ -182,25 +191,26 @@ func New(parent context.Context, network bsnet.BitSwapNetwork,
 	notif := notifications.New()
 	sm = bssm.New(ctx, sessionFactory, sim, sessionPeerManagerFactory, bpm, pm, notif, network.Self())
 
-	bs := &Bitswap{
-		blockstore:              bstore,
-		network:                 network,
-		process:                 px,
-		newBlocks:               make(chan cid.Cid, HasBlockBufferSize),
-		provideKeys:             make(chan cid.Cid, provideKeysBufferSize),
-		pm:                      pm,
-		pqm:                     pqm,
-		sm:                      sm,
-		sim:                     sim,
-		notif:                   notif,
-		counters:                new(counters),
-		dupMetric:               dupHist,
-		allMetric:               allHist,
-		sentHistogram:           sentHistogram,
-		provideEnabled:          true,
-		provSearchDelay:         defaultProvSearchDelay,
-		rebroadcastDelay:        delay.Fixed(time.Minute),
-		engineBstoreWorkerCount: defaulEngineBlockstoreWorkerCount,
+	bs = &Bitswap{
+		blockstore:                 bstore,
+		network:                    network,
+		process:                    px,
+		newBlocks:                  make(chan cid.Cid, HasBlockBufferSize),
+		provideKeys:                make(chan cid.Cid, provideKeysBufferSize),
+		pm:                         pm,
+		pqm:                        pqm,
+		sm:                         sm,
+		sim:                        sim,
+		notif:                      notif,
+		counters:                   new(counters),
+		dupMetric:                  dupHist,
+		allMetric:                  allHist,
+		sentHistogram:              sentHistogram,
+		provideEnabled:             true,
+		provSearchDelay:            defaultProvSearchDelay,
+		rebroadcastDelay:           delay.Fixed(time.Minute),
+		engineBstoreWorkerCount:    defaulEngineBlockstoreWorkerCount,
+		simulateDontHavesOnTimeout: true,
 	}
 
 	// apply functional options before starting and running bitswap
@@ -293,6 +303,9 @@ type Bitswap struct {
 
 	// the score ledger used by the decision engine
 	engineScoreLedger deciface.ScoreLedger
+
+	// whether we should actually simulate dont haves on request timeout
+	simulateDontHavesOnTimeout bool
 }
 
 type counters struct {
diff --git a/go.mod b/go.mod
index faa0386..893606c 100644
--- a/go.mod
+++ b/go.mod
@@ -1,6 +1,7 @@
 module github.com/ipfs/go-bitswap
 
 require (
+	github.com/benbjohnson/clock v1.1.0
 	github.com/cskr/pubsub v1.0.2
 	github.com/gogo/protobuf v1.3.1
 	github.com/google/uuid v1.1.2
diff --git a/go.sum b/go.sum
index 633f37d..1793451 100644
--- a/go.sum
+++ b/go.sum
@@ -7,6 +7,8 @@ github.com/Kubuxu/go-os-helper v0.0.1/go.mod h1:N8B+I7vPCT80IcP58r50u4+gEEcsZETF
 github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU=
 github.com/aead/siphash v1.0.1/go.mod h1:Nywa3cDsYNNK3gaciGTWPwHt0wlpNV15vwmswBAUSII=
 github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6/go.mod h1:grANhF5doyWs3UAsr3K4I6qtAmlQcZDesFNEHPZAzj8=
+github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8=
+github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/btcsuite/btcd v0.0.0-20190213025234-306aecffea32/go.mod h1:DrZx5ec/dmnfpw9KyYoQyYo7d0KEvTkk/5M/vbZjAr8=
 github.com/btcsuite/btcd v0.0.0-20190523000118-16327141da8c/go.mod h1:3J08xEfcugPacsc34/LKRU2yO7YmuT8yt28J8k2+rrI=
 github.com/btcsuite/btcd v0.0.0-20190824003749-130ea5bddde3/go.mod h1:3J08xEfcugPacsc34/LKRU2yO7YmuT8yt28J8k2+rrI=
diff --git a/internal/decision/scoreledger.go b/internal/decision/scoreledger.go
index b9f1dfb..188c998 100644
--- a/internal/decision/scoreledger.go
+++ b/internal/decision/scoreledger.go
@@ -4,6 +4,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/benbjohnson/clock"
 	peer "github.com/libp2p/go-libp2p-core/peer"
 )
 
@@ -55,6 +56,8 @@ type scoreledger struct {
 
 	// the record lock
 	lock sync.RWMutex
+
+	clock clock.Clock
 }
 
 // Receipt is a summary of the ledger for a given peer
@@ -73,7 +76,7 @@ func (l *scoreledger) AddToSentBytes(n int) {
 	l.lock.Lock()
 	defer l.lock.Unlock()
 	l.exchangeCount++
-	l.lastExchange = time.Now()
+	l.lastExchange = l.clock.Now()
 	l.bytesSent += uint64(n)
 }
 
@@ -82,7 +85,7 @@ func (l *scoreledger) AddToReceivedBytes(n int) {
 	l.lock.Lock()
 	defer l.lock.Unlock()
 	l.exchangeCount++
-	l.lastExchange = time.Now()
+	l.lastExchange = l.clock.Now()
 	l.bytesRecv += uint64(n)
 }
 
@@ -114,6 +117,7 @@ type DefaultScoreLedger struct {
 	peerSampleInterval time.Duration
 	// used by the tests to detect when a sample is taken
 	sampleCh chan struct{}
+	clock    clock.Clock
 }
 
 // scoreWorker keeps track of how "useful" our peers are, updating scores in the
@@ -134,7 +138,7 @@ type DefaultScoreLedger struct {
 // adjust it ±25% based on our debt ratio. Peers that have historically been
 // more useful to us than we are to them get the highest score.
 func (dsl *DefaultScoreLedger) scoreWorker() {
-	ticker := time.NewTicker(dsl.peerSampleInterval)
+	ticker := dsl.clock.Ticker(dsl.peerSampleInterval)
 	defer ticker.Stop()
 
 	type update struct {
@@ -236,9 +240,10 @@ func (dsl *DefaultScoreLedger) find(p peer.ID) *scoreledger {
 }
 
 // Returns a new scoreledger.
-func newScoreLedger(p peer.ID) *scoreledger {
+func newScoreLedger(p peer.ID, clock clock.Clock) *scoreledger {
 	return &scoreledger{
 		partner: p,
+		clock:   clock,
 	}
 }
 
@@ -255,7 +260,7 @@ func (dsl *DefaultScoreLedger) findOrCreate(p peer.ID) *scoreledger {
 	defer dsl.lock.Unlock()
 	l, ok := dsl.ledgerMap[p]
 	if !ok {
-		l = newScoreLedger(p)
+		l = newScoreLedger(p, dsl.clock)
 		dsl.ledgerMap[p] = l
 	}
 	return l
@@ -315,7 +320,7 @@ func (dsl *DefaultScoreLedger) PeerConnected(p peer.ID) {
 	defer dsl.lock.Unlock()
 	_, ok := dsl.ledgerMap[p]
 	if !ok {
-		dsl.ledgerMap[p] = newScoreLedger(p)
+		dsl.ledgerMap[p] = newScoreLedger(p, dsl.clock)
 	}
 }
 
@@ -333,14 +338,16 @@ func NewDefaultScoreLedger() *DefaultScoreLedger {
 		ledgerMap:          make(map[peer.ID]*scoreledger),
 		closing:            make(chan struct{}),
 		peerSampleInterval: shortTerm,
+		clock:              clock.New(),
 	}
 }
 
 // Creates a new instance of the default score ledger with testing
 // parameters.
-func NewTestScoreLedger(peerSampleInterval time.Duration, sampleCh chan struct{}) *DefaultScoreLedger {
+func NewTestScoreLedger(peerSampleInterval time.Duration, sampleCh chan struct{}, clock clock.Clock) *DefaultScoreLedger {
 	dsl := NewDefaultScoreLedger()
 	dsl.peerSampleInterval = peerSampleInterval
 	dsl.sampleCh = sampleCh
+	dsl.clock = clock
 	return dsl
 }
diff --git a/internal/messagequeue/donthavetimeoutmgr.go b/internal/messagequeue/donthavetimeoutmgr.go
index 14e70c0..4e3aae8 100644
--- a/internal/messagequeue/donthavetimeoutmgr.go
+++ b/internal/messagequeue/donthavetimeoutmgr.go
@@ -5,6 +5,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/benbjohnson/clock"
 	cid "github.com/ipfs/go-cid"
 	"github.com/libp2p/go-libp2p/p2p/protocol/ping"
 )
@@ -60,6 +61,7 @@ type pendingWant struct {
 // we ping the peer to estimate latency. If we receive a response from the
 // peer we use the response latency.
 type dontHaveTimeoutMgr struct {
+	clock                      clock.Clock
 	ctx                        context.Context
 	shutdown                   func()
 	peerConn                   PeerConnection
@@ -83,14 +85,16 @@ type dontHaveTimeoutMgr struct {
 	// ewma of message latency (time from message sent to response received)
 	messageLatency *latencyEwma
 	// timer used to wait until want at front of queue expires
-	checkForTimeoutsTimer *time.Timer
+	checkForTimeoutsTimer *clock.Timer
+	// used for testing -- timeoutsTriggered when a scheduled dont have timeouts were triggered
+	timeoutsTriggered chan struct{}
 }
 
 // newDontHaveTimeoutMgr creates a new dontHaveTimeoutMgr
 // onDontHaveTimeout is called when pending keys expire (not cancelled before timeout)
 func newDontHaveTimeoutMgr(pc PeerConnection, onDontHaveTimeout func([]cid.Cid)) *dontHaveTimeoutMgr {
 	return newDontHaveTimeoutMgrWithParams(pc, onDontHaveTimeout, dontHaveTimeout, maxTimeout,
-		pingLatencyMultiplier, messageLatencyMultiplier, maxExpectedWantProcessTime)
+		pingLatencyMultiplier, messageLatencyMultiplier, maxExpectedWantProcessTime, clock.New(), nil)
 }
 
 // newDontHaveTimeoutMgrWithParams is used by the tests
@@ -101,10 +105,13 @@ func newDontHaveTimeoutMgrWithParams(
 	maxTimeout time.Duration,
 	pingLatencyMultiplier int,
 	messageLatencyMultiplier int,
-	maxExpectedWantProcessTime time.Duration) *dontHaveTimeoutMgr {
+	maxExpectedWantProcessTime time.Duration,
+	clock clock.Clock,
+	timeoutsTriggered chan struct{}) *dontHaveTimeoutMgr {
 
 	ctx, shutdown := context.WithCancel(context.Background())
 	mqp := &dontHaveTimeoutMgr{
+		clock:                      clock,
 		ctx:                        ctx,
 		shutdown:                   shutdown,
 		peerConn:                   pc,
@@ -117,6 +124,7 @@ func newDontHaveTimeoutMgrWithParams(
 		messageLatencyMultiplier:   messageLatencyMultiplier,
 		maxExpectedWantProcessTime: maxExpectedWantProcessTime,
 		onDontHaveTimeout:          onDontHaveTimeout,
+		timeoutsTriggered:          timeoutsTriggered,
 	}
 
 	return mqp
@@ -214,6 +222,7 @@ func (dhtm *dontHaveTimeoutMgr) measurePingLatency() {
 // checkForTimeouts checks pending wants to see if any are over the timeout.
 // Note: this function should only be called within the lock.
 func (dhtm *dontHaveTimeoutMgr) checkForTimeouts() {
+
 	if len(dhtm.wantQueue) == 0 {
 		return
 	}
@@ -228,7 +237,7 @@ func (dhtm *dontHaveTimeoutMgr) checkForTimeouts() {
 		if pw.active {
 			// The queue is in order from earliest to latest, so if we
 			// didn't find an expired entry we can stop iterating
-			if time.Since(pw.sent) < dhtm.timeout {
+			if dhtm.clock.Since(pw.sent) < dhtm.timeout {
 				break
 			}
 
@@ -259,20 +268,29 @@ func (dhtm *dontHaveTimeoutMgr) checkForTimeouts() {
 	// Schedule the next check for the moment when the oldest pending want will
 	// timeout
 	oldestStart := dhtm.wantQueue[0].sent
-	until := time.Until(oldestStart.Add(dhtm.timeout))
+	until := oldestStart.Add(dhtm.timeout).Sub(dhtm.clock.Now())
 	if dhtm.checkForTimeoutsTimer == nil {
-		dhtm.checkForTimeoutsTimer = time.AfterFunc(until, func() {
-			dhtm.lk.Lock()
-			defer dhtm.lk.Unlock()
-
-			dhtm.checkForTimeouts()
-		})
+		dhtm.checkForTimeoutsTimer = dhtm.clock.Timer(until)
+		go dhtm.consumeTimeouts()
 	} else {
 		dhtm.checkForTimeoutsTimer.Stop()
 		dhtm.checkForTimeoutsTimer.Reset(until)
 	}
 }
 
+func (dhtm *dontHaveTimeoutMgr) consumeTimeouts() {
+	for {
+		select {
+		case <-dhtm.ctx.Done():
+			return
+		case <-dhtm.checkForTimeoutsTimer.C:
+			dhtm.lk.Lock()
+			dhtm.checkForTimeouts()
+			dhtm.lk.Unlock()
+		}
+	}
+}
+
 // AddPending adds the given keys that will expire if not cancelled before
 // the timeout
 func (dhtm *dontHaveTimeoutMgr) AddPending(ks []cid.Cid) {
@@ -280,7 +298,7 @@ func (dhtm *dontHaveTimeoutMgr) AddPending(ks []cid.Cid) {
 		return
 	}
 
-	start := time.Now()
+	start := dhtm.clock.Now()
 
 	dhtm.lk.Lock()
 	defer dhtm.lk.Unlock()
@@ -331,6 +349,11 @@ func (dhtm *dontHaveTimeoutMgr) fireTimeout(pending []cid.Cid) {
 
 	// Fire the timeout
 	dhtm.onDontHaveTimeout(pending)
+
+	// signal a timeout fired
+	if dhtm.timeoutsTriggered != nil {
+		dhtm.timeoutsTriggered <- struct{}{}
+	}
 }
 
 // calculateTimeoutFromPingLatency calculates a reasonable timeout derived from latency
diff --git a/network/ipfs_impl.go b/network/ipfs_impl.go
index fc48ef6..b05ce55 100644
--- a/network/ipfs_impl.go
+++ b/network/ipfs_impl.go
@@ -403,8 +403,8 @@ func (bsnet *impl) handleNewStream(s network.Stream) {
 		ctx := context.Background()
 		log.Debugf("bitswap net handleNewStream from %s", s.Conn().RemotePeer())
 		bsnet.connectEvtMgr.OnMessage(s.Conn().RemotePeer())
-		bsnet.receiver.ReceiveMessage(ctx, p, received)
 		atomic.AddUint64(&bsnet.stats.MessagesRecvd, 1)
+		bsnet.receiver.ReceiveMessage(ctx, p, received)
 	}
 }
 


Test patch:
diff --git a/bitswap_with_sessions_test.go b/bitswap_with_sessions_test.go
index f710879..4417453 100644
--- a/bitswap_with_sessions_test.go
+++ b/bitswap_with_sessions_test.go
@@ -74,7 +74,7 @@ func TestSessionBetweenPeers(t *testing.T) {
 	defer cancel()
 
 	vnet := tn.VirtualNetwork(mockrouting.NewServer(), delay.Fixed(time.Millisecond))
-	ig := testinstance.NewTestInstanceGenerator(vnet, nil, nil)
+	ig := testinstance.NewTestInstanceGenerator(vnet, nil, []bitswap.Option{bitswap.SetSimulateDontHavesOnTimeout(false)})
 	defer ig.Close()
 	bgen := blocksutil.NewBlockGenerator()
 
diff --git a/internal/decision/engine_test.go b/internal/decision/engine_test.go
index ac370d0..f7a7525 100644
--- a/internal/decision/engine_test.go
+++ b/internal/decision/engine_test.go
@@ -10,6 +10,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/benbjohnson/clock"
 	"github.com/ipfs/go-bitswap/internal/testutil"
 	message "github.com/ipfs/go-bitswap/message"
 	pb "github.com/ipfs/go-bitswap/message/pb"
@@ -91,13 +92,13 @@ type engineSet struct {
 }
 
 func newTestEngine(ctx context.Context, idStr string) engineSet {
-	return newTestEngineWithSampling(ctx, idStr, shortTerm, nil)
+	return newTestEngineWithSampling(ctx, idStr, shortTerm, nil, clock.New())
 }
 
-func newTestEngineWithSampling(ctx context.Context, idStr string, peerSampleInterval time.Duration, sampleCh chan struct{}) engineSet {
+func newTestEngineWithSampling(ctx context.Context, idStr string, peerSampleInterval time.Duration, sampleCh chan struct{}, clock clock.Clock) engineSet {
 	fpt := &fakePeerTagger{}
 	bs := blockstore.NewBlockstore(dssync.MutexWrap(ds.NewMapDatastore()))
-	e := newEngine(bs, 4, fpt, "localhost", 0, NewTestScoreLedger(peerSampleInterval, sampleCh))
+	e := newEngine(bs, 4, fpt, "localhost", 0, NewTestScoreLedger(peerSampleInterval, sampleCh, clock))
 	e.StartWorkers(ctx, process.WithTeardown(func() error { return nil }))
 	return engineSet{
 		Peer: peer.ID(idStr),
@@ -184,7 +185,7 @@ func peerIsPartner(p peer.ID, e *Engine) bool {
 
 func TestOutboxClosedWhenEngineClosed(t *testing.T) {
 	t.SkipNow() // TODO implement *Engine.Close
-	e := newEngine(blockstore.NewBlockstore(dssync.MutexWrap(ds.NewMapDatastore())), 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil))
+	e := newEngine(blockstore.NewBlockstore(dssync.MutexWrap(ds.NewMapDatastore())), 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil, clock.New()))
 	e.StartWorkers(context.Background(), process.WithTeardown(func() error { return nil }))
 	var wg sync.WaitGroup
 	wg.Add(1)
@@ -512,7 +513,7 @@ func TestPartnerWantHaveWantBlockNonActive(t *testing.T) {
 		testCases = onlyTestCases
 	}
 
-	e := newEngine(bs, 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil))
+	e := newEngine(bs, 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil, clock.New()))
 	e.StartWorkers(context.Background(), process.WithTeardown(func() error { return nil }))
 	for i, testCase := range testCases {
 		t.Logf("Test case %d:", i)
@@ -668,7 +669,7 @@ func TestPartnerWantHaveWantBlockActive(t *testing.T) {
 		testCases = onlyTestCases
 	}
 
-	e := newEngine(bs, 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil))
+	e := newEngine(bs, 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil, clock.New()))
 	e.StartWorkers(context.Background(), process.WithTeardown(func() error { return nil }))
 
 	var next envChan
@@ -853,7 +854,7 @@ func TestPartnerWantsThenCancels(t *testing.T) {
 	ctx := context.Background()
 	for i := 0; i < numRounds; i++ {
 		expected := make([][]string, 0, len(testcases))
-		e := newEngine(bs, 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil))
+		e := newEngine(bs, 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil, clock.New()))
 		e.StartWorkers(ctx, process.WithTeardown(func() error { return nil }))
 		for _, testcase := range testcases {
 			set := testcase[0]
@@ -878,7 +879,7 @@ func TestSendReceivedBlocksToPeersThatWantThem(t *testing.T) {
 	partner := libp2ptest.RandPeerIDFatal(t)
 	otherPeer := libp2ptest.RandPeerIDFatal(t)
 
-	e := newEngine(bs, 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil))
+	e := newEngine(bs, 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil, clock.New()))
 	e.StartWorkers(context.Background(), process.WithTeardown(func() error { return nil }))
 
 	blks := testutil.GenerateBlocksOfSize(4, 8*1024)
@@ -922,7 +923,7 @@ func TestSendDontHave(t *testing.T) {
 	partner := libp2ptest.RandPeerIDFatal(t)
 	otherPeer := libp2ptest.RandPeerIDFatal(t)
 
-	e := newEngine(bs, 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil))
+	e := newEngine(bs, 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil, clock.New()))
 	e.StartWorkers(context.Background(), process.WithTeardown(func() error { return nil }))
 
 	blks := testutil.GenerateBlocksOfSize(4, 8*1024)
@@ -986,7 +987,7 @@ func TestWantlistForPeer(t *testing.T) {
 	partner := libp2ptest.RandPeerIDFatal(t)
 	otherPeer := libp2ptest.RandPeerIDFatal(t)
 
-	e := newEngine(bs, 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil))
+	e := newEngine(bs, 4, &fakePeerTagger{}, "localhost", 0, NewTestScoreLedger(shortTerm, nil, clock.New()))
 	e.StartWorkers(context.Background(), process.WithTeardown(func() error { return nil }))
 
 	blks := testutil.GenerateBlocksOfSize(4, 8*1024)
@@ -1044,13 +1045,15 @@ func TestTaggingPeers(t *testing.T) {
 }
 
 func TestTaggingUseful(t *testing.T) {
-	peerSampleInterval := 1 * time.Millisecond
+	peerSampleIntervalHalf := 10 * time.Millisecond
 
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 
 	sampleCh := make(chan struct{})
-	me := newTestEngineWithSampling(ctx, "engine", peerSampleInterval, sampleCh)
+	mockClock := clock.NewMock()
+	me := newTestEngineWithSampling(ctx, "engine", peerSampleIntervalHalf*2, sampleCh, mockClock)
+	mockClock.Add(1 * time.Millisecond)
 	friend := peer.ID("friend")
 
 	block := blocks.NewBlock([]byte("foobar"))
@@ -1061,18 +1064,18 @@ func TestTaggingUseful(t *testing.T) {
 		if untagged := me.PeerTagger.count(me.Engine.tagUseful); untagged != 0 {
 			t.Fatalf("%d peers should be untagged but weren't", untagged)
 		}
-
+		mockClock.Add(peerSampleIntervalHalf)
 		me.Engine.MessageSent(friend, msg)
 
-		for j := 0; j < 2; j++ {
-			<-sampleCh
-		}
+		mockClock.Add(peerSampleIntervalHalf)
+		<-sampleCh
 
 		if tagged := me.PeerTagger.count(me.Engine.tagUseful); tagged != 1 {
 			t.Fatalf("1 peer should be tagged, but %d were", tagged)
 		}
 
 		for j := 0; j < longTermRatio; j++ {
+			mockClock.Add(peerSampleIntervalHalf * 2)
 			<-sampleCh
 		}
 	}
@@ -1082,6 +1085,7 @@ func TestTaggingUseful(t *testing.T) {
 	}
 
 	for j := 0; j < longTermRatio; j++ {
+		mockClock.Add(peerSampleIntervalHalf * 2)
 		<-sampleCh
 	}
 
@@ -1090,6 +1094,7 @@ func TestTaggingUseful(t *testing.T) {
 	}
 
 	for j := 0; j < longTermRatio; j++ {
+		mockClock.Add(peerSampleIntervalHalf * 2)
 		<-sampleCh
 	}
 
diff --git a/internal/messagequeue/donthavetimeoutmgr_test.go b/internal/messagequeue/donthavetimeoutmgr_test.go
index cc0ebb3..61023f0 100644
--- a/internal/messagequeue/donthavetimeoutmgr_test.go
+++ b/internal/messagequeue/donthavetimeoutmgr_test.go
@@ -7,6 +7,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/benbjohnson/clock"
 	"github.com/ipfs/go-bitswap/internal/testutil"
 	cid "github.com/ipfs/go-cid"
 	"github.com/libp2p/go-libp2p/p2p/protocol/ping"
@@ -16,10 +17,13 @@ type mockPeerConn struct {
 	err       error
 	latency   time.Duration
 	latencies []time.Duration
+	clock     clock.Clock
+	pinged    chan struct{}
 }
 
 func (pc *mockPeerConn) Ping(ctx context.Context) ping.Result {
-	timer := time.NewTimer(pc.latency)
+	timer := pc.clock.Timer(pc.latency)
+	pc.pinged <- struct{}{}
 	select {
 	case <-timer.C:
 		if pc.err != nil {
@@ -75,19 +79,21 @@ func TestDontHaveTimeoutMgrTimeout(t *testing.T) {
 	latMultiplier := 2
 	expProcessTime := 5 * time.Millisecond
 	expectedTimeout := expProcessTime + latency*time.Duration(latMultiplier)
-	pc := &mockPeerConn{latency: latency}
+	clock := clock.NewMock()
+	pinged := make(chan struct{})
+	pc := &mockPeerConn{latency: latency, clock: clock, pinged: pinged}
 	tr := timeoutRecorder{}
-
+	timeoutsTriggered := make(chan struct{})
 	dhtm := newDontHaveTimeoutMgrWithParams(pc, tr.onTimeout,
-		dontHaveTimeout, maxTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime)
+		dontHaveTimeout, maxTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime, clock, timeoutsTriggered)
 	dhtm.Start()
 	defer dhtm.Shutdown()
-
+	<-pinged
 	// Add first set of keys
 	dhtm.AddPending(firstks)
 
 	// Wait for less than the expected timeout
-	time.Sleep(expectedTimeout - 10*time.Millisecond)
+	clock.Add(expectedTimeout - 10*time.Millisecond)
 
 	// At this stage no keys should have timed out
 	if tr.timedOutCount() > 0 {
@@ -98,18 +104,21 @@ func TestDontHaveTimeoutMgrTimeout(t *testing.T) {
 	dhtm.AddPending(secondks)
 
 	// Wait until after the expected timeout
-	time.Sleep(20 * time.Millisecond)
+	clock.Add(20 * time.Millisecond)
+
+	<-timeoutsTriggered
 
 	// At this stage first set of keys should have timed out
 	if tr.timedOutCount() != len(firstks) {
 		t.Fatal("expected timeout", tr.timedOutCount(), len(firstks))
 	}
-
 	// Clear the recorded timed out keys
 	tr.clear()
 
 	// Sleep until the second set of keys should have timed out
-	time.Sleep(expectedTimeout + 10*time.Millisecond)
+	clock.Add(expectedTimeout + 10*time.Millisecond)
+
+	<-timeoutsTriggered
 
 	// At this stage all keys should have timed out. The second set included
 	// the first set of keys, but they were added before the first set timed
@@ -125,24 +134,29 @@ func TestDontHaveTimeoutMgrCancel(t *testing.T) {
 	latMultiplier := 1
 	expProcessTime := time.Duration(0)
 	expectedTimeout := latency
-	pc := &mockPeerConn{latency: latency}
+	clock := clock.NewMock()
+	pinged := make(chan struct{})
+	pc := &mockPeerConn{latency: latency, clock: clock, pinged: pinged}
 	tr := timeoutRecorder{}
-
+	timeoutsTriggered := make(chan struct{})
 	dhtm := newDontHaveTimeoutMgrWithParams(pc, tr.onTimeout,
-		dontHaveTimeout, maxTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime)
+		dontHaveTimeout, maxTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime, clock, timeoutsTriggered)
 	dhtm.Start()
 	defer dhtm.Shutdown()
+	<-pinged
 
 	// Add keys
 	dhtm.AddPending(ks)
-	time.Sleep(5 * time.Millisecond)
+	clock.Add(5 * time.Millisecond)
 
 	// Cancel keys
 	cancelCount := 1
 	dhtm.CancelPending(ks[:cancelCount])
 
 	// Wait for the expected timeout
-	time.Sleep(expectedTimeout)
+	clock.Add(expectedTimeout)
+
+	<-timeoutsTriggered
 
 	// At this stage all non-cancelled keys should have timed out
 	if tr.timedOutCount() != len(ks)-cancelCount {
@@ -156,30 +170,36 @@ func TestDontHaveTimeoutWantCancelWant(t *testing.T) {
 	latMultiplier := 1
 	expProcessTime := time.Duration(0)
 	expectedTimeout := latency
-	pc := &mockPeerConn{latency: latency}
+	clock := clock.NewMock()
+	pinged := make(chan struct{})
+	pc := &mockPeerConn{latency: latency, clock: clock, pinged: pinged}
 	tr := timeoutRecorder{}
+	timeoutsTriggered := make(chan struct{})
 
 	dhtm := newDontHaveTimeoutMgrWithParams(pc, tr.onTimeout,
-		dontHaveTimeout, maxTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime)
+		dontHaveTimeout, maxTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime, clock, timeoutsTriggered)
 	dhtm.Start()
 	defer dhtm.Shutdown()
+	<-pinged
 
 	// Add keys
 	dhtm.AddPending(ks)
 
 	// Wait for a short time
-	time.Sleep(expectedTimeout - 10*time.Millisecond)
+	clock.Add(expectedTimeout - 10*time.Millisecond)
 
 	// Cancel two keys
 	dhtm.CancelPending(ks[:2])
 
-	time.Sleep(5 * time.Millisecond)
+	clock.Add(5 * time.Millisecond)
 
 	// Add back one cancelled key
 	dhtm.AddPending(ks[:1])
 
 	// Wait till after initial timeout
-	time.Sleep(10 * time.Millisecond)
+	clock.Add(10 * time.Millisecond)
+
+	<-timeoutsTriggered
 
 	// At this stage only the key that was never cancelled should have timed out
 	if tr.timedOutCount() != 1 {
@@ -187,7 +207,9 @@ func TestDontHaveTimeoutWantCancelWant(t *testing.T) {
 	}
 
 	// Wait till after added back key should time out
-	time.Sleep(latency)
+	clock.Add(latency)
+
+	<-timeoutsTriggered
 
 	// At this stage the key that was added back should also have timed out
 	if tr.timedOutCount() != 2 {
@@ -200,13 +222,17 @@ func TestDontHaveTimeoutRepeatedAddPending(t *testing.T) {
 	latency := time.Millisecond * 5
 	latMultiplier := 1
 	expProcessTime := time.Duration(0)
-	pc := &mockPeerConn{latency: latency}
+	clock := clock.NewMock()
+	pinged := make(chan struct{})
+	pc := &mockPeerConn{latency: latency, clock: clock, pinged: pinged}
 	tr := timeoutRecorder{}
+	timeoutsTriggered := make(chan struct{})
 
 	dhtm := newDontHaveTimeoutMgrWithParams(pc, tr.onTimeout,
-		dontHaveTimeout, maxTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime)
+		dontHaveTimeout, maxTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime, clock, timeoutsTriggered)
 	dhtm.Start()
 	defer dhtm.Shutdown()
+	<-pinged
 
 	// Add keys repeatedly
 	for _, c := range ks {
@@ -214,7 +240,9 @@ func TestDontHaveTimeoutRepeatedAddPending(t *testing.T) {
 	}
 
 	// Wait for the expected timeout
-	time.Sleep(latency + 5*time.Millisecond)
+	clock.Add(latency + 5*time.Millisecond)
+
+	<-timeoutsTriggered
 
 	// At this stage all keys should have timed out
 	if tr.timedOutCount() != len(ks) {
@@ -228,14 +256,17 @@ func TestDontHaveTimeoutMgrMessageLatency(t *testing.T) {
 	latMultiplier := 1
 	expProcessTime := time.Duration(0)
 	msgLatencyMultiplier := 1
-	pc := &mockPeerConn{latency: latency}
+	clock := clock.NewMock()
+	pinged := make(chan struct{})
+	pc := &mockPeerConn{latency: latency, clock: clock, pinged: pinged}
 	tr := timeoutRecorder{}
+	timeoutsTriggered := make(chan struct{})
 
 	dhtm := newDontHaveTimeoutMgrWithParams(pc, tr.onTimeout,
-		dontHaveTimeout, maxTimeout, latMultiplier, msgLatencyMultiplier, expProcessTime)
+		dontHaveTimeout, maxTimeout, latMultiplier, msgLatencyMultiplier, expProcessTime, clock, timeoutsTriggered)
 	dhtm.Start()
 	defer dhtm.Shutdown()
-
+	<-pinged
 	// Add keys
 	dhtm.AddPending(ks)
 
@@ -245,7 +276,7 @@ func TestDontHaveTimeoutMgrMessageLatency(t *testing.T) {
 	// = 40ms
 
 	// Wait for less than the expected timeout
-	time.Sleep(25 * time.Millisecond)
+	clock.Add(25 * time.Millisecond)
 
 	// Receive two message latency updates
 	dhtm.UpdateMessageLatency(time.Millisecond * 20)
@@ -259,7 +290,9 @@ func TestDontHaveTimeoutMgrMessageLatency(t *testing.T) {
 	// the keys should have timed out
 
 	// Give the queue some time to process the updates
-	time.Sleep(5 * time.Millisecond)
+	clock.Add(5 * time.Millisecond)
+
+	<-timeoutsTriggered
 
 	if tr.timedOutCount() != len(ks) {
 		t.Fatal("expected keys to timeout")
@@ -268,16 +301,19 @@ func TestDontHaveTimeoutMgrMessageLatency(t *testing.T) {
 
 func TestDontHaveTimeoutMgrMessageLatencyMax(t *testing.T) {
 	ks := testutil.GenerateCids(2)
-	pc := &mockPeerConn{latency: time.Second} // ignored
+	clock := clock.NewMock()
+	pinged := make(chan struct{})
+	pc := &mockPeerConn{latency: time.Second, clock: clock, pinged: pinged}
 	tr := timeoutRecorder{}
 	msgLatencyMultiplier := 1
 	testMaxTimeout := time.Millisecond * 10
+	timeoutsTriggered := make(chan struct{})
 
 	dhtm := newDontHaveTimeoutMgrWithParams(pc, tr.onTimeout,
-		dontHaveTimeout, testMaxTimeout, pingLatencyMultiplier, msgLatencyMultiplier, maxExpectedWantProcessTime)
+		dontHaveTimeout, testMaxTimeout, pingLatencyMultiplier, msgLatencyMultiplier, maxExpectedWantProcessTime, clock, timeoutsTriggered)
 	dhtm.Start()
 	defer dhtm.Shutdown()
-
+	<-pinged
 	// Add keys
 	dhtm.AddPending(ks)
 
@@ -286,7 +322,9 @@ func TestDontHaveTimeoutMgrMessageLatencyMax(t *testing.T) {
 	dhtm.UpdateMessageLatency(testMaxTimeout * 4)
 
 	// Sleep until just after the maximum timeout
-	time.Sleep(testMaxTimeout + 5*time.Millisecond)
+	clock.Add(testMaxTimeout + 5*time.Millisecond)
+
+	<-timeoutsTriggered
 
 	// Keys should have timed out
 	if tr.timedOutCount() != len(ks) {
@@ -302,18 +340,22 @@ func TestDontHaveTimeoutMgrUsesDefaultTimeoutIfPingError(t *testing.T) {
 	defaultTimeout := 10 * time.Millisecond
 	expectedTimeout := expProcessTime + defaultTimeout
 	tr := timeoutRecorder{}
-	pc := &mockPeerConn{latency: latency, err: fmt.Errorf("ping error")}
+	clock := clock.NewMock()
+	pinged := make(chan struct{})
+	pc := &mockPeerConn{latency: latency, clock: clock, pinged: pinged, err: fmt.Errorf("ping error")}
+	timeoutsTriggered := make(chan struct{})
 
 	dhtm := newDontHaveTimeoutMgrWithParams(pc, tr.onTimeout,
-		defaultTimeout, dontHaveTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime)
+		defaultTimeout, dontHaveTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime, clock, timeoutsTriggered)
 	dhtm.Start()
 	defer dhtm.Shutdown()
+	<-pinged
 
 	// Add keys
 	dhtm.AddPending(ks)
 
 	// Sleep for less than the expected timeout
-	time.Sleep(expectedTimeout - 5*time.Millisecond)
+	clock.Add(expectedTimeout - 5*time.Millisecond)
 
 	// At this stage no timeout should have happened yet
 	if tr.timedOutCount() > 0 {
@@ -321,7 +363,9 @@ func TestDontHaveTimeoutMgrUsesDefaultTimeoutIfPingError(t *testing.T) {
 	}
 
 	// Sleep until after the expected timeout
-	time.Sleep(10 * time.Millisecond)
+	clock.Add(10 * time.Millisecond)
+
+	<-timeoutsTriggered
 
 	// Now the keys should have timed out
 	if tr.timedOutCount() != len(ks) {
@@ -335,19 +379,23 @@ func TestDontHaveTimeoutMgrUsesDefaultTimeoutIfLatencyLonger(t *testing.T) {
 	latMultiplier := 1
 	expProcessTime := time.Duration(0)
 	defaultTimeout := 10 * time.Millisecond
+	clock := clock.NewMock()
+	pinged := make(chan struct{})
+	pc := &mockPeerConn{latency: latency, clock: clock, pinged: pinged}
 	tr := timeoutRecorder{}
-	pc := &mockPeerConn{latency: latency}
+	timeoutsTriggered := make(chan struct{})
 
 	dhtm := newDontHaveTimeoutMgrWithParams(pc, tr.onTimeout,
-		defaultTimeout, dontHaveTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime)
+		defaultTimeout, dontHaveTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime, clock, timeoutsTriggered)
 	dhtm.Start()
 	defer dhtm.Shutdown()
+	<-pinged
 
 	// Add keys
 	dhtm.AddPending(ks)
 
 	// Sleep for less than the default timeout
-	time.Sleep(defaultTimeout - 5*time.Millisecond)
+	clock.Add(defaultTimeout - 5*time.Millisecond)
 
 	// At this stage no timeout should have happened yet
 	if tr.timedOutCount() > 0 {
@@ -355,7 +403,9 @@ func TestDontHaveTimeoutMgrUsesDefaultTimeoutIfLatencyLonger(t *testing.T) {
 	}
 
 	// Sleep until after the default timeout
-	time.Sleep(defaultTimeout * 2)
+	clock.Add(defaultTimeout * 2)
+
+	<-timeoutsTriggered
 
 	// Now the keys should have timed out
 	if tr.timedOutCount() != len(ks) {
@@ -368,25 +418,29 @@ func TestDontHaveTimeoutNoTimeoutAfterShutdown(t *testing.T) {
 	latency := time.Millisecond * 10
 	latMultiplier := 1
 	expProcessTime := time.Duration(0)
+	clock := clock.NewMock()
+	pinged := make(chan struct{})
+	pc := &mockPeerConn{latency: latency, clock: clock, pinged: pinged}
 	tr := timeoutRecorder{}
-	pc := &mockPeerConn{latency: latency}
+	timeoutsTriggered := make(chan struct{})
 
 	dhtm := newDontHaveTimeoutMgrWithParams(pc, tr.onTimeout,
-		dontHaveTimeout, maxTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime)
+		dontHaveTimeout, maxTimeout, latMultiplier, messageLatencyMultiplier, expProcessTime, clock, timeoutsTriggered)
 	dhtm.Start()
 	defer dhtm.Shutdown()
+	<-pinged
 
 	// Add keys
 	dhtm.AddPending(ks)
 
 	// Wait less than the timeout
-	time.Sleep(latency - 5*time.Millisecond)
+	clock.Add(latency - 5*time.Millisecond)
 
 	// Shutdown the manager
 	dhtm.Shutdown()
 
 	// Wait for the expected timeout
-	time.Sleep(10 * time.Millisecond)
+	clock.Add(10 * time.Millisecond)
 
 	// Manager was shut down so timeout should not have fired
 	if tr.timedOutCount() != 0 {