feat: add PushTasksTruncated which only push a limited amount of tasks

Author: Jorropo

go.mod

@@ -4,7 +4,7 @@ go 1.19
 
 require (
 	github.com/benbjohnson/clock v1.3.0
-	github.com/ipfs/go-ipfs-pq v0.0.2
+	github.com/ipfs/go-ipfs-pq v0.0.3
 	github.com/libp2p/go-libp2p v0.22.0
 )
 

go.sum

@@ -8,8 +8,8 @@ github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
 github.com/ipfs/go-cid v0.2.0 h1:01JTiihFq9en9Vz0lc0VDWvZe/uBonGpzo4THP0vcQ0=
 github.com/ipfs/go-cid v0.2.0/go.mod h1:P+HXFDF4CVhaVayiEb4wkAy7zBHxBwsJyt0Y5U6MLro=
-github.com/ipfs/go-ipfs-pq v0.0.2 h1:e1vOOW6MuOwG2lqxcLA+wEn93i/9laCY8sXAw76jFOY=
-github.com/ipfs/go-ipfs-pq v0.0.2/go.mod h1:LWIqQpqfRG3fNc5XsnIhz/wQ2XXGyugQwls7BgUmUfY=
+github.com/ipfs/go-ipfs-pq v0.0.3 h1:YpoHVJB+jzK15mr/xsWC574tyDLkezVrDNeaalQBsTE=
+github.com/ipfs/go-ipfs-pq v0.0.3/go.mod h1:btNw5hsHBpRcSSgZtiNm/SLj5gYIZ18AKtv3kERkRb4=
 github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
 github.com/klauspost/cpuid/v2 v2.0.4/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=

peertaskqueue.go

@@ -1,6 +1,7 @@
 package peertaskqueue
 
 import (
+	"math"
 	"sync"
 
 	pq "github.com/ipfs/go-ipfs-pq"
@@ -204,6 +205,13 @@ func (ptq *PeerTaskQueue) PeerTopics(p peer.ID) *peertracker.PeerTrackerTopics {
 
 // PushTasks adds a new group of tasks for the given peer to the queue
 func (ptq *PeerTaskQueue) PushTasks(to peer.ID, tasks ...peertask.Task) {
+	ptq.PushTasksTruncated(math.MaxUint, to, tasks...)
+}
+
+// PushTasksTruncated is like PushTasks but it will not grow that peers's queue beyond n.
+// When truncation happen we will keep older tasks in the queue to avoid some infinite
+// tasks rotations if we are continously receiving work faster than we process it.
+func (ptq *PeerTaskQueue) PushTasksTruncated(n uint, to peer.ID, tasks ...peertask.Task) {
 	ptq.lock.Lock()
 	defer ptq.lock.Unlock()
 
@@ -219,7 +227,7 @@ func (ptq *PeerTaskQueue) PushTasks(to peer.ID, tasks ...peertask.Task) {
 		ptq.callHooks(to, peerAdded)
 	}
 
-	peerTracker.PushTasks(tasks...)
+	peerTracker.PushTasksTruncated(n, tasks...)
 	ptq.pQueue.Update(peerTracker.Index())
 }
 

peertracker/peertracker.go

@@ -1,6 +1,8 @@
 package peertracker
 
 import (
+	"math"
+	"math/bits"
 	"sync"
 
 	"github.com/benbjohnson/clock"
@@ -205,11 +207,32 @@ func (p *PeerTracker) SetIndex(i int) {
 
 // PushTasks adds a group of tasks onto a peer's queue
 func (p *PeerTracker) PushTasks(tasks ...peertask.Task) {
+	p.PushTasksTruncated(math.MaxUint, tasks...)
+}
+
+// PushTasksTruncated is like PushTasks but it will never grow the queue more than n.
+// When truncation happen we will keep older tasks in the queue to avoid some infinite
+// tasks rotations if we are continously receiving work faster than we process it.
+func (p *PeerTracker) PushTasksTruncated(n uint, tasks ...peertask.Task) {
 	now := clockInstance.Now()
 
 	p.activelk.Lock()
 	defer p.activelk.Unlock()
 
+	l := p.taskQueue.Len()
+	if l < 0 {
+		panic("negative length")
+	}
+
+	if wouldBe := uint(l + len(tasks)); wouldBe > n {
+		available, o := bits.Sub(n, uint(l), 0)
+		if o != 0 {
+			// happen if you mix Truncated and Untrucated or varies n.
+			available = 0
+		}
+		tasks = tasks[:available]
+	}
+
 	for _, task := range tasks {
 		// If the new task doesn't add any more information over what we
 		// already have in the active queue, then we can skip the new task

peertracker/peertracker_test.go

@@ -35,6 +35,13 @@ func TestPushPop(t *testing.T) {
 		},
 	}
 	tracker.PushTasks(tasks...)
+
+	tracker.PushTasksTruncated(1, peertask.Task{
+		Topic:    "2",
+		Priority: 2,
+		Work:     20,
+	})
+
 	popped, _ := tracker.PopTasks(100)
 	if len(popped) != 1 {
 		t.Fatal("Expected 1 task")