func GenerateSts()

in deploy/kubernetes/operator/pkg/controller/sync/shuffleserver/shuffleserver.go [123:240]
101 lines of code
8 McCabe index (conditional complexity)

func GenerateSts(kubeClient kubernetes.Interface, rss *unifflev1alpha1.RemoteShuffleService) *appsv1.StatefulSet {
	name := GenerateName(rss)
	replicas := getReplicas(kubeClient, rss)

	podSpec := corev1.PodSpec{
		SecurityContext:    rss.Spec.ShuffleServer.SecurityContext,
		HostNetwork:        *rss.Spec.ShuffleServer.HostNetwork,
		ServiceAccountName: GenerateName(rss),
		Tolerations:        rss.Spec.ShuffleServer.Tolerations,
		Volumes:            rss.Spec.ShuffleServer.Volumes,
		NodeSelector:       rss.Spec.ShuffleServer.NodeSelector,
		Affinity:           rss.Spec.ShuffleServer.Affinity,
		ImagePullSecrets:   rss.Spec.ImagePullSecrets,
	}

	configurationVolume := corev1.Volume{
		Name: controllerconstants.ConfigurationVolumeName,
		VolumeSource: corev1.VolumeSource{
			ConfigMap: &corev1.ConfigMapVolumeSource{
				LocalObjectReference: corev1.LocalObjectReference{
					Name: rss.Spec.ConfigMapName,
				},
				DefaultMode: pointer.Int32(0777),
			},
		},
	}
	podSpec.Volumes = append(podSpec.Volumes, configurationVolume)
	if podSpec.HostNetwork {
		podSpec.DNSPolicy = corev1.DNSClusterFirstWithHostNet
	}

	defaultLabels := utils.GenerateShuffleServerLabels(rss)
	sts := &appsv1.StatefulSet{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: rss.Namespace,
			Labels:    defaultLabels,
		},
		Spec: appsv1.StatefulSetSpec{
			Selector: &metav1.LabelSelector{
				MatchLabels: defaultLabels,
			},
			PodManagementPolicy: rss.Spec.ShuffleServer.PodManagementPolicy,
			UpdateStrategy: appsv1.StatefulSetUpdateStrategy{
				Type: appsv1.RollingUpdateStatefulSetStrategyType,
				RollingUpdate: &appsv1.RollingUpdateStatefulSetStrategy{
					Partition: replicas,
				},
			},
			ServiceName: generateHeadlessSVCName(rss),
			Replicas:    replicas,
			Template: corev1.PodTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{
					Labels: make(map[string]string),
				},
				Spec: podSpec,
			},
		},
	}
	for k, v := range rss.Spec.ShuffleServer.Labels {
		sts.Spec.Template.Labels[k] = v
	}
	for k, v := range defaultLabels {
		sts.Spec.Template.Labels[k] = v
	}

	// set runtimeClassName
	if rss.Spec.ShuffleServer.RuntimeClassName != nil {
		sts.Spec.Template.Spec.RuntimeClassName = rss.Spec.ShuffleServer.RuntimeClassName
	}

	// add VolumeClaimTemplates, support cloud storage
	sts.Spec.VolumeClaimTemplates = make([]corev1.PersistentVolumeClaim, 0, len(rss.Spec.ShuffleServer.VolumeClaimTemplates))
	for _, pvcTemplate := range rss.Spec.ShuffleServer.VolumeClaimTemplates {
		sts.Spec.VolumeClaimTemplates = append(sts.Spec.VolumeClaimTemplates, corev1.PersistentVolumeClaim{
			ObjectMeta: metav1.ObjectMeta{
				Name: *pvcTemplate.VolumeNameTemplate,
			},
			Spec: pvcTemplate.Spec,
		})
	}

	// add custom annotation, and default annotation used by rss
	reservedAnnotations := map[string]string{
		constants.AnnotationRssName: rss.Name,
		constants.AnnotationRssUID:  string(rss.UID),
		constants.AnnotationMetricsServerPort: fmt.Sprintf("%v",
			*rss.Spec.ShuffleServer.HTTPPort),
		constants.AnnotationShuffleServerPort: fmt.Sprintf("%v",
			*rss.Spec.ShuffleServer.RPCPort),
	}

	annotations := map[string]string{}
	for key, value := range rss.Spec.ShuffleServer.Annotations {
		annotations[key] = value
	}
	// reserved annotations have higher preference.
	for key, value := range reservedAnnotations {
		annotations[key] = value
	}

	sts.Spec.Template.Annotations = annotations

	// add init containers, the main container and other containers.
	sts.Spec.Template.Spec.InitContainers = util.GenerateInitContainers(rss.Spec.ShuffleServer.RSSPodSpec)
	containers := []corev1.Container{*generateMainContainer(rss)}
	containers = append(containers, rss.Spec.ShuffleServer.SidecarContainers...)
	sts.Spec.Template.Spec.Containers = containers

	// add hostPath volumes for shuffle servers.
	hostPathMounts := rss.Spec.ShuffleServer.HostPathMounts
	logHostPath := rss.Spec.ShuffleServer.LogHostPath
	sts.Spec.Template.Spec.Volumes = append(sts.Spec.Template.Spec.Volumes,
		util.GenerateHostPathVolumes(hostPathMounts, logHostPath, name)...)

	util.AddOwnerReference(&sts.ObjectMeta, rss)
	return sts
}