Add readiness retries support for Constraint Templates

Signed-off-by: juliankatz <juliankatz@google.com>

pkg/controller/constrainttemplate/constrainttemplate_controller.go

@@ -324,7 +324,7 @@ func (r *ReconcileConstraintTemplate) Reconcile(ctx context.Context, request rec
 	}
 	unversionedProposedCRD, err := r.opa.CreateCRD(context.Background(), unversionedCT)
 	if err != nil {
-		r.tracker.CancelTemplate(unversionedCT) // Don't track templates that failed compilation
+		r.tracker.TryCancelTemplate(unversionedCT) // Don't track templates that failed compilation
 		r.metrics.registry.add(request.NamespacedName, metrics.ErrorStatus)
 		var createErr *v1beta1.CreateCRDError
 		if parseErrs, ok := err.(ast.Errors); ok {
@@ -347,7 +347,7 @@ func (r *ReconcileConstraintTemplate) Reconcile(ctx context.Context, request rec
 
 	proposedCRD := &apiextensionsv1beta1.CustomResourceDefinition{}
 	if err := r.scheme.Convert(unversionedProposedCRD, proposedCRD, nil); err != nil {
-		r.tracker.CancelTemplate(unversionedCT) // Don't track templates that failed compilation
+		r.tracker.TryCancelTemplate(unversionedCT) // Don't track templates that failed compilation
 		r.metrics.registry.add(request.NamespacedName, metrics.ErrorStatus)
 		log.Error(err, "conversion error")
 		logError(request.NamespacedName.Name)
@@ -419,7 +419,7 @@ func (r *ReconcileConstraintTemplate) handleUpdate(
 			log.Error(err, "failed to report constraint template ingestion duration")
 		}
 		err := r.reportErrorOnCTStatus("ingest_error", "Could not ingest Rego", status, err)
-		r.tracker.CancelTemplate(unversionedCT) // Don't track templates that failed compilation
+		r.tracker.TryCancelTemplate(unversionedCT) // Don't track templates that failed compilation
 		return reconcile.Result{}, err
 	}
 

pkg/readiness/noop_expectations.go

@@ -24,7 +24,9 @@ func (n noopExpectations) Expect(o runtime.Object) {}
 
 func (n noopExpectations) CancelExpect(o runtime.Object) {}
 
-func (n noopExpectations) TryCancelExpect(o runtime.Object) {}
+func (n noopExpectations) TryCancelExpect(o runtime.Object) bool {
+	return false
+}
 
 func (n noopExpectations) ExpectationsDone() {}
 

pkg/readiness/object_tracker.go

@@ -39,7 +39,7 @@ var readinessRetries = flag.Int("readiness-retries", 0, "The number of resource
 type Expectations interface {
 	Expect(o runtime.Object)
 	CancelExpect(o runtime.Object)
-	TryCancelExpect(o runtime.Object)
+	TryCancelExpect(o runtime.Object) bool
 	ExpectationsDone()
 	Observe(o runtime.Object)
 	Satisfied() bool
@@ -146,19 +146,22 @@ func (t *objectTracker) CancelExpect(o runtime.Object) {
 	t.cancelExpectNoLock(k)
 }
 
-func (t *objectTracker) TryCancelExpect(o runtime.Object) {
+// TryCancelExpect will check the readinessRetries left on an Object, and cancel
+// the expectation for that object if no retries remain.  Returns True if the
+// expectation was cancelled.
+func (t *objectTracker) TryCancelExpect(o runtime.Object) bool {
 	t.mu.Lock()
 	defer t.mu.Unlock()
 
 	// Respect circuit-breaker.
 	if t.allSatisfied {
-		return
+		return false
 	}
 
 	k, err := objKeyFromObject(o)
 	if err != nil {
 		log.Error(err, "skipping")
-		return
+		return false
 	}
 
 	// Check if it's time to delete an expectation or just decrement its allotted retries
@@ -170,9 +173,13 @@ func (t *objectTracker) TryCancelExpect(o runtime.Object) {
 	shouldDel := obj.decrementRetries()
 	t.tryCancelled[k] = obj // set the changed obj back to the map, as the value is not a pointer
 
+	fmt.Printf("SHOULDDEL? %v\n", shouldDel)
+
 	if shouldDel {
 		t.cancelExpectNoLock(k)
 	}
+
+	return shouldDel
 }
 
 // ExpectationsDone tells the tracker to stop accepting new expectations.

pkg/readiness/object_tracker_test.go

@@ -358,7 +358,6 @@ func Test_ObjectTracker_TryCancelExpect_CancelBeforeExpected(t *testing.T) {
 	ot.TryCancelExpect(ct) // 0 retries --> DELETE
 
 	g.Expect(ot.Satisfied()).To(gomega.BeTrue(), "should be satisfied")
-
 }
 
 // Verify that unexpected observations do not prevent the tracker from reaching its satisfied state.

pkg/readiness/ready_tracker.go

@@ -72,20 +72,24 @@ type Tracker struct {
 
 // NewTracker creates a new Tracker and initializes the internal trackers
 func NewTracker(lister Lister, mutationEnabled bool) *Tracker {
+	return newTracker(lister, mutationEnabled, nil)
+}
+
+func newTracker(lister Lister, mutationEnabled bool, fn objDataFactory) *Tracker {
 	tracker := Tracker{
 		lister:             lister,
-		templates:          newObjTracker(v1beta1.SchemeGroupVersion.WithKind("ConstraintTemplate"), nil),
-		config:             newObjTracker(configv1alpha1.GroupVersion.WithKind("Config"), nil),
-		constraints:        newTrackerMap(),
-		data:               newTrackerMap(),
+		templates:          newObjTracker(v1beta1.SchemeGroupVersion.WithKind("ConstraintTemplate"), fn),
+		config:             newObjTracker(configv1alpha1.GroupVersion.WithKind("Config"), fn),
+		constraints:        newTrackerMap(fn),
+		data:               newTrackerMap(fn),
 		ready:              make(chan struct{}),
 		constraintTrackers: &syncutil.SingleRunner{},
 
 		mutationEnabled: mutationEnabled,
 	}
 	if mutationEnabled {
-		tracker.assignMetadata = newObjTracker(mutationv1alpha.GroupVersion.WithKind("AssignMetadata"), nil)
-		tracker.assign = newObjTracker(mutationv1alpha.GroupVersion.WithKind("Assign"), nil)
+		tracker.assignMetadata = newObjTracker(mutationv1alpha.GroupVersion.WithKind("AssignMetadata"), fn)
+		tracker.assign = newObjTracker(mutationv1alpha.GroupVersion.WithKind("Assign"), fn)
 	}
 	return &tracker
 }
@@ -138,16 +142,30 @@ func (t *Tracker) ForData(gvk schema.GroupVersionKind) Expectations {
 	return t.data.Get(gvk)
 }
 
-// CancelTemplate stops expecting the provided ConstraintTemplate and associated Constraints.
-func (t *Tracker) CancelTemplate(ct *templates.ConstraintTemplate) {
-	log.V(1).Info("cancel tracking for template", "namespace", ct.GetNamespace(), "name", ct.GetName())
-	t.templates.CancelExpect(ct)
+func (t *Tracker) templateCleanup(ct *templates.ConstraintTemplate) {
 	gvk := constraintGVK(ct)
 	t.constraints.Remove(gvk)
 	<-t.ready // constraintTrackers are setup in Run()
 	t.constraintTrackers.Cancel(gvk.String())
 }
 
+// CancelTemplate stops expecting the provided ConstraintTemplate and associated Constraints.
+func (t *Tracker) CancelTemplate(ct *templates.ConstraintTemplate) {
+	log.V(1).Info("cancel tracking for template", "namespace", ct.GetNamespace(), "name", ct.GetName())
+	t.templates.CancelExpect(ct)
+	t.templateCleanup(ct)
+}
+
+// TryCancelTemplate will check the readiness retries left on a CT and
+// cancel the expectation for that CT and its associated Constraints if
+// no retries remain.
+func (t *Tracker) TryCancelTemplate(ct *templates.ConstraintTemplate) {
+	log.V(1).Info("try to cancel tracking for template", "namespace", ct.GetNamespace(), "name", ct.GetName())
+	if t.templates.TryCancelExpect(ct) {
+		t.templateCleanup(ct)
+	}
+}
+
 // CancelData stops expecting data for the specified resource kind.
 func (t *Tracker) CancelData(gvk schema.GroupVersionKind) {
 	log.V(1).Info("cancel tracking for data", "gvk", gvk)
@@ -196,6 +214,7 @@ func (t *Tracker) Satisfied(ctx context.Context) bool {
 
 	t.mu.Lock()
 	defer t.mu.Unlock()
+
 	t.satisfied = true
 	return true
 }
@@ -262,6 +281,17 @@ func (t *Tracker) Run(ctx context.Context) error {
 	return nil
 }
 
+func (t *Tracker) Populated() bool {
+	mutationPopulated := true
+	if t.mutationEnabled {
+		// If !t.mutationEnabled and we call this, it yields a null pointer exception
+		mutationPopulated = t.assignMetadata.Populated() && t.assign.Populated()
+	}
+
+	ispop := t.templates.Populated() && t.config.Populated() && mutationPopulated && t.constraints.Populated() && t.data.Populated()
+	return ispop
+}
+
 // collectForObjectTracker identifies objects that are unsatisfied for the provided
 // `es`, which must be an objectTracker, and removes those expectations
 func (t *Tracker) collectForObjectTracker(ctx context.Context, es Expectations, cleanup func(schema.GroupVersionKind)) error {

pkg/readiness/ready_tracker_unit_test.go

@@ -0,0 +1,128 @@
+/*
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package readiness
+
+import (
+	"context"
+	"testing"
+
+	"github.com/onsi/gomega"
+	"github.com/open-policy-agent/frameworks/constraint/pkg/apis/templates/v1beta1"
+	"github.com/open-policy-agent/frameworks/constraint/pkg/core/templates"
+	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+// Stub out the lister
+type dummy_lister struct{}
+
+var scheme *runtime.Scheme
+
+func init() {
+	scheme = runtime.NewScheme()
+	if err := v1beta1.AddToScheme(scheme); err != nil {
+		panic(err)
+	}
+}
+
+var testConstraintTemplate = templates.ConstraintTemplate{
+	ObjectMeta: v1.ObjectMeta{
+		Name: "test-contraint-template",
+	},
+	Spec: templates.ConstraintTemplateSpec{
+		CRD: templates.CRD{
+			Spec: templates.CRDSpec{
+				Names: templates.Names{
+					Kind: "test-constraint",
+				},
+			},
+		},
+	},
+}
+
+func (dl dummy_lister) List(ctx context.Context, list client.ObjectList, opts ...client.ListOption) error {
+	switch l := list.(type) {
+	case *v1beta1.ConstraintTemplateList:
+		i := v1beta1.ConstraintTemplate{}
+		if err := scheme.Convert(&testConstraintTemplate, &i, nil); err != nil {
+			// These failures will be swallowed by readiness.retryAll
+			return err
+		}
+		l.Items = []v1beta1.ConstraintTemplate{i}
+	}
+	return nil
+}
+
+// Verify that TryCancelTemplate functions the same as regular CancelTemplate if readinessRetries is set to 0
+func Test_ReadyTracker_TryCancelTemplate_No_Retries(t *testing.T) {
+	g := gomega.NewWithT(t)
+
+	l := dummy_lister{}
+	rt := newTracker(l, false, func() objData {
+		return objData{retries: 0}
+	})
+
+	// Run kicks off all the tracking
+	ctx, cancel := context.WithCancel(context.Background())
+
+	go rt.Run(ctx)
+	defer cancel()
+
+	g.Eventually(func() bool {
+		return rt.Populated()
+	}, "10s").Should(gomega.BeTrue())
+
+	g.Expect(rt.Satisfied(nil)).NotTo(gomega.BeTrue(), "tracker with 0 retries should not be satisfied")
+
+	rt.TryCancelTemplate(&testConstraintTemplate) // 0 retries --> DELETE
+
+	g.Expect(rt.Satisfied(nil)).To(gomega.BeTrue(), "tracker with 0 retries and cancellation should be satisfied")
+}
+
+// Verify that TryCancelTemplate must be called enough times to remove all retries before cancelling a template
+func Test_ReadyTracker_TryCancelTemplate_Retries(t *testing.T) {
+	g := gomega.NewWithT(t)
+
+	l := dummy_lister{}
+	rt := newTracker(l, false, func() objData {
+		return objData{retries: 2}
+	})
+
+	// Run kicks off all the tracking
+	ctx, cancel := context.WithCancel(context.Background())
+
+	go rt.Run(ctx)
+	defer cancel()
+
+	g.Eventually(func() bool {
+		return rt.Populated()
+	}, "10s").Should(gomega.BeTrue())
+
+	g.Expect(rt.Satisfied(nil)).NotTo(gomega.BeTrue(), "tracker with 2 retries should not be satisfied")
+
+	rt.TryCancelTemplate(&testConstraintTemplate) // 2 --> 1 retries
+
+	g.Expect(rt.Satisfied(nil)).NotTo(gomega.BeTrue(), "tracker with 1 retries should not be satisfied")
+
+	rt.TryCancelTemplate(&testConstraintTemplate) // 1 --> 0 retries
+
+	g.Expect(rt.Satisfied(nil)).NotTo(gomega.BeTrue(), "tracker with 0 retries should not be satisfied")
+
+	rt.TryCancelTemplate(&testConstraintTemplate) // 0 retries --> DELETE
+
+	g.Expect(rt.Satisfied(nil)).To(gomega.BeTrue(), "tracker with 0 retries and cancellation should be satisfied")
+}

pkg/readiness/tracker_map.go

@@ -25,12 +25,14 @@ type trackerMap struct {
 	mu      sync.RWMutex
 	m       map[schema.GroupVersionKind]*objectTracker
 	removed map[schema.GroupVersionKind]struct{}
+	fn      objDataFactory
 }
 
-func newTrackerMap() *trackerMap {
+func newTrackerMap(fn objDataFactory) *trackerMap {
 	return &trackerMap{
 		m:       make(map[schema.GroupVersionKind]*objectTracker),
 		removed: make(map[schema.GroupVersionKind]struct{}),
+		fn:      fn,
 	}
 }
 
@@ -64,7 +66,7 @@ func (t *trackerMap) Get(gvk schema.GroupVersionKind) Expectations {
 
 	t.mu.Lock()
 	defer t.mu.Unlock()
-	entry := newObjTracker(gvk, nil)
+	entry := newObjTracker(gvk, t.fn)
 	t.m[gvk] = entry
 	return entry
 }
@@ -100,3 +102,15 @@ func (t *trackerMap) Satisfied() bool {
 	}
 	return true
 }
+
+// Populated returns true if all objectTrackers are populated
+func (t *trackerMap) Populated() bool {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	for _, ot := range t.m {
+		if !ot.Populated() {
+			return false
+		}
+	}
+	return true
+}