Mutation cache warming for Assign (#1161)

* Mutation cache warming for Assign

Signed-off-by: mmirecki <mmirecki@redhat.com>

* Add Assign readiness test

Signed-off-by: mmirecki <mmirecki@redhat.com>

Co-authored-by: Max Smythe <smythe@google.com>

pkg/controller/assign/assign_controller.go

@@ -29,6 +29,7 @@ import (
 	"github.com/open-policy-agent/gatekeeper/pkg/watch"
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/controller"
@@ -43,14 +44,21 @@ var (
 	log = logf.Log.WithName("controller").WithValues(logging.Process, "assign_controller")
 )
 
+var gvkAssign = schema.GroupVersionKind{
+	Group:   mutationsv1alpha1.GroupVersion.Group,
+	Version: mutationsv1alpha1.GroupVersion.Version,
+	Kind:    "Assign",
+}
+
 type Adder struct {
 	MutationCache *mutation.System
+	Tracker       *readiness.Tracker
 }
 
 // Add creates a new Assign Controller and adds it to the Manager. The Manager will set fields on the Controller
 // and Start it when the Manager is Started.
 func (a *Adder) Add(mgr manager.Manager) error {
-	r := newReconciler(mgr, a.MutationCache)
+	r := newReconciler(mgr, a.MutationCache, a.Tracker)
 	return add(mgr, r)
 }
 
@@ -60,15 +68,18 @@ func (a *Adder) InjectWatchManager(w *watch.Manager) {}
 
 func (a *Adder) InjectControllerSwitch(cs *watch.ControllerSwitch) {}
 
-func (a *Adder) InjectTracker(t *readiness.Tracker) {}
+func (a *Adder) InjectTracker(t *readiness.Tracker) {
+	a.Tracker = t
+
+}
 
 func (a *Adder) InjectMutationCache(mutationCache *mutation.System) {
 	a.MutationCache = mutationCache
 }
 
 // newReconciler returns a new reconcile.Reconciler
-func newReconciler(mgr manager.Manager, mutationCache *mutation.System) *Reconciler {
-	r := &Reconciler{system: mutationCache, Client: mgr.GetClient()}
+func newReconciler(mgr manager.Manager, mutationCache *mutation.System, tracker *readiness.Tracker) *Reconciler {
+	r := &Reconciler{system: mutationCache, Client: mgr.GetClient(), tracker: tracker}
 	return r
 }
 
@@ -96,7 +107,8 @@ func add(mgr manager.Manager, r reconcile.Reconciler) error {
 // Reconciler reconciles a Assign object
 type Reconciler struct {
 	client.Client
-	system *mutation.System
+	system  *mutation.System
+	tracker *readiness.Tracker
 }
 
 // +kubebuilder:rbac:groups=mutations.gatekeeper.sh,resources=*,verbs=get;list;watch;create;update;patch;delete
@@ -125,27 +137,32 @@ func (r *Reconciler) Reconcile(ctx context.Context, request reconcile.Request) (
 		}
 	}
 	deleted = deleted || !assign.GetDeletionTimestamp().IsZero()
+	tracker := r.tracker.For(gvkAssign)
 
 	if deleted {
 		id, err := types.MakeID(assign)
 		if err != nil {
 			log.Error(err, "Failed to get id out of assign")
-			return ctrl.Result{}, nil
-		}
-
-		if err := r.system.Remove(id); err != nil {
-			log.Error(err, "Remove failed", "resource", request.NamespacedName)
+		} else {
+			if err := r.system.Remove(id); err != nil {
+				log.Error(err, "Remove failed", "resource", request.NamespacedName)
+			}
 		}
+		tracker.CancelExpect(assign)
 		return ctrl.Result{}, nil
 	}
 
 	mutator, err := mutation.MutatorForAssign(assign)
 	if err != nil {
 		log.Error(err, "Creating mutator for resource failed", "resource", request.NamespacedName)
+		tracker.CancelExpect(assign)
 	}
+
 	if err := r.system.Upsert(mutator); err != nil {
 		log.Error(err, "Insert failed", "resource", request.NamespacedName)
+		tracker.TryCancelExpect(assign)
+	} else {
+		tracker.Observe(assign)
 	}
-
 	return ctrl.Result{}, nil
 }

pkg/readiness/ready_tracker.go

@@ -60,6 +60,7 @@ type Tracker struct {
 	templates      *objectTracker
 	config         *objectTracker
 	assignMetadata *objectTracker
+	assign         *objectTracker
 	constraints    *trackerMap
 	data           *trackerMap
 
@@ -84,6 +85,7 @@ func NewTracker(lister Lister, mutationEnabled bool) *Tracker {
 	}
 	if mutationEnabled {
 		tracker.assignMetadata = newObjTracker(mutationv1alpha.GroupVersion.WithKind("AssignMetadata"), nil)
+		tracker.assign = newObjTracker(mutationv1alpha.GroupVersion.WithKind("Assign"), nil)
 	}
 	return &tracker
 }
@@ -109,6 +111,11 @@ func (t *Tracker) For(gvk schema.GroupVersionKind) Expectations {
 			return t.assignMetadata
 		}
 		return noopExpectations{}
+	case gvk.GroupVersion() == mutationv1alpha.GroupVersion && gvk.Kind == "Assign":
+		if t.mutationEnabled {
+			return t.assign
+		}
+		return noopExpectations{}
 	}
 
 	// Avoid new constraint trackers after templates have been populated.
@@ -158,10 +165,11 @@ func (t *Tracker) Satisfied(ctx context.Context) bool {
 	}
 
 	if t.mutationEnabled {
-		if !t.assignMetadata.Satisfied() {
+		if !t.assignMetadata.Satisfied() || !t.assign.Satisfied() {
 			return false
 		}
 		log.V(1).Info("all expectations satisfied", "tracker", "assignMetadata")
+		log.V(1).Info("all expectations satisfied", "tracker", "assign")
 	}
 
 	if !t.templates.Satisfied() {
@@ -206,6 +214,9 @@ func (t *Tracker) Run(ctx context.Context) error {
 		grp.Go(func() error {
 			return t.trackAssignMetadata(gctx)
 		})
+		grp.Go(func() error {
+			return t.trackAssign(gctx)
+		})
 	}
 	grp.Go(func() error {
 		return t.trackConstraintTemplates(gctx)
@@ -385,6 +396,31 @@ func (t *Tracker) trackAssignMetadata(ctx context.Context) error {
 	return nil
 }
 
+func (t *Tracker) trackAssign(ctx context.Context) error {
+	defer func() {
+		t.assign.ExpectationsDone()
+		log.V(1).Info("Assign expectations populated")
+		_ = t.constraintTrackers.Wait()
+	}()
+
+	if !t.mutationEnabled {
+		return nil
+	}
+
+	assignList := &mutationv1alpha.AssignList{}
+	lister := retryLister(t.lister, retryAll)
+	if err := lister.List(ctx, assignList); err != nil {
+		return fmt.Errorf("listing Assign: %w", err)
+	}
+	log.V(1).Info("setting expectations for Assign", "Assign Count", len(assignList.Items))
+
+	for index := range assignList.Items {
+		log.V(1).Info("expecting Assign", "name", assignList.Items[index].GetName())
+		t.assign.Expect(&assignList.Items[index])
+	}
+	return nil
+}
+
 func (t *Tracker) trackConstraintTemplates(ctx context.Context) error {
 	defer func() {
 		t.templates.ExpectationsDone()
@@ -636,6 +672,7 @@ func (t *Tracker) statsPrinter(ctx context.Context) {
 		}
 		if t.mutationEnabled {
 			logUnsatisfiedAssignMetadata(t)
+			logUnsatisfiedAssign(t)
 		}
 	}
 }
@@ -646,6 +683,12 @@ func logUnsatisfiedAssignMetadata(t *Tracker) {
 	}
 }
 
+func logUnsatisfiedAssign(t *Tracker) {
+	for _, amKey := range t.assign.unsatisfied() {
+		log.Info("unsatisfied Assign", "name", amKey.namespacedName)
+	}
+}
+
 // Returns the constraint GVK that would be generated by a template.
 func constraintGVK(ct *templates.ConstraintTemplate) schema.GroupVersionKind {
 	return schema.GroupVersionKind{

pkg/readiness/ready_tracker_test.go

@@ -174,6 +174,54 @@ func Test_AssignMetadata(t *testing.T) {
 	}
 }
 
+func Test_Assign(t *testing.T) {
+	g := gomega.NewWithT(t)
+
+	defer func() {
+		mutationEnabled := false
+		mutation.MutationEnabled = &mutationEnabled
+	}()
+
+	mutationEnabled := true
+	mutation.MutationEnabled = &mutationEnabled
+
+	os.Setenv("POD_NAME", "no-pod")
+	podstatus.DisablePodOwnership()
+
+	// Apply fixtures *before* the controllers are setup.
+	err := applyFixtures("testdata")
+	g.Expect(err).NotTo(gomega.HaveOccurred(), "applying fixtures")
+
+	// Wire up the rest.
+	mgr, wm := setupManager(t)
+	opaClient := setupOpa(t)
+	mutationCache := mutation.NewSystem()
+	if err := setupController(mgr, wm, opaClient, mutationCache); err != nil {
+		t.Fatalf("setupControllers: %v", err)
+	}
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	mgrStopped := StartTestManager(ctx, mgr, g)
+	defer func() {
+		cancelFunc()
+		mgrStopped.Wait()
+	}()
+
+	g.Eventually(func() (bool, error) {
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		return probeIsReady(ctx)
+	}, 300*time.Second, 1*time.Second).Should(gomega.BeTrue())
+
+	// Verify that the Assign is present in the cache
+	for _, am := range testAssign {
+		id, err := mutationtypes.MakeID(am)
+		g.Expect(err).NotTo(gomega.HaveOccurred(), "can not create Assign id")
+		exptectedMutator := mutationCache.Get(id)
+		g.Expect(exptectedMutator).NotTo(gomega.BeNil(), "expected mutator was not found")
+	}
+}
+
 // Test_Tracker verifies that once an initial set of fixtures are loaded into OPA,
 // the readiness probe reflects that Gatekeeper is ready to enforce policy. Adding
 // additional constraints afterwards will not change the readiness state.

pkg/readiness/testdata/99-assign.yaml

@@ -0,0 +1,14 @@
+apiVersion: mutations.gatekeeper.sh/v1alpha1
+kind: Assign
+metadata:
+  name: demo
+spec:
+    match:
+      scope: Namespaced
+      kinds:
+        - apiGroups: ["*"]
+          kinds: ["Pod"]
+    location: "spec.dnsPolicy"
+    parameters:
+      assign:
+        value:  "None"

pkg/readiness/testdata_test.go

@@ -44,6 +44,10 @@ var testAssignMetadata = []*mutationsv1alpha1.AssignMetadata{
 	makeAssignMetadata("demo"),
 }
 
+var testAssign = []*mutationsv1alpha1.Assign{
+	makeAssign("demo"),
+}
+
 func makeTemplate(name string) *templates.ConstraintTemplate {
 	return &templates.ConstraintTemplate{
 		ObjectMeta: metav1.ObjectMeta{
@@ -74,3 +78,18 @@ func makeAssignMetadata(name string) *mutationsv1alpha1.AssignMetadata {
 		},
 	}
 }
+
+func makeAssign(name string) *mutationsv1alpha1.Assign {
+	return &mutationsv1alpha1.Assign{
+		TypeMeta: metav1.TypeMeta{
+			APIVersion: "mutations.gatekeeper.sh/v1alpha1",
+			Kind:       "Assign",
+		},
+		ObjectMeta: metav1.ObjectMeta{
+			Name: name,
+		},
+		Spec: mutationsv1alpha1.AssignSpec{
+			Location: "spec.dnsPolicy",
+		},
+	}
+}