Fix metrics controller races

go.mod

@@ -15,6 +15,7 @@ require (
 	github.com/onsi/gomega v1.18.1
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/prometheus/client_golang v1.12.1
+	github.com/prometheus/client_model v0.2.0
 	go.uber.org/multierr v1.7.0
 	go.uber.org/zap v1.20.0
 	golang.org/x/time v0.0.0-20210723032227-1f47c861a9ac
@@ -62,7 +63,6 @@ require (
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/nxadm/tail v1.4.8 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
-	github.com/prometheus/client_model v0.2.0 // indirect
 	github.com/prometheus/common v0.32.1 // indirect
 	github.com/prometheus/procfs v0.7.3 // indirect
 	github.com/prometheus/statsd_exporter v0.21.0 // indirect

pkg/controllers/metrics/node/controller.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"sync"
 
 	"knative.dev/pkg/logging"
 
@@ -128,15 +129,14 @@ func labelNames() []string {
 }
 
 type Controller struct {
-	KubeClient      client.Client
-	LabelCollection map[types.NamespacedName][]prometheus.Labels
+	kubeClient      client.Client
+	labelCollection sync.Map
 }
 
 // NewController constructs a controller instance
 func NewController(kubeClient client.Client) *Controller {
 	return &Controller{
-		KubeClient:      kubeClient,
-		LabelCollection: make(map[types.NamespacedName][]prometheus.Labels),
+		kubeClient: kubeClient,
 	}
 }
 
@@ -147,7 +147,7 @@ func (c *Controller) Reconcile(ctx context.Context, req reconcile.Request) (reco
 	c.cleanup(req.NamespacedName)
 	// Retrieve node from reconcile request
 	node := &v1.Node{}
-	if err := c.KubeClient.Get(ctx, req.NamespacedName, node); err != nil {
+	if err := c.kubeClient.Get(ctx, req.NamespacedName, node); err != nil {
 		if errors.IsNotFound(err) {
 			return reconcile.Result{}, nil
 		}
@@ -170,7 +170,7 @@ func (c *Controller) Register(ctx context.Context, m manager.Manager) error {
 			&source.Kind{Type: &v1alpha5.Provisioner{}},
 			handler.EnqueueRequestsFromMapFunc(func(o client.Object) (requests []reconcile.Request) {
 				nodes := &v1.NodeList{}
-				if err := c.KubeClient.List(ctx, nodes, client.MatchingLabels(map[string]string{v1alpha5.ProvisionerNameLabelKey: o.GetName()})); err != nil {
+				if err := c.kubeClient.List(ctx, nodes, client.MatchingLabels(map[string]string{v1alpha5.ProvisionerNameLabelKey: o.GetName()})); err != nil {
 					logging.FromContext(ctx).Errorf("Failed to list nodes when mapping expiration watch events, %s", err)
 					return requests
 				}
@@ -194,8 +194,8 @@ func (c *Controller) Register(ctx context.Context, m manager.Manager) error {
 }
 
 func (c *Controller) cleanup(nodeNamespacedName types.NamespacedName) {
-	if labelSet, ok := c.LabelCollection[nodeNamespacedName]; ok {
-		for _, labels := range labelSet {
+	if labelSet, ok := c.labelCollection.Load(nodeNamespacedName); ok {
+		for _, labels := range labelSet.([]prometheus.Labels) {
 			allocatableGaugeVec.Delete(labels)
 			podRequestsGaugeVec.Delete(labels)
 			podLimitsGaugeVec.Delete(labels)
@@ -204,7 +204,7 @@ func (c *Controller) cleanup(nodeNamespacedName types.NamespacedName) {
 			overheadGaugeVec.Delete(labels)
 		}
 	}
-	c.LabelCollection[nodeNamespacedName] = []prometheus.Labels{}
+	c.labelCollection.Store(nodeNamespacedName, []prometheus.Labels{})
 }
 
 // labels creates the labels using the current state of the pod
@@ -231,7 +231,7 @@ func (c *Controller) labels(node *v1.Node, resourceTypeName string) prometheus.L
 
 func (c *Controller) record(ctx context.Context, node *v1.Node) error {
 	podlist := &v1.PodList{}
-	if err := c.KubeClient.List(ctx, podlist, client.MatchingFields{"spec.nodeName": node.Name}); err != nil {
+	if err := c.kubeClient.List(ctx, podlist, client.MatchingFields{"spec.nodeName": node.Name}); err != nil {
 		return fmt.Errorf("listing pods on node %s, %w", node.Name, err)
 	}
 	var daemons, pods []*v1.Pod
@@ -287,7 +287,11 @@ func (c *Controller) set(resourceList v1.ResourceList, node *v1.Node, gaugeVec *
 		labels := c.labels(node, resourceTypeName)
 		// Register the set of labels that are generated for node
 		nodeNamespacedName := types.NamespacedName{Name: node.Name}
-		c.LabelCollection[nodeNamespacedName] = append(c.LabelCollection[nodeNamespacedName], labels)
+
+		existingLabels, _ := c.labelCollection.LoadOrStore(nodeNamespacedName, []prometheus.Labels{})
+		existingLabels = append(existingLabels.([]prometheus.Labels), labels)
+		c.labelCollection.Store(nodeNamespacedName, existingLabels)
+
 		gauge, err := gaugeVec.GetMetricWith(labels)
 		if err != nil {
 			return fmt.Errorf("generate new gauge: %w", err)

pkg/controllers/metrics/node/suite_test.go

@@ -0,0 +1,83 @@
+/*
+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 node_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/aws/karpenter/pkg/cloudprovider/fake"
+	"github.com/aws/karpenter/pkg/cloudprovider/registry"
+	"github.com/aws/karpenter/pkg/controllers/metrics/node"
+	"github.com/aws/karpenter/pkg/test"
+	. "github.com/aws/karpenter/pkg/test/expectations"
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	. "knative.dev/pkg/logging/testing"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+var controller *node.Controller
+var ctx context.Context
+var env *test.Environment
+
+func TestAPIs(t *testing.T) {
+	ctx = TestContextWithLogger(t)
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Controllers/Metrics/Node")
+}
+
+var _ = BeforeSuite(func() {
+	env = test.NewEnvironment(ctx, func(e *test.Environment) {
+		cloudProvider := &fake.CloudProvider{}
+		registry.RegisterOrDie(ctx, cloudProvider)
+		controller = node.NewController(env.Client)
+	})
+	Expect(env.Start()).To(Succeed(), "Failed to start environment")
+})
+
+var _ = Describe("Node Metrics", func() {
+	It("should update the allocatable metric", func() {
+		node := test.Node(test.NodeOptions{
+			Allocatable: v1.ResourceList{
+				v1.ResourcePods:   resource.MustParse("100"),
+				v1.ResourceCPU:    resource.MustParse("5000"),
+				v1.ResourceMemory: resource.MustParse("32Gi"),
+			},
+		})
+		ExpectCreated(ctx, env.Client, node)
+		ExpectReconcileSucceeded(ctx, controller, client.ObjectKeyFromObject(node))
+
+		// metrics should now be tracking the allocatable capacity of our single node
+		nodeAllocation := ExpectMetric("karpenter_nodes_allocatable")
+
+		expectedValues := map[string]float64{
+			"cpu":    5000.0,
+			"pods":   100.0,
+			"memory": 32 * 1024 * 1024 * 1024,
+		}
+
+		for _, m := range nodeAllocation.Metric {
+			for _, l := range m.Label {
+				if l.GetName() == "resource_type" {
+					Expect(m.GetGauge().GetValue()).To(Equal(expectedValues[l.GetValue()]),
+						fmt.Sprintf("%s, %f to equal %f", l.GetValue(), m.GetGauge().GetValue(),
+							expectedValues[l.GetValue()]))
+				}
+			}
+		}
+	})
+})

pkg/controllers/metrics/pod/controller.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"sync"
 
 	"knative.dev/pkg/logging"
 
@@ -60,8 +61,8 @@ var (
 
 // Controller for the resource
 type Controller struct {
-	KubeClient client.Client
-	LabelsMap  map[types.NamespacedName]prometheus.Labels
+	kubeClient client.Client
+	labelsMap  sync.Map
 }
 
 func init() {
@@ -86,21 +87,20 @@ func labelNames() []string {
 // NewController constructs a controller instance
 func NewController(kubeClient client.Client) *Controller {
 	return &Controller{
-		KubeClient: kubeClient,
-		LabelsMap:  make(map[types.NamespacedName]prometheus.Labels),
+		kubeClient: kubeClient,
 	}
 }
 
 // Reconcile executes a termination control loop for the resource
 func (c *Controller) Reconcile(ctx context.Context, req reconcile.Request) (reconcile.Result, error) {
 	ctx = logging.WithLogger(ctx, logging.FromContext(ctx).Named("podmetrics").With("pod", req.Name))
 	// Remove the previous gauge after pod labels are updated
-	if labels, ok := c.LabelsMap[req.NamespacedName]; ok {
-		podGaugeVec.Delete(labels)
+	if labels, ok := c.labelsMap.Load(req.NamespacedName); ok {
+		podGaugeVec.Delete(labels.(prometheus.Labels))
 	}
 	// Retrieve pod from reconcile request
 	pod := &v1.Pod{}
-	if err := c.KubeClient.Get(ctx, req.NamespacedName, pod); err != nil {
+	if err := c.kubeClient.Get(ctx, req.NamespacedName, pod); err != nil {
 		if errors.IsNotFound(err) {
 			return reconcile.Result{}, nil
 		}
@@ -113,7 +113,7 @@ func (c *Controller) Reconcile(ctx context.Context, req reconcile.Request) (reco
 func (c *Controller) record(ctx context.Context, pod *v1.Pod) {
 	labels := c.labels(ctx, pod)
 	podGaugeVec.With(labels).Set(float64(1))
-	c.LabelsMap[client.ObjectKeyFromObject(pod)] = labels
+	c.labelsMap.Store(client.ObjectKeyFromObject(pod), labels)
 }
 
 func (c *Controller) Register(ctx context.Context, m manager.Manager) error {
@@ -141,7 +141,7 @@ func (c *Controller) labels(ctx context.Context, pod *v1.Pod) prometheus.Labels
 	metricLabels[podHostName] = pod.Spec.NodeName
 	metricLabels[podPhase] = string(pod.Status.Phase)
 	node := &v1.Node{}
-	if err := c.KubeClient.Get(ctx, types.NamespacedName{Name: pod.Spec.NodeName}, node); err != nil {
+	if err := c.kubeClient.Get(ctx, types.NamespacedName{Name: pod.Spec.NodeName}, node); err != nil {
 		metricLabels[podHostZone] = "N/A"
 		metricLabels[podHostArchitecture] = "N/A"
 		metricLabels[podHostCapacityType] = "N/A"

pkg/controllers/metrics/pod/suite_test.go

@@ -0,0 +1,74 @@
+/*
+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 pod_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/aws/karpenter/pkg/cloudprovider/fake"
+	"github.com/aws/karpenter/pkg/cloudprovider/registry"
+	"github.com/aws/karpenter/pkg/controllers/metrics/pod"
+	"github.com/aws/karpenter/pkg/test"
+	. "github.com/aws/karpenter/pkg/test/expectations"
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+	prometheus "github.com/prometheus/client_model/go"
+	. "knative.dev/pkg/logging/testing"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+var controller *pod.Controller
+var ctx context.Context
+var env *test.Environment
+
+func TestAPIs(t *testing.T) {
+	ctx = TestContextWithLogger(t)
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Controllers/Metrics/Node")
+}
+
+var _ = BeforeSuite(func() {
+	env = test.NewEnvironment(ctx, func(e *test.Environment) {
+		cloudProvider := &fake.CloudProvider{}
+		registry.RegisterOrDie(ctx, cloudProvider)
+		controller = pod.NewController(env.Client)
+	})
+	Expect(env.Start()).To(Succeed(), "Failed to start environment")
+})
+
+var _ = Describe("Pod Metrics", func() {
+	It("should update the pod state metrics", func() {
+		p := test.Pod()
+		ExpectCreated(ctx, env.Client, p)
+		ExpectReconcileSucceeded(ctx, controller, client.ObjectKeyFromObject(p))
+
+		podState := ExpectMetric("karpenter_pods_state")
+		ExpectMetricLabel(podState, "name", p.GetName())
+		ExpectMetricLabel(podState, "namespace", p.GetNamespace())
+	})
+})
+
+func ExpectMetricLabel(mf *prometheus.MetricFamily, name string, value string) {
+	found := false
+	for _, m := range mf.Metric {
+		for _, l := range m.Label {
+			if l.GetName() == name {
+				Expect(l.GetValue()).To(Equal(value), fmt.Sprintf("expected metrics %s = %s", name, value))
+				found = true
+			}
+		}
+	}
+	Expect(found).To(BeTrue())
+}

pkg/test/expectations/expectations.go

@@ -22,6 +22,8 @@ import (
 
 	//nolint:revive,stylecheck
 	. "github.com/onsi/gomega"
+	prometheus "github.com/prometheus/client_model/go"
+	"sigs.k8s.io/controller-runtime/pkg/metrics"
 
 	appsv1 "k8s.io/api/apps/v1"
 	v1 "k8s.io/api/core/v1"
@@ -196,3 +198,16 @@ func ExpectReconcileSucceeded(ctx context.Context, reconciler reconcile.Reconcil
 	Expect(err).ToNot(HaveOccurred())
 	return result
 }
+
+func ExpectMetric(prefix string) *prometheus.MetricFamily {
+	metrics, err := metrics.Registry.Gather()
+	Expect(err).To(BeNil())
+	var selected *prometheus.MetricFamily
+	for _, mf := range metrics {
+		if mf.GetName() == prefix {
+			selected = mf
+		}
+	}
+	Expect(selected).ToNot(BeNil(), fmt.Sprintf("expected to find a '%s' metric", prefix))
+	return selected
+}