[yugabyte#28010] docdb: Account for snapshots in per-tablet on-disk size metrics

Summary:
The current `OnDiskSizeInfo` struct for a tablet does not account for snapshots, which can potentially have a lot of data. We need to account for snapshot size so that the cluster balancer can estimate how long remote bootstraps will take (that work will come in future diffs). This diff enhances the existing `ts_data_size` metric updater, which runs every 60s and updates each tserver's total on-disk size, to also set the on-disk size for each tablet, accounting for snapshots, MANIFEST files, the retryable requests file, and anything else in the data directory. It still accounts for hard links, so creating a snapshot when there is 1 GB of SSTs in the tablet will not change the total tablet size (because the snapshot is a hardlink to the SST in the active rocksdb).

The total on-disk size is stored in-memory in each tablet peer. This is passed back to the master via `StorageMetadata` (the same way we currently pass back the SST and WAL sizes), but is not currently used.

Also changed the metrics heartbeat updater to use a common function for getting tablet sizes.

**Upgrade/Downgrade Safety**
Added `total_size` field to `TabletDriveStorageMetadataPB` which is passed back as part of the metrics heartbeat from tserver to master. On upgrade/downgrade, master code is always newer than tserver code and thus may receive a value of 0 for a tablet. Future master code that uses this field must be tolerant to this.

Also added `total_on_disk_size` to the TabletStatusPB proto. This field is read on the metrics heartbeat path and paths where we call ListTablets (like the tablets UI).

The tablets UI has been changed to use the size with snapshots under "Total" (the previous field was just a sum of the sizes of consensus metadata, WALs, and SSTs). If the size with snapshots has not been calculated, "N/A" is displayed. Also made some UI readability modifications: the uncompressed SST size has been removed and the sub-components are now a sublist of Total:
Old UI:
{F377256}

New UI:
{F377270}

During upgrades/downgrades, some tservers may have the field and some may not. Both master and tserver code should be tolerant to getting a value of 0 or a stale value from any tserver.
Jira: DB-17628

Test Plan:
`./yb_build.sh --cxx-test master_heartbeat-itest --gtest_filter=DriveInfoTest.DriveInfo`
`./yb_build.sh --cxx-test integration-tests_tserver_metrics_heartbeat_data_provider-itest --gtest_filter TServerDriveInfoMetricsHeartbeatDataProviderITest.DriveInfo`

Reviewers: mhaddad

Reviewed By: mhaddad

Subscribers: ybase

Differential Revision: https://phorge.dev.yugabyte.com/D45458

Author: SrivastavaAnubhav

src/yb/integration-tests/master_heartbeat-itest.cc

@@ -65,6 +65,8 @@ DECLARE_int32(TEST_mini_cluster_registration_wait_time_sec);
 DECLARE_int32(tserver_unresponsive_timeout_ms);
 DECLARE_bool(persist_tserver_registry);
 DECLARE_bool(master_enable_universe_uuid_heartbeat_check);
+DECLARE_int32(data_size_metric_updater_interval_sec);
+DECLARE_int32(tserver_heartbeat_metrics_interval_ms);
 
 namespace yb::integration_tests {
 
@@ -461,6 +463,41 @@ TEST_F(MasterHeartbeatITest, TestRegistrationThroughRaftPersisted) {
       << "TS registered through raft config should be unresponsive, not live";
 }
 
+class DriveInfoTest : public MasterHeartbeatITest {
+ public:
+  void SetUp() override {
+    ANNOTATE_UNPROTECTED_WRITE(FLAGS_data_size_metric_updater_interval_sec) = 1;
+    ANNOTATE_UNPROTECTED_WRITE(FLAGS_tserver_heartbeat_metrics_interval_ms) = 1000;
+    MasterHeartbeatITest::SetUp();
+  }
+};
+
+TEST_F(DriveInfoTest, DriveInfo) {
+  CreateTable();
+  auto table = table_name();
+  auto& catalog_mgr = ASSERT_RESULT(mini_cluster_->GetLeaderMiniMaster())->catalog_manager();
+  auto table_info = catalog_mgr.GetTableInfoFromNamespaceNameAndTableName(
+      table.namespace_type(), table.namespace_name(), table.table_name());
+  auto tablets = ASSERT_RESULT(table_info->GetTablets());
+
+  // Insert 1000 rows and flush to an SST.
+  for (int i = 0; i < 1000; ++i) {
+    PutKeyValue(Format("k$0", i), Format("v$0", i));
+  }
+  ASSERT_OK(mini_cluster_->CompactTablets());
+  ASSERT_OK(WaitFor([&]() {
+    for (const auto& tablet : tablets) {
+      for (auto& replica : *tablet->GetReplicaLocations()) {
+        if (replica.second.drive_info.sst_files_size == 0) return false;
+        if (replica.second.drive_info.wal_files_size == 0) return false;
+        if (replica.second.drive_info.uncompressed_sst_file_size == 0) return false;
+        if (replica.second.drive_info.total_size == 0) return false;
+      }
+    }
+    return true;
+  }, 30s, "Waiting for drive info to be populated for all tablets"));
+}
+
 class PersistTabletServerRegistryUpgradeTest : public MasterHeartbeatITest {
  public:
   void SetUp() override {

src/yb/integration-tests/minicluster-snapshot-test.cc

@@ -41,6 +41,7 @@
 
 #include "yb/client/client.h"
 #include "yb/client/schema.h"
+#include "yb/client/snapshot_test_util.h"
 #include "yb/client/table_creator.h"
 #include "yb/client/table_info.h"
 #include "yb/client/transaction_manager.h"
@@ -83,6 +84,7 @@
 
 DECLARE_int32(cleanup_split_tablets_interval_sec);
 DECLARE_int32(data_size_metric_updater_interval_sec);
+DECLARE_int32(timestamp_history_retention_interval_sec);
 DECLARE_bool(enable_db_clone);
 DECLARE_int32(load_balancer_initial_delay_secs);
 DECLARE_bool(master_auto_run_initdb);
@@ -195,8 +197,8 @@ Status DeleteSnapshotSchedule(MasterBackupProxy* proxy, const SnapshotScheduleId
 
 Result<TxnSnapshotId> WaitNewSnapshot(MasterBackupProxy* proxy, const SnapshotScheduleId& id) {
   LOG(INFO) << "WaitNewSnapshot, schedule id: " << id;
-  std::string last_snapshot_id;
-  std::string new_snapshot_id;
+  TxnSnapshotId last_snapshot_id;
+  TxnSnapshotId new_snapshot_id;
   RETURN_NOT_OK(WaitFor(
       [&proxy, &id, &last_snapshot_id, &new_snapshot_id]() -> Result<bool> {
         // If there's a master leader failover then we should wait for the next cycle.
@@ -205,9 +207,10 @@ Result<TxnSnapshotId> WaitNewSnapshot(MasterBackupProxy* proxy, const SnapshotSc
         if (snapshots.empty()) {
           return false;
         }
-        auto snapshot_id = snapshots[snapshots.size() - 1].id();
+        auto snapshot_id = VERIFY_RESULT(FullyDecodeTxnSnapshotId(
+            snapshots[snapshots.size() - 1].id()));
         LOG(INFO) << "WaitNewSnapshot, last snapshot id: " << snapshot_id;
-        if (last_snapshot_id.empty()) {
+        if (last_snapshot_id.IsNil()) {
           last_snapshot_id = snapshot_id;
           return false;
         }
@@ -219,7 +222,7 @@ Result<TxnSnapshotId> WaitNewSnapshot(MasterBackupProxy* proxy, const SnapshotSc
         }
       },
       kInterval * 5, "Wait new schedule snapshot"));
-  return FullyDecodeTxnSnapshotId(new_snapshot_id);
+  return new_snapshot_id;
 }
 
 Status WaitForRestoration(
@@ -450,6 +453,16 @@ class PostgresMiniClusterTest : public pgwrapper::PgMiniTestBase {
 
   MiniCluster* mini_cluster() { return cluster_.get(); }
 
+  Result<TableInfoPtr> GetTable(const std::string& table_name, const std::string& db_name) {
+    auto leader_master = VERIFY_RESULT(cluster_->GetLeaderMiniMaster());
+    for (const auto& table : leader_master->catalog_manager_impl().GetTables(GetTablesMode::kAll)) {
+      if (table->name() == table_name && table->namespace_name() == db_name) {
+        return table;
+      }
+    }
+    return STATUS_FORMAT(NotFound, "Table $0 not found", table_name);
+  }
+
   Status CreateDatabase(
       const std::string& namespace_name,
       master::YsqlColocationConfig colocated = master::YsqlColocationConfig::kNotColocated) {
@@ -653,16 +666,6 @@ class PgCloneInitiallyEmptyDBTest : public PostgresMiniClusterTest {
     PostgresMiniClusterTest::DoTearDown();
   }
 
-  Result<TableInfoPtr> GetTable(const std::string& table_name, const std::string& db_name) {
-    auto leader_master = VERIFY_RESULT(cluster_->GetLeaderMiniMaster());
-    for (const auto& table : leader_master->catalog_manager_impl().GetTables(GetTablesMode::kAll)) {
-      if (table->name() == table_name && table->namespace_name() == db_name) {
-        return table;
-      }
-    }
-    return STATUS_FORMAT(NotFound, "Table $0 not found", table_name);
-  }
-
   Status SplitTablet(const TabletId& tablet_id) {
     SplitTabletRequestPB req;
     SplitTabletResponsePB resp;
@@ -847,6 +850,70 @@ TEST_F(PgCloneTest, CloneVectorIndex) {
   ASSERT_OK(source_conn_->ExecuteFormat("DROP DATABASE $0", kTargetNamespaceName2));
 }
 
+class TabletDataSizeMetricsTest : public PostgresMiniClusterTest {
+ protected:
+  void SetUp() override {
+    ANNOTATE_UNPROTECTED_WRITE(FLAGS_data_size_metric_updater_interval_sec) = 1;
+    PostgresMiniClusterTest::SetUp();
+    messenger_ = ASSERT_RESULT(rpc::MessengerBuilder("test-msgr").set_num_reactors(1).Build());
+    proxy_cache_ = std::make_unique<rpc::ProxyCache>(messenger_.get());
+    snapshot_util_ = std::make_unique<client::SnapshotTestUtil>();
+    snapshot_util_->SetProxy(&client_->proxy_cache());
+    snapshot_util_->SetCluster(cluster_.get());
+  }
+
+  void DoTearDown() override {
+    messenger_->Shutdown();
+    PostgresMiniClusterTest::DoTearDown();
+  }
+
+  std::unique_ptr<client::SnapshotTestUtil> snapshot_util_;
+  std::unique_ptr<rpc::Messenger> messenger_;
+  std::unique_ptr<rpc::ProxyCache> proxy_cache_;
+};
+
+TEST_F(TabletDataSizeMetricsTest, TotalOnDiskSizeIncludesSnapshots) {
+  ASSERT_OK(CreateDatabase("testdb"));
+  auto conn = ASSERT_RESULT(ConnectToDB("testdb"));
+  ASSERT_OK(conn.Execute(
+      "CREATE TABLE t1 (key INT PRIMARY KEY, value INT) SPLIT INTO 1 TABLETS"));
+  auto table = ASSERT_RESULT(GetTable("t1", "testdb"));
+  auto tablet_id = ASSERT_RESULT(table->GetTablets())[0]->tablet_id();
+  auto tablet_peer = ASSERT_RESULT(
+      cluster_->mini_tablet_server(0)->server()->tablet_peer_lookup()->GetServingTablet(tablet_id));
+  LOG(INFO) << "Tablet id: " << tablet_id;
+
+  // Write some data and create a snapshot (this should create an SST file).
+  ASSERT_OK(conn.Execute(
+      "INSERT INTO t1 VALUES (generate_series(1,1000), generate_series(1,1000))"));
+  ASSERT_RESULT(snapshot_util_->CreateSnapshot(table->id()));
+
+  // Get the on-disk size of the tablet.
+  tablet::TabletOnDiskSizeInfo size_before;
+  ASSERT_OK(WaitFor([&]() {
+    size_before = tablet_peer->GetOnDiskSizeInfo();
+    return size_before.sst_files_disk_size > 0 &&
+           size_before.total_on_disk_size > size_before.sst_files_disk_size;
+  }, 30s, "Wait for on-disk size to include the SST"));
+
+  // Delete the data in the tablet and re-compact. The SST file size should drop to 0.
+  ASSERT_OK(conn.Execute("DELETE FROM t1"));
+  ASSERT_OK(WaitFor([&]() {
+    // Retrigger compactions until this peer has no SST files. This is required because if the
+    // delete has only applied on the other two peers, the compaction would not do anything on this
+    // peer.
+    FlushAndCompactTablets();
+    return tablet_peer->GetOnDiskSizeInfo().sst_files_disk_size == 0;
+  }, 30s, "Wait for on-disk size to drop after delete"));
+
+  // Once the on-disk size metric updater has run, check that the total on-disk size has not
+  // decrease by a significant portion of the size of the SST we deleted.
+  SleepFor(FLAGS_data_size_metric_updater_interval_sec * 2s);
+  auto size_after = tablet_peer->GetOnDiskSizeInfo();
+  auto size_diff = size_before.total_on_disk_size - size_after.total_on_disk_size;
+  ASSERT_LT(size_diff, size_before.sst_files_disk_size * 0.9);
+}
+
 class TsDataSizeMetricsTest : public PgCloneTest {
  public:
   uint64_t GetTsDataSize() {

src/yb/master/catalog_entity_info.h

@@ -145,11 +145,12 @@ struct TabletReplicaDriveInfo {
   uint64 wal_files_size = 0;
   uint64 uncompressed_sst_file_size = 0;
   bool may_have_orphaned_post_split_data = true;
+  uint64 total_size = 0;
 
   std::string ToString() const {
     return YB_STRUCT_TO_STRING(
         sst_files_size, wal_files_size, uncompressed_sst_file_size,
-        may_have_orphaned_post_split_data);
+        may_have_orphaned_post_split_data, total_size);
   }
 };
 

src/yb/master/master_heartbeat.proto

@@ -104,6 +104,7 @@ message TabletDriveStorageMetadataPB {
   optional uint64 wal_file_size = 3;
   optional uint64 uncompressed_sst_file_size = 4;
   optional bool may_have_orphaned_post_split_data = 5 [default = true];
+  optional uint64 total_size = 6;
 }
 
 message TabletLeaderMetricsPB {

src/yb/master/master_heartbeat_service.cc

@@ -1438,6 +1438,7 @@ void MasterHeartbeatServiceImpl::ProcessTabletMetadata(
     .wal_files_size = storage_metadata.wal_file_size(),
     .uncompressed_sst_file_size = storage_metadata.uncompressed_sst_file_size(),
     .may_have_orphaned_post_split_data = storage_metadata.may_have_orphaned_post_split_data(),
+    .total_size = storage_metadata.total_size(),
   };
   tablet->UpdateReplicaInfo(ts_uuid, drive_info, leader_lease_info);
 }

src/yb/tablet/tablet.proto

@@ -52,7 +52,10 @@ message TabletStatusPB {
   optional bytes start_key = 5;
   // DEPRECATED.
   optional bytes end_key = 6;
-  optional int64 estimated_on_disk_size = 7;
+  // Estimated size of the tablet on disk. Includes consensus metadata, WALs, and SSTs, but does
+  // not include snapshots, retryable requests, MANIFEST, and other structures.
+  // Unlike total_on_disk_size, this is always up-to-date.
+  optional int64 active_on_disk_size = 7;
   optional int64 consensus_metadata_disk_size = 11;
   optional int64 wal_files_disk_size = 12;
   optional int64 sst_files_disk_size = 13;
@@ -67,6 +70,9 @@ message TabletStatusPB {
   optional string pgschema_name = 20;
   repeated bytes vector_index_finished_backfills = 21;
   optional OpIdPB last_op_id = 22;
+  // Estimated size of the tablet on disk, including snapshots, retryable requests, MANIFEST,
+  // and other files in those directories. This is updated periodically, so may be stale or 0.
+  optional int64 total_on_disk_size = 23;
 }
 
 // Used to present the maintenance manager's internal state.

src/yb/tablet/tablet_peer.cc

@@ -798,7 +798,7 @@ void TabletPeer::GetTabletStatusPB(TabletStatusPB* status_pb_out) {
     std::lock_guard lock(lock_);
     DCHECK(status_pb_out != nullptr);
     DCHECK(status_listener_.get() != nullptr);
-    const auto disk_size_info = GetOnDiskSizeInfo();
+    const auto disk_size_info = GetOnDiskSizeInfoUnlocked();
     status_pb_out->set_tablet_id(status_listener_->tablet_id());
     status_pb_out->set_namespace_name(status_listener_->namespace_name());
     status_pb_out->set_table_name(status_listener_->table_name());
@@ -1585,6 +1585,11 @@ void TabletPeer::UnregisterMaintenanceOps() {
 }
 
 TabletOnDiskSizeInfo TabletPeer::GetOnDiskSizeInfo() const {
+  std::lock_guard lock(lock_);
+  return GetOnDiskSizeInfoUnlocked();
+}
+
+TabletOnDiskSizeInfo TabletPeer::GetOnDiskSizeInfoUnlocked() const {
   TabletOnDiskSizeInfo info;
 
   if (consensus_) {
@@ -1602,10 +1607,15 @@ TabletOnDiskSizeInfo TabletPeer::GetOnDiskSizeInfo() const {
     info.wal_files_disk_size = log->OnDiskSize();
   }
 
+  info.total_on_disk_size = total_on_disk_size_;
   info.RecomputeTotalSize();
   return info;
 }
 
+void TabletPeer::SetTabletOnDiskSize(size_t total_on_disk_size) {
+  total_on_disk_size_ = total_on_disk_size;
+}
+
 size_t TabletPeer::GetNumLogSegments() const {
   auto log = log_atomic_.load(std::memory_order_acquire);
   return log ? log->num_segments() : 0;

src/yb/tablet/tablet_peer.h

@@ -84,7 +84,14 @@ struct TabletOnDiskSizeInfo {
   int64_t wal_files_disk_size = 0;
   int64_t sst_files_disk_size = 0;
   int64_t uncompressed_sst_files_disk_size = 0;
-  int64_t sum_on_disk_size = 0;
+
+  // Sum of consensus metadata, WALs, and SSTs. Excludes snapshots, retryable requests, MANIFEST,
+  // and other files in those directories. This is always up-to-date.
+  int64_t active_on_disk_size = 0;
+
+  // Estimated size of the tablet on disk, including snapshots, retryable requests, MANIFEST,
+  // and other files in those directories. This is updated periodically, so may be stale or 0.
+  int64_t total_on_disk_size = 0;
 
   template <class PB>
   static TabletOnDiskSizeInfo FromPB(const PB& pb) {
@@ -93,7 +100,8 @@ struct TabletOnDiskSizeInfo {
       .wal_files_disk_size = pb.wal_files_disk_size(),
       .sst_files_disk_size = pb.sst_files_disk_size(),
       .uncompressed_sst_files_disk_size = pb.uncompressed_sst_files_disk_size(),
-      .sum_on_disk_size = pb.estimated_on_disk_size()
+      .active_on_disk_size = pb.active_on_disk_size(),
+      .total_on_disk_size = pb.total_on_disk_size(),
     };
   }
 
@@ -103,19 +111,21 @@ struct TabletOnDiskSizeInfo {
     pb->set_wal_files_disk_size(wal_files_disk_size);
     pb->set_sst_files_disk_size(sst_files_disk_size);
     pb->set_uncompressed_sst_files_disk_size(uncompressed_sst_files_disk_size);
-    pb->set_estimated_on_disk_size(sum_on_disk_size);
+    pb->set_active_on_disk_size(active_on_disk_size);
+    pb->set_total_on_disk_size(total_on_disk_size);
   }
 
   void operator+=(const TabletOnDiskSizeInfo& other) {
     consensus_metadata_disk_size += other.consensus_metadata_disk_size;
     wal_files_disk_size += other.wal_files_disk_size;
     sst_files_disk_size += other.sst_files_disk_size;
     uncompressed_sst_files_disk_size += other.uncompressed_sst_files_disk_size;
-    sum_on_disk_size += other.sum_on_disk_size;
+    active_on_disk_size += other.active_on_disk_size;
+    total_on_disk_size += other.total_on_disk_size;
   }
 
   void RecomputeTotalSize() {
-    sum_on_disk_size =
+    active_on_disk_size =
         consensus_metadata_disk_size +
         sst_files_disk_size +
         wal_files_disk_size;
@@ -495,6 +505,10 @@ class TabletPeer : public std::enable_shared_from_this<TabletPeer>,
 
   bool HasSufficientDiskSpaceForWrite();
 
+  TabletOnDiskSizeInfo GetOnDiskSizeInfo() const EXCLUDES(lock_);
+
+  void SetTabletOnDiskSize(size_t total_on_disk_size);
+
  protected:
   friend class RefCountedThreadSafe<TabletPeer>;
   friend class TabletPeerTest;
@@ -598,7 +612,7 @@ class TabletPeer : public std::enable_shared_from_this<TabletPeer>,
   Status CheckOperationAllowed(
       const OpId& op_id, consensus::OperationType op_type) override;
   // Return granular types of on-disk size of this tablet replica, in bytes.
-  TabletOnDiskSizeInfo GetOnDiskSizeInfo() const REQUIRES(lock_);
+  TabletOnDiskSizeInfo GetOnDiskSizeInfoUnlocked() const REQUIRES(lock_);
 
   bool FlushBootstrapStateEnabled() const;
 
@@ -625,6 +639,11 @@ class TabletPeer : public std::enable_shared_from_this<TabletPeer>,
   std::atomic<bool> flush_bootstrap_state_enabled_{false};
   std::shared_ptr<TabletBootstrapStateFlusher> bootstrap_state_flusher_;
 
+  // The size of the tablet on disk in bytes, including snapshots, retryable requests, MANIFEST,
+  // and other files in those directories. This can be stale as it is only updated every
+  // FLAGS_data_size_metric_updater_interval_sec seconds.
+  std::atomic<size_t> total_on_disk_size_{0};
+
   DISALLOW_COPY_AND_ASSIGN(TabletPeer);
 };