Skip to content

Commit 5e77606

Browse files
artemcmartemcm
authored
Merge pull request #84633 from KushalP/TaskQueueTest
Add tests for `lib/Basic/TaskQueue.cpp`
2 parents 5060e8f + 1075d8c commit 5e77606

File tree

2 files changed

+305
-0
lines changed

2 files changed

+305
-0
lines changed

‎unittests/Basic/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ add_swift_unittest(SwiftBasicTests
3737
StringExtrasTest.cpp
3838
SuccessorMapTest.cpp
3939
TaggedUnionTest.cpp
40+
TaskQueueTest.cpp
4041
ThreadSafeRefCntPointerTest.cpp
4142
TransformRangeTest.cpp
4243
TypeLookupError.cpp

‎unittests/Basic/TaskQueueTest.cpp‎

Lines changed: 304 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,304 @@
1+
//===--- TaskQueueTest.cpp tests -----------------------------------------===//
2+
//
3+
// This source file is part of the Swift.org open source project
4+
//
5+
// Copyright (c) 2025 Apple Inc. and the Swift project authors
6+
// Licensed under Apache License v2.0 with Runtime Library Exception
7+
//
8+
// See https://swift.org/LICENSE.txt for license information
9+
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#include "swift/Basic/TaskQueue.h"
14+
15+
#include "llvm/ADT/ArrayRef.h"
16+
#include "llvm/ADT/StringRef.h"
17+
#include "gtest/gtest.h"
18+
19+
#include <algorithm>
20+
#include <chrono>
21+
#include <mutex>
22+
#include <thread>
23+
24+
#if __is_target_os(darwin) || __is_target_os(linux)
25+
#include <signal.h>
26+
27+
using namespace swift::sys;
28+
29+
TEST(TaskQueueTest, LargeOutput) {
30+
TaskQueue TQ(1);
31+
32+
EXPECT_TRUE(TQ.supportsBufferingOutput());
33+
EXPECT_TRUE(TQ.supportsParallelExecution());
34+
EXPECT_EQ(1U, TQ.getNumberOfParallelTasks());
35+
36+
int OutputSize = 0;
37+
auto TaskFinished = [&](ProcessId Pid, int Result, llvm::StringRef Output,
38+
llvm::StringRef Errors, TaskProcessInformation ProcInfo,
39+
void *Context) -> TaskFinishedResponse {
40+
EXPECT_EQ(0, Result);
41+
OutputSize = Output.size();
42+
return TaskFinishedResponse::ContinueExecution;
43+
};
44+
45+
const char *Args[] = {"-c", "seq 1 100 | while read i; do echo 'This is line number $i with some additional content to make it longer than usual'; done", nullptr};
46+
TQ.addTask("/bin/sh", Args, llvm::ArrayRef<const char *>(), nullptr, false);
47+
48+
EXPECT_FALSE(TQ.execute(nullptr, TaskFinished, nullptr));
49+
EXPECT_GT(OutputSize, 1024U) << "Should produce at least 1KB";
50+
}
51+
52+
TEST(TaskQueueTest, ErrorHandling) {
53+
TaskQueue TQ(1);
54+
int ReceivedResult = 0;
55+
56+
auto TaskFinished = [&](ProcessId Pid, int Result, llvm::StringRef Output,
57+
llvm::StringRef Errors, TaskProcessInformation ProcInfo,
58+
void *Context) -> TaskFinishedResponse {
59+
ReceivedResult = Result;
60+
return TaskFinishedResponse::ContinueExecution;
61+
};
62+
63+
const char *Args[] = {"false", nullptr};
64+
TQ.addTask("/usr/bin/false", Args, llvm::ArrayRef<const char *>(), nullptr, false);
65+
66+
EXPECT_FALSE(TQ.execute(nullptr, TaskFinished, nullptr));
67+
EXPECT_NE(0, ReceivedResult);
68+
}
69+
70+
TEST(TaskQueueTest, SeparateErrorStream) {
71+
TaskQueue TQ(1);
72+
std::string StdoutContent, StderrContent;
73+
74+
auto TaskFinished = [&](ProcessId Pid, int Result, llvm::StringRef Output,
75+
llvm::StringRef Errors, TaskProcessInformation ProcInfo,
76+
void *Context) -> TaskFinishedResponse {
77+
StdoutContent = Output.str();
78+
StderrContent = Errors.str();
79+
return TaskFinishedResponse::ContinueExecution;
80+
};
81+
82+
const char *Args[] = {"-c", "echo 'stdout message'; echo 'stderr message' >&2", nullptr};
83+
TQ.addTask("/bin/sh", Args, llvm::ArrayRef<const char *>(), nullptr, true);
84+
85+
EXPECT_FALSE(TQ.execute(nullptr, TaskFinished, nullptr));
86+
EXPECT_NE(StdoutContent.find("stdout message"), std::string::npos);
87+
EXPECT_NE(StderrContent.find("stderr message"), std::string::npos);
88+
}
89+
90+
TEST(TaskQueueTest, TaskSignalHandling) {
91+
TaskQueue TQ(1);
92+
93+
bool TaskSignalled = false;
94+
int ReceivedSignal = 0;
95+
ProcessId ChildPid = 0;
96+
97+
auto TaskBegan = [&](ProcessId Pid, void *Context) {
98+
ChildPid = Pid;
99+
};
100+
101+
auto TaskSignalledCallback = [&](ProcessId Pid, llvm::StringRef ErrorMsg,
102+
llvm::StringRef Output, llvm::StringRef Errors,
103+
void *Context, std::optional<int> Signal,
104+
TaskProcessInformation ProcInfo) -> TaskFinishedResponse {
105+
TaskSignalled = true;
106+
if (Signal.has_value()) {
107+
ReceivedSignal = Signal.value();
108+
}
109+
return TaskFinishedResponse::ContinueExecution;
110+
};
111+
112+
const char *Args[] = {"-c", "sleep 10", nullptr};
113+
TQ.addTask("/bin/sh", Args, llvm::ArrayRef<const char *>(), nullptr, false);
114+
115+
// Start execution in a separate thread and kill the process
116+
std::thread executor([&] {
117+
TQ.execute(TaskBegan, nullptr, TaskSignalledCallback);
118+
});
119+
120+
std::this_thread::sleep_for(std::chrono::milliseconds(100));
121+
122+
if (ChildPid > 0) {
123+
EXPECT_EQ(0, kill(ChildPid, SIGTERM)) << "Should kill the specific child process we spawned";
124+
}
125+
126+
executor.join();
127+
128+
EXPECT_TRUE(TaskSignalled);
129+
EXPECT_EQ(SIGTERM, ReceivedSignal);
130+
}
131+
132+
TEST(TaskQueueTest, HighConcurrency) {
133+
TaskQueue TQ(10);
134+
135+
int TasksCompleted = 0;
136+
bool AnyTaskFailed = false;
137+
std::mutex CompletionMutex;
138+
139+
auto TaskFinished = [&](ProcessId Pid, int Result, llvm::StringRef Output,
140+
llvm::StringRef Errors, TaskProcessInformation ProcInfo,
141+
void *Context) -> TaskFinishedResponse {
142+
std::lock_guard<std::mutex> lock(CompletionMutex);
143+
TasksCompleted++;
144+
if (Result != 0) {
145+
AnyTaskFailed = true;
146+
} else {
147+
EXPECT_EQ("test", Output.rtrim().str());
148+
}
149+
return TaskFinishedResponse::ContinueExecution;
150+
};
151+
152+
for (int i = 0; i < 50; i++) {
153+
const char *Args[] = {"test", nullptr};
154+
TQ.addTask("/bin/echo", Args, llvm::ArrayRef<const char *>(), nullptr, false);
155+
}
156+
157+
bool ExecutionFailed = TQ.execute(nullptr, TaskFinished, nullptr);
158+
159+
EXPECT_FALSE(ExecutionFailed);
160+
EXPECT_EQ(50, TasksCompleted);
161+
EXPECT_FALSE(AnyTaskFailed);
162+
EXPECT_EQ(10U, TQ.getNumberOfParallelTasks());
163+
}
164+
165+
TEST(TaskQueueTest, TaskBeganCallback) {
166+
TaskQueue TQ(2);
167+
168+
std::vector<ProcessId> StartedTasks;
169+
std::vector<ProcessId> FinishedTasks;
170+
std::mutex TaskMutex;
171+
172+
auto TaskBegan = [&](ProcessId Pid, void *Context) {
173+
std::lock_guard<std::mutex> lock(TaskMutex);
174+
StartedTasks.push_back(Pid);
175+
};
176+
177+
auto TaskFinished = [&](ProcessId Pid, int Result, llvm::StringRef Output,
178+
llvm::StringRef Errors, TaskProcessInformation ProcInfo,
179+
void *Context) -> TaskFinishedResponse {
180+
std::lock_guard<std::mutex> lock(TaskMutex);
181+
FinishedTasks.push_back(Pid);
182+
return TaskFinishedResponse::ContinueExecution;
183+
};
184+
185+
for (int i = 0; i < 3; i++) {
186+
const char *Args[] = {"echo", "test", nullptr};
187+
TQ.addTask("/bin/echo", Args, llvm::ArrayRef<const char *>(), nullptr, false);
188+
}
189+
190+
bool ExecutionFailed = TQ.execute(TaskBegan, TaskFinished, nullptr);
191+
192+
EXPECT_FALSE(ExecutionFailed);
193+
EXPECT_EQ(3U, StartedTasks.size());
194+
EXPECT_EQ(3U, FinishedTasks.size());
195+
196+
std::sort(StartedTasks.begin(), StartedTasks.end());
197+
std::sort(FinishedTasks.begin(), FinishedTasks.end());
198+
EXPECT_EQ(StartedTasks, FinishedTasks);
199+
}
200+
201+
TEST(TaskQueueTest, StopExecutionOnFailure) {
202+
TaskQueue TQ(2);
203+
204+
int TasksCompleted = 0;
205+
int TasksStarted = 0;
206+
207+
auto TaskBegan = [&](ProcessId Pid, void *Context) {
208+
TasksStarted++;
209+
};
210+
211+
auto TaskFinished = [&](ProcessId Pid, int Result, llvm::StringRef Output,
212+
llvm::StringRef Errors, TaskProcessInformation ProcInfo,
213+
void *Context) -> TaskFinishedResponse {
214+
TasksCompleted++;
215+
// Stop execution after the first task completes
216+
if (TasksCompleted >= 1) {
217+
return TaskFinishedResponse::StopExecution;
218+
}
219+
return TaskFinishedResponse::ContinueExecution;
220+
};
221+
222+
for (int i = 0; i < 10; i++) {
223+
const char *Args[] = {"test", nullptr};
224+
TQ.addTask("/bin/echo", Args, llvm::ArrayRef<const char *>(), nullptr, false);
225+
}
226+
227+
bool ExecutionFailed = TQ.execute(TaskBegan, TaskFinished, nullptr);
228+
229+
EXPECT_TRUE(ExecutionFailed);
230+
EXPECT_LE(TasksStarted, 2U) << "Should have started up to 2 tasks";
231+
EXPECT_LT(TasksCompleted, 10) << "Should have completed fewer than 10 tasks";
232+
EXPECT_GE(TasksCompleted, 1) << "At least one task should have completed (the one that triggered StopExecution)";
233+
}
234+
235+
TEST(TaskQueueTest, DummyTaskQueueBasicOperation) {
236+
DummyTaskQueue TQ(2);
237+
238+
int TasksStarted = 0;
239+
int TasksCompleted = 0;
240+
int Context1 = 42;
241+
int Context2 = 100;
242+
std::vector<void*> ReceivedContexts;
243+
244+
auto TaskBegan = [&](ProcessId Pid, void *Context) {
245+
TasksStarted++;
246+
ReceivedContexts.push_back(Context);
247+
};
248+
249+
auto TaskFinished = [&](ProcessId Pid, int Result, llvm::StringRef Output,
250+
llvm::StringRef Errors, TaskProcessInformation ProcInfo,
251+
void *Context) -> TaskFinishedResponse {
252+
TasksCompleted++;
253+
EXPECT_EQ(0, Result);
254+
EXPECT_EQ("Output placeholder\n", Output.str());
255+
EXPECT_NE(nullptr, Context);
256+
if (Context == &Context1) {
257+
EXPECT_EQ(42, *static_cast<int*>(Context));
258+
} else if (Context == &Context2) {
259+
EXPECT_EQ(100, *static_cast<int*>(Context));
260+
}
261+
return TaskFinishedResponse::ContinueExecution;
262+
};
263+
264+
const char *Args[] = {"test", nullptr};
265+
TQ.addTask("/dummy/path1", Args, llvm::ArrayRef<const char *>(), &Context1, false);
266+
TQ.addTask("/dummy/path2", Args, llvm::ArrayRef<const char *>(), &Context2, false);
267+
268+
bool ExecutionFailed = TQ.execute(TaskBegan, TaskFinished, nullptr);
269+
270+
EXPECT_FALSE(ExecutionFailed);
271+
EXPECT_EQ(2, TasksStarted);
272+
EXPECT_EQ(2, TasksCompleted);
273+
EXPECT_EQ(&Context1, ReceivedContexts[0]);
274+
EXPECT_EQ(&Context2, ReceivedContexts[1]);
275+
}
276+
277+
TEST(TaskQueueTest, DummyTaskQueueSeparateErrors) {
278+
DummyTaskQueue TQ(1);
279+
280+
bool TaskExecuted = false;
281+
std::string ReceivedOutput;
282+
std::string ReceivedErrors;
283+
284+
auto TaskFinished = [&](ProcessId Pid, int Result, llvm::StringRef Output,
285+
llvm::StringRef Errors, TaskProcessInformation ProcInfo,
286+
void *Context) -> TaskFinishedResponse {
287+
TaskExecuted = true;
288+
ReceivedOutput = Output.str();
289+
ReceivedErrors = Errors.str();
290+
return TaskFinishedResponse::ContinueExecution;
291+
};
292+
293+
const char *Args[] = {"test", nullptr};
294+
TQ.addTask("/dummy/path", Args, llvm::ArrayRef<const char *>(), nullptr, true);
295+
296+
bool ExecutionFailed = TQ.execute(nullptr, TaskFinished, nullptr);
297+
298+
EXPECT_FALSE(ExecutionFailed);
299+
EXPECT_TRUE(TaskExecuted);
300+
EXPECT_EQ("Output placeholder\n", ReceivedOutput);
301+
EXPECT_EQ("Error placeholder\n", ReceivedErrors);
302+
}
303+
304+
#endif

0 commit comments

Comments
 (0)