Skip to content

Commit 287a975

Browse files
Rupam-ItRupam-It
committed
feat(controller): add annotation-based checkpoint trigger
Signed-off-by: Rupam-It <mannarupam3@gmail.com>
1 parent bd7ba15 commit 287a975

4 files changed

Lines changed: 236 additions & 8 deletions

File tree

‎config/rbac/role.yaml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ rules:
1919
verbs:
2020
- get
2121
- list
22+
- patch
2223
- watch
2324
- apiGroups:
2425
- criu.org

‎internal/controller/annotation_trigger.go‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
package controller
2+
3+
import (
4+
"context"
5+
"time"
6+
7+
v1 "github.com/checkpoint-restore/checkpoint-restore-operator/api/v1"
8+
"sigs.k8s.io/controller-runtime/pkg/client"
9+
"sigs.k8s.io/controller-runtime/pkg/log"
10+
)
11+
12+
const CheckpointTriggerAnnotation = "checkpoint.criu.org/trigger"
13+
14+
type AnnotationTrigger struct {
15+
client client.Client
16+
creator Checkpointer
17+
schedule *v1.CheckpointSchedule
18+
stopCh chan struct{}
19+
interval time.Duration
20+
}
21+
22+
func NewAnnotationTrigger(c client.Client, creator Checkpointer, schedule *v1.CheckpointSchedule) *AnnotationTrigger {
23+
return &AnnotationTrigger{
24+
client: c,
25+
creator: creator,
26+
schedule: schedule,
27+
stopCh: make(chan struct{}),
28+
interval: 30 * time.Second,
29+
}
30+
}
31+
32+
func (at *AnnotationTrigger) Start(ctx context.Context) {
33+
logger := log.FromContext(ctx)
34+
logger.Info("annotation trigger started", "interval", at.interval)
35+
36+
go func() {
37+
ticker := time.NewTicker(at.interval)
38+
defer ticker.Stop()
39+
for {
40+
select {
41+
case <-ticker.C:
42+
at.run(ctx)
43+
case <-at.stopCh:
44+
return
45+
case <-ctx.Done():
46+
return
47+
}
48+
}
49+
}()
50+
}
51+
52+
func (at *AnnotationTrigger) run(ctx context.Context) {
53+
logger := log.FromContext(ctx)
54+
55+
pods, err := getMatchingPods(ctx, at.client, at.schedule.Spec.Namespace, &at.schedule.Spec.Selector)
56+
if err != nil {
57+
logger.Error(err, "annotation trigger: failed to list pods")
58+
return
59+
}
60+
61+
for _, pod := range pods {
62+
if pod.Annotations[CheckpointTriggerAnnotation] != "true" {
63+
continue
64+
}
65+
66+
for _, c := range pod.Spec.Containers {
67+
if err := at.creator.createCheckpoint(ctx, pod.Namespace, pod.Name, c.Name, pod.Spec.NodeName); err != nil {
68+
logger.Error(err, "annotation trigger: checkpoint failed", "pod", pod.Name, "container", c.Name)
69+
}
70+
}
71+
72+
// clear the annotation so we don't checkpoint again next poll
73+
patch := client.MergeFrom(pod.DeepCopy())
74+
delete(pod.Annotations, CheckpointTriggerAnnotation)
75+
if err := at.client.Patch(ctx, &pod, patch); err != nil {
76+
logger.Error(err, "annotation trigger: failed to clear annotation", "pod", pod.Name)
77+
}
78+
}
79+
}
80+
81+
func (at *AnnotationTrigger) Stop() {
82+
close(at.stopCh)
83+
}

