educode-controller/internal/k8s/workspace.go

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 }