feat: add task completion

queue.go

@@ -46,6 +46,7 @@ type QueueTask struct {
 	ctx      context.Context
 	payload  *Task
 	priority int
+	index    int // The index in the heap
 }
 
 type PriorityQueue []*QueueTask
@@ -56,17 +57,25 @@ func (pq PriorityQueue) Less(i, j int) bool {
 	return pq[i].priority > pq[j].priority
 }
 
-func (pq PriorityQueue) Swap(i, j int) { pq[i], pq[j] = pq[j], pq[i] }
+func (pq PriorityQueue) Swap(i, j int) {
+	pq[i], pq[j] = pq[j], pq[i]
+	pq[i].index = i
+	pq[j].index = j
+}
 
 func (pq *PriorityQueue) Push(x interface{}) {
-	item := x.(*QueueTask)
-	*pq = append(*pq, item)
+	n := len(*pq)
+	task := x.(*QueueTask)
+	task.index = n
+	*pq = append(*pq, task)
 }
 
 func (pq *PriorityQueue) Pop() interface{} {
 	old := *pq
 	n := len(old)
-	item := old[n-1]
+	task := old[n-1]
+	old[n-1] = nil  // avoid memory leak
+	task.index = -1 // for safety
 	*pq = old[0 : n-1]
-	return item
+	return task
 }