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

@@ -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
 }
 

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 {

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
+}

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)