Repository: dragonflyoss/Dragonfly2
Pull request #1224 (2022-04-02 10:43:56)

feat: scheduler blocks steal peers
Signed-off-by: Gaius <gaius.qi@gmail.com>

<!--- Provide a general summary of your changes in the Title above -->

## Description
- Scheduler blocks steal peers.
<!--- Describe your changes in detail -->

## Related Issue

<!--- This project only accepts pull requests related to open issues -->
<!--- If suggesting a new feature or change, please discuss it in an issue first -->
<!--- If fixing a bug, there should be an issue describing it with steps to reproduce -->
<!--- Please link to the issue here: -->

## Motivation and Context

<!--- Why is this change required? What problem does it solve? -->

## Screenshots (if appropriate)

## Types of changes

<!--- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->

- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation Update (if none of the other choices apply)

## Checklist

<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->

- [x] Code compiles correctly.
- [x] My code follows the code style of this project.
- [ ] My change requires a change to the documentation.
- [ ] I have updated the documentation accordingly.
- [x] I have read the **CONTRIBUTING** document.
- [x] I have added tests to cover my changes.
- [x] All new and existing tests passed.



Patch:
diff --git a/pkg/container/set/safe_set.go b/pkg/container/set/safe_set.go
index 5055ff41..20d21deb 100644
--- a/pkg/container/set/safe_set.go
+++ b/pkg/container/set/safe_set.go
@@ -27,6 +27,7 @@ type SafeSet interface {
 	Contains(...interface{}) bool
 	Len() uint
 	Range(func(interface{}) bool)
+	Clear()
 }
 
 type safeSet struct {
@@ -99,3 +100,9 @@ func (s *safeSet) Range(fn func(interface{}) bool) {
 		}
 	}
 }
+
+func (s *safeSet) Clear() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.data = make(map[interface{}]struct{})
+}
diff --git a/pkg/container/set/set.go b/pkg/container/set/set.go
index 89a5c20e..c3f2f65c 100644
--- a/pkg/container/set/set.go
+++ b/pkg/container/set/set.go
@@ -23,6 +23,7 @@ type Set interface {
 	Contains(...interface{}) bool
 	Len() uint
 	Range(func(interface{}) bool)
+	Clear()
 }
 
 type set map[interface{}]struct{}
@@ -76,3 +77,7 @@ func (s *set) Range(fn func(interface{}) bool) {
 		}
 	}
 }
+
+func (s *set) Clear() {
+	*s = set{}
+}
diff --git a/scheduler/resource/peer.go b/scheduler/resource/peer.go
index b5fab2b8..b0dd1563 100644
--- a/scheduler/resource/peer.go
+++ b/scheduler/resource/peer.go
@@ -143,6 +143,9 @@ type Peer struct {
 	// ChildCount is child count
 	ChildCount *atomic.Int32
 
+	// StealPeers is steal peer ids
+	StealPeers set.SafeSet
+
 	// BlockPeers is bad peer ids
 	BlockPeers set.SafeSet
 
@@ -172,6 +175,7 @@ func NewPeer(id string, task *Task, host *Host, options ...PeerOption) *Peer {
 		Parent:     &atomic.Value{},
 		Children:   &sync.Map{},
 		ChildCount: atomic.NewInt32(0),
+		StealPeers: set.NewSafeSet(),
 		BlockPeers: set.NewSafeSet(),
 		CreateAt:   atomic.NewTime(time.Now()),
 		UpdateAt:   atomic.NewTime(time.Now()),
diff --git a/scheduler/scheduler/scheduler.go b/scheduler/scheduler/scheduler.go
index 00002490..11be1975 100644
--- a/scheduler/scheduler/scheduler.go
+++ b/scheduler/scheduler/scheduler.go
@@ -186,8 +186,16 @@ func (s *scheduler) NotifyAndFindParent(ctx context.Context, peer *resource.Peer
 		return []*resource.Peer{}, false
 	}
 
+	// Add steal peers to current peer
+	peer.StealPeers.Clear()
+	for _, parent := range parents[1:] {
+		peer.StealPeers.Add(parent.ID)
+	}
+
+	// Replace peer's parent with scheduled parent
 	peer.ReplaceParent(parents[0])
-	peer.Log.Infof("schedule parent successful, replace parent to %s", parents[0].ID)
+	peer.Log.Infof("schedule parent successful, replace parent to %s and steal peers is %v",
+		parents[0].ID, peer.StealPeers.Values())
 	return parents, true
 }
 
@@ -239,6 +247,11 @@ func (s *scheduler) filterParents(peer *resource.Peer, blocklist set.SafeSet) []
 			return true
 		}
 
