//go:build linux

package cgroup

import (
	"context"
	"errors"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
	"time"

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

// amadan #272: once cgroup.kill has been written to a cgroup, every later
// clone3(CLONE_INTO_CGROUP) into that same cgroup is SIGKILLed at once. The
// fixture has no kernel behind it, so these tests pin the bookkeeping: a
// cgroup that took a kill is removed and made fresh (its sentinel gone, its limit
// files written again) before the next spawn into its path, and a removal
// that fails refuses the spawn.

// fakeCgroupRmdir stands in for rmdir on cgroupfs, where the interface files
// do not count against it: the fixture's are ordinary files, so they go
// first. fail, while non-nil, is returned instead. It records each removal.
func fakeCgroupRmdir(t *testing.T, fail *error) *[]string {
	t.Helper()
	var removed []string
	orig := rmdirCgroup
	rmdirCgroup = func(path string) error {
		if fail != nil && *fail != nil {
			return *fail
		}
		entries, err := os.ReadDir(path)
		if err != nil {
			return err
		}
		for _, e := range entries {
			if !e.IsDir() {
				if err := os.Remove(filepath.Join(path, e.Name())); err != nil {
					return err
				}
			}
		}
		if err := unix.Rmdir(path); err != nil {
			return err
		}
		removed = append(removed, path)
		return nil
	}
	t.Cleanup(func() { rmdirCgroup = orig })
	return &removed
}

// forgetKillsAfter clears the kill marks a test leaves, so a later test's
// temp dir can never meet one.
func forgetKillsAfter(t *testing.T) {
	t.Helper()
	t.Cleanup(forgetKills)
}

func init() {
	forgetKills = func() {
		killedCgroups.Lock()
		clear(killedCgroups.ids)
		killedCgroups.Unlock()
	}
}

// sentinelName is a file production code never creates in a cgroup. A test
// plants it in a directory to learn later whether that very directory was
// removed and made again (the sentinel is gone) or left alone (it is still
// there). Freshness is proved this way, not by a changed inode number:
// ext4 and xfs hand a directory's inode straight back after rmdir+mkdir,
// so an inode comparison passes or fails by filesystem. Real cgroupfs does
// not reuse cgroup inode ids, which is why the production marks can key on
// identity.
const sentinelName = "test-sentinel"

func plantSentinel(t *testing.T, dir string) {
	t.Helper()
	writeFixture(t, filepath.Join(dir, sentinelName), "planted")
}

func hasSentinel(t *testing.T, dir string) bool {
	t.Helper()
	_, err := os.Lstat(filepath.Join(dir, sentinelName))
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("stat sentinel in %s: %v", dir, err)
	}
	return err == nil
}

// killFixture gives dir the files a cleanup handle reads, empty, and writes
// cgroup.kill through a cleanup handle on host (child "" is the host's own
// cgroup), as KillHost and KillSidecar do.
func killFixture(t *testing.T, m *Manager, host, child, dir string) {
	t.Helper()
	for name, body := range map[string]string{"cgroup.events": "populated 0\n", "cgroup.procs": "", "cgroup.kill": ""} {
		writeFixture(t, filepath.Join(dir, name), body)
	}
	c, err := m.openCleanupAt(host, child, func(*os.File) error { return nil })
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(t.Context(), time.Second)
	defer cancel()
	if err := c.KillAndEmpty(ctx); err != nil {
		t.Fatal(err)
	}
	if err := c.Close(); err != nil {
		t.Fatal(err)
	}
	if got := readFixture(t, filepath.Join(dir, "cgroup.kill")); got != "1" {
		t.Fatalf("cgroup.kill = %q, want 1", got)
	}
}

