distributor/internal/app/prune.go

package app

import (
	"context"
	"fmt"
	"time"

	"gitea.maximumdirect.net/eric/distributor/internal/config"
	"gitea.maximumdirect.net/eric/distributor/internal/state"
	"gitea.maximumdirect.net/eric/distributor/internal/storage"
)

type PruneOptions struct {
	ConfigPath    string
	PipelineID    string
	DestinationID string
	DryRun        bool
	Now           time.Time
}

type PrunePlanOptions struct {
	PipelineID    string
	DestinationID string
	Now           time.Time
}

type PrunePlanReport struct {
	PipelineID       string              `json:"pipeline_id"`
	DestinationID    string              `json:"destination_id"`
	OwnerScope       PruneOwnerScope     `json:"owner_scope"`
	Enabled          bool                `json:"enabled"`
	CheckedCount     int                 `json:"checked_count"`
	PrunedOutputs    []PruneOutputRecord `json:"pruned_outputs"`
	PreservedOutputs []PruneOutputRecord `json:"preserved_outputs"`
}

type PruneReport struct {
	PipelineID       string              `json:"pipeline_id"`
	DestinationID    string              `json:"destination_id"`
	Backend          string              `json:"backend"`
	RootPath         string              `json:"root_path"`
	OwnerScope       PruneOwnerScope     `json:"owner_scope"`
	Enabled          bool                `json:"enabled"`
	CheckedCount     int                 `json:"checked_count"`
	PlannedOutputs   []PruneOutputRecord `json:"planned_outputs"`
	DeletedOutputs   []PruneOutputRecord `json:"deleted_outputs"`
	PreservedOutputs []PruneOutputRecord `json:"preserved_outputs"`
	FailedOutput     *PruneOutputRecord  `json:"failed_output,omitempty"`
	StateChanged     bool                `json:"state_changed"`
	WouldChange      bool                `json:"would_change"`
	DryRun           bool                `json:"dry_run"`
}

type PruneOwnerScope struct {
	PipelineID    string `json:"pipeline_id"`
	DestinationID string `json:"destination_id"`
}

type PruneOutputRecord struct {
	Path      string           `json:"path"`
	UpdatedAt string           `json:"updated_at"`
	Owner     *PruneOwnerScope `json:"owner,omitempty"`
}

func Prune(ctx context.Context, options PruneOptions) (PruneReport, error) {
	if err := ctx.Err(); err != nil {
		return PruneReport{}, err
	}

	setup, err := loadRuntimeSetup(options.ConfigPath)
	if err != nil {
		return PruneReport{}, err
	}
	return pruneSetup(ctx, setup, options)
}

func pruneConfigWithBackendFactory(ctx context.Context, cfg config.Config, options PruneOptions, provider backendFactoryProvider) (PruneReport, error) {
	setup, err := runtimeSetupFromConfig("", cfg)
	if err != nil {
		return PruneReport{}, err
	}
	return pruneSetupWithBackendFactory(ctx, setup, options, provider)
}

func pruneSetup(ctx context.Context, setup runtimeSetup, options PruneOptions) (PruneReport, error) {
	return pruneSetupWithBackendFactory(ctx, setup, options, newBackendFactoryWithEnvironment)
}

func pruneSetupWithBackendFactory(ctx context.Context, setup runtimeSetup, options PruneOptions, provider backendFactoryProvider) (PruneReport, error) {
	if err := requirePruneScope(options); err != nil {
		return PruneReport{}, err
	}
	pipeline, ok := findPipeline(setup.Config, options.PipelineID)
	if !ok {
		return PruneReport{}, PipelineNotFoundError{ID: options.PipelineID}
	}
	destination, ok := findDestination(pipeline, options.DestinationID)
	if !ok {
		return PruneReport{}, fmt.Errorf("pipeline %s destination %s not found", options.PipelineID, options.DestinationID)
	}

	backends := provider(setup.Environment)
	destinationBackend, err := backends.openDestination(ctx, destination)
	if err != nil {
		return PruneReport{}, err
	}
	defer closeBackend(destinationBackend)

	return executePrune(ctx, destinationBackend, pipeline, destination, options)
}

func requirePruneScope(options PruneOptions) error {
	if options.PipelineID == "" {
		return fmt.Errorf("pipeline id is required")
	}
	if options.DestinationID == "" {
		return fmt.Errorf("destination id is required")
	}
	return nil
}

