fix(dataprotection): protect the cluster restore lifecycle

cmd/dataprotection/main.go

@@ -340,6 +340,14 @@ func main() {
 		os.Exit(1)
 	}
 
+	if err = (&dpcontrollers.ClusterRestoreReconciler{
+		Client:   mgr.GetClient(),
+		Recorder: mgr.GetEventRecorderFor("cluster-restore-controller"),
+	}).SetupWithManager(mgr); err != nil {
+		setupLog.Error(err, "unable to create controller", "controller", "ClusterRestore")
+		os.Exit(1)
+	}
+
 	if err = (&dpcontrollers.BackupScheduleReconciler{
 		Client:   dputils.NewCompatClient(mgr.GetClient()),
 		Scheme:   mgr.GetScheme(),

controllers/dataprotection/cluster_restore_controller.go

@@ -0,0 +1,198 @@
+/*
+Copyright (C) 2022-2026 ApeCloud Co., Ltd
+
+This file is part of KubeBlocks project
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+package dataprotection
+
+import (
+	"context"
+
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/meta"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/tools/record"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
+	"sigs.k8s.io/controller-runtime/pkg/handler"
+	"sigs.k8s.io/controller-runtime/pkg/log"
+	"sigs.k8s.io/controller-runtime/pkg/reconcile"
+
+	appsv1 "github.com/apecloud/kubeblocks/apis/apps/v1"
+	dpv1alpha1 "github.com/apecloud/kubeblocks/apis/dataprotection/v1alpha1"
+	"github.com/apecloud/kubeblocks/pkg/constant"
+	intctrlutil "github.com/apecloud/kubeblocks/pkg/controllerutil"
+	dprestore "github.com/apecloud/kubeblocks/pkg/dataprotection/restore"
+	dptypes "github.com/apecloud/kubeblocks/pkg/dataprotection/types"
+)
+
+// ClusterRestoreReconciler coordinates the Cluster-level restore lifecycle.
+type ClusterRestoreReconciler struct {
+	client.Client
+	Recorder record.EventRecorder
+}
+
+// +kubebuilder:rbac:groups=apps.kubeblocks.io,resources=clusters,verbs=get;list;watch;patch;update
+// +kubebuilder:rbac:groups=apps.kubeblocks.io,resources=clusters/finalizers,verbs=update;patch
+// +kubebuilder:rbac:groups=dataprotection.kubeblocks.io,resources=restores,verbs=get;list;watch
+// +kubebuilder:rbac:groups=core,resources=persistentvolumeclaims,verbs=get;list;watch
+
+func (r *ClusterRestoreReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
+	reqCtx := intctrlutil.RequestCtx{
+		Ctx: ctx, Req: req,
+		Log:      log.FromContext(ctx).WithValues("cluster-restore", req.NamespacedName),
+		Recorder: r.Recorder,
+	}
+	cluster := &appsv1.Cluster{}
+	if err := r.Client.Get(ctx, req.NamespacedName, cluster); err != nil {
+		return intctrlutil.CheckedRequeueWithError(err, reqCtx.Log, "")
+	}
+
+	restoring, err := r.isClusterRestoring(ctx, cluster)
+	if err != nil {
+		return intctrlutil.CheckedRequeueWithError(err, reqCtx.Log, "failed to determine Cluster restore state")
+	}
+	if !restoring {
+		return r.releaseClusterRestoreProtection(reqCtx, cluster)
+	}
+	if !isClusterRestoreProtected(cluster) {
+		// If deletion began without our finalizer, this controller cannot acquire
+		// protection and must not delay Cluster deletion.
+		if !cluster.DeletionTimestamp.IsZero() {
+			return intctrlutil.Reconciled()
+		}
+		return r.protectClusterRestore(reqCtx, cluster)
+	}
+	if !cluster.DeletionTimestamp.IsZero() {
+		return intctrlutil.RequeueAfter(reconcileInterval, reqCtx.Log,
+			"waiting for restore owners to finish Cluster termination")
+	}
+	return intctrlutil.Reconciled()
+}
+
+func (r *ClusterRestoreReconciler) SetupWithManager(mgr ctrl.Manager) error {
+	return intctrlutil.NewControllerManagedBy(mgr).
+		For(&appsv1.Cluster{}).
+		Watches(&dpv1alpha1.Restore{}, handler.EnqueueRequestsFromMapFunc(r.mapObjectToCluster)).
+		Watches(&corev1.PersistentVolumeClaim{}, handler.EnqueueRequestsFromMapFunc(r.mapObjectToCluster)).
+		Complete(r)
+}
+
+func (r *ClusterRestoreReconciler) mapObjectToCluster(_ context.Context, obj client.Object) []reconcile.Request {
+	clusterName := obj.GetLabels()[constant.AppInstanceLabelKey]
+	if clusterName == "" {
+		return nil
+	}
+	return []reconcile.Request{{NamespacedName: client.ObjectKey{Namespace: obj.GetNamespace(), Name: clusterName}}}
+}
+
+func (r *ClusterRestoreReconciler) isClusterRestoring(ctx context.Context,
+	cluster *appsv1.Cluster) (bool, error) {
+	if cluster.DeletionTimestamp.IsZero() && clusterAllowsRestoreProgress(cluster) {
+		return true, nil
+	}
+
+	// Inspect PVCs before Restores. VP releases temporary target protection only
+	// after observing postReady Restore, so this order closes the handoff window.
+	pvcs := &corev1.PersistentVolumeClaimList{}
+	if err := r.Client.List(ctx, pvcs, client.InNamespace(cluster.Namespace), client.MatchingLabels{
+		constant.AppInstanceLabelKey: cluster.Name,
+	}); err != nil {
+		return false, err
+	}
+	for i := range pvcs.Items {
+		pvc := &pvcs.Items[i]
+		if pvc.Labels[dptypes.ClusterUIDLabelKey] != string(cluster.UID) {
+			continue
+		}
+		if isClusterRestoreHelperPVC(pvc) ||
+			(isClusterRestoreTargetPVC(pvc) && controllerutil.ContainsFinalizer(pvc, dptypes.DataProtectionFinalizerName)) {
+			return true, nil
+		}
+	}
+
+	restores := &dpv1alpha1.RestoreList{}
+	if err := r.Client.List(ctx, restores, client.InNamespace(cluster.Namespace), client.MatchingLabels{
+		constant.AppInstanceLabelKey: cluster.Name,
+	}); err != nil {
+		return false, err
+	}
+	for i := range restores.Items {
+		restore := &restores.Items[i]
+		if restore.Labels[dprestore.DataProtectionRestoreLabelKey] != restore.Name {
+			continue
+		}
+		owned := restore.Labels[dptypes.ClusterUIDLabelKey] == string(cluster.UID)
+		terminal := restore.Status.Phase == dpv1alpha1.RestorePhaseCompleted ||
+			restore.Status.Phase == dpv1alpha1.RestorePhaseFailed
+		if owned && (!cluster.DeletionTimestamp.IsZero() || !terminal || !restore.DeletionTimestamp.IsZero()) {
+			return true, nil
+		}
+	}
+	return false, nil
+}
+
+func clusterAllowsRestoreProgress(cluster *appsv1.Cluster) bool {
+	if cluster.Spec.Restore == nil {
+		return false
+	}
+	condition := meta.FindStatusCondition(cluster.Status.Conditions, appsv1.ConditionTypeRestore)
+	// Restore=False is terminal for status aggregation, but the Cluster still
+	// has initial-restore intent. Keep the lifecycle active until deletion so
+	// PVC restores cannot lose protection while converging on the failure.
+	return condition == nil || condition.Status != metav1.ConditionTrue
+}
+
+func isClusterRestoreProtected(cluster *appsv1.Cluster) bool {
+	return controllerutil.ContainsFinalizer(cluster, dptypes.RestoreProtectionFinalizerName)
+}
+
+func (r *ClusterRestoreReconciler) protectClusterRestore(reqCtx intctrlutil.RequestCtx,
+	cluster *appsv1.Cluster) (ctrl.Result, error) {
+	if isClusterRestoreProtected(cluster) {
+		return intctrlutil.Reconciled()
+	}
+	patch := client.MergeFromWithOptions(cluster.DeepCopy(), client.MergeFromWithOptimisticLock{})
+	controllerutil.AddFinalizer(cluster, dptypes.RestoreProtectionFinalizerName)
+	if err := r.Client.Patch(reqCtx.Ctx, cluster, patch); err != nil {
+		return intctrlutil.CheckedRequeueWithError(err, reqCtx.Log, "failed to add Cluster restore-protection finalizer")
+	}
+	return intctrlutil.Reconciled()
+}
+
+func (r *ClusterRestoreReconciler) releaseClusterRestoreProtection(reqCtx intctrlutil.RequestCtx,
+	cluster *appsv1.Cluster) (ctrl.Result, error) {
+	if !isClusterRestoreProtected(cluster) {
+		return intctrlutil.Reconciled()
+	}
+	patch := client.MergeFromWithOptions(cluster.DeepCopy(), client.MergeFromWithOptimisticLock{})
+	controllerutil.RemoveFinalizer(cluster, dptypes.RestoreProtectionFinalizerName)
+	if err := r.Client.Patch(reqCtx.Ctx, cluster, patch); err != nil {
+		return intctrlutil.CheckedRequeueWithError(err, reqCtx.Log, "failed to remove Cluster restore-protection finalizer")
+	}
+	return intctrlutil.Reconciled()
+}
+
+func isClusterRestoreHelperPVC(pvc *corev1.PersistentVolumeClaim) bool {
+	return pvc.Labels[dprestore.DataProtectionPopulatePVCLabelKey] != ""
+}
+
+func isClusterRestoreTargetPVC(pvc *corev1.PersistentVolumeClaim) bool {
+	return pvc.Spec.DataSourceRef != nil && pvc.Spec.DataSourceRef.APIGroup != nil &&
+		*pvc.Spec.DataSourceRef.APIGroup == dptypes.DataprotectionAPIGroup
+}

controllers/dataprotection/cluster_restore_controller_test.go

@@ -0,0 +1,149 @@
+/*
+Copyright (C) 2022-2026 ApeCloud Co., Ltd
+
+This file is part of KubeBlocks project
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+package dataprotection
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
+
+	appsv1 "github.com/apecloud/kubeblocks/apis/apps/v1"
+	dpv1alpha1 "github.com/apecloud/kubeblocks/apis/dataprotection/v1alpha1"
+	dptypes "github.com/apecloud/kubeblocks/pkg/dataprotection/types"
+)
+
+func TestClusterRestoreProtectionLifecycle(t *testing.T) {
+	for _, tc := range []struct {
+		name      string
+		deleting  bool
+		protected bool
+		status    metav1.ConditionStatus
+		resource  string
+		wantKeep  bool
+		wantWait  bool
+	}{
+		{"initial intent", false, false, metav1.ConditionUnknown, "", true, false},
+		{"deletion without protection", true, false, metav1.ConditionUnknown, "target", false, false},
+		{"failed restore", false, true, metav1.ConditionFalse, "", true, false},
+		{"successful restore", false, true, metav1.ConditionTrue, "completed restore", false, false},
+		{"target still protected", false, true, metav1.ConditionTrue, "target", true, false},
+		{"helper remains", false, true, metav1.ConditionTrue, "helper", true, false},
+		{"execution still running", false, true, metav1.ConditionTrue, "running restore", true, false},
+		{"deletion waits for target", true, true, metav1.ConditionTrue, "target", true, true},
+		{"deletion waits for helper", true, true, metav1.ConditionTrue, "helper", true, true},
+		{"deletion waits for failed restore", true, true, metav1.ConditionTrue, "failed restore", true, true},
+		{"deletion waits for completed restore", true, true, metav1.ConditionTrue, "completed restore", true, true},
+		{"deletion cleanup finished", true, true, metav1.ConditionTrue, "", false, false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := context.Background()
+			scheme, cluster, _, _, target := parentRestoreObjects(t)
+			cluster.Finalizers = []string{"example.io/app-owner"}
+			if tc.protected {
+				cluster.Finalizers = append(cluster.Finalizers, dptypes.RestoreProtectionFinalizerName)
+			}
+			if tc.deleting {
+				now := metav1.Now()
+				cluster.DeletionTimestamp = &now
+			}
+			cluster.Status.Conditions = []metav1.Condition{{
+				Type: appsv1.ConditionTypeRestore, Status: tc.status,
+			}}
+			var resource client.Object
+			switch tc.resource {
+			case "target":
+				target.Finalizers = []string{dptypes.DataProtectionFinalizerName}
+				resource = target
+			case "helper":
+				resource = restoreHelperForTarget(target, cluster)
+			case "running restore", "failed restore", "completed restore":
+				restore := executionRestoreForTarget(target, cluster)
+				switch tc.resource {
+				case "failed restore":
+					restore.Status.Phase = dpv1alpha1.RestorePhaseFailed
+				case "completed restore":
+					restore.Status.Phase = dpv1alpha1.RestorePhaseCompleted
+				default:
+					restore.Status.Phase = dpv1alpha1.RestorePhaseRunning
+				}
+				restore.Finalizers = []string{dptypes.DataProtectionFinalizerName}
+				resource = restore
+			}
+			objects := []client.Object{cluster}
+			if resource != nil {
+				objects = append(objects, resource)
+			}
+			cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(objects...).Build()
+			if resource != nil {
+				require.NoError(t, cli.Get(ctx, client.ObjectKeyFromObject(resource), resource))
+			}
+			reconciler := &ClusterRestoreReconciler{Client: cli}
+
+			result, err := reconciler.Reconcile(ctx, ctrl.Request{NamespacedName: client.ObjectKeyFromObject(cluster)})
+			require.NoError(t, err)
+			require.Equal(t, tc.wantWait, result.RequeueAfter > 0)
+			require.NoError(t, cli.Get(ctx, client.ObjectKeyFromObject(cluster), cluster))
+			expected := []string{"example.io/app-owner"}
+			if tc.wantKeep {
+				expected = append(expected, dptypes.RestoreProtectionFinalizerName)
+			}
+			require.ElementsMatch(t, expected, cluster.Finalizers)
+			if resource != nil {
+				current := resource.DeepCopyObject().(client.Object)
+				require.NoError(t, cli.Get(ctx, client.ObjectKeyFromObject(resource), current))
+				require.Equal(t, resource, current)
+			}
+		})
+	}
+}
+
+func TestClusterRestoreControllerIgnoresResourcesWithoutExactClusterUID(t *testing.T) {
+	for _, uid := range []string{"", "another-cluster-uid"} {
+		t.Run("uid="+uid, func(t *testing.T) {
+			ctx := context.Background()
+			scheme, cluster, _, _, target := parentRestoreObjects(t)
+			now := metav1.Now()
+			cluster.DeletionTimestamp = &now
+			cluster.Finalizers = append(cluster.Finalizers, "example.io/app-owner")
+			target.Finalizers = []string{dptypes.DataProtectionFinalizerName}
+			helper := restoreHelperForTarget(target, cluster)
+			restore := executionRestoreForTarget(target, cluster)
+			for _, obj := range []client.Object{target, helper, restore} {
+				obj.GetLabels()[dptypes.ClusterUIDLabelKey] = uid
+			}
+			cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster, target, helper, restore).Build()
+			reconciler := &ClusterRestoreReconciler{Client: cli}
+
+			_, err := reconciler.Reconcile(ctx, ctrl.Request{NamespacedName: client.ObjectKeyFromObject(cluster)})
+			require.NoError(t, err)
+			require.NoError(t, cli.Get(ctx, client.ObjectKeyFromObject(cluster), cluster))
+			require.Equal(t, []string{"example.io/app-owner"}, cluster.Finalizers)
+			require.NoError(t, cli.Get(ctx, client.ObjectKeyFromObject(target), &corev1.PersistentVolumeClaim{}))
+			require.NoError(t, cli.Get(ctx, client.ObjectKeyFromObject(helper), &corev1.PersistentVolumeClaim{}))
+			require.NoError(t, cli.Get(ctx, client.ObjectKeyFromObject(restore), &dpv1alpha1.Restore{}))
+		})
+	}
+}

controllers/dataprotection/suite_test.go

@@ -187,6 +187,12 @@ var _ = BeforeSuite(func() {
 	}).SetupWithManager(k8sManager)
 	Expect(err).ToNot(HaveOccurred())
 
+	err = (&ClusterRestoreReconciler{
+		Client:   k8sClient,
+		Recorder: k8sManager.GetEventRecorderFor("cluster-restore-controller"),
+	}).SetupWithManager(k8sManager)
+	Expect(err).ToNot(HaveOccurred())
+
 	err = (&BackupScheduleReconciler{
 		Client:   k8sClient,
 		Scheme:   k8sManager.GetScheme(),