+		if parent.StealPeers.Contains(peer.ID) {
+			peer.Log.Debugf("parent %s is not selected because it is in steal peers", parent.ID)
+			return true
+		}
+
 		if parent.ID == peer.ID {
 			peer.Log.Debug("parent is not selected because it is same")
 			return true


Test patch:
diff --git a/pkg/container/set/safe_set_test.go b/pkg/container/set/safe_set_test.go
index c6f30f3d..13a747ed 100644
--- a/pkg/container/set/safe_set_test.go
+++ b/pkg/container/set/safe_set_test.go
@@ -31,12 +31,12 @@ func TestSafeSetAdd(t *testing.T) {
 	tests := []struct {
 		name   string
 		value  interface{}
-		expect func(t *testing.T, ok bool, s Set, value interface{})
+		expect func(t *testing.T, ok bool, s SafeSet, value interface{})
 	}{
 		{
 			name:  "add value",
 			value: "foo",
-			expect: func(t *testing.T, ok bool, s Set, value interface{}) {
+			expect: func(t *testing.T, ok bool, s SafeSet, value interface{}) {
 				assert := assert.New(t)
 				assert.Equal(ok, true)
 				assert.Equal(s.Values(), []interface{}{value})
@@ -45,7 +45,7 @@ func TestSafeSetAdd(t *testing.T) {
 		{
 			name:  "add value failed",
 			value: "foo",
-			expect: func(t *testing.T, _ bool, s Set, value interface{}) {
+			expect: func(t *testing.T, _ bool, s SafeSet, value interface{}) {
 				assert := assert.New(t)
 				ok := s.Add("foo")
 				assert.Equal(ok, false)
@@ -89,12 +89,12 @@ func TestSafeSetDelete(t *testing.T) {
 	tests := []struct {
 		name   string
 		value  interface{}
-		expect func(t *testing.T, s Set, value interface{})
+		expect func(t *testing.T, s SafeSet, value interface{})
 	}{
 		{
 			name:  "delete value",
 			value: "foo",
-			expect: func(t *testing.T, s Set, value interface{}) {
+			expect: func(t *testing.T, s SafeSet, value interface{}) {
 				assert := assert.New(t)
 				s.Delete(value)
 				assert.Equal(s.Len(), uint(0))
@@ -103,7 +103,7 @@ func TestSafeSetDelete(t *testing.T) {
 		{
 			name:  "delete value does not exist",
 			value: "foo",
-			expect: func(t *testing.T, s Set, _ interface{}) {
+			expect: func(t *testing.T, s SafeSet, _ interface{}) {
 				assert := assert.New(t)
 				s.Delete("bar")
 				assert.Equal(s.Len(), uint(1))
@@ -148,12 +148,12 @@ func TestSafeSetContains(t *testing.T) {
 	tests := []struct {
 		name   string
 		value  interface{}
-		expect func(t *testing.T, s Set, value interface{})
+		expect func(t *testing.T, s SafeSet, value interface{})
 	}{
 		{
 			name:  "contains value",
 			value: "foo",
-			expect: func(t *testing.T, s Set, value interface{}) {
+			expect: func(t *testing.T, s SafeSet, value interface{}) {
 				assert := assert.New(t)
 				assert.Equal(s.Contains(value), true)
 			},
@@ -161,7 +161,7 @@ func TestSafeSetContains(t *testing.T) {
 		{
 			name:  "contains value does not exist",
 			value: "foo",
-			expect: func(t *testing.T, s Set, _ interface{}) {
+			expect: func(t *testing.T, s SafeSet, _ interface{}) {
 				assert := assert.New(t)
 				assert.Equal(s.Contains("bar"), false)
 			},
@@ -202,11 +202,11 @@ func TestSafeSetContains_Concurrent(t *testing.T) {
 func TestSetSafeLen(t *testing.T) {
 	tests := []struct {
 		name   string
-		expect func(t *testing.T, s Set)
+		expect func(t *testing.T, s SafeSet)
 	}{
 		{
 			name: "get length",
-			expect: func(t *testing.T, s Set) {
+			expect: func(t *testing.T, s SafeSet) {
 				assert := assert.New(t)
 				s.Add("foo")
 				assert.Equal(s.Len(), uint(1))
@@ -214,7 +214,7 @@ func TestSetSafeLen(t *testing.T) {
 		},
 		{
 			name: "get empty set length",
-			expect: func(t *testing.T, s Set) {
+			expect: func(t *testing.T, s SafeSet) {
 				assert := assert.New(t)
 				assert.Equal(s.Len(), uint(0))
 			},
@@ -256,11 +256,11 @@ func TestSafeSetLen_Concurrent(t *testing.T) {
 func TestSafeSetValues(t *testing.T) {
 	tests := []struct {
 		name   string
-		expect func(t *testing.T, s Set)
+		expect func(t *testing.T, s SafeSet)
 	}{
 		{
 			name: "get values",
-			expect: func(t *testing.T, s Set) {
+			expect: func(t *testing.T, s SafeSet) {
 				assert := assert.New(t)
 				s.Add("foo")
 				assert.Equal(s.Values(), []interface{}{"foo"})
@@ -268,14 +268,14 @@ func TestSafeSetValues(t *testing.T) {
 		},
 		{
 			name: "get empty values",
-			expect: func(t *testing.T, s Set) {
+			expect: func(t *testing.T, s SafeSet) {
 				assert := assert.New(t)
 				assert.Equal(s.Values(), []interface{}(nil))
 			},
 		},
 		{
 			name: "get multi values",
-			expect: func(t *testing.T, s Set) {
+			expect: func(t *testing.T, s SafeSet) {
 				assert := assert.New(t)
 				s.Add("foo")
 				s.Add("bar")
@@ -320,11 +320,11 @@ func TestSafeSetValues_Concurrent(t *testing.T) {
 func TestSafeSetRange(t *testing.T) {
 	tests := []struct {
 		name   string
-		expect func(t *testing.T, s Set)
+		expect func(t *testing.T, s SafeSet)
 	}{
 		{
 			name: "range values",
-			expect: func(t *testing.T, s Set) {
+			expect: func(t *testing.T, s SafeSet) {
 				assert := assert.New(t)
 				s.Add("foo")
 				s.Range(func(v interface{}) bool {
@@ -335,7 +335,7 @@ func TestSafeSetRange(t *testing.T) {
 		},
 		{
 			name: "range values failed",
-			expect: func(t *testing.T, s Set) {
+			expect: func(t *testing.T, s SafeSet) {
 				assert := assert.New(t)
 				s.Add("foo")
 				s.Add("bar")
@@ -380,3 +380,61 @@ func TestSafeSetRange_Concurrent(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+func TestSafeSetClear(t *testing.T) {
+	tests := []struct {
+		name   string
+		expect func(t *testing.T, s SafeSet)
+	}{
+		{
+			name: "clear empty set",
+			expect: func(t *testing.T, s SafeSet) {
+				assert := assert.New(t)
+				s.Clear()
+				assert.Equal(s.Values(), []interface{}(nil))
+			},
+		},
+		{
+			name: "clear set",
+			expect: func(t *testing.T, s SafeSet) {
+				assert := assert.New(t)
+				assert.Equal(s.Add("foo"), true)
+				s.Clear()
+				assert.Equal(s.Values(), []interface{}(nil))
+				assert.Equal(s.Add("foo"), true)
+				assert.Equal(s.Values(), []interface{}{"foo"})
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			s := NewSafeSet()
+			tc.expect(t, s)
+		})
+	}
+}
+
+func TestSafeSetClear_Concurrent(t *testing.T) {
+	runtime.GOMAXPROCS(2)
+
+	s := NewSafeSet()
+	nums := rand.Perm(N)
+
+	var wg sync.WaitGroup
+	wg.Add(len(nums))
+	for i := 0; i < len(nums); i++ {
+		go func(i int) {
+			s.Add(i)
+			s.Clear()
+			wg.Done()
+		}(i)
+	}
+
+	wg.Wait()
+	for _, n := range nums {
+		if s.Contains(n) {
+			t.Errorf("SafeSet contains element: %v", n)
+		}
+	}
+}
diff --git a/pkg/container/set/set_test.go b/pkg/container/set/set_test.go
index 4e7af857..25e10c58 100644
--- a/pkg/container/set/set_test.go
+++ b/pkg/container/set/set_test.go
@@ -232,3 +232,37 @@ func TestSetRange(t *testing.T) {
 		})
 	}
 }
+
+func TestSetClear(t *testing.T) {
+	tests := []struct {
+		name   string
+		expect func(t *testing.T, s Set)
+	}{
+		{
+			name: "clear empty set",
+			expect: func(t *testing.T, s Set) {
+				assert := assert.New(t)
+				s.Clear()
+				assert.Equal(s.Values(), []interface{}(nil))
+			},
+		},
+		{
+			name: "clear set",
+			expect: func(t *testing.T, s Set) {
+				assert := assert.New(t)
+				assert.Equal(s.Add("foo"), true)
+				s.Clear()
+				assert.Equal(s.Values(), []interface{}(nil))
+				assert.Equal(s.Add("foo"), true)
+				assert.Equal(s.Values(), []interface{}{"foo"})
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			s := New()
+			tc.expect(t, s)
+		})
+	}
+}
diff --git a/scheduler/scheduler/scheduler_test.go b/scheduler/scheduler/scheduler_test.go
index fe8c8353..1e984684 100644
--- a/scheduler/scheduler/scheduler_test.go
+++ b/scheduler/scheduler/scheduler_test.go
@@ -449,14 +449,14 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 	tests := []struct {
 		name   string
 		mock   func(peer *resource.Peer, mockPeer *resource.Peer, blocklist set.SafeSet, stream rpcscheduler.Scheduler_ReportPieceResultServer, dynconfig config.DynconfigInterface, ms *rpcschedulermocks.MockScheduler_ReportPieceResultServerMockRecorder, md *configmocks.MockDynconfigInterfaceMockRecorder)
-		expect func(t *testing.T, parents []*resource.Peer, ok bool)
+		expect func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool)
 	}{
 		{
 			name: "peer state is PeerStatePending",
 			mock: func(peer *resource.Peer, mockPeer *resource.Peer, blocklist set.SafeSet, stream rpcscheduler.Scheduler_ReportPieceResultServer, dynconfig config.DynconfigInterface, ms *rpcschedulermocks.MockScheduler_ReportPieceResultServerMockRecorder, md *configmocks.MockDynconfigInterfaceMockRecorder) {
 				peer.FSM.SetState(resource.PeerStatePending)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -466,7 +466,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 			mock: func(peer *resource.Peer, mockPeer *resource.Peer, blocklist set.SafeSet, stream rpcscheduler.Scheduler_ReportPieceResultServer, dynconfig config.DynconfigInterface, ms *rpcschedulermocks.MockScheduler_ReportPieceResultServerMockRecorder, md *configmocks.MockDynconfigInterfaceMockRecorder) {
 				peer.FSM.SetState(resource.PeerStateReceivedSmall)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -476,7 +476,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 			mock: func(peer *resource.Peer, mockPeer *resource.Peer, blocklist set.SafeSet, stream rpcscheduler.Scheduler_ReportPieceResultServer, dynconfig config.DynconfigInterface, ms *rpcschedulermocks.MockScheduler_ReportPieceResultServerMockRecorder, md *configmocks.MockDynconfigInterfaceMockRecorder) {
 				peer.FSM.SetState(resource.PeerStateReceivedNormal)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -486,7 +486,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 			mock: func(peer *resource.Peer, mockPeer *resource.Peer, blocklist set.SafeSet, stream rpcscheduler.Scheduler_ReportPieceResultServer, dynconfig config.DynconfigInterface, ms *rpcschedulermocks.MockScheduler_ReportPieceResultServerMockRecorder, md *configmocks.MockDynconfigInterfaceMockRecorder) {
 				peer.FSM.SetState(resource.PeerStateBackToSource)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -496,7 +496,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 			mock: func(peer *resource.Peer, mockPeer *resource.Peer, blocklist set.SafeSet, stream rpcscheduler.Scheduler_ReportPieceResultServer, dynconfig config.DynconfigInterface, ms *rpcschedulermocks.MockScheduler_ReportPieceResultServerMockRecorder, md *configmocks.MockDynconfigInterfaceMockRecorder) {
 				peer.FSM.SetState(resource.PeerStateSucceeded)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -506,7 +506,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 			mock: func(peer *resource.Peer, mockPeer *resource.Peer, blocklist set.SafeSet, stream rpcscheduler.Scheduler_ReportPieceResultServer, dynconfig config.DynconfigInterface, ms *rpcschedulermocks.MockScheduler_ReportPieceResultServerMockRecorder, md *configmocks.MockDynconfigInterfaceMockRecorder) {
 				peer.FSM.SetState(resource.PeerStateFailed)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -516,7 +516,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 			mock: func(peer *resource.Peer, mockPeer *resource.Peer, blocklist set.SafeSet, stream rpcscheduler.Scheduler_ReportPieceResultServer, dynconfig config.DynconfigInterface, ms *rpcschedulermocks.MockScheduler_ReportPieceResultServerMockRecorder, md *configmocks.MockDynconfigInterfaceMockRecorder) {
 				peer.FSM.SetState(resource.PeerStateLeave)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -528,7 +528,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 
 				md.GetSchedulerClusterConfig().Return(types.SchedulerClusterConfig{}, false).Times(1)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -541,7 +541,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 
 				md.GetSchedulerClusterConfig().Return(types.SchedulerClusterConfig{}, false).Times(1)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -555,7 +555,21 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 
 				md.GetSchedulerClusterConfig().Return(types.SchedulerClusterConfig{}, false).Times(1)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
+				assert := assert.New(t)
+				assert.False(ok)
+			},
+		},
+		{
+			name: "peer is in steal peers",
+			mock: func(peer *resource.Peer, mockPeer *resource.Peer, blocklist set.SafeSet, stream rpcscheduler.Scheduler_ReportPieceResultServer, dynconfig config.DynconfigInterface, ms *rpcschedulermocks.MockScheduler_ReportPieceResultServerMockRecorder, md *configmocks.MockDynconfigInterfaceMockRecorder) {
+				peer.FSM.SetState(resource.PeerStateRunning)
+				peer.Task.StorePeer(mockPeer)
+				peer.StealPeers.Add(mockPeer.ID)
+
+				md.GetSchedulerClusterConfig().Return(types.SchedulerClusterConfig{}, false).Times(1)
+			},
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -567,7 +581,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 				peer.FSM.SetState(resource.PeerStateFailed)
 				peer.Task.StorePeer(mockPeer)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -582,7 +596,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 
 				md.GetSchedulerClusterConfig().Return(types.SchedulerClusterConfig{}, false).Times(1)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -597,7 +611,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 
 				md.GetSchedulerClusterConfig().Return(types.SchedulerClusterConfig{}, false).Times(1)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -612,7 +626,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 
 				md.GetSchedulerClusterConfig().Return(types.SchedulerClusterConfig{}, false).Times(1)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -627,7 +641,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 
 				md.GetSchedulerClusterConfig().Return(types.SchedulerClusterConfig{}, false).Times(1)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -648,7 +662,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 					ms.Send(gomock.Any()).Return(errors.New("foo")).Times(1),
 				)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
 				assert.False(ok)
 			},
@@ -658,7 +672,12 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 			mock: func(peer *resource.Peer, mockPeer *resource.Peer, blocklist set.SafeSet, stream rpcscheduler.Scheduler_ReportPieceResultServer, dynconfig config.DynconfigInterface, ms *rpcschedulermocks.MockScheduler_ReportPieceResultServerMockRecorder, md *configmocks.MockDynconfigInterfaceMockRecorder) {
 				peer.FSM.SetState(resource.PeerStateRunning)
 				mockPeer.FSM.SetState(resource.PeerStateRunning)
+				mockHost := resource.NewHost(mockRawHost)
+				mockTask := resource.NewTask(mockTaskID, mockTaskURL, mockTaskBackToSourceLimit, mockTaskURLMeta)
+				stealPeer := resource.NewPeer(idgen.PeerID("127.0.0.1"), mockTask, mockHost)
+				stealPeer.FSM.SetState(resource.PeerStateRunning)
 				peer.Task.StorePeer(mockPeer)
+				peer.Task.StorePeer(stealPeer)
 				mockPeer.Pieces.Set(0)
 				peer.StoreStream(stream)
 				gomock.InOrder(
@@ -669,9 +688,10 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 					ms.Send(gomock.Any()).Return(nil).Times(1),
 				)
 			},
-			expect: func(t *testing.T, parents []*resource.Peer, ok bool) {
+			expect: func(t *testing.T, peer *resource.Peer, parents []*resource.Peer, ok bool) {
 				assert := assert.New(t)
-				assert.Equal(len(parents), 1)
+				assert.Equal(len(parents), 2)
+				assert.Equal(peer.StealPeers.Len(), uint(1))
 				assert.True(ok)
 			},
 		},
@@ -692,7 +712,7 @@ func TestScheduler_NotifyAndFindParent(t *testing.T) {
 			tc.mock(peer, mockPeer, blocklist, stream, dynconfig, stream.EXPECT(), dynconfig.EXPECT())
 			scheduler := New(mockSchedulerConfig, dynconfig, mockPluginDir)
 			parents, ok := scheduler.NotifyAndFindParent(context.Background(), peer, blocklist)
-			tc.expect(t, parents, ok)
+			tc.expect(t, peer, parents, ok)
 		})
 	}
 }