Add a way to keep meta data blocks in memory

Instead of writing meta data blocks directly to disk, the writer can
now alternatively keep the blocks in memory until explicitly told to
write to disk.

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

5 files changed, 108 insertions, 34 deletions

diff --git a/include/meta_writer.h b/include/meta_writer.h
index df0c70c..c10dd16 100644
--- a/include/meta_writer.h
+++ b/include/meta_writer.h
@@ -22,8 +22,11 @@ typedef struct {
 typedef struct meta_writer_t meta_writer_t;
 
 /* Create a meta data reader using a given compressor to compress data.
-   Internally prints error message to stderr on failure. */
-meta_writer_t *meta_writer_create(int fd, compressor_t *cmp);
+   Internally prints error message to stderr on failure.
+   If keep_in_mem is true, the blocks are collected in memory and must
+   be explicitly flushed to disk using meta_write_write_to_file.
+*/
+meta_writer_t *meta_writer_create(int fd, compressor_t *cmp, bool keep_in_mem);
 
 void meta_writer_destroy(meta_writer_t *m);
 
@@ -41,6 +44,11 @@ void meta_writer_get_position(const meta_writer_t *m, uint64_t *block_start,
 /* Reset all internal state, including the current block start position. */
 void meta_writer_reset(meta_writer_t *m);
 
+/* If created with keep_in_mem true, write the collected blocks to disk.
+   Does not flush the current block. Writes error messages to stderr and
+   returns non-zero on failure. */
+int meta_write_write_to_file(meta_writer_t *m);
+
 /*
   High level helper function that writes squashfs directory entries to
   a meta data writer.
diff --git a/lib/sqfs/meta_writer.c b/lib/sqfs/meta_writer.c
index 2dacaa0..ef869e7 100644
--- a/lib/sqfs/meta_writer.c
+++ b/lib/sqfs/meta_writer.c
@@ -8,6 +8,13 @@
 #include <unistd.h>
 #include <stdio.h>
 
+typedef struct meta_block_t {
+	struct meta_block_t *next;
+
+	/* possibly compressed data with 2 byte header */
+	uint8_t data[SQFS_META_BLOCK_SIZE + 2];
+} meta_block_t;
+
 struct meta_writer_t {
 	/* A byte offset into the uncompressed data of the current block */
 	size_t offset;
@@ -22,13 +29,36 @@ struct meta_writer_t {
 	compressor_t *cmp;
 
 	/* The raw data chunk that data is appended to */
-	uint8_t data[SQFS_META_BLOCK_SIZE + 2];
+	uint8_t data[SQFS_META_BLOCK_SIZE];
 
-	/* Scratch buffer for compressing data */
-	uint8_t scratch[SQFS_META_BLOCK_SIZE + 2];
+	bool keep_in_mem;
+	meta_block_t *list;
+	meta_block_t *list_end;
 };
 
-meta_writer_t *meta_writer_create(int fd, compressor_t *cmp)
+static int write_block(int fd, meta_block_t *outblk)
+{
+	size_t count;
+	ssize_t ret;
+
+	count = le16toh(((uint16_t *)outblk->data)[0]) & 0x7FFF;
+
+	ret = write_retry(fd, outblk->data, count + 2);
+
+	if (ret < 0) {
+		perror("writing meta data block");
+		return -1;
+	}
+
+	if ((size_t)ret < count) {
+		fputs("meta data written to file was truncated\n", stderr);
+		return -1;
+	}
+
+	return 0;
+}
+
+meta_writer_t *meta_writer_create(int fd, compressor_t *cmp, bool keep_in_mem)
 {
 	meta_writer_t *m = calloc(1, sizeof(*m));
 
@@ -39,53 +69,71 @@ meta_writer_t *meta_writer_create(int fd, compressor_t *cmp)
 
 	m->cmp = cmp;
 	m->outfd = fd;
+	m->keep_in_mem = keep_in_mem;
 	return m;
 }
 
 void meta_writer_destroy(meta_writer_t *m)
 {
+	meta_block_t *blk;
+
+	while (m->list != NULL) {
+		blk = m->list;
+		m->list = blk->next;
+		free(blk);
+	}
+
 	free(m);
 }
 
 int meta_writer_flush(meta_writer_t *m)
 {
-	ssize_t ret, count;
-	void *ptr;
+	meta_block_t *outblk;
+	size_t count;
+	ssize_t ret;
 
 	if (m->offset == 0)
 		return 0;
 
-	ret = m->cmp->do_block(m->cmp, m->data + 2, m->offset,
-			       m->scratch + 2, sizeof(m->scratch) - 2);
-	if (ret < 0)
+	outblk = calloc(1, sizeof(*outblk));
+	if (outblk == NULL) {
+		perror("generating meta data block");
+		return -1;
+	}
+
+	ret = m->cmp->do_block(m->cmp, m->data, m->offset,
+			       outblk->data + 2, sizeof(outblk->data) - 2);
+	if (ret < 0) {
+		free(outblk);
 		return -1;
+	}
 
 	if (ret > 0) {
-		((uint16_t *)m->scratch)[0] = htole16(ret);
+		((uint16_t *)outblk->data)[0] = htole16(ret);
 		count = ret + 2;
-		ptr = m->scratch;
 	} else {
-		((uint16_t *)m->data)[0] = htole16(m->offset | 0x8000);
+		((uint16_t *)outblk->data)[0] = htole16(m->offset | 0x8000);
+		memcpy(outblk->data + 2, m->data, m->offset);
 		count = m->offset + 2;
-		ptr = m->data;
-	}
-
-	ret = write_retry(m->outfd, ptr, count);
-
-	if (ret < 0) {
-		perror("writing meta data");
-		return -1;
 	}
 
-	if (ret < count) {
-		fputs("meta data written to file was truncated\n", stderr);
-		return -1;
+	if (m->keep_in_mem) {
+		if (m->list == NULL) {
+			m->list = outblk;
+		} else {
+			m->list_end->next = outblk;
+		}
+		m->list_end = outblk;
+		ret = 0;
+	} else {
+		ret = write_block(m->outfd, outblk);
+		free(outblk);
 	}
 
 	memset(m->data, 0, sizeof(m->data));
 	m->offset = 0;
 	m->block_offset += count;
-	return 0;
+	return ret;
 }
 
 int meta_writer_append(meta_writer_t *m, const void *data, size_t size)
@@ -93,24 +141,24 @@ int meta_writer_append(meta_writer_t *m, const void *data, size_t size)
 	size_t diff;
 
 	while (size != 0) {
-		diff = sizeof(m->data) - 2 - m->offset;
+		diff = sizeof(m->data) - m->offset;
 
 		if (diff == 0) {
 			if (meta_writer_flush(m))
 				return -1;
-			diff = sizeof(m->data) - 2;
+			diff = sizeof(m->data);
 		}
 
 		if (diff > size)
 			diff = size;
 
-		memcpy(m->data + 2 + m->offset, data, diff);
+		memcpy(m->data + m->offset, data, diff);
 		m->offset += diff;
 		size -= diff;
 		data = (const char *)data + diff;
 	}
 
-	if (m->offset == (sizeof(m->data) - 2))
+	if (m->offset == sizeof(m->data))
 		return meta_writer_flush(m);
 
 	return 0;
@@ -128,3 +176,21 @@ void meta_writer_reset(meta_writer_t *m)
 	m->block_offset = 0;
 	m->offset = 0;
 }
+
+int meta_write_write_to_file(meta_writer_t *m)
+{
+	meta_block_t *blk;
+
+	while (m->list != NULL) {
+		blk = m->list;
+
+		if (write_block(m->outfd, blk))
+			return -1;
+
+		m->list = blk->next;
+		free(blk);
+	}
+
+	m->list_end = NULL;
+	return 0;
+}
diff --git a/lib/sqfs/serialize_fstree.c b/lib/sqfs/serialize_fstree.c
index 44aeca4..f3e7bba 100644
--- a/lib/sqfs/serialize_fstree.c
+++ b/lib/sqfs/serialize_fstree.c
@@ -26,11 +26,11 @@ int sqfs_serialize_fstree(int outfd, sqfs_super_t *super, fstree_t *fs,
 
 	tmpfd = fileno(tmp);
 
-	im = meta_writer_create(outfd, cmp);
+	im = meta_writer_create(outfd, cmp, false);
 	if (im == NULL)
 		goto fail_tmp;
 
-	dm = meta_writer_create(tmpfd, cmp);
+	dm = meta_writer_create(tmpfd, cmp, false);
 	if (dm == NULL)
 		goto fail_im;
 
diff --git a/lib/sqfs/table.c b/lib/sqfs/table.c
index abc95bd..fb6bddc 100644
--- a/lib/sqfs/table.c
+++ b/lib/sqfs/table.c
@@ -19,7 +19,7 @@ int sqfs_write_table(int outfd, sqfs_super_t *super, const void *data,
 
 	/* Write actual data. Whenever we cross a block boundary, remember
 	   the block start offset */
-	m = meta_writer_create(outfd, cmp);
+	m = meta_writer_create(outfd, cmp, false);
 	if (m == NULL)
 		return -1;
 
diff --git a/lib/sqfs/write_xattr.c b/lib/sqfs/write_xattr.c
index e81f677..0e11757 100644
--- a/lib/sqfs/write_xattr.c
+++ b/lib/sqfs/write_xattr.c
@@ -110,7 +110,7 @@ int write_xattr(int outfd, fstree_t *fs, sqfs_super_t *super,
 	if (fs->xattr == NULL)
 		return 0;
 
-	mw = meta_writer_create(outfd, cmp);
+	mw = meta_writer_create(outfd, cmp, false);
 	if (mw == NULL)
 		return -1;