1 files changed, 102 insertions, 0 deletions
diff --git a/src/rocksdb/db/blob/blob_garbage_meter.h b/src/rocksdb/db/blob/blob_garbage_meter.h
new file mode 100644
index 000000000..a6c04b0b2
--- /dev/null
+++ b/src/rocksdb/db/blob/blob_garbage_meter.h
@@ -0,0 +1,102 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#pragma once
+
+#include <cassert>
+#include <cstdint>
+#include <unordered_map>
+
+#include "db/blob/blob_constants.h"
+#include "rocksdb/rocksdb_namespace.h"
+#include "rocksdb/status.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+class Slice;
+
+// A class that can be used to compute the amount of additional garbage
+// generated by a compaction. It parses the keys and blob references in the
+// input and output of a compaction, and aggregates the "inflow" and "outflow"
+// on a per-blob file basis. The amount of additional garbage for any given blob
+// file can then be computed by subtracting the outflow from the inflow.
+class BlobGarbageMeter {
+ public:
+  // A class to store the number and total size of blobs on a per-blob file
+  // basis.
+  class BlobStats {
+   public:
+    void Add(uint64_t bytes) {
+      ++count_;
+      bytes_ += bytes;
+    }
+    void Add(uint64_t count, uint64_t bytes) {
+      count_ += count;
+      bytes_ += bytes;
+    }
+
+    uint64_t GetCount() const { return count_; }
+    uint64_t GetBytes() const { return bytes_; }
+
+   private:
+    uint64_t count_ = 0;
+    uint64_t bytes_ = 0;
+  };
+
+  // A class to keep track of the "inflow" and the "outflow" and to compute the
+  // amount of additional garbage for a given blob file.
+  class BlobInOutFlow {
+   public:
+    void AddInFlow(uint64_t bytes) {
+      in_flow_.Add(bytes);
+      assert(IsValid());
+    }
+    void AddOutFlow(uint64_t bytes) {
+      out_flow_.Add(bytes);
+      assert(IsValid());
+    }
+
+    const BlobStats& GetInFlow() const { return in_flow_; }
+    const BlobStats& GetOutFlow() const { return out_flow_; }
+
+    bool IsValid() const {
+      return in_flow_.GetCount() >= out_flow_.GetCount() &&
+             in_flow_.GetBytes() >= out_flow_.GetBytes();
+    }
+    bool HasGarbage() const {
+      assert(IsValid());
+      return in_flow_.GetCount() > out_flow_.GetCount();
+    }
+    uint64_t GetGarbageCount() const {
+      assert(IsValid());
+      assert(HasGarbage());
+      return in_flow_.GetCount() - out_flow_.GetCount();
+    }
+    uint64_t GetGarbageBytes() const {
+      assert(IsValid());
+      assert(HasGarbage());
+      return in_flow_.GetBytes() - out_flow_.GetBytes();
+    }
+
+   private:
+    BlobStats in_flow_;
+    BlobStats out_flow_;
+  };
+
+  Status ProcessInFlow(const Slice& key, const Slice& value);
+  Status ProcessOutFlow(const Slice& key, const Slice& value);
+
+  const std::unordered_map<uint64_t, BlobInOutFlow>& flows() const {
+    return flows_;
+  }
+
+ private:
+  static Status Parse(const Slice& key, const Slice& value,
+                      uint64_t* blob_file_number, uint64_t* bytes);
+
+  std::unordered_map<uint64_t, BlobInOutFlow> flows_;
+};
+
+}  // namespace ROCKSDB_NAMESPACE