Skip to content

Commit 385d276

Browse files
committed
TS SDK: reject a Dag whose order-only edges form a cycle
1 parent 4166797 commit 385d276

5 files changed

Lines changed: 399 additions & 1 deletion

File tree

‎airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -344,6 +344,8 @@ Tasks and groups share one id namespace, as they do in Python, so a Dag cannot h
344344
group called ``staging``. A ``.`` is what separates a group from what it holds, so it cannot appear in
345345
an id of either.
346346

347+
.. warning:: A cycle is not allowed. The Dag is rejected when it is read, naming the tasks on the cycle.
348+
347349
``new Dag`` and ``dag.task`` both take a trailing spec of Airflow options:
348350
``{ schedule: "@daily", tags: ["etl"] }`` for the Dag, ``{ retries: 2, retryDelay: 30 }`` for a task.
349351

‎ts-sdk/src/sdk/cycle.ts‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
/*!
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
// Finds a cycle in a Dag's task graph.
21+
//
22+
// Only the edges `before`/`after` draw can form one. A cycle through *wiring*
23+
// is unrepresentable rather than rejected: a `TaskRef` exists only once its
24+
// producing call has returned, so there is no way to write one. Both kinds of
25+
// edge are searched together all the same, because a cycle can run through one
26+
// of each — `transform({ extracted })` and then `extract.after(transform)`.
27+
28+
/** A directed edge between two task IDs. */
29+
export interface TaskEdge {
30+
readonly upstream: string;
31+
readonly downstream: string;
32+
}
33+
34+
/**
35+
* The tasks on a cycle, or undefined when there is none.
36+
*
37+
* The returned path starts and ends on the same task, so it reads as the cycle
38+
* it describes: `["a", "b", "a"]` is `a >> b >> a`.
39+
*
40+
* Iterative rather than recursive: a Dag of a few thousand chained tasks is
41+
* within reach for a generated Dag file, and would overflow the stack.
42+
*/
43+
export function findTaskCycle(
44+
taskIds: Iterable<string>,
45+
edges: Iterable<TaskEdge>,
46+
): string[] | undefined {
47+
const downstreamOf = new Map<string, string[]>();
48+
for (const taskId of taskIds) downstreamOf.set(taskId, []);
49+
for (const { upstream, downstream } of edges) {
50+
// An edge naming a task the Dag does not hold cannot be part of a cycle
51+
// among the tasks it does, and the caller reports it in its own terms.
52+
downstreamOf.get(upstream)?.push(downstream);
53+
}
54+
55+
// Three states, as a depth-first cycle search needs: unvisited, on the
56+
// current path, and finished. Finding a task already on the path is the
57+
// cycle; finding a finished one only means the search has been there before.
58+
const onPath = new Set<string>();
59+
const finished = new Set<string>();
60+
const path: string[] = [];
61+
62+
for (const root of downstreamOf.keys()) {
63+
if (finished.has(root)) continue;
64+
// Each frame is a task and how far its downstream list has been walked.
65+
const stack: { taskId: string; next: number }[] = [{ taskId: root, next: 0 }];
66+
onPath.add(root);
67+
path.push(root);
68+
69+
while (stack.length > 0) {
70+
const frame = stack[stack.length - 1]!;
71+
const downstream = downstreamOf.get(frame.taskId) ?? [];
72+
if (frame.next === downstream.length) {
73+
finished.add(frame.taskId);
74+
onPath.delete(frame.taskId);
75+
path.pop();
76+
stack.pop();
77+
continue;
78+
}
79+
const next = downstream[frame.next]!;
80+
frame.next += 1;
81+
if (onPath.has(next)) return [...path.slice(path.indexOf(next)), next];
82+
if (finished.has(next) || !downstreamOf.has(next)) continue;
83+
onPath.add(next);
84+
path.push(next);
85+
stack.push({ taskId: next, next: 0 });
86+
}
87+
}
88+
return undefined;
89+
}

‎ts-sdk/src/sdk/dag.ts‎

Lines changed: 59 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import {
3030
} from "../generated/dag-schema-fields.js";
3131
import { argListValues, isArgList, type ArgList } from "./arg-list.js";
3232
import { brand, DUPLICATE_COPY_HINT, hasBrand } from "./brand.js";
33+
import { findTaskCycle, type TaskEdge } from "./cycle.js";
3334
import type { JsonValue } from "./client-types.js";
3435
import type { TaskFunction } from "./task.js";
3536

@@ -891,10 +892,67 @@ export class Dag {
891892
);
892893
}
893894
}
894-
// Last, so a Dag that failed the check reports that same failure again
895+
const cycle = findTaskCycle(this.#tasks.keys(), this.#taskEdges());
896+
if (cycle) {
897+
throw new Error(
898+
`Dag "${this.dagId}" has a cycle: ${cycle.join(" >> ")}. An edge drawn with ` +
899+
"before() or after() runs in one direction, so a Dag cannot come back to a task it " +
900+
"has already run.",
901+
);
902+
}
903+
// Last, so a Dag that failed a check reports that same failure again
895904
// rather than reporting itself as already read.
896905
this.#finalized = true;
897906
}
907+
908+
/**
909+
* Every edge of this Dag, between task IDs.
910+
*
911+
* Both kinds are flattened onto the same graph: the wiring a factory call
912+
* recorded, and the order-only edges, with a group endpoint standing for
913+
* every task the group holds. A cycle can run through one of each, so
914+
* searching either kind alone would miss it.
915+
*/
916+
*#taskEdges(): Generator<TaskEdge> {
917+
for (const [downstream, inputs] of this.#inputs) {
918+
for (const value of Object.values(inputs)) {
919+
if (isTaskRef(value)) yield { upstream: value.taskId, downstream };
920+
}
921+
}
922+
for (const { upstream, downstream } of this.#orderEdges.values()) {
923+
for (const from of this.#tasksOf(upstream)) {
924+
for (const to of this.#tasksOf(downstream)) {
925+
yield { upstream: from, downstream: to };
926+
}
927+
}
928+
}
929+
}
930+
931+
/**
932+
* The tasks an edge endpoint stands for: the task itself, or every task a
933+
* group holds, nested groups included.
934+
*
935+
* An empty group stands for no task, so an edge to one constrains nothing —
936+
* which is also why it cannot put a Dag in a cycle.
937+
*/
938+
#tasksOf(nodeId: string): readonly string[] {
939+
const group = this.#groups.get(nodeId);
940+
if (group === undefined) return [nodeId];
941+
const tasks: string[] = [];
942+
// Breadth-first over the group tree, the queue growing as children are
943+
// found. It terminates because a child group is only ever created after
944+
// its parent, so the tree cannot close on itself.
945+
const pending = [group];
946+
for (let i = 0; i < pending.length; i += 1) {
947+
const current = pending[i]!;
948+
tasks.push(...current.taskIds);
949+
for (const childId of current.childGroupIds) {
950+
const child = this.#groups.get(childId);
951+
if (child) pending.push(child);
952+
}
953+
}
954+
return tasks;
955+
}
898956
}
899957

