ThothII/tools/thothctl/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 restarting again")
	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")
)

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}, ErrConfirmationRequired
	}
	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 state, err := readState(request.StatePath); err == nil && state.MutationStarted {
		return RestartResult{StatePath: request.StatePath}, errInterruptedRestart
	} else if err != nil && !errors.Is(err, os.ErrNotExist) {
		return RestartResult{StatePath: request.StatePath}, err
	}
	if err := setMaintenance(ctx, runner, true); 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 := waitForInactiveSessions(ctx, runner, request.Drain, hooks.sleep); err != nil {
		return RestartResult{StatePath: request.StatePath}, 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
	state = State{
		Transaction: lifecycleTransaction(request.StatePath),
		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
	}
	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}, ErrActiveSessions
	}
	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 := recreateCore(ctx, runner); err != nil {
		return RestartResult{StatePath: request.StatePath, Version: version}, recoveryRequired("Pi restart core recreation failed", err)
	}
	if err := ensureMaintenance(ctx, runner); 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, runner, 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(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 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 ErrConfirmationRequired
	}
	lock, err := acquireLock(restartStatePath)
	if err != nil {
		return err
	}
	defer lock.Release()

	restartState, restartErr := readState(restartStatePath)
	if restartErr == nil {
		if restartState.MutationStarted {
			if restartState.Target.Source != "restart" || restartState.Target.Version == "" {
				return errors.New("restart recovery state is incomplete")
			}
			if err := ensureMaintenance(ctx, runner); err != nil {
				return recoveryRequired("Pi restart maintenance recovery failed", err)
			}
			if err := verifyRestart(ctx, runner, restartState.Target.Version, restartState.Previous); err != nil {
				return recoveryRequired("Pi restart recovery verification failed", 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 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
}