func executePrune(ctx context.Context, backend storage.Backend, pipeline config.Pipeline, destination config.Destination, options PruneOptions) (PruneReport, error) {
	now := options.Now
	if now.IsZero() {
		now = time.Now().UTC()
	} else {
		now = now.UTC()
	}

	statePath, err := storage.StatePath("")
	if err != nil {
		return PruneReport{}, err
	}
	data, err := backend.ReadFile(ctx, statePath)
	if err != nil {
		return PruneReport{}, err
	}
	document, err := state.ParseDocument(data)
	if err != nil {
		return PruneReport{}, err
	}

	plan, err := PlanPrune(document, destination.Retention.Prune, PrunePlanOptions{
		PipelineID:    pipeline.ID,
		DestinationID: destination.ID,
		Now:           now,
	})
	if err != nil {
		return PruneReport{}, err
	}
	report := PruneReport{
		PipelineID:       pipeline.ID,
		DestinationID:    destination.ID,
		Backend:          destination.Backend,
		RootPath:         destinationRootPath(destination),
		OwnerScope:       plan.OwnerScope,
		Enabled:          plan.Enabled,
		CheckedCount:     plan.CheckedCount,
		PlannedOutputs:   plan.PrunedOutputs,
		DeletedOutputs:   []PruneOutputRecord{},
		PreservedOutputs: plan.PreservedOutputs,
		DryRun:           options.DryRun,
	}
	report.WouldChange = options.DryRun && len(report.PlannedOutputs) > 0
	if options.DryRun || len(report.PlannedOutputs) == 0 {
		return report, nil
	}

	deletedPaths := make([]string, 0, len(report.PlannedOutputs))
	for _, output := range report.PlannedOutputs {
		err := backend.DeleteManagedOutputs(ctx, "", []string{output.Path}, storage.DeleteOptions{
			IgnoreMissing:  true,
			PruneEmptyDirs: true,
		})
		if err != nil {
			failed := output
			report.FailedOutput = &failed
			if len(deletedPaths) > 0 {
				changed, writeErr := removePrunedStateRecords(ctx, backend, statePath, document, state.CurrentOwnerScope(pipeline.ID, destination.ID), deletedPaths, now)
				report.StateChanged = changed
				report.DeletedOutputs = report.PlannedOutputs[:len(deletedPaths)]
				if writeErr != nil {
					return report, writeErr
				}
			}
			return report, err
		}
		deletedPaths = append(deletedPaths, output.Path)
	}

	changed, err := removePrunedStateRecords(ctx, backend, statePath, document, state.CurrentOwnerScope(pipeline.ID, destination.ID), deletedPaths, now)
	report.StateChanged = changed
	report.DeletedOutputs = report.PlannedOutputs
	return report, err
}

func removePrunedStateRecords(ctx context.Context, backend storage.Backend, statePath string, document state.StateDocument, scope state.OwnerScope, paths []string, now time.Time) (bool, error) {
	if len(paths) == 0 {
		return false, nil
	}
	if document.SingleOwner != nil {
		next, changed := state.RemoveMissingOutputs(*document.SingleOwner, paths)
		if !changed {
			return false, nil
		}
		next.UpdatedAt = now
		if err := state.Validate(next); err != nil {
			return false, err
		}
		return true, writeRepairedState(ctx, backend, statePath, next)
	}
	if document.SharedRoot != nil {
		next, changed := state.RemoveMissingSharedRootOwnerOutputs(*document.SharedRoot, scope, paths)
		if !changed {
			return false, nil
		}
		next.UpdatedAt = now
		if err := state.ValidateSharedRoot(next); err != nil {
			return false, err
		}
		return true, writeRepairedState(ctx, backend, statePath, next)
	}
	return false, fmt.Errorf("destination state document is empty")
}

func PlanPrune(document state.StateDocument, policy config.PrunePolicy, options PrunePlanOptions) (PrunePlanReport, error) {
	scope := state.CurrentOwnerScope(options.PipelineID, options.DestinationID)
	report := PrunePlanReport{
		PipelineID:       options.PipelineID,
		DestinationID:    options.DestinationID,
		OwnerScope:       PruneOwnerScope{PipelineID: scope.PipelineID, DestinationID: scope.DestinationID},
		Enabled:          policy.Enabled,
		PrunedOutputs:    []PruneOutputRecord{},
		PreservedOutputs: []PruneOutputRecord{},
	}
	if !policy.Enabled {
		return report, nil
	}

	candidates, err := pruneCandidatesForDocument(document, scope)
	if err != nil {
		return PrunePlanReport{}, err
	}
	report.CheckedCount = len(candidates)
	plan := state.PlanPrune(candidates, state.PrunePlanOptions{
		Now:        options.Now,
		OlderThan:  pruneOlderThan(policy),
		KeepLatest: policy.KeepLatest,
	})
	report.PrunedOutputs = pruneOutputRecords(plan.Pruned)
	report.PreservedOutputs = pruneOutputRecords(plan.Preserved)
	return report, nil
}

func pruneCandidatesForDocument(document state.StateDocument, scope state.OwnerScope) ([]state.PruneCandidate, error) {
	if document.SingleOwner != nil {
		singleOwner := *document.SingleOwner
		if singleOwner.PipelineID != scope.PipelineID || singleOwner.DestinationID != scope.DestinationID {
			return nil, fmt.Errorf("state owner is %s/%s, not %s/%s", singleOwner.PipelineID, singleOwner.DestinationID, scope.PipelineID, scope.DestinationID)
		}
		return state.SingleOwnerPruneCandidates(singleOwner), nil
	}
	if document.SharedRoot != nil {
		return state.SharedRootPruneCandidates(*document.SharedRoot, scope), nil
	}
	return nil, fmt.Errorf("destination state document is empty")
}

func pruneOlderThan(policy config.PrunePolicy) *time.Duration {
	if policy.OlderThan == nil {
		return nil
	}
	duration := policy.OlderThan.AsDuration()
	return &duration
}

func pruneOutputRecords(candidates []state.PruneCandidate) []PruneOutputRecord {
	records := make([]PruneOutputRecord, 0, len(candidates))
	for _, candidate := range candidates {
		var owner *PruneOwnerScope
		if candidate.Owner != nil {
			owner = &PruneOwnerScope{
				PipelineID:    candidate.Owner.PipelineID,
				DestinationID: candidate.Owner.DestinationID,
			}
		}
		records = append(records, PruneOutputRecord{
			Path:      candidate.Path,
			UpdatedAt: candidate.UpdatedAt.UTC().Format(time.RFC3339),
			Owner:     owner,
		})
	}
	return records
}