btrfs: zstd: introduce zstd_compress_bio() helper

The new helper has the following enhancements against the existing
zstd_compress_folios()

- Much smaller parameter list

  No more shared IN/OUT members, no need to pre-allocate a
  compressed_folios[] array.

  Just a workspace and compressed_bio pointer, everything we need can be
  extracted from that @cb pointer.

- Ready-to-be-submitted compressed bio

  Although the caller still needs to do some common works like
  rounding up and zeroing the tailing part of the last fs block.

Overall the workflow is the same as zstd_compress_folios(), but with
some minor changes:

- @start/@len is now constant
  For the current input file offset, use @start + @tot_in instead.

  The original change of @start and @len makes it pretty hard to know
  what value we're really comparing to.

- No more @cur_len
  It's only utilized when switching input buffer.
  Directly use btrfs_calc_input_length() instead.

Reviewed-by: Boris Burkov <boris@bur.io>
Signed-off-by: Qu Wenruo <wqu@suse.com>
Reviewed-by: David Sterba <dsterba@suse.com>
Signed-off-by: David Sterba <dsterba@suse.com>

2 files changed, 187 insertions, 0 deletions

diff --git a/fs/btrfs/compression.h b/fs/btrfs/compression.h
index 4b63d7e4a9ad..454c8e0461b4 100644
--- a/fs/btrfs/compression.h
+++ b/fs/btrfs/compression.h
@@ -172,6 +172,7 @@ void lzo_free_workspace(struct list_head *ws);
 int zstd_compress_folios(struct list_head *ws, struct btrfs_inode *inode,
 			 u64 start, struct folio **folios, unsigned long *out_folios,
 		unsigned long *total_in, unsigned long *total_out);
+int zstd_compress_bio(struct list_head *ws, struct compressed_bio *cb);
 int zstd_decompress_bio(struct list_head *ws, struct compressed_bio *cb);
 int zstd_decompress(struct list_head *ws, const u8 *data_in,
 		struct folio *dest_folio, unsigned long dest_pgoff, size_t srclen,
diff --git a/fs/btrfs/zstd.c b/fs/btrfs/zstd.c
index 7fad1e299c7a..135b0b32579f 100644
--- a/fs/btrfs/zstd.c
+++ b/fs/btrfs/zstd.c
@@ -585,6 +585,192 @@ out:
 	return ret;
 }
 
