nekonest-cloud/relay/internal/migration/manager.go

package migration

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"github.com/klarkxy/nekonest-cloud/relay/internal/controlplane"
	"github.com/klarkxy/nekonest-cloud/relay/internal/tenantbackup"
	"github.com/klarkxy/nekonest-cloud/relay/internal/tenantfs"
)

type Registry interface {
	Quiesce(tenantID string, generation int64) (tenantfs.Paths, error)
}

type ControlPlane interface {
	AdvanceMigration(context.Context, controlplane.MigrationAdvance) error
}

type Config struct {
	DataRoot      string
	BackupRoot    string
	Registry      Registry
	ControlPlane  ControlPlane
	MaxConcurrent int
	Now           func() time.Time
	Logger        *slog.Logger
}

type Manager struct {
	config    Config
	mu        sync.Mutex
	inFlight  map[string]struct{}
	semaphore chan struct{}
	wait      sync.WaitGroup
}

func New(config Config) (*Manager, error) {
	if config.DataRoot == "" || config.BackupRoot == "" || config.Registry == nil || config.ControlPlane == nil {
		return nil, errors.New("migration manager requires data, backup, registry, and control-plane ports")
	}
	if config.MaxConcurrent <= 0 {
		config.MaxConcurrent = 2
	}
	if config.MaxConcurrent > 8 {
		return nil, errors.New("migration concurrency is unreasonably high")
	}
	if config.Now == nil {
		config.Now = time.Now
	}
	if config.Logger == nil {
		config.Logger = slog.Default()
	}
	return &Manager{
		config: config, inFlight: make(map[string]struct{}),
		semaphore: make(chan struct{}, config.MaxConcurrent),
	}, nil
}

// Handle starts bounded, de-duplicated migration work returned by a trusted
// heartbeat. A later heartbeat safely retries any stage that did not commit.
func (manager *Manager) Handle(ctx context.Context, assignments []controlplane.MigrationAssignment) {
	for _, assignment := range assignments {
		manager.mu.Lock()
		if _, exists := manager.inFlight[assignment.MigrationID]; exists {
			manager.mu.Unlock()
			continue
		}
		select {
		case manager.semaphore <- struct{}{}:
			manager.inFlight[assignment.MigrationID] = struct{}{}
			manager.wait.Add(1)
			manager.mu.Unlock()
			go func(assignment controlplane.MigrationAssignment) {
				defer func() {
					<-manager.semaphore
					manager.mu.Lock()
					delete(manager.inFlight, assignment.MigrationID)
					manager.mu.Unlock()
					manager.wait.Done()
				}()
				if err := manager.process(ctx, assignment); err != nil && ctx.Err() == nil {
					manager.config.Logger.Error("Relay migration stage failed",
						"migration_id", assignment.MigrationID, "tenant_id", assignment.TenantID,
						"state", assignment.State, "error", err)
				}
			}(assignment)
		default:
			manager.mu.Unlock()
		}
	}
}

func (manager *Manager) process(ctx context.Context, assignment controlplane.MigrationAssignment) error {
	if assignment.MigrationID == "" || assignment.TenantID == "" {
		return errors.New("migration assignment is incomplete")
	}
	switch assignment.State {
	case "quiescing":
		if assignment.Role != "source" {
			return errors.New("quiescing assignment did not target the source")
		}
		if _, err := manager.config.Registry.Quiesce(assignment.TenantID, assignment.SourceGeneration); err != nil {
			return manager.fail(ctx, assignment, "relay_quiesce_failed", err)
		}
		backup, err := tenantbackup.Create(
			ctx, manager.config.DataRoot, manager.config.BackupRoot,
			assignment.TenantID, assignment.SourceGeneration, manager.config.Now(),
		)
		if err != nil {
			return manager.fail(ctx, assignment, "relay_backup_failed", err)
		}
		return manager.config.ControlPlane.AdvanceMigration(ctx, controlplane.MigrationAdvance{
			MigrationID: assignment.MigrationID, Action: "quiesced",
			BackupRef: backup.BackupRef, ManifestSHA256: backup.ManifestSHA256,
		})
	case "copying":
		if assignment.Role != "target" || assignment.BackupRef == "" || assignment.ManifestSHA256 == "" {
			return errors.New("copying assignment is incomplete")
		}
		backupPath, err := tenantbackup.ResolveReference(manager.config.BackupRoot, assignment.BackupRef)
		if err != nil {
			return manager.fail(ctx, assignment, "relay_backup_unavailable", err)
		}
		if _, err := tenantbackup.Restore(
			ctx, backupPath, manager.config.DataRoot, assignment.TenantID,
			assignment.SourceGeneration, assignment.ManifestSHA256,
		); err != nil {
			return manager.fail(ctx, assignment, "relay_restore_failed", err)
		}
		return manager.config.ControlPlane.AdvanceMigration(ctx, controlplane.MigrationAdvance{
			MigrationID: assignment.MigrationID, Action: "copied",
			BackupRef: assignment.BackupRef, ManifestSHA256: assignment.ManifestSHA256,
		})
	case "switching":
		if assignment.Role != "target" {
			return errors.New("switching assignment did not target the destination")
		}
		return manager.config.ControlPlane.AdvanceMigration(ctx, controlplane.MigrationAdvance{
			MigrationID: assignment.MigrationID, Action: "switched",
		})
	case "draining":
		if assignment.Role != "target" || assignment.FinalizeAfter == "" {
			return errors.New("draining assignment is incomplete")
		}
		finalizeAfter, err := time.Parse(time.RFC3339Nano, assignment.FinalizeAfter)
		if err != nil {
			return fmt.Errorf("invalid finalize fence: %w", err)
		}
		if manager.config.Now().Before(finalizeAfter) {
			return nil
		}
		return manager.config.ControlPlane.AdvanceMigration(ctx, controlplane.MigrationAdvance{
			MigrationID: assignment.MigrationID, Action: "finalized",
		})
	default:
		return fmt.Errorf("unsupported migration state %q", assignment.State)
	}
}

func (manager *Manager) fail(ctx context.Context, assignment controlplane.MigrationAssignment, code string, cause error) error {
	advanceErr := manager.config.ControlPlane.AdvanceMigration(ctx, controlplane.MigrationAdvance{
		MigrationID: assignment.MigrationID, Action: "failed", ErrorCode: code,
	})
	if advanceErr != nil {
		return errors.Join(cause, fmt.Errorf("report migration failure: %w", advanceErr))
	}
	return cause
}

func (manager *Manager) Wait() { manager.wait.Wait() }