falcon/main.go

229 lines
4.9 KiB
Go
Raw Normal View History

2024-09-02 22:09:26 +02:00
package main
import (
"bufio"
"context"
"encoding/json"
"fmt"
"os"
"os/exec"
"path/filepath"
"sort"
"strconv"
"strings"
"golang.org/x/sync/errgroup"
)
type CPUInfo struct {
CPU int `json:"cpu"`
MaxMHz float64 `json:"maxmhz"`
}
type LSCPUOutput struct {
CPUs []CPUInfo `json:"cpus"`
}
type CoreFreq struct {
Core int
Freq int
}
func main() {
args := os.Args[1:]
if len(args) == 0 {
fmt.Println("Please provide a command to run")
os.Exit(1)
}
eg, ctx := errgroup.WithContext(context.Background())
canPerf, cores, err := getBestCores()
if err != nil {
eg.Go(func() error {
return runCommand(ctx, false, []int{}, args)
})
err = eg.Wait()
if err != nil {
fmt.Println(err)
os.Exit(1)
}
os.Exit(0)
}
eg.Go(func() error {
return runCommand(ctx, canPerf, cores, args)
})
err = eg.Wait()
if err != nil {
fmt.Println(err)
os.Exit(1)
}
os.Exit(0)
}
func runCommand(ctx context.Context, canPerf bool, cores []int, args []string) error {
commandToRun := make([]string, 0)
if canPerf {
commandToRun = append(commandToRun, "powerprofilesctl", "launch", "-p", "performance")
}
if len(cores) > 0 {
coreVals := make([]string, len(cores))
for i, core := range cores {
coreVals[i] = fmt.Sprint(core)
}
commandToRun = append(commandToRun, "taskset", "-c", strings.Join(coreVals, ","))
}
commandToRun = append(commandToRun, args...)
cmd := exec.CommandContext(ctx, commandToRun[0], commandToRun[1:]...)
cmd.Env = append(cmd.Environ(), "POSIXLY_CORRECT=1")
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
cmd.Stdin = os.Stdin
return cmd.Run()
}
func getBestCores() (bool, []int, error) {
perfMode, amdPrefCores, amdFlag, highestCores, err := GetPPStatusAndHighestCores()
if err != nil {
fmt.Println(err)
os.Exit(1)
}
if !amdFlag {
return perfMode, highestCores, nil
}
numAmdCores := len(amdPrefCores)
switch numAmdCores {
case 32:
return perfMode, getHighestCores(16, amdPrefCores), nil
case 24:
return perfMode, getHighestCores(12, amdPrefCores), nil
case 16:
return perfMode, getHighestCores(8, amdPrefCores), nil
default:
return perfMode, getHighestCores(numAmdCores, amdPrefCores), nil
}
}
func getHighestCores(n int, coreFreqs []CoreFreq) []int {
if n <= 0 {
return []int{}
}
sort.Slice(coreFreqs, func(i, j int) bool {
return coreFreqs[i].Freq > coreFreqs[j].Freq
})
result := make([]int, 0, n)
for i := 0; i < n && i < len(coreFreqs); i++ {
result = append(result, coreFreqs[i].Core)
}
sort.Stable(sort.IntSlice(result))
return result
}
func GetPPStatusAndHighestCores() (bool, []CoreFreq, bool, []int, error) {
var (
performanceMode bool
amdPstateFreqs []CoreFreq
amdFlag bool
highestCores []int
)
g, ctx := errgroup.WithContext(context.Background())
g.Go(func() error {
cmd := exec.CommandContext(ctx, "powerprofilesctl", "list")
output, err := cmd.Output()
if err != nil {
performanceMode = false
return nil
}
performanceMode = strings.Contains(string(output), "performance")
return nil
})
g.Go(func() error {
pattern := "/sys/devices/system/cpu/cpu*/cpufreq/amd_pstate_prefcore_ranking"
matches, err := filepath.Glob(pattern)
if err != nil {
return fmt.Errorf("failed to glob AMD pstate files: %w", err)
}
if len(matches) > 0 {
amdFlag = true
}
for _, match := range matches {
file, err := os.Open(match)
if err != nil {
return fmt.Errorf("failed to open file %s: %w", match, err)
}
defer file.Close()
parts := strings.Split(match, "/")
cpuPart := parts[len(parts)-3]
coreNum, err := strconv.Atoi(strings.TrimPrefix(cpuPart, "cpu"))
if err != nil {
return fmt.Errorf("failed to parse core number from path %s: %w", match, err)
}
scanner := bufio.NewScanner(file)
if scanner.Scan() {
value, err := strconv.Atoi(strings.TrimSpace(scanner.Text()))
if err != nil {
return fmt.Errorf("failed to parse value from file %s: %w", match, err)
}
amdPstateFreqs = append(amdPstateFreqs, CoreFreq{Core: coreNum, Freq: value})
}
if err := scanner.Err(); err != nil {
return fmt.Errorf("error reading file %s: %w", match, err)
}
}
return nil
})
g.Go(func() error {
cmd := exec.CommandContext(ctx, "lscpu", "-e", "-J")
output, err := cmd.Output()
if err != nil {
return fmt.Errorf("failed to run lscpu command: %w", err)
}
var lscpuOutput LSCPUOutput
if err := json.Unmarshal(output, &lscpuOutput); err != nil {
return fmt.Errorf("failed to unmarshal lscpu output: %w", err)
}
var maxMHz float64
for _, cpu := range lscpuOutput.CPUs {
if cpu.MaxMHz > maxMHz {
maxMHz = cpu.MaxMHz
}
}
for _, cpu := range lscpuOutput.CPUs {
if cpu.MaxMHz == maxMHz {
highestCores = append(highestCores, cpu.CPU)
}
}
return nil
})
if err := g.Wait(); err != nil {
return false, nil, false, nil, err
}
return performanceMode, amdPstateFreqs, amdFlag, highestCores, nil
}