+int zstd_compress_bio(struct list_head *ws, struct compressed_bio *cb)
+{
+	struct btrfs_inode *inode = cb->bbio.inode;
+	struct btrfs_fs_info *fs_info = inode->root->fs_info;
+	struct workspace *workspace = list_entry(ws, struct workspace, list);
+	struct address_space *mapping = inode->vfs_inode.i_mapping;
+	struct bio *bio = &cb->bbio.bio;
+	zstd_cstream *stream;
+	int ret = 0;
+	/* The current folio to read. */
+	struct folio *in_folio = NULL;
+	/* The current folio to write to. */
+	struct folio *out_folio = NULL;
+	unsigned long tot_in = 0;
+	unsigned long tot_out = 0;
+	const u64 start = cb->start;
+	const u32 len = cb->len;
+	const u64 end = start + len;
+	const u32 blocksize = fs_info->sectorsize;
+	const u32 min_folio_size = btrfs_min_folio_size(fs_info);
+
+	workspace->params = zstd_get_btrfs_parameters(workspace->req_level, len);
+
+	/* Initialize the stream. */
+	stream = zstd_init_cstream(&workspace->params, len, workspace->mem, workspace->size);
+	if (unlikely(!stream)) {
+		btrfs_err(fs_info,
+	"zstd compression init level %d failed, root %llu inode %llu offset %llu",
+			  workspace->req_level, btrfs_root_id(inode->root),
+			  btrfs_ino(inode), start);
+		ret = -EIO;
+		goto out;
+	}
+
+	/* Map in the first page of input data. */
+	ret = btrfs_compress_filemap_get_folio(mapping, start, &in_folio);
+	if (ret < 0)
+		goto out;
+	workspace->in_buf.src = kmap_local_folio(in_folio, offset_in_folio(in_folio, start));
+	workspace->in_buf.pos = 0;
+	workspace->in_buf.size = btrfs_calc_input_length(in_folio, end, start);
+
+	/* Allocate and map in the output buffer. */
+	out_folio = btrfs_alloc_compr_folio(fs_info);
+	if (out_folio == NULL) {
+		ret = -ENOMEM;
+		goto out;
+	}
+	workspace->out_buf.dst = folio_address(out_folio);
+	workspace->out_buf.pos = 0;
+	workspace->out_buf.size = min_folio_size;
+
+	while (1) {
+		size_t ret2;
+
+		ret2 = zstd_compress_stream(stream, &workspace->out_buf, &workspace->in_buf);
+		if (unlikely(zstd_is_error(ret2))) {
+			btrfs_warn(fs_info,
+"zstd compression level %d failed, error %d root %llu inode %llu offset %llu",
+				   workspace->req_level, zstd_get_error_code(ret2),
+				   btrfs_root_id(inode->root), btrfs_ino(inode),
+				   start + tot_in);
+			ret = -EIO;
+			goto out;
+		}
+
+		/* Check to see if we are making it bigger. */
+		if (tot_in + workspace->in_buf.pos > blocksize * 2 &&
+		    tot_in + workspace->in_buf.pos < tot_out + workspace->out_buf.pos) {
+			ret = -E2BIG;
+			goto out;
+		}
+
+		/* Check if we need more output space. */
+		if (workspace->out_buf.pos >= workspace->out_buf.size) {
+			tot_out += min_folio_size;
+			if (tot_out >= len) {
+				ret = -E2BIG;
+				goto out;
+			}
+			/* Queue the current foliot into the bio. */
+			if (!bio_add_folio(bio, out_folio, folio_size(out_folio), 0)) {
+				ret = -E2BIG;
+				goto out;
+			}
+
+			out_folio = btrfs_alloc_compr_folio(fs_info);
+			if (out_folio == NULL) {
+				ret = -ENOMEM;
+				goto out;
+			}
+			workspace->out_buf.dst = folio_address(out_folio);
+			workspace->out_buf.pos = 0;
+			workspace->out_buf.size = min_folio_size;
+		}
+
+		/* We've reached the end of the input. */
+		if (tot_in + workspace->in_buf.pos >= len) {
+			tot_in += workspace->in_buf.pos;
+			break;
+		}
+
+		/* Check if we need more input. */
+		if (workspace->in_buf.pos >= workspace->in_buf.size) {
+			u64 cur;
+
+			tot_in += workspace->in_buf.size;
+			cur = start + tot_in;
+
+			kunmap_local(workspace->in_buf.src);
+			workspace->in_buf.src = NULL;
+			folio_put(in_folio);
+
+			ret = btrfs_compress_filemap_get_folio(mapping, cur, &in_folio);
+			if (ret < 0)
+				goto out;
+			workspace->in_buf.src = kmap_local_folio(in_folio,
+							 offset_in_folio(in_folio, cur));
+			workspace->in_buf.pos = 0;
+			workspace->in_buf.size = btrfs_calc_input_length(in_folio, end, cur);
+		}
+	}
+
+	while (1) {
+		size_t ret2;
+
+		ret2 = zstd_end_stream(stream, &workspace->out_buf);
+		if (unlikely(zstd_is_error(ret2))) {
+			btrfs_err(fs_info,
+"zstd compression end level %d failed, error %d root %llu inode %llu offset %llu",
+				  workspace->req_level, zstd_get_error_code(ret2),
+				  btrfs_root_id(inode->root), btrfs_ino(inode),
+				  start + tot_in);
+			ret = -EIO;
+			goto out;
+		}
+		/* Queue the remaining part of the output folio into bio. */
+		if (ret2 == 0) {
+			tot_out += workspace->out_buf.pos;
+			if (tot_out >= len) {
+				ret = -E2BIG;
+				goto out;
+			}
+			if (!bio_add_folio(bio, out_folio, workspace->out_buf.pos, 0)) {
+				ret = -E2BIG;
+				goto out;
+			}
+			out_folio = NULL;
+			break;
+		}
+		tot_out += min_folio_size;
+		if (tot_out >= len) {
+			ret = -E2BIG;
+			goto out;
+		}
+		if (!bio_add_folio(bio, out_folio, folio_size(out_folio), 0)) {
+			ret = -E2BIG;
+			goto out;
+		}
+		out_folio = btrfs_alloc_compr_folio(fs_info);
+		if (out_folio == NULL) {
+			ret = -ENOMEM;
+			goto out;
+		}
+		workspace->out_buf.dst = folio_address(out_folio);
+		workspace->out_buf.pos = 0;
+		workspace->out_buf.size = min_folio_size;
+	}
+
+	if (tot_out >= tot_in) {
+		ret = -E2BIG;
+		goto out;
+	}
+
+	ret = 0;
+	ASSERT(tot_out == bio->bi_iter.bi_size);
+out:
+	if (out_folio)
+		btrfs_free_compr_folio(out_folio);
+	if (workspace->in_buf.src) {
+		kunmap_local(workspace->in_buf.src);
+		folio_put(in_folio);
+	}
+	return ret;
+}
+
 int zstd_decompress_bio(struct list_head *ws, struct compressed_bio *cb)
 {
 	struct btrfs_fs_info *fs_info = cb_to_fs_info(cb);