CO-3733 Support timeout for EvictPod (#191)

Author: pijusn-cast

internal/actions/evict_pod_handler.go

@@ -49,11 +49,17 @@ func (h *EvictPodHandler) Handle(ctx context.Context, action *castai.ClusterActi
 		"namespace":      req.Namespace,
 		"pod":            req.PodName,
 	})
-	return h.handle(ctx, log, req)
+	return h.handle(ctx, log, action, req)
 }
 
-func (h *EvictPodHandler) handle(ctx context.Context, log logrus.FieldLogger, req *castai.ActionEvictPod) error {
-	log.Infof("evicting pod")
+func (h *EvictPodHandler) handle(ctx context.Context, log logrus.FieldLogger, action *castai.ClusterAction, req *castai.ActionEvictPod) error {
+	deadline, deadlineStr := h.computeDeadline(action, req)
+	if !deadline.IsZero() {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithDeadline(ctx, deadline)
+		defer cancel()
+	}
+	log.WithField("deadline", deadlineStr).Infof("evicting pod")
 	err := h.evictPod(ctx, log, req.Namespace, req.PodName)
 	if err != nil {
 		return fmt.Errorf("evict pod: %w", err)
@@ -160,3 +166,12 @@ func (h *EvictPodHandler) isPodDeleted(ctx context.Context, namespace, name stri
 	}
 	return false, p.Status.Phase, nil
 }
+
+func (h *EvictPodHandler) computeDeadline(action *castai.ClusterAction, req *castai.ActionEvictPod) (time.Time, string) {
+	if req.TimeoutSeconds > 0 {
+		deadline := action.CreatedAt.Add(time.Duration(req.TimeoutSeconds) * time.Second).UTC()
+		deadlineStr := deadline.Format(time.RFC3339)
+		return deadline, deadlineStr
+	}
+	return time.Time{}, "-"
+}

internal/actions/evict_pod_handler_test.go

@@ -42,15 +42,39 @@ func TestEvictPodHandler(t *testing.T) {
 		)
 
 		action := newEvictPodAction(&castai.ActionEvictPod{
-			Namespace: pod.Namespace,
-			PodName:   pod.Name,
+			Namespace:      pod.Namespace,
+			PodName:        pod.Name,
+			TimeoutSeconds: 120,
 		})
 		err := h.Handle(ctx, action)
 		require.NoError(t, err)
 
 		requirePodNotExists(t, ctx, clientset, pod.Namespace, pod.Name)
 	})
 
+	t.Run("stops trying after the timeout", func(t *testing.T) {
+		t.Parallel()
+
+		clientset, pod := setupFakeClientWithPodEviction()
+
+		reactionErr := &apierrors.StatusError{ErrStatus: metav1.Status{Reason: metav1.StatusReasonTooManyRequests}}
+		reaction := failingPodEvictionReaction(reactionErr, 1000, nil)
+		prependPodEvictionReaction(clientset, reaction)
+
+		h := NewEvictPodHandler(
+			log,
+			clientset,
+		)
+
+		action := newEvictPodAction(&castai.ActionEvictPod{
+			Namespace:      pod.Namespace,
+			PodName:        pod.Name,
+			TimeoutSeconds: 1,
+		})
+		err := h.Handle(ctx, action)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+
 	t.Run("retries on TooManyRequests status", func(t *testing.T) {
 		t.Parallel()
 

internal/castai/types.go

@@ -160,8 +160,9 @@ type ActionDrainNode struct {
 }
 
 type ActionEvictPod struct {
-	Namespace string `json:"namespace"`
-	PodName   string `json:"podName"`
+	Namespace      string `json:"namespace"`
+	PodName        string `json:"podName"`
+	TimeoutSeconds int    `json:"timeoutSeconds"`
 }
 
 type ActionPatchNode struct {