9c7a442ee64fabcbd8a3a581a3f79e85ebbac725
[ccan] / tools / ccanlint / async.c
1 #include "ccanlint.h"
2 #include "../tools.h"
3 #include <sys/types.h>
4 #include <sys/time.h>
5 #include <sys/resource.h>
6 #include <sys/wait.h>
7 #include <sys/stat.h>
8 #include <fcntl.h>
9 #include <stdlib.h>
10 #include <unistd.h>
11 #include <err.h>
12 #include <assert.h>
13 #include <ccan/lbalance/lbalance.h>
14 #include <ccan/tlist/tlist.h>
15 #include <ccan/grab_file/grab_file.h>
16 #include <ccan/time/time.h>
17 #include <ccan/talloc/talloc.h>
18
19 static struct lbalance *lb;
20 TLIST_TYPE(command, struct command);
21 static struct tlist_command pending = TLIST_INIT(pending);
22 static struct tlist_command running = TLIST_INIT(running);
23 static unsigned int num_running = 0;
24 static struct tlist_command done = TLIST_INIT(done);
25
26 struct command {
27         struct list_node list;
28         char *command;
29         pid_t pid;
30         int output_fd;
31         unsigned int time_ms;
32         struct lbalance_task *task;
33         int status;
34         char *output;
35         bool done;
36         const void *ctx;
37 };
38
39 static void killme(int sig)
40 {
41         kill(-getpid(), SIGKILL);
42 }
43
44 static void run_more(void)
45 {
46         struct command *c;
47
48         while (num_running < lbalance_target(lb)) {
49                 int p[2];
50                 c = tlist_top(&pending, struct command, list);
51                 if (!c)
52                         break;
53
54                 if (pipe(p) != 0)
55                         err(1, "Pipe failed");
56                 c->pid = fork();
57                 if (c->pid == -1)
58                         err(1, "Fork failed");
59                 if (c->pid == 0) {
60                         struct itimerval itim;
61
62                         if (dup2(p[1], STDOUT_FILENO) != STDOUT_FILENO
63                             || dup2(p[1], STDERR_FILENO) != STDERR_FILENO
64                             || close(p[0]) != 0
65                             || close(STDIN_FILENO) != 0
66                             || open("/dev/null", O_RDONLY) != STDIN_FILENO)
67                                 exit(128);
68
69                         signal(SIGALRM, killme);
70                         itim.it_interval.tv_sec = itim.it_interval.tv_usec = 0;
71                         itim.it_value = time_from_msec(c->time_ms);
72                         setitimer(ITIMER_REAL, &itim, NULL);
73
74                         c->status = system(c->command);
75                         if (WIFEXITED(c->status))
76                                 exit(WEXITSTATUS(c->status));
77                         /* Here's a hint... */
78                         exit(128 + WTERMSIG(c->status));
79                 }
80
81                 if (tools_verbose)
82                         printf("Running async: %s => %i\n", c->command, c->pid);
83
84                 close(p[1]);
85                 c->output_fd = p[0];
86                 c->task = lbalance_task_new(lb);
87                 tlist_del_from(&pending, c, list);
88                 tlist_add_tail(&running, c, list);
89                 num_running++;
90         }
91 }
92
93 static int destroy_command(struct command *command)
94 {
95         if (!command->done && command->pid) {
96                 kill(-command->pid, SIGKILL);
97                 close(command->output_fd);
98                 num_running--;
99         }
100
101         tlist_del(command, list);
102         return 0;
103 }
104
105 void run_command_async(const void *ctx, unsigned int time_ms,
106                        const char *fmt, ...)
107 {
108         struct command *command;
109         va_list ap;
110
111         assert(ctx);
112
113         if (!lb)
114                 lb = lbalance_new();
115
116         command = talloc(ctx, struct command);
117         command->ctx = ctx;
118         command->time_ms = time_ms;
119         command->pid = 0;
120         command->output = talloc_strdup(command, "");
121         va_start(ap, fmt);
122         command->command = talloc_vasprintf(command, fmt, ap);
123         va_end(ap);
124         tlist_add_tail(&pending, command, list);
125         command->done = false;
126         talloc_set_destructor(command, destroy_command);
127
128         run_more();
129 }
130
131 static void reap_output(void)
132 {
133         fd_set in;
134         struct command *c, *next;
135         int max_fd = 0;
136
137         FD_ZERO(&in);
138
139         tlist_for_each(&running, c, list) {
140                 FD_SET(c->output_fd, &in);
141                 if (c->output_fd > max_fd)
142                         max_fd = c->output_fd;
143         }
144
145         if (select(max_fd+1, &in, NULL, NULL, NULL) < 0)
146                 err(1, "select failed");
147
148         tlist_for_each_safe(&running, c, next, list) {
149                 if (FD_ISSET(c->output_fd, &in)) {
150                         int old_len, len;
151                         old_len = talloc_array_length(c->output);
152                         c->output = talloc_realloc(c, c->output, char,
153                                                    old_len + 1024);
154                         len = read(c->output_fd, c->output + old_len, 1024);
155                         if (len < 0)
156                                 err(1, "Reading from async command");
157                         c->output = talloc_realloc(c, c->output, char,
158                                                    old_len + len);
159                         if (len == 0) {
160                                 struct rusage ru;
161                                 wait4(c->pid, &c->status, 0, &ru);
162                                 if (tools_verbose)
163                                         printf("Finished async %i: %s %u\n",
164                                                c->pid,
165                                                WIFEXITED(c->status)
166                                                ? "exit status"
167                                                : "killed by signal",
168                                                WIFEXITED(c->status)
169                                                ? WEXITSTATUS(c->status)
170                                                : WTERMSIG(c->status));
171                                 lbalance_task_free(c->task, &ru);
172                                 c->task = NULL;
173                                 c->done = true;
174                                 close(c->output_fd);
175                                 tlist_del_from(&running, c, list);
176                                 tlist_add_tail(&done, c, list);
177                                 num_running--;
178                         }
179                 }
180         }
181 }
182
183 void *collect_command(bool *ok, char **output)
184 {
185         struct command *c;
186         const void *ctx;
187
188         while ((c = tlist_top(&done, struct command, list)) == NULL) {
189                 if (tlist_empty(&pending) && tlist_empty(&running))
190                         return NULL;
191                 reap_output();
192                 run_more();
193         }
194
195         *ok = (WIFEXITED(c->status) && WEXITSTATUS(c->status) == 0);
196         ctx = c->ctx;
197         *output = talloc_steal(ctx, c->output);
198         talloc_free(c);
199         return (void *)ctx;
200 }
201
202 /* Compile and link single C file, with object files, async. */
203 void compile_and_link_async(const void *ctx, unsigned int time_ms,
204                             const char *cfile, const char *ccandir,
205                             const char *objs, const char *compiler,
206                             const char *cflags,
207                             const char *libs, const char *outfile)
208 {
209         if (compile_verbose)
210                 printf("Compiling and linking (async) %s\n", outfile);
211         run_command_async(ctx, time_ms,
212                           "%s %s -I%s -o %s %s %s %s",
213                           compiler, cflags,
214                           ccandir, outfile, cfile, objs, libs);
215 }