From 1ccdcb5b6c1cfd176a86df4f115a88accc81b4d0 Mon Sep 17 00:00:00 2001
From: Jason Dillaman <dillaman@redhat.com>
Date: Thu, 28 Apr 2016 15:53:59 -0400
Subject: [PATCH] rbd: disk usage now includes all snapshots by default

Added a new "--from-snap" optional to restrict the calculation to
specific subsets of snapshots and the HEAD revision.

Signed-off-by: Jason Dillaman <dillaman@redhat.com>
---
 src/test/cli/rbd/help.t           |  2 ++
 src/tools/rbd/action/DiskUsage.cc | 32 ++++++++++++++++++++++++++-----
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/src/test/cli/rbd/help.t b/src/test/cli/rbd/help.t
index f68a486616f2..97ff90a48970 100644
--- a/src/test/cli/rbd/help.t
+++ b/src/test/cli/rbd/help.t
@@ -282,6 +282,7 @@
   rbd help disk-usage
   usage: rbd disk-usage [--pool <pool>] [--image <image>] [--snap <snap>] 
                         [--format <format>] [--pretty-format] 
+                        [--from-snap <from-snap>] 
                         <image-or-snap-spec> 
   
   Show disk usage stats for pool, image or snapshot
@@ -296,6 +297,7 @@
     --snap arg            snapshot name
     --format arg          output format [plain, json, or xml]
     --pretty-format       pretty formatting (json and xml)
+    --from-snap arg       snapshot starting point
   
   rbd help export
   usage: rbd export [--pool <pool>] [--image <image>] [--snap <snap>] 
diff --git a/src/tools/rbd/action/DiskUsage.cc b/src/tools/rbd/action/DiskUsage.cc
index e694dd742d99..9ef2e7e93a4e 100644
--- a/src/tools/rbd/action/DiskUsage.cc
+++ b/src/tools/rbd/action/DiskUsage.cc
@@ -86,8 +86,8 @@ static int compute_image_disk_usage(const std::string& name,
 }
 
 static int do_disk_usage(librbd::RBD &rbd, librados::IoCtx &io_ctx,
-                        const char *imgname, const char *snapname,
-                        Formatter *f) {
+                         const char *imgname, const char *snapname,
+                         const char *from_snapname, Formatter *f) {
   std::vector<std::string> names;
   int r = rbd.list(io_ctx, names);
   if (r == -ENOENT) {
@@ -106,6 +106,7 @@ static int do_disk_usage(librbd::RBD &rbd, librados::IoCtx &io_ctx,
     tbl.define_column("USED", TextTable::RIGHT, TextTable::RIGHT);
   }
 
+  uint32_t count = 0;
   uint64_t used_size = 0;
   uint64_t total_prov = 0;
   uint64_t total_used = 0;
@@ -152,6 +153,7 @@ static int do_disk_usage(librbd::RBD &rbd, librados::IoCtx &io_ctx,
       continue;
     }
 
+    bool found_from_snap = (from_snapname == nullptr);
     std::string last_snap_name;
     std::sort(snap_list.begin(), snap_list.end(),
               boost::bind(&librbd::snap_info_t::id, _1) <
@@ -167,7 +169,8 @@ static int do_disk_usage(librbd::RBD &rbd, librados::IoCtx &io_ctx,
         goto out;
       }
 
-      if (imgname == NULL || (snapname != NULL && snap->name == snapname)) {
+      if (imgname == nullptr || found_from_snap ||
+         (found_from_snap && snapname != nullptr && snap->name == snapname)) {
         r = compute_image_disk_usage(*name, snap->name, last_snap_name,
                                      snap_image, snap->size, tbl, f,
                                      &used_size);
@@ -179,6 +182,15 @@ static int do_disk_usage(librbd::RBD &rbd, librados::IoCtx &io_ctx,
           total_prov += snap->size;
         }
         total_used += used_size;
+        ++count;
+      }
+
+      if (!found_from_snap && from_snapname != nullptr &&
+          snap->name == from_snapname) {
+        found_from_snap = true;
+      }
+      if (snapname != nullptr && snap->name == snapname) {
+        break;
       }
       last_snap_name = snap->name;
     }
@@ -191,6 +203,7 @@ static int do_disk_usage(librbd::RBD &rbd, librados::IoCtx &io_ctx,
       }
       total_prov += info.size;
       total_used += used_size;
+      ++count;
     }
   }
 
@@ -204,7 +217,7 @@ out:
     f->close_section();
     f->flush(std::cout);
   } else {
-    if (imgname == NULL) {
+    if (count > 1) {
       tbl << "<TOTAL>"
           << stringify(si_t(total_prov))
           << stringify(si_t(total_used))
@@ -221,6 +234,9 @@ void get_arguments(po::options_description *positional,
   at::add_image_or_snap_spec_options(positional, options,
                                      at::ARGUMENT_MODIFIER_NONE);
   at::add_format_options(options);
+  options->add_options()
+    (at::FROM_SNAPSHOT_NAME.c_str(), po::value<std::string>(),
+     "snapshot starting point");
 }
 
 int execute(const po::variables_map &vm) {
@@ -231,11 +247,16 @@ int execute(const po::variables_map &vm) {
   int r = utils::get_pool_image_snapshot_names(
     vm, at::ARGUMENT_MODIFIER_NONE, &arg_index, &pool_name, &image_name,
     &snap_name, utils::SNAPSHOT_PRESENCE_PERMITTED,
-    utils::SPEC_VALIDATION_NONE, false);
+    utils::SPEC_VALIDATION_NONE, vm.count(at::FROM_SNAPSHOT_NAME));
   if (r < 0) {
     return r;
   }
 
+  std::string from_snap_name;
+  if (vm.count(at::FROM_SNAPSHOT_NAME)) {
+    from_snap_name = vm[at::FROM_SNAPSHOT_NAME].as<std::string>();
+  }
+
   at::Format::Formatter formatter;
   r = utils::get_formatter(vm, &formatter);
   if (r < 0) {
@@ -253,6 +274,7 @@ int execute(const po::variables_map &vm) {
   r = do_disk_usage(rbd, io_ctx,
                     image_name.empty() ? nullptr: image_name.c_str() ,
                     snap_name.empty() ? nullptr : snap_name.c_str(),
+                    from_snap_name.empty() ? nullptr : from_snap_name.c_str(),
                     formatter.get());
   if (r < 0) {
     std::cerr << "du failed: " << cpp_strerror(r) << std::endl;
-- 
2.47.3