From: Boris Burkov <bo...@bur.io>

In send stream v2, write commands can now be an arbitrary size. For that
reason, we can no longer allocate a fixed array in sctx for read_cmd.
Instead, read_cmd dynamically allocates sctx->read_buf. To avoid
needless reallocations, we reuse read_buf between read_cmd calls by also
keeping track of the size of the allocated buffer in sctx->read_buf_sz.

We do the first allocation of the old default size at the start of
processing the stream, and we only reallocate if we encounter a command
that needs a larger buffer.

Signed-off-by: Boris Burkov <bo...@bur.io>
---
 common/send-stream.c | 55 ++++++++++++++++++++++++++++----------------
 send.h               |  2 +-
 2 files changed, 36 insertions(+), 21 deletions(-)

diff --git a/common/send-stream.c b/common/send-stream.c
index cd5aa311..3d3585c3 100644
--- a/common/send-stream.c
+++ b/common/send-stream.c
@@ -35,11 +35,11 @@ struct btrfs_send_attribute {
 };
 
 struct btrfs_send_stream {
-       char read_buf[BTRFS_SEND_BUF_SIZE];
+       char *read_buf;
+       size_t read_buf_sz;
        int fd;
 
        int cmd;
-       struct btrfs_cmd_header *cmd_hdr;
        struct btrfs_send_attribute cmd_attrs[BTRFS_SEND_A_MAX + 1];
        u32 version;
 
@@ -111,11 +111,12 @@ static int read_cmd(struct btrfs_send_stream *sctx)
        u32 pos;
        u32 crc;
        u32 crc2;
+       struct btrfs_cmd_header *cmd_hdr;
+       size_t buf_len;
 
        memset(sctx->cmd_attrs, 0, sizeof(sctx->cmd_attrs));
 
-       ASSERT(sizeof(*sctx->cmd_hdr) <= sizeof(sctx->read_buf));
-       ret = read_buf(sctx, sctx->read_buf, sizeof(*sctx->cmd_hdr));
+       ret = read_buf(sctx, sctx->read_buf, sizeof(*cmd_hdr));
        if (ret < 0)
                goto out;
        if (ret) {
@@ -124,18 +125,22 @@ static int read_cmd(struct btrfs_send_stream *sctx)
                goto out;
        }
 
-       sctx->cmd_hdr = (struct btrfs_cmd_header *)sctx->read_buf;
-       cmd = le16_to_cpu(sctx->cmd_hdr->cmd);
-       cmd_len = le32_to_cpu(sctx->cmd_hdr->len);
-
-       if (cmd_len + sizeof(*sctx->cmd_hdr) >= sizeof(sctx->read_buf)) {
-               ret = -EINVAL;
-               error("command length %u too big for buffer %zu",
-                               cmd_len, sizeof(sctx->read_buf));
-               goto out;
+       cmd_hdr = (struct btrfs_cmd_header *)sctx->read_buf;
+       cmd_len = le32_to_cpu(cmd_hdr->len);
+       cmd = le16_to_cpu(cmd_hdr->cmd);
+       buf_len = sizeof(*cmd_hdr) + cmd_len;
+       if (sctx->read_buf_sz < buf_len) {
+               sctx->read_buf = realloc(sctx->read_buf, buf_len);
+               if (!sctx->read_buf) {
+                       ret = -ENOMEM;
+                       error("failed to reallocate read buffer for cmd");
+                       goto out;
+               }
+               sctx->read_buf_sz = buf_len;
+               /* We need to reset cmd_hdr after realloc of sctx->read_buf */
+               cmd_hdr = (struct btrfs_cmd_header *)sctx->read_buf;
        }
-
-       data = sctx->read_buf + sizeof(*sctx->cmd_hdr);
+       data = sctx->read_buf + sizeof(*cmd_hdr);
        ret = read_buf(sctx, data, cmd_len);
        if (ret < 0)
                goto out;
@@ -145,11 +150,12 @@ static int read_cmd(struct btrfs_send_stream *sctx)
                goto out;
        }
 
-       crc = le32_to_cpu(sctx->cmd_hdr->crc);
-       sctx->cmd_hdr->crc = 0;
+       crc = le32_to_cpu(cmd_hdr->crc);
+       /* in send, crc is computed with header crc = 0, replicate that */
+       cmd_hdr->crc = 0;
 
        crc2 = crc32c(0, (unsigned char*)sctx->read_buf,
-                       sizeof(*sctx->cmd_hdr) + cmd_len);
+                       sizeof(*cmd_hdr) + cmd_len);
 
        if (crc != crc2) {
                ret = -EINVAL;
@@ -524,19 +530,28 @@ int btrfs_read_and_process_send_stream(int fd,
                goto out;
        }
 
+       sctx.read_buf = malloc(BTRFS_SEND_BUF_SIZE_V1);
+       if (!sctx.read_buf) {
+               ret = -ENOMEM;
+               error("unable to allocate send stream read buffer");
+               goto out;
+       }
+       sctx.read_buf_sz = BTRFS_SEND_BUF_SIZE_V1;
+
        while (1) {
                ret = read_and_process_cmd(&sctx);
                if (ret < 0) {
                        last_err = ret;
                        errors++;
                        if (max_errors > 0 && errors >= max_errors)
-                               goto out;
+                               break;
                } else if (ret > 0) {
                        if (!honor_end_cmd)
                                ret = 0;
-                       goto out;
+                       break;
                }
        }
+       free(sctx.read_buf);
 
 out:
        if (last_err && !ret)
diff --git a/send.h b/send.h
index 8dd865ec..228928a0 100644
--- a/send.h
+++ b/send.h
@@ -33,7 +33,7 @@ extern "C" {
 #define BTRFS_SEND_STREAM_MAGIC "btrfs-stream"
 #define BTRFS_SEND_STREAM_VERSION 1
 
-#define BTRFS_SEND_BUF_SIZE SZ_64K
+#define BTRFS_SEND_BUF_SIZE_V1 SZ_64K
 #define BTRFS_SEND_READ_SIZE (1024 * 48)
 
 enum btrfs_tlv_type {
-- 
2.30.2

Reply via email to