bugabinga.net

package rollout

import (
	"context"
	"errors"
	"fmt"
	"time"

	"bugabinga.net/toad/internal/enroll"
	"bugabinga.net/toad/internal/store"
)

type ContainerState struct {
	Running     bool
	ImageDigest string
}

type UnitState struct {
	LoadState   string
	ActiveState string
	Result      string
	Restarts    int
}

// Host is the only path from a rollout to the machine.
type Host interface {
	Pull(ctx context.Context, reference string) (string, error)
	InstallUnit(fileName, content string) error
	DaemonReload(ctx context.Context) error
	Stop(ctx context.Context, unit string) error
	Start(ctx context.Context, unit string) error
	Container(ctx context.Context, name string) (ContainerState, error)
	RestartCount(ctx context.Context, unit string) (int, error)
	Unit(ctx context.Context, unit string) (UnitState, error)
	Probe(ctx context.Context, url string) error
}

type Engine struct {
	Store       store.Store
	Host        Host
	Poll        time.Duration
	StopTimeout time.Duration
	Now         func() time.Time
	Sleep       func(context.Context, time.Duration) error
}

var errProbation = errors.New("probation failed")

// Run drives one rollout to a terminal phase and always persists the outcome.
// Cancelling ctx aborts probation; recovery from that abort still runs to completion.
func (e Engine) Run(ctx context.Context, service enroll.Service, operation store.Operation) {
	err := e.rollout(ctx, service, &operation)
	if err == nil {
		return
	}
	e.log(operation, "rollout failed: "+err.Error())
	recovery := context.WithoutCancel(ctx)
	if operation.Rollback == nil {
		e.stopCandidate(recovery, service, &operation)
		e.finish(&operation, store.PhaseFailed, err.Error())
		return
	}
	e.rollback(recovery, service, &operation, err)
}

func (e Engine) rollout(ctx context.Context, service enroll.Service, operation *store.Operation) error {
	if !e.Now().Before(operation.Deadline) {
		return fmt.Errorf("%w: rollout deadline reached before activation", errProbation)
	}
	e.phase(operation, store.PhasePreparing, "")
	e.log(*operation, "pulling "+operation.Candidate.Reference)
	pulled, err := e.Host.Pull(ctx, operation.Candidate.Reference)
	if err != nil {
		return fmt.Errorf("pull: %w", err)
	}
	if pulled != operation.Candidate.Digest {
		return fmt.Errorf("pulled digest %s does not match requested %s", pulled, operation.Candidate.Digest)
	}

	e.phase(operation, store.PhaseActivating, "")
	if err := e.Host.InstallUnit(service.UnitFileName(), operation.Candidate.Unit); err != nil {
		return fmt.Errorf("install unit: %w", err)
	}
	if err := e.Host.DaemonReload(ctx); err != nil {
		return fmt.Errorf("daemon reload: %w", err)
	}
	if err := e.stopVerified(ctx, service, *operation); err != nil {
		e.restorePreviousUnit(ctx, service, *operation)
		return err
	}
	if err := e.Host.Start(ctx, service.UnitName()); err != nil {
		return fmt.Errorf("start: %w", err)
	}

	e.phase(operation, store.PhaseProbation, "")
	if err := e.probation(ctx, service, operation); err != nil {
		return err
	}

	e.phase(operation, store.PhasePromoting, "")
	return e.promote(operation)
}

// stopVerified refuses to continue while the singleton workload may still run.
func (e Engine) stopVerified(ctx context.Context, service enroll.Service, operation store.Operation) error {
	e.log(operation, "stopping "+service.UnitName())
	if err := e.Host.Stop(ctx, service.UnitName()); err != nil {
		return fmt.Errorf("stop: %w", err)
	}
	limit := e.Now().Add(e.StopTimeout)
	for {
		state, err := e.Host.Container(ctx, service.ContainerName)
		if err != nil {
			return fmt.Errorf("verify stop: %w", err)
		}
		if !state.Running {
			e.log(operation, "previous instance terminated")
			return nil
		}
		if !e.Now().Before(limit) {
			return errors.New("previous instance still running after stop timeout")
		}
		if err := e.Sleep(ctx, e.Poll); err != nil {
			return err
		}
	}
}

func (e Engine) probation(ctx context.Context, service enroll.Service, operation *store.Operation) error {
	baseline, err := e.Host.RestartCount(ctx, service.UnitName())
	if err != nil {
		return fmt.Errorf("restart count: %w", err)
	}
	var readySince time.Time
	stableFor := service.Probation.StableForDuration()
	for {
		now := e.Now()
		restarts, err := e.Host.RestartCount(ctx, service.UnitName())
		if err != nil {
			return fmt.Errorf("restart count: %w", err)
		}
		if restarts-baseline != operation.Restarts {
			operation.Restarts = restarts - baseline
			readySince = time.Time{}
			e.log(*operation, fmt.Sprintf("candidate restarted, restarts=%d", operation.Restarts))
			e.save(*operation)
		}
		if operation.Restarts > service.Probation.RestartBudget {
			return fmt.Errorf("%w: restart budget %d exceeded", errProbation, service.Probation.RestartBudget)
		}

		if e.ready(ctx, service, operation.Candidate.Digest) {
			if readySince.IsZero() {
				readySince = now
				e.log(*operation, "candidate ready")
			} else if now.Sub(readySince) >= stableFor {
				return nil
			}
		} else if !readySince.IsZero() {
			readySince = time.Time{}
			e.log(*operation, "candidate readiness lost")
		}

		if !now.Before(operation.Deadline) {
			return fmt.Errorf("%w: rollout deadline reached", errProbation)
		}
		if err := e.Sleep(ctx, e.Poll); err != nil {
			return fmt.Errorf("%w: probation interrupted: %v", errProbation, err)
		}
	}
}

