diff --git a/pkg/payload/task_graph.go b/pkg/payload/task_graph.go
index a37e024ed..6378928ec 100644
--- a/pkg/payload/task_graph.go
+++ b/pkg/payload/task_graph.go
@@ -514,9 +514,8 @@ func RunGraph(ctx context.Context, graph *TaskGraph, maxParallelism int, fn func
 				inflight--
 			case <-ctx.Done():
 				select {
-				case runTask := <-workCh: // workers canceled, so remove any work from the queue ourselves
+				case <-workCh: // workers canceled, so remove any work from the queue ourselves
 					inflight--
-					submitted[runTask.index] = false
 				default:
 				}
 			}
diff --git a/pkg/payload/task_graph_test.go b/pkg/payload/task_graph_test.go
index c8a779676..af60fa89b 100644
--- a/pkg/payload/task_graph_test.go
+++ b/pkg/payload/task_graph_test.go
@@ -949,7 +949,6 @@ func TestRunGraph(t *testing.T) {
 			callbacks: map[string]callbackFn{
 				"a": func(t *testing.T, name string, ctx context.Context, cancelFn func()) error {
 					cancelFn()
-					time.Sleep(time.Second)
 					return nil
 				},
 				"*": func(t *testing.T, name string, ctx context.Context, cancelFn func()) error {