[k8s, k8s multicluster plugin] Implement health state for statefulset

pkg/app/pipedv1/plugin/kubernetes/provider/health_status.go

@@ -30,6 +30,12 @@ func (m Manifest) calculateHealthStatus() (sdk.ResourceHealthStatus, string) {
 			return sdk.ResourceHealthStateUnknown, ""
 		}
 		return deploymentHealthStatus(obj)
+	case m.IsStatefulSet():
+		obj := &appsv1.StatefulSet{}
+		if err := m.ConvertToStructuredObject(obj); err != nil {
+			return sdk.ResourceHealthStateUnknown, ""
+		}
+		return statefulSetHealthStatus(obj)
 	default:
 		// TODO: Implement health status calculation for other resource types.
 		return sdk.ResourceHealthStateUnknown, fmt.Sprintf("Unimplemented or unknown resource: %s", m.body.GroupVersionKind())
@@ -66,3 +72,35 @@ func deploymentHealthStatus(obj *appsv1.Deployment) (sdk.ResourceHealthStatus, s
 	}
 	return sdk.ResourceHealthStateHealthy, ""
 }
+
+func statefulSetHealthStatus(obj *appsv1.StatefulSet) (sdk.ResourceHealthStatus, string) {
+	// Referred to:
+	//   https://github.com/kubernetes/kubernetes/blob/7942dca975b7be9386540df3c17e309c3cb2de60/staging/src/k8s.io/kubectl/pkg/polymorphichelpers/rollout_status.go#L130-L149
+	if obj.Status.ObservedGeneration == 0 || obj.Generation > obj.Status.ObservedGeneration {
+		return sdk.ResourceHealthStateUnhealthy, "Waiting for statefulset spec update to be observed"
+	}
+
+	if obj.Spec.Replicas == nil {
+		return sdk.ResourceHealthStateUnhealthy, "The number of desired replicas is unspecified"
+	}
+	if *obj.Spec.Replicas != obj.Status.ReadyReplicas {
+		return sdk.ResourceHealthStateUnhealthy, fmt.Sprintf("The number of ready replicas (%d) is different from the desired number (%d)", obj.Status.ReadyReplicas, *obj.Spec.Replicas)
+	}
+
+	// Check if the partitioned roll out is in progress.
+	if obj.Spec.UpdateStrategy.Type == appsv1.RollingUpdateStatefulSetStrategyType && obj.Spec.UpdateStrategy.RollingUpdate != nil {
+		if obj.Spec.Replicas != nil && obj.Spec.UpdateStrategy.RollingUpdate.Partition != nil {
+			if obj.Status.UpdatedReplicas < (*obj.Spec.Replicas - *obj.Spec.UpdateStrategy.RollingUpdate.Partition) {
+				return sdk.ResourceHealthStateUnhealthy, fmt.Sprintf("Waiting for partitioned roll out to finish because %d out of %d new pods have been updated",
+					obj.Status.UpdatedReplicas, (*obj.Spec.Replicas - *obj.Spec.UpdateStrategy.RollingUpdate.Partition))
+			}
+		}
+		return sdk.ResourceHealthStateHealthy, ""
+	}
+
+	if obj.Status.UpdateRevision != obj.Status.CurrentRevision {
+		return sdk.ResourceHealthStateUnhealthy, fmt.Sprintf("Waiting for statefulset rolling update to complete %d pods at revision %s", obj.Status.UpdatedReplicas, obj.Status.UpdateRevision)
+	}
+
+	return sdk.ResourceHealthStateHealthy, ""
+}

pkg/app/pipedv1/plugin/kubernetes/provider/health_status_test.go

@@ -125,3 +125,113 @@ func TestDeploymentHealthStatus(t *testing.T) {
 		})
 	}
 }
