ThothII/tools/tht/internal/pi/restart.go

package pi

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
)

var (
	ErrInterruptedRestart        = errors.New("a previous Pi restart is incomplete; recover lifecycle maintenance before another Pi lifecycle operation")
	errRestartImageDrift         = errors.New("core image changed during Pi restart")
	errRestartConfigurationDrift = errors.New("external endpoint configuration changed during Pi restart")
	errRestartMountDrift         = errors.New("core persistence mount contract changed during Pi restart")
	errRestartConfirmation       = &restartDiagnosticError{
		message: "restart requires --yes after reviewing the planned Pi core recreation",
		cause:   ErrConfirmationRequired,
	}
	errRestartActiveSessions = &restartDiagnosticError{
		message: "active sessions must be drained before restarting Pi; use --drain only after they are complete",
		cause:   ErrActiveSessions,
	}
)

type restartDiagnosticError struct {
	message string
	cause   error
}

func (e *restartDiagnosticError) Error() string { return e.message }

func (e *restartDiagnosticError) Unwrap() error { return e.cause }

type RestartRequest struct {
	StatePath       string
	UpdateStatePath string
	Confirm         bool
	Drain           bool
}

type RestartResult struct {
	StatePath string
	Version   string
}

func Restart(ctx context.Context, runner Runner, request RestartRequest) (RestartResult, error) {
	return restartWithHooks(ctx, runner, request, defaultLifecycleHooks)
}

func restartWithHooks(
	ctx context.Context,
	runner Runner,
	request RestartRequest,
	hooks lifecycleHooks,
) (result RestartResult, retErr error) {
	if err := validateRestartStatePaths(request.StatePath, request.UpdateStatePath); err != nil {
		return RestartResult{}, err
	}
	lock, err := acquireLock(request.StatePath)
	if err != nil {
		return RestartResult{StatePath: request.StatePath}, err
	}
	defer lock.Release()
	if !request.Confirm {
		return RestartResult{StatePath: request.StatePath}, errRestartConfirmation
	}
	if state, err := readState(request.UpdateStatePath); err == nil && stateNeedsRecovery(state) {
		return RestartResult{StatePath: request.StatePath}, ErrInterruptedUpdate
	} else if err != nil && !errors.Is(err, os.ErrNotExist) {
		return RestartResult{StatePath: request.StatePath}, err
	}
	if err := prepareLifecycleMutation(request.StatePath, hooks.removeFile); err != nil {
		return RestartResult{StatePath: request.StatePath}, err
	}
	clearMaintenance := true
	mutationStarted := false
	var state State
	defer func() {
		if !clearMaintenance {
			return
		}
		if clearErr := setMaintenance(context.Background(), runner, false); clearErr != nil {
			result.StatePath = request.StatePath
			if mutationStarted {
				clearMaintenance = false
				cause := clearErr
				if writeErr := hooks.writeState(request.StatePath, state); writeErr != nil {
					cause = errors.Join(cause, fmt.Errorf("restart recovery state could not be restored: %w", writeErr))
				}
				retErr = errors.Join(retErr, recoveryRequired("Pi restart maintenance cleanup failed", cause))
				return
			}
			retErr = errors.Join(retErr, fmt.Errorf("maintenance admission gate could not be cleared: %w", clearErr))
		}
	}()
	if err := setMaintenance(ctx, runner, true); err != nil {
		return RestartResult{StatePath: request.StatePath}, err
	}

	if err := waitForInactiveSessions(ctx, runner, request.Drain, hooks.sleep); err != nil {
		return RestartResult{StatePath: request.StatePath}, restartDiagnostic(err)
	}
	if err := Doctor(ctx, runner); err != nil {
		return RestartResult{StatePath: request.StatePath}, err
	}
	version, err := Status(ctx, runner)
	if err != nil {
		return RestartResult{StatePath: request.StatePath}, err
	}
	configured, err := renderedCore(ctx, runner)
	if err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, err
	}
	previous, err := runningImage(ctx, runner, configured.Reference)
	if err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, err
	}
	previous.ConfigurationSHA = configured.ConfigurationSHA
	transaction := lifecycleTransaction(request.StatePath)
	previous.Reference = lifecycleImageTag(transaction, "restart")
	if err := tagImage(ctx, runner, previous.ID, previous.Reference, "restart image pin"); err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, err
	}
	state = State{
		Transaction: transaction,
		Phase:       PhasePreflight,
		Target:      Target{Version: version, Source: "restart"},
		Previous:    previous,
	}
	if err := hooks.writeState(request.StatePath, state); err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, err
	}
	overridePath := lifecycleOverridePath(request.StatePath, transaction)
	if err := writeLifecycleOverride(overridePath, previous.Reference); err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, err
	}
	lifecycle := composeOverrideRunner{Runner: runner, path: overridePath}
	if running, err := activeSessions(ctx, runner); err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, err
	} else if running {
		return RestartResult{StatePath: request.StatePath, Version: version}, errRestartActiveSessions
	}
	state.MutationStarted = true
	if err := hooks.writeState(request.StatePath, state); err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, err
	}
	mutationStarted = true
	clearMaintenance = false
	if err := recreateCoreWithoutImageChanges(ctx, lifecycle); err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, recoveryRequired("Pi restart core recreation failed", err)
	}
	if err := ensureMaintenance(ctx, lifecycle); err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, recoveryRequired("Pi restart maintenance proof failed", err)
	}
	state.Phase = PhaseRecreated
	if err := hooks.writeState(request.StatePath, state); err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, recoveryRequired("Pi restart recreation state could not be recorded", err)
	}
	if err := verifyRestart(ctx, lifecycle, version, previous); err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, recoveryRequired("Pi restart verification failed", err)
	}
	state.Phase = PhaseVerified
	if err := hooks.writeState(request.StatePath, state); err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, recoveryRequired("Pi restart verification state could not be recorded", err)
	}
	if err := hooks.removeFile(overridePath); err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, recoveryRequired("Pi restart image override could not be removed", err)
	}
	if err := hooks.removeFile(request.StatePath); err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, recoveryRequired("Pi restart recovery state could not be removed", err)
	}
	clearMaintenance = true
	return RestartResult{StatePath: request.StatePath, Version: version}, nil
}

