fix: add unit tests for AnnouncePeers

Signed-off-by: BruceAko <chongzhi@hust.edu.cn>
dragonflyoss · Aug 29, 2024 · 84ae6eb · 84ae6eb
1 parent 5826a1f
commit 84ae6eb
Show file tree

Hide file tree

Showing 2 changed files with 281 additions and 20 deletions.
diff --git a/scheduler/service/service_v2.go b/scheduler/service/service_v2.go
@@ -93,7 +93,7 @@ func (v *V2) AnnouncePeer(stream schedulerv2.Scheduler_AnnouncePeerServer) error
 	for {
 		select {
 		case <-ctx.Done():
-			logger.Info("context was done")
+			logger.Info("announce peer context was done")
 			return ctx.Err()
 		default:
 		}
@@ -141,29 +141,29 @@ func (v *V2) AnnouncePeer(stream schedulerv2.Scheduler_AnnouncePeerServer) error
 		case *schedulerv2.AnnouncePeerRequest_DownloadPeerFinishedRequest:
 			downloadPeerFinishedRequest := announcePeerRequest.DownloadPeerFinishedRequest
 			log.Infof("receive DownloadPeerFinishedRequest, content length: %d, piece count: %d", downloadPeerFinishedRequest.GetContentLength(), downloadPeerFinishedRequest.GetPieceCount())
-			// Notice: Handler uses context.Background() to avoid stream cancel by dfdameon.
+			// Notice: Handler uses context.Background() to avoid stream cancel by dfdaemon.
 			if err := v.handleDownloadPeerFinishedRequest(context.Background(), req.GetPeerId()); err != nil {
 				log.Error(err)
 				return err
 			}
 		case *schedulerv2.AnnouncePeerRequest_DownloadPeerBackToSourceFinishedRequest:
 			downloadPeerBackToSourceFinishedRequest := announcePeerRequest.DownloadPeerBackToSourceFinishedRequest
 			log.Infof("receive DownloadPeerBackToSourceFinishedRequest, content length: %d, piece count: %d", downloadPeerBackToSourceFinishedRequest.GetContentLength(), downloadPeerBackToSourceFinishedRequest.GetPieceCount())
-			// Notice: Handler uses context.Background() to avoid stream cancel by dfdameon.
+			// Notice: Handler uses context.Background() to avoid stream cancel by dfdaemon.
 			if err := v.handleDownloadPeerBackToSourceFinishedRequest(context.Background(), req.GetPeerId(), downloadPeerBackToSourceFinishedRequest); err != nil {
 				log.Error(err)
 				return err
 			}
 		case *schedulerv2.AnnouncePeerRequest_DownloadPeerFailedRequest:
 			log.Infof("receive DownloadPeerFailedRequest, description: %s", announcePeerRequest.DownloadPeerFailedRequest.GetDescription())
-			// Notice: Handler uses context.Background() to avoid stream cancel by dfdameon.
+			// Notice: Handler uses context.Background() to avoid stream cancel by dfdaemon.
 			if err := v.handleDownloadPeerFailedRequest(context.Background(), req.GetPeerId()); err != nil {
 				log.Error(err)
 				return err
 			}
 		case *schedulerv2.AnnouncePeerRequest_DownloadPeerBackToSourceFailedRequest:
 			log.Infof("receive DownloadPeerBackToSourceFailedRequest, description: %s", announcePeerRequest.DownloadPeerBackToSourceFailedRequest.GetDescription())
-			// Notice: Handler uses context.Background() to avoid stream cancel by dfdameon.
+			// Notice: Handler uses context.Background() to avoid stream cancel by dfdaemon.
 			if err := v.handleDownloadPeerBackToSourceFailedRequest(context.Background(), req.GetPeerId()); err != nil {
 				log.Error(err)
 				return err
@@ -867,28 +867,33 @@ func (v *V2) AnnouncePeers(stream schedulerv2.Scheduler_AnnouncePeersServer) err
 	ctx, cancel := context.WithCancel(stream.Context())
 	defer cancel()
 
+	announcePeersCount := 0
+
 	for {
 		select {
 		case <-ctx.Done():
-			logger.Info("context was done")
+			logger.Info("announce peers context was done")
 			return ctx.Err()
 		default:
 		}
 
 		request, err := stream.Recv()
 		if err != nil {
 			if err == io.EOF {
+				logger.Infof("announce %d peers", announcePeersCount)
 				return nil
 			}
 
 			logger.Errorf("receive error: %s", err.Error())
 			return err
 		}
 
-		if _, err := v.handleAnnouncePeersRequest(ctx, request); err != nil {
+		peers, err := v.handleAnnouncePeersRequest(ctx, request)
+		if err != nil {
 			logger.Error(err)
 			return err
 		}
+		announcePeersCount += len(peers)
 	}
 }
 
@@ -1484,16 +1489,17 @@ func (v *V2) downloadTaskBySeedPeer(ctx context.Context, taskID string, download
 }
 
 // handleAnnouncePeersRequest handles AnnouncePeersRequest.
-func (v *V2) handleAnnouncePeersRequest(ctx context.Context, request *schedulerv2.AnnouncePeersRequest) ([]*resource.Peer, error) {
-	var peers []*resource.Peer
-
+func (v *V2) handleAnnouncePeersRequest(ctx context.Context, request *schedulerv2.AnnouncePeersRequest) (peers []*resource.Peer, err error) {
 	for _, p := range request.Peers {
 		hostID := p.GetHost().GetId()
 		peerTask := p.GetTask()
+		if peerTask == nil {
+			return nil, status.Error(codes.InvalidArgument, "request is invalid and doesn't contain a task")
+		}
 		taskID := peerTask.GetId()
 		peerID := p.GetId()
 		download := &commonv2.Download{
-			PieceLength:         peerTask.GetPieceLength(),
+			PieceLength:         &peerTask.PieceLength,
 			Digest:              peerTask.Digest,
 			Url:                 peerTask.GetUrl(),
 			Tag:                 peerTask.Tag,
@@ -1528,11 +1534,7 @@ func (v *V2) handleAnnouncePeersRequest(ctx context.Context, request *schedulerv
 		// advance the task state to TaskStateSucceeded.
 		if !task.FSM.Is(resource.TaskStateSucceeded) {
 			if task.FSM.Can(resource.TaskEventDownload) {
-				if err := task.FSM.Event(ctx, resource.TaskEventDownload); err != nil {
-					msg := fmt.Sprintf("task fsm event failed: %s", err.Error())
-					peer.Log.Error(msg)
-					return nil, status.Error(codes.Internal, err.Error())
-				}
+				task.FSM.SetState(resource.TaskEventDownload)
 			}
 
 			// Construct piece.
@@ -1571,9 +1573,7 @@ func (v *V2) handleAnnouncePeersRequest(ctx context.Context, request *schedulerv
 			if peer.Range == nil && !peer.Task.FSM.Is(resource.TaskStateSucceeded) {
 				peer.Task.ContentLength.Store(int64(len(p.Pieces)))
 				peer.Task.TotalPieceCount.Store(int32(task.ContentLength.Load()))
-				if err := peer.Task.FSM.Event(ctx, resource.TaskEventDownloadSucceeded); err != nil {
-					return nil, status.Error(codes.Internal, err.Error())
-				}
+				peer.Task.FSM.SetState(resource.TaskEventDownloadSucceeded)
 			}
 		}