wenshu/k8s-orchestrator/internal/controller/sparkapplication_cleaner.go

package controller

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	orcv1 "github.com/cyrus/spark-orchestrator/api/v1alpha1"
	"github.com/cyrus/spark-orchestrator/internal/k8s"
	"github.com/cyrus/spark-orchestrator/internal/service"
	"github.com/cyrus/spark-orchestrator/internal/util"
)

// SparkApplicationCleanerReconciler periodically archives completed SparkApplication logs
// and deletes SparkApplications after logs have been persisted or the retry grace expires.
type SparkApplicationCleanerReconciler struct {
	client.Client
	TTLSeconds                  int64
	ArchiveGraceSeconds         int64
	BackendBaseURL              string
	BackendTimeout              time.Duration
	BackendInternalHeader       string
	BackendInternalToken        string
	OrchestratorSecretNamespace string
}

func (r *SparkApplicationCleanerReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	logger := log.FromContext(ctx)

	var namespaces []string
	if err := r.listManagedNamespaces(ctx, &namespaces); err != nil {
		logger.Error(err, "failed to list managed namespaces")
		return ctrl.Result{RequeueAfter: 1 * time.Minute}, nil
	}

	totalDeleted := 0
	for _, ns := range namespaces {
		deleted, err := r.cleanupNamespace(ctx, ns)
		if err != nil {
			logger.Error(err, "failed to cleanup namespace", "namespace", ns)
			continue
		}
		totalDeleted += deleted
	}

	if totalDeleted > 0 {
		logger.Info("SparkApplication cleanup completed", "deleted", totalDeleted)
	}

	return ctrl.Result{RequeueAfter: 5 * time.Minute}, nil
}

func (r *SparkApplicationCleanerReconciler) cleanupNamespace(ctx context.Context, namespace string) (int, error) {
	completed, err := k8s.ListCompletedSparkApplications(ctx, r.Client, namespace)
	if err != nil {
		return 0, err
	}

	grace := r.archiveGrace()
	orchestratorSvc := service.NewOrchestratorService(r.Client, r.orchestratorSecretNamespace())
	deleted := 0

	for _, app := range completed {
		terminationTime, ok := r.getTerminationTime(&app)
		if ok == false {
			continue
		}

		ttl := r.cleanupTTLForApp(&app)
		age := time.Since(terminationTime)
		if age < ttl {
			continue
		}

		archived, archiveErr := r.archiveSparkApplication(ctx, orchestratorSvc, &app)
		if archived {
			if err := r.deleteSparkApplication(ctx, &app); err == nil {
				deleted++
			} else if apierrors.IsNotFound(err) {
				deleted++
			}
			continue
		}

		if age < ttl+grace {
			if archiveErr != nil {
				log.FromContext(ctx).Info("spark log archive will be retried", "namespace", namespace, "name", app.GetName(), "error", archiveErr.Error())
			}
			continue
		}

		if err := r.deleteSparkApplication(ctx, &app); err == nil || apierrors.IsNotFound(err) {
			deleted++
		}
	}

	return deleted, nil
}

