cmd/k8s-operator: add support for priorityClassName

Updates #8155

Signed-off-by: Gabriel Martinez <gabrielmartinez@sisti.pt>

cmd/k8s-operator/operator.go

@@ -65,6 +65,7 @@ func main() {
 		clientIDPath       = defaultEnv("CLIENT_ID_FILE", "")
 		clientSecretPath   = defaultEnv("CLIENT_SECRET_FILE", "")
 		image              = defaultEnv("PROXY_IMAGE", "tailscale/tailscale:latest")
+		priorityClassName  = defaultEnv("PROXY_PRIORITY_CLASS_NAME", "")
 		tags               = defaultEnv("PROXY_TAGS", "tag:k8s")
 		shouldRunAuthProxy = defaultBool("AUTH_PROXY", false)
 	)
@@ -201,12 +202,13 @@ waitOnline:
 	}
 
 	sr := &ServiceReconciler{
-		Client:            mgr.GetClient(),
+		Client:                 mgr.GetClient(),
-		tsClient:          tsClient,
+		tsClient:               tsClient,
-		defaultTags:       strings.Split(tags, ","),
+		defaultTags:            strings.Split(tags, ","),
-		operatorNamespace: tsNamespace,
+		operatorNamespace:      tsNamespace,
-		proxyImage:        image,
+		proxyImage:             image,
-		logger:            zlog.Named("service-reconciler"),
+		proxyPriorityClassName: priorityClassName,
+		logger:                 zlog.Named("service-reconciler"),
 	}
 
 	reconcileFilter := handler.EnqueueRequestsFromMapFunc(func(o client.Object) []reconcile.Request {
@@ -279,11 +281,12 @@ const (
 // ServiceReconciler is a simple ControllerManagedBy example implementation.
 type ServiceReconciler struct {
 	client.Client
-	tsClient          tsClient
+	tsClient               tsClient
-	defaultTags       []string
+	defaultTags            []string
-	operatorNamespace string
+	operatorNamespace      string
-	proxyImage        string
+	proxyImage             string
-	logger            *zap.SugaredLogger
+	proxyPriorityClassName string
+	logger                 *zap.SugaredLogger
 }
 
 type tsClient interface {
@@ -638,6 +641,7 @@ func (a *ServiceReconciler) reconcileSTS(ctx context.Context, logger *zap.Sugare
 	ss.Spec.Template.ObjectMeta.Labels = map[string]string{
 		"app": string(parentSvc.UID),
 	}
+	ss.Spec.Template.Spec.PriorityClassName = a.proxyPriorityClassName
 	logger.Debugf("reconciling statefulset %s/%s", ss.GetNamespace(), ss.GetName())
 	return createOrUpdate(ctx, a.Client, a.operatorNamespace, &ss, func(s *appsv1.StatefulSet) { s.Spec = ss.Spec })
 }

cmd/k8s-operator/operator_test.go

@@ -64,7 +64,7 @@ func TestLoadBalancerClass(t *testing.T) {
 
 	expectEqual(t, fc, expectedSecret(fullName))
 	expectEqual(t, fc, expectedHeadlessService(shortName))
-	expectEqual(t, fc, expectedSTS(shortName, fullName, "default-test"))
+	expectEqual(t, fc, expectedSTS(shortName, fullName, "default-test", ""))
 
 	// Normally the Tailscale proxy pod would come up here and write its info
 	// into the secret. Simulate that, then verify reconcile again and verify
@@ -185,7 +185,7 @@ func TestAnnotations(t *testing.T) {
 
 	expectEqual(t, fc, expectedSecret(fullName))
 	expectEqual(t, fc, expectedHeadlessService(shortName))
-	expectEqual(t, fc, expectedSTS(shortName, fullName, "default-test"))
+	expectEqual(t, fc, expectedSTS(shortName, fullName, "default-test", ""))
 	want := &corev1.Service{
 		TypeMeta: metav1.TypeMeta{
 			Kind:       "Service",
@@ -282,7 +282,7 @@ func TestAnnotationIntoLB(t *testing.T) {
 
 	expectEqual(t, fc, expectedSecret(fullName))
 	expectEqual(t, fc, expectedHeadlessService(shortName))
-	expectEqual(t, fc, expectedSTS(shortName, fullName, "default-test"))
+	expectEqual(t, fc, expectedSTS(shortName, fullName, "default-test", ""))
 
 	// Normally the Tailscale proxy pod would come up here and write its info
 	// into the secret. Simulate that, since it would have normally happened at
@@ -326,7 +326,7 @@ func TestAnnotationIntoLB(t *testing.T) {
 	expectReconciled(t, sr, "default", "test")
 	// None of the proxy machinery should have changed...
 	expectEqual(t, fc, expectedHeadlessService(shortName))
-	expectEqual(t, fc, expectedSTS(shortName, fullName, "default-test"))
+	expectEqual(t, fc, expectedSTS(shortName, fullName, "default-test", ""))
 	// ... but the service should have a LoadBalancer status.
 
 	want = &corev1.Service{
@@ -398,7 +398,7 @@ func TestLBIntoAnnotation(t *testing.T) {
 
 	expectEqual(t, fc, expectedSecret(fullName))
 	expectEqual(t, fc, expectedHeadlessService(shortName))
-	expectEqual(t, fc, expectedSTS(shortName, fullName, "default-test"))
+	expectEqual(t, fc, expectedSTS(shortName, fullName, "default-test", ""))
 
 	// Normally the Tailscale proxy pod would come up here and write its info
 	// into the secret. Simulate that, then verify reconcile again and verify
@@ -455,7 +455,7 @@ func TestLBIntoAnnotation(t *testing.T) {
 	expectReconciled(t, sr, "default", "test")
 
 	expectEqual(t, fc, expectedHeadlessService(shortName))
-	expectEqual(t, fc, expectedSTS(shortName, fullName, "default-test"))
+	expectEqual(t, fc, expectedSTS(shortName, fullName, "default-test", ""))
 
 	want = &corev1.Service{
 		TypeMeta: metav1.TypeMeta{
@@ -522,7 +522,7 @@ func TestCustomHostname(t *testing.T) {
 
 	expectEqual(t, fc, expectedSecret(fullName))
 	expectEqual(t, fc, expectedHeadlessService(shortName))
-	expectEqual(t, fc, expectedSTS(shortName, fullName, "reindeer-flotilla"))
+	expectEqual(t, fc, expectedSTS(shortName, fullName, "reindeer-flotilla", ""))
 	want := &corev1.Service{
 		TypeMeta: metav1.TypeMeta{
 			Kind:       "Service",
@@ -581,6 +581,51 @@ func TestCustomHostname(t *testing.T) {
 	expectEqual(t, fc, want)
 }
 
+func TestCustomPriorityClassName(t *testing.T) {
+	fc := fake.NewFakeClient()
+	ft := &fakeTSClient{}
+	zl, err := zap.NewDevelopment()
+	if err != nil {
+		t.Fatal(err)
+	}
+	sr := &ServiceReconciler{
+		Client:                 fc,
+		tsClient:               ft,
+		defaultTags:            []string{"tag:k8s"},
+		operatorNamespace:      "operator-ns",
+		proxyImage:             "tailscale/tailscale",
+		proxyPriorityClassName: "tailscale-critical",
+		logger:                 zl.Sugar(),
+	}
+
+	// Create a service that we should manage, and check that the initial round
+	// of objects looks right.
+	mustCreate(t, fc, &corev1.Service{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test",
+			Namespace: "default",
+			// The apiserver is supposed to set the UID, but the fake client
+			// doesn't. So, set it explicitly because other code later depends
+			// on it being set.
+			UID: types.UID("1234-UID"),
+			Annotations: map[string]string{
+				"tailscale.com/expose":   "true",
+				"tailscale.com/hostname": "custom-priority-class-name",
+			},
+		},
+		Spec: corev1.ServiceSpec{
+			ClusterIP: "10.20.30.40",
+			Type:      corev1.ServiceTypeClusterIP,
+		},
+	})
+
+	expectReconciled(t, sr, "default", "test")
+
+	fullName, shortName := findGenName(t, fc, "default", "test")
+
+	expectEqual(t, fc, expectedSTS(shortName, fullName, "custom-priority-class-name", "tailscale-critical"))
+}
+
 func expectedSecret(name string) *corev1.Secret {
 	return &corev1.Secret{
 		TypeMeta: metav1.TypeMeta{
@@ -629,7 +674,7 @@ func expectedHeadlessService(name string) *corev1.Service {
 	}
 }
 
-func expectedSTS(stsName, secretName, hostname string) *appsv1.StatefulSet {
+func expectedSTS(stsName, secretName, hostname, priorityClassName string) *appsv1.StatefulSet {
 	return &appsv1.StatefulSet{
 		TypeMeta: metav1.TypeMeta{
 			Kind:       "StatefulSet",
@@ -658,6 +703,7 @@ func expectedSTS(stsName, secretName, hostname string) *appsv1.StatefulSet {
 				},
 				Spec: corev1.PodSpec{
 					ServiceAccountName: "proxies",
+					PriorityClassName:  priorityClassName,
 					InitContainers: []corev1.Container{
 						{
 							Name:    "sysctler",