makarna/pkg/backend/cpu/affinity_linux.go

//go:build linux

package cpu

import (
	"os"
	"runtime"
	"sort"
	"strconv"
	"strings"

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

var (
	affinityEnabled bool
	affinitySet     unix.CPUSet
)

func init() {
	// Optional: enable via env to keep default behavior unchanged.
	//
	// - MAKARNA_CPUSET: e.g. "0-15,32-47"
	// - MAKARNA_NUMA_NODE: e.g. "0" (reads /sys/devices/system/node/node0/cpulist)
	if err := initAffinityFromEnv(); err != nil {
		affinityEnabled = false
	}
}

func initAffinityFromEnv() error {
	if cpus, ok := cpusFromEnv("MAKARNA_CPUSET"); ok {
		return setAffinityCPUs(cpus)
	}
	if nodeStr := strings.TrimSpace(os.Getenv("MAKARNA_NUMA_NODE")); nodeStr != "" {
		node, err := strconv.Atoi(nodeStr)
		if err != nil || node < 0 {
			return err
		}
		data, err := os.ReadFile("/sys/devices/system/node/node" + strconv.Itoa(node) + "/cpulist")
		if err != nil {
			return err
		}
		cpus, err := parseCPUList(strings.TrimSpace(string(data)))
		if err != nil {
			return err
		}
		return setAffinityCPUs(cpus)
	}
	return nil
}

func cpusFromEnv(key string) ([]int, bool) {
	raw := strings.TrimSpace(os.Getenv(key))
	if raw == "" {
		return nil, false
	}
	cpus, err := parseCPUList(raw)
	if err != nil {
		return nil, false
	}
	return cpus, true
}

func parseCPUList(s string) ([]int, error) {
	if s == "" {
		return nil, strconv.ErrSyntax
	}
	var cpus []int
	for _, part := range strings.Split(s, ",") {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}
		if lo, hi, ok := strings.Cut(part, "-"); ok {
			start, err := strconv.Atoi(strings.TrimSpace(lo))
			if err != nil {
				return nil, err
			}
			end, err := strconv.Atoi(strings.TrimSpace(hi))
			if err != nil {
				return nil, err
			}
			if start < 0 || end < start {
				return nil, strconv.ErrSyntax
			}
			for c := start; c <= end; c++ {
				cpus = append(cpus, c)
			}
			continue
		}
		v, err := strconv.Atoi(part)
		if err != nil {
			return nil, err
		}
		if v < 0 {
			return nil, strconv.ErrSyntax
		}
		cpus = append(cpus, v)
	}
	if len(cpus) == 0 {
		return nil, strconv.ErrSyntax
	}
	sort.Ints(cpus)
	cpus = compactInts(cpus)
	return cpus, nil
}

func compactInts(xs []int) []int {
	if len(xs) < 2 {
		return xs
	}
	out := xs[:1]
	for _, v := range xs[1:] {
		if v == out[len(out)-1] {
			continue
		}
		out = append(out, v)
	}
	return out
}

func setAffinityCPUs(cpus []int) error {
	var set unix.CPUSet
	set.Zero()
	for _, c := range cpus {
		set.Set(c)
	}
	affinitySet = set
	affinityEnabled = true
	return nil
}

// WithPinnedThread runs fn on a locked OS thread with optional CPU affinity set
// (enabled via env). If affinity is not configured, it behaves like fn().
func WithPinnedThread(fn func()) {
	if !affinityEnabled {
		fn()
		return
	}
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	_ = unix.SchedSetaffinity(0, &affinitySet)
	fn()
}