+
+func int32Ptr(i int32) *int32 { return &i }
+
+func Test_statefulSetHealthStatus(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name    string
+		obj     *appsv1.StatefulSet
+		want    sdk.ResourceHealthStatus
+		wantMsg string
+	}{
+		{
+			name: "ObservedGeneration is zero",
+			obj: &appsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{Generation: 2},
+				Status:     appsv1.StatefulSetStatus{ObservedGeneration: 0},
+			},
+			want:    sdk.ResourceHealthStateUnhealthy,
+			wantMsg: "Waiting for statefulset spec update to be observed",
+		},
+		{
+			name: "Generation > ObservedGeneration",
+			obj: &appsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{Generation: 2},
+				Status:     appsv1.StatefulSetStatus{ObservedGeneration: 1},
+			},
+			want:    sdk.ResourceHealthStateUnhealthy,
+			wantMsg: "Waiting for statefulset spec update to be observed",
+		},
+		{
+			name: "Replicas is nil",
+			obj: &appsv1.StatefulSet{
+				Status: appsv1.StatefulSetStatus{ObservedGeneration: 1},
+			},
+			want:    sdk.ResourceHealthStateUnhealthy,
+			wantMsg: "The number of desired replicas is unspecified",
+		},
+		{
+			name: "ReadyReplicas != Spec.Replicas",
+			obj: &appsv1.StatefulSet{
+				Spec:   appsv1.StatefulSetSpec{Replicas: int32Ptr(3)},
+				Status: appsv1.StatefulSetStatus{ObservedGeneration: 1, ReadyReplicas: 2},
+			},
+			want:    sdk.ResourceHealthStateUnhealthy,
+			wantMsg: "The number of ready replicas (2) is different from the desired number (3)",
+		},
+		{
+			name: "Partitioned rollout in progress",
+			obj: &appsv1.StatefulSet{
+				Spec: appsv1.StatefulSetSpec{
+					Replicas: int32Ptr(5),
+					UpdateStrategy: appsv1.StatefulSetUpdateStrategy{
+						Type: appsv1.RollingUpdateStatefulSetStrategyType,
+						RollingUpdate: &appsv1.RollingUpdateStatefulSetStrategy{
+							Partition: int32Ptr(2),
+						},
+					},
+				},
+				Status: appsv1.StatefulSetStatus{
+					ObservedGeneration: 1,
+					ReadyReplicas:      5,
+					UpdatedReplicas:    2,
+				},
+			},
+			want:    sdk.ResourceHealthStateUnhealthy,
+			wantMsg: "Waiting for partitioned roll out to finish because 2 out of 3 new pods have been updated",
+		},
+		{
+			name: "UpdateRevision != CurrentRevision",
+			obj: &appsv1.StatefulSet{
+				Spec: appsv1.StatefulSetSpec{Replicas: int32Ptr(2)},
+				Status: appsv1.StatefulSetStatus{
+					ObservedGeneration: 1,
+					ReadyReplicas:      2,
+					UpdateRevision:     "rev2",
+					CurrentRevision:    "rev1",
+					UpdatedReplicas:    2,
+				},
+			},
+			want:    sdk.ResourceHealthStateUnhealthy,
+			wantMsg: "Waiting for statefulset rolling update to complete 2 pods at revision rev2",
+		},
+		{
+			name: "Healthy statefulset",
+			obj: &appsv1.StatefulSet{
+				Spec: appsv1.StatefulSetSpec{Replicas: int32Ptr(2)},
+				Status: appsv1.StatefulSetStatus{
+					ObservedGeneration: 1,
+					ReadyReplicas:      2,
+					UpdateRevision:     "rev1",
+					CurrentRevision:    "rev1",
+					UpdatedReplicas:    2,
+				},
+			},
+			want:    sdk.ResourceHealthStateHealthy,
+			wantMsg: "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			got, gotMsg := statefulSetHealthStatus(tt.obj)
+			assert.Equal(t, tt.want, got)
+			assert.Equal(t, tt.wantMsg, gotMsg)
+		})
+	}
+}

pkg/app/pipedv1/plugin/kubernetes/provider/manifest.go

@@ -145,7 +145,7 @@ func (m Manifest) IsWorkload() bool {
 	}
 
 	switch m.body.GetKind() {
-	case KindDeployment, KindReplicaSet, KindDaemonSet, KindPod:
+	case KindDeployment, KindReplicaSet, KindDaemonSet, KindPod, KindStatefulSet:
 		return true
 	default:
 		return false
@@ -166,6 +166,13 @@ func (m Manifest) IsDeployment() bool {
 	return isBuiltinAPIGroup(m.body.GroupVersionKind().Group) && m.body.GetKind() == KindDeployment
 }
 
+// IsStatefulSet returns true if the manifest is a StatefulSet.
+// It checks the API group and the kind of the manifest.
+func (m Manifest) IsStatefulSet() bool {
+	// TODO: check the API group more strictly.
+	return isBuiltinAPIGroup(m.body.GroupVersionKind().Group) && m.body.GetKind() == KindStatefulSet
+}
+
 // IsSecret returns true if the manifest is a Secret.
 // It checks the API group and the kind of the manifest.
 func (m Manifest) IsSecret() bool {

pkg/app/pipedv1/plugin/kubernetes/provider/manifest_test.go

@@ -631,6 +631,91 @@ spec:
 	}
 }
 
