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,94 @@
+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/client-go/rest"
+	"k8s.io/test-infra/prow/interrupts"
+	"sigs.k8s.io/controller-runtime/pkg/manager"
+	"sigs.k8s.io/controller-runtime/pkg/webhook"
+	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
+
+	buildv1 "github.com/openshift/api/build/v1"
+	buildclientset "github.com/openshift/client-go/build/clientset/versioned/typed/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) {
+	logger := logrus.WithField("component", "admission")
+	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.")
+	}
+	mgr, err := manager.New(restConfig, manager.Options{Port: port})
+	if err != nil {
+		logrus.WithError(err).Fatal("Could not start manager.")
+	}
+	server := mgr.GetWebhookServer()
+	server.Register("/pods", &webhook.Admission{Handler: &podMutator{logger: logger, client: client}})
+	if err := mgr.Start(interrupts.Context()); err != nil {
+		logrus.WithError(err).Error("Failed to serve admission responses.")
+	}
+	logger.Debug("Ready to serve HTTP requests.")
+}
+
+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.")
+	}
+	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
+		}
+	}
+
+	marshaledPod, err := json.Marshal(pod)
+	if err != nil {
+		return admission.Errored(http.StatusInternalServerError, err)
+	}
+
+	return admission.PatchResponseFromRaw(req.Object.Raw, marshaledPod)
+}
+
+//nolint:unparam
+func (m *podMutator) InjectDecoder(d *admission.Decoder) error {
+	m.decoder = d
+	return nil
+}

cmd/pod-scaler/main.go

@@ -47,7 +47,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 +63,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 +112,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 +154,7 @@ func mainProduce(opts *options, cache cache) {
 
 	go produce(clients, cache)
 }
+
+func mainAdmission(opts *options) {
+	go admit(opts.port)
+}

cmd/pod-scaler/testdata/zz_fixture_TestMutatePods_pod_associated_with_a_build_that_has_no_labels.yaml

@@ -0,0 +1 @@
+{"response":{"uid":"705ab4f5-6393-11e8-b7cc-42010a800002","allowed":true}}

cmd/pod-scaler/testdata/zz_fixture_TestMutatePods_pod_associated_with_a_build_with_labels.yaml

@@ -0,0 +1 @@
+{"response":{"uid":"705ab4f5-6393-11e8-b7cc-42010a800002","allowed":true,"patch":"eyJtZXRhZGF0YSI6eyJsYWJlbHMiOnsiY2kub3BlbnNoaWZ0LmlvL21ldGFkYXRhLmJyYW5jaCI6ImJyYW5jaCIsImNpLm9wZW5zaGlmdC5pby9tZXRhZGF0YS5vcmciOiJvcmciLCJjaS5vcGVuc2hpZnQuaW8vbWV0YWRhdGEucmVwbyI6InJlcG8iLCJjaS5vcGVuc2hpZnQuaW8vbWV0YWRhdGEuc3RlcCI6InN0ZXAiLCJjaS5vcGVuc2hpZnQuaW8vbWV0YWRhdGEudGFyZ2V0IjoidGFyZ2V0IiwiY2kub3BlbnNoaWZ0LmlvL21ldGFkYXRhLnZhcmlhbnQiOiJ2YXJpYW50In19fQ==","patchType":"JSONPatch"}}

cmd/pod-scaler/testdata/zz_fixture_TestMutatePods_pod_not_associated_with_a_build.yaml

@@ -0,0 +1 @@
+{"response":{"uid":"705ab4f5-6393-11e8-b7cc-42010a800002","allowed":true}}