//go:build linux

package cgroup

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"golang.org/x/sys/unix"
)

// A cgroup that took cgroup.kill is never spawned into again (amadan #272).
// Found on the dev box (arm64, kernel 6.12.105, systemd 257): once
// cgroup.kill has been written to a cgroup, every later
// clone3(CLONE_INTO_CGROUP) into that same cgroup is SIGKILLed at once,
// while a cgroup made fresh at the same path is not. The Manager's starts
// now place through the exec-shim, a cgroup.procs write rather than a
// clone3, but the strict claims path still clones, and a fresh cgroup is
// also the one place a kept cgroup's limits are written again, so every
// start still makes a killed cgroup fresh first. Generation kills
// (KillGen, CloseRemove) remove their child and never meet it. The kills
// that keep their cgroup are a worker's host cgroup (KillHost) and the
// legacy sidecar's role child (KillSidecar); a strict group's own kill
// keeps its directory too.
//
// Every kill marks the cgroup it is about to kill, by identity (device and
// inode), before it writes cgroup.kill. The next start into that path
// (Start, StartChild, PrepareStrict) removes the killed cgroup, and any
// empty cgroups left below it, and then creates it as a first start does,
// so it gets the same limits and controller settings from the same code,
// under the same parent. The recreation happens at the start, not right
// after the kill, for three reasons. The start is the one place that knows
// what the cgroup must carry (its tier, or the parent's), so a recreation
// cannot drift from a creation. One check there covers every kill, a strict
// one included, without each kill site having to remember it. And a strict
// handle may still pin the killed directory when the kill returns (it is the
// handle's identity), so removing it then would strand the handle; at the
// start a pinned cgroup is refused instead, and removed once the handle lets
// go.
//
// A removal that fails (EBUSY: something is in it again) refuses the start
// with ErrKilledCgroupKept. The mark stays, so every later start refuses too
// until the cgroup can be removed. Marks are by identity, so a cgroup made
// fresh at the same path is never mistaken for the killed one, and a
// cgroup someone else removed leaves only a harmless stale mark (cgroupfs
// inode numbers are not reused).
//
// The marks live in this process only. An agent that restarts with Setup
// sweeps the empty carlos-exec@ directories a previous run left, killed
// ones included, so it starts with fresh cgroups.

// killedCgroups is every cgroup this process wrote cgroup.kill to that may
// still exist, by identity. Package-wide, not per Manager: a strict handle's
// kill has no Manager, and an identity names one cgroup on the box.
var killedCgroups = struct {
	sync.Mutex
	ids map[strictIdentity]bool
}{ids: make(map[strictIdentity]bool)}

// markKilled records that id is about to take cgroup.kill. It is called
// before the write, so a kill that fails half way still counts.
func markKilled(id strictIdentity) {
	killedCgroups.Lock()
	killedCgroups.ids[id] = true
	killedCgroups.Unlock()
}

func forgetKilled(id strictIdentity) {
	killedCgroups.Lock()
	delete(killedCgroups.ids, id)
	killedCgroups.Unlock()
}

func wasKilled(id strictIdentity) bool {
	killedCgroups.Lock()
	defer killedCgroups.Unlock()
	return killedCgroups.ids[id]
}

// anyKilled lets a start skip every stat while nothing was ever killed: on
// a box that never kills a kept cgroup (the flag off, before Task 15's F6
// path) the start path is the one it always was.
func anyKilled() bool {
	killedCgroups.Lock()
	defer killedCgroups.Unlock()
	return len(killedCgroups.ids) > 0
}

func reservedStrict(id strictIdentity) bool {
	strictReservations.Lock()
	defer strictReservations.Unlock()
	return strictReservations.held[id]
}

// rmdirCgroup removes one cgroup directory. A seam so the temp-dir fixture,
// whose interface files are ordinary files, can stand in for cgroupfs.
var rmdirCgroup = func(path string) error { return unix.Rmdir(path) }

// pathIdentity is the identity of the directory at path, not following a
// symlink. ok is false when nothing is there.
func pathIdentity(path string) (id strictIdentity, ok bool, err error) {
	var st unix.Stat_t
	if err := unix.Lstat(path, &st); err != nil {
		if errors.Is(err, unix.ENOENT) {
			return strictIdentity{}, false, nil
		}
		return strictIdentity{}, false, err
	}
	return strictIdentity{uint64(st.Dev), st.Ino}, true, nil
}

// freshen makes sure no start into dir (the host's cgroup, or one of its
// children) lands in a cgroup that took a kill: the host's cgroup first,
// since a kill there killed everything below it, then dir. It runs under
// the host's start lock (hostLock), so two starts of one host never both
// remove: the first removes the killed cgroup, and the second finds the
// fresh one, whose identity took no kill, and keeps it. The lock is held
// only for the check and the removal, not across the start's placement: a
// start's own directory is never removed by another start, because only a
// killed identity is removed and every start checks under this lock first.
// What it does not order is a kill landing while a start is between its
// check and its placement; the activator never does that (a kept cgroup is
// killed only once its leader has exited, and nothing of that host or role
// is started until the proof returns).
func (m *Manager) freshen(host, dir string) error {
	if !anyKilled() {
		return nil
	}
	mu := m.hostLock(host)
	mu.Lock()
	defer mu.Unlock()
	hostDir := m.hostDir(host)
	if err := m.freshenOne(host, hostDir); err != nil {
		return err
	}
	if dir == hostDir {
		return nil
	}
	return m.freshenOne(host, dir)
}