‎internal/controller/annotation_trigger_test.go‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
package controller
2+
3+
import (
4+
"context"
5+
6+
. "github.com/onsi/ginkgo/v2"
7+
. "github.com/onsi/gomega"
8+
corev1 "k8s.io/api/core/v1"
9+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
10+
"k8s.io/client-go/kubernetes/scheme"
11+
"sigs.k8s.io/controller-runtime/pkg/client"
12+
"sigs.k8s.io/controller-runtime/pkg/client/fake"
13+
14+
v1 "github.com/checkpoint-restore/checkpoint-restore-operator/api/v1"
15+
)
16+
17+
func makeAnnotationTrigger(mock *mockCheckpointer, pods ...*corev1.Pod) *AnnotationTrigger {
18+
builder := fake.NewClientBuilder().WithScheme(scheme.Scheme)
19+
for _, p := range pods {
20+
builder = builder.WithObjects(p).WithStatusSubresource(p)
21+
}
22+
sched := &v1.CheckpointSchedule{
23+
Spec: v1.CheckpointScheduleSpec{
24+
Namespace: "default",
25+
Selector: metav1.LabelSelector{MatchLabels: map[string]string{"app": "test"}},
26+
},
27+
}
28+
return NewAnnotationTrigger(builder.Build(), mock, sched)
29+
}
30+
31+
var _ = Describe("AnnotationTrigger.run", func() {
32+
var mock *mockCheckpointer
33+
34+
BeforeEach(func() {
35+
mock = &mockCheckpointer{}
36+
})
37+
38+
It("checkpoints a pod that has the trigger annotation and clears it", func() {
39+
pod := &corev1.Pod{
40+
ObjectMeta: metav1.ObjectMeta{
41+
Name: "pod-1", Namespace: "default",
42+
Labels: map[string]string{"app": "test"},
43+
Annotations: map[string]string{CheckpointTriggerAnnotation: "true"},
44+
},
45+
Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "app"}}},
46+
Status: corev1.PodStatus{Phase: corev1.PodRunning},
47+
}
48+
trigger := makeAnnotationTrigger(mock, pod)
49+
trigger.run(context.Background())
50+
51+
Expect(mock.calls).To(HaveLen(1))
52+
Expect(mock.calls[0].container).To(Equal("app"))
53+
54+
// annotation must be cleared
55+
updated := &corev1.Pod{}
56+
Expect(trigger.client.Get(context.Background(), client.ObjectKeyFromObject(pod), updated)).To(Succeed())
57+
Expect(updated.Annotations).NotTo(HaveKey(CheckpointTriggerAnnotation))
58+
})
59+
60+
It("skips pods that do not have the trigger annotation", func() {
61+
pod := &corev1.Pod{
62+
ObjectMeta: metav1.ObjectMeta{
63+
Name: "pod-1", Namespace: "default",
64+
Labels: map[string]string{"app": "test"},
65+
},
66+
Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "app"}}},
67+
Status: corev1.PodStatus{Phase: corev1.PodRunning},
68+
}
69+
trigger := makeAnnotationTrigger(mock, pod)
70+
trigger.run(context.Background())
71+
72+
Expect(mock.calls).To(BeEmpty())
73+
})
74+
})

‎internal/controller/checkpointschedule_controller.go‎

Lines changed: 78 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,47 +18,117 @@ package controller
1818

1919
import (
2020
"context"
21+
"sync"
2122

2223
v1 "github.com/checkpoint-restore/checkpoint-restore-operator/api/v1"
2324
"k8s.io/apimachinery/pkg/runtime"
2425
"k8s.io/client-go/rest"
2526
ctrl "sigs.k8s.io/controller-runtime"
2627
"sigs.k8s.io/controller-runtime/pkg/client"
28+
"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
2729
"sigs.k8s.io/controller-runtime/pkg/log"
2830
)
2931

3032
//+kubebuilder:rbac:groups=criu.org,resources=checkpointschedules,verbs=get;list;watch;create;update;patch;delete
3133
//+kubebuilder:rbac:groups=criu.org,resources=checkpointschedules/status,verbs=get;update;patch
3234
//+kubebuilder:rbac:groups=criu.org,resources=checkpointschedules/finalizers,verbs=update
35+
//+kubebuilder:rbac:groups="",resources=pods,verbs=get;list;watch;patch
36+
37+
const checkpointScheduleFinalizer = "criu.org/checkpoint-schedule-finalizer"
38+
39+
type Stoppable interface {
40+
Stop()
41+
}
3342

