package cgroup

import (
	"context"
	"path/filepath"
	"strings"
	"sync"
)

// The in-flight registry: target cgroup directories that a shim is being
// placed into, from just before its fork until it answers placed or the
// start fails. With clone3 a child appeared in its cgroup atomically; a
// shim spends a few milliseconds in the agent's cgroup first, so a proof
// that a cgroup is empty must count these too (Ruling 6, Task 2). The flat
// PID-move fallback (startAndMoveIn) registers its target as well, from
// its start until its PID write has returned. Readers: Cleanup's
// population and members (cleanup_linux.go) and ensureParent (gen.go).
//
// A registration ends only once its process is in the kernel's view of its
// cgroup (or never will be), so a placement is always visible to one of the
// two: the registry, or the kernel. But a reader that asks them one after
// the other can miss a placement that moves from one to the other between
// the two reads (whole-branch review, finding 1). Two tools close that:
//
//   - readQuiet brackets one kernel read: it answers quiet only when no
//     placement into the directory or below it was in flight at any moment
//     of that read (none when it began, none registered or ended during it,
//     none at its end). A reader that is not quiet treats the directory as
//     populated, or its member list as unproven. Cleanup's reads use it.
//   - holdPlacements stops any placement into the directory or below it
//     from registering until it is released, so a decision taken under it
//     (ensureParent's cold parent) holds until its writes are done. It is
//     held across cgroupfs reads and writes only, never across a before
//     callback, a shim wait or a start.
//
// No lock of the registry is ever held across a before callback, a shim
// wait or a start: placeMu guards only the maps, and a registration blocked
// by a hold waits on a channel with placeMu released.

// placementReadHook, when set, runs right after each kernel read that
// readQuiet brackets, and after ensureParent's population read. Only test
// files set it, to register or end a placement at exactly that point.
var placementReadHook func(dir string)

// placementHeldHook, when set, runs when a registration finds a hold over
// its target and is about to wait for it. Only test files set it.
var placementHeldHook func(dir string)

// placeWindow is one readQuiet in progress: busy means a placement under
// dir was in flight when it began; changed, that one registered or ended
// under dir since.
type placeWindow struct {
	dir     string
	busy    bool
	changed bool
}

// beginPlacement registers one placement into dir. done ends it, and is
// idempotent: a second call does nothing. A nil manager registers nothing.
// While a hold covers dir (holdPlacements), it waits for the hold to end.
func (m *Manager) beginPlacement(dir string) (done func()) {
	if m == nil {
		return func() {}
	}
	dir = filepath.Clean(dir)
	m.placeMu.Lock()
	for m.heldLocked(dir) {
		wake := m.wakeLocked()
		m.placeMu.Unlock()
		if placementHeldHook != nil {
			placementHeldHook(dir)
		}
		<-wake
		m.placeMu.Lock()
	}
	if m.placing == nil {
		m.placing = make(map[string]int)
	}
	m.placing[dir]++
	m.touchLocked(dir)
	m.placeMu.Unlock()
	var once sync.Once
	return func() {
		once.Do(func() {
			m.placeMu.Lock()
			defer m.placeMu.Unlock()
			if m.placing[dir]--; m.placing[dir] <= 0 {
				delete(m.placing, dir)
			}
			m.touchLocked(dir)
			// Wake every waiter; each rechecks its own subtree.
			m.broadcastLocked()
		})
	}
}

// readQuiet runs read, a read of the kernel's view of dir, and reports
// whether no placement into dir or below it was in flight at any moment
// while read ran. Only a quiet read covers every process bound for dir: a
// placement that ended before read began is in the kernel's view, and one
// that registers after read ended had not forked while it ran. A nil
// manager has no placements: always quiet.
func (m *Manager) readQuiet(dir string, read func()) (quiet bool) {
	dir = filepath.Clean(dir)
	if m == nil {
		read()
		return true
	}
	w := &placeWindow{dir: dir}
	m.placeMu.Lock()
	w.busy = m.countUnderLocked(dir) > 0
	if m.windows == nil {
		m.windows = make(map[*placeWindow]struct{})
	}
	m.windows[w] = struct{}{}
	m.placeMu.Unlock()
	read()
	if placementReadHook != nil {
		placementReadHook(dir)
	}
	m.placeMu.Lock()
	defer m.placeMu.Unlock()
	delete(m.windows, w)
	return !w.busy && !w.changed && m.countUnderLocked(dir) == 0
}