func (m *Manager) freshenOne(host, dir string) error {
	id, ok, err := pathIdentity(dir)
	if err != nil {
		// Unknown is not fresh.
		return fmt.Errorf("%w: %s: %w", ErrKilledCgroupKept, dir, err)
	}
	if !ok || !wasKilled(id) {
		return nil
	}
	groups := 0
	if err := removeKilledTree(dir, 0, &groups); err != nil {
		m.logf("cgroup: %s: %s took cgroup.kill and could not be removed, so nothing is started in it: %v", host, filepath.Base(dir), err)
		return fmt.Errorf("%w: %s: %w", ErrKilledCgroupKept, dir, err)
	}
	m.logf("cgroup: %s: %s took cgroup.kill; removed so the next spawn gets a fresh cgroup", host, filepath.Base(dir))
	return nil
}

// removeKilledTree removes dir and every cgroup below it, deepest first:
// rmdir refuses a cgroup with child cgroups, empty or not, and a tenant
// running as the agent's uid may have made some. The kernel refuses to
// remove a cgroup with a process in it, so a removal that succeeds also
// proves it empty. A cgroup a strict handle still pins is refused rather
// than removed under it, and each directory is checked before anything
// below it is touched, so a pinned killed parent loses none of its
// children. The walk is bounded as the cleanup walk is.
func removeKilledTree(dir string, depth int, groups *int) error {
	id, ok, err := pathIdentity(dir)
	if err != nil || !ok {
		return err
	}
	if reservedStrict(id) {
		return fmt.Errorf("%w: a strict handle still holds %s", ErrStrictState, filepath.Base(dir))
	}
	entries, err := os.ReadDir(dir)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return err
	}
	for _, e := range entries {
		if !e.IsDir() {
			continue
		}
		if depth+1 > maxCleanupDepth {
			return fmt.Errorf("%w: cgroup tree deeper than %d", ErrStrictRefused, maxCleanupDepth)
		}
		if *groups++; *groups > maxCleanupGroups {
			return fmt.Errorf("%w: more than %d cgroups below one host", ErrStrictRefused, maxCleanupGroups)
		}
		if err := removeKilledTree(filepath.Join(dir, e.Name()), depth+1, groups); err != nil {
			return err
		}
	}
	if err := rmdirCgroup(dir); err != nil && !errors.Is(err, unix.ENOENT) {
		return err
	}
	forgetKilled(id)
	return nil
}

// removeEmptyDir is CleanupHost's and RemoveChild's os.Remove, which also
// drops a kill mark the removed cgroup carried, so the set does not grow.
func removeEmptyDir(path string) error {
	id, ok, _ := pathIdentity(path)
	err := os.Remove(path)
	if err == nil && ok {
		forgetKilled(id)
	}
	return err
}

// refuseKilledChild is pinChild's check on a child that already exists: a
// strict child that took a kill is never pinned again for a spawn. Generation
// numbers are never reused, and no production path re-pins a role child, so
// a refusal costs nothing there; it keeps the rule whole.
func refuseKilledChild(parent *Prepared, name string) error {
	var st unix.Stat_t
	if err := unix.Fstatat(int(parent.dir.Fd()), name, &st, unix.AT_SYMLINK_NOFOLLOW); err != nil {
		return ErrStrictRefused
	}
	if wasKilled(strictIdentity{uint64(st.Dev), st.Ino}) {
		return ErrStrictRefused
	}
	return nil
}

// seedKilled is a retaining setup's stand-in for the sweep it skips
// (Task 19 fix round 1). The marks are this process's own, so a host
// cgroup or a role child that an earlier run killed and kept is unknown to
// it, and claim-entry mode (SetupRetaining) does not sweep the empty
// carlos-exec@ directories as Setup does. Every one that holds no process
// is marked instead, and each of its child directories that holds none, so
// the first start into it makes it fresh. Being marked costs an empty
// cgroup only a recreation. A directory that holds a process, or whose
// population cannot be read, is never marked: removing it is not ours to
// do. The claim containment groups (IsClaimContainment: carlos-exec@claim-<id>.local, unique
// to one controller) are left alone: they are the retained evidence the
// retaining mode exists to keep, and nothing starts into them again.
func (m *Manager) seedKilled() {
	entries, err := os.ReadDir(m.root)
	if err != nil {
		return
	}
	seeded := 0
	for _, e := range entries {
		if !e.IsDir() || !IsExecName(e.Name()) || IsClaimContainment(e.Name()) {
			continue
		}
		dir := filepath.Join(m.root, e.Name())
		if kids, err := os.ReadDir(dir); err == nil {
			for _, k := range kids {
				if k.IsDir() && seedIfEmpty(filepath.Join(dir, k.Name())) {
					seeded++
				}
			}
		}
		if seedIfEmpty(dir) {
			seeded++
		}
	}
	if seeded > 0 {
		m.logf("cgroup: %d retained empty cgroup(s) will be made fresh before their next spawn", seeded)
	}
}

// seedIfEmpty marks dir when its cgroup.events reads "populated 0" (no
// process in it or below it) and its own cgroup.procs, when present, lists
// none. It reports whether it marked.
func seedIfEmpty(dir string) bool {
	b, err := os.ReadFile(filepath.Join(dir, "cgroup.events"))
	if err != nil {
		return false
	}
	populated := true
	for _, line := range strings.Split(string(b), "\n") {
		if f := strings.Fields(line); len(f) == 2 && f[0] == "populated" {
			populated = f[1] != "0"
		}
	}
	if populated {
		return false
	}
	if procs, err := os.ReadFile(filepath.Join(dir, "cgroup.procs")); err == nil && len(strings.Fields(string(procs))) > 0 {
		return false
	} else if err != nil && !errors.Is(err, os.ErrNotExist) {
		return false
	}
	id, ok, err := pathIdentity(dir)
	if err != nil || !ok {
		return false
	}
	markKilled(id)
	return true
}