diff --git a/meta/hibernate.go b/meta/hibernate.go index 8db1c2f..517fedc 100644 --- a/meta/hibernate.go +++ b/meta/hibernate.go @@ -36,3 +36,15 @@ func ReadRestoreFromHibernate(pod *corev1.Pod) bool { } return pod.Annotations[AnnotationRestoreFromHibernate] == annotationTrue } + +// MarkRestoreFromHibernate flags a freshly-built pod to restore its VM from the +// :hibernate snapshot instead of cloning from the base image. The operator sets +// this when it (re)creates a pod for an agent that is currently hibernated (see +// AnnotationRestoreFromHibernate); symmetric with ReadRestoreFromHibernate. +func MarkRestoreFromHibernate(pod *corev1.Pod) { + if pod == nil { + return + } + a := ensurePodAnnotations(pod) + a[AnnotationRestoreFromHibernate] = annotationTrue +} diff --git a/meta/hibernate_test.go b/meta/hibernate_test.go index 4350334..7987a85 100644 --- a/meta/hibernate_test.go +++ b/meta/hibernate_test.go @@ -65,6 +65,21 @@ func TestReadHibernateStateNilPod(t *testing.T) { } } +func TestMarkRestoreFromHibernate(t *testing.T) { + pod := &corev1.Pod{} + if ReadRestoreFromHibernate(pod) { + t.Fatal("fresh pod should not be flagged for restore") + } + MarkRestoreFromHibernate(pod) + if !ReadRestoreFromHibernate(pod) { + t.Errorf("MarkRestoreFromHibernate should round-trip through ReadRestoreFromHibernate") + } +} + +func TestMarkRestoreFromHibernateNilPod(t *testing.T) { + MarkRestoreFromHibernate(nil) // must not panic +} + func TestHibernateSnapshotTagConstant(t *testing.T) { if HibernateSnapshotTag != "hibernate" { t.Errorf("HibernateSnapshotTag = %q, want %q", HibernateSnapshotTag, "hibernate")