Mike Klein | 154e6da | 2017-07-26 15:13:47 -0400 | [diff] [blame] | 1 | /* |
| 2 | * Copyright 2017 Google Inc. |
| 3 | * |
| 4 | * Use of this source code is governed by a BSD-style license that can be |
| 5 | * found in the LICENSE file. |
| 6 | */ |
| 7 | |
| 8 | #include "ok.h" |
Mike Klein | 5f13bef | 2017-09-01 09:13:01 -0400 | [diff] [blame] | 9 | #include <stdlib.h> |
Mike Klein | 154e6da | 2017-07-26 15:13:47 -0400 | [diff] [blame] | 10 | |
| 11 | struct SerialEngine : Engine { |
| 12 | static std::unique_ptr<Engine> Factory(Options) { |
| 13 | SerialEngine engine; |
| 14 | return move_unique(engine); |
| 15 | } |
| 16 | |
| 17 | bool crashproof() override { return false; } |
| 18 | |
| 19 | std::future<Status> spawn(std::function<Status(void)> fn) override { |
| 20 | return std::async(std::launch::deferred, fn); |
| 21 | } |
| 22 | }; |
| 23 | static Register serial("serial", |
| 24 | "Run tasks serially on the main thread of a single process.", |
| 25 | SerialEngine::Factory); |
| 26 | |
| 27 | struct ThreadEngine : Engine { |
| 28 | static std::unique_ptr<Engine> Factory(Options) { |
| 29 | ThreadEngine engine; |
| 30 | return move_unique(engine); |
| 31 | } |
| 32 | |
| 33 | bool crashproof() override { return false; } |
| 34 | |
| 35 | std::future<Status> spawn(std::function<Status(void)> fn) override { |
| 36 | return std::async(std::launch::async, fn); |
| 37 | } |
| 38 | }; |
| 39 | static Register thread("thread", |
| 40 | "Run each task on its own thread of a single process.", |
| 41 | ThreadEngine::Factory); |
| 42 | |
| 43 | #if !defined(_MSC_VER) |
| 44 | #include <sys/wait.h> |
| 45 | #include <unistd.h> |
| 46 | |
| 47 | struct ForkEngine : Engine { |
Mike Klein | 5f13bef | 2017-09-01 09:13:01 -0400 | [diff] [blame] | 48 | int limit; // How many concurrent subprocesses do we allow to run at max? |
| 49 | int alive = 0; // How many concurrent subprocesses do we have running right now? |
| 50 | |
| 51 | static std::unique_ptr<Engine> Factory(Options options) { |
Mike Klein | 154e6da | 2017-07-26 15:13:47 -0400 | [diff] [blame] | 52 | ForkEngine engine; |
Mike Klein | 5f13bef | 2017-09-01 09:13:01 -0400 | [diff] [blame] | 53 | engine.limit = atoi(options("limit", "0").c_str()); |
| 54 | if (engine.limit < 1) { |
| 55 | engine.limit = std::thread::hardware_concurrency(); |
| 56 | } |
Mike Klein | 154e6da | 2017-07-26 15:13:47 -0400 | [diff] [blame] | 57 | return move_unique(engine); |
| 58 | } |
| 59 | |
| 60 | bool crashproof() override { return true; } |
| 61 | |
| 62 | std::future<Status> spawn(std::function<Status(void)> fn) override { |
Mike Klein | 5f13bef | 2017-09-01 09:13:01 -0400 | [diff] [blame] | 63 | if (alive == limit) { |
| 64 | // The caller will wait for a child process to finish then try again. |
| 65 | return std::future<Status>(); |
| 66 | } |
| 67 | |
Mike Klein | 154e6da | 2017-07-26 15:13:47 -0400 | [diff] [blame] | 68 | switch (fork()) { |
| 69 | case 0: |
| 70 | // We are the spawned child process. |
| 71 | // Run fn() and exit() with its Status as our return code. |
| 72 | _exit((int)fn()); |
| 73 | |
| 74 | case -1: |
| 75 | // The OS won't let us fork() another process right now. |
| 76 | // We'll need to wait for at least one live task to finish and try again. |
| 77 | return std::future<Status>(); |
| 78 | |
| 79 | default: |
| 80 | // We succesfully spawned a child process! |
| 81 | // This will wait for any spawned process to finish and return its Status. |
Mike Klein | 5f13bef | 2017-09-01 09:13:01 -0400 | [diff] [blame] | 82 | alive++; |
| 83 | return std::async(std::launch::deferred, [&] { |
Mike Klein | 154e6da | 2017-07-26 15:13:47 -0400 | [diff] [blame] | 84 | do { |
| 85 | int status; |
| 86 | if (wait(&status) > 0) { |
Mike Klein | 5f13bef | 2017-09-01 09:13:01 -0400 | [diff] [blame] | 87 | alive--; |
Mike Klein | 154e6da | 2017-07-26 15:13:47 -0400 | [diff] [blame] | 88 | return WIFEXITED(status) ? (Status)WEXITSTATUS(status) |
| 89 | : Status::Crashed; |
| 90 | } |
| 91 | } while (errno == EINTR); |
| 92 | return Status::None; |
| 93 | }); |
| 94 | } |
| 95 | } |
| 96 | }; |
| 97 | static Register _fork("fork", |
Mike Klein | 5f13bef | 2017-09-01 09:13:01 -0400 | [diff] [blame] | 98 | "Run each task in an independent process with fork(), limit=ncpus.", |
Mike Klein | 154e6da | 2017-07-26 15:13:47 -0400 | [diff] [blame] | 99 | ForkEngine::Factory); |
| 100 | #endif |