fix: content length is zero when task succeed (#1732)

Signed-off-by: Gaius <gaius.qi@gmail.com>
dragonflyoss · Oct 9, 2022 · 1dde9af · 1dde9af
1 parent 9247017
commit 1dde9af
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 7 deletions.
diff --git a/scheduler/resource/task.go b/scheduler/resource/task.go
@@ -148,7 +148,7 @@ func NewTask(id, url string, taskType commonv1.TaskType, meta *commonv1.UrlMeta,
 		URL:               url,
 		Type:              taskType,
 		URLMeta:           meta,
-		ContentLength:     atomic.NewInt64(0),
+		ContentLength:     atomic.NewInt64(-1),
 		TotalPieceCount:   atomic.NewInt32(0),
 		BackToSourceLimit: atomic.NewInt32(0),
 		BackToSourcePeers: set.NewSafeSet[string](),

diff --git a/scheduler/resource/task_test.go b/scheduler/resource/task_test.go
@@ -74,7 +74,7 @@ func TestTask_NewTask(t *testing.T) {
 				assert.Equal(task.URL, mockTaskURL)
 				assert.EqualValues(task.URLMeta, mockTaskURLMeta)
 				assert.Empty(task.DirectPiece)
-				assert.Equal(task.ContentLength.Load(), int64(0))
+				assert.Equal(task.ContentLength.Load(), int64(-1))
 				assert.Equal(task.TotalPieceCount.Load(), int32(0))
 				assert.Equal(task.BackToSourceLimit.Load(), int32(200))
 				assert.Equal(task.BackToSourcePeers.Len(), uint(0))

diff --git a/scheduler/service/service.go b/scheduler/service/service.go
@@ -806,14 +806,14 @@ func (s *Service) handleTaskSuccess(ctx context.Context, task *resource.Task, re
 		return
 	}
 
+	// Update task's resource total piece count and content length.
+	task.TotalPieceCount.Store(result.TotalPieceCount)
+	task.ContentLength.Store(result.ContentLength)
+
 	if err := task.FSM.Event(resource.TaskEventDownloadSucceeded); err != nil {
 		task.Log.Errorf("task fsm event failed: %s", err.Error())
 		return
 	}
-
-	// Update task's resource total piece count and content length.
-	task.TotalPieceCount.Store(result.TotalPieceCount)
-	task.ContentLength.Store(result.ContentLength)
 }
 
 // Conditions for the task to switch to the TaskStateSucceeded are:

diff --git a/scheduler/service/service_test.go b/scheduler/service/service_test.go
@@ -279,6 +279,7 @@ func TestService_RegisterPeerTask(t *testing.T) {
 			) {
 				mockPeer.Task.FSM.SetState(resource.TaskStateSucceeded)
 				mockPeer.Task.StorePeer(mockSeedPeer)
+				mockPeer.Task.ContentLength.Store(0)
 				gomock.InOrder(
 					mr.TaskManager().Return(taskManager).Times(1),
 					mt.LoadOrStore(gomock.Any()).Return(mockPeer.Task, true).Times(1),
@@ -1266,7 +1267,7 @@ func TestService_StatTask(t *testing.T) {
 				assert.EqualValues(task, &schedulerv1.Task{
 					Id:               mockTaskID,
 					Type:             commonv1.TaskType_Normal,
-					ContentLength:    0,
+					ContentLength:    -1,
 					TotalPieceCount:  0,
 					State:            resource.TaskStatePending,
 					PeerCount:        0,