// ready counts only observations attributable to the expected container and digest.
func (e Engine) ready(ctx context.Context, service enroll.Service, digest string) bool {
	state, err := e.Host.Container(ctx, service.ContainerName)
	if err != nil || !state.Running || state.ImageDigest != digest {
		return false
	}
	return e.Host.Probe(ctx, service.ReadinessURL) == nil
}

func (e Engine) promote(operation *store.Operation) error {
	state, err := e.Store.ReadState(operation.Service)
	if err != nil {
		return fmt.Errorf("read state: %w", err)
	}
	candidate := operation.Candidate
	state.Rollback = state.Current
	state.Current = &candidate
	state.Generation++
	state.OperationID = operation.ID
	if err := e.Store.WriteState(operation.Service, state); err != nil {
		return fmt.Errorf("write state: %w", err)
	}
	operation.Generation = state.Generation
	e.log(*operation, fmt.Sprintf("promoted generation %d", state.Generation))
	e.finish(operation, store.PhasePromoted, "")
	return nil
}

func (e Engine) rollback(ctx context.Context, service enroll.Service, operation *store.Operation, cause error) {
	e.phase(operation, store.PhaseRollingBack, cause.Error())
	target := *operation.Rollback
	e.log(*operation, "rolling back to "+target.Reference)

	if err := e.restore(ctx, service, *operation, target); err != nil {
		e.log(*operation, "rollback failed: "+err.Error())
		e.finish(operation, store.PhaseFailed, "rollback failed after "+cause.Error()+": "+err.Error())
		return
	}
	if err := e.advanceGeneration(operation); err != nil {
		e.log(*operation, "rollback bookkeeping failed: "+err.Error())
		e.finish(operation, store.PhaseFailed, err.Error())
		return
	}
	e.finish(operation, store.PhaseRolledBack, cause.Error())
}

func (e Engine) restore(ctx context.Context, service enroll.Service, operation store.Operation, target store.Revision) error {
	if err := e.Host.InstallUnit(service.UnitFileName(), target.Unit); err != nil {
		return err
	}
	if err := e.Host.DaemonReload(ctx); err != nil {
		return err
	}
	if err := e.stopVerified(ctx, service, operation); err != nil {
		return err
	}
	if err := e.Host.Start(ctx, service.UnitName()); err != nil {
		return err
	}
	limit := e.Now().Add(service.Probation.StableForDuration() + e.StopTimeout)
	for {
		if e.ready(ctx, service, target.Digest) {
			return nil
		}
		if !e.Now().Before(limit) {
			return errors.New("restored revision did not become ready")
		}
		if err := e.Sleep(ctx, e.Poll); err != nil {
			return err
		}
	}
}

func (e Engine) restorePreviousUnit(ctx context.Context, service enroll.Service, operation store.Operation) {
	if operation.Rollback == nil {
		return
	}
	if err := e.Host.InstallUnit(service.UnitFileName(), operation.Rollback.Unit); err != nil {
		e.log(operation, "restoring previous unit failed: "+err.Error())
		return
	}
	if err := e.Host.DaemonReload(ctx); err != nil {
		e.log(operation, "daemon reload after unit restore failed: "+err.Error())
	}
}

func (e Engine) stopCandidate(ctx context.Context, service enroll.Service, operation *store.Operation) {
	if err := e.Host.Stop(ctx, service.UnitName()); err != nil {
		e.log(*operation, "stopping failed candidate failed: "+err.Error())
	}
}

func (e Engine) advanceGeneration(operation *store.Operation) error {
	state, err := e.Store.ReadState(operation.Service)
	if err != nil {
		return fmt.Errorf("read state: %w", err)
	}
	state.Generation++
	state.OperationID = operation.ID
	if err := e.Store.WriteState(operation.Service, state); err != nil {
		return fmt.Errorf("write state: %w", err)
	}
	operation.Generation = state.Generation
	return nil
}

func (e Engine) phase(operation *store.Operation, phase store.Phase, reason string) {
	operation.Phase = phase
	if reason != "" {
		operation.Reason = reason
	}
	e.save(*operation)
	e.log(*operation, "phase "+string(phase))
}

func (e Engine) finish(operation *store.Operation, phase store.Phase, reason string) {
	ended := e.Now().UTC()
	operation.Phase = phase
	operation.Reason = reason
	operation.EndedAt = &ended
	e.save(*operation)
	e.log(*operation, "finished "+string(phase))
}

func (e Engine) save(operation store.Operation) {
	if err := e.Store.WriteOperation(operation); err != nil {
		_ = e.Store.AppendLog(operation.Service, operation.ID, "persisting operation failed: "+err.Error())
	}
}

func (e Engine) log(operation store.Operation, line string) {
	_ = e.Store.AppendLog(operation.Service, operation.ID, line)
}