wip

api/v1beta2/statefulset_webhhok.go

@@ -0,0 +1,93 @@
+package v1beta2
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/cybozu-go/moco/pkg/constants"
+	admissionv1 "k8s.io/api/admission/v1"
+	appsv1 "k8s.io/api/apps/v1"
+	"k8s.io/apimachinery/pkg/api/equality"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/utils/ptr"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
+)
+
+func SetupStatefulSetWebhookWithManager(mgr ctrl.Manager) error {
+	return ctrl.NewWebhookManagedBy(mgr).
+		For(&appsv1.StatefulSet{}).
+		WithDefaulter(&StatefulSetDefaulter{}).
+		Complete()
+}
+
+//+kubebuilder:webhook:path=/mutate-apps-v1-statefulset,mutating=true,failurePolicy=fail,sideEffects=None,groups=apps,resources=statefulsets,verbs=create;update,versions=v1,name=statefulset.kb.io,admissionReviewVersions=v1
+
+type StatefulSetDefaulter struct{}
+
+var _ admission.CustomDefaulter = &StatefulSetDefaulter{}
+
+// Default implements webhook.Defaulter so a webhook will be registered for the type
+func (*StatefulSetDefaulter) Default(ctx context.Context, obj runtime.Object) error {
+	sts, ok := obj.(*appsv1.StatefulSet)
+	if !ok {
+		return fmt.Errorf("unknown obj type %T", obj)
+	}
+
+	req, err := admission.RequestFromContext(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to get admission request from context: %w", err)
+	}
+
+	if sts.Annotations[constants.AnnForceRollingUpdate] == "true" &&
+		sts.Spec.UpdateStrategy.RollingUpdate != nil &&
+		sts.Spec.UpdateStrategy.RollingUpdate.Partition != nil {
+		sts.Spec.UpdateStrategy.RollingUpdate = nil
+		return nil
+	}
+
+	if sts.Spec.UpdateStrategy.RollingUpdate == nil || sts.Spec.UpdateStrategy.RollingUpdate.Partition == nil {
+		sts.Spec.UpdateStrategy.RollingUpdate = &appsv1.RollingUpdateStatefulSetStrategy{
+			Partition: ptr.To[int32](*sts.Spec.Replicas),
+		}
+		return nil
+	}
+
+	if req.Operation != admissionv1.Update {
+		return nil
+	}
+
+	oldSts, err := readStatefulSet(req.OldObject.Raw)
+	if err != nil {
+		return fmt.Errorf("failed to read old statefulset: %w", err)
+	}
+
+	partition := *sts.Spec.UpdateStrategy.RollingUpdate.Partition
+	oldPartition := *oldSts.Spec.UpdateStrategy.RollingUpdate.Partition
+
+	newSts := sts.DeepCopy()
+	newSts.Spec.UpdateStrategy = oldSts.Spec.UpdateStrategy
+
+	if partition != oldPartition && equality.Semantic.DeepEqual(newSts, oldSts) {
+		return nil
+	}
+
+	sts.Spec.UpdateStrategy.RollingUpdate = &appsv1.RollingUpdateStatefulSetStrategy{
+		Partition: ptr.To[int32](partition),
+	}
+
+	return nil
+}
+
+func readStatefulSet(raw []byte) (*appsv1.StatefulSet, error) {
+	var sts appsv1.StatefulSet
+
+	if _, _, err := unstructured.UnstructuredJSONScheme.Decode(raw, nil, &sts); err != nil {
+		return nil, err
+	}
+
+	sts.TypeMeta.APIVersion = appsv1.SchemeGroupVersion.Group + "/" + appsv1.SchemeGroupVersion.Version
+
+	return &sts, nil
+}

charts/moco/templates/generated/generated.yaml

@@ -155,6 +155,12 @@ rules:
       - patch
       - update
       - watch
+  - apiGroups:
+      - ""
+    resources:
+      - pods/status
+    verbs:
+      - get
   - apiGroups:
       - ""
     resources:
@@ -460,6 +466,26 @@ webhooks:
         resources:
           - mysqlclusters
     sideEffects: None
+  - admissionReviewVersions:
+      - v1
+    clientConfig:
+      service:
+        name: moco-webhook-service
+        namespace: '{{ .Release.Namespace }}'
+        path: /mutate-apps-v1-statefulset
+    failurePolicy: Fail
+    name: statefulset.kb.io
+    rules:
+      - apiGroups:
+          - apps
+        apiVersions:
+          - v1
+        operations:
+          - CREATE
+          - UPDATE
+        resources:
+          - statefulsets
+    sideEffects: None
 ---
 apiVersion: admissionregistration.k8s.io/v1
 kind: ValidatingWebhookConfiguration

cmd/moco-controller/cmd/run.go

@@ -113,6 +113,15 @@ func subMain(ns, addr string, port int) error {
 		return err
 	}
 
+	if err = (&controllers.StatefulSetPartitionReconciler{
+		Client:                  mgr.GetClient(),
+		Recorder:                mgr.GetEventRecorderFor("moco-controller"),
+		MaxConcurrentReconciles: config.maxConcurrentReconciles,
+	}).SetupWithManager(mgr); err != nil {
+		setupLog.Error(err, "unable to create controller", "controller", "Partition")
+		return err
+	}
+
 	if err = (&controllers.PodWatcher{
 		Client:                  mgr.GetClient(),
 		ClusterManager:          clusterMgr,
@@ -132,6 +141,11 @@ func subMain(ns, addr string, port int) error {
 		return err
 	}
 
+	if err := mocov1beta2.SetupStatefulSetWebhookWithManager(mgr); err != nil {
+		setupLog.Error(err, "unable to setup webhook", "webhook", "StatefulSet")
+		return err
+	}
+
 	if err := mgr.AddHealthzCheck("health", healthz.Ping); err != nil {
 		setupLog.Error(err, "unable to set up health check")
 		return err

config/rbac/role.yaml

@@ -50,6 +50,12 @@ rules:
   - patch
   - update
   - watch
+- apiGroups:
+  - ""
+  resources:
+  - pods/status
+  verbs:
+  - get
 - apiGroups:
   - ""
   resources:

config/webhook/manifests.yaml

@@ -24,6 +24,26 @@ webhooks:
     resources:
     - mysqlclusters
   sideEffects: None
+- admissionReviewVersions:
+  - v1
+  clientConfig:
+    service:
+      name: webhook-service
+      namespace: system
+      path: /mutate-apps-v1-statefulset
+  failurePolicy: Fail
+  name: statefulset.kb.io
+  rules:
+  - apiGroups:
+    - apps
+    apiVersions:
+    - v1
+    operations:
+    - CREATE
+    - UPDATE
+    resources:
+    - statefulsets
+  sideEffects: None
 ---
 apiVersion: admissionregistration.k8s.io/v1
 kind: ValidatingWebhookConfiguration