[Qemu-devel] [PATCH v3 06/22] qcow2: Helper for refcount array reallocation

[Max Reitz]

Add a helper function for reallocating a refcount array, independent of
the refcount order. The newly allocated space is zeroed and the function
handles failed reallocations gracefully.

The helper function will always align the buffer size to a cluster
boundary; if storing the refcounts in such an array in big endian byte
order, this makes it possible to write parts of the array directly as
refcount blocks into the image file.

Signed-off-by: Max Reitz <address@hidden>
Reviewed-by: Eric Blake <address@hidden>
---
 block/qcow2-refcount.c | 135 +++++++++++++++++++++++++++++++------------------
 1 file changed, 86 insertions(+), 49 deletions(-)

diff --git a/block/qcow2-refcount.c b/block/qcow2-refcount.c
index 66c78c0..3cd540c 100644
--- a/block/qcow2-refcount.c
+++ b/block/qcow2-refcount.c
@@ -1108,6 +1108,68 @@ fail:
 /* refcount checking functions */
 
 
+static size_t refcount_array_byte_size(BDRVQcowState *s, uint64_t entries)
+{
+    if (s->refcount_order < 3) {
+        /* sub-byte width */
+        int shift = 3 - s->refcount_order;
+        return (entries + (1 << shift) - 1) >> shift;
+    } else if (s->refcount_order == 3) {
+        /* byte width */
+        return entries;
+    } else {
+        /* multiple bytes wide */
+
+        /* This assertion holds because there is no way we can address more 
than
+         * 2^(64 - 9) clusters at once (with cluster size 512 = 2^9, and 
because
+         * offsets have to be representable in bytes); due to every cluster
+         * corresponding to one refcount entry and because refcount_order has 
to
+         * be below 7, we are far below that limit */
+        assert(!(entries >> (64 - (s->refcount_order - 3))));
+
+        return entries << (s->refcount_order - 3);
+    }
+}
+
+/**
+ * Reallocates *array so that it can hold new_size entries. *size must contain
+ * the current number of entries in *array. If the reallocation fails, *array
+ * and *size will not be modified and -errno will be returned. If the
+ * reallocation is successful, *array will be set to the new buffer and *size
+ * will be set to new_size. The size of the reallocated refcount array buffer
+ * will be aligned to a cluster boundary, and the newly allocated area will be
+ * zeroed.
+ */
+static int realloc_refcount_array(BDRVQcowState *s, uint16_t **array,
+                                  int64_t *size, int64_t new_size)
+{
+    /* Round to clusters so the array can be directly written to disk */
+    size_t old_byte_size = ROUND_UP(refcount_array_byte_size(s, *size),
+                                    s->cluster_size);
+    size_t new_byte_size = ROUND_UP(refcount_array_byte_size(s, new_size),
+                                    s->cluster_size);
+    uint16_t *new_ptr;
+
+    if (new_byte_size <= old_byte_size) {
+        *size = new_size;
+        return 0;
+    }
+
+    assert(new_byte_size > 0);
+
+    new_ptr = g_try_realloc(*array, new_byte_size);
+    if (!new_ptr) {
+        return -ENOMEM;
+    }
+
+    memset((void *)((uintptr_t)new_ptr + old_byte_size), 0,
+           new_byte_size - old_byte_size);
+
+    *array = new_ptr;
+    *size  = new_size;
+
+    return 0;
+}
 
 /*
  * Increases the refcount for a range of clusters in a given refcount table.
@@ -1124,6 +1186,7 @@ static int inc_refcounts(BlockDriverState *bs,
 {
     BDRVQcowState *s = bs->opaque;
     uint64_t start, last, cluster_offset, k;
+    int ret;
 
     if (size <= 0) {
         return 0;
@@ -1135,23 +1198,12 @@ static int inc_refcounts(BlockDriverState *bs,
         cluster_offset += s->cluster_size) {
         k = cluster_offset >> s->cluster_bits;
         if (k >= *refcount_table_size) {
-            int64_t old_refcount_table_size = *refcount_table_size;
-            uint16_t *new_refcount_table;
-
-            *refcount_table_size = k + 1;
-            new_refcount_table = g_try_realloc(*refcount_table,
-                                               *refcount_table_size *
-                                               sizeof(**refcount_table));
-            if (!new_refcount_table) {
-                *refcount_table_size = old_refcount_table_size;
+            ret = realloc_refcount_array(s, refcount_table,
+                                         refcount_table_size, k + 1);
+            if (ret < 0) {
                 res->check_errors++;
-                return -ENOMEM;
+                return ret;
             }
-            *refcount_table = new_refcount_table;
-
-            memset(*refcount_table + old_refcount_table_size, 0,
-                   (*refcount_table_size - old_refcount_table_size) *
-                   sizeof(**refcount_table));
         }
 
         if (++(*refcount_table)[k] == 0) {
@@ -1518,8 +1570,7 @@ static int check_refblocks(BlockDriverState *bs, 
BdrvCheckResult *res,
                     fix & BDRV_FIX_ERRORS ? "Repairing" : "ERROR", i);
 
             if (fix & BDRV_FIX_ERRORS) {
-                int64_t old_nb_clusters = *nb_clusters;
-                uint16_t *new_refcount_table;
+                int64_t new_nb_clusters;
 
                 if (offset > INT64_MAX - s->cluster_size) {
                     ret = -EINVAL;
@@ -1536,22 +1587,15 @@ static int check_refblocks(BlockDriverState *bs, 
BdrvCheckResult *res,
                     goto resize_fail;
                 }
 
-                *nb_clusters = size_to_clusters(s, size);
-                assert(*nb_clusters >= old_nb_clusters);
+                new_nb_clusters = size_to_clusters(s, size);
+                assert(new_nb_clusters >= *nb_clusters);
 
-                new_refcount_table = g_try_realloc(*refcount_table,
-                                                   *nb_clusters *
-                                                   sizeof(**refcount_table));
-                if (!new_refcount_table) {
-                    *nb_clusters = old_nb_clusters;
+                ret = realloc_refcount_array(s, refcount_table,
+                                             nb_clusters, new_nb_clusters);
+                if (ret < 0) {
                     res->check_errors++;
-                    return -ENOMEM;
+                    return ret;
                 }
-                *refcount_table = new_refcount_table;
-
-                memset(*refcount_table + old_nb_clusters, 0,
-                       (*nb_clusters - old_nb_clusters) *
-                       sizeof(**refcount_table));
 
                 if (cluster >= *nb_clusters) {
                     ret = -EINVAL;
@@ -1611,10 +1655,12 @@ static int calculate_refcounts(BlockDriverState *bs, 
BdrvCheckResult *res,
     int ret;
 
     if (!*refcount_table) {
-        *refcount_table = g_try_new0(uint16_t, *nb_clusters);
-        if (*nb_clusters && *refcount_table == NULL) {
+        int64_t old_size = 0;
+        ret = realloc_refcount_array(s, refcount_table,
+                                     &old_size, *nb_clusters);
+        if (ret < 0) {
             res->check_errors++;
-            return -ENOMEM;
+            return ret;
         }
     }
 
@@ -1746,6 +1792,7 @@ static int64_t alloc_clusters_imrt(BlockDriverState *bs,
     int64_t cluster = *first_free_cluster, i;
     bool first_gap = true;
     int contiguous_free_clusters;
+    int ret;
 
     /* Starting at *first_free_cluster, find a range of at least cluster_count
      * continuously free clusters */
