fix: infinite loop in peer.Ancestors() (#1469)

Signed-off-by: Gaius <gaius.qi@gmail.com>
dragonflyoss · Jul 19, 2022 · 311ca60 · 311ca60
1 parent bed8e7a
commit 311ca60
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 25 deletions.
diff --git a/pkg/slices/slices.go b/pkg/slices/slices.go
@@ -26,3 +26,18 @@ func Contains[T comparable](s []T, e T) bool {
 
 	return false
 }
+
+// FindDuplicate returns duplicate element in a collection.
+func FindDuplicate[T comparable](s []T) (T, bool) {
+	visited := make(map[T]struct{})
+	for _, v := range s {
+		if _, ok := visited[v]; ok {
+			return v, true
+		}
+
+		visited[v] = struct{}{}
+	}
+
+	var zero T
+	return zero, false
+}
diff --git a/pkg/slices/slices_test.go b/pkg/slices/slices_test.go
@@ -29,3 +29,25 @@ func TestContains(t *testing.T) {
 	assert.True(Contains([]string{"a", "b", "c"}, "a"))
 	assert.False(Contains([]string{"a", "b", "c"}, "d"))
 }
+
+func TestFindDuplicate(t *testing.T) {
+	assert := assert.New(t)
+	var (
+		dupi  int
+		dups  string
+		found bool
+	)
+	dupi, found = FindDuplicate([]int{1, 2, 1, 3})
+	assert.True(found)
+	assert.Equal(dupi, 1)
+
+	_, found = FindDuplicate([]int{1, 2, 3})
+	assert.False(found)
+
+	dups, found = FindDuplicate([]string{"a", "b", "c", "b"})
+	assert.True(found)
+	assert.Equal(dups, "b")
+
+	_, found = FindDuplicate([]string{"a", "b", "c"})
+	assert.False(found)
+}
diff --git a/scheduler/resource/peer.go b/scheduler/resource/peer.go
@@ -33,6 +33,7 @@ import (
 	logger "d7y.io/dragonfly/v2/internal/dflog"
 	"d7y.io/dragonfly/v2/pkg/container/set"
 	"d7y.io/dragonfly/v2/pkg/rpc/scheduler"
+	"d7y.io/dragonfly/v2/pkg/slices"
 )
 
 const (
@@ -373,11 +374,11 @@ func (p *Peer) Depth() int {
 	p.mu.RLock()
 	defer p.mu.RUnlock()
 
-	node := p
-	var depth int
+	var (
+		node      = p
+		ancestors = []string{p.ID}
+	)
 	for node != nil {
-		depth++
-
 		if node.Host.Type != HostTypeNormal {
 			break
 		}
@@ -387,74 +388,84 @@ func (p *Peer) Depth() int {
 			break
 		}
 
+		ancestors = append(ancestors, parent.ID)
+
 		// Prevent traversal tree from infinite loop.
-		if p.ID == parent.ID {
+		if _, found := slices.FindDuplicate(ancestors); found {
 			p.Log.Error("tree structure produces an infinite loop")
 			break
 		}
 
 		node = parent
 	}
 
-	return depth
+	return len(ancestors)
 }
 
 // Ancestors returns peer's ancestors.
 func (p *Peer) Ancestors() []string {
 	p.mu.RLock()
 	defer p.mu.RUnlock()
 
-	var ancestors []string
-	node := p
+	var (
+		node      = p
+		ancestors = []string{p.ID}
+	)
 	for node != nil {
 		parent, ok := node.LoadParent()
 		if !ok {
 			return ancestors
 		}
 
+		ancestors = append(ancestors, parent.ID)
+
 		// Prevent traversal tree from infinite loop.
-		if parent.ID == node.ID {
+		if _, found := slices.FindDuplicate(ancestors); found {
 			p.Log.Error("tree structure produces an infinite loop")
-			return ancestors
+			break
 		}
 
 		node = parent
-		ancestors = append(ancestors, node.ID)
 	}
 
 	return ancestors
 }
 
 // IsDescendant determines whether it is ancestor of peer.
 func (p *Peer) IsDescendant(ancestor *Peer) bool {
-	return p.isDescendant(ancestor, p)
+	return ancestor.isDescendant(p)
 }
 
 // IsAncestor determines whether it is descendant of peer.
 func (p *Peer) IsAncestor(descendant *Peer) bool {
-	return p.isDescendant(p, descendant)
+	return p.isDescendant(descendant)
 }
 
 // isDescendant determines whether it is ancestor of peer.
-func (p *Peer) isDescendant(ancestor, descendant *Peer) bool {
+func (p *Peer) isDescendant(descendant *Peer) bool {
 	p.mu.RLock()
 	defer p.mu.RUnlock()
 
-	node := descendant
+	var (
+		node      = descendant
+		ancestors = []string{descendant.ID}
+	)
 	for node != nil {
 		parent, ok := node.LoadParent()
 		if !ok {
 			return false
 		}
 
-		if parent.ID == ancestor.ID {
+		if parent.ID == p.ID {
 			return true
 		}
 
+		ancestors = append(ancestors, parent.ID)
+
 		// Prevent traversal tree from infinite loop.
-		if parent.ID == descendant.ID {
+		if _, found := slices.FindDuplicate(ancestors); found {
 			p.Log.Error("tree structure produces an infinite loop")
-			return true
+			break
 		}
 
 		node = parent

diff --git a/scheduler/resource/peer_test.go b/scheduler/resource/peer_test.go
@@ -526,7 +526,7 @@ func TestPeer_Depth(t *testing.T) {
 				peer.StoreParent(peer)
 
 				assert := assert.New(t)
-				assert.Equal(peer.Depth(), 1)
+				assert.Equal(peer.Depth(), 2)
 			},
 		},
 	}
@@ -557,17 +557,17 @@ func TestPeer_Ancestors(t *testing.T) {
 			expect: func(t *testing.T, peer *Peer, mockChildPeer *Peer) {
 				assert := assert.New(t)
 				peer.StoreChild(mockChildPeer)
-				assert.Equal(len(mockChildPeer.Ancestors()), 1)
-				assert.EqualValues(mockChildPeer.Ancestors(), []string{peer.ID})
+				assert.Equal(len(mockChildPeer.Ancestors()), 2)
+				assert.EqualValues(mockChildPeer.Ancestors(), []string{mockChildPeer.ID, peer.ID})
 			},
 		},
 		{
 			name:    "child has no parent",
 			childID: idgen.PeerID("127.0.0.1"),
 			expect: func(t *testing.T, peer *Peer, mockChildPeer *Peer) {
 				assert := assert.New(t)
-				assert.Equal(len(mockChildPeer.Ancestors()), 0)
-				assert.EqualValues(mockChildPeer.Ancestors(), []string(nil))
+				assert.Equal(len(mockChildPeer.Ancestors()), 1)
+				assert.EqualValues(mockChildPeer.Ancestors(), []string{mockChildPeer.ID})
 			},
 		},
 		{
@@ -576,8 +576,8 @@ func TestPeer_Ancestors(t *testing.T) {
 			expect: func(t *testing.T, peer *Peer, mockChildPeer *Peer) {
 				assert := assert.New(t)
 				peer.StoreChild(peer)
-				assert.Equal(len(peer.Ancestors()), 0)
-				assert.Equal(peer.Ancestors(), []string(nil))
+				assert.Equal(len(peer.Ancestors()), 2)
+				assert.Equal(peer.Ancestors(), []string{peer.ID, peer.ID})
 			},
 		},
 	}