Repository: dragonflyoss/Dragonfly2
Pull request #1876 (2022-11-29 10:14:52)

feat: support reregister peer task in client
<!--- Provide a general summary of your changes in the Title above -->

## Description

<!--- 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)
- [ ] 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! -->

- [ ] My change requires a change to the documentation.
- [ ] I have updated the documentation accordingly.
- [ ] I have read the **CONTRIBUTING** document.
- [ ] I have added tests to cover my changes.



Patch:
diff --git a/client/daemon/peer/peertask_conductor.go b/client/daemon/peer/peertask_conductor.go
index 47a69c2e..57eacd7c 100644
--- a/client/daemon/peer/peertask_conductor.go
+++ b/client/daemon/peer/peertask_conductor.go
@@ -738,7 +738,12 @@ loop:
 			break loop
 		}
 		if err != nil {
-			pt.confirmReceivePeerPacketError(err)
+			// some errors, like commonv1.Code_SchedReregister, after reregister success,
+			// we can continue receive peer packet from the new scheduler
+			cont := pt.confirmReceivePeerPacketError(err)
+			if cont {
+				continue
+			}
 			if !firstPacketReceived {
 				firstPeerSpan.RecordError(err)
 			}
@@ -833,12 +838,12 @@ func (pt *peerTaskConductor) updateSynchronizer(lastNum int32, p *schedulerv1.Pe
 	return desiredPiece
 }
 
-func (pt *peerTaskConductor) confirmReceivePeerPacketError(err error) {
+func (pt *peerTaskConductor) confirmReceivePeerPacketError(err error) (cont bool) {
 	select {
 	case <-pt.successCh:
-		return
+		return false
 	case <-pt.failCh:
-		return
+		return false
 	default:
 	}
 	var (
@@ -846,19 +851,30 @@ func (pt *peerTaskConductor) confirmReceivePeerPacketError(err error) {
 		failedReason string
 	)
 	de, ok := err.(*dferrors.DfError)
-	if ok && de.Code == commonv1.Code_SchedNeedBackSource {
-		pt.markBackSource()
-		pt.Infof("receive back source code")
-		return
-	} else if ok && de.Code != commonv1.Code_SchedNeedBackSource {
-		failedCode = de.Code
-		failedReason = de.Message
-		pt.Errorf("receive peer packet failed: %s", pt.failedReason)
+	if ok {
+		switch de.Code {
+		case commonv1.Code_SchedNeedBackSource:
+			pt.markBackSource()
+			pt.Infof("receive back source code")
+			return false
+		case commonv1.Code_SchedReregister:
+			pt.Infof("receive reregister code")
+			rerr := pt.register()
+			if rerr == nil {
+				pt.Infof("reregister ok")
+				return true
+			}
+			pt.Errorf("reregister to scheduler error:%s", err)
+		default:
+			failedCode = de.Code
+			failedReason = de.Message
+			pt.Errorf("receive peer packet failed: %s", pt.failedReason)
+		}
 	} else {
 		pt.Errorf("receive peer packet failed: %s", err)
 	}
 	pt.cancel(failedCode, failedReason)
-	return
+	return false
 }
 
 func (pt *peerTaskConductor) isExitPeerPacketCode(pp *schedulerv1.PeerPacket) bool {


Test patch:
diff --git a/client/daemon/peer/peertask_manager_test.go b/client/daemon/peer/peertask_manager_test.go
index c4f309b2..90d02852 100644
--- a/client/daemon/peer/peertask_manager_test.go
+++ b/client/daemon/peer/peertask_manager_test.go
@@ -84,6 +84,7 @@ type componentsOption struct {
 	scope              commonv1.SizeScope
 	content            []byte
 	getPieceTasks      bool
+	reregister         bool
 }
 
 func setupPeerTaskManagerComponents(ctrl *gomock.Controller, opt componentsOption) (
@@ -189,8 +190,13 @@ func setupPeerTaskManagerComponents(ctrl *gomock.Controller, opt componentsOptio
 		sent       = make(chan struct{}, 1)
 	)
 	sent <- struct{}{}
+	var reregistered bool
 	pps.EXPECT().Recv().AnyTimes().DoAndReturn(
 		func() (*schedulerv1.PeerPacket, error) {
+			if opt.reregister && !reregistered {
+				reregistered = true
+				return nil, dferrors.New(commonv1.Code_SchedReregister, "reregister")
+			}
 			if len(opt.peerPacketDelay) > delayCount {
 				if delay := opt.peerPacketDelay[delayCount]; delay > 0 {
 					time.Sleep(delay)
@@ -349,6 +355,7 @@ type testSpec struct {
 	peerID             string
 	url                string
 	legacyFeature      bool
+	reregister         bool
 	// when urlGenerator is not nil, use urlGenerator instead url
 	// it's useful for httptest server
 	urlGenerator func(ts *testSpec) string
@@ -399,6 +406,18 @@ func TestPeerTaskManager_TaskSuite(t *testing.T) {
 			mockPieceDownloader:  commonPieceDownloader,
 			mockHTTPSourceClient: nil,
 		},
+		{
+			name:                 "normal size scope - p2p - reregister",
+			taskData:             testBytes,
+			pieceParallelCount:   4,
+			pieceSize:            1024,
+			peerID:               "normal-size-peer",
+			url:                  "http://localhost/test/data",
+			sizeScope:            commonv1.SizeScope_NORMAL,
+			mockPieceDownloader:  commonPieceDownloader,
+			mockHTTPSourceClient: nil,
+			reregister:           true,
+		},
 		{
 			name:                 "small size scope - p2p",
 			taskData:             testBytes,
@@ -663,6 +682,7 @@ func TestPeerTaskManager_TaskSuite(t *testing.T) {
 							peerPacketDelay:    tc.peerPacketDelay,
 							backSource:         tc.backSource,
 							getPieceTasks:      tc.legacyFeature,
+							reregister:         tc.reregister,
 						}
 						// keep peer task running in enough time to check "getOrCreatePeerTaskConductor" always return same
 						if tc.taskType == taskTypeConductor {