Add attr WaitForReady in ExecRestoreHook

Signed-off-by: Ripolin <[email protected]>
vmware-tanzu · Oct 6, 2023 · eb1ff04 · eb1ff04
1 parent 668e516
commit eb1ff04
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 6 deletions.
diff --git a/internal/hook/wait_exec_hook_handler.go b/internal/hook/wait_exec_hook_handler.go
@@ -29,6 +29,7 @@ import (
 
 	velerov1api "github.com/vmware-tanzu/velero/pkg/apis/velero/v1"
 	"github.com/vmware-tanzu/velero/pkg/podexec"
+	"github.com/vmware-tanzu/velero/pkg/util/boolptr"
 	"github.com/vmware-tanzu/velero/pkg/util/kube"
 )
 
@@ -126,8 +127,8 @@ func (e *DefaultWaitExecHookHandler) HandleHooks(
 		}
 
 		for containerName, hooks := range byContainer {
-			if !isContainerReady(newPod, containerName) {
-				podLog.Infof("Container %s is not ready: post-restore hooks will not yet be executed", containerName)
+			if !isContainerUp(newPod, containerName, hooks) {
+				podLog.Infof("Container %s is not ready or running: post-restore hooks will not yet be executed", containerName)
 				continue
 			}
 			podMap, err := runtime.DefaultUnstructuredConverter.ToUnstructured(newPod)
@@ -243,15 +244,25 @@ func podHasContainer(pod *v1.Pod, containerName string) bool {
 	return false
 }
 
-func isContainerReady(pod *v1.Pod, containerName string) bool {
+func isContainerUp(pod *v1.Pod, containerName string, hooks []PodExecRestoreHook) bool {
 	if pod == nil {
 		return false
 	}
+	var waitForReady bool
+	for _, hook := range hooks {
+		if boolptr.IsSetToTrue(hook.Hook.WaitForReady) {
+			waitForReady = true
+			break
+		}
+	}
 	for _, cs := range pod.Status.ContainerStatuses {
 		if cs.Name != containerName {
 			continue
 		}
-		return cs.Ready
+		if waitForReady {
+			return cs.Ready
+		}
+		return cs.State.Running != nil
 	}
 
 	return false

diff --git a/internal/hook/wait_exec_hook_handler_test.go b/internal/hook/wait_exec_hook_handler_test.go
@@ -35,6 +35,7 @@ import (
 	velerov1api "github.com/vmware-tanzu/velero/pkg/apis/velero/v1"
 	"github.com/vmware-tanzu/velero/pkg/builder"
 	velerotest "github.com/vmware-tanzu/velero/pkg/test"
+	"github.com/vmware-tanzu/velero/pkg/util/boolptr"
 )
 
 type fakeListWatchFactory struct {
@@ -790,12 +791,13 @@ func TestPodHasContainer(t *testing.T) {
 	}
 }
 
-func TestIsContainerReady(t *testing.T) {
+func TestIsContainerUp(t *testing.T) {
 	tests := []struct {
 		name      string
 		pod       *v1.Pod
 		container string
 		expect    bool
+		hooks     []PodExecRestoreHook
 	}{
 		{
 			name:      "should return true when running",
@@ -809,6 +811,49 @@ func TestIsContainerReady(t *testing.T) {
 					},
 				}).
 				Result(),
+			hooks: []PodExecRestoreHook{},
+		},
+		{
+			name:      "should return false when running but not ready",
+			container: "container1",
+			expect:    false,
+			pod: builder.ForPod("default", "my-pod").
+				ContainerStatuses(&v1.ContainerStatus{
+					Name: "container1",
+					State: v1.ContainerState{
+						Running: &v1.ContainerStateRunning{},
+					},
+					Ready: false,
+				}).
+				Result(),
+			hooks: []PodExecRestoreHook{
+				{
+					Hook: velerov1api.ExecRestoreHook{
+						WaitForReady: boolptr.True(),
+					},
+				},
+			},
+		},
+		{
+			name:      "should return true when running and ready",
+			container: "container1",
+			expect:    true,
+			pod: builder.ForPod("default", "my-pod").
+				ContainerStatuses(&v1.ContainerStatus{
+					Name: "container1",
+					State: v1.ContainerState{
+						Running: &v1.ContainerStateRunning{},
+					},
+					Ready: true,
+				}).
+				Result(),
+			hooks: []PodExecRestoreHook{
+				{
+					Hook: velerov1api.ExecRestoreHook{
+						WaitForReady: boolptr.True(),
+					},
+				},
+			},
 		},
 		{
 			name:      "should return false when no state is set",
@@ -820,6 +865,7 @@ func TestIsContainerReady(t *testing.T) {
 					State: v1.ContainerState{},
 				}).
 				Result(),
+			hooks: []PodExecRestoreHook{},
 		},
 		{
 			name:      "should return false when waiting",
@@ -833,6 +879,7 @@ func TestIsContainerReady(t *testing.T) {
 					},
 				}).
 				Result(),
+			hooks: []PodExecRestoreHook{},
 		},
 		{
 			name:      "should return true when running and first container is terminated",
@@ -852,11 +899,12 @@ func TestIsContainerReady(t *testing.T) {
 						},
 					}).
 				Result(),
+			hooks: []PodExecRestoreHook{},
 		},
 	}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			actual := isContainerReady(test.pod, test.container)
+			actual := isContainerUp(test.pod, test.container, test.hooks)
 			assert.Equal(t, actual, test.expect)
 		})
 	}

diff --git a/pkg/apis/velero/v1/restore_types.go b/pkg/apis/velero/v1/restore_types.go
@@ -214,6 +214,11 @@ type ExecRestoreHook struct {
 	// before attempting to run the command.
 	// +optional
 	WaitTimeout metav1.Duration `json:"waitTimeout,omitempty"`
+
+	// WaitForReady ensure command will be launch when container is Ready instead of Running.
+	// +optional
+	// +nullable
+	WaitForReady *bool `json:"waitForReady,omitempty"`
 }
 
 // InitRestoreHook is a hook that adds an init container to a PodSpec to run commands before the