1 | // Copyright 2020 The Marl Authors. |
2 | // |
3 | // Licensed under the Apache License, Version 2.0 (the "License"); |
4 | // you may not use this file except in compliance with the License. |
5 | // You may obtain a copy of the License at |
6 | // |
7 | // https://www.apache.org/licenses/LICENSE-2.0 |
8 | // |
9 | // Unless required by applicable law or agreed to in writing, software |
10 | // distributed under the License is distributed on an "AS IS" BASIS, |
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | // See the License for the specific language governing permissions and |
13 | // limitations under the License. |
14 | |
15 | #include "marl/dag.h" |
16 | |
17 | #include "marl_test.h" |
18 | |
19 | using namespace testing; |
20 | |
21 | namespace { |
22 | |
23 | struct Data { |
24 | std::mutex mutex; |
25 | std::vector<std::string> order; |
26 | |
27 | void push(std::string&& s) { |
28 | std::unique_lock<std::mutex> lock(mutex); |
29 | order.emplace_back(std::move(s)); |
30 | } |
31 | }; |
32 | |
33 | template <typename T> |
34 | std::vector<T> slice(const std::vector<T>& in, size_t from, size_t to) { |
35 | return {in.begin() + from, in.begin() + to}; |
36 | } |
37 | |
38 | } // namespace |
39 | |
40 | // [A] --> [B] --> [C] | |
41 | TEST_P(WithBoundScheduler, DAGChainNoArg) { |
42 | marl::DAG<>::Builder builder; |
43 | |
44 | Data data; |
45 | builder.root() |
46 | .then([&] { data.push("A" ); }) |
47 | .then([&] { data.push("B" ); }) |
48 | .then([&] { data.push("C" ); }); |
49 | |
50 | auto dag = builder.build(); |
51 | dag->run(); |
52 | |
53 | ASSERT_THAT(data.order, ElementsAre("A" , "B" , "C" )); |
54 | } |
55 | |
56 | // [A] --> [B] --> [C] | |
57 | TEST_P(WithBoundScheduler, DAGChain) { |
58 | marl::DAG<Data&>::Builder builder; |
59 | |
60 | builder.root() |
61 | .then([](Data& data) { data.push("A" ); }) |
62 | .then([](Data& data) { data.push("B" ); }) |
63 | .then([](Data& data) { data.push("C" ); }); |
64 | |
65 | auto dag = builder.build(); |
66 | |
67 | Data data; |
68 | dag->run(data); |
69 | |
70 | ASSERT_THAT(data.order, ElementsAre("A" , "B" , "C" )); |
71 | } |
72 | |
73 | // [A] --> [B] --> [C] | |
74 | TEST_P(WithBoundScheduler, DAGRunRepeat) { |
75 | marl::DAG<Data&>::Builder builder; |
76 | |
77 | builder.root() |
78 | .then([](Data& data) { data.push("A" ); }) |
79 | .then([](Data& data) { data.push("B" ); }) |
80 | .then([](Data& data) { data.push("C" ); }); |
81 | |
82 | auto dag = builder.build(); |
83 | |
84 | Data dataA, dataB; |
85 | dag->run(dataA); |
86 | dag->run(dataB); |
87 | dag->run(dataA); |
88 | |
89 | ASSERT_THAT(dataA.order, ElementsAre("A" , "B" , "C" , "A" , "B" , "C" )); |
90 | ASSERT_THAT(dataB.order, ElementsAre("A" , "B" , "C" )); |
91 | } |
92 | |
93 | // /--> [A] | |
94 | // [root] --|--> [B] | |
95 | // \--> [C] | |
96 | TEST_P(WithBoundScheduler, DAGFanOutFromRoot) { |
97 | marl::DAG<Data&>::Builder builder; |
98 | |
99 | auto root = builder.root(); |
100 | root.then([](Data& data) { data.push("A" ); }); |
101 | root.then([](Data& data) { data.push("B" ); }); |
102 | root.then([](Data& data) { data.push("C" ); }); |
103 | |
104 | auto dag = builder.build(); |
105 | |
106 | Data data; |
107 | dag->run(data); |
108 | |
109 | ASSERT_THAT(data.order, UnorderedElementsAre("A" , "B" , "C" )); |
110 | } |
111 | |
112 | // /--> [A] | |
113 | // [root] -->[N]--|--> [B] | |
114 | // \--> [C] | |
115 | TEST_P(WithBoundScheduler, DAGFanOutFromNonRoot) { |
116 | marl::DAG<Data&>::Builder builder; |
117 | |
118 | auto root = builder.root(); |
119 | auto node = root.then([](Data& data) { data.push("N" ); }); |
120 | node.then([](Data& data) { data.push("A" ); }); |
121 | node.then([](Data& data) { data.push("B" ); }); |
122 | node.then([](Data& data) { data.push("C" ); }); |
123 | |
124 | auto dag = builder.build(); |
125 | |
126 | Data data; |
127 | dag->run(data); |
128 | |
129 | ASSERT_THAT(data.order, UnorderedElementsAre("N" , "A" , "B" , "C" )); |
130 | ASSERT_EQ(data.order[0], "N" ); |
131 | ASSERT_THAT(slice(data.order, 1, 4), UnorderedElementsAre("A" , "B" , "C" )); |
132 | } |
133 | |
134 | // /--> [A0] --\ /--> [C0] --\ /--> [E0] --\ | |
135 | // [root] --|--> [A1] --|-->[B]--|--> [C1] --|-->[D]--|--> [E1] --|-->[F] | |
136 | // \--> [C2] --/ |--> [E2] --| | |
137 | // \--> [E3] --/ | |
138 | TEST_P(WithBoundScheduler, DAGFanOutFanIn) { |
139 | marl::DAG<Data&>::Builder builder; |
140 | |
141 | auto root = builder.root(); |
142 | auto a0 = root.then([](Data& data) { data.push("A0" ); }); |
143 | auto a1 = root.then([](Data& data) { data.push("A1" ); }); |
144 | |
145 | auto b = builder.node([](Data& data) { data.push("B" ); }, {a0, a1}); |
146 | |
147 | auto c0 = b.then([](Data& data) { data.push("C0" ); }); |
148 | auto c1 = b.then([](Data& data) { data.push("C1" ); }); |
149 | auto c2 = b.then([](Data& data) { data.push("C2" ); }); |
150 | |
151 | auto d = builder.node([](Data& data) { data.push("D" ); }, {c0, c1, c2}); |
152 | |
153 | auto e0 = d.then([](Data& data) { data.push("E0" ); }); |
154 | auto e1 = d.then([](Data& data) { data.push("E1" ); }); |
155 | auto e2 = d.then([](Data& data) { data.push("E2" ); }); |
156 | auto e3 = d.then([](Data& data) { data.push("E3" ); }); |
157 | |
158 | builder.node([](Data& data) { data.push("F" ); }, {e0, e1, e2, e3}); |
159 | |
160 | auto dag = builder.build(); |
161 | |
162 | Data data; |
163 | dag->run(data); |
164 | |
165 | ASSERT_THAT(data.order, |
166 | UnorderedElementsAre("A0" , "A1" , "B" , "C0" , "C1" , "C2" , "D" , "E0" , |
167 | "E1" , "E2" , "E3" , "F" )); |
168 | ASSERT_THAT(slice(data.order, 0, 2), UnorderedElementsAre("A0" , "A1" )); |
169 | ASSERT_THAT(data.order[2], "B" ); |
170 | ASSERT_THAT(slice(data.order, 3, 6), UnorderedElementsAre("C0" , "C1" , "C2" )); |
171 | ASSERT_THAT(data.order[6], "D" ); |
172 | ASSERT_THAT(slice(data.order, 7, 11), |
173 | UnorderedElementsAre("E0" , "E1" , "E2" , "E3" )); |
174 | ASSERT_THAT(data.order[11], "F" ); |
175 | } |
176 | |