#include "dentry.h"
#include "lookup_table.h"
#include "xml.h"
+#include "lzx.h"
+#include "xpress.h"
#include <unistd.h>
+#include <semaphore.h>
+#include <pthread.h>
+#include <errno.h>
+
+#ifdef WITH_NTFS_3G
+#include <time.h>
+#include <ntfs-3g/attrib.h>
+#include <ntfs-3g/inode.h>
+#include <ntfs-3g/dir.h>
+#endif
+
+
+#ifdef HAVE_ALLOCA_H
+#include <alloca.h>
+#endif
/* Reopens the FILE* for a WIM read-write. */
return WIMLIB_ERR_RENAME;
}
+ if (write_flags & WIMLIB_WRITE_FLAG_VERBOSE)
+ printf("Successfully renamed `%s' to `%s'\n", tmpfile, wimfile_name);
+
return 0;
}
}
-/* Write the file resources for the current image. */
-static int write_file_resources(WIMStruct *w)
+/* Chunk table that's located at the beginning of each compressed resource in
+ * the WIM. (This is not the on-disk format; the on-disk format just has an
+ * array of offsets.) */
+struct chunk_table {
+ off_t file_offset;
+ u64 num_chunks;
+ u64 original_resource_size;
+ u64 bytes_per_chunk_entry;
+ u64 table_disk_size;
+ u64 cur_offset;
+ u64 *cur_offset_p;
+ u64 offsets[0];
+};
+
+/*
+ * Allocates and initializes a chunk table, and reserves space for it in the
+ * output file.
+ */
+static int
+begin_wim_resource_chunk_tab(const struct lookup_table_entry *lte,
+ FILE *out_fp,
+ off_t file_offset,
+ struct chunk_table **chunk_tab_ret)
+{
+ u64 size = wim_resource_size(lte);
+ u64 num_chunks = (size + WIM_CHUNK_SIZE - 1) / WIM_CHUNK_SIZE;
+ size_t alloc_size = sizeof(struct chunk_table) + num_chunks * sizeof(u64);
+ struct chunk_table *chunk_tab = CALLOC(1, alloc_size);
+ int ret;
+
+ if (!chunk_tab) {
+ ERROR("Failed to allocate chunk table for %"PRIu64" byte "
+ "resource", size);
+ ret = WIMLIB_ERR_NOMEM;
+ goto out;
+ }
+ chunk_tab->file_offset = file_offset;
+ chunk_tab->num_chunks = num_chunks;
+ chunk_tab->original_resource_size = size;
+ chunk_tab->bytes_per_chunk_entry = (size >= (1ULL << 32)) ? 8 : 4;
+ chunk_tab->table_disk_size = chunk_tab->bytes_per_chunk_entry *
+ (num_chunks - 1);
+ chunk_tab->cur_offset = 0;
+ chunk_tab->cur_offset_p = chunk_tab->offsets;
+
+ if (fwrite(chunk_tab, 1, chunk_tab->table_disk_size, out_fp) !=
+ chunk_tab->table_disk_size) {
+ ERROR_WITH_ERRNO("Failed to write chunk table in compressed "
+ "file resource");
+ ret = WIMLIB_ERR_WRITE;
+ goto out;
+ }
+
+ ret = 0;
+out:
+ *chunk_tab_ret = chunk_tab;
+ return ret;
+}
+
+typedef int (*compress_func_t)(const void *, unsigned, void *, unsigned *);
+
+compress_func_t get_compress_func(int out_ctype)
+{
+ if (out_ctype == WIM_COMPRESSION_TYPE_LZX)
+ return lzx_compress;
+ else
+ return xpress_compress;
+}
+
+
+/*
+ * Compresses a chunk of a WIM resource.
+ *
+ * @chunk: Uncompressed data of the chunk.
+ * @chunk_size: Size of the uncompressed chunk in bytes.
+ * @compressed_chunk: Pointer to output buffer of size at least
+ * (@chunk_size - 1) bytes.
+ * @compressed_chunk_len_ret: Pointer to an unsigned int into which the size
+ * of the compressed chunk will be
+ * returned.
+ * @ctype: Type of compression to use. Must be WIM_COMPRESSION_TYPE_LZX
+ * or WIM_COMPRESSION_TYPE_XPRESS.
+ *
+ * Returns zero if compressed succeeded, and nonzero if the chunk could not be
+ * compressed to any smaller than @chunk_size. This function cannot fail for
+ * any other reasons.
+ */
+static int compress_chunk(const u8 chunk[], unsigned chunk_size,
+ u8 compressed_chunk[],
+ unsigned *compressed_chunk_len_ret,
+ int ctype)
+{
+ compress_func_t compress = get_compress_func(ctype);
+ return (*compress)(chunk, chunk_size, compressed_chunk,
+ compressed_chunk_len_ret);
+}
+
+/*
+ * Writes a chunk of a WIM resource to an output file.
+ *
+ * @chunk: Uncompressed data of the chunk.
+ * @chunk_size: Size of the chunk (<= WIM_CHUNK_SIZE)
+ * @out_fp: FILE * to write tho chunk to.
+ * @out_ctype: Compression type to use when writing the chunk (ignored if no
+ * chunk table provided)
+ * @chunk_tab: Pointer to chunk table being created. It is updated with the
+ * offset of the chunk we write.
+ *
+ * Returns 0 on success; nonzero on failure.
+ */
+static int write_wim_resource_chunk(const u8 chunk[], unsigned chunk_size,
+ FILE *out_fp, int out_ctype,
+ struct chunk_table *chunk_tab)
+{
+ const u8 *out_chunk;
+ unsigned out_chunk_size;
+
+ wimlib_assert(chunk_size <= WIM_CHUNK_SIZE);
+
+ if (!chunk_tab) {
+ out_chunk = chunk;
+ out_chunk_size = chunk_size;
+ } else {
+ u8 *compressed_chunk = alloca(chunk_size);
+ int ret;
+
+ ret = compress_chunk(chunk, chunk_size, compressed_chunk,
+ &out_chunk_size, out_ctype);
+ if (ret == 0) {
+ out_chunk = compressed_chunk;
+ } else {
+ out_chunk = chunk;
+ out_chunk_size = chunk_size;
+ }
+ *chunk_tab->cur_offset_p++ = chunk_tab->cur_offset;
+ chunk_tab->cur_offset += out_chunk_size;
+ }
+
+ if (fwrite(out_chunk, 1, out_chunk_size, out_fp) != out_chunk_size) {
+ ERROR_WITH_ERRNO("Failed to write WIM resource chunk");
+ return WIMLIB_ERR_WRITE;
+ }
+ return 0;
+}
+
+/*
+ * Finishes a WIM chunk tale and writes it to the output file at the correct
+ * offset.
+ *
+ * The final size of the full compressed resource is returned in the
+ * @compressed_size_p.
+ */
+static int
+finish_wim_resource_chunk_tab(struct chunk_table *chunk_tab,
+ FILE *out_fp, u64 *compressed_size_p)
+{
+ size_t bytes_written;
+ if (fseeko(out_fp, chunk_tab->file_offset, SEEK_SET) != 0) {
+ ERROR_WITH_ERRNO("Failed to seek to byte %"PRIu64" of output "
+ "WIM file", chunk_tab->file_offset);
+ return WIMLIB_ERR_WRITE;
+ }
+
+ if (chunk_tab->bytes_per_chunk_entry == 8) {
+ array_cpu_to_le64(chunk_tab->offsets, chunk_tab->num_chunks);
+ } else {
+ for (u64 i = 0; i < chunk_tab->num_chunks; i++)
+ ((u32*)chunk_tab->offsets)[i] =
+ cpu_to_le32(chunk_tab->offsets[i]);
+ }
+ bytes_written = fwrite((u8*)chunk_tab->offsets +
+ chunk_tab->bytes_per_chunk_entry,
+ 1, chunk_tab->table_disk_size, out_fp);
+ if (bytes_written != chunk_tab->table_disk_size) {
+ ERROR_WITH_ERRNO("Failed to write chunk table in compressed "
+ "file resource");
+ return WIMLIB_ERR_WRITE;
+ }
+ if (fseeko(out_fp, 0, SEEK_END) != 0) {
+ ERROR_WITH_ERRNO("Failed to seek to end of output WIM file");
+ return WIMLIB_ERR_WRITE;
+ }
+ *compressed_size_p = chunk_tab->cur_offset + chunk_tab->table_disk_size;
+ return 0;
+}
+
+static int prepare_resource_for_read(struct lookup_table_entry *lte
+
+ #ifdef WITH_NTFS_3G
+ , ntfs_inode **ni_ret
+ #endif
+ )
+{
+ if (lte->resource_location == RESOURCE_IN_FILE_ON_DISK
+ && !lte->file_on_disk_fp)
+ {
+ wimlib_assert(lte->file_on_disk);
+ lte->file_on_disk_fp = fopen(lte->file_on_disk, "rb");
+ if (!lte->file_on_disk_fp) {
+ ERROR_WITH_ERRNO("Failed to open the file `%s' for "
+ "reading", lte->file_on_disk);
+ return WIMLIB_ERR_OPEN;
+ }
+ }
+#ifdef WITH_NTFS_3G
+ else if (lte->resource_location == RESOURCE_IN_NTFS_VOLUME
+ && !lte->attr)
+ {
+ struct ntfs_location *loc = lte->ntfs_loc;
+ ntfs_inode *ni;
+ wimlib_assert(loc);
+ ni = ntfs_pathname_to_inode(*loc->ntfs_vol_p, NULL, loc->path_utf8);
+ if (!ni) {
+ ERROR_WITH_ERRNO("Failed to open inode `%s' in NTFS "
+ "volume", loc->path_utf8);
+ return WIMLIB_ERR_NTFS_3G;
+ }
+ lte->attr = ntfs_attr_open(ni,
+ loc->is_reparse_point ? AT_REPARSE_POINT : AT_DATA,
+ (ntfschar*)loc->stream_name_utf16,
+ loc->stream_name_utf16_num_chars);
+ if (!lte->attr) {
+ ERROR_WITH_ERRNO("Failed to open attribute of `%s' in "
+ "NTFS volume", loc->path_utf8);
+ ntfs_inode_close(ni);
+ return WIMLIB_ERR_NTFS_3G;
+ }
+ *ni_ret = ni;
+ }
+#endif
+ return 0;
+}
+
+static void end_wim_resource_read(struct lookup_table_entry *lte
+ #ifdef WITH_NTFS_3G
+ , ntfs_inode *ni
+ #endif
+ )
{
- DEBUG("Writing file resources for image %u.", w->current_image);
- return for_dentry_in_tree(wim_root_dentry(w), write_dentry_resources, w);
+ if (lte->resource_location == RESOURCE_IN_FILE_ON_DISK
+ && lte->file_on_disk_fp) {
+ fclose(lte->file_on_disk_fp);
+ lte->file_on_disk_fp = NULL;
+ }
+#ifdef WITH_NTFS_3G
+ else if (lte->resource_location == RESOURCE_IN_NTFS_VOLUME) {
+ if (lte->attr) {
+ ntfs_attr_close(lte->attr);
+ lte->attr = NULL;
+ }
+ if (ni)
+ ntfs_inode_close(ni);
+ }
+#endif
+}
+
+/*
+ * Writes a WIM resource to a FILE * opened for writing. The resource may be
+ * written uncompressed or compressed depending on the @out_ctype parameter.
+ *
+ * If by chance the resource compresses to more than the original size (this may
+ * happen with random data or files than are pre-compressed), the resource is
+ * instead written uncompressed (and this is reflected in the @out_res_entry by
+ * removing the WIM_RESHDR_FLAG_COMPRESSED flag).
+ *
+ * @lte: The lookup table entry for the WIM resource.
+ * @out_fp: The FILE * to write the resource to.
+ * @out_ctype: The compression type of the resource to write. Note: if this is
+ * the same as the compression type of the WIM resource we
+ * need to read, we simply copy the data (i.e. we do not
+ * uncompress it, then compress it again).
+ * @out_res_entry: If non-NULL, a resource entry that is filled in with the
+ * offset, original size, compressed size, and compression flag
+ * of the output resource.
+ *
+ * Returns 0 on success; nonzero on failure.
+ */
+int write_wim_resource(struct lookup_table_entry *lte,
+ FILE *out_fp, int out_ctype,
+ struct resource_entry *out_res_entry,
+ int flags)
+{
+ u64 bytes_remaining;
+ u64 original_size;
+ u64 old_compressed_size;
+ u64 new_compressed_size;
+ u64 offset;
+ int ret;
+ struct chunk_table *chunk_tab = NULL;
+ bool raw;
+ off_t file_offset;
+#ifdef WITH_NTFS_3G
+ ntfs_inode *ni = NULL;
+#endif
+
+ wimlib_assert(lte);
+
+ /* Original size of the resource */
+ original_size = wim_resource_size(lte);
+
+ /* Compressed size of the resource (as it exists now) */
+ old_compressed_size = wim_resource_compressed_size(lte);
+
+ /* Current offset in output file */
+ file_offset = ftello(out_fp);
+ if (file_offset == -1) {
+ ERROR_WITH_ERRNO("Failed to get offset in output "
+ "stream");
+ return WIMLIB_ERR_WRITE;
+ }
+
+ /* Are the compression types the same? If so, do a raw copy (copy
+ * without decompressing and recompressing the data). */
+ raw = (wim_resource_compression_type(lte) == out_ctype
+ && out_ctype != WIM_COMPRESSION_TYPE_NONE);
+
+ if (raw) {
+ flags |= WIMLIB_RESOURCE_FLAG_RAW;
+ bytes_remaining = old_compressed_size;
+ } else {
+ flags &= ~WIMLIB_RESOURCE_FLAG_RAW;
+ bytes_remaining = original_size;
+ }
+
+ /* Empty resource; nothing needs to be done, so just return success. */
+ if (bytes_remaining == 0)
+ return 0;
+
+ /* Buffer for reading chunks for the resource */
+ u8 buf[min(WIM_CHUNK_SIZE, bytes_remaining)];
+
+ /* If we are writing a compressed resource and not doing a raw copy, we
+ * need to initialize the chunk table */
+ if (out_ctype != WIM_COMPRESSION_TYPE_NONE && !raw) {
+ ret = begin_wim_resource_chunk_tab(lte, out_fp, file_offset,
+ &chunk_tab);
+ if (ret != 0)
+ goto out;
+ }
+
+ /* If the WIM resource is in an external file, open a FILE * to it so we
+ * don't have to open a temporary one in read_wim_resource() for each
+ * chunk. */
+#ifdef WITH_NTFS_3G
+ ret = prepare_resource_for_read(lte, &ni);
+#else
+ ret = prepare_resource_for_read(lte);
+#endif
+ if (ret != 0)
+ goto out;
+
+ /* If we aren't doing a raw copy, we will compute the SHA1 message
+ * digest of the resource as we read it, and verify it's the same as the
+ * hash given in the lookup table entry once we've finished reading the
+ * resource. */
+ SHA_CTX ctx;
+ if (!raw)
+ sha1_init(&ctx);
+
+ /* While there are still bytes remaining in the WIM resource, read a
+ * chunk of the resource, update SHA1, then write that chunk using the
+ * desired compression type. */
+ offset = 0;
+ do {
+ u64 to_read = min(bytes_remaining, WIM_CHUNK_SIZE);
+ ret = read_wim_resource(lte, buf, to_read, offset, flags);
+ if (ret != 0)
+ goto out_fclose;
+ if (!raw)
+ sha1_update(&ctx, buf, to_read);
+ ret = write_wim_resource_chunk(buf, to_read, out_fp,
+ out_ctype, chunk_tab);
+ if (ret != 0)
+ goto out_fclose;
+ bytes_remaining -= to_read;
+ offset += to_read;
+ } while (bytes_remaining);
+
+ /* Raw copy: The new compressed size is the same as the old compressed
+ * size
+ *
+ * Using WIM_COMPRESSION_TYPE_NONE: The new compressed size is the
+ * original size
+ *
+ * Using a different compression type: Call
+ * finish_wim_resource_chunk_tab() and it will provide the new
+ * compressed size.
+ */
+ if (raw) {
+ new_compressed_size = old_compressed_size;
+ } else {
+ if (out_ctype == WIM_COMPRESSION_TYPE_NONE)
+ new_compressed_size = original_size;
+ else {
+ ret = finish_wim_resource_chunk_tab(chunk_tab, out_fp,
+ &new_compressed_size);
+ if (ret != 0)
+ goto out_fclose;
+ }
+ }
+
+ /* Verify SHA1 message digest of the resource, unless we are doing a raw
+ * write (in which case we never even saw the uncompressed data). Or,
+ * if the hash we had before is all 0's, just re-set it to be the new
+ * hash. */
+ if (!raw) {
+ u8 md[SHA1_HASH_SIZE];
+ sha1_final(md, &ctx);
+ if (is_zero_hash(lte->hash)) {
+ copy_hash(lte->hash, md);
+ } else if (!hashes_equal(md, lte->hash)) {
+ ERROR("WIM resource has incorrect hash!");
+ if (lte->resource_location == RESOURCE_IN_FILE_ON_DISK) {
+ ERROR("We were reading it from `%s'; maybe it changed "
+ "while we were reading it.",
+ lte->file_on_disk);
+ }
+ ret = WIMLIB_ERR_INVALID_RESOURCE_HASH;
+ goto out_fclose;
+ }
+ }
+
+ if (!raw && new_compressed_size >= original_size &&
+ out_ctype != WIM_COMPRESSION_TYPE_NONE)
+ {
+ /* Oops! We compressed the resource to larger than the original
+ * size. Write the resource uncompressed instead. */
+ if (fseeko(out_fp, file_offset, SEEK_SET) != 0) {
+ ERROR_WITH_ERRNO("Failed to seek to byte %"PRIu64" "
+ "of output WIM file", file_offset);
+ ret = WIMLIB_ERR_WRITE;
+ goto out_fclose;
+ }
+ ret = write_wim_resource(lte, out_fp, WIM_COMPRESSION_TYPE_NONE,
+ out_res_entry, flags);
+ if (ret != 0)
+ goto out_fclose;
+ if (fflush(out_fp) != 0) {
+ ERROR_WITH_ERRNO("Failed to flush output WIM file");
+ ret = WIMLIB_ERR_WRITE;
+ goto out_fclose;
+ }
+ if (ftruncate(fileno(out_fp), file_offset + out_res_entry->size) != 0) {
+ ERROR_WITH_ERRNO("Failed to truncate output WIM file");
+ ret = WIMLIB_ERR_WRITE;
+ goto out_fclose;
+ }
+ } else {
+ if (out_res_entry) {
+ out_res_entry->size = new_compressed_size;
+ out_res_entry->original_size = original_size;
+ out_res_entry->offset = file_offset;
+ out_res_entry->flags = lte->resource_entry.flags
+ & ~WIM_RESHDR_FLAG_COMPRESSED;
+ if (out_ctype != WIM_COMPRESSION_TYPE_NONE)
+ out_res_entry->flags |= WIM_RESHDR_FLAG_COMPRESSED;
+ }
+ }
+ ret = 0;
+out_fclose:
+#ifdef WITH_NTFS_3G
+ end_wim_resource_read(lte, ni);
+#else
+ end_wim_resource_read(lte);
+#endif
+out:
+ FREE(chunk_tab);
+ return ret;
+}
+
+
+struct shared_queue {
+ sem_t filled_slots;
+ sem_t empty_slots;
+ pthread_mutex_t lock;
+ unsigned front;
+ unsigned back;
+ void **array;
+ unsigned size;
+};
+
+static int shared_queue_init(struct shared_queue *q, unsigned size)
+{
+ q->array = CALLOC(sizeof(q->array[0]), size);
+ if (!q->array)
+ return WIMLIB_ERR_NOMEM;
+
+ sem_init(&q->filled_slots, 0, 0);
+ sem_init(&q->empty_slots, 0, size);
+ pthread_mutex_init(&q->lock, NULL);
+ q->front = 0;
+ q->back = size - 1;
+ q->size = size;
+ return 0;
+}
+
+static void shared_queue_destroy(struct shared_queue *q)
+{
+ sem_destroy(&q->filled_slots);
+ sem_destroy(&q->empty_slots);
+ pthread_mutex_destroy(&q->lock);
+ FREE(q->array);
+}
+
+static void shared_queue_put(struct shared_queue *q, void *obj)
+{
+ sem_wait(&q->empty_slots);
+ pthread_mutex_lock(&q->lock);
+
+ q->back = (q->back + 1) % q->size;
+ q->array[q->back] = obj;
+
+ sem_post(&q->filled_slots);
+ pthread_mutex_unlock(&q->lock);
+}
+
+static void *shared_queue_get(struct shared_queue *q)
+{
+ sem_wait(&q->filled_slots);
+ pthread_mutex_lock(&q->lock);
+
+ void *obj = q->array[q->front];
+ q->array[q->front] = NULL;
+ q->front = (q->front + 1) % q->size;
+
+ sem_post(&q->empty_slots);
+ pthread_mutex_unlock(&q->lock);
+ return obj;
+}
+
+static inline int shared_queue_get_filled(struct shared_queue *q)
+{
+ int sval;
+ sem_getvalue(&q->filled_slots, &sval);
+ return sval;
+}
+
+struct compressor_thread_params {
+ struct shared_queue *res_to_compress_queue;
+ struct shared_queue *compressed_res_queue;
+ compress_func_t compress;
+};
+
+#define MAX_CHUNKS_PER_MSG 2
+
+struct message {
+ struct lookup_table_entry *lte;
+ u8 *uncompressed_chunks[MAX_CHUNKS_PER_MSG];
+ u8 *out_compressed_chunks[MAX_CHUNKS_PER_MSG];
+ u8 *compressed_chunks[MAX_CHUNKS_PER_MSG];
+ unsigned uncompressed_chunk_sizes[MAX_CHUNKS_PER_MSG];
+ unsigned compressed_chunk_sizes[MAX_CHUNKS_PER_MSG];
+ unsigned num_chunks;
+ struct list_head list;
+ bool complete;
+ u64 begin_chunk;
+};
+
+static void compress_chunks(struct message *msg, compress_func_t compress)
+{
+ for (unsigned i = 0; i < msg->num_chunks; i++) {
+ DEBUG2("compress chunk %u of %u", i, msg->num_chunks);
+ int ret = compress(msg->uncompressed_chunks[i],
+ msg->uncompressed_chunk_sizes[i],
+ msg->compressed_chunks[i],
+ &msg->compressed_chunk_sizes[i]);
+ if (ret == 0) {
+ msg->out_compressed_chunks[i] = msg->compressed_chunks[i];
+ } else {
+ msg->out_compressed_chunks[i] = msg->uncompressed_chunks[i];
+ msg->compressed_chunk_sizes[i] = msg->uncompressed_chunk_sizes[i];
+ }
+ }
+}
+
+static void *compressor_thread_proc(void *arg)
+{
+ struct compressor_thread_params *params = arg;
+ struct shared_queue *res_to_compress_queue = params->res_to_compress_queue;
+ struct shared_queue *compressed_res_queue = params->compressed_res_queue;
+ compress_func_t compress = params->compress;
+ struct message *msg;
+
+ DEBUG("Compressor thread ready");
+ while ((msg = shared_queue_get(res_to_compress_queue)) != NULL) {
+ compress_chunks(msg, compress);
+ shared_queue_put(compressed_res_queue, msg);
+ }
+ DEBUG("Compressor thread terminating");
+}
+
+static int write_stream_list_serial(struct list_head *stream_list,
+ FILE *out_fp, int out_ctype,
+ int write_flags)
+{
+ struct lookup_table_entry *lte;
+ int ret;
+
+ list_for_each_entry(lte, stream_list, staging_list) {
+ ret = write_wim_resource(lte, out_fp, out_ctype,
+ <e->output_resource_entry, 0);
+ if (ret != 0)
+ return ret;
+ }
+ return 0;
+}
+
+static int write_wim_chunks(struct message *msg, FILE *out_fp,
+ struct chunk_table *chunk_tab)
+{
+ for (unsigned i = 0; i < msg->num_chunks; i++) {
+ unsigned chunk_csize = msg->compressed_chunk_sizes[i];
+
+ DEBUG2("Write wim chunk %u of %u (csize = %u)",
+ i, msg->num_chunks, chunk_csize);
+
+ if (fwrite(msg->out_compressed_chunks[i], 1, chunk_csize, out_fp)
+ != chunk_csize)
+ {
+ ERROR_WITH_ERRNO("Failed to write WIM");
+ return WIMLIB_ERR_WRITE;
+ }
+
+ *chunk_tab->cur_offset_p++ = chunk_tab->cur_offset;
+ chunk_tab->cur_offset += chunk_csize;
+ }
+ return 0;
+}
+
+static int main_writer_thread_proc(struct list_head *stream_list,
+ FILE *out_fp,
+ int out_ctype,
+ struct shared_queue *res_to_compress_queue,
+ struct shared_queue *compressed_res_queue,
+ size_t queue_size)
+{
+ int ret;
+
+
+ struct message msgs[queue_size];
+ ZERO_ARRAY(msgs);
+
+ LIST_HEAD(available_msgs);
+
+ for (size_t i = 0; i < ARRAY_LEN(msgs); i++)
+ list_add(&msgs[i].list, &available_msgs);
+
+
+ struct list_head *next_resource = stream_list->next;
+ struct lookup_table_entry *next_lte = container_of(next_resource,
+ struct lookup_table_entry,
+ staging_list);
+ next_resource = next_resource->next;
+ u64 next_chunk = 0;
+ u64 next_num_chunks = wim_resource_chunks(next_lte);
+ INIT_LIST_HEAD(&next_lte->msg_list);
+ SHA_CTX next_sha_ctx;
+ sha1_init(&next_sha_ctx);
+
+ u8 next_hash[SHA1_HASH_SIZE];
+
+ // Resources owning chunks that have been sent off for compression
+ LIST_HEAD(outstanding_resources);
+
+ // Resources that are going to be written by the main thread
+ LIST_HEAD(my_resources);
+
+ list_add_tail(&next_lte->staging_list, &outstanding_resources);
+
+ struct lookup_table_entry *cur_lte = next_lte;
+ struct chunk_table *cur_chunk_tab = NULL;
+ struct lookup_table_entry *lte;
+ struct message *msg;
+
+#ifdef WITH_NTFS_3G
+ ntfs_inode *ni = NULL;
+#endif
+
+#ifdef WITH_NTFS_3G
+ ret = prepare_resource_for_read(next_lte, &ni);
+#else
+ ret = prepare_resource_for_read(next_lte);
+#endif
+
+ DEBUG("Initializing buffers for uncompressed "
+ "and compressed data (%zu bytes needed)",
+ queue_size * MAX_CHUNKS_PER_MSG * WIM_CHUNK_SIZE * 2);
+
+ for (size_t i = 0; i < ARRAY_LEN(msgs); i++) {
+ for (size_t j = 0; j < MAX_CHUNKS_PER_MSG; j++) {
+ msgs[i].compressed_chunks[j] = MALLOC(WIM_CHUNK_SIZE);
+ msgs[i].uncompressed_chunks[j] = MALLOC(WIM_CHUNK_SIZE);
+ if (msgs[i].compressed_chunks[j] == NULL ||
+ msgs[i].uncompressed_chunks[j] == NULL)
+ {
+ ERROR("Could not allocate enough memory for "
+ "multi-threaded compression");
+ ret = WIMLIB_ERR_NOMEM;
+ goto out;
+ }
+ }
+ }
+
+ while (1) {
+ // Send chunks to the compressor threads until either (a) there
+ // are no more messages available since they were all sent off,
+ // or (b) there are no more resources that need to be
+ // compressed.
+ while (!list_empty(&available_msgs) && next_lte != NULL) {
+
+ wimlib_assert(next_chunk < next_num_chunks);
+
+ // Get a message from the available messages
+ // list
+ msg = container_of(available_msgs.next,
+ struct message,
+ list);
+
+ // ... and delete it from the available messages
+ // list
+ list_del(&msg->list);
+
+ // Initialize the message with the chunks to
+ // compress.
+ msg->num_chunks = min(next_num_chunks - next_chunk,
+ MAX_CHUNKS_PER_MSG);
+ msg->lte = next_lte;
+ msg->complete = false;
+ msg->begin_chunk = next_chunk;
+
+ unsigned size = WIM_CHUNK_SIZE;
+ for (unsigned i = 0; i < msg->num_chunks; i++) {
+ if (next_chunk == next_num_chunks - 1 &&
+ wim_resource_size(next_lte) % WIM_CHUNK_SIZE != 0)
+ {
+ size = wim_resource_size(next_lte) % WIM_CHUNK_SIZE;
+ }
+
+ DEBUG2("Read resource (size=%u, offset=%zu)",
+ size, next_chunk * WIM_CHUNK_SIZE);
+
+ msg->uncompressed_chunk_sizes[i] = size;
+
+ ret = read_wim_resource(next_lte,
+ msg->uncompressed_chunks[i],
+ size,
+ next_chunk * WIM_CHUNK_SIZE,
+ 0);
+ if (ret != 0)
+ goto out;
+ sha1_update(&next_sha_ctx,
+ msg->uncompressed_chunks[i], size);
+ next_chunk++;
+ }
+
+ // Send the compression request
+ list_add_tail(&msg->list, &next_lte->msg_list);
+ shared_queue_put(res_to_compress_queue, msg);
+ DEBUG2("Compression request sent");
+
+ if (next_chunk != next_num_chunks)
+ // More chunks to send for this resource
+ continue;
+
+ // Done sending compression requests for a resource!
+ // Check the SHA1 message digest.
+ DEBUG2("Finalize SHA1 md (next_num_chunks=%zu)", next_num_chunks);
+ sha1_final(next_hash, &next_sha_ctx);
+ if (!hashes_equal(next_lte->hash, next_hash)) {
+ ERROR("WIM resource has incorrect hash!");
+ if (next_lte->resource_location == RESOURCE_IN_FILE_ON_DISK) {
+ ERROR("We were reading it from `%s'; maybe it changed "
+ "while we were reading it.",
+ next_lte->file_on_disk);
+ }
+ ret = WIMLIB_ERR_INVALID_RESOURCE_HASH;
+ goto out;
+ }
+
+ // Advance to the next resource.
+ //
+ // If the next resource needs no compression,
+ // just write it with this thread (not now
+ // though--- we could be in the middle of
+ // writing another resource.) Keep doing this
+ // until we either get to the end of the
+ // resources list, or we get to a resource that
+ // needs compression.
+
+ while (1) {
+ if (next_resource == stream_list) {
+ next_lte = NULL;
+ break;
+ }
+ #ifdef WITH_NTFS_3G
+ end_wim_resource_read(next_lte, ni);
+ ni = NULL;
+ #else
+ end_wim_resource_read(next_lte);
+ #endif
+
+ next_lte = container_of(next_resource,
+ struct lookup_table_entry,
+ staging_list);
+ next_resource = next_resource->next;
+ if ((next_lte->resource_location == RESOURCE_IN_WIM
+ && wimlib_get_compression_type(next_lte->wim) == out_ctype)
+ || wim_resource_size(next_lte) == 0)
+ {
+ list_add_tail(&next_lte->staging_list,
+ &my_resources);
+ } else {
+ list_add_tail(&next_lte->staging_list,
+ &outstanding_resources);
+ next_chunk = 0;
+ next_num_chunks = wim_resource_chunks(next_lte);
+ sha1_init(&next_sha_ctx);
+ INIT_LIST_HEAD(&next_lte->msg_list);
+ #ifdef WITH_NTFS_3G
+ ret = prepare_resource_for_read(next_lte, &ni);
+ #else
+ ret = prepare_resource_for_read(next_lte);
+ #endif
+ if (ret != 0)
+ goto out;
+ DEBUG2("Updated next_lte");
+ break;
+ }
+ }
+ }
+
+ // If there are no outstanding resources, there are no more
+ // resources to compress.
+ if (list_empty(&outstanding_resources)) {
+ DEBUG("No outstanding resources! Done");
+ ret = 0;
+ goto out;
+ }
+
+ // Get the next message from the queue and process it.
+ // The message will contain 1 or more data chunks that have been
+ // compressed.
+ DEBUG2("Waiting for message");
+ msg = shared_queue_get(compressed_res_queue);
+ msg->complete = true;
+
+ DEBUG2("Received msg (begin_chunk=%"PRIu64")", msg->begin_chunk);
+
+ list_for_each_entry(msg, &cur_lte->msg_list, list) {
+ DEBUG2("complete=%d", msg->complete);
+ }
+
+ // Is this the next chunk in the current resource? If it's not
+ // (i.e., an earlier chunk in a same or different resource
+ // hasn't been compressed yet), do nothing, and keep this
+ // message around until all earlier chunks are received.
+ //
+ // Otherwise, write all the chunks we can.
+ while (!list_empty(&cur_lte->msg_list)
+ && (msg = container_of(cur_lte->msg_list.next,
+ struct message,
+ list))->complete)
+ {
+ DEBUG2("Complete msg (begin_chunk=%"PRIu64")", msg->begin_chunk);
+ if (msg->begin_chunk == 0) {
+ DEBUG2("Begin chunk tab");
+ // This is the first set of chunks.
+ // Leave space for the chunk table in
+ // the output file.
+ off_t cur_offset = ftello(out_fp);
+ if (cur_offset == -1) {
+ ret = WIMLIB_ERR_WRITE;
+ goto out;
+ }
+ FREE(cur_chunk_tab);
+ cur_chunk_tab = NULL;
+ ret = begin_wim_resource_chunk_tab(cur_lte,
+ out_fp,
+ cur_offset,
+ &cur_chunk_tab);
+ if (ret != 0)
+ goto out;
+ }
+ ret = write_wim_chunks(msg, out_fp, cur_chunk_tab);
+ if (ret != 0)
+ goto out;
+
+ list_del(&msg->list);
+ list_add(&msg->list, &available_msgs);
+
+ if (list_empty(&cur_lte->msg_list) &&
+ msg->begin_chunk + msg->num_chunks == cur_chunk_tab->num_chunks)
+ {
+ DEBUG2("Finish wim chunk tab");
+ u64 res_csize;
+ ret = finish_wim_resource_chunk_tab(cur_chunk_tab,
+ out_fp,
+ &res_csize);
+ if (ret != 0)
+ return ret;
+
+ cur_lte->output_resource_entry.size =
+ res_csize;
+
+ cur_lte->output_resource_entry.original_size =
+ cur_lte->resource_entry.original_size;
+
+ cur_lte->output_resource_entry.offset =
+ cur_chunk_tab->file_offset;
+
+ cur_lte->output_resource_entry.flags =
+ cur_lte->resource_entry.flags |
+ WIM_RESHDR_FLAG_COMPRESSED;
+
+ struct list_head *next = cur_lte->staging_list.next;
+ list_del(&cur_lte->staging_list);
+
+ if (next == &outstanding_resources) {
+ DEBUG("No more outstanding resources");
+ ret = 0;
+ goto out;
+ } else {
+ cur_lte = container_of(cur_lte->staging_list.next,
+ struct lookup_table_entry,
+ staging_list);
+ }
+
+ struct lookup_table_entry *tmp;
+ list_for_each_entry_safe(lte,
+ tmp,
+ &my_resources,
+ staging_list)
+ {
+ ret = write_wim_resource(lte,
+ out_fp,
+ out_ctype,
+ <e->output_resource_entry,
+ 0);
+ list_del(<e->staging_list);
+ if (ret != 0)
+ goto out;
+ }
+ }
+ }
+ }
+
+out:
+#ifdef WITH_NTFS_3G
+ end_wim_resource_read(cur_lte, ni);
+#else
+ end_wim_resource_read(cur_lte);
+#endif
+ if (ret == 0) {
+ list_for_each_entry(lte, &my_resources, staging_list) {
+ ret = write_wim_resource(lte, out_fp,
+ out_ctype,
+ <e->output_resource_entry,
+ 0);
+ if (ret != 0)
+ break;
+ }
+ } else {
+ size_t num_available_msgs = 0;
+ struct list_head *cur;
+
+ list_for_each(cur, &available_msgs) {
+ num_available_msgs++;
+ }
+
+ while (num_available_msgs < ARRAY_LEN(msgs)) {
+ msg = shared_queue_get(compressed_res_queue);
+ num_available_msgs++;
+ }
+ }
+
+ DEBUG("Freeing messages");
+
+ for (size_t i = 0; i < ARRAY_LEN(msgs); i++) {
+ for (size_t j = 0; j < MAX_CHUNKS_PER_MSG; j++) {
+ FREE(msgs[i].compressed_chunks[j]);
+ FREE(msgs[i].uncompressed_chunks[j]);
+ }
+ }
+
+ if (cur_chunk_tab != NULL)
+ FREE(cur_chunk_tab);
+ return ret;
+}
+
+static int write_stream_list_parallel(struct list_head *stream_list,
+ FILE *out_fp, int out_ctype,
+ int write_flags)
+{
+ int ret;
+ long nthreads;
+ struct shared_queue res_to_compress_queue;
+ struct shared_queue compressed_res_queue;
+
+ nthreads = sysconf(_SC_NPROCESSORS_ONLN);
+ if (nthreads < 1) {
+ WARNING("Could not determine number of processors! Assuming 1");
+ goto out_serial;
+ }
+
+ wimlib_assert(stream_list->next != stream_list);
+
+ {
+ pthread_t compressor_threads[nthreads];
+
+ static const double MESSAGES_PER_THREAD = 2.0;
+ size_t queue_size = (size_t)(nthreads * MESSAGES_PER_THREAD);
+
+ DEBUG("Initializing shared queues (queue_size=%zu)", queue_size);
+
+ ret = shared_queue_init(&res_to_compress_queue, queue_size);
+ if (ret != 0)
+ goto out_serial;
+
+ ret = shared_queue_init(&compressed_res_queue, queue_size);
+ if (ret != 0)
+ goto out_destroy_res_to_compress_queue;
+
+
+ struct compressor_thread_params params;
+ params.res_to_compress_queue = &res_to_compress_queue;
+ params.compressed_res_queue = &compressed_res_queue;
+ params.compress = get_compress_func(out_ctype);
+
+ for (long i = 0; i < nthreads; i++) {
+ DEBUG("pthread_create thread %ld", i);
+ ret = pthread_create(&compressor_threads[i], NULL,
+ compressor_thread_proc, ¶ms);
+ if (ret != 0) {
+ ERROR_WITH_ERRNO("Failed to create compressor "
+ "thread %ld", i);
+ nthreads = i;
+ goto out_join;
+ }
+ }
+
+ if (write_flags & WIMLIB_WRITE_FLAG_VERBOSE) {
+ printf("Writing compressed data using %ld threads...\n",
+ nthreads);
+ }
+
+ ret = main_writer_thread_proc(stream_list,
+ out_fp,
+ out_ctype,
+ &res_to_compress_queue,
+ &compressed_res_queue,
+ queue_size);
+
+ out_join:
+ for (long i = 0; i < nthreads; i++)
+ shared_queue_put(&res_to_compress_queue, NULL);
+
+ for (long i = 0; i < nthreads; i++) {
+ if (pthread_join(compressor_threads[i], NULL)) {
+ WARNING("Failed to join compressor thread %ld: %s",
+ i, strerror(errno));
+ }
+ }
+ }
+ shared_queue_destroy(&compressed_res_queue);
+out_destroy_res_to_compress_queue:
+ shared_queue_destroy(&res_to_compress_queue);
+ if (ret >= 0 && ret != WIMLIB_ERR_NOMEM)
+ return ret;
+out_serial:
+ WARNING("Falling back to single-threaded compression");
+ return write_stream_list_serial(stream_list, out_fp,
+ out_ctype, write_flags);
+}
+
+static int write_stream_list(struct list_head *stream_list, FILE *out_fp,
+ int out_ctype, int write_flags)
+{
+ struct lookup_table_entry *lte;
+ size_t num_streams = 0;
+ u64 total_size = 0;
+ bool compression_needed = false;
+
+ list_for_each_entry(lte, stream_list, staging_list) {
+ num_streams++;
+ total_size += wim_resource_size(lte);
+ if (!compression_needed
+ && out_ctype != WIM_COMPRESSION_TYPE_NONE
+ && (lte->resource_location != RESOURCE_IN_WIM
+ || wimlib_get_compression_type(lte->wim) != out_ctype)
+ && wim_resource_size(lte) != 0)
+ compression_needed = true;
+ }
+
+ if (write_flags & WIMLIB_WRITE_FLAG_VERBOSE) {
+ printf("Preparing to write %zu streams "
+ "(%"PRIu64" total bytes uncompressed)\n",
+ num_streams, total_size);
+ printf("Using compression type %s\n",
+ wimlib_get_compression_type_string(out_ctype));
+ }
+
+ if (compression_needed && total_size >= 0) { // XXX
+ return write_stream_list_parallel(stream_list, out_fp,
+ out_ctype, write_flags);
+ } else {
+ if (write_flags & WIMLIB_WRITE_FLAG_VERBOSE) {
+ puts("Using 1 thread (no compression needed)");
+ }
+
+ return write_stream_list_serial(stream_list, out_fp,
+ out_ctype, write_flags);
+ }
+}
+
+
+static int dentry_find_streams_to_write(struct dentry *dentry,
+ void *wim)
+{
+ WIMStruct *w = wim;
+ struct list_head *stream_list = w->private;
+ struct lookup_table_entry *lte;
+ for (unsigned i = 0; i <= dentry->d_inode->num_ads; i++) {
+ lte = inode_stream_lte(dentry->d_inode, i, w->lookup_table);
+ if (lte && ++lte->out_refcnt == 1)
+ list_add(<e->staging_list, stream_list);
+ }
+ return 0;
+}
+
+static int find_streams_to_write(WIMStruct *w)
+{
+ return for_dentry_in_tree(wim_root_dentry(w),
+ dentry_find_streams_to_write, w);
+}
+
+static int write_wim_streams(WIMStruct *w, int image, int write_flags)
+{
+
+ LIST_HEAD(stream_list);
+
+ w->private = &stream_list;
+ for_image(w, image, find_streams_to_write);
+ return write_stream_list(&stream_list, w->out_fp,
+ wimlib_get_compression_type(w), write_flags);
}
/*
if (fseeko(out, 0, SEEK_SET) != 0)
return WIMLIB_ERR_WRITE;
- return write_header(&hdr, out);
+ ret = write_header(&hdr, out);
+ if (ret != 0)
+ return ret;
+
+ DEBUG("Closing output file.");
+ wimlib_assert(w->out_fp != NULL);
+ if (fclose(w->out_fp) != 0) {
+ ERROR_WITH_ERRNO("Failed to close the WIM file");
+ ret = WIMLIB_ERR_WRITE;
+ }
+ w->out_fp = NULL;
+ return ret;
}
/* Open file stream and write dummy header for WIM. */
else
mode = "wb";
+ if (w->out_fp)
+ fclose(w->out_fp);
+
w->out_fp = fopen(path, mode);
if (!w->out_fp) {
ERROR_WITH_ERRNO("Failed to open the file `%s' for writing",
ret = begin_write(w, path, write_flags);
if (ret != 0)
- goto out;
+ return ret;
for_lookup_table_entry(w->lookup_table, lte_zero_out_refcnt, NULL);
- w->write_flags = write_flags;
+ ret = write_wim_streams(w, image, write_flags);
- ret = for_image(w, image, write_file_resources);
if (ret != 0) {
- ERROR("Failed to write WIM file resources to `%s'", path);
- goto out;
+ /*ERROR("Failed to write WIM file resources to `%s'", path);*/
+ return ret;
}
ret = for_image(w, image, write_metadata_resource);
if (ret != 0) {
/*ERROR("Failed to write WIM image metadata to `%s'", path);*/
- goto out;
+ return ret;
}
-
ret = finish_write(w, image, write_flags);
+ if (ret != 0)
+ return ret;
-out:
- DEBUG("Closing output file.");
- if (w->out_fp != NULL) {
- if (fclose(w->out_fp) != 0) {
- ERROR_WITH_ERRNO("Failed to close the file `%s'", path);
- ret = WIMLIB_ERR_WRITE;
- }
- w->out_fp = NULL;
- }
- return ret;
+ if (write_flags & WIMLIB_WRITE_FLAG_VERBOSE)
+ printf("Successfully wrote `%s'\n", path);
+ return 0;
}
git://wimlib.net / wimlib / blobdiff