package k8s import ( "context" "fmt" "time" appsv1 "k8s.io/api/apps/v1" corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/util/intstr" "k8s.io/client-go/kubernetes" ) const ( DefaultImage = "ghcr.io/dreamstarsky/educode:cpu-latest" ) var DefaultResourceLimits = ResourceLimits{ CPU: "500m", Memory: "1Gi", Storage: "5Gi", } type WorkspaceRequest struct { Image string Env map[string]string ResourceLimits *ResourceLimits WorkspaceID string Clientset *kubernetes.Clientset Namespace string } type ResourceLimits struct { CPU string `json:"cpu"` Memory string `json:"memory"` Storage string `json:"storage"` } func CreateWorkspace(req *WorkspaceRequest) error { // Define StatefulSet sts := &appsv1.StatefulSet{ ObjectMeta: metav1.ObjectMeta{ Name: req.WorkspaceID, Namespace: req.Namespace, Annotations: map[string]string{ "educode/expires-at": time.Now().Add(1 * time.Hour).Format(time.RFC3339), }, }, Spec: appsv1.StatefulSetSpec{ Selector: &metav1.LabelSelector{ MatchLabels: map[string]string{"app": req.WorkspaceID}, }, ServiceName: req.WorkspaceID, Replicas: int32Ptr(1), Template: corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{ Labels: map[string]string{"app": req.WorkspaceID}, }, Spec: corev1.PodSpec{ Containers: []corev1.Container{ { Name: "code-server", Image: req.Image, Ports: []corev1.ContainerPort{ {ContainerPort: 8080}, // code-server port {ContainerPort: 22}, // SSH port }, Env: getEnvVars(req.Env), Resources: getResourceRequirements(req.ResourceLimits), VolumeMounts: []corev1.VolumeMount{ {Name: "workspace-storage", MountPath: "/home/ubuntu/workspace"}, }, }, }, }, }, VolumeClaimTemplates: []corev1.PersistentVolumeClaim{ { ObjectMeta: metav1.ObjectMeta{Name: "workspace-storage"}, Spec: corev1.PersistentVolumeClaimSpec{ AccessModes: []corev1.PersistentVolumeAccessMode{corev1.ReadWriteMany}, Resources: corev1.VolumeResourceRequirements{ Requests: corev1.ResourceList{ corev1.ResourceStorage: resource.MustParse(req.ResourceLimits.Storage), }, }, }, }, }, }, } // Create resources sts, err := req.Clientset.AppsV1().StatefulSets(req.Namespace).Create(context.TODO(), sts, metav1.CreateOptions{}) if err != nil { return fmt.Errorf("failed to create statefulset: %w", err) } // Define Service svc := &corev1.Service{ ObjectMeta: metav1.ObjectMeta{ Name: req.WorkspaceID, Namespace: req.Namespace, OwnerReferences: []metav1.OwnerReference{ { APIVersion: "apps/v1", Kind: "StatefulSet", Name: req.WorkspaceID, UID: sts.UID, Controller: boolPrt(true), }, }, }, Spec: corev1.ServiceSpec{ Selector: map[string]string{"app": req.WorkspaceID}, Ports: []corev1.ServicePort{ {Name: "code-server", Port: 8080, TargetPort: intstr.FromInt(8080)}, {Name: "ssh", Port: 22, TargetPort: intstr.FromInt(22)}, }, Type: corev1.ServiceTypeClusterIP, }, } _, err = req.Clientset.CoreV1().Services(req.Namespace).Create(context.TODO(), svc, metav1.CreateOptions{}) if err != nil { // Cleanup StatefulSet if Service creation fails deletePolicy := metav1.DeletePropagationBackground req.Clientset.AppsV1().StatefulSets(req.Namespace).Delete(context.TODO(), req.WorkspaceID, metav1.DeleteOptions{ PropagationPolicy: &deletePolicy, }) return fmt.Errorf("failed to create service: %w", err) } return nil } func DeleteWorkspace(clientset *kubernetes.Clientset, namespace, workspaceID string) error { // Delete StatefulSet deletePolicy := metav1.DeletePropagationBackground // TODO // Maybe not delete pvc? err := clientset.AppsV1().StatefulSets(namespace).Delete(context.TODO(), workspaceID, metav1.DeleteOptions{ PropagationPolicy: &deletePolicy, }) if err != nil { return fmt.Errorf("failed to delete statefulset: %w", err) } return nil } func getEnvVars(env map[string]string) []corev1.EnvVar { var envVars []corev1.EnvVar for k, v := range env { envVars = append(envVars, corev1.EnvVar{Name: k, Value: v}) } return envVars } func getResourceRequirements(limits *ResourceLimits) corev1.ResourceRequirements { req := corev1.ResourceRequirements{ Requests: corev1.ResourceList{ corev1.ResourceCPU: resource.MustParse(limits.CPU), corev1.ResourceMemory: resource.MustParse(limits.Memory), }, } return req } func int32Ptr(i int32) *int32 { return &i } func boolPrt(b bool) *bool { return &b }