+func TestManifest_IsStatefulSet(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name     string
+		manifest string
+		want     bool
+	}{
+		{
+			name: "is statefulset",
+			manifest: `
+apiVersion: apps/v1
+kind: StatefulSet
+metadata:
+  name: my-statefulset
+  namespace: default
+spec:
+  serviceName: "nginx"
+  replicas: 3
+  selector:
+    matchLabels:
+      app: nginx
+  template:
+    metadata:
+      labels:
+        app: nginx
+    spec:
+      containers:
+      - name: nginx
+        image: nginx:1.19.3
+`,
+			want: true,
+		},
+		{
+			name: "is not statefulset",
+			manifest: `
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: nginx-deployment
+spec:
+  template:
+    spec:
+      containers:
+      - name: nginx
+        image: nginx:1.19.3
+`,
+			want: false,
+		},
+		{
+			name: "is not statefulset with custom apigroup",
+			manifest: `
+apiVersion: custom.io/v1
+kind: StatefulSet
+metadata:
+  name: custom-statefulset
+spec:
+  serviceName: "custom"
+  replicas: 1
+  selector:
+    matchLabels:
+      app: custom
+  template:
+    metadata:
+      labels:
+        app: custom
+    spec:
+      containers:
+      - name: custom
+        image: custom:1.0.0
+`,
+			want: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			manifest := mustParseManifests(t, strings.TrimSpace(tt.manifest))[0]
+			got := manifest.IsStatefulSet()
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}
+
 func TestManifest_IsSecret(t *testing.T) {
 	t.Parallel()
 

pkg/app/pipedv1/plugin/kubernetes/provider/resource.go

@@ -27,10 +27,11 @@ const (
 	KindService = "Service"
 
 	// Workload
-	KindDeployment = "Deployment"
-	KindReplicaSet = "ReplicaSet"
-	KindDaemonSet  = "DaemonSet"
-	KindPod        = "Pod"
+	KindDeployment  = "Deployment"
+	KindReplicaSet  = "ReplicaSet"
+	KindDaemonSet   = "DaemonSet"
+	KindPod         = "Pod"
+	KindStatefulSet = "StatefulSet"
 
 	// ConfigMap and Secret
 	KindSecret    = "Secret"

pkg/app/pipedv1/plugin/kubernetes_multicluster/provider/health_status.go

@@ -30,6 +30,12 @@ func (m Manifest) calculateHealthStatus() (sdk.ResourceHealthStatus, string) {
 			return sdk.ResourceHealthStateUnknown, ""
 		}
 		return deploymentHealthStatus(obj)
+	case m.IsStatefulSet():
+		obj := &appsv1.StatefulSet{}
+		if err := m.ConvertToStructuredObject(obj); err != nil {
+			return sdk.ResourceHealthStateUnknown, ""
+		}
+		return statefulSetHealthStatus(obj)
 	default:
 		// TODO: Implement health status calculation for other resource types.
 		return sdk.ResourceHealthStateUnknown, fmt.Sprintf("Unimplemented or unknown resource: %s", m.body.GroupVersionKind())
@@ -66,3 +72,35 @@ func deploymentHealthStatus(obj *appsv1.Deployment) (sdk.ResourceHealthStatus, s
 	}
 	return sdk.ResourceHealthStateHealthy, ""
 }
+
+func statefulSetHealthStatus(obj *appsv1.StatefulSet) (sdk.ResourceHealthStatus, string) {
+	// Referred to:
+	//   https://github.com/kubernetes/kubernetes/blob/7942dca975b7be9386540df3c17e309c3cb2de60/staging/src/k8s.io/kubectl/pkg/polymorphichelpers/rollout_status.go#L130-L149
+	if obj.Status.ObservedGeneration == 0 || obj.Generation > obj.Status.ObservedGeneration {
+		return sdk.ResourceHealthStateUnhealthy, "Waiting for statefulset spec update to be observed"
+	}
+
+	if obj.Spec.Replicas == nil {
+		return sdk.ResourceHealthStateUnhealthy, "The number of desired replicas is unspecified"
+	}
+	if *obj.Spec.Replicas != obj.Status.ReadyReplicas {
+		return sdk.ResourceHealthStateUnhealthy, fmt.Sprintf("The number of ready replicas (%d) is different from the desired number (%d)", obj.Status.ReadyReplicas, *obj.Spec.Replicas)
+	}
+
+	// Check if the partitioned roll out is in progress.
+	if obj.Spec.UpdateStrategy.Type == appsv1.RollingUpdateStatefulSetStrategyType && obj.Spec.UpdateStrategy.RollingUpdate != nil {
+		if obj.Spec.Replicas != nil && obj.Spec.UpdateStrategy.RollingUpdate.Partition != nil {
+			if obj.Status.UpdatedReplicas < (*obj.Spec.Replicas - *obj.Spec.UpdateStrategy.RollingUpdate.Partition) {
+				return sdk.ResourceHealthStateUnhealthy, fmt.Sprintf("Waiting for partitioned roll out to finish because %d out of %d new pods have been updated",
+					obj.Status.UpdatedReplicas, (*obj.Spec.Replicas - *obj.Spec.UpdateStrategy.RollingUpdate.Partition))
+			}
+		}
+		return sdk.ResourceHealthStateHealthy, ""
+	}
+
+	if obj.Status.UpdateRevision != obj.Status.CurrentRevision {
+		return sdk.ResourceHealthStateUnhealthy, fmt.Sprintf("Waiting for statefulset rolling update to complete %d pods at revision %s", obj.Status.UpdatedReplicas, obj.Status.UpdateRevision)
+	}
+
+	return sdk.ResourceHealthStateHealthy, ""
+}