func recreateCoreWithoutImageChanges(ctx context.Context, runner Runner) error {
	result, err := runCompose(
		ctx,
		runner,
		"up",
		"--detach",
		"--wait",
		"--wait-timeout",
		"45",
		"--no-deps",
		"--force-recreate",
		"--no-build",
		"--pull",
		"never",
		"core",
	)
	if err != nil {
		return commandError("core recreation", result, err)
	}
	return nil
}

func verifyRestart(ctx context.Context, runner Runner, wanted string, previous Image) error {
	if err := Doctor(ctx, runner); err != nil {
		return err
	}
	version, err := Status(ctx, runner)
	if err != nil {
		return err
	}
	if version != wanted {
		return errors.New("Pi version changed during core restart")
	}
	configured, err := renderedCore(ctx, runner)
	if err != nil {
		return err
	}
	after, err := runningImage(ctx, runner, configured.Reference)
	if err != nil {
		return err
	}
	if after.ID != previous.ID {
		return errRestartImageDrift
	}
	if configured.ConfigurationSHA != previous.ConfigurationSHA {
		return errRestartConfigurationDrift
	}
	if !sameMounts(previous.Mounts, after.Mounts) {
		return errRestartMountDrift
	}
	return nil
}

func RecoverLifecycleMaintenance(
	ctx context.Context,
	runner Runner,
	updateStatePath string,
	restartStatePath string,
	confirm bool,
) error {
	if err := validateRestartStatePaths(restartStatePath, updateStatePath); err != nil {
		return err
	}
	if !confirm {
		return errRestartConfirmation
	}
	lock, err := acquireLock(restartStatePath)
	if err != nil {
		return err
	}
	defer lock.Release()

	restartState, restartErr := readState(restartStatePath)
	if restartErr == nil {
		if err := validateRestartRecoveryState(restartState); err != nil {
			return err
		}
		restartOverride := lifecycleOverridePath(restartStatePath, restartState.Transaction)
		if restartState.MutationStarted {
			if err := tagImage(ctx, runner, restartState.Previous.ID, restartState.Previous.Reference, "restart recovery image pin"); err != nil {
				return recoveryRequired("Pi restart recovery image pin could not be restored", err)
			}
			if err := writeLifecycleOverride(restartOverride, restartState.Previous.Reference); err != nil {
				return recoveryRequired("Pi restart recovery image override could not be restored", err)
			}
			lifecycle := composeOverrideRunner{Runner: runner, path: restartOverride}
			if err := ensureMaintenance(ctx, lifecycle); err != nil {
				return recoveryRequired("Pi restart maintenance recovery failed", err)
			}
			if err := verifyRestart(ctx, lifecycle, restartState.Target.Version, restartState.Previous); err != nil {
				return recoveryRequired("Pi restart recovery verification failed", err)
			}
		}
		if err := durableRemove(restartOverride); err != nil {
			return recoveryRequired("Pi restart recovery image override could not be removed", err)
		}
		if err := durableRemove(restartStatePath); err != nil {
			return recoveryRequired("Pi restart recovery state could not be removed", err)
		}
	} else if !errors.Is(restartErr, os.ErrNotExist) {
		return restartErr
	}
	return recoverMaintenanceLocked(ctx, runner, updateStatePath)
}

