Support set image and rollout commands for ads (#50)

Signed-off-by: Siyu Wang <FillZpp.pub@gmail.com>

pkg/cmd/rollout/rollout_restart.go

@@ -158,7 +158,7 @@ func (o RestartOptions) RunRestart() error {
 	}
 
 	switch infos[0].Object.(type) {
-	case *kruiseappsv1alpha1.CloneSet:
+	case *kruiseappsv1alpha1.CloneSet, *kruiseappsv1beta1.StatefulSet, *kruiseappsv1alpha1.DaemonSet:
 
 		obj, err := resource.
 			NewHelper(infos[0].Client, infos[0].Mapping).
@@ -166,37 +166,11 @@ func (o RestartOptions) RunRestart() error {
 		if err != nil {
 			return err
 		}
-		res := obj.(*kruiseappsv1alpha1.CloneSet)
+		internalpolymorphichelpers.UpdateResourceEnv(obj)
-		internalpolymorphichelpers.UpdateResourceEnv(res)
 
 		_, err = resource.
 			NewHelper(infos[0].Client, infos[0].Mapping).
-			Replace(infos[0].Namespace, infos[0].Name, true, res)
+			Replace(infos[0].Namespace, infos[0].Name, true, obj)
-		if err != nil {
-			return err
-		}
-		printer, err := o.ToPrinter("restarted")
-		if err != nil {
-			allErrs = append(allErrs, err)
-		}
-		if err = printer.PrintObj(infos[0].Object, o.Out); err != nil {
-			allErrs = append(allErrs, err)
-		}
-		return utilerrors.NewAggregate(allErrs)
-
-	case *kruiseappsv1beta1.StatefulSet:
-		obj, err := resource.
-			NewHelper(infos[0].Client, infos[0].Mapping).
-			Get(infos[0].Namespace, infos[0].Name)
-		if err != nil {
-			return err
-		}
-		res := obj.(*kruiseappsv1beta1.StatefulSet)
-		internalpolymorphichelpers.UpdateResourceEnv(res)
-
-		_, err = resource.
-			NewHelper(infos[0].Client, infos[0].Mapping).
-			Replace(infos[0].Namespace, infos[0].Name, true, res)
 		if err != nil {
 			return err
 		}

pkg/internal/apps/kind_visitor.go

@@ -36,6 +36,7 @@ type KindVisitor interface {
 	VisitCronJob(kind GroupKindElement)
 	VisitCloneSet(kind GroupKindElement)
 	VisitAdvancedStatefulSet(kind GroupKindElement)
+	VisitAdvancedDaemonSet(kind GroupKindElement)
 }
 
 // GroupKindElement defines a Kubernetes API group elem
@@ -64,6 +65,8 @@ func (elem GroupKindElement) Accept(visitor KindVisitor) error {
 		visitor.VisitCloneSet(elem)
 	case elem.GroupMatch("apps.kruise.io") && elem.Kind == "StatefulSet":
 		visitor.VisitAdvancedStatefulSet(elem)
+	case elem.GroupMatch("apps.kruise.io") && elem.Kind == "DaemonSet":
+		visitor.VisitAdvancedDaemonSet(elem)
 	default:
 		return fmt.Errorf("no visitor method exists for %v", elem)
 	}

pkg/internal/polymorphichelpers/helpers.go

@@ -144,6 +144,18 @@ func SelectorsForObject(object runtime.Object) (namespace string, selector label
 		if err != nil {
 			return "", nil, fmt.Errorf("invalid label selector:%v", err)
 		}
+	case *kruiseappsv1beta1.StatefulSet:
+		namespace = t.Namespace
+		selector, err = metav1.LabelSelectorAsSelector(t.Spec.Selector)
+		if err != nil {
+			return "", nil, fmt.Errorf("invalid label selector:%v", err)
+		}
+	case *kruiseappsv1alpha1.DaemonSet:
+		namespace = t.Namespace
+		selector, err = metav1.LabelSelectorAsSelector(t.Spec.Selector)
+		if err != nil {
+			return "", nil, fmt.Errorf("invalid label selector:%v", err)
+		}
 	case *appsv1.DaemonSet:
 		namespace = t.Namespace
 		selector, err = metav1.LabelSelectorAsSelector(t.Spec.Selector)
@@ -257,6 +269,11 @@ func UpdateResourceEnv(object runtime.Object) {
 			tmp := &obj.Spec.Template.Spec.Containers[i]
 			tmp.Env = updateEnv(tmp.Env, addingEnvs, []string{})
 		}
+	case *kruiseappsv1alpha1.DaemonSet:
+		for i := range obj.Spec.Template.Spec.Containers {
+			tmp := &obj.Spec.Template.Spec.Containers[i]
+			tmp.Env = updateEnv(tmp.Env, addingEnvs, []string{})
+		}
 	}
 
 }

pkg/internal/polymorphichelpers/history.go

@@ -116,6 +116,11 @@ type AdvancedStatefulSetHistoryViewer struct {
 	kc kruiseclientsets.Interface
 }
 
+type AdvancedDaemonSetHistoryViewer struct {
+	k  kubernetes.Interface
+	kc kruiseclientsets.Interface
+}
+
 func (v *HistoryVisitor) VisitCloneSet(kind internalapps.GroupKindElement) {
 	v.result = &CloneSetHistoryViewer{v.clientset, v.kruiseclientset}
 }
@@ -124,6 +129,10 @@ func (v *HistoryVisitor) VisitAdvancedStatefulSet(kind internalapps.GroupKindEle
 	v.result = &AdvancedStatefulSetHistoryViewer{v.clientset, v.kruiseclientset}
 }
 
+func (v *HistoryVisitor) VisitAdvancedDaemonSet(kind internalapps.GroupKindElement) {
+	v.result = &AdvancedDaemonSetHistoryViewer{v.clientset, v.kruiseclientset}
+}
+
 // TODO impl ViewHistory func for CloneSet
 func (h *CloneSetHistoryViewer) ViewHistory(namespace, name string, revision int64) (string, error) {
 
@@ -133,11 +142,11 @@ func (h *CloneSetHistoryViewer) ViewHistory(namespace, name string, revision int
 	}
 
 	return printHistory(history, revision, func(history *appsv1.ControllerRevision) (*corev1.PodTemplateSpec, error) {
-		stsOfHistory, err := applyCloneSetHistory(cs, history)
+		cloneSetOfHistory, err := applyCloneSetHistory(cs, history)
 		if err != nil {
 			return nil, err
 		}
-		return &stsOfHistory.Spec.Template, err
+		return &cloneSetOfHistory.Spec.Template, err
 	})
 }
 
@@ -155,6 +164,20 @@ func (h *AdvancedStatefulSetHistoryViewer) ViewHistory(namespace, name string, r
 	})
 }
 
