Resize the shadow disk when the size of the source disk changes

ydb-platform · Sep 27, 2024 · cebf5f9 · cebf5f9
1 parent 9c55e94
commit cebf5f9
Show file tree

Hide file tree

Showing 3 changed files with 135 additions and 3 deletions.
diff --git a/cloud/blockstore/libs/storage/disk_registry/disk_registry_state.cpp b/cloud/blockstore/libs/storage/disk_registry/disk_registry_state.cpp
@@ -2010,6 +2010,8 @@ NProto::TError TDiskRegistryState::AllocateDisk(
     auto& disk = Disks[params.DiskId];
     auto error = ValidateAllocateDiskParams(disk, params);
 
+    auto originalDiskSize = GetDiskSize(params.DiskId);
+
     if (HasError(error)) {
         if (disk.Devices.empty() && !disk.ReplicaCount) {
             Disks.erase(params.DiskId);
@@ -2019,11 +2021,19 @@ NProto::TError TDiskRegistryState::AllocateDisk(
         return error;
     }
 
-    if (params.ReplicaCount) {
-        return AllocateMirroredDisk(now, db, params, disk, result);
+    NProto::TError allocateError =
+        params.ReplicaCount
+            ? AllocateMirroredDisk(now, db, params, disk, result)
+            : AllocateSimpleDisk(now, db, params, {}, disk, result);
+
+    // If disk size changed reallocate checkpoints too.
+    if (!HasError(allocateError) && originalDiskSize &&
+        originalDiskSize != GetDiskSize(params.DiskId))
+    {
+        ReallocateCheckpointByDisk(now, db, params.DiskId);
     }
 
-    return AllocateSimpleDisk(now, db, params, {}, disk, result);
+    return allocateError;
 }
 
 NProto::TError TDiskRegistryState::AllocateCheckpoint(
@@ -2818,6 +2828,24 @@ void TDiskRegistryState::DeleteCheckpointByDisk(
     }
 }
 
+void TDiskRegistryState::ReallocateCheckpointByDisk(
+    TInstant now,
+    TDiskRegistryDatabase& db,
+    const TDiskId& diskId)
+{
+    TVector<TCheckpointId> checkpointsToReallocate;
+    for (const auto& [id, checkpointInfo]: Checkpoints) {
+        if (checkpointInfo.SourceDiskId == diskId) {
+            checkpointsToReallocate.push_back(checkpointInfo.CheckpointId);
+        }
+    }
+
+    for (const auto& checkpointId: checkpointsToReallocate) {
+        TAllocateCheckpointResult result;
+        AllocateCheckpoint(now, db, diskId, checkpointId, &result);
+    }
+}
+
 void TDiskRegistryState::AddToBrokenDisks(
     TInstant now,
     TDiskRegistryDatabase& db,
@@ -7240,4 +7268,14 @@ TVector<NProto::TAgentInfo> TDiskRegistryState::QueryAgentsInfo() const
     return ret;
 }
 
+std::optional<ui64> TDiskRegistryState::GetDiskSize(const TDiskId& diskId) const
+{
+    TDiskInfo diskInfo;
+    auto error = GetDiskInfo(diskId, diskInfo);
+    if (HasError(error)) {
+        return std::nullopt;
+    }
+    return diskInfo.GetBlocksCount();
+}
+
 }   // namespace NCloud::NBlockStore::NStorage
diff --git a/cloud/blockstore/libs/storage/disk_registry/disk_registry_state.h b/cloud/blockstore/libs/storage/disk_registry/disk_registry_state.h
@@ -1125,6 +1125,10 @@ class TDiskRegistryState
     void DeleteCheckpointByDisk(
         TDiskRegistryDatabase& db,
         const TDiskId& diskId);
+    void ReallocateCheckpointByDisk(
+        TInstant now,
+        TDiskRegistryDatabase& db,
+        const TDiskId& diskId);
 
     void ForgetDevices(
         TDiskRegistryDatabase& db,
@@ -1258,6 +1262,8 @@ class TDiskRegistryState
     struct TConfigUpdateEffect;
     TResultOrError<TConfigUpdateEffect> CalcConfigUpdateEffect(
         const NProto::TDiskRegistryConfig& newConfig) const;
+
+    std::optional<ui64> GetDiskSize(const TDiskId& diskId) const;
 };
 
 }   // namespace NCloud::NBlockStore::NStorage
diff --git a/cloud/blockstore/libs/storage/disk_registry/disk_registry_ut_checkpoint.cpp b/cloud/blockstore/libs/storage/disk_registry/disk_registry_ut_checkpoint.cpp
@@ -413,6 +413,94 @@ Y_UNIT_TEST_SUITE(TDiskRegistryTest)
         }
     }
 
+    Y_UNIT_TEST(ShouldReallocateShadowDiskOnSourceDiskReallocation)
+    {
+        const auto agent1 = CreateAgentConfig(
+            "agent-1",
+            {
+                Device("dev-1", "uuid-1", "rack-1", 10_GB),
+                Device("dev-2", "uuid-2", "rack-1", 10_GB),
+            });
+
+        const auto agent2 = CreateAgentConfig(
+            "agent-2",
+            {
+                Device("dev-1", "uuid-4", "rack-1", 10_GB),
+                Device("dev-2", "uuid-5", "rack-1", 10_GB),
+            });
+
+        auto runtime =
+            TTestRuntimeBuilder().WithAgents({agent1, agent2}).Build();
+
+        TDiskRegistryClient diskRegistry(*runtime);
+        diskRegistry.WaitReady();
+        diskRegistry.SetWritableState(true);
+
+        diskRegistry.UpdateConfig(CreateRegistryConfig(0, {agent1, agent2}));
+
+        RegisterAgents(*runtime, 2);
+        WaitForAgents(*runtime, 2);
+        WaitForSecureErase(*runtime, {agent1, agent2});
+
+        // Allocate disk.
+        {
+            auto response = diskRegistry.AllocateDisk("disk-1", 10_GB);
+            UNIT_ASSERT_VALUES_EQUAL_C(
+                S_OK,
+                response->GetStatus(),
+                response->GetError());
+        }
+
+        TString shadowDiskId;
+        // Create checkpoint
+        {
+            auto response =
+                diskRegistry.AllocateCheckpoint("disk-1", "checkpoint-1");
+
+            UNIT_ASSERT_VALUES_EQUAL_C(
+                S_OK,
+                response->GetStatus(),
+                response->GetError());
+
+            shadowDiskId = response->Record.GetShadowDiskId();
+        }
+
+        // Check the size of the shadow disk
+        {
+            auto response = diskRegistry.DescribeDisk(shadowDiskId);
+            UNIT_ASSERT_VALUES_EQUAL_C(
+                S_OK,
+                response->GetStatus(),
+                response->GetError());
+
+            auto checkpointSize = response->Record.GetBlocksCount() *
+                                  response->Record.GetBlockSize();
+            UNIT_ASSERT_VALUES_EQUAL(10_GB, checkpointSize);
+        }
+
+        // Reallocate disk
+        {
+            auto response = diskRegistry.AllocateDisk("disk-1", 20_GB);
+            UNIT_ASSERT_VALUES_EQUAL_C(
+                S_OK,
+                response->GetStatus(),
+                response->GetError());
+        }
+
+        // Check that the size of the shadow disk has also changed
+        {
+            auto response = diskRegistry.DescribeDisk(shadowDiskId);
+            UNIT_ASSERT_VALUES_EQUAL_C(
+                S_OK,
+                response->GetStatus(),
+                response->GetError());
+
+            auto checkpointSize = response->Record.GetBlocksCount() *
+                                  response->Record.GetBlockSize();
+            UNIT_ASSERT_VALUES_EQUAL(20_GB, checkpointSize);
+        }
+    }
+
     Y_UNIT_TEST(ShouldChangeStateForCheckpoint)
     {
         const auto agent1 = CreateAgentConfig(