pod-scaler: add labels to Build Pods via admission

We need labels on Pods to associate individual executions with their
semantic category. Today, we label Builds (as that's what we create) but
the Build subsystem does not carry those labels forward to the Pods that
implement the Build. This patch adds a new mutating admission webhook
server that will add the labels where necessary.

Signed-off-by: Steve Kuznetsov <skuznets@redhat.com>

cmd/pod-scaler/admission.go

@@ -0,0 +1,84 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+
+	"github.com/sirupsen/logrus"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/test-infra/prow/interrupts"
+	"k8s.io/test-infra/prow/pjutil"
+	"sigs.k8s.io/controller-runtime/pkg/webhook"
+	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
+
+	buildv1 "github.com/openshift/api/build/v1"
+	buildclientv1 "github.com/openshift/client-go/build/clientset/versioned/typed/build/v1"
+
+	"github.com/openshift/ci-tools/pkg/steps"
+)
+
+func admit(port int, client buildclientv1.BuildV1Interface) {
+	logger := logrus.WithField("component", "admission")
+	health := pjutil.NewHealth()
+	health.ServeReady()
+	httpServer := webhook.Server{Port: port}
+	httpServer.Register("/pods", &webhook.Admission{Handler: &podMutator{logger: logger, client: client}})
+	if err := httpServer.StartStandalone(interrupts.Context(), nil); err != nil {
+		logrus.WithError(err).Error("Failed to serve admission webhooks.")
+	}
+}
+
+type podMutator struct {
+	logger  *logrus.Entry
+	client  buildclientv1.BuildV1Interface
+	decoder *admission.Decoder
+}
+
+func (m *podMutator) Handle(ctx context.Context, req admission.Request) admission.Response {
+	pod := &corev1.Pod{}
+
+	err := m.decoder.Decode(req, pod)
+	if err != nil {
+		return admission.Errored(http.StatusBadRequest, err)
+	}
+	buildName, isBuildPod := pod.Labels[buildv1.BuildLabel]
+	if !isBuildPod {
+		return admission.Allowed("Not a Pod implementing a Build.")
+	}
+	logger := m.logger.WithField("build", buildName)
+	logger.Debug("Handling labels on Pod created for a Build.")
+	build, err := m.client.Builds(pod.Namespace).Get(ctx, buildName, metav1.GetOptions{})
+	if err != nil {
+		logger.WithError(err).Error("Could not get Build for Pod.")
+		return admission.Allowed("Could not get Build for Pod, ignoring.")
+	}
+	mutatePod(pod, build)
+
+	marshaledPod, err := json.Marshal(pod)
+	if err != nil {
+		return admission.Errored(http.StatusInternalServerError, err)
+	}
+
+	return admission.PatchResponseFromRaw(req.Object.Raw, marshaledPod)
+}
+
+func mutatePod(pod *corev1.Pod, build *buildv1.Build) {
+	if pod.Labels == nil {
+		pod.Labels = map[string]string{}
+	}
+	for _, label := range []string{steps.LabelMetadataOrg, steps.LabelMetadataRepo, steps.LabelMetadataBranch, steps.LabelMetadataVariant, steps.LabelMetadataTarget, steps.LabelMetadataStep} {
+		buildValue, buildHas := build.Labels[label]
+		_, podHas := pod.Labels[label]
+		if buildHas && !podHas {
+			pod.Labels[label] = buildValue
+		}
+	}
+}
+
+//nolint:unparam
+func (m *podMutator) InjectDecoder(d *admission.Decoder) error {
+	m.decoder = d
+	return nil
+}

cmd/pod-scaler/admission_test.go

