diff --git a/peertaskqueue.go b/peertaskqueue.go index 115a897..40200bd 100644 --- a/peertaskqueue.go +++ b/peertaskqueue.go @@ -169,19 +169,20 @@ func (ptq *PeerTaskQueue) Remove(identifier peertask.Identifier, p peer.ID) { ptq.lock.Lock() peerTracker, ok := ptq.peerTrackers[p] if ok { - peerTracker.Remove(identifier) - // we now also 'freeze' that partner. If they sent us a cancel for a - // block we were about to send them, we should wait a short period of time - // to make sure we receive any other in-flight cancels before sending - // them a block they already potentially have - if !ptq.ignoreFreezing { - if !peerTracker.IsFrozen() { - ptq.frozenPeers[p] = struct{}{} + if peerTracker.Remove(identifier) { + // we now also 'freeze' that partner. If they sent us a cancel for a + // block we were about to send them, we should wait a short period of time + // to make sure we receive any other in-flight cancels before sending + // them a block they already potentially have + if !ptq.ignoreFreezing { + if !peerTracker.IsFrozen() { + ptq.frozenPeers[p] = struct{}{} + } + + peerTracker.Freeze() } - - peerTracker.Freeze() + ptq.pQueue.Update(peerTracker.Index()) } - ptq.pQueue.Update(peerTracker.Index()) } ptq.lock.Unlock() } diff --git a/peertaskqueue_test.go b/peertaskqueue_test.go index 0cf5a6c..5a8dde3 100644 --- a/peertaskqueue_test.go +++ b/peertaskqueue_test.go @@ -92,7 +92,7 @@ func TestFreezeUnfreeze(t *testing.T) { // now, pop off four tasks, there should be one from each matchNTasks(t, ptq, 4, a.Pretty(), b.Pretty(), c.Pretty(), d.Pretty()) - ptq.Remove(peertask.Task{Identifier: "1"}, b) + ptq.Remove("1", b) // b should be frozen, causing it to get skipped in the rotation matchNTasks(t, ptq, 3, a.Pretty(), c.Pretty(), d.Pretty()) @@ -101,6 +101,12 @@ func TestFreezeUnfreeze(t *testing.T) { matchNTasks(t, ptq, 1, b.Pretty()) + // remove none existent task + ptq.Remove("-1", b) + + // b should not be frozen + matchNTasks(t, ptq, 4, a.Pretty(), b.Pretty(), c.Pretty(), d.Pretty()) + } func TestFreezeUnfreezeNoFreezingOption(t *testing.T) { @@ -124,7 +130,7 @@ func TestFreezeUnfreezeNoFreezingOption(t *testing.T) { // now, pop off four tasks, there should be one from each matchNTasks(t, ptq, 4, a.Pretty(), b.Pretty(), c.Pretty(), d.Pretty()) - ptq.Remove(peertask.Task{Identifier: "1"}, b) + ptq.Remove("1", b) // b should be frozen, causing it to get skipped in the rotation matchNTasks(t, ptq, 4, a.Pretty(), b.Pretty(), c.Pretty(), d.Pretty()) diff --git a/peertracker/peertracker.go b/peertracker/peertracker.go index 24fb33f..3b7451b 100644 --- a/peertracker/peertracker.go +++ b/peertracker/peertracker.go @@ -179,12 +179,13 @@ func (p *PeerTracker) PopBlock() *peertask.TaskBlock { } // Remove removes the task with the given identifier from this peers queue -func (p *PeerTracker) Remove(identifier peertask.Identifier) { +func (p *PeerTracker) Remove(identifier peertask.Identifier) bool { taskBlock, ok := p.taskMap[identifier] if ok { taskBlock.MarkPrunable(identifier) p.numTasks-- } + return ok } // Freeze increments the freeze value for this peer. While a peer is frozen