@@ -1775,28 +1822,18 @@ static int64_t alloc_clusters_imrt(BlockDriverState *bs,
     /* If no such range could be found, grow the in-memory refcount table
      * accordingly to append free clusters at the end of the image */
     if (contiguous_free_clusters < cluster_count) {
-        int64_t old_imrt_nb_clusters = *imrt_nb_clusters;
-        uint16_t *new_refcount_table;
-
         /* contiguous_free_clusters clusters are already empty at the image 
end;
          * we need cluster_count clusters; therefore, we have to allocate
          * cluster_count - contiguous_free_clusters new clusters at the end of
          * the image (which is the current value of cluster; note that cluster
          * may exceed old_imrt_nb_clusters if *first_free_cluster pointed 
beyond
          * the image end) */
-        *imrt_nb_clusters = cluster + cluster_count - contiguous_free_clusters;
-        new_refcount_table = g_try_realloc(*refcount_table,
-                                           *imrt_nb_clusters *
-                                           sizeof(**refcount_table));
-        if (!new_refcount_table) {
-            *imrt_nb_clusters = old_imrt_nb_clusters;
-            return -ENOMEM;
-        }
-        *refcount_table = new_refcount_table;
-
-        memset(*refcount_table + old_imrt_nb_clusters, 0,
-               (*imrt_nb_clusters - old_imrt_nb_clusters) *
-               sizeof(**refcount_table));
+        ret = realloc_refcount_array(s, refcount_table, imrt_nb_clusters,
+                                     cluster + cluster_count
+                                     - contiguous_free_clusters);
+        if (ret < 0) {
+            return ret;
+        }
     }
 
     /* Go back to the first free cluster */
-- 
1.9.3