// A worker's host cgroup (KillHost's) is killed and kept. The next start
// into it removes it and creates it again, with its limit files written as
// a first creation writes them, and a start that follows no kill removes
// nothing.
func TestAKilledHostCgroupIsRecreatedBeforeTheNextStart(t *testing.T) {
	forgetKillsAfter(t)
	quietOOM(t)
	m, err := SetupAt(newDelegatedRoot(t), t.Logf)
	if err != nil {
		t.Fatal(err)
	}
	removed := fakeCgroupRmdir(t, nil)
	calls := fakeStart(t, func(*exec.Cmd) error { return nil })
	const host = "worker.example.com"
	dir := m.hostDir(host)
	for range 2 {
		if _, err := m.Start(host, TierDefault, buildTrue); err != nil {
			t.Fatal(err)
		}
	}
	if len(*removed) != 0 {
		t.Fatalf("starts that followed no kill removed %v", *removed)
	}
	killFixture(t, m, host, "", dir)
	plantSentinel(t, dir)
	// A limit file left behind would say nothing about the recreation;
	// scribble on it so only a fresh write can restore it.
	writeFixture(t, filepath.Join(dir, "memory.max"), "max")
	if _, err := m.Start(host, TierDefault, buildTrue); err != nil {
		t.Fatal(err)
	}
	if len(*calls) != 3 {
		t.Fatalf("started %d times, want 3", len(*calls))
	}
	if len(*removed) != 1 || (*removed)[0] != dir {
		t.Fatalf("removed %v, want only %s", *removed, dir)
	}
	if hasSentinel(t, dir) {
		t.Fatal("the respawn reused the killed cgroup (the sentinel planted in it survived)")
	}
	for file, want := range map[string]string{
		"memory.max":       "536870912",
		"memory.high":      "483183819", // HighOf: max/10*9
		"memory.swap.max":  "536870912",
		"memory.oom.group": "1",
	} {
		if got := readFixture(t, filepath.Join(dir, file)); got != want {
			t.Errorf("%s = %q, want %q", file, got, want)
		}
	}
	// The fresh cgroup took no kill: the next start keeps it.
	if _, err := m.Start(host, TierDefault, buildTrue); err != nil {
		t.Fatal(err)
	}
	if len(*removed) != 1 {
		t.Fatalf("a start after the recreation removed again: %v", *removed)
	}
}

// A killed cgroup that will not go (EBUSY: something is in it again) is not
// proven fresh, so the start refuses and execs nothing. Once it can be
// removed, the next start proceeds into a fresh one.
func TestAKilledCgroupThatCannotBeRemovedRefusesTheNextStart(t *testing.T) {
	forgetKillsAfter(t)
	quietOOM(t)
	m, err := SetupAt(newDelegatedRoot(t), t.Logf)
	if err != nil {
		t.Fatal(err)
	}
	fail := error(unix.EBUSY)
	removed := fakeCgroupRmdir(t, &fail)
	calls := fakeStart(t, func(*exec.Cmd) error { return nil })
	const host = "worker.example.com"
	dir := m.hostDir(host)
	if _, err := m.Start(host, TierDefault, buildTrue); err != nil {
		t.Fatal(err)
	}
	killFixture(t, m, host, "", dir)
	plantSentinel(t, dir)
	built := 0
	cmd, err := m.Start(host, TierDefault, func() *exec.Cmd { built++; return buildTrue() })
	if !errors.Is(err, ErrKilledCgroupKept) || !errors.Is(err, unix.EBUSY) || cmd != nil {
		t.Fatalf("start into a killed cgroup that would not go = %v, %v; want ErrKilledCgroupKept (EBUSY)", cmd, err)
	}
	if built != 0 || len(*calls) != 1 {
		t.Fatalf("the refused start built %d command(s) and started %d time(s) in all, want 0 and 1", built, len(*calls))
	}
	if !hasSentinel(t, dir) {
		t.Fatal("the refused start replaced the cgroup")
	}
	// Still refused while it will not go: the mark outlives the failure.
	if _, err := m.Start(host, TierDefault, buildTrue); !errors.Is(err, ErrKilledCgroupKept) {
		t.Fatalf("second start = %v, want ErrKilledCgroupKept", err)
	}
	fail = nil
	if _, err := m.Start(host, TierDefault, buildTrue); err != nil {
		t.Fatal(err)
	}
	if len(*removed) != 1 || hasSentinel(t, dir) {
		t.Fatalf("removed %v; the start after the removal succeeded must be in a fresh cgroup", *removed)
	}
}