+func (h *AdvancedDaemonSetHistoryViewer) ViewHistory(namespace, name string, revision int64) (string, error) {
+	ads, history, err := advancedDaemonSetHistory(h.k.AppsV1(), h.kc.AppsV1alpha1(), namespace, name)
+	if err != nil {
+		return "", err
+	}
+	return printHistory(history, revision, func(history *appsv1.ControllerRevision) (*corev1.PodTemplateSpec, error) {
+		adsOfHistory, err := applyAdvancedDaemonSetHistory(ads, history)
+		if err != nil {
+			return nil, err
+		}
+		return &adsOfHistory.Spec.Template, err
+	})
+}
+
 // ViewHistory returns a revision-to-replicaset map as the revision history of a deployment
 // TODO: this should be a describer
 func (h *DeploymentHistoryViewer) ViewHistory(namespace, name string, revision int64) (string, error) {
@@ -382,6 +405,29 @@ func daemonSetHistory(
 	return ds, history, nil
 }
 
+// advancedDaemonSetHistory returns the Advanced DaemonSet named name in namespace and all ControllerRevisions in its history.
+func advancedDaemonSetHistory(
+	apps clientappsv1.AppsV1Interface, appsv1alpha1 kruiseclientappsv1alpha1.AppsV1alpha1Interface,
+	namespace, name string) (*kruiseappsv1alpha1.DaemonSet, []*appsv1.ControllerRevision, error) {
+	ds, err := appsv1alpha1.DaemonSets(namespace).Get(context.TODO(), name, metav1.GetOptions{})
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to retrieve DaemonSet %s: %v", name, err)
+	}
+	selector, err := metav1.LabelSelectorAsSelector(ds.Spec.Selector)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create selector for DaemonSet %s: %v", ds.Name, err)
+	}
+	accessor, err := meta.Accessor(ds)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create accessor for DaemonSet %s: %v", ds.Name, err)
+	}
+	history, err := controlledHistoryV1(apps, ds.Namespace, selector, accessor)
+	if err != nil {
+		return nil, nil, fmt.Errorf("unable to find history controlled by DaemonSet %s: %v", ds.Name, err)
+	}
+	return ds, history, nil
+}
+
 func clonesetHistory(
 	apps clientappsv1.AppsV1Interface, appsv1alpha1 kruiseclientappsv1alpha1.AppsV1alpha1Interface,
 	namespace, name string) (*kruiseappsv1alpha1.CloneSet, []*appsv1.ControllerRevision, error) {
@@ -507,6 +553,7 @@ func applyCloneSetHistory(cs *kruiseappsv1alpha1.CloneSet,
 	}
 	return result, nil
 }