func restartDiagnostic(err error) error {
	if errors.Is(err, ErrActiveSessions) {
		return errRestartActiveSessions
	}
	return err
}

func validateRestartRecoveryState(state State) error {
	invalid := func(reason string) error {
		return fmt.Errorf("%w: restart recovery state is invalid: %s", ErrInvalidRequest, reason)
	}
	if state.Transaction == "" {
		return invalid("transaction is missing")
	}
	if state.Target.Source != "restart" || state.Target.Version == "" {
		return invalid("target is not a restart with a recorded version")
	}
	if state.Previous.ConfigurationSHA == "" {
		return invalid("previous external configuration identity is missing")
	}
	if state.Candidate.ID != "" || state.Candidate.Reference != "" || len(state.Candidate.Mounts) != 0 ||
		state.Candidate.MountFingerprint != "" || state.Candidate.ConfigurationSHA != "" {
		return invalid("restart state contains image candidate metadata")
	}
	switch state.Phase {
	case PhasePreflight:
		return nil
	case PhaseRecreated, PhaseVerified:
		if state.MutationStarted {
			return nil
		}
		return invalid("post-recreation phase has no mutation marker")
	default:
		return invalid("phase is not valid for restart")
	}
}

func validateRestartStatePaths(restartStatePath, updateStatePath string) error {
	if restartStatePath == "" {
		return errors.New("restart state path is required")
	}
	if updateStatePath == "" {
		return errors.New("update state path is required")
	}
	if filepath.Clean(restartStatePath) == filepath.Clean(updateStatePath) {
		return fmt.Errorf("%w: restart and update state paths must remain separate", ErrInvalidRequest)
	}
	if lifecycleLockPath(restartStatePath) != lifecycleLockPath(updateStatePath) {
		return fmt.Errorf("%w: restart and update state paths must share one lifecycle control directory", ErrInvalidRequest)
	}
	return nil
}

func pairedRestartStatePath(updateStatePath string) string {
	return filepath.Join(filepath.Dir(updateStatePath), "restart-state.json")
}

func prepareLifecycleMutation(restartStatePath string, removeFile func(string) error) error {
	state, err := readState(restartStatePath)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("restart recovery state could not be validated: %w", err)
	}
	if err := validateRestartRecoveryState(state); err != nil {
		return err
	}
	if state.Phase != PhaseVerified || !state.MutationStarted {
		return ErrInterruptedRestart
	}
	if err := removeFile(lifecycleOverridePath(restartStatePath, state.Transaction)); err != nil {
		return recoveryRequired("verified restart override could not be cleaned up", err)
	}
	if err := removeFile(restartStatePath); err != nil {
		return recoveryRequired("verified restart state could not be cleaned up", err)
	}
	return nil
}