900958
// What Airflow accepts as a task id, and so what a name taken from a handler

‎ts-sdk/tests/sdk/cycle.test.ts‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
/*!
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
import { describe, expect, it } from "vitest";
21+
import { findTaskCycle, type TaskEdge } from "../../src/sdk/cycle.js";
22+
23+
/** `"a>b c>d"` as edges, so a graph reads as the shape it is. */
24+
function edges(spec: string): TaskEdge[] {
25+
if (spec === "") return [];
26+
return spec.split(" ").map((pair) => {
27+
const [upstream, downstream] = pair.split(">");
28+
return { upstream: upstream!, downstream: downstream! };
29+
});
30+
}
31+
32+
describe("findTaskCycle", () => {
33+
it.each([
34+
["no tasks at all", [], ""],
35+
["one task and no edges", ["a"], ""],
36+
["a chain", ["a", "b", "c"], "a>b b>c"],
37+
["a diamond, where one task is reached twice", ["a", "b", "c", "d"], "a>b a>c b>d c>d"],
38+
["a fan-out", ["a", "b", "c"], "a>b a>c"],
39+
["a fan-in", ["a", "b", "c"], "a>c b>c"],
40+
["two disconnected components", ["a", "b", "c", "d"], "a>b c>d"],
41+
["an isolated task beside a chain", ["a", "b", "lonely"], "a>b"],
42+
])("finds none in %s", (_label, taskIds, spec) => {
43+
expect(findTaskCycle(taskIds, edges(spec))).toBeUndefined();
44+
});
45+
46+
it.each([
47+
["a self-edge", ["a"], "a>a", ["a", "a"]],
48+
["a two-task cycle", ["a", "b"], "a>b b>a", ["a", "b", "a"]],
49+
["a three-task cycle", ["a", "b", "c"], "a>b b>c c>a", ["a", "b", "c", "a"]],
50+
])("reports the tasks on %s", (_label, taskIds, spec, expected) => {
51+
expect(findTaskCycle(taskIds, edges(spec))).toEqual(expected);
52+
});
53+
54+
it("reports only the cycle, not the path that led into it", () => {
55+
// `entry` reaches the cycle but is not on it, so naming it would send a
56+
// reader looking for an edge that does not exist.
57+
expect(findTaskCycle(["entry", "a", "b"], edges("entry>a a>b b>a"))).toEqual(["a", "b", "a"]);
58+
});
59+
60+
it("finds a cycle in a component the search reaches second", () => {
61+
expect(findTaskCycle(["a", "b", "x", "y"], edges("a>b x>y y>x"))).toEqual(["x", "y", "x"]);
62+
});
63+
64+
it("ignores an edge naming a task the Dag does not hold", () => {
65+
// The Dag reports an unknown endpoint in its own terms; it cannot be part
66+
// of a cycle among the tasks that do exist.
67+
expect(findTaskCycle(["a"], edges("a>ghost ghost>a"))).toBeUndefined();
68+
});
69+
70+
it("handles a chain far longer than the call stack would take", () => {
71+
const taskIds = Array.from({ length: 50_000 }, (_, i) => `t${i}`);
72+
const chain = taskIds.slice(0, -1).map((upstream, i) => ({
73+
upstream,
74+
downstream: taskIds[i + 1]!,
75+
}));
76+
77+
expect(findTaskCycle(taskIds, chain)).toBeUndefined();
78+
expect(
79+
findTaskCycle(taskIds, [...chain, { upstream: "t49999", downstream: "t0" }]),
80+
).toHaveLength(taskIds.length + 1);
81+
});
82+
83+
it("does not re-walk a task the search has already finished with", () => {
84+
// A wide diamond lattice: exponential if finished tasks were revisited.
85+
const taskIds = ["start"];
86+
const wide: TaskEdge[] = [];
87+
let previous = "start";
88+
for (let layer = 0; layer < 30; layer += 1) {
89+
const left = `l${layer}`;
90+
const right = `r${layer}`;
91+
const join = `j${layer}`;
92+
taskIds.push(left, right, join);
93+
wide.push(
94+
{ upstream: previous, downstream: left },
95+
{ upstream: previous, downstream: right },
96+
{ upstream: left, downstream: join },
97+
{ upstream: right, downstream: join },
98+
);
99+
previous = join;
100+
}
101+
102+
expect(findTaskCycle(taskIds, wide)).toBeUndefined();
103+
});
104+
});

0 commit comments

Comments
 (0)