@@ -0,0 +1,191 @@
+package main
+
+import (
+	"context"
+	"sort"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	buildv1 "github.com/openshift/api/build/v1"
+	"github.com/openshift/ci-tools/pkg/testhelper"
+	fakebuildv1client "github.com/openshift/client-go/build/clientset/versioned/fake"
+	"github.com/sirupsen/logrus"
+	admissionv1 "k8s.io/api/admission/v1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/client-go/kubernetes/scheme"
+	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
+)
+
+func TestMutatePods(t *testing.T) {
+	client := fakebuildv1client.NewSimpleClientset(
+		&buildv1.Build{
+			TypeMeta: metav1.TypeMeta{
+				Kind:       "Build",
+				APIVersion: "build.openshift.io/v1",
+			},
+			ObjectMeta: metav1.ObjectMeta{
+				Namespace: "namespace",
+				Name:      "withoutlabels",
+				Labels:    map[string]string{},
+			},
+		},
+		&buildv1.Build{
+			TypeMeta: metav1.TypeMeta{
+				Kind:       "Build",
+				APIVersion: "build.openshift.io/v1",
+			},
+			ObjectMeta: metav1.ObjectMeta{
+				Namespace: "namespace",
+				Name:      "withlabels",
+				Labels: map[string]string{
+					"ci.openshift.io/metadata.org":     "org",
+					"ci.openshift.io/metadata.repo":    "repo",
+					"ci.openshift.io/metadata.branch":  "branch",
+					"ci.openshift.io/metadata.variant": "variant",
+					"ci.openshift.io/metadata.target":  "target",
+					"ci.openshift.io/metadata.step":    "step",
+				},
+			},
+		},
+	)
+	decoder, err := admission.NewDecoder(scheme.Scheme)
+	if err != nil {
+		t.Fatalf("failed to create decoder from scheme: %v", err)
+	}
+	mutator := podMutator{
+		logger:  logrus.WithField("test", t.Name()),
+		client:  client.BuildV1(),
+		decoder: decoder,
+	}
+
+	var testCases = []struct {
+		name    string
+		request admission.Request
+	}{
+		{
+			name: "not a pod",
+			request: admission.Request{
+				AdmissionRequest: admissionv1.AdmissionRequest{
+					UID:      "705ab4f5-6393-11e8-b7cc-42010a800002",
+					Kind:     metav1.GroupVersionKind{Group: "", Version: "v1", Kind: "Secret"},
+					Resource: metav1.GroupVersionResource{Group: "", Version: "v1", Resource: "secrets"},
+					Object:   runtime.RawExtension{Raw: []byte(`{"apiVersion": "v1","kind": "Secret","metadata": {"name": "somethingelse","namespace": "namespace"}}`)},
+				},
+			},
+		},
+		{
+			name: "pod not associated with a build",
+			request: admission.Request{
+				AdmissionRequest: admissionv1.AdmissionRequest{
+					UID:      "705ab4f5-6393-11e8-b7cc-42010a800002",
+					Kind:     metav1.GroupVersionKind{Group: "", Version: "v1", Kind: "Pod"},
+					Resource: metav1.GroupVersionResource{Group: "", Version: "v1", Resource: "pods"},
+					Object:   runtime.RawExtension{Raw: []byte(`{"apiVersion": "v1","kind": "Pod","metadata": {"creationTimestamp": null, "name": "somethingelse","namespace": "namespace"}, "spec":{"containers":[]}, "status":{}}`)},
+				},
+			},
+		},
+		{
+			name: "pod associated with a build that has no labels",
+			request: admission.Request{
+				AdmissionRequest: admissionv1.AdmissionRequest{
+					UID:      "705ab4f5-6393-11e8-b7cc-42010a800002",
+					Kind:     metav1.GroupVersionKind{Group: "", Version: "v1", Kind: "Pod"},
+					Resource: metav1.GroupVersionResource{Group: "", Version: "v1", Resource: "pods"},
+					Object:   runtime.RawExtension{Raw: []byte(`{"apiVersion": "v1","kind": "Pod","metadata": {"creationTimestamp": null, "labels": {"openshift.io/build.name": "withoutlabels"}, "name": "withoutlabels-build","namespace": "namespace"}, "spec":{"containers":[]}, "status":{}}`)},
+				},
+			},
+		},
+		{
+			name: "pod associated with a build with labels",
+			request: admission.Request{
+				AdmissionRequest: admissionv1.AdmissionRequest{
+					UID:      "705ab4f5-6393-11e8-b7cc-42010a800002",
+					Kind:     metav1.GroupVersionKind{Group: "", Version: "v1", Kind: "Pod"},
+					Resource: metav1.GroupVersionResource{Group: "", Version: "v1", Resource: "pods"},
+					Object:   runtime.RawExtension{Raw: []byte(`{"apiVersion": "v1","kind": "Pod","metadata": {"creationTimestamp": null, "labels": {"openshift.io/build.name": "withlabels"}, "name": "withoutlabels-build","namespace": "namespace"}, "spec":{"containers":[]}, "status":{}}`)},
+				},
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			response := mutator.Handle(context.Background(), testCase.request)
+			sort.Slice(response.Patches, func(i, j int) bool {
+				return response.Patches[i].Path < response.Patches[j].Path
+			})
+			testhelper.CompareWithFixture(t, response)
+		})
+	}
+}
+
+func TestMutatePod(t *testing.T) {
+	var testCases = []struct {
+		name     string
+		build    *buildv1.Build
+		pod      *corev1.Pod
+		expected *corev1.Pod
+	}{
+		{
+			name:     "no labels to add",
+			build:    &buildv1.Build{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{}}},
+			pod:      &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{}}},
+			expected: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{}}},
+		},
+		{
+			name: "many labels to add",
+			build: &buildv1.Build{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{
+				"ci.openshift.io/metadata.org":     "org",
+				"ci.openshift.io/metadata.repo":    "repo",
+				"ci.openshift.io/metadata.branch":  "branch",
+				"ci.openshift.io/metadata.variant": "variant",
+				"ci.openshift.io/metadata.target":  "target",
+				"ci.openshift.io/metadata.step":    "step",
+			}}},
+			pod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{}}},
+			expected: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{
+				"ci.openshift.io/metadata.org":     "org",
+				"ci.openshift.io/metadata.repo":    "repo",
+				"ci.openshift.io/metadata.branch":  "branch",
+				"ci.openshift.io/metadata.variant": "variant",
+				"ci.openshift.io/metadata.target":  "target",
+				"ci.openshift.io/metadata.step":    "step",
+			}}},
+		},
+		{
+			name: "some labels to add",
+			build: &buildv1.Build{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{
+				"ci.openshift.io/metadata.org":     "org",
+				"ci.openshift.io/metadata.repo":    "repo",
+				"ci.openshift.io/metadata.branch":  "branch",
+				"ci.openshift.io/metadata.variant": "variant",
+				"ci.openshift.io/metadata.target":  "target",
+				"ci.openshift.io/metadata.step":    "step",
+			}}},
+			pod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{
+				"ci.openshift.io/metadata.org":  "org",
+				"ci.openshift.io/metadata.repo": "repo",
+				"ci.openshift.io/metadata.step": "step",
+			}}},
+			expected: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{
+				"ci.openshift.io/metadata.org":     "org",
+				"ci.openshift.io/metadata.repo":    "repo",
+				"ci.openshift.io/metadata.branch":  "branch",
+				"ci.openshift.io/metadata.variant": "variant",
+				"ci.openshift.io/metadata.target":  "target",
+				"ci.openshift.io/metadata.step":    "step",
+			}}},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			mutatePod(testCase.pod, testCase.build)
+			if diff := cmp.Diff(testCase.pod, testCase.expected); diff != "" {
+				t.Errorf("%s: got incorrect pod after mutation: %v", testCase.name, diff)
+			}
+		})
+	}
+}