// The legacy sidecar's role child (KillSidecar's) is killed and kept. The
// respawn gets a fresh role-sidecar under the very same parent, whose limit
// and accounting are left as they were.
func TestAKilledRoleSidecarIsRecreatedUnderTheSameParent(t *testing.T) {
	forgetKillsAfter(t)
	const host = "sc.example.com"
	m, parent := genManager(t, host)
	removed := fakeCgroupRmdir(t, nil)
	fakeStart(t, func(*exec.Cmd) error { return nil })
	child := RoleChild("sidecar")
	dir := filepath.Join(parent, child)
	if _, err := m.StartChild(host, child, TierDefault, Limit{}, buildTrue); err != nil {
		t.Fatal(err)
	}
	killFixture(t, m, host, child, dir)
	plantSentinel(t, dir)
	plantSentinel(t, parent)
	writeFixture(t, filepath.Join(parent, "memory.max"), "1073741824") // a raise the respawn must not undo
	writeFixture(t, filepath.Join(parent, "cgroup.events"), "populated 1\n")
	if _, err := m.StartChild(host, child, TierDefault, Limit{}, buildTrue); err != nil {
		t.Fatal(err)
	}
	if len(*removed) != 1 || (*removed)[0] != dir {
		t.Fatalf("removed %v, want only %s", *removed, dir)
	}
	if hasSentinel(t, dir) {
		t.Fatal("the sidecar respawn reused the killed role child")
	}
	if !hasSentinel(t, parent) {
		t.Fatal("the sidecar respawn replaced the parent")
	}
	if got := readFixture(t, filepath.Join(parent, "memory.max")); got != "1073741824" {
		t.Fatalf("parent memory.max = %q, the respawn touched it", got)
	}
	if _, err := os.Stat(filepath.Join(dir, "memory.max")); !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("a role child carries no limit of its own: %v", err)
	}
}

// KillHost on a generation parent kills everything under it. The next start
// of any child removes the whole killed tree, children first, and builds
// the parent again as a cold start does: controllers delegated, the tier
// written with memory.oom.group=0.
func TestAKilledParentIsRecreatedWithItsChildren(t *testing.T) {
	forgetKillsAfter(t)
	const host = "gp.example.com"
	m, parent := genManager(t, host)
	removed := fakeCgroupRmdir(t, nil)
	fakeStart(t, func(*exec.Cmd) error { return nil })
	sidecar := RoleChild("sidecar")
	if _, err := m.StartChild(host, sidecar, TierDefault, Limit{}, buildTrue); err != nil {
		t.Fatal(err)
	}
	if _, err := m.StartChild(host, "g4", TierDefault, TierDefault, buildTrue); err != nil {
		t.Fatal(err)
	}
	killFixture(t, m, host, "", parent)
	plantSentinel(t, parent)
	plantSentinel(t, filepath.Join(parent, sidecar))
	plantSentinel(t, filepath.Join(parent, "g4"))
	// The fixture's parent would lose its controller list with it; the
	// kernel lists them on every new cgroup, so the fake rmdir's recreation
	// needs it back.
	orig := rmdirCgroup
	rmdirCgroup = func(path string) error {
		if err := orig(path); err != nil {
			return err
		}
		if path == parent {
			if err := os.Mkdir(parent, 0o755); err != nil {
				return err
			}
			writeFixture(t, filepath.Join(parent, "cgroup.controllers"), "cpu io memory pids\n")
		}
		return nil
	}
	t.Cleanup(func() { rmdirCgroup = orig })
	if _, err := m.StartChild(host, sidecar, TierDefault, Limit{}, buildTrue); err != nil {
		t.Fatal(err)
	}
	want := map[string]bool{filepath.Join(parent, sidecar): true, filepath.Join(parent, "g4"): true}
	if len(*removed) != 3 || !want[(*removed)[0]] || !want[(*removed)[1]] || (*removed)[2] != parent {
		t.Fatalf("removed %v, want both children then %s", *removed, parent)
	}
	if hasSentinel(t, parent) {
		t.Fatal("the respawn reused the killed parent")
	}
	if _, err := os.Stat(filepath.Join(parent, "g4")); !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("the killed generation child survived: %v", err)
	}
	for file, want := range map[string]string{
		"cgroup.subtree_control": "+cpu +memory +io +pids",
		"memory.max":             "536870912",
		"memory.oom.group":       "0",
	} {
		if got := readFixture(t, filepath.Join(parent, file)); got != want {
			t.Errorf("parent %s = %q, want %q", file, got, want)
		}
	}
}

