Implement simple, fork() based parallel unpacking in rdsquashfs

Signed-off-by: David Oberhollenzer <david.oberhollenzer@sigma-star.at>

4 files changed, 119 insertions, 6 deletions

diff --git a/unpack/fill_files.c b/unpack/fill_files.c
index 6584699..865d050 100644
--- a/unpack/fill_files.c
+++ b/unpack/fill_files.c
@@ -33,7 +33,98 @@ static int fill_files(data_reader_t *data, file_info_t *list, int flags)
 	return 0;
 }
 
-int fill_unpacked_files(fstree_t *fs, data_reader_t *data, int flags)
+static file_info_t *split_list(file_info_t *list, uint64_t threashold)
 {
-	return fill_files(data, fs->files, flags);
+	file_info_t *it, *new = NULL;
+	uint64_t size = 0;
+
+	for (it = list; it != NULL; it = it->next) {
+		if (it->input_file == NULL)
+			continue;
+
+		size += it->size - it->sparse;
+
+		if (size >= threashold) {
+			new = it->next;
+			it->next = NULL;
+			break;
+		}
+	}
+
+	return new;
+}
+
+static uint64_t total_size(file_info_t *list)
+{
+	uint64_t size = 0;
+	file_info_t *it;
+
+	for (it = list; it != NULL; it = it->next) {
+		if (it->input_file != NULL)
+			size += it->size - it->sparse;
+	}
+
+	return size;
+}
+
+int fill_unpacked_files(fstree_t *fs, data_reader_t *data, int flags,
+			unsigned int num_jobs)
+{
+	file_info_t *sublists[num_jobs], *it;
+	int exitstatus, status = 0;
+	uint64_t threshold;
+	unsigned int i;
+	pid_t pid;
+
+	if (num_jobs <= 1) {
+		status = fill_files(data, fs->files, flags);
+
+		for (it = fs->files; it != NULL; it = it->next)
+			free(it->input_file);
+
+		return status;
+	}
+
+	threshold = total_size(fs->files) / num_jobs;
+
+	for (i = 0; i < num_jobs; ++i) {
+		sublists[i] = fs->files;
+
+		fs->files = split_list(fs->files, threshold);
+	}
+
+	for (i = 0; i < num_jobs; ++i) {
+		pid = fork();
+
+		if (pid == 0) {
+			if (fill_files(data, sublists[i], flags))
+				exit(EXIT_FAILURE);
+			exit(EXIT_SUCCESS);
+		}
+
+		if (pid < 0) {
+			perror("fork");
+			status = -1;
+			num_jobs = i;
+			break;
+		}
+	}
+
+	for (i = 0; i < num_jobs; ++i) {
+		do {
+			pid = waitpid(-1, &exitstatus, 0);
+			if (errno == ECHILD)
+				goto out;
+		} while (pid < 0);
+
+		if (!WIFEXITED(exitstatus) ||
+		    WEXITSTATUS(exitstatus) != EXIT_SUCCESS) {
+			status = -1;
+		}
+
+		for (it = sublists[i]; it != NULL; it = it->next)
+			free(it->input_file);
+	}
+out:
+	return status;
 }
