Merge pull request #16 from ctrox/container-status-labels

feat: add container status labels

README.md

@@ -304,12 +304,34 @@ node driver to patch pods. To deploy this, simply uncomment the
 This will add the flag and the required permissions when building the
 kustomization.
 
+#### Status Labels
+
+To reflect the container scaling status in the k8s API, the manager can set
+status labels on a pod. This requires the flag `-status-labels=true`, which is
+set by default in the production deployment.
+
+The resulting labels have the following structure:
+
+```yaml
+status.zeropod.ctrox.dev/<container name>: <container status>
+```
+
+So if our pod has two containers, one of them running and one in scaled-down
+state, the labels would be set like this:
+
+```yaml
+labels:
+  status.zeropod.ctrox.dev/container1: RUNNING
+  status.zeropod.ctrox.dev/container2: SCALED_DOWN
+```
+
 #### Flags
 
 ```
 -metrics-addr=":8080"    sets the address of the metrics server
 -debug                   enables debug logging
 -in-place-scaling=false  enable in-place resource scaling, requires InPlacePodVerticalScaling feature flag
+-status-labels=false     update pod labels to reflect container status
 ```
 
 ## Metrics

cmd/manager/main.go

@@ -19,6 +19,7 @@ var (
 	debug          = flag.Bool("debug", false, "enable debug logs")
 	inPlaceScaling = flag.Bool("in-place-scaling", false,
 		"enable in-place resource scaling, requires InPlacePodVerticalScaling feature flag")
+	statusLabels = flag.Bool("status-labels", false, "update pod labels to reflect container status")
 )
 
 func main() {
@@ -43,17 +44,15 @@ func main() {
 		os.Exit(1)
 	}
 
-	subscribers := []manager.StatusHandler{}
+	podHandlers := []manager.PodHandler{}
+	if *statusLabels {
+		podHandlers = append(podHandlers, manager.NewPodLabeller())
+	}
 	if *inPlaceScaling {
-		podScaler, err := manager.NewPodScaler()
-		if err != nil {
-			slog.Error("podScaler init", "err", err)
-			os.Exit(1)
-		}
-		subscribers = append(subscribers, podScaler)
+		podHandlers = append(podHandlers, manager.NewPodScaler())
 	}
 
-	if err := manager.StartSubscribers(ctx, subscribers...); err != nil {
+	if err := manager.StartSubscribers(ctx, podHandlers...); err != nil {
 		slog.Error("starting subscribers", "err", err)
 		os.Exit(1)
 	}

config/in-place-scaling/kustomization.yaml

@@ -1,7 +1,5 @@
 apiVersion: kustomize.config.k8s.io/v1alpha1
 kind: Component
-resources:
-  - rbac.yaml
 patches:
   - patch: |-
       - op: add

config/kind/kustomization.yaml

@@ -2,6 +2,8 @@ resources:
   - ../base
 components:
   - ../in-place-scaling
+  - ../pod-updater
+  - ../status-labels
 images:
   - name: manager
     newName: ghcr.io/ctrox/zeropod-manager

config/pod-updater/kustomization.yaml

@@ -0,0 +1,4 @@
+apiVersion: kustomize.config.k8s.io/v1alpha1
+kind: Component
+resources:
+  - rbac.yaml

config/production/kustomization.yaml

@@ -1,8 +1,11 @@
 resources:
 - ../base
-# uncommment to enable in-place scaling
-# components:
-#  - ../in-place-scaling
+# pod-updater is required if status-labels or in-place-scaling is enabled
+components:
+- ../pod-updater
+- ../status-labels
+# uncommment to enable in-place-scaling
+# - ../in-place-scaling
 images:
 - name: installer
   newName: ghcr.io/ctrox/zeropod-installer

config/status-labels/kustomization.yaml

@@ -0,0 +1,9 @@
+apiVersion: kustomize.config.k8s.io/v1alpha1
+kind: Component
+patches:
+  - patch: |-
+      - op: add
+        path: /spec/template/spec/containers/0/args/-
+        value: -status-labels=true
+    target:
+      kind: DaemonSet

e2e/e2e_test.go

@@ -5,10 +5,12 @@ import (
 	"fmt"
 	"net/http"
 	"runtime"
+	"strings"
 	"sync"
 	"testing"
 	"time"
 
+	v1 "github.com/ctrox/zeropod/api/shim/v1"
 	"github.com/ctrox/zeropod/manager"
 	"github.com/ctrox/zeropod/zeropod"
 	"github.com/prometheus/client_golang/prometheus"
@@ -240,6 +242,29 @@ func TestE2E(t *testing.T) {
 		}, time.Minute, time.Second)
 	})
 
