fix: top N selection

filecoin-saturn · Sep 18, 2023 · ea1d62b · ea1d62b
1 parent 61c82da
commit ea1d62b
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/node_heap.go b/node_heap.go
@@ -51,13 +51,22 @@ func (nh *NodeHeap) PeekRandom() *Node {
 
 func (nh *NodeHeap) TopN(n int) []*Node {
 	m := make([]*Node, 0, n)
+	temp := make([]*Node, 0, n)
 	nh.lk.Lock()
 	defer nh.lk.Unlock()
+
 	heap.Init(nh)
-	for i := 0; i < n && i < len(nh.Nodes); i++ {
-		node := nh.Nodes[i]
+	for i := 0; i < n && nh.Len() > 0; i++ {
+		item := heap.Pop(nh)
+		node := item.(*Node)
 		m = append(m, node)
+		temp = append(temp, node)
+	}
+
+	for _, node := range temp {
+		heap.Push(nh, node)
 	}
+
 	return m
 }