// A killed cgroup that a strict handle still pins (its reservation is
// held) is not removed under it: removing it would strand the handle, whose
// identity is that directory. The start refuses until the handle lets go.
func TestAKilledCgroupAStrictHandleHoldsIsNotRemoved(t *testing.T) {
	forgetKillsAfter(t)
	quietOOM(t)
	m, err := SetupAt(newDelegatedRoot(t), t.Logf)
	if err != nil {
		t.Fatal(err)
	}
	removed := fakeCgroupRmdir(t, nil)
	fakeStart(t, func(*exec.Cmd) error { return nil })
	const host = "pinned.example.com"
	dir := m.hostDir(host)
	if _, err := m.Start(host, TierDefault, buildTrue); err != nil {
		t.Fatal(err)
	}
	killFixture(t, m, host, "", dir)
	var st unix.Stat_t
	if err := unix.Lstat(dir, &st); err != nil {
		t.Fatal(err)
	}
	id := strictIdentity{uint64(st.Dev), st.Ino}
	strictReservations.Lock()
	strictReservations.held[id] = true
	strictReservations.Unlock()
	release := func() {
		strictReservations.Lock()
		delete(strictReservations.held, id)
		strictReservations.Unlock()
	}
	t.Cleanup(release)
	if _, err := m.Start(host, TierDefault, buildTrue); !errors.Is(err, ErrStrictState) || !errors.Is(err, ErrKilledCgroupKept) {
		t.Fatalf("start into a pinned killed cgroup = %v, want ErrKilledCgroupKept and ErrStrictState", err)
	}
	if len(*removed) != 0 {
		t.Fatalf("removed %v under a strict handle", *removed)
	}
	release()
	if _, err := m.Start(host, TierDefault, buildTrue); err != nil {
		t.Fatal(err)
	}
	if len(*removed) != 1 {
		t.Fatalf("removed %v once the handle let go, want the killed cgroup", *removed)
	}
}

// A strict role child that took a kill is never pinned again for a spawn.
func TestAKilledStrictRoleChildIsNotPinnedAgain(t *testing.T) {
	forgetKillsAfter(t)
	p := strictFixture(t)
	makeRoleFixture(t, p, "seed", "populated 0\n")
	child, err := p.PrepareRoleChild("seed")
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(t.Context(), time.Second)
	defer cancel()
	if err := child.KillAndEmpty(ctx); err != nil {
		t.Fatal(err)
	}
	if err := child.Close(); err != nil {
		t.Fatal(err)
	}
	again, err := p.PrepareRoleChild("seed")
	if err == nil {
		again.Close()
		t.Fatal("a killed role child was pinned again for a spawn")
	}
	if !errors.Is(err, ErrStrictRefused) {
		t.Fatalf("err = %v, want ErrStrictRefused", err)
	}
}

// A killed parent that a strict handle still pins keeps everything below it
// too: the reservation is checked before the walk touches a child (Task 19
// fix round 1).
func TestAPinnedKilledParentLosesNoChild(t *testing.T) {
	forgetKillsAfter(t)
	const host = "pinnedparent.example.com"
	m, parent := genManager(t, host)
	removed := fakeCgroupRmdir(t, nil)
	fakeStart(t, func(*exec.Cmd) error { return nil })
	sidecar := RoleChild("sidecar")
	if _, err := m.StartChild(host, sidecar, TierDefault, Limit{}, buildTrue); err != nil {
		t.Fatal(err)
	}
	killFixture(t, m, host, "", parent)
	var st unix.Stat_t
	if err := unix.Lstat(parent, &st); err != nil {
		t.Fatal(err)
	}
	id := strictIdentity{uint64(st.Dev), st.Ino}
	strictReservations.Lock()
	strictReservations.held[id] = true
	strictReservations.Unlock()
	t.Cleanup(func() {
		strictReservations.Lock()
		delete(strictReservations.held, id)
		strictReservations.Unlock()
	})
	if _, err := m.StartChild(host, sidecar, TierDefault, Limit{}, buildTrue); !errors.Is(err, ErrStrictState) {
		t.Fatalf("start under a pinned killed parent = %v, want ErrStrictState", err)
	}
	if len(*removed) != 0 {
		t.Fatalf("removed %v under a pinned parent", *removed)
	}
	if _, err := os.Stat(filepath.Join(parent, sidecar)); err != nil {
		t.Fatalf("the pinned parent's child went: %v", err)
	}
}

