pod-scaler: add labels to Build Pods via admission

cmd/pod-scaler/admission.go

@@ -0,0 +1,85 @@
+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,194 @@
+package main
+
+import (
+	"context"
+	"sort"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"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"
+
+	buildv1 "github.com/openshift/api/build/v1"
+	fakebuildv1client "github.com/openshift/client-go/build/clientset/versioned/fake"
+
+	"github.com/openshift/ci-tools/pkg/testhelper"
+)
+
+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

@@ -14,9 +14,11 @@ import (
 	"gopkg.in/fsnotify.v1"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/rest"
 	"k8s.io/client-go/transport"
 	"k8s.io/test-infra/prow/interrupts"
 
+	buildclientset "github.com/openshift/client-go/build/clientset/versioned/typed/build/v1"
 	routeclientset "github.com/openshift/client-go/route/clientset/versioned/typed/route/v1"
 
 	"github.com/openshift/ci-tools/pkg/util"
@@ -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: ""