3443
type CheckpointScheduleReconciler struct {
3544
client.Client
36-
Scheme *runtime.Scheme
37-
RestConfig *rest.Config
45+
Scheme *runtime.Scheme
46+
RestConfig *rest.Config
47+
activeTriggers map[string][]Stoppable
48+
mu sync.Mutex
3849
}
3950

4051
func (r *CheckpointScheduleReconciler) Reconcile(
4152
ctx context.Context,
4253
req ctrl.Request,
4354
) (ctrl.Result, error) {
4455
logger := log.FromContext(ctx)
45-
// fetch the CR
56+
4657
schedule := &v1.CheckpointSchedule{}
4758
if err := r.Get(ctx, req.NamespacedName, schedule); err != nil {
4859
return ctrl.Result{}, client.IgnoreNotFound(err)
4960
}
61+
62+
key := req.NamespacedName.String()
63+
64+
// CR is being deleted — stop triggers and remove finalizer
65+
if !schedule.DeletionTimestamp.IsZero() {
66+
r.stopTriggers(key)
67+
if controllerutil.ContainsFinalizer(schedule, checkpointScheduleFinalizer) {
68+
controllerutil.RemoveFinalizer(schedule, checkpointScheduleFinalizer)
69+
if err := r.Update(ctx, schedule); err != nil {
70+
return ctrl.Result{}, err
71+
}
72+
}
73+
return ctrl.Result{}, nil
74+
}
75+
76+
// add finalizer if not present
77+
if !controllerutil.ContainsFinalizer(schedule, checkpointScheduleFinalizer) {
78+
controllerutil.AddFinalizer(schedule, checkpointScheduleFinalizer)
79+
if err := r.Update(ctx, schedule); err != nil {
80+
return ctrl.Result{}, err
81+
}
82+
return ctrl.Result{}, nil
83+
}
84+
85+
logger.Info("reconciling", "name", schedule.Name)
86+
87+
// already running — don't start duplicates
88+
r.mu.Lock()
89+
_, running := r.activeTriggers[key]
90+
r.mu.Unlock()
91+
if running {
92+
return ctrl.Result{}, nil
93+
}
94+
5095
creator := NewCheckpointCreator(r.Client, r.RestConfig)
96+
var triggers []Stoppable
97+
98+
if schedule.Spec.Triggers.Schedule != "" {
99+
t := NewScheduleTrigger(r.Client, creator, schedule)
100+
if err := t.Start(ctx); err != nil {
101+
logger.Error(err, "failed to start schedule trigger")
102+
return ctrl.Result{}, err
103+
}
104+
triggers = append(triggers, t)
105+
}
51106

52-
trigger := NewScheduleTrigger(r.Client, creator, schedule)
53-
if err := trigger.Start(ctx); err != nil {
54-
logger.Error(err, "failed to start schedule trigger")
55-
return ctrl.Result{}, err
107+
if schedule.Spec.Triggers.OnAnnotation {
108+
t := NewAnnotationTrigger(r.Client, creator, schedule)
109+
t.Start(ctx)
110+
triggers = append(triggers, t)
56111
}
57112

58-
logger.Info("schedule trigger started", "name", schedule.Name)
113+
r.mu.Lock()
114+
if r.activeTriggers == nil {
115+
r.activeTriggers = make(map[string][]Stoppable)
116+
}
117+
r.activeTriggers[key] = triggers
118+
r.mu.Unlock()
119+
59120
return ctrl.Result{}, nil
121+
}
60122

123+
func (r *CheckpointScheduleReconciler) stopTriggers(key string) {
124+
r.mu.Lock()
125+
defer r.mu.Unlock()
126+
for _, t := range r.activeTriggers[key] {
127+
t.Stop()
128+
}
129+
delete(r.activeTriggers, key)
61130
}
131+
62132
func (r *CheckpointScheduleReconciler) SetupWithManager(mgr ctrl.Manager) error {
63133
return ctrl.NewControllerManagedBy(mgr).
64134
For(&v1.CheckpointSchedule{}).

0 commit comments

Comments
 (0)