URBM/internal/platform/workloads.go

package platform

import (
	"context"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/backupper-unraid/backupper/internal/model"
)

type WorkloadManager struct{ RuntimeDir string }

type Prepared struct {
	Sources     []string
	Stopped     []model.Source
	MetadataDir string
	Kind        model.JobType
}

func (w *WorkloadManager) Prepare(ctx context.Context, job model.Job) (Prepared, error) {
	prepared := Prepared{Kind: job.Type}
	for _, source := range job.Sources {
		if source.Path != "" {
			prepared.Sources = append(prepared.Sources, source.Path)
		}
	}
	if job.Type != model.JobDocker && job.Type != model.JobVM {
		return prepared, nil
	}
	metadataDir, err := os.MkdirTemp(w.RuntimeDir, "metadata-"+job.ID+"-")
	if err != nil {
		return prepared, err
	}
	prepared.MetadataDir = metadataDir
	prepared.Sources = append(prepared.Sources, metadataDir)
	if job.Type == model.JobDocker {
		if _, err := os.Stat("/boot/config/plugins/dockerMan/templates-user"); err == nil {
			prepared.Sources = append(prepared.Sources, "/boot/config/plugins/dockerMan/templates-user")
		}
	}
	for _, source := range job.Sources {
		if source.WorkloadID == "" {
			continue
		}
		discovered, err := w.captureMetadata(ctx, job.Type, source.WorkloadID, metadataDir)
		if err != nil {
			w.Cleanup(context.Background(), prepared)
			return Prepared{}, err
		}
		prepared.Sources = append(prepared.Sources, discovered...)
		if job.Consistency.Mode == "stop" {
			if err := w.stop(ctx, job.Type, source.WorkloadID, time.Duration(job.ShutdownSecs)*time.Second); err != nil {
				w.Cleanup(context.Background(), prepared)
				return Prepared{}, err
			}
			prepared.Stopped = append(prepared.Stopped, source)
		}
	}
	prepared.Sources = uniqueStrings(prepared.Sources)
	return prepared, nil
}

func (w *WorkloadManager) Cleanup(ctx context.Context, prepared Prepared) error {
	var first error
	for i := len(prepared.Stopped) - 1; i >= 0; i-- {
		source := prepared.Stopped[i]
		if err := w.start(ctx, prepared.Kind, source.WorkloadID); err != nil && first == nil {
			first = err
		}
	}
	if prepared.MetadataDir != "" {
		_ = os.RemoveAll(prepared.MetadataDir)
	}
	return first
}

func (w *WorkloadManager) captureMetadata(ctx context.Context, kind model.JobType, id, dir string) ([]string, error) {
	var cmd *exec.Cmd
	var name string
	if kind == model.JobDocker {
		cmd, name = exec.CommandContext(ctx, "docker", "inspect", id), "docker-"+safeName(id)+".json"
	} else {
		cmd, name = exec.CommandContext(ctx, "virsh", "dumpxml", id), "vm-"+safeName(id)+".xml"
	}
	output, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("source: capture metadata for %s: %w", id, err)
	}
	if err := os.WriteFile(filepath.Join(dir, name), output, 0600); err != nil {
		return nil, err
	}
	if kind == model.JobDocker {
		var inspected []struct {
			Mounts []struct {
				Type   string `json:"Type"`
				Source string `json:"Source"`
			} `json:"Mounts"`
		}
		if err := json.Unmarshal(output, &inspected); err != nil {
			return nil, fmt.Errorf("source: decode docker metadata: %w", err)
		}
		var paths []string
		if len(inspected) > 0 {
			for _, mount := range inspected[0].Mounts {
				if mount.Type == "bind" && filepath.IsAbs(mount.Source) {
					paths = append(paths, mount.Source)
				}
			}
		}
		return paths, nil
	}
	var domain struct {
		OS struct {
			NVRAM string `xml:"nvram"`
		} `xml:"os"`
		Devices struct {
			Disks []struct {
				Device string `xml:"device,attr"`
				Source struct {
					File string `xml:"file,attr"`
					Dev  string `xml:"dev,attr"`
				} `xml:"source"`
			} `xml:"disk"`
		} `xml:"devices"`
	}
	if err := xml.Unmarshal(output, &domain); err != nil {
		return nil, fmt.Errorf("source: decode vm metadata: %w", err)
	}
	paths := []string{}
	if filepath.IsAbs(strings.TrimSpace(domain.OS.NVRAM)) {
		paths = append(paths, strings.TrimSpace(domain.OS.NVRAM))
	}
	for _, disk := range domain.Devices.Disks {
		if disk.Device != "disk" && disk.Device != "lun" {
			continue
		}
		path := disk.Source.File
		if path == "" {
			path = disk.Source.Dev
		}
		if filepath.IsAbs(path) {
			paths = append(paths, path)
		}
	}
	return paths, nil
}

func (w *WorkloadManager) stop(ctx context.Context, kind model.JobType, id string, timeout time.Duration) error {
	if timeout <= 0 {
		timeout = 120 * time.Second
	}
	if kind == model.JobDocker {
		seconds := fmt.Sprint(int(timeout.Seconds()))
		if err := exec.CommandContext(ctx, "docker", "stop", "--time", seconds, id).Run(); err != nil {
			return fmt.Errorf("source: stop docker %s: %w", id, err)
		}
		return nil
	}
	if err := exec.CommandContext(ctx, "virsh", "shutdown", id).Run(); err != nil {
		return fmt.Errorf("source: shutdown vm %s: %w", id, err)
	}
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		output, _ := exec.CommandContext(ctx, "virsh", "domstate", id).Output()
		if strings.Contains(strings.ToLower(string(output)), "shut off") {
			return nil
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(2 * time.Second):
		}
	}
	return fmt.Errorf("source: vm %s did not stop before timeout", id)
}

func (w *WorkloadManager) start(ctx context.Context, kind model.JobType, id string) error {
	if kind == model.JobDocker {
		return exec.CommandContext(ctx, "docker", "start", id).Run()
	}
	return exec.CommandContext(ctx, "virsh", "start", id).Run()
}

func safeName(value string) string {
	value = strings.Map(func(r rune) rune {
		if r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '-' || r == '_' {
			return r
		}
		return '_'
	}, value)
	if value == "" {
		return "workload"
	}
	return value
}

func uniqueStrings(values []string) []string {
	seen := make(map[string]struct{}, len(values))
	result := make([]string, 0, len(values))
	for _, value := range values {
		value = filepath.Clean(value)
		if _, exists := seen[value]; exists {
			continue
		}
		seen[value] = struct{}{}
		result = append(result, value)
	}
	return result
}