core/auth/ui/e2e/watchdog.go

package e2e

import (
	"fmt"
	"os"
	"sort"
	"strings"
	"sync"
	"time"
)

const (
	// The harness deadline includes setup and leaves one minute before Make's
	// exact 10-minute Go timeout for cleanup and the emergency exit bound.
	defaultSuiteWatchdogTimeout = 9 * time.Minute
	watchdogCleanupExitTimeout  = 45 * time.Second
	emergencyCleanupExitGrace   = 5 * time.Second
	suiteWatchdogTimeoutEnv     = "AUTH_UI_E2E_WATCHDOG_TIMEOUT"
)

type sessionClaim uint8

const (
	sessionUnclaimed sessionClaim = iota
	sessionNormal
	sessionTimeout
)

type sessionEntry struct {
	claim sessionClaim
	done  chan struct{}
	err   error
}

type sessionRegistry struct {
	mu       sync.Mutex
	sessions map[*browserSession]*sessionEntry
	timedOut bool
}

func newSessionRegistry() *sessionRegistry {
	return &sessionRegistry{sessions: make(map[*browserSession]*sessionEntry)}
}

func (r *sessionRegistry) add(session *browserSession) bool {
	if r == nil || session == nil {
		return false
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.timedOut {
		return false
	}
	r.sessions[session] = &sessionEntry{done: make(chan struct{})}
	return true
}

func (r *sessionRegistry) finalizeNormal(session *browserSession) error {
	if r == nil {
		return session.finalizeClaimed(false)
	}
	r.mu.Lock()
	entry := r.sessions[session]
	if entry == nil {
		r.mu.Unlock()
		return session.finalizeClaimed(false)
	}
	owner := entry.claim == sessionUnclaimed
	if owner {
		entry.claim = sessionNormal
	}
	r.mu.Unlock()
	if owner {
		err := session.finalizeClaimed(false)
		r.complete(session, entry, err)
		return err
	}
	<-entry.done
	return entry.err
}

func (r *sessionRegistry) complete(session *browserSession, entry *sessionEntry, err error) {
	r.mu.Lock()
	entry.err = err
	delete(r.sessions, session)
	close(entry.done)
	r.mu.Unlock()
}

func (r *sessionRegistry) finalizeAfterTimeout() error {
	if r == nil {
		return nil
	}
	r.mu.Lock()
	r.timedOut = true
	type claimedSession struct {
		session *browserSession
		entry   *sessionEntry
		owner   bool
	}
	claimed := make([]claimedSession, 0, len(r.sessions))
	for session, entry := range r.sessions {
		owner := entry.claim == sessionUnclaimed
		if owner {
			entry.claim = sessionTimeout
		}
		claimed = append(claimed, claimedSession{session: session, entry: entry, owner: owner})
	}
	r.mu.Unlock()
	sort.Slice(claimed, func(i, j int) bool { return claimed[i].session.dir < claimed[j].session.dir })
	var errs []string
	for _, claim := range claimed {
		if claim.owner {
			err := claim.session.finalizeClaimed(true)
			r.complete(claim.session, claim.entry, err)
		}
		<-claim.entry.done
		if claim.entry.err != nil {
			errs = append(errs, fmt.Sprintf("%s: %v", filepathBase(claim.session.dir), claim.entry.err))
		}
	}
	if len(errs) > 0 {
		return fmt.Errorf("finalize timed-out browser sessions: %s", strings.Join(errs, "; "))
	}
	return nil
}

// filepathBase avoids putting full retained paths into aggregate errors.
func filepathBase(path string) string {
	for len(path) > 1 && os.IsPathSeparator(path[len(path)-1]) {
		path = path[:len(path)-1]
	}
	if i := strings.LastIndexAny(path, `/\\`); i >= 0 {
		return path[i+1:]
	}
	return path
}

type ownedCloser interface{ Close() error }

type closerRegistry struct {
	mu      sync.Mutex
	closers map[ownedCloser]struct{}
	closing bool
}

func newCloserRegistry() *closerRegistry {
	return &closerRegistry{closers: make(map[ownedCloser]struct{})}
}

func (r *closerRegistry) add(closer ownedCloser) bool {
	if r == nil || closer == nil {
		return false
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.closing {
		return false
	}
	r.closers[closer] = struct{}{}
	return true
}

func (r *closerRegistry) remove(closer ownedCloser) {
	if r == nil || closer == nil {
		return
	}
	r.mu.Lock()
	delete(r.closers, closer)
	r.mu.Unlock()
}

func (r *closerRegistry) closeAll() error {
	if r == nil {
		return nil
	}
	r.mu.Lock()
	r.closing = true
	closers := make([]ownedCloser, 0, len(r.closers))
	for closer := range r.closers {
		closers = append(closers, closer)
		delete(r.closers, closer)
	}
	r.mu.Unlock()
	var errs []string
	for _, closer := range closers {
		if err := closer.Close(); err != nil {
			errs = append(errs, err.Error())
		}
	}
	if len(errs) > 0 {
		return fmt.Errorf("close owned listeners: %s", strings.Join(errs, "; "))
	}
	return nil
}

type terminalReason uint8

const (
	terminalRunning terminalReason = iota
	terminalNormal
	terminalTimeout
)

type suiteLifecycle struct {
	mu                 sync.Mutex
	reason             terminalReason
	timeoutExitClaimed bool
	stack              *Stack
	browser            *browserRuntime
	sessions           *sessionRegistry
	closers            *closerRegistry

	finishOnce sync.Once
	finishErr  error
}

// Installed before tagged-suite setup and retained until process exit.
var activeSuiteLifecycle *suiteLifecycle

func newSuiteLifecycle(stack *Stack, browser *browserRuntime) *suiteLifecycle {
	return &suiteLifecycle{stack: stack, browser: browser, sessions: newSessionRegistry(), closers: newCloserRegistry()}
}

func (l *suiteLifecycle) setStack(stack *Stack) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.reason != terminalRunning {
		return false
	}
	l.stack = stack
	return true
}

func (l *suiteLifecycle) setBrowser(browser *browserRuntime) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.reason != terminalRunning {
		return false
	}
	l.browser = browser
	return true
}