+
 func applyAdvancedStatefulSetHistory(asts *kruiseappsv1beta1.StatefulSet,
 	history *appsv1.ControllerRevision) (*kruiseappsv1beta1.StatefulSet, error) {
 	astsBytes, err := json.Marshal(asts)
@@ -525,6 +572,24 @@ func applyAdvancedStatefulSetHistory(asts *kruiseappsv1beta1.StatefulSet,
 	return result, nil
 }
 
+func applyAdvancedDaemonSetHistory(ads *kruiseappsv1alpha1.DaemonSet,
+	history *appsv1.ControllerRevision) (*kruiseappsv1alpha1.DaemonSet, error) {
+	adsBytes, err := json.Marshal(ads)
+	if err != nil {
+		return nil, err
+	}
+	patched, err := strategicpatch.StrategicMergePatch(adsBytes, history.Data.Raw, ads)
+	if err != nil {
+		return nil, err
+	}
+	result := &kruiseappsv1alpha1.DaemonSet{}
+	err = json.Unmarshal(patched, result)
+	if err != nil {
+		return nil, err
+	}
+	return result, nil
+}
+
 // TODO: copied here until this becomes a describer
 func tabbedString(f func(io.Writer) error) (string, error) {
 	out := new(tabwriter.Writer)

pkg/internal/polymorphichelpers/objectrestarter.go

@@ -114,6 +114,7 @@ func defaultObjectRestarter(obj runtime.Object) ([]byte, error) {
 		}
 		obj.Spec.Template.ObjectMeta.Annotations["kubectl.kubernetes.io/restartedAt"] = time.Now().Format(time.RFC3339)
 		return runtime.Encode(scheme.Codecs.LegacyCodec(appsv1beta2.SchemeGroupVersion), obj)
+
 	case *kruiseappsv1alpha1.CloneSet:
 		if obj.Spec.Template.ObjectMeta.Annotations == nil {
 			obj.Spec.Template.ObjectMeta.Annotations = make(map[string]string)

pkg/internal/polymorphichelpers/rollback.go

@@ -84,6 +84,9 @@ func (v *RollbackVisitor) VisitCronJob(kind internalapps.GroupKindElement)
 func (v *RollbackVisitor) VisitAdvancedStatefulSet(kind internalapps.GroupKindElement) {
 	v.result = &AdvancedStatefulSetRollbacker{k: v.clientset, kc: v.kruiseclientset}
 }
+func (v *RollbackVisitor) VisitAdvancedDaemonSet(kind internalapps.GroupKindElement) {
+	v.result = &AdvancedDaemonSetRollbacker{k: v.clientset, kc: v.kruiseclientset}
+}
 
 // RollbackerFor returns an implementation of Rollbacker interface for the given schema kind
 func RollbackerFor(kind schema.GroupKind, c kubernetes.Interface, kc kruiseclientsets.Interface) (Rollbacker, error) {
@@ -532,6 +535,64 @@ func (r *AdvancedStatefulSetRollbacker) Rollback(obj runtime.Object,
 	return rollbackSuccess, nil
 }
 
+type AdvancedDaemonSetRollbacker struct {
+	k  kubernetes.Interface
+	kc kruiseclientsets.Interface
+}
+
+func (r *AdvancedDaemonSetRollbacker) Rollback(obj runtime.Object,
+	updatedAnnotations map[string]string,
+	toRevision int64,
+	dryRunStrategy cmdutil.DryRunStrategy) (string, error) {
+	if toRevision < 0 {
+		return "", revisionNotFoundErr(toRevision)
+	}
+	accessor, err := meta.Accessor(obj)
+	if err != nil {
+		return "", fmt.Errorf("failed to create accessor for kind %v: %s", obj.GetObjectKind(), err.Error())
+	}
+	ads, history, err := advancedDaemonSetHistory(r.k.AppsV1(), r.kc.AppsV1alpha1(), accessor.GetNamespace(), accessor.GetName())
+	if err != nil {
+		return "", err
+	}
+	if toRevision == 0 && len(history) <= 1 {
+		return "", fmt.Errorf("no last revision to roll back to")
+	}
+
+	toHistory := findHistory(toRevision, history)
+	if toHistory == nil {
+		return "", revisionNotFoundErr(toRevision)
+	}
+
+	if dryRunStrategy == cmdutil.DryRunClient {
+		appliedDS, err := applyAdvancedDaemonSetHistory(ads, toHistory)
+		if err != nil {
+			return "", err
+		}
+		return printPodTemplate(&appliedDS.Spec.Template)
+	}
+
+	// Skip if the revision already matches current DaemonSet
+	done, err := advancedDaemonSetMatch(ads, toHistory)
+	if err != nil {
+		return "", err
+	}
+	if done {
+		return fmt.Sprintf("%s (current template already matches revision %d)", rollbackSkipped, toRevision), nil
+	}
+
+	patchOptions := metav1.PatchOptions{}
+	if dryRunStrategy == cmdutil.DryRunServer {
+		patchOptions.DryRun = []string{metav1.DryRunAll}
+	}
+	// Restore revision
+	if _, err = r.kc.AppsV1alpha1().CloneSets(ads.Namespace).Patch(context.TODO(), ads.Name, types.MergePatchType, toHistory.Data.Raw, patchOptions); err != nil {
+		return "", fmt.Errorf("failed restoring revision %d: %v", toRevision, err)
+	}
+
+	return rollbackSuccess, nil
+}
+
 var appsCodec = scheme.Codecs.LegacyCodec(appsv1.SchemeGroupVersion)
 
 // applyRevision returns a new StatefulSet constructed by restoring the state in revision to set. If the returned error
@@ -549,13 +610,13 @@ func applyRevision(set *appsv1.StatefulSet, revision *appsv1.ControllerRevision)
 	return result, nil
 }
 
-var kruiseAppsCodec = scheme.Codecs.LegacyCodec(kruiseappsv1alpha1.SchemeGroupVersion)
+var kruiseAppsv1alpha1Codec = scheme.Codecs.LegacyCodec(kruiseappsv1alpha1.SchemeGroupVersion)
 
 // applyCloneSetRevision returns a new CloneSet constructed by restoring the state in revision to set. If the returned error
 // is nil, the returned CloneSet is valid.
 func applyCloneSetRevision(cs *kruiseappsv1alpha1.CloneSet,
 	revision *appsv1.ControllerRevision) (*kruiseappsv1alpha1.CloneSet, error) {
-	patched, err := strategicpatch.StrategicMergePatch([]byte(runtime.EncodeOrDie(kruiseAppsCodec, cs)), revision.Data.Raw, cs)
+	patched, err := strategicpatch.StrategicMergePatch([]byte(runtime.EncodeOrDie(kruiseAppsv1alpha1Codec, cs)), revision.Data.Raw, cs)
 	if err != nil {
 		return nil, err
 	}
@@ -567,13 +628,13 @@ func applyCloneSetRevision(cs *kruiseappsv1alpha1.CloneSet,
 	return result, nil
 }
 
-var kruisev1beta1AppsCodec = scheme.Codecs.LegacyCodec(kruiseappsv1beta1.SchemeGroupVersion)
+var kruiseAppsv1beta1Codec = scheme.Codecs.LegacyCodec(kruiseappsv1beta1.SchemeGroupVersion)
 
 // apply  applyAdvancedStatefulSetRevision returns a new Advanced StatefulSet constructed by restoring the state in revision to set.
 // If the returned error is nil, the returned Advanced StatefulSet is valid.
 func applyAdvancedStatefulSetRevision(asts *kruiseappsv1beta1.StatefulSet,
 	revision *appsv1.ControllerRevision) (*kruiseappsv1beta1.StatefulSet, error) {
-	patched, err := strategicpatch.StrategicMergePatch([]byte(runtime.EncodeOrDie(kruisev1beta1AppsCodec, asts)), revision.Data.Raw, asts)
+	patched, err := strategicpatch.StrategicMergePatch([]byte(runtime.EncodeOrDie(kruiseAppsv1beta1Codec, asts)), revision.Data.Raw, asts)
 	if err != nil {
 		return nil, err
 	}
@@ -612,6 +673,15 @@ func cloneSetMatch(cs *kruiseappsv1alpha1.CloneSet, history *appsv1.ControllerRe
 	return bytes.Equal(patch, history.Data.Raw), nil
 }
 
+// advancedDaemonSetMatch check if the given Advanced DaemonSet's template matches the template stored in the given history.
+func advancedDaemonSetMatch(ads *kruiseappsv1alpha1.DaemonSet, history *appsv1.ControllerRevision) (bool, error) {
+	patch, err := getAdvancedDaemonSetPatch(ads)
+	if err != nil {
+		return false, err
+	}
+	return bytes.Equal(patch, history.Data.Raw), nil
+}
+
 // getStatefulSetPatch returns a strategic merge patch that can be applied to restore a StatefulSet to a
 // previous version. If the returned error is nil the patch is valid. The current state that we save is just the
 // PodSpecTemplate. We can modify this later to encompass more state (or less) and remain compatible with previously
@@ -639,7 +709,7 @@ func getStatefulSetPatch(set *appsv1.StatefulSet) ([]byte, error) {
 // getCloneSetPatch returns a strategic merge patch that can be applied to restore a CloneSet to a
 // previous version.
 func getCloneSetPatch(cs *kruiseappsv1alpha1.CloneSet) ([]byte, error) {
-	str, err := runtime.Encode(kruiseAppsCodec, cs)
+	str, err := runtime.Encode(kruiseAppsv1alpha1Codec, cs)
 	if err != nil {
 		return nil, err
 	}
@@ -662,7 +732,7 @@ func getCloneSetPatch(cs *kruiseappsv1alpha1.CloneSet) ([]byte, error) {
 // getAdvancedStatefulSetPatch returns a strategic merge patch that can be applied to restore a Advanced StatefulSet to
 // a previous version
 func getAdvancedStatefulSetPatch(asts *kruiseappsv1beta1.StatefulSet) ([]byte, error) {
-	str, err := runtime.Encode(kruisev1beta1AppsCodec, asts)
+	str, err := runtime.Encode(kruiseAppsv1beta1Codec, asts)
 	if err != nil {
 		return nil, err
 	}
@@ -681,6 +751,27 @@ func getAdvancedStatefulSetPatch(asts *kruiseappsv1beta1.StatefulSet) ([]byte, e
 	return patch, err
 }
 
+func getAdvancedDaemonSetPatch(ads *kruiseappsv1alpha1.DaemonSet) ([]byte, error) {
+	str, err := runtime.Encode(kruiseAppsv1alpha1Codec, ads)
+	if err != nil {
+		return nil, err
+	}
+	var raw map[string]interface{}
+	if err := json.Unmarshal(str, &raw); err != nil {
+		return nil, err
+	}
+	objCopy := make(map[string]interface{})
+	specCopy := make(map[string]interface{})
+	// Create a patch of the DaemonSet that replaces spec.template
+	spec := raw["spec"].(map[string]interface{})
+	template := spec["template"].(map[string]interface{})
+	specCopy["template"] = template
+	template["$patch"] = "replace"
+	objCopy["spec"] = specCopy
+	patch, err := json.Marshal(objCopy)
+	return patch, err
+}
+
 // findHistory returns a controllerrevision of a specific revision from the given controllerrevisions.
 // It returns nil if no such controllerrevision exists.
 // If toRevision is 0, the last previously used history is returned.

pkg/internal/polymorphichelpers/updatepodspec.go

@@ -39,6 +39,8 @@ func updatePodSpecForObject(obj runtime.Object, fn func(*v1.PodSpec) error) (boo
 		return true, fn(&t.Spec.Template.Spec)
 	case *kruiseappsv1beta1.StatefulSet:
 		return true, fn(&t.Spec.Template.Spec)
+	case *kruiseappsv1alpha1.DaemonSet:
+		return true, fn(&t.Spec.Template.Spec)
 	case *v1.Pod:
 		return true, fn(&t.Spec)
 		// ReplicationController