func (r *SparkApplicationCleanerReconciler) archiveSparkApplication(
	ctx context.Context,
	orchestratorSvc *service.OrchestratorService,
	app *unstructured.Unstructured,
) (bool, error) {
	if app == nil {
		return false, fmt.Errorf("spark application is nil")
	}
	if orchestratorSvc == nil {
		return false, fmt.Errorf("orchestrator service is nil")
	}
	if strings.TrimSpace(r.BackendBaseURL) == "" {
		return false, fmt.Errorf("backend base url is empty")
	}

	operationID, err := r.findSparkOperationID(ctx, app)
	if err != nil {
		return false, err
	}

	result, err := orchestratorSvc.GetSparkJobResultByOperation(ctx, operationID)
	if err != nil {
		return false, err
	}

	payload, err := json.Marshal(result)
	if err != nil {
		return false, err
	}

	url := strings.TrimRight(strings.TrimSpace(r.BackendBaseURL), "/")
	url += "/api/data/internal/task-exec/spark-jobs/" + operationID + "/archive"

	httpClient := &http.Client{Timeout: r.backendTimeout()}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
	if err != nil {
		return false, err
	}
	req.Header.Set("Content-Type", "application/json")
	if token := strings.TrimSpace(r.BackendInternalToken); token != "" {
		headerName := strings.TrimSpace(r.BackendInternalHeader)
		if headerName == "" {
			headerName = "X-Internal-Token"
		}
		req.Header.Set(headerName, token)
	}

	resp, err := httpClient.Do(req)
	if err != nil {
		return false, err
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return false, fmt.Errorf("archive request failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(body)))
	}
	return true, nil
}

func (r *SparkApplicationCleanerReconciler) findSparkOperationID(ctx context.Context, app *unstructured.Unstructured) (string, error) {
	if app == nil {
		return "", fmt.Errorf("spark application is nil")
	}
	namespace := strings.TrimSpace(app.GetNamespace())
	name := strings.TrimSpace(app.GetName())
	if namespace == "" || name == "" {
		return "", fmt.Errorf("spark application namespace or name is empty")
	}

	var operations orcv1.ClusterOperationList
	if err := r.Client.List(ctx, &operations, client.MatchingLabels{
		util.LabelOperationType: string(orcv1.OperationSubmitSparkJob),
	}); err != nil {
		return "", err
	}

	targetResourceRef := namespace + "/" + name
	for i := range operations.Items {
		op := &operations.Items[i]
		if op.Spec.Type != orcv1.OperationSubmitSparkJob {
			continue
		}
		if strings.TrimSpace(op.Status.ResourceRef) != targetResourceRef {
			continue
		}
		if strings.TrimSpace(op.Name) == "" {
			continue
		}
		return op.Name, nil
	}

	return "", fmt.Errorf("spark operation not found for resourceRef=%s", targetResourceRef)
}

func (r *SparkApplicationCleanerReconciler) deleteSparkApplication(ctx context.Context, app *unstructured.Unstructured) error {
	if app == nil {
		return fmt.Errorf("spark application is nil")
	}
	return r.Client.Delete(ctx, app)
}

func (r *SparkApplicationCleanerReconciler) cleanupTTLForApp(app *unstructured.Unstructured) time.Duration {
	ttl := r.cleanupTTL()
	if app == nil {
		return ttl
	}
	annotations := app.GetAnnotations()
	if ttlStr, ok := annotations["spark.cyrus.io/ttl-seconds"]; ok {
		var customTTL int64
		if _, err := fmt.Sscanf(ttlStr, "%d", &customTTL); err == nil && customTTL > 0 {
			return time.Duration(customTTL) * time.Second
		}
	}
	return ttl
}

func (r *SparkApplicationCleanerReconciler) cleanupTTL() time.Duration {
	ttl := time.Duration(r.TTLSeconds) * time.Second
	if ttl <= 0 {
		ttl = time.Duration(k8s.SparkApplicationTTL) * time.Second
	}
	return ttl
}

func (r *SparkApplicationCleanerReconciler) archiveGrace() time.Duration {
	grace := time.Duration(r.ArchiveGraceSeconds) * time.Second
	if grace <= 0 {
		grace = time.Duration(k8s.SparkApplicationTTL) * time.Second
	}
	return grace
}

func (r *SparkApplicationCleanerReconciler) backendTimeout() time.Duration {
	if r.BackendTimeout <= 0 {
		return 15 * time.Second
	}
	return r.BackendTimeout
}

func (r *SparkApplicationCleanerReconciler) orchestratorSecretNamespace() string {
	if strings.TrimSpace(r.OrchestratorSecretNamespace) == "" {
		return "orchestrator-system"
	}
	return strings.TrimSpace(r.OrchestratorSecretNamespace)
}

func (r *SparkApplicationCleanerReconciler) getTerminationTime(app *unstructured.Unstructured) (time.Time, bool) {
	if app == nil {
		return time.Time{}, false
	}

	status, ok := app.Object["status"].(map[string]interface{})
	if !ok {
		return time.Time{}, false
	}

	terminationTimeStr, _ := status["terminationTime"].(string)
	if terminationTimeStr == "" {
		appState, _ := status["applicationState"].(map[string]interface{})
		if appState != nil {
			terminationTimeStr, _ = appState["stateTransitionTime"].(string)
		}
	}
	if terminationTimeStr == "" {
		return time.Time{}, false
	}

	terminationTime, err := time.Parse(time.RFC3339, terminationTimeStr)
	if err != nil {
		return time.Time{}, false
	}
	return terminationTime, true
}

func (r *SparkApplicationCleanerReconciler) listManagedNamespaces(ctx context.Context, namespaces *[]string) error {
	list := &unstructured.UnstructuredList{
		Object: map[string]interface{}{
			"apiVersion": k8s.SparkApplicationAPIVersion,
			"kind":       k8s.SparkApplicationKind,
		},
	}

	if err := r.Client.List(ctx, list); err != nil {
		return err
	}

	seen := make(map[string]struct{})
	for _, app := range list.Items {
		ns := app.GetNamespace()
		if _, exists := seen[ns]; !exists {
			seen[ns] = struct{}{}
			*namespaces = append(*namespaces, ns)
		}
	}

	return nil
}

func (r *SparkApplicationCleanerReconciler) SetupWithManager(mgr ctrl.Manager) error {
	if r.TTLSeconds == 0 {
		r.TTLSeconds = k8s.SparkApplicationTTL
	}
	if r.ArchiveGraceSeconds == 0 {
		r.ArchiveGraceSeconds = k8s.SparkApplicationTTL
	}
	if r.BackendTimeout == 0 {
		r.BackendTimeout = 15 * time.Second
	}
	if strings.TrimSpace(r.BackendInternalHeader) == "" {
		r.BackendInternalHeader = "X-Internal-Token"
	}
	if strings.TrimSpace(r.OrchestratorSecretNamespace) == "" {
		r.OrchestratorSecretNamespace = "orchestrator-system"
	}

	return ctrl.NewControllerManagedBy(mgr).
		For(&unstructured.Unstructured{
			Object: map[string]interface{}{
				"apiVersion": k8s.SparkApplicationAPIVersion,
				"kind":       k8s.SparkApplicationKind,
			},
		}).
		Complete(r)
}