func (l *suiteLifecycle) claimTimeout() bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.reason != terminalRunning {
		return false
	}
	l.reason = terminalTimeout
	return true
}

// claimAfterRun stops and claims normal cleanup atomically. If the timer is
// already firing, this caller claims timeout completion so timeout status and
// artifacts cannot be lost while the callback is waiting to run.
func (l *suiteLifecycle) claimAfterRun(watchdog *suiteWatchdog) terminalReason {
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.reason != terminalRunning {
		return l.reason
	}
	if watchdog != nil && !watchdog.timer.Stop() {
		l.reason = terminalTimeout
	} else {
		l.reason = terminalNormal
	}
	return l.reason
}

func (l *suiteLifecycle) claimTimeoutExit() bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.reason != terminalTimeout || l.timeoutExitClaimed {
		return false
	}
	l.timeoutExitClaimed = true
	return true
}

func (l *suiteLifecycle) currentReason() terminalReason {
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.reason
}

func (l *suiteLifecycle) diagnosticPaths() (workspace, artifacts string) {
	l.mu.Lock()
	stack := l.stack
	l.mu.Unlock()
	if stack != nil {
		return stack.Workspace, stack.ArtifactDir
	}
	return "unavailable (setup incomplete)", "unavailable (setup incomplete)"
}

func (l *suiteLifecycle) finish(success bool) error {
	if l == nil {
		return nil
	}
	l.finishOnce.Do(func() {
		l.mu.Lock()
		reason, stack, browser := l.reason, l.stack, l.browser
		l.mu.Unlock()
		timedOut := reason == terminalTimeout
		var errs []string
		if timedOut && stack != nil {
			stack.recordRun("watchdog_timeout", fmt.Errorf("internal E2E suite deadline exceeded"))
		}
		if timedOut {
			if err := l.sessions.finalizeAfterTimeout(); err != nil {
				errs = append(errs, err.Error())
			}
		}
		if err := l.closers.closeAll(); err != nil {
			errs = append(errs, err.Error())
		}
		if stack != nil {
			if err := stack.StopServices(); err != nil {
				errs = append(errs, "services: "+err.Error())
			}
		}
		if browser != nil {
			if err := browser.Close(); err != nil {
				errs = append(errs, "browser: "+err.Error())
			}
		}
		finalSuccess := success && !timedOut && len(errs) == 0
		if stack != nil {
			if err := stack.Finalize(finalSuccess); err != nil {
				errs = append(errs, "artifacts: "+err.Error())
			}
		}
		if len(errs) > 0 {
			l.finishErr = fmt.Errorf("E2E suite cleanup: %s", strings.Join(errs, "; "))
		}
	})
	return l.finishErr
}

func (l *suiteLifecycle) emergencyTerminate() error {
	l.mu.Lock()
	stack := l.stack
	l.mu.Unlock()
	var groups []int
	if stack != nil {
		groups = stack.emergencyProcessGroups()
	}
	return emergencyTerminateOwnedProcesses(os.Getpid(), groups)
}

func suiteWatchdogTimeout() (time.Duration, error) {
	value := os.Getenv(suiteWatchdogTimeoutEnv)
	if value == "" {
		return defaultSuiteWatchdogTimeout, nil
	}
	duration, err := time.ParseDuration(value)
	if err != nil || duration <= 0 {
		return 0, fmt.Errorf("%s must be a positive Go duration", suiteWatchdogTimeoutEnv)
	}
	return duration, nil
}

type suiteWatchdog struct{ timer *time.Timer }

func startSuiteWatchdog(timeout time.Duration, fire func()) *suiteWatchdog {
	return &suiteWatchdog{timer: time.AfterFunc(timeout, fire)}
}

func (w *suiteWatchdog) stop() bool { return w == nil || w.timer.Stop() }

func stopCleanupTimerOrWait(stop func() bool, wait func()) {
	if !stop() {
		wait()
	}
}

func waitForEmergencyExit() { select {} }