Improve logic

apache · Sep 30, 2024 · 312195e · 312195e
1 parent 338e4ce
commit 312195e
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/pkg/cache/task.go b/pkg/cache/task.go
@@ -496,7 +496,7 @@ func (task *Task) beforeTaskCompleted() {
 func (task *Task) releaseAllocation() {
 	// We need to get the termination type from the task, when it is "" we default to the normal termination.
 	// We should log it correctly instead of "" in the logs.
-	terminationType := string(common.GetTerminationTypeFromString(task.terminationType))
+	terminationType := common.GetTerminationTypeFromString(task.terminationType)
 
 	// scheduler api might be nil in some tests
 	if task.context.apiProvider.GetAPIs().SchedulerAPI != nil {
@@ -506,7 +506,7 @@ func (task *Task) releaseAllocation() {
 			zap.String("taskAlias", task.alias),
 			zap.String("allocationKey", task.allocationKey),
 			zap.String("task", task.GetTaskState()),
-			zap.String("terminationType", terminationType))
+			zap.String("terminationType", string(terminationType)))
 
 		// send an AllocationReleaseRequest
 		var releaseRequest *si.AllocationRequest

diff --git a/pkg/common/si_helper.go b/pkg/common/si_helper.go
@@ -121,13 +121,13 @@ func GetTerminationTypeFromString(terminationTypeStr string) si.TerminationType
 	return si.TerminationType_STOPPED_BY_RM
 }
 
-func CreateReleaseRequestForTask(appID, taskID, partition, terminationType string) *si.AllocationRequest {
+func CreateReleaseRequestForTask(appID, taskID, partition string, terminationType si.TerminationType) *si.AllocationRequest {
 	allocToRelease := make([]*si.AllocationRelease, 1)
 	allocToRelease[0] = &si.AllocationRelease{
 		ApplicationID:   appID,
 		AllocationKey:   taskID,
 		PartitionName:   partition,
-		TerminationType: GetTerminationTypeFromString(terminationType),
+		TerminationType: terminationType,
 		Message:         "task completed",
 	}
 

diff --git a/pkg/common/si_helper_test.go b/pkg/common/si_helper_test.go
@@ -32,7 +32,7 @@ const nodeID = "node-01"
 
 func TestCreateReleaseRequestForTask(t *testing.T) {
 	// with allocationKey
-	request := CreateReleaseRequestForTask("app01", "task01", "default", "STOPPED_BY_RM")
+	request := CreateReleaseRequestForTask("app01", "task01", "default", si.TerminationType_STOPPED_BY_RM)
 	assert.Assert(t, request.Releases != nil)
 	assert.Assert(t, request.Releases.AllocationsToRelease != nil)
 	assert.Equal(t, len(request.Releases.AllocationsToRelease), 1)
@@ -41,7 +41,7 @@ func TestCreateReleaseRequestForTask(t *testing.T) {
 	assert.Equal(t, request.Releases.AllocationsToRelease[0].PartitionName, "default")
 	assert.Equal(t, request.Releases.AllocationsToRelease[0].TerminationType, si.TerminationType_STOPPED_BY_RM)
 
-	request = CreateReleaseRequestForTask("app01", "task01", "default", "UNKNOWN_TERMINATION_TYPE")
+	request = CreateReleaseRequestForTask("app01", "task01", "default", si.TerminationType_UNKNOWN_TERMINATION_TYPE)
 	assert.Assert(t, request.Releases != nil)
 	assert.Assert(t, request.Releases.AllocationsToRelease != nil)
 	assert.Equal(t, len(request.Releases.AllocationsToRelease), 1)