ayatori/internal/controller/job_controller.go

package controller

import (
	"context"
	"fmt"
	"slices"
	"time"

	executionv1alpha1 "git.ddupan.top/panxiao81/ayatori/api/execution/v1alpha1"
	kubernetesadapter "git.ddupan.top/panxiao81/ayatori/internal/adapter/kubernetes"
	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/types"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

const (
	jobFinalizer        = "execution.ayatori.ddupan.top/job-cleanup"
	reasonResultUnknown = "ResultUnknown"
)

// JobReconciler executes Ayatori Jobs using supported adapters.
type JobReconciler struct {
	client.Client
	Now func() time.Time
}

// +kubebuilder:rbac:groups=execution.ayatori.ddupan.top,resources=jobs,verbs=get;list;watch;update;patch
// +kubebuilder:rbac:groups=execution.ayatori.ddupan.top,resources=jobs/status,verbs=get;update;patch
// +kubebuilder:rbac:groups=execution.ayatori.ddupan.top,resources=jobs/finalizers,verbs=update
// +kubebuilder:rbac:groups=execution.ayatori.ddupan.top,resources=jobclasses;kubernetesexecutionparameters,verbs=get;list;watch
// +kubebuilder:rbac:groups=batch,resources=jobs,verbs=get;list;watch;create;delete
// +kubebuilder:rbac:groups="",resources=namespaces;serviceaccounts,verbs=get;list;watch

func (r *JobReconciler) Reconcile(ctx context.Context, request ctrl.Request) (ctrl.Result, error) {
	logger := log.FromContext(ctx)
	job := &executionv1alpha1.Job{}
	if err := r.Get(ctx, request.NamespacedName, job); err != nil {
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}

	if !job.DeletionTimestamp.IsZero() {
		return ctrl.Result{}, r.finalize(ctx, job)
	}
	if isTerminal(job) {
		return ctrl.Result{}, nil
	}
	if job.Spec.DesiredState == executionv1alpha1.JobDesiredStateCancelled {
		return ctrl.Result{}, r.cancel(ctx, job)
	}

	if !containsString(job.Finalizers, jobFinalizer) {
		job.Finalizers = append(job.Finalizers, jobFinalizer)
		if err := r.Update(ctx, job); err != nil {
			return ctrl.Result{}, err
		}
		return ctrl.Result{}, nil
	}
	if job.Status.Execution != nil {
		return ctrl.Result{}, r.observeExisting(ctx, job)
	}

	class, parameters, resources, waiting, err := r.resolve(ctx, job)
	if err != nil {
		return ctrl.Result{}, err
	}
	if waiting {
		return ctrl.Result{RequeueAfter: 30 * time.Second}, nil
	}

	backend := &batchv1.Job{}
	key := types.NamespacedName{Namespace: job.Namespace, Name: job.Name}
	err = r.Get(ctx, key, backend)
	if apierrors.IsNotFound(err) {
		backend = kubernetesadapter.BuildJob(job, parameters, resources)
		if err := r.Create(ctx, backend); err != nil {
			return ctrl.Result{}, err
		}
		logger.Info("Created Kubernetes backend Job", "backend", key)
		return ctrl.Result{}, r.markScheduled(ctx, job, class, parameters, resources, backend)
	}
	if err != nil {
		return ctrl.Result{}, err
	}
	if err := kubernetesadapter.ValidateOwnership(job, backend); err != nil {
		return ctrl.Result{}, r.setCondition(ctx, job, metav1.Condition{
			Type: executionv1alpha1.JobConditionScheduled, Status: metav1.ConditionFalse,
			Reason: "BackendConflict", Message: err.Error(),
		})
	}
	if !conditionTrue(job.Status.Conditions, executionv1alpha1.JobConditionScheduled) {
		return ctrl.Result{}, r.markScheduled(ctx, job, class, parameters, resources, backend)
	}
	return ctrl.Result{}, r.observe(ctx, job, backend)
}

func (r *JobReconciler) observeExisting(ctx context.Context, job *executionv1alpha1.Job) error {
	if job.Status.Execution.Adapter != "kubernetes" {
		return r.setCondition(ctx, job, metav1.Condition{
			Type: executionv1alpha1.JobConditionSucceeded, Status: metav1.ConditionUnknown,
			Reason: reasonResultUnknown, Message: fmt.Sprintf("adapter %q is not available", job.Status.Execution.Adapter),
		})
	}
	backend := &batchv1.Job{}
	key := types.NamespacedName{Namespace: job.Namespace, Name: job.Name}
	if err := r.Get(ctx, key, backend); err != nil {
		if apierrors.IsNotFound(err) {
			return r.setCondition(ctx, job, metav1.Condition{
				Type: executionv1alpha1.JobConditionSucceeded, Status: metav1.ConditionUnknown,
				Reason: reasonResultUnknown, Message: "Kubernetes backend Job is missing",
			})
		}
		return err
	}
	if err := kubernetesadapter.ValidateOwnership(job, backend); err != nil {
		return r.setCondition(ctx, job, metav1.Condition{
			Type: executionv1alpha1.JobConditionSucceeded, Status: metav1.ConditionUnknown,
			Reason: reasonResultUnknown, Message: err.Error(),
		})
	}
	return r.observe(ctx, job, backend)
}

func (r *JobReconciler) resolve(
	ctx context.Context,
	job *executionv1alpha1.Job,
) (*executionv1alpha1.JobClass, *executionv1alpha1.KubernetesExecutionParameters, executionv1alpha1.ExecutionResourceRequirements, bool, error) {
	if job.Spec.JobClassName == "" {
		return nil, nil, executionv1alpha1.ExecutionResourceRequirements{}, true, r.reject(ctx, job, "NoDefaultJobClass", "spec.jobClassName is required in the first implementation slice")
	}

	class := &executionv1alpha1.JobClass{}
	if err := r.Get(ctx, types.NamespacedName{Name: job.Spec.JobClassName}, class); err != nil {
		if apierrors.IsNotFound(err) {
			return nil, nil, executionv1alpha1.ExecutionResourceRequirements{}, true, r.reject(ctx, job, "JobClassNotFound", fmt.Sprintf("JobClass %q does not exist", job.Spec.JobClassName))
		}
		return nil, nil, executionv1alpha1.ExecutionResourceRequirements{}, false, err
	}
	if class.Spec.ControllerName != kubernetesadapter.ControllerName {
		return nil, nil, executionv1alpha1.ExecutionResourceRequirements{}, true, r.reject(ctx, job, "UnsupportedController", fmt.Sprintf("controller %q is not supported", class.Spec.ControllerName))
	}
	ref := class.Spec.ParametersRef
	if ref.Group != executionv1alpha1.GroupVersion.Group || ref.Kind != "KubernetesExecutionParameters" {
		return nil, nil, executionv1alpha1.ExecutionResourceRequirements{}, true, r.reject(ctx, job, "InvalidParametersReference", "JobClass must reference KubernetesExecutionParameters")
	}
	if allowed, err := r.namespaceAllowed(ctx, job.Namespace, class.Spec.AllowedNamespaces); err != nil {
		return nil, nil, executionv1alpha1.ExecutionResourceRequirements{}, false, err
	} else if !allowed {
		return nil, nil, executionv1alpha1.ExecutionResourceRequirements{}, true, r.reject(ctx, job, "NamespaceNotAllowed", fmt.Sprintf("namespace %q is not allowed by JobClass %q", job.Namespace, class.Name))
	}

	parameters := &executionv1alpha1.KubernetesExecutionParameters{}
	if err := r.Get(ctx, types.NamespacedName{Name: ref.Name}, parameters); err != nil {
		if apierrors.IsNotFound(err) {
			return nil, nil, executionv1alpha1.ExecutionResourceRequirements{}, true, r.reject(ctx, job, "ParametersNotFound", fmt.Sprintf("KubernetesExecutionParameters %q does not exist", ref.Name))
		}
		return nil, nil, executionv1alpha1.ExecutionResourceRequirements{}, false, err
	}
	serviceAccount := &corev1.ServiceAccount{}
	if err := r.Get(ctx, types.NamespacedName{Namespace: job.Namespace, Name: parameters.Spec.ServiceAccountName}, serviceAccount); err != nil {
		if apierrors.IsNotFound(err) {
			return nil, nil, executionv1alpha1.ExecutionResourceRequirements{}, true, r.reject(ctx, job, "ServiceAccountNotFound", fmt.Sprintf("ServiceAccount %q does not exist", parameters.Spec.ServiceAccountName))
		}
		return nil, nil, executionv1alpha1.ExecutionResourceRequirements{}, false, err
	}

	resources := applyResourceDefaults(job.Spec.Resources, class.Spec.Resources.Defaults)
	if err := validateResources(resources); err != nil {
		return nil, nil, executionv1alpha1.ExecutionResourceRequirements{}, true, r.reject(ctx, job, "InvalidResources", err.Error())
	}
	if err := r.accept(ctx, job, class, parameters, resources); err != nil {
		return nil, nil, executionv1alpha1.ExecutionResourceRequirements{}, false, err
	}
	return class, parameters, resources, false, nil
}

func (r *JobReconciler) namespaceAllowed(ctx context.Context, namespace string, selector *metav1.LabelSelector) (bool, error) {
	if selector == nil {
		return true, nil
	}
	ns := &corev1.Namespace{}
	if err := r.Get(ctx, types.NamespacedName{Name: namespace}, ns); err != nil {
		return false, err
	}
	compiled, err := metav1.LabelSelectorAsSelector(selector)
	if err != nil {
		return false, err
	}
	return compiled.Matches(labels.Set(ns.Labels)), nil
}

func (r *JobReconciler) accept(ctx context.Context, job *executionv1alpha1.Job, class *executionv1alpha1.JobClass, parameters *executionv1alpha1.KubernetesExecutionParameters, resources executionv1alpha1.ExecutionResourceRequirements) error {
	job.Status.ResolvedJobClass = &executionv1alpha1.ResolvedJobClassReference{
		Name: class.Name, UID: class.UID, ControllerName: class.Spec.ControllerName,
		ParametersRef: executionv1alpha1.ParametersReference{
			Group: class.Spec.ParametersRef.Group, Kind: class.Spec.ParametersRef.Kind,
			Name: parameters.Name, UID: parameters.UID,
		},
	}
	job.Status.EffectiveResources = resources
	return r.setCondition(ctx, job, metav1.Condition{
		Type: executionv1alpha1.JobConditionAccepted, Status: metav1.ConditionTrue,
		Reason: "Accepted", Message: fmt.Sprintf("JobClass %q accepted", class.Name),
	})
}

func (r *JobReconciler) reject(ctx context.Context, job *executionv1alpha1.Job, reason, message string) error {
	return r.setCondition(ctx, job, metav1.Condition{
		Type: executionv1alpha1.JobConditionAccepted, Status: metav1.ConditionFalse,
		Reason: reason, Message: message,
	})
}

func (r *JobReconciler) markScheduled(ctx context.Context, job *executionv1alpha1.Job, class *executionv1alpha1.JobClass, parameters *executionv1alpha1.KubernetesExecutionParameters, resources executionv1alpha1.ExecutionResourceRequirements, backend *batchv1.Job) error {
	job.Status.ResolvedJobClass = &executionv1alpha1.ResolvedJobClassReference{
		Name: class.Name, UID: class.UID, ControllerName: class.Spec.ControllerName,
		ParametersRef: executionv1alpha1.ParametersReference{Group: class.Spec.ParametersRef.Group, Kind: class.Spec.ParametersRef.Kind, Name: parameters.Name, UID: parameters.UID},
	}
	job.Status.EffectiveResources = resources
	job.Status.Execution = &executionv1alpha1.ExecutionStatus{
		Adapter:    "kubernetes",
		References: []executionv1alpha1.ExecutionReference{{Type: kubernetesadapter.ReferenceType, ID: string(backend.UID)}},
	}
	meta.SetStatusCondition(&job.Status.Conditions, condition(job, executionv1alpha1.JobConditionAccepted, metav1.ConditionTrue, "Accepted", "Job accepted"))
	meta.SetStatusCondition(&job.Status.Conditions, condition(job, executionv1alpha1.JobConditionScheduled, metav1.ConditionTrue, "BackendCreated", "Kubernetes Job created"))
	meta.SetStatusCondition(&job.Status.Conditions, condition(job, executionv1alpha1.JobConditionSucceeded, metav1.ConditionUnknown, "Pending", "Waiting for task to start"))
	job.Status.ObservedGeneration = job.Generation
	return r.Status().Update(ctx, job)
}

func (r *JobReconciler) observe(ctx context.Context, job *executionv1alpha1.Job, backend *batchv1.Job) error {
	if job.Status.StartTime == nil && backend.Status.StartTime != nil {
		job.Status.StartTime = backend.Status.StartTime.DeepCopy()
	}
	for _, backendCondition := range backend.Status.Conditions {
		switch {
		case backendCondition.Type == batchv1.JobComplete && backendCondition.Status == corev1.ConditionTrue:
			completion := backend.Status.CompletionTime
			if completion == nil {
				now := metav1.NewTime(r.now())
				completion = &now
			}
			job.Status.CompletionTime = completion.DeepCopy()
			job.Status.Result = &executionv1alpha1.JobResult{Reason: "Completed"}
			return r.setCondition(ctx, job, metav1.Condition{Type: executionv1alpha1.JobConditionSucceeded, Status: metav1.ConditionTrue, Reason: "Completed", Message: backendCondition.Message})
		case backendCondition.Type == batchv1.JobFailed && backendCondition.Status == corev1.ConditionTrue:
			completion := metav1.NewTime(r.now())
			job.Status.CompletionTime = &completion
			job.Status.Result = &executionv1alpha1.JobResult{Reason: "ProcessFailed"}
			return r.setCondition(ctx, job, metav1.Condition{Type: executionv1alpha1.JobConditionSucceeded, Status: metav1.ConditionFalse, Reason: "ProcessFailed", Message: backendCondition.Message})
		}
	}
	reason := "Pending"
	message := "Waiting for task to start"
	if backend.Status.StartTime != nil || backend.Status.Active > 0 {
		reason = "Running"
		message = "Task is running"
	}
	return r.setCondition(ctx, job, metav1.Condition{Type: executionv1alpha1.JobConditionSucceeded, Status: metav1.ConditionUnknown, Reason: reason, Message: message})
}

func (r *JobReconciler) cancel(ctx context.Context, job *executionv1alpha1.Job) error {
	backend := &batchv1.Job{}
	key := types.NamespacedName{Namespace: job.Namespace, Name: job.Name}
	err := r.Get(ctx, key, backend)
	if err == nil {
		if err := kubernetesadapter.ValidateOwnership(job, backend); err != nil {
			return err
		}
		for _, backendCondition := range backend.Status.Conditions {
			if (backendCondition.Type == batchv1.JobComplete || backendCondition.Type == batchv1.JobFailed) &&
				backendCondition.Status == corev1.ConditionTrue {
				return r.observe(ctx, job, backend)
			}
		}
		if err := r.Delete(ctx, backend, client.PropagationPolicy(metav1.DeletePropagationBackground)); err != nil && !apierrors.IsNotFound(err) {
			return err
		}
		return nil
	}
	if !apierrors.IsNotFound(err) {
		return err
	}
	now := metav1.NewTime(r.now())
	job.Status.CompletionTime = &now
	job.Status.Result = &executionv1alpha1.JobResult{Reason: "Cancelled"}
	return r.setCondition(ctx, job, metav1.Condition{Type: executionv1alpha1.JobConditionSucceeded, Status: metav1.ConditionFalse, Reason: "Cancelled", Message: "Execution cancelled"})
}

func (r *JobReconciler) finalize(ctx context.Context, job *executionv1alpha1.Job) error {
	if !containsString(job.Finalizers, jobFinalizer) {
		return nil
	}
	backend := &batchv1.Job{}
	key := types.NamespacedName{Namespace: job.Namespace, Name: job.Name}
	if err := r.Get(ctx, key, backend); err == nil {
		if err := kubernetesadapter.ValidateOwnership(job, backend); err != nil {
			return err
		}
		if err := r.Delete(ctx, backend, client.PropagationPolicy(metav1.DeletePropagationBackground)); err != nil && !apierrors.IsNotFound(err) {
			return err
		}
		return nil
	} else if !apierrors.IsNotFound(err) {
		return err
	}
	job.Finalizers = removeString(job.Finalizers, jobFinalizer)
	return r.Update(ctx, job)
}

func (r *JobReconciler) setCondition(ctx context.Context, job *executionv1alpha1.Job, next metav1.Condition) error {
	meta.SetStatusCondition(&job.Status.Conditions, condition(job, next.Type, next.Status, next.Reason, next.Message))
	job.Status.ObservedGeneration = job.Generation
	return r.Status().Update(ctx, job)
}

func condition(job *executionv1alpha1.Job, conditionType string, status metav1.ConditionStatus, reason, message string) metav1.Condition {
	return metav1.Condition{Type: conditionType, Status: status, Reason: reason, Message: message, ObservedGeneration: job.Generation}
}

func conditionTrue(conditions []metav1.Condition, conditionType string) bool {
	current := meta.FindStatusCondition(conditions, conditionType)
	return current != nil && current.Status == metav1.ConditionTrue
}

func isTerminal(job *executionv1alpha1.Job) bool {
	current := meta.FindStatusCondition(job.Status.Conditions, executionv1alpha1.JobConditionSucceeded)
	return current != nil && (current.Status == metav1.ConditionTrue || current.Status == metav1.ConditionFalse)
}

func applyResourceDefaults(requested, defaults executionv1alpha1.ExecutionResourceRequirements) executionv1alpha1.ExecutionResourceRequirements {
	result := requested.DeepCopy()
	if result.Requests.CPU == nil && defaults.Requests.CPU != nil {
		result.Requests.CPU = copyQuantity(defaults.Requests.CPU)
	}
	if result.Requests.Memory == nil && defaults.Requests.Memory != nil {
		result.Requests.Memory = copyQuantity(defaults.Requests.Memory)
	}
	if result.Limits.CPU == nil && defaults.Limits.CPU != nil {
		result.Limits.CPU = copyQuantity(defaults.Limits.CPU)
	}
	if result.Limits.Memory == nil && defaults.Limits.Memory != nil {
		result.Limits.Memory = copyQuantity(defaults.Limits.Memory)
	}
	return *result
}

func copyQuantity(value *resource.Quantity) *resource.Quantity {
	copy := value.DeepCopy()
	return &copy
}

func validateResources(resources executionv1alpha1.ExecutionResourceRequirements) error {
	if resources.Requests.CPU != nil && resources.Limits.CPU != nil && resources.Requests.CPU.Cmp(*resources.Limits.CPU) > 0 {
		return fmt.Errorf("CPU request must not exceed limit")
	}
	if resources.Requests.Memory != nil && resources.Limits.Memory != nil && resources.Requests.Memory.Cmp(*resources.Limits.Memory) > 0 {
		return fmt.Errorf("memory request must not exceed limit")
	}
	return nil
}

func containsString(values []string, target string) bool {
	return slices.Contains(values, target)
}

func removeString(values []string, target string) []string {
	result := values[:0]
	for _, value := range values {
		if value != target {
			result = append(result, value)
		}
	}
	return result
}

func (r *JobReconciler) now() time.Time {
	if r.Now != nil {
		return r.Now()
	}
	return time.Now()
}

func (r *JobReconciler) SetupWithManager(manager ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(manager).
		For(&executionv1alpha1.Job{}).
		Owns(&batchv1.Job{}).
		Named("execution-job").
		Complete(r)
}