// holdPlacements keeps every placement into dir or below it from
// registering until release is called, and answers the placements already
// in flight: at, into dir itself, and under, into dir or below it. With
// the hold in place those counts can only fall, and a count of zero stays
// zero, so a kernel read taken under the hold, with both zero, sees every
// process bound for dir. release is idempotent. Hold it only across
// cgroupfs reads and writes.
func (m *Manager) holdPlacements(dir string) (at, under int, release func()) {
	if m == nil {
		return 0, 0, func() {}
	}
	dir = filepath.Clean(dir)
	m.placeMu.Lock()
	if m.holds == nil {
		m.holds = make(map[string]int)
	}
	m.holds[dir]++
	at, under = m.placing[dir], m.countUnderLocked(dir)
	m.placeMu.Unlock()
	var once sync.Once
	return at, under, func() {
		once.Do(func() {
			m.placeMu.Lock()
			defer m.placeMu.Unlock()
			if m.holds[dir]--; m.holds[dir] <= 0 {
				delete(m.holds, dir)
			}
			m.broadcastLocked()
		})
	}
}

// placementsUnder counts placements into dir or anywhere below it.
func (m *Manager) placementsUnder(dir string) int {
	if m == nil {
		return 0
	}
	m.placeMu.Lock()
	defer m.placeMu.Unlock()
	return m.countUnderLocked(filepath.Clean(dir))
}

// placementsAt counts placements into dir itself.
func (m *Manager) placementsAt(dir string) int {
	if m == nil {
		return 0
	}
	m.placeMu.Lock()
	defer m.placeMu.Unlock()
	return m.placing[filepath.Clean(dir)]
}

// waitPlacementsUnder returns nil once no placement is in flight in dir or
// below it, or ctx's error if ctx ends first.
func (m *Manager) waitPlacementsUnder(ctx context.Context, dir string) error {
	if m == nil {
		return nil
	}
	dir = filepath.Clean(dir)
	for {
		m.placeMu.Lock()
		if m.countUnderLocked(dir) == 0 {
			m.placeMu.Unlock()
			return nil
		}
		wake := m.wakeLocked()
		m.placeMu.Unlock()
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-wake:
		}
	}
}

// wakeLocked is the channel the next broadcastLocked closes.
func (m *Manager) wakeLocked() chan struct{} {
	if m.placeWake == nil {
		m.placeWake = make(chan struct{})
	}
	return m.placeWake
}

// broadcastLocked wakes every waiter: a registration ended, or a hold did.
func (m *Manager) broadcastLocked() {
	if m.placeWake != nil {
		close(m.placeWake)
		m.placeWake = nil
	}
}

// touchLocked marks every open readQuiet over target as changed.
func (m *Manager) touchLocked(target string) {
	for w := range m.windows {
		if isUnder(target, w.dir) {
			w.changed = true
		}
	}
}

// heldLocked reports whether a hold covers target.
func (m *Manager) heldLocked(target string) bool {
	for h := range m.holds {
		if isUnder(target, h) {
			return true
		}
	}
	return false
}

func (m *Manager) countUnderLocked(dir string) int {
	n := 0
	for target, c := range m.placing {
		if isUnder(target, dir) {
			n += c
		}
	}
	return n
}

// isUnder reports whether target is dir or below it. Both are clean.
func isUnder(target, dir string) bool {
	prefix := dir
	if !strings.HasSuffix(prefix, "/") {
		prefix += "/"
	}
	return target == dir || strings.HasPrefix(target, prefix)
}