diff --git a/unpack/options.c b/unpack/options.c
index c42f16a..e7911e5 100644
--- a/unpack/options.c
+++ b/unpack/options.c
@@ -12,6 +12,7 @@ static struct option long_opts[] = {
 	{ "no-slink", no_argument, NULL, 'L' },
 	{ "no-empty-dir", no_argument, NULL, 'E' },
 	{ "no-sparse", no_argument, NULL, 'Z' },
+	{ "jobs", required_argument, NULL, 'j' },
 	{ "describe", no_argument, NULL, 'd' },
 	{ "chmod", no_argument, NULL, 'C' },
 	{ "chown", no_argument, NULL, 'O' },
@@ -20,7 +21,7 @@ static struct option long_opts[] = {
 	{ "version", no_argument, NULL, 'V' },
 };
 
-static const char *short_opts = "l:c:u:p:DSFLCOEZdqhV";
+static const char *short_opts = "l:c:u:p:DSFLCOEZj:dqhV";
 
 static const char *help_string =
 "Usage: %s [OPTIONS] <squashfs-file>\n"
@@ -50,6 +51,7 @@ static const char *help_string =
 "                            empty after applying the above rules.\n"
 "  --no-sparse, -Z           Do not create sparse files, always write zero\n"
 "                            blocks to disk.\n"
+"  --jobs, -j <count>        Number of parallel unpacking jobs to start.\n"
 "  --chmod, -C               Change permission flags of unpacked files to\n"
 "                            those store in the squashfs image.\n"
 "  --chown, -O               Change ownership of unpacked files to the\n"
@@ -85,7 +87,7 @@ static char *get_path(char *old, const char *arg)
 
 void process_command_line(options_t *opt, int argc, char **argv)
 {
-	int i;
+	int i, j;
 
 	opt->op = OP_NONE;
 	opt->rdtree_flags = 0;
@@ -93,6 +95,7 @@ void process_command_line(options_t *opt, int argc, char **argv)
 	opt->cmdpath = NULL;
 	opt->unpack_root = NULL;
 	opt->image_name = NULL;
+	opt->num_jobs = 1;
 
 	for (;;) {
 		i = getopt_long(argc, argv, short_opts, long_opts, NULL);
@@ -124,6 +127,17 @@ void process_command_line(options_t *opt, int argc, char **argv)
 		case 'Z':
 			opt->flags |= UNPACK_NO_SPARSE;
 			break;
+		case 'j':
+			for (j = 0; optarg[j] != '\0'; ++j) {
+				if (j > 6 || !isdigit(optarg[j]))
+					goto fail_num_jobs;
+			}
+
+			opt->num_jobs = atoi(optarg);
+
+			if (opt->num_jobs < 1)
+				goto fail_num_jobs;
+			break;
 		case 'c':
 			opt->op = OP_CAT;
 			opt->cmdpath = get_path(opt->cmdpath, optarg);
@@ -176,4 +190,7 @@ fail_arg:
 	fprintf(stderr, "Try `%s --help' for more information.\n", __progname);
 	free(opt->cmdpath);
 	exit(EXIT_FAILURE);
+fail_num_jobs:
+	fputs("Expected a positive integer for --jobs.\n", stderr);
+	goto fail_arg;
 }
diff --git a/unpack/rdsquashfs.c b/unpack/rdsquashfs.c
index a3618e3..f165b90 100644
--- a/unpack/rdsquashfs.c
+++ b/unpack/rdsquashfs.c
@@ -89,7 +89,7 @@ int main(int argc, char **argv)
 		if (data == NULL)
 			goto out_fs;
 
-		if (fill_unpacked_files(&fs, data, opt.flags))
+		if (fill_unpacked_files(&fs, data, opt.flags, opt.num_jobs))
 			goto out_fs;
 
 		if (update_tree_attribs(n, opt.flags))
diff --git a/unpack/rdsquashfs.h b/unpack/rdsquashfs.h
index 45c91db..34404b0 100644
--- a/unpack/rdsquashfs.h
+++ b/unpack/rdsquashfs.h
@@ -14,10 +14,13 @@
 #include "util.h"
 
 #include <sys/sysmacros.h>
+#include <sys/types.h>
+#include <sys/wait.h>
 #include <string.h>
 #include <stdlib.h>
 #include <unistd.h>
 #include <getopt.h>
+#include <ctype.h>
 #include <errno.h>
 #include <fcntl.h>
 #include <stdio.h>
@@ -41,6 +44,7 @@ typedef struct {
 	int op;
 	int rdtree_flags;
 	int flags;
+	unsigned int num_jobs;
 	char *cmdpath;
 	const char *unpack_root;
 	const char *image_name;
@@ -52,7 +56,8 @@ int restore_fstree(tree_node_t *root, int flags);
 
 int update_tree_attribs(tree_node_t *root, int flags);
 
-int fill_unpacked_files(fstree_t *fs, data_reader_t *data, int flags);
+int fill_unpacked_files(fstree_t *fs, data_reader_t *data, int flags,
+			unsigned int num_jobs);
 
 void describe_tree(tree_node_t *root, const char *unpack_root);