bug; abort called for not started nodes (#140)

flyteorg · May 30, 2020 · 7c5df0d · 7c5df0d
1 parent b18dca0
commit 7c5df0d
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 0 deletions.
diff --git a/pkg/controller/nodes/executor.go b/pkg/controller/nodes/executor.go
@@ -669,6 +669,12 @@ func (c *nodeExecutor) FinalizeHandler(ctx context.Context, execContext executor
 	nodeStatus := nl.GetNodeExecutionStatus(ctx, currentNode.GetID())
 	nodePhase := nodeStatus.GetPhase()
 
+	if nodePhase == v1alpha1.NodePhaseNotYetStarted {
+		logger.Infof(ctx, "Node not yet started, will not finalize")
+		// Nothing to be aborted
+		return nil
+	}
+
 	if canHandleNode(nodePhase) {
 		ctx = contextutils.WithNodeID(ctx, currentNode.GetID())
 
@@ -721,6 +727,13 @@ func (c *nodeExecutor) FinalizeHandler(ctx context.Context, execContext executor
 func (c *nodeExecutor) AbortHandler(ctx context.Context, execContext executors.ExecutionContext, dag executors.DAGStructure, nl executors.NodeLookup, currentNode v1alpha1.ExecutableNode, reason string) error {
 	nodeStatus := nl.GetNodeExecutionStatus(ctx, currentNode.GetID())
 	nodePhase := nodeStatus.GetPhase()
+
+	if nodePhase == v1alpha1.NodePhaseNotYetStarted {
+		logger.Infof(ctx, "Node not yet started, will not finalize")
+		// Nothing to be aborted
+		return nil
+	}
+
 	if canHandleNode(nodePhase) {
 		ctx = contextutils.WithNodeID(ctx, currentNode.GetID())
 

diff --git a/pkg/controller/nodes/executor_test.go b/pkg/controller/nodes/executor_test.go
@@ -1523,3 +1523,35 @@ func Test_nodeExecutor_abort(t *testing.T) {
 		assert.True(t, called)
 	})
 }
+
+func TestNodeExecutor_AbortHandler(t *testing.T) {
+	ctx := context.Background()
+	exec := nodeExecutor{}
+
+	t.Run("not-yet-started", func(t *testing.T) {
+		id := "id"
+		n := &mocks.ExecutableNode{}
+		n.OnGetID().Return(id)
+		nl := &mocks4.NodeLookup{}
+		ns := &mocks.ExecutableNodeStatus{}
+		ns.OnGetPhase().Return(v1alpha1.NodePhaseNotYetStarted)
+		nl.OnGetNodeExecutionStatusMatch(mock.Anything, id).Return(ns)
+		assert.NoError(t, exec.AbortHandler(ctx, nil, nil, nl, n, "aborting"))
+	})
+}
+
+func TestNodeExecutor_FinalizeHandler(t *testing.T) {
+	ctx := context.Background()
+	exec := nodeExecutor{}
+
+	t.Run("not-yet-started", func(t *testing.T) {
+		id := "id"
+		n := &mocks.ExecutableNode{}
+		n.OnGetID().Return(id)
+		nl := &mocks4.NodeLookup{}
+		ns := &mocks.ExecutableNodeStatus{}
+		ns.OnGetPhase().Return(v1alpha1.NodePhaseNotYetStarted)
+		nl.OnGetNodeExecutionStatusMatch(mock.Anything, id).Return(ns)
+		assert.NoError(t, exec.FinalizeHandler(ctx, nil, nil, nl, n))
+	})
+}