From 3712944d8687a989e1f3f32dec24c83855706297 Mon Sep 17 00:00:00 2001
From: Dennis Trautwein <git@dtrautwein.eu>
Date: Mon, 18 Sep 2023 17:09:55 +0200
Subject: [PATCH] fix: avoid panic when node is re-added to probe list

---
 v2/coord/routing/probe.go      | 19 +++++++++++++------
 v2/coord/routing/probe_test.go | 22 +++++++++++++++++++++-
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/v2/coord/routing/probe.go b/v2/coord/routing/probe.go
index b8f78c60..45e5881f 100644
--- a/v2/coord/routing/probe.go
+++ b/v2/coord/routing/probe.go
@@ -334,11 +334,12 @@ type nodeValue[K kad.Key[K], N kad.NodeID[K]] struct {
 
 type nodeValueEntry[K kad.Key[K], N kad.NodeID[K]] struct {
 	nv    *nodeValue[K, N]
-	index int // the index of the item in the ordering
+	index int // the index of the item in the ordering, set to -1 when the item is popped from the heap
 }
 
 type nodeValueList[K kad.Key[K], N kad.NodeID[K]] struct {
-	nodes   map[string]*nodeValueEntry[K, N]
+	nodes map[string]*nodeValueEntry[K, N]
+	// pending is a list of nodes ordered by the time of the next check
 	pending *nodeValuePendingList[K, N]
 	// ongoing is a list of nodes with ongoing/in-progress probes, loosely ordered earliest to most recent
 	ongoing []N
@@ -359,14 +360,20 @@ func (l *nodeValueList[K, N]) Put(nv *nodeValue[K, N]) {
 	nve, exists := l.nodes[mk]
 	if !exists {
 		nve = &nodeValueEntry[K, N]{
-			nv: nv,
+			nv:    nv,
+			index: -1,
 		}
+		l.nodes[mk] = nve
 	} else {
 		nve.nv = nv
-		heap.Remove(l.pending, nve.index)
 	}
-	heap.Push(l.pending, nve)
-	l.nodes[mk] = nve
+
+	// nve.index is -1 when the node is not already in the pending list
+	// this could be because it is new or if there is an ongoing check
+	if nve.index == -1 {
+		heap.Push(l.pending, nve)
+	}
+
 	heap.Fix(l.pending, nve.index)
 	l.removeFromOngoing(nv.NodeID)
 }
diff --git a/v2/coord/routing/probe_test.go b/v2/coord/routing/probe_test.go
index d4106c51..e07d6445 100644
--- a/v2/coord/routing/probe_test.go
+++ b/v2/coord/routing/probe_test.go
@@ -512,7 +512,7 @@ func TestNodeValueList(t *testing.T) {
 		require.Equal(t, 0, l.OngoingCount())
 		require.Equal(t, 1, l.NodeCount())
 
-		l.MarkOngoing(tiny.NewNode(5), clk.Now().Add(time.Minute))
+		l.MarkOngoing(nv1.NodeID, clk.Now().Add(time.Minute))
 		require.Equal(t, 0, l.PendingCount())
 		require.Equal(t, 1, l.OngoingCount())
 		require.Equal(t, 1, l.NodeCount())
@@ -582,6 +582,26 @@ func TestNodeValueList(t *testing.T) {
 		require.Equal(t, 0, l.OngoingCount())
 		require.Equal(t, 1, l.NodeCount())
 	})
+
+	t.Run("mark ongoing pending mixed", func(t *testing.T) {
+		t.Parallel()
+
+		clk := clock.NewMock()
+		l := NewNodeValueList[tiny.Key, tiny.Node]()
+		nv1 := &nodeValue[tiny.Key, tiny.Node]{
+			NodeID:       tiny.NewNode(5),
+			NextCheckDue: clk.Now().Add(time.Minute),
+		}
+		nv2 := &nodeValue[tiny.Key, tiny.Node]{
+			NodeID:       tiny.NewNode(6),
+			NextCheckDue: clk.Now().Add(time.Minute),
+		}
+
+		l.Put(nv1)
+		l.Put(nv2)
+		l.MarkOngoing(nv1.NodeID, clk.Now().Add(time.Minute))
+		l.Put(nv1)
+	})
 }
 
 func TestProbeConnectivityCheckSuccess(t *testing.T) {