// retainedTree builds what an earlier run left under a delegated root:
// carlos-exec@<host> with cgroup.events and cgroup.procs, and children
// (name to events body) with the same files, empty.
func retainedTree(t *testing.T, root, host, events, procs string, kids map[string]string) string {
	t.Helper()
	dir := filepath.Join(root, ExecName(host))
	if err := os.Mkdir(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	writeFixture(t, filepath.Join(dir, "cgroup.events"), events)
	writeFixture(t, filepath.Join(dir, "cgroup.procs"), procs)
	writeFixture(t, filepath.Join(dir, "cgroup.controllers"), "cpu io memory pids\n")
	for kid, body := range kids {
		k := filepath.Join(dir, kid)
		if err := os.Mkdir(k, 0o755); err != nil {
			t.Fatal(err)
		}
		writeFixture(t, filepath.Join(k, "cgroup.events"), body)
		writeFixture(t, filepath.Join(k, "cgroup.procs"), "")
	}
	return dir
}

// Claim-entry mode (SetupRetaining) skips the sweep, and the kill marks die
// with the process that made them, so a cgroup an earlier run killed and
// kept would be cloned into again. The retaining setup marks every empty
// exec cgroup (and every empty child of one) instead: the first start into
// each makes it fresh with its limits written again. One that holds a
// process, and a claim containment group, are left exactly as they are
// (Task 19 fix round 1).
func TestARetainingSetupMakesEmptyLeftoversFresh(t *testing.T) {
	forgetKillsAfter(t)
	quietOOM(t)
	root := newDelegatedRoot(t)
	empty := retainedTree(t, root, "worker.example.com", "populated 0\n", "", nil)
	busy := retainedTree(t, root, "busy.example.com", "populated 1\n", "4242\n", nil)
	// A generation parent whose live g3 survived, beside an empty sidecar
	// child a KillSidecar left.
	live := retainedTree(t, root, "live.example.com", "populated 1\n", "", map[string]string{"g3": "populated 1\n", "role-sidecar": "populated 0\n"})
	writeFixture(t, filepath.Join(live, "cgroup.subtree_control"), "cpu memory io pids\n")
	claim := retainedTree(t, root, "claim-0123abcd.local", "populated 0\n", "", nil)
	for _, d := range []string{empty, busy, live, filepath.Join(live, "g3"), filepath.Join(live, "role-sidecar")} {
		plantSentinel(t, d)
	}
	m, err := SetupAtRetaining(root, t.Logf)
	if err != nil {
		t.Fatal(err)
	}
	removed := fakeCgroupRmdir(t, nil)
	fakeStart(t, func(*exec.Cmd) error { return nil })
	FakePlacementForTest(t) // a temp dir is no cgroupfs, and a role child is never moved by PID

	if _, err := m.Start("worker.example.com", TierDefault, buildTrue); err != nil {
		t.Fatal(err)
	}
	if hasSentinel(t, empty) {
		t.Fatal("the first start reused an empty retained host cgroup")
	}
	if got := readFixture(t, filepath.Join(empty, "memory.max")); got != "536870912" {
		t.Fatalf("recreated memory.max = %q, want the tier", got)
	}

	if _, err := m.Start("busy.example.com", TierDefault, buildTrue); err != nil {
		t.Fatal(err)
	}
	if !hasSentinel(t, busy) {
		t.Fatal("a retained cgroup holding a process was replaced")
	}

	if _, err := m.StartChild("live.example.com", RoleChild("sidecar"), TierDefault, Limit{}, buildTrue); err != nil {
		t.Fatal(err)
	}
	if hasSentinel(t, filepath.Join(live, "role-sidecar")) {
		t.Fatal("the sidecar respawn reused an empty retained role child")
	}
	if !hasSentinel(t, live) || !hasSentinel(t, filepath.Join(live, "g3")) {
		t.Fatal("a populated parent or its live generation was replaced")
	}

	if id, ok, err := pathIdentity(claim); err != nil || !ok || wasKilled(id) {
		t.Fatalf("a claim containment group was marked (ok=%v err=%v)", ok, err)
	}
	want := map[string]bool{empty: true, filepath.Join(live, "role-sidecar"): true}
	if len(*removed) != 2 || !want[(*removed)[0]] || !want[(*removed)[1]] {
		t.Fatalf("removed %v, want only the empty host cgroup and the empty role child", *removed)
	}
}