cmd/pod-scaler/main.go

@@ -7,11 +7,13 @@ import (
 	"os"
 
 	"cloud.google.com/go/storage"
+	buildclientset "github.com/openshift/client-go/build/clientset/versioned/typed/build/v1"
 	prometheusclient "github.com/prometheus/client_golang/api"
 	prometheusapi "github.com/prometheus/client_golang/api/prometheus/v1"
 	"github.com/sirupsen/logrus"
 	"google.golang.org/api/option"
 	"gopkg.in/fsnotify.v1"
+	"k8s.io/client-go/rest"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/transport"
@@ -47,7 +49,7 @@ func bindOptions(fs *flag.FlagSet) *options {
 	o := options{producerOptions: producerOptions{}}
 	fs.StringVar(&o.mode, "mode", "", "Which mode to run in.")
 	fs.StringVar(&o.kubeconfig, "kubeconfig", "", "Path to a ~/.kube/config to use for querying Prometheuses. Each context will be considered a cluster to query.")
-	fs.IntVar(&o.port, "port", 0, "Port to serve requirements on.")
+	fs.IntVar(&o.port, "port", 0, "Port to serve admission webhooks on.")
 	fs.IntVar(&o.uiPort, "ui-port", 0, "Port to serve frontend on.")
 	fs.StringVar(&o.loglevel, "loglevel", "debug", "Logging level.")
 	fs.StringVar(&o.cacheDir, "cache-dir", "", "Local directory holding cache data (for development mode).")
@@ -63,15 +65,16 @@ func (o *options) validate() error {
 		if o.kubeconfig == "" && !kubeconfigSet {
 			return errors.New("--kubeconfig or $KUBECONFIG is required")
 		}
-	case "consumer":
-		if o.port == 0 {
-			return errors.New("--port is required")
-		}
+	case "consumer.ui":
 		if o.uiPort == 0 {
 			return errors.New("--ui-port is required")
 		}
+	case "consumer.admission":
+		if o.port == 0 {
+			return errors.New("--port is required")
+		}
 	default:
-		return errors.New("--mode must be either \"producer\" or \"consumer\"")
+		return errors.New("--mode must be either \"producer\", \"consumer.ui\", or \"consumer.admission\"")
 	}
 	if o.cacheDir == "" {
 		if o.cacheBucket == "" {
@@ -111,8 +114,10 @@ func main() {
 	switch opts.mode {
 	case "producer":
 		mainProduce(opts, cache)
-	case "consumer":
+	case "consumer.ui":
 		// TODO
+	case "consumer.admission":
+		mainAdmission(opts)
 	}
 	interrupts.WaitForGracefulShutdown()
 }
@@ -151,3 +156,15 @@ func mainProduce(opts *options, cache cache) {
 
 	go produce(clients, cache)
 }
+
+func mainAdmission(opts *options) {
+	restConfig, err := rest.InClusterConfig()
+	if err != nil {
+		logrus.WithError(err).Fatal("Failed to load in-cluster config.")
+	}
+	client, err := buildclientset.NewForConfig(restConfig)
+	if err != nil {
+		logrus.WithError(err).Fatal("Failed to construct client.")
+	}
+	go admit(opts.port, client)
+}

cmd/pod-scaler/testdata/zz_fixture_TestMutatePods_not_a_pod.yaml

@@ -0,0 +1,7 @@
+Patches: null
+allowed: false
+status:
+  code: 400
+  message: unable to decode /v1, Kind=Secret into *v1.Pod
+  metadata: {}
+uid: ""

cmd/pod-scaler/testdata/zz_fixture_TestMutatePods_pod_associated_with_a_build_that_has_no_labels.yaml

@@ -0,0 +1,3 @@
+Patches: []
+allowed: true
+uid: ""

cmd/pod-scaler/testdata/zz_fixture_TestMutatePods_pod_associated_with_a_build_with_labels.yaml

@@ -0,0 +1,22 @@
+Patches:
+- op: add
+  path: /metadata/labels/ci.openshift.io~1metadata.branch
+  value: branch
+- op: add
+  path: /metadata/labels/ci.openshift.io~1metadata.org
+  value: org
+- op: add
+  path: /metadata/labels/ci.openshift.io~1metadata.repo
+  value: repo
+- op: add
+  path: /metadata/labels/ci.openshift.io~1metadata.step
+  value: step
+- op: add
+  path: /metadata/labels/ci.openshift.io~1metadata.target
+  value: target
+- op: add
+  path: /metadata/labels/ci.openshift.io~1metadata.variant
+  value: variant
+allowed: true
+patchType: JSONPatch
+uid: ""

cmd/pod-scaler/testdata/zz_fixture_TestMutatePods_pod_not_associated_with_a_build.yaml

@@ -0,0 +1,7 @@
+Patches: null
+allowed: true
+status:
+  code: 200
+  metadata: {}
+  reason: Not a Pod implementing a Build.
+uid: ""