+	t.Run("status labels", func(t *testing.T) {
+		pod := testPod(scaleDownAfter(0), agnContainer("agn", 8080), agnContainer("agn2", 8081))
+
+		cleanupPod := createPodAndWait(t, ctx, client, pod)
+		defer cleanupPod()
+		require.Eventually(t, func() bool {
+			if err := client.Get(ctx, objectName(pod), pod); err != nil {
+				return false
+			}
+			labelCount := 0
+			expectedLabels := 2
+			for k, v := range pod.GetLabels() {
+				if strings.HasPrefix(k, manager.StatusLabelKeyPrefix) {
+					if v == v1.ContainerPhase_SCALED_DOWN.String() {
+						labelCount++
+					}
+				}
+			}
+
+			return labelCount == expectedLabels
+		}, time.Minute, time.Second)
+	})
+
 	t.Run("metrics", func(t *testing.T) {
 		// create two pods to test metric merging
 		runningPod := testPod(scaleDownAfter(time.Hour))

manager/pod_labeller.go

@@ -0,0 +1,40 @@
+package manager
+
+import (
+	"context"
+	"log/slog"
+	"path"
+
+	v1 "github.com/ctrox/zeropod/api/shim/v1"
+	corev1 "k8s.io/api/core/v1"
+)
+
+const (
+	StatusLabelKeyPrefix = "status.zeropod.ctrox.dev"
+)
+
+type PodLabeller struct {
+	log *slog.Logger
+}
+
+func NewPodLabeller() *PodLabeller {
+	log := slog.With("component", "podupdater")
+	log.Info("init")
+	return &PodLabeller{log: log}
+}
+
+func (pl *PodLabeller) Handle(ctx context.Context, status *v1.ContainerStatus, pod *corev1.Pod) error {
+	clog := pl.log.With("container", status.Name, "pod", status.PodName,
+		"namespace", status.PodNamespace, "phase", status.Phase)
+	clog.Info("status event")
+
+	pl.setLabel(pod, status)
+	return nil
+}
+
+func (pu *PodLabeller) setLabel(pod *corev1.Pod, status *v1.ContainerStatus) {
+	if pod.Labels == nil {
+		pod.Labels = map[string]string{}
+	}
+	pod.Labels[path.Join(StatusLabelKeyPrefix, status.Name)] = status.Phase.String()
+}

manager/pod_labeller_test.go

@@ -0,0 +1,74 @@
+package manager
+
+import (
+	"context"
+	"log/slog"
+	"testing"
+
+	v1 "github.com/ctrox/zeropod/api/shim/v1"
+	"github.com/stretchr/testify/assert"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+)
+
+func TestPodLabeller(t *testing.T) {
+	slog.SetLogLoggerLevel(slog.LevelDebug)
+
+	scheme := runtime.NewScheme()
+	if err := corev1.AddToScheme(scheme); err != nil {
+		t.Fatal(err)
+	}
+	cases := map[string]struct {
+		statusEventPhase v1.ContainerPhase
+		beforeEvent      map[string]string
+		expected         map[string]string
+	}{
+		"no labels set": {
+			statusEventPhase: v1.ContainerPhase_RUNNING,
+			beforeEvent:      nil,
+			expected: map[string]string{
+				"status.zeropod.ctrox.dev/first-container": v1.ContainerPhase_RUNNING.String(),
+			},
+		},
+		"existing labels are kept": {
+			statusEventPhase: v1.ContainerPhase_RUNNING,
+			beforeEvent:      map[string]string{"existing": "label"},
+			expected: map[string]string{
+				"existing": "label",
+				"status.zeropod.ctrox.dev/first-container": v1.ContainerPhase_RUNNING.String(),
+			},
+		},
+		"status label is updated": {
+			statusEventPhase: v1.ContainerPhase_SCALED_DOWN,
+			beforeEvent: map[string]string{
+				"status.zeropod.ctrox.dev/first-container": v1.ContainerPhase_RUNNING.String(),
+			},
+			expected: map[string]string{
+				"status.zeropod.ctrox.dev/first-container": v1.ContainerPhase_SCALED_DOWN.String(),
+			},
+		},
+	}
+
+	for name, tc := range cases {
+		tc := tc
+		t.Run(name, func(t *testing.T) {
+			pod := newPod(nil)
+			pod.SetLabels(tc.beforeEvent)
+
+			if err := NewPodLabeller().Handle(
+				context.Background(),
+				&v1.ContainerStatus{
+					Name:         pod.Spec.Containers[0].Name,
+					PodName:      pod.Name,
+					PodNamespace: pod.Namespace,
+					Phase:        tc.statusEventPhase,
+				},
+				pod,
+			); err != nil {
+				t.Fatal(err)
+			}
+
+			assert.Equal(t, pod.GetLabels(), tc.expected)
+		})
+	}
+}

manager/pod_scaler.go

@@ -8,11 +8,7 @@ import (
 
 	v1 "github.com/ctrox/zeropod/api/shim/v1"
 	corev1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/api/resource"
-	"k8s.io/apimachinery/pkg/types"
-	"sigs.k8s.io/controller-runtime/pkg/client"
-	"sigs.k8s.io/controller-runtime/pkg/client/config"
 )
 
 const (
@@ -28,33 +24,20 @@ var (
 type containerResource map[string]resource.Quantity
 
 type PodScaler struct {
-	client client.Client
-	log    *slog.Logger
+	log *slog.Logger
 }
 
-func NewPodScaler() (*PodScaler, error) {
+func NewPodScaler() *PodScaler {
 	log := slog.With("component", "podscaler")
 	log.Info("init")
-	cfg, err := config.GetConfig()
-	if err != nil {
-		return nil, err
-	}
-	c, err := client.New(cfg, client.Options{})
-	return &PodScaler{client: c, log: log}, err
+	return &PodScaler{log: log}
 }
 
-func (ps *PodScaler) Handle(ctx context.Context, status *v1.ContainerStatus) error {
+func (ps *PodScaler) Handle(ctx context.Context, status *v1.ContainerStatus, pod *corev1.Pod) error {
 	clog := ps.log.With("container", status.Name, "pod", status.PodName,
 		"namespace", status.PodNamespace, "phase", status.Phase)
 	clog.Info("status event")
 
-	pod := &corev1.Pod{}
-	podName := types.NamespacedName{Name: status.PodName, Namespace: status.PodNamespace}
-	if err := ps.client.Get(ctx, podName, pod); err != nil {
-		return err
-	}
-
-	updatePod := false
 	for i, container := range pod.Spec.Containers {
 		if container.Name != status.Name {
 			continue
@@ -85,19 +68,6 @@ func (ps *PodScaler) Handle(ctx context.Context, status *v1.ContainerStatus) err
 		new := ps.newRequests(initial, current, status)
 		pod.Spec.Containers[i].Resources.Requests = new
 		clog.Debug("container needs to be updated", "current", printResources(current), "new", printResources(new))
-		updatePod = true
-	}
-
-	if !updatePod {
-		return nil
-	}
-
-	if err := ps.updateRequests(ctx, pod); err != nil {
-		if errors.IsInvalid(err) {
-			clog.Error("in-place scaling failed, ensure InPlacePodVerticalScaling feature flag is enabled")
-			return nil
-		}
-		return err
 	}
 
 	return nil
@@ -187,10 +157,6 @@ func (ps *PodScaler) setAnnotations(pod *corev1.Pod) error {
 	return nil
 }
 
-func (ps *PodScaler) updateRequests(ctx context.Context, pod *corev1.Pod) error {
-	return ps.client.Update(ctx, pod)
-}
-
 func printResources(res corev1.ResourceList) string {
 	cpu := res[corev1.ResourceCPU]
 	memory := res[corev1.ResourceMemory]