feat(controller): add configurable restart timeout and exclude restarting pods from allocation #452

- Add restart-timeout flag to controller with default 90s for pod restart operations
- Pass restartTimeout value to restartTracker for managing pod lifecycle
- Modify restartTracker to use configurable restartTimeout instead of constant
- Exclude pods in restarting state from allocator’s available pod list
- Add unit test to verify allocator excludes restarting pods during scheduling
- Update e2e test to deploy controller with restart-timeout=10s for timeout testing
- Add setup and teardown steps in e2e test for namespace, CRDs, and controller deployment
- Reduce pod restart timeout wait in e2e test from 4 minutes to 1 minute for faster feedback

Author: fengcone

kubernetes/cmd/controller/main.go

@@ -19,6 +19,7 @@ import (
 	"flag"
 	"os"
 	"path/filepath"
+	"time"
 
 	// Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.)
 	// to ensure that exec-entrypoint and run can make use of them.
@@ -78,6 +79,9 @@ func main() {
 	var kubeClientQPS float64
 	var kubeClientBurst int
 
+	// Restart timeout configuration
+	var restartTimeout time.Duration
+
 	flag.StringVar(&metricsAddr, "metrics-bind-address", "0", "The address the metrics endpoint binds to. "+
 		"Use :8443 for HTTPS or :8080 for HTTP, or leave as 0 to disable the metrics service.")
 	flag.StringVar(&probeAddr, "health-probe-bind-address", ":8081", "The address the probe endpoint binds to.")
@@ -105,6 +109,7 @@ func main() {
 	flag.BoolVar(&logCompress, "log-compress", true, "Compress determines if the rotated log files should be compressed using gzip")
 	flag.Float64Var(&kubeClientQPS, "kube-client-qps", 100, "QPS for Kubernetes client rate limiter.")
 	flag.IntVar(&kubeClientBurst, "kube-client-burst", 200, "Burst for Kubernetes client rate limiter.")
+	flag.DurationVar(&restartTimeout, "restart-timeout", 90*time.Second, "Timeout for Pod restart operations. If a Pod fails to restart within this duration, it will be deleted.")
 
 	opts := zap.Options{}
 	opts.BindFlags(flag.CommandLine)
@@ -261,7 +266,7 @@ func main() {
 		os.Exit(1)
 	}
 	kubeClient := kubernetes.NewForConfigOrDie(mgr.GetConfig())
-	restartTracker := controller.NewRestartTracker(mgr.GetClient(), kubeClient, mgr.GetConfig())
+	restartTracker := controller.NewRestartTracker(mgr.GetClient(), kubeClient, mgr.GetConfig(), restartTimeout)
 	if err := (&controller.PoolReconciler{
 		Client:         mgr.GetClient(),
 		Scheme:         mgr.GetScheme(),

kubernetes/internal/controller/allocator.go

@@ -227,6 +227,9 @@ func (allocator *defaultAllocator) Schedule(ctx context.Context, spec *AllocSpec
 		if pod.Status.Phase != corev1.PodRunning {
 			continue
 		}
+		if isRestarting(pod) {
+			continue
+		}
 		availablePods = append(availablePods, pod.Name)
 	}
 	log.V(1).Info("Schedule init", "existingAllocations", len(status.PodAllocation), "availablePods", len(availablePods))

kubernetes/internal/controller/allocator_test.go

@@ -16,6 +16,7 @@ package controller
 
 import (
 	"context"
+	"encoding/json"
 	"reflect"
 	"testing"
 
@@ -494,3 +495,70 @@ func TestSyncSandboxAllocationError(t *testing.T) {
 	err := allocator.SyncSandboxAllocation(context.Background(), sandbox, pods)
 	assert.Error(t, err)
 }
+
+func TestScheduleExcludesRestartingPods(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	store := NewMockAllocationStore(ctrl)
+	syncer := NewMockAllocationSyncer(ctrl)
+	allocator := &defaultAllocator{
+		store:  store,
+		syncer: syncer,
+	}
+	replica1 := int32(1)
+
+	// Create pods: one normal, one restarting (should be excluded from allocation)
+	restartingMeta := PodRecycleMeta{
+		State:       RecycleStateRestarting,
+		TriggeredAt: 1234567890,
+	}
+	restartingMetaJSON, _ := json.Marshal(restartingMeta)
+
+	pods := []*corev1.Pod{
+		{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: "pod-normal",
+			},
+			Status: corev1.PodStatus{Phase: corev1.PodRunning},
+		},
+		{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: "pod-restarting",
+				Annotations: map[string]string{
+					AnnoPodRecycleMeta: string(restartingMetaJSON),
+				},
+			},
+			Status: corev1.PodStatus{Phase: corev1.PodRunning},
+		},
+	}
+	sandboxes := []*sandboxv1alpha1.BatchSandbox{
+		{
+			ObjectMeta: metav1.ObjectMeta{Name: "sbx1"},
+			Spec:       sandboxv1alpha1.BatchSandboxSpec{Replicas: &replica1},
+		},
+		{
+			ObjectMeta: metav1.ObjectMeta{Name: "sbx2"},
+			Spec:       sandboxv1alpha1.BatchSandboxSpec{Replicas: &replica1},
+		},
+	}
+	spec := &AllocSpec{
+		Pods:      pods,
+		Sandboxes: sandboxes,
+		Pool:      &sandboxv1alpha1.Pool{ObjectMeta: metav1.ObjectMeta{Name: "pool1"}},
+	}
+
+	store.EXPECT().GetAllocation(gomock.Any(), gomock.Any()).Return(&PoolAllocation{PodAllocation: map[string]string{}}, nil).Times(1)
+	syncer.EXPECT().GetAllocation(gomock.Any(), gomock.Any()).Return(&SandboxAllocation{Pods: []string{}}, nil).Times(2)
+	syncer.EXPECT().GetRelease(gomock.Any(), gomock.Any()).Return(&AllocationRelease{Pods: []string{}}, nil).Times(2)
+
+	status, pendingSyncs, poolDirty, err := allocator.Schedule(context.Background(), spec)
+
+	assert.NoError(t, err)
+	assert.True(t, poolDirty)
+	// Only the normal pod should be allocated, sbx2 should have no pod
+	assert.Contains(t, status.PodAllocation, "pod-normal")
+	assert.NotContains(t, status.PodAllocation, "pod-restarting")
+	// sbx2 should need supplement since restarting pod is excluded
+	assert.Equal(t, int32(1), status.PodSupplement)
+	assert.Len(t, pendingSyncs, 1)
+}

kubernetes/internal/controller/restart_tracker.go

@@ -33,8 +33,8 @@ import (
 
 // Restart timeout configurations
 const (
-	restartTimeout = 90 * time.Second
-	killTimeout    = 30 * time.Second
+	defaultRestartTimeout = 90 * time.Second
+	killTimeout           = 30 * time.Second
 )
 
 // restartTracker manages the Pod restart lifecycle as part of the PoolReconciler.
@@ -48,22 +48,28 @@ const (
 //	    all restarted & ready → None (clear annotation, reuse)
 //	    timeout / CrashLoop   → delete Pod
 type restartTracker struct {
-	client     client.Client
-	kubeClient kubernetes.Interface
-	restConfig *rest.Config
+	client         client.Client
+	kubeClient     kubernetes.Interface
+	restConfig     *rest.Config
+	restartTimeout time.Duration
 }
 
 type RestartTracker interface {
 	HandleRestart(ctx context.Context, pod *corev1.Pod) error
 }
 
-// NewRestartTracker creates a new restartTracker.
-func NewRestartTracker(c client.Client, kubeClient kubernetes.Interface, restConfig *rest.Config) RestartTracker {
-	return &restartTracker{
-		client:     c,
-		kubeClient: kubeClient,
-		restConfig: restConfig,
+// NewRestartTracker creates a new restartTracker with custom restart timeout.
+func NewRestartTracker(c client.Client, kubeClient kubernetes.Interface, restConfig *rest.Config, restartTimeout time.Duration) RestartTracker {
+	r := &restartTracker{
+		client:         c,
+		kubeClient:     kubeClient,
+		restConfig:     restConfig,
+		restartTimeout: restartTimeout,
 	}
+	if r.restartTimeout == 0 {
+		r.restartTimeout = defaultRestartTimeout
+	}
+	return r
 }
 
 // HandleRestart handles the Restart recycle policy for a Pod.
@@ -206,9 +212,9 @@ func (t *restartTracker) checkRestartStatus(ctx context.Context, pod *corev1.Pod
 		return t.client.Delete(ctx, pod)
 	}
 
-	if elapsed > restartTimeout {
+	if elapsed > t.restartTimeout {
 		log.Info("Pod restart timeout, deleting", "pod", pod.Name,
-			"elapsed", elapsed, "timeout", restartTimeout,
+			"elapsed", elapsed, "timeout", t.restartTimeout,
 			"allRestarted", allRestarted, "allReady", allReady)
 		return t.client.Delete(ctx, pod)
 	}

kubernetes/internal/controller/restart_tracker_test.go

@@ -134,7 +134,7 @@ func TestRestartTracker_CheckRestartStatus_Timeout(t *testing.T) {
 
 	meta := &PodRecycleMeta{
 		State:       RecycleStateRestarting,
-		TriggeredAt: time.Now().UnixMilli() - (restartTimeout.Milliseconds() + 1000),
+		TriggeredAt: time.Now().UnixMilli() - (defaultRestartTimeout.Milliseconds() + 1000),
 		InitialRestartCounts: map[string]int32{
 			"c1": 0,
 		},
@@ -143,7 +143,8 @@ func TestRestartTracker_CheckRestartStatus_Timeout(t *testing.T) {
 
 	client := fake.NewClientBuilder().WithScheme(scheme).WithObjects(pod).Build()
 	tracker := &restartTracker{
-		client: client,
+		client:         client,
+		restartTimeout: defaultRestartTimeout,
 	}
 
 	err := tracker.checkRestartStatus(context.Background(), pod)
@@ -229,7 +230,8 @@ func TestRestartTracker_CheckRestartStatus_StillRestarting(t *testing.T) {
 
 	client := fake.NewClientBuilder().WithScheme(scheme).WithObjects(pod).Build()
 	tracker := &restartTracker{
-		client: client,
+		client:         client,
+		restartTimeout: defaultRestartTimeout,
 	}
 
 	err := tracker.checkRestartStatus(context.Background(), pod)

kubernetes/test/e2e/pod_recycle_policy_test.go

@@ -35,6 +35,39 @@ var _ = Describe("Pod Recycle Policy", Ordered, func() {
 	const testNamespace = "default"
 
 	BeforeAll(func() {
+		By("creating manager namespace")
+		cmd := exec.Command("kubectl", "create", "ns", namespace)
+		_, err := utils.Run(cmd)
+		Expect(err).NotTo(HaveOccurred(), "Failed to create namespace")
+
+		By("labeling the namespace to enforce the restricted security policy")
+		cmd = exec.Command("kubectl", "label", "--overwrite", "ns", namespace,
+			"pod-security.kubernetes.io/enforce=restricted")
+		_, err = utils.Run(cmd)
+		Expect(err).NotTo(HaveOccurred(), "Failed to label namespace with restricted policy")
+
+		By("installing CRDs")
+		cmd = exec.Command("make", "install")
+		_, err = utils.Run(cmd)
+		Expect(err).NotTo(HaveOccurred(), "Failed to install CRDs")
+
+		By("deploying the controller-manager")
+		cmd = exec.Command("make", "deploy", fmt.Sprintf("CONTROLLER_IMG=%s", utils.ControllerImage))
+		_, err = utils.Run(cmd)
+		Expect(err).NotTo(HaveOccurred(), "Failed to deploy the controller-manager")
+
+		By("patching controller deployment with restart-timeout for testing")
+		cmd = exec.Command("kubectl", "patch", "deployment", "opensandbox-controller-manager", "-n", namespace,
+			"--type", "json", "-p",
+			`[{"op": "add", "path": "/spec/template/spec/containers/0/args/-", "value": "--restart-timeout=10s"}]`)
+		_, err = utils.Run(cmd)
+		Expect(err).NotTo(HaveOccurred(), "Failed to patch controller deployment")
+
+		By("waiting for controller rollout to complete")
+		cmd = exec.Command("kubectl", "rollout", "status", "deployment/opensandbox-controller-manager", "-n", namespace, "--timeout=60s")
+		_, err = utils.Run(cmd)
+		Expect(err).NotTo(HaveOccurred(), "Failed to wait for controller rollout")
+
 		By("waiting for controller to be ready")
 		Eventually(func(g Gomega) {
 			cmd := exec.Command("kubectl", "get", "pods", "-l", "control-plane=controller-manager",
@@ -45,6 +78,20 @@ var _ = Describe("Pod Recycle Policy", Ordered, func() {
 		}, 2*time.Minute).Should(Succeed())
 	})
 
+	AfterAll(func() {
+		By("undeploying the controller-manager")
+		cmd := exec.Command("make", "undeploy")
+		_, _ = utils.Run(cmd)
+
+		By("uninstalling CRDs")
+		cmd = exec.Command("make", "uninstall")
+		_, _ = utils.Run(cmd)
+
+		By("removing manager namespace")
+		cmd = exec.Command("kubectl", "delete", "ns", namespace)
+		_, _ = utils.Run(cmd)
+	})
+
 	SetDefaultEventuallyTimeout(3 * time.Minute)
 	SetDefaultEventuallyPollingInterval(2 * time.Second)
 
@@ -411,7 +458,7 @@ var _ = Describe("Pod Recycle Policy", Ordered, func() {
 				_, err := utils.Run(cmd)
 				g.Expect(err).To(HaveOccurred(), "Pod should be deleted after timeout")
 				g.Expect(err.Error()).To(ContainSubstring("not found"))
-			}, 4*time.Minute).Should(Succeed()) // restartTimeout (90s) + buffer
+			}, 1*time.Minute).Should(Succeed())
 
 			By("cleaning up")
 			cmd = exec.Command("kubectl", "delete", "pool", poolName, "-n", testNamespace)