btllib: include/btllib/data_stream.hpp Source File

btllib

data_stream.hpp
1 #ifndef BTLLIB_DATA_STREAM_HPP
2 #define BTLLIB_DATA_STREAM_HPP
3
4 #include "status.hpp"
5 #include "util.hpp"
6
7 #include <algorithm>
8 #include <cassert>
9 #include <cerrno>
10 #include <csignal>
11 #include <cstdarg>
12 #include <cstdio>
13 #include <cstdlib>
14 #include <cstring>
15 #include <map>
16 #include <mutex>
17 #include <string>
18 #include <vector>
19
20 #include <dlfcn.h>
21 #include <fcntl.h>
22 #include <sys/stat.h>
23 #include <sys/types.h>
24 #include <sys/wait.h>
25 #include <unistd.h>
26
27 namespace btllib {
28
29 static const int PIPE_READ_END = 0;
30 static const int PIPE_WRITE_END = 1;
31 static const int COMM_BUFFER_SIZE = 1024;
32 static const mode_t PIPE_PERMISSIONS = 0666;
33
34 using PipeId = unsigned long;
35 class _Pipeline;
36
37 inline bool&
38 process_spawner_initialized()
39 {
40 static bool _process_spawner_initialized;
41 return _process_spawner_initialized;
42 }
43 inline int*
44 process_spawner_parent2child_fd()
45 {
46 static int _process_spawner_parent2child_fd[2];
47 return _process_spawner_parent2child_fd;
48 }
49 inline int*
50 process_spawner_child2parent_fd()
51 {
52 static int _process_spawner_child2parent_fd[2];
53 return _process_spawner_child2parent_fd;
54 }
55 inline std::mutex&
56 process_spawner_comm_mutex()
57 {
58 static std::mutex _process_spawner_comm_mutex;
59 return _process_spawner_comm_mutex;
60 };
61 inline std::vector<pid_t>&
62 may_fail()
63 {
64 static std::vector<pid_t> _may_fail;
65 return _may_fail;
66 }
67 inline std::mutex&
68 may_fail_mutex()
69 {
70 static std::mutex _may_fail_mutex;
71 return _may_fail_mutex;
72 }
73 inline PipeId
74 new_pipe_id()
75 {
76 static PipeId _last_pipe_id = 0;
77 return _last_pipe_id++;
78 }
79 inline std::map<std::string, _Pipeline>&
80 pipeline_map()
81 {
82 static std::map<std::string, _Pipeline> _pipeline_map;
83 return _pipeline_map;
84 }
85
86 static inline std::string
87 get_pipepath(const PipeId id)
88 {
89 return "btllib-" + std::to_string(getpid()) + "-" + std::to_string(id);
90 }
91
92 class DataStream
93 {
94 public:
95 enum Operation
96 {
97 READ,
98 WRITE,
99 APPEND,
100 CLOSE
101 };
102
103 DataStream(const std::string& path, Operation op);
104 ~DataStream() { close(); }
105 void close();
106
107 FILE* operator*() const { return file; }
108 FILE* operator->() const { return file; }
109 operator FILE*() const { return file; }
110
111 protected:
112 std::string streampath;
113 Operation op;
114 std::string pipepath;
115 FILE* file = nullptr;
116 bool closed = false;
117 };
118
119 class DataSource : public DataStream
120 {
121
122 public:
123 DataSource(const std::string& path)
124 : DataStream(path, READ)
125 {}
126 };
127
128 class DataSink : public DataStream
129 {
130
131 public:
132 DataSink(const std::string& path, bool append = false)
133 : DataStream(path, append ? APPEND : WRITE)
134 {}
135 };
136
137 inline DataStream::DataStream(const std::string& path, Operation op)
138 : streampath(path)
139 , op(op)
140 {
141 std::unique_lock<std::mutex> lock(process_spawner_comm_mutex());
142
143 write(process_spawner_parent2child_fd()[PIPE_WRITE_END], &op, sizeof(op));
144
145 size_t pathlen = path.size() + 1;
146 check_error(pathlen > COMM_BUFFER_SIZE,
147 "Stream path length too large for the buffer.");
148 write(process_spawner_parent2child_fd()[PIPE_WRITE_END],
149 &pathlen,
150 sizeof(pathlen));
151 write(
152 process_spawner_parent2child_fd()[PIPE_WRITE_END], path.c_str(), pathlen);
153
154 char buf[COMM_BUFFER_SIZE];
155 read(process_spawner_child2parent_fd()[PIPE_READ_END],
156 &pathlen,
157 sizeof(pathlen));
158 read(process_spawner_child2parent_fd()[PIPE_READ_END], buf, pathlen);
159 pipepath = buf;
160
161 file = fopen(pipepath.c_str(), op == READ ? "r" : "w");
162 unlink(pipepath.c_str());
163 }
164
165 inline void
166 DataStream::close()
167 {
168 if (!closed) {
169 std::unique_lock<std::mutex> lock(process_spawner_comm_mutex());
170
171 if (op == READ) {
172 op = CLOSE;
173 if (file != stdin) {
174 write(
175 process_spawner_parent2child_fd()[PIPE_WRITE_END], &op, sizeof(op));
176
177 size_t pathlen = pipepath.size() + 1;
178 check_error(pathlen > COMM_BUFFER_SIZE,
179 "Stream path length too large for the buffer.");
180 write(process_spawner_parent2child_fd()[PIPE_WRITE_END],
181 &pathlen,
182 sizeof(pathlen));
183 write(process_spawner_parent2child_fd()[PIPE_WRITE_END],
184 pipepath.c_str(),
185 pathlen);
186
187 read(process_spawner_child2parent_fd()[PIPE_READ_END], &op, 1);
188
189 std::fclose(file);
190 }
191 } else if (op == WRITE || op == APPEND) {
192 op = CLOSE;
193 if (file != stdout) {
194 std::fclose(file);
195
196 write(
197 process_spawner_parent2child_fd()[PIPE_WRITE_END], &op, sizeof(op));
198
199 size_t pathlen = pipepath.size() + 1;
200 check_error(pathlen > COMM_BUFFER_SIZE,
201 "Stream path length too large for the buffer.");
202 write(process_spawner_parent2child_fd()[PIPE_WRITE_END],
203 &pathlen,
204 sizeof(pathlen));
205 write(process_spawner_parent2child_fd()[PIPE_WRITE_END],
206 pipepath.c_str(),
207 pathlen);
208
209 read(process_spawner_child2parent_fd()[PIPE_READ_END], &op, 1);
210 }
211 }
212
213 closed = true;
214 }
215 }
216
217 class _Pipeline
218 {
219
220 public:
221 enum Direction
222 {
223 SOURCE,
224 SINK
225 };
226
227 _Pipeline() {}
228
229 _Pipeline(std::string pipepath,
230 Direction direction,
231 pid_t pid_first,
232 pid_t pid_last)
233 : pipepath(std::move(pipepath))
234 , direction(direction)
235 , pid_first(pid_first)
236 , pid_last(pid_last)
237 {}
238
239 void finish();
240
241 std::string pipepath;
242 Direction direction = SOURCE;
243 pid_t pid_first = -1;
244 pid_t pid_last = -1;
245 bool closed = false;
246 };
247
248 inline void
249 _Pipeline::finish()
250 {
251 if (!closed) {
252 if (direction == SOURCE) {
253 {
254 std::unique_lock<std::mutex> lock(may_fail_mutex());
255 may_fail().push_back(pid_first);
256 }
257 kill(pid_first, SIGTERM);
258 int status;
259 waitpid(pid_last, &status, 0);
260 } else if (direction == SINK) {
261 int status;
262 waitpid(pid_last, &status, 0);
263 }
264
265 closed = true;
266 }
267 }
268
269 static inline bool
270 process_spawner_init();
271
272 static const bool process_spawner_initializer = process_spawner_init();
273
274 static inline void
275 sigchld_handler(const int sig)
276 {
277 assert(sig == SIGCHLD);
278 (void)sig;
279
280 pid_t pid;
281 int status;
282 while ((pid = waitpid(-1, &status, WNOHANG)) > 0) {
283 if (status != 0) {
284 {
285 std::unique_lock<std::mutex> lock(may_fail_mutex());
286 auto it = std::find(may_fail().begin(), may_fail().end(), pid);
287 if (it != may_fail().end()) {
288 may_fail().erase(it);
289 return;
290 }
291 }
292
293 if (WIFEXITED(status)) { // NOLINT
294 std::cerr << "PID " << pid << " exited with status "
295 << WEXITSTATUS(status) << std::endl; // NOLINT
296 } else if (WIFSIGNALED(status)) { // NOLINT
297 std::cerr << "PID " << pid << " killed by signal "
298 << WTERMSIG(status) // NOLINT
299 << std::endl;
300 } else {
301 std::cerr << "PID " << pid << " exited with code " << status
302 << std::endl;
303 }
304 std::exit(EXIT_FAILURE);
305 }
306 }
307 if (pid == -1 && errno != ECHILD) {
308 std::perror("waitpid");
309 std::exit(EXIT_FAILURE);
310 }
311 }
312
313 static inline std::string
314 get_pipeline_cmd(const std::string& path, DataStream::Operation op);
315
316 static inline _Pipeline
317 run_pipeline_cmd(const std::string& cmd, DataStream::Operation op);
318
319 static inline bool
320 process_spawner_init()
321 {
322 if (!process_spawner_initialized()) {
323 process_spawner_initialized() = true;
324
325 process_spawner_parent2child_fd()[PIPE_READ_END] = -1;
326 process_spawner_parent2child_fd()[PIPE_WRITE_END] = -1;
327 process_spawner_child2parent_fd()[PIPE_READ_END] = -1;
328 process_spawner_child2parent_fd()[PIPE_WRITE_END] = -1;
329 check_error(pipe(process_spawner_parent2child_fd()) == -1,
330 "Error opening a pipe.");
331 check_error(pipe(process_spawner_child2parent_fd()) == -1,
332 "Error opening a pipe.");
333
334 pid_t pid = fork();
335 if (pid == 0) {
336 close(process_spawner_parent2child_fd()[PIPE_WRITE_END]);
337 close(process_spawner_child2parent_fd()[PIPE_READ_END]);
338
339 struct sigaction action; // NOLINT
340 action.sa_handler = sigchld_handler;
341 sigemptyset(&action.sa_mask);
342 action.sa_flags = SA_RESTART;
343 sigaction(SIGCHLD, &action, nullptr);
344
345 DataStream::Operation op;
346 char buf[COMM_BUFFER_SIZE];
347 size_t pathlen;
348 _Pipeline pipeline;
349 for (;;) {
350 if (read(process_spawner_parent2child_fd()[PIPE_READ_END],
351 &op,
352 sizeof(op)) <= 0) {
353 kill(0, SIGTERM);
354 std::exit(EXIT_SUCCESS);
355 }
356
357 read(process_spawner_parent2child_fd()[PIPE_READ_END],
358 &pathlen,
359 sizeof(pathlen));
360 read(process_spawner_parent2child_fd()[PIPE_READ_END], buf, pathlen);
361
362 switch (op) {
363 case DataStream::Operation::READ:
364 case DataStream::Operation::WRITE:
365 case DataStream::Operation::APPEND:
366 pipeline = run_pipeline_cmd(get_pipeline_cmd(buf, op), op);
367
368 pathlen = pipeline.pipepath.size() + 1;
369 check_error(pathlen > COMM_BUFFER_SIZE,
370 "Stream path length too large for the buffer.");
371 write(process_spawner_child2parent_fd()[PIPE_WRITE_END],
372 &pathlen,
373 sizeof(pathlen));
374 write(process_spawner_child2parent_fd()[PIPE_WRITE_END],
375 pipeline.pipepath.c_str(),
376 pathlen);
377
378 pipeline_map()[pipeline.pipepath] = pipeline;
379 break;
380 case DataStream::Operation::CLOSE:
381 pipeline = pipeline_map()[std::string(buf)];
382 pipeline.finish();
383 pipeline_map().erase(std::string(buf));
384 write(process_spawner_child2parent_fd()[PIPE_WRITE_END], &op, 1);
385 break;
386 default:
387 log_error("Invalid stream operation.");
388 std::exit(EXIT_FAILURE);
389 }
390 }
391 }
392 close(process_spawner_parent2child_fd()[PIPE_READ_END]);
393 close(process_spawner_child2parent_fd()[PIPE_WRITE_END]);
394 }
395 return true;
396 }
397
398 static inline std::string
399 get_pipeline_cmd(const std::string& path, DataStream::Operation op)
400 {
401 struct Datatype
402 {
403 std::vector<std::string> prefixes;
404 std::vector<std::string> suffixes;
405 std::vector<std::string> cmds_check_existence;
406 std::vector<std::string> read_cmds;
407 std::vector<std::string> write_cmds;
408 std::vector<std::string> append_cmds;
409 };
410
411 // clang-format off
412 static const Datatype DATATYPES[]{
413 { { "http://", "https://", "ftp://" }, {}, { "which wget" }, { "wget -O-" }, { "" }, { "" } },
414 { {}, { ".url" }, { "which wget" }, { "wget -O- -i" }, { "" }, { "" } },
415 { {}, { ".ar" }, { "which ar" }, { "ar -p" }, { "" }, { "" } },
416 { {}, { ".tar" }, { "which tar" }, { "tar -xOf" }, { "" }, { "" } },
417 { {}, { ".tgz" }, { "which tar" }, { "tar -zxOf" }, { "" }, { "" } },
418 { {}, { ".gz", ".z" }, { "which pigz", "which gzip" }, { "pigz -dc", "gzip -dc" }, { "pigz >", "gzip >" }, { "pigz >>", "gzip >>" } },
419 { {}, { ".bz2" }, { "which bzip2" }, { "bunzip2 -dc" }, { "bzip2 >" }, { "bzip2 >>" } },
420 { {}, { ".xz" }, { "which xz" }, { "unxz -dc" }, { "xz -T0 >" }, { "xz -T0 >>" } },
421 { {}, { ".7z" }, { "which 7z" }, { "7z -so e" }, { "7z -si a" }, { "7z -si a" } },
422 { {}, { ".zip" }, { "which zip" }, { "unzip -p" }, { "" }, { "" } },
423 { {}, { ".bam", ".cram" }, { "which samtools" }, { "samtools view -h" }, { "samtools -Sb - >" }, { "samtools -Sb - >>" } },
424 };
425 // clang-format on
426 std::string default_cmd = "cat";
427 if (op == DataStream::Operation::WRITE) {
428 default_cmd += " >";
429 } else if (op == DataStream::Operation::APPEND) {
430 default_cmd += " >>";
431 }
432
433 std::string path_trimmed = path;
434 std::vector<std::string> cmd_layers;
435 for (;;) {
436 bool found_datatype = false;
437 for (const auto& datatype : DATATYPES) {
438 size_t trim_start = 0, trim_end = 0;
439 bool this_datatype = false;
440 for (const auto& prefix : datatype.prefixes) {
441 if (starts_with(path_trimmed, prefix)) {
442 this_datatype = true;
443 trim_start += prefix.size();
444 break;
445 }
446 }
447 for (const auto& suffix : datatype.suffixes) {
448 if (ends_with(path_trimmed, suffix)) {
449 this_datatype = true;
450 trim_end += suffix.size();
451 break;
452 }
453 }
454
455 if (this_datatype) {
456 found_datatype = true;
457 bool found_cmd = false;
458 int cmd_idx = 0;
459 for (const auto& existence_cmd : datatype.cmds_check_existence) {
460 bool good = true;
461 auto sub_cmds = split(existence_cmd, "&&");
462 std::for_each(sub_cmds.begin(), sub_cmds.end(), trim);
463 for (const auto& sub_cmd : sub_cmds) {
464 auto args = split(sub_cmd, " ");
465 std::for_each(args.begin(), args.end(), trim);
466
467 char* const* argv = new char*[args.size() + 2];
468 ((char*&)(argv[0])) = (char*)(args[0].c_str());
469 for (size_t i = 0; i < args.size(); i++) {
470 ((char*&)(argv[i + 1])) = (char*)(args[i].c_str());
471 }
472 ((char*&)(argv[args.size() + 1])) = nullptr;
473
474 pid_t pid;
475 {
476 std::unique_lock<std::mutex> lock(may_fail_mutex());
477 pid = fork();
478 may_fail().push_back(pid);
479 }
480 if (pid == 0) {
481 int null_fd = open("/dev/null", O_WRONLY, 0);
482 dup2(null_fd, STDOUT_FILENO);
483 dup2(null_fd, STDERR_FILENO);
484 close(null_fd);
485
486 execvp(argv[0], argv + 1);
487 log_error("exec failed.");
488 std::exit(EXIT_FAILURE);
489 } else {
490 delete[] argv;
491 check_error(pid == -1, "Error on fork.");
492 int status;
493 waitpid(pid, &status, 0);
494 if (WIFSIGNALED(status) ||
495 (WIFEXITED(status) && WEXITSTATUS(status) != 0)) { // NOLINT
496 good = false;
497 break;
498 }
499 {
500 std::unique_lock<std::mutex> lock(may_fail_mutex());
501 auto it = std::find(may_fail().begin(), may_fail().end(), pid);
502 if (it != may_fail().end()) {
503 may_fail().erase(it);
504 }
505 }
506 }
507 }
508 if (good) {
509 found_cmd = true;
510 break;
511 }
512 cmd_idx++;
513 }
514
515 if (found_cmd) {
516 std::string cmd;
517 switch (op) {
518 case DataStream::Operation::READ:
519 cmd = datatype.read_cmds[cmd_idx];
520 break;
521 case DataStream::Operation::WRITE:
522 cmd = datatype.write_cmds[cmd_idx];
523 break;
524 case DataStream::Operation::APPEND:
525 cmd = datatype.append_cmds[cmd_idx];
526 break;
527 default:
528 log_error("Invalid operation");
529 std::exit(EXIT_FAILURE);
530 }
531 if (cmd.empty()) {
532 log_warning("Filetype recognized for '" + path +
533 "', but no tool available to work with it.");
534 } else {
535 cmd_layers.push_back(cmd);
536 }
537 } else {
538 log_warning("Filetype recognized for '" + path +
539 "', but no tool available to work with it.");
540 }
541 path_trimmed.erase(0, trim_start);
542 path_trimmed.erase(path_trimmed.size() - trim_end);
543 }
544 }
545 if (!found_datatype) {
546 break;
547 }
548 }
549 if (cmd_layers.empty()) {
550 cmd_layers.push_back(default_cmd);
551 }
552 if (op == DataStream::Operation::WRITE ||
553 op == DataStream::Operation::APPEND) {
554 std::reverse(cmd_layers.begin(), cmd_layers.end());
555 }
556
557 std::string result_cmd;
558 for (size_t i = 0; i < cmd_layers.size(); i++) {
559 auto& cmd = cmd_layers[i];
560 if (op == DataStream::Operation::WRITE ||
561 op == DataStream::Operation::APPEND) {
562 if (i == cmd_layers.size() - 1) {
563 if (cmd.back() == '>') {
564 cmd += path;
565 } else {
566 cmd += " ";
567 cmd += path;
568 }
569 } else {
570 if (cmd.back() == '>') {
571 while (cmd.back() == '>' || cmd.back() == ' ') {
572 cmd.pop_back();
573 }
574 } else {
575 cmd += " -";
576 }
577 }
578 } else {
579 if (i == 0) {
580 cmd += " ";
581 cmd += path;
582 } else {
583 cmd += " -";
584 }
585 }
586 if (i > 0) {
587 result_cmd += " | ";
588 }
589 result_cmd += cmd;
590 }
591
592 check_error(result_cmd.empty(),
593 (op == DataStream::Operation::READ ? "Error loading from "
594 : "Error saving to ") +
595 path);
596 return result_cmd;
597 }
598
599 static inline _Pipeline
600 run_pipeline_cmd(const std::string& cmd, DataStream::Operation op)
601 {
602 std::string pipepath = get_pipepath(new_pipe_id());
603 unlink(pipepath.c_str());
604 mkfifo(pipepath.c_str(), PIPE_PERMISSIONS);
605
606 auto individual_cmds = split(cmd, " | ");
607 check_error(individual_cmds.empty(),
608 "Error processing data stream commands.");
609 std::reverse(individual_cmds.begin(), individual_cmds.end());
610
611 std::vector<pid_t> pids;
612
613 int input_fd[2], output_fd[2];
614 input_fd[PIPE_READ_END] = -1;
615 input_fd[PIPE_WRITE_END] = -1;
616 output_fd[PIPE_READ_END] = -1;
617 output_fd[PIPE_WRITE_END] = -1;
618
619 size_t i = 0;
620 for (const auto& individual_cmd : individual_cmds) {
621 auto args = split(individual_cmd, " ");
622 std::for_each(args.begin(), args.end(), trim);
623
624 std::string stdout_to_file;
625 decltype(args)::iterator it;
626 for (it = args.begin(); it != args.end(); ++it) {
627 if (it->front() == '>') {
628 stdout_to_file = it->substr(1);
629 break;
630 }
631 }
632 if (it != args.end()) {
633 args.erase(it);
634 }
635
636 char* const* argv = new char*[args.size() + 2];
637 ((char*&)(argv[0])) = (char*)(args[0].c_str());
638 for (size_t i = 0; i < args.size(); i++) {
639 ((char*&)(argv[i + 1])) = (char*)(args[i].c_str());
640 }
641 ((char*&)(argv[args.size() + 1])) = nullptr;
642
643 if (i < individual_cmds.size() - 1) {
644 check_error(pipe(input_fd) == -1, "Error opening a pipe.");
645 fcntl(input_fd[PIPE_READ_END], F_SETFD, FD_CLOEXEC);
646 fcntl(input_fd[PIPE_WRITE_END], F_SETFD, FD_CLOEXEC);
647 }
648
649 pid_t pid = fork();
650 if (pid == 0) {
651 if (op == DataStream::Operation::READ) {
652 if (i == 0) {
653 int fd = open(pipepath.c_str(), O_WRONLY);
654 dup2(fd, STDOUT_FILENO);
655 close(fd);
656 } else {
657 dup2(output_fd[PIPE_WRITE_END], STDOUT_FILENO);
658 close(output_fd[PIPE_READ_END]);
659 close(output_fd[PIPE_WRITE_END]);
660 }
661
662 if (i < individual_cmds.size() - 1) {
663 dup2(input_fd[PIPE_READ_END], STDIN_FILENO);
664 close(input_fd[PIPE_READ_END]);
665 close(input_fd[PIPE_WRITE_END]);
666 }
667
668 execvp(argv[0], argv + 1);
669 log_error("exec failed.");
670 std::exit(EXIT_FAILURE);
671 } else {
672 if (i == individual_cmds.size() - 1) {
673 int fd = open(pipepath.c_str(), O_RDONLY);
674 dup2(fd, STDIN_FILENO);
675 close(fd);
676 } else {
677 dup2(input_fd[PIPE_READ_END], STDIN_FILENO);
678 close(input_fd[PIPE_READ_END]);
679 close(input_fd[PIPE_WRITE_END]);
680 }
681
682 if (!stdout_to_file.empty()) {
683 int outfd =
684 open(stdout_to_file.c_str(),
685 O_WRONLY | O_CREAT |
686 (op == DataStream::Operation::APPEND ? O_APPEND : 0),
687 S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH | S_IWOTH);
688 dup2(outfd, STDOUT_FILENO);
689 close(outfd);
690 } else if (i > 0) {
691 dup2(output_fd[PIPE_WRITE_END], STDOUT_FILENO);
692 close(output_fd[PIPE_READ_END]);
693 close(output_fd[PIPE_WRITE_END]);
694 }
695
696 execvp(argv[0], argv + 1);
697 log_error("exec failed.");
698 exit(EXIT_FAILURE);
699 }
700 }
701 check_error(pid == -1, "Error on fork.");
702
703 delete[] argv;
704
705 pids.push_back(pid);
706
707 if (i > 0) {
708 close(output_fd[PIPE_READ_END]);
709 close(output_fd[PIPE_WRITE_END]);
710 }
711
712 if (i < individual_cmds.size() - 1) {
713 output_fd[PIPE_READ_END] = input_fd[PIPE_READ_END];
714 output_fd[PIPE_WRITE_END] = input_fd[PIPE_WRITE_END];
715 }
716
717 i++;
718 }
719
720 return _Pipeline(pipepath,
721 op == DataStream::Operation::READ
722 ? _Pipeline::Direction::SOURCE
723 : _Pipeline::Direction::SINK,
724 pids.back(),
725 pids.front());
726 }
727
728 } // namespace btllib
729
730 #endif
btllib::_Pipeline
Definition: data_stream.hpp:218
btllib::DataSink
Definition: data_stream.hpp:129
btllib::DataSource
Definition: data_stream.hpp:120
btllib::DataStream
Definition: data_stream.hpp:93