From 6a4a38869cbaee333d85c9dc4f8de9b5b29c09ba Mon Sep 17 00:00:00 2001
From: Rakshith R <rar@redhat.com>
Date: Wed, 13 Nov 2024 17:29:33 +0530
Subject: [PATCH] rbd: consolidate snapshot flatten logic in
 PrepareVolumeForSnapshot()

This commit consolidates flatten logic checks for cloneDepth
and snapshotLimit in PrepareVolumeForSnapshot. This allows
the function to be called for both CreateSnapshot and
CreateVolumeGroupSnapshot.
Clone Depth check and flattening of grand parent image
now occurs before creation of snapshot starts.
This aligns better with how PVC-PVC clone and
PVC-restore process occurs currently.
Flattening the grandparent image once prevents
flattening of every newly created snapshot.
Snapshot in above para refers to k8s VolumeSnapshot
(which is backed by a rbd image).

Signed-off-by: Rakshith R <rar@redhat.com>
(cherry picked from commit 9936033283f697ce10d6bbb47696929de3797517)
---
 internal/rbd/controllerserver.go       |  7 +------
 internal/rbd/group_controllerserver.go | 23 +++++++++++++++++++++++
 internal/rbd/rbd_util.go               | 25 +++++++++++++++++++++++++
 internal/rbd/types/volume.go           |  4 ++++
 4 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/internal/rbd/controllerserver.go b/internal/rbd/controllerserver.go
index eb37a8063..5ccec05f9 100644
--- a/internal/rbd/controllerserver.go
+++ b/internal/rbd/controllerserver.go
@@ -1157,7 +1157,7 @@ func (cs *ControllerServer) CreateSnapshot(
 		return cloneFromSnapshot(ctx, rbdVol, rbdSnap, cr, req.GetParameters())
 	}
 
-	err = flattenTemporaryClonedImages(ctx, rbdVol, cr)
+	err = rbdVol.PrepareVolumeForSnapshot(ctx, cr)
 	if err != nil {
 		return nil, err
 	}
@@ -1376,11 +1376,6 @@ func (cs *ControllerServer) doSnapshotClone(
 		return cloneRbd, err
 	}
 
-	err = cloneRbd.flattenRbdImage(ctx, false, rbdHardMaxCloneDepth, rbdSoftMaxCloneDepth)
-	if err != nil {
-		return cloneRbd, err
-	}
-
 	return cloneRbd, nil
 }
 
diff --git a/internal/rbd/group_controllerserver.go b/internal/rbd/group_controllerserver.go
index 4e4d3c408..8b447b921 100644
--- a/internal/rbd/group_controllerserver.go
+++ b/internal/rbd/group_controllerserver.go
@@ -37,6 +37,8 @@ import (
 // group snapshot and remove all images from the group again. This leaves the
 // group and its snapshot around, the group snapshot can be inspected to list
 // the snapshots of the images.
+//
+//nolint:gocyclo,cyclop // TODO: reduce complexity.
 func (cs *ControllerServer) CreateVolumeGroupSnapshot(
 	ctx context.Context,
 	req *csi.CreateVolumeGroupSnapshotRequest,
@@ -130,6 +132,27 @@ func (cs *ControllerServer) CreateVolumeGroupSnapshot(
 			"failed to get existing one with name %q: %v", vgsName, err)
 	}
 
+	creds, err := util.NewUserCredentials(req.GetSecrets())
+	if err != nil {
+		return nil, status.Error(codes.Internal, err.Error())
+	}
+	defer creds.DeleteCredentials()
+
+	errList := make([]error, 0)
+	for _, volume := range volumes {
+		err = volume.PrepareVolumeForSnapshot(ctx, creds)
+		if err != nil {
+			errList = append(errList, err)
+		}
+	}
+	if len(errList) > 0 {
+		// FIXME: we should probably choose a error code that has more priority.
+		return nil, status.Errorf(
+			status.Code(errList[0]),
+			"failed to prepare volumes for snapshot: %v",
+			errList)
+	}
+
 	// create a temporary VolumeGroup with a different name
 	vg, err = mgr.CreateVolumeGroup(ctx, vgName)
 	if err != nil {
diff --git a/internal/rbd/rbd_util.go b/internal/rbd/rbd_util.go
index 343d63708..6dc4d139a 100644
--- a/internal/rbd/rbd_util.go
+++ b/internal/rbd/rbd_util.go
@@ -2277,3 +2277,28 @@ func (ri *rbdImage) GetClusterID(ctx context.Context) (string, error) {
 
 	return ri.ClusterID, nil
 }
+
+func (rv *rbdVolume) PrepareVolumeForSnapshot(ctx context.Context, cr *util.Credentials) error {
+	hardLimit := rbdHardMaxCloneDepth
+	softLimit := rbdSoftMaxCloneDepth
+	err := flattenTemporaryClonedImages(ctx, rv, cr)
+	if err != nil {
+		return err
+	}
+
+	// choosing 2, since snapshot adds one depth and we'll be flattening the parent.
+	const depthToAvoidFlatten = 2
+	if rbdHardMaxCloneDepth > depthToAvoidFlatten {
+		hardLimit = rbdHardMaxCloneDepth - depthToAvoidFlatten
+	}
+	if rbdSoftMaxCloneDepth > depthToAvoidFlatten {
+		softLimit = rbdSoftMaxCloneDepth - depthToAvoidFlatten
+	}
+
+	err = rv.flattenParent(ctx, hardLimit, softLimit)
+	if err != nil {
+		return getGRPCErrorForCreateVolume(err)
+	}
+
+	return nil
+}
diff --git a/internal/rbd/types/volume.go b/internal/rbd/types/volume.go
index 1f235bf16..fde0125b1 100644
--- a/internal/rbd/types/volume.go
+++ b/internal/rbd/types/volume.go
@@ -58,6 +58,10 @@ type Volume interface {
 	// if the parent image is in trash, it returns an error.
 	// if the parent image exists and is not enabled for mirroring, it returns an error.
 	HandleParentImageExistence(ctx context.Context, flattenMode FlattenMode) error
+	// PrepareVolumeForSnapshot prepares the volume for snapshot by
+	// checking snapshots limit and clone depth limit and flatten it
+	// if required.
+	PrepareVolumeForSnapshot(ctx context.Context, cr *util.Credentials) error
 
 	// ToMirror converts the Volume to a Mirror.
 	ToMirror() (Mirror, error)