ma/agent/main.go

package main

import (
	"io"
	"log/slog"
	"math"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"
)

type cpuHandler struct {
	corePaths []string
	maxFreq   int
	minFreq   int
	logger    *slog.Logger
}

// readIntFromFile reads the given file until the end
// and converts the contained string into an integer
func readIntFromFile(path string) (int, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return 0, err
	}
	return strconv.Atoi(strings.TrimSpace(string(b)))
}

func enumerateCPUCores(logger *slog.Logger) ([]string, int, int, error) {
	files, err := os.ReadDir("/sys/devices/system/cpu/")
	if err != nil {
		logger.Error("reading /sys/devices/system/cpu/", "err", err)
		return nil, -1, -1, err
	}
	r, _ := regexp.Compile("cpu([0-9]+)")

	freqPaths := make([]string, 0)
	maxF := math.MaxInt
	minF := math.MinInt

	// we look for cpuN directories
	for _, f := range files {
		if !f.IsDir() || !r.MatchString(f.Name()) {
			continue
		}

		// read files for min & max freqs
		dir := "/sys/devices/system/cpu/" + f.Name() + "/cpufreq/"
		minFreq, err := readIntFromFile(dir + "cpuinfo_min_freq")
		if err != nil {
			logger.Error("failed to read cpu min freq", "path", dir+"cpuinfo_min_freq", "error", err)
			continue
		}
		maxFreq, err := readIntFromFile(dir + "cpuinfo_max_freq")
		if err != nil {
			logger.Error("failed to read cpu max freq", "path", dir+"cpuinfo_min_freq", "error", err)
			continue
		}
		freqPaths = append(freqPaths, dir+"scaling_max_freq")
		maxF = min(maxF, maxFreq)
		minF = max(minF, minFreq)
	}
	return freqPaths, minF, maxF, nil
}

func (s *cpuHandler) handleMin(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		writeResponse("invalid method", http.StatusBadRequest, w, s.logger)
		return
	}
	writeResponse(strconv.Itoa(s.minFreq), http.StatusOK, w, s.logger)
}

func (s *cpuHandler) handleMax(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		writeResponse("invalid method", http.StatusBadRequest, w, s.logger)
		return
	}
	writeResponse(strconv.Itoa(s.maxFreq), http.StatusOK, w, s.logger)
}

func (s *cpuHandler) handleCurrent(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		maxFreq := math.MinInt
		for _, core := range s.corePaths {
			freq, err := readIntFromFile(core)
			if err != nil {
				s.logger.Error("error while reading core frequency", "core", core, "error", err)
				writeResponse("internal server error", http.StatusInternalServerError, w, s.logger)
				return
			}
			s.logger.Debug("core frequency", "core", core, "frequency", freq)
			maxFreq = max(maxFreq, freq)
		}
		s.logger.Debug("reported frequency", "freq", maxFreq)
		writeResponse(strconv.Itoa(maxFreq), http.StatusOK, w, s.logger)
	case "PUT":
		body, err := io.ReadAll(r.Body)
		if err != nil {
			s.logger.Error("failed to read the request body", "err", err)
			writeResponse("internal server error", http.StatusInternalServerError, w, s.logger)
		}
		bodyStr := strings.TrimSpace(string(body))
		// Check if it is an integer
		freqInt, err := strconv.Atoi(bodyStr)
		if err != nil {
			s.logger.Error("failed to convert body to int", "error", err)
			writeResponse("bad request", http.StatusBadRequest, w, s.logger)
		}
		if freqInt < s.minFreq || freqInt > s.maxFreq {
			s.logger.Warn("requested frequency is out-of-bounds", "freq", freqInt, "min", s.minFreq, "max", s.maxFreq)
			freqInt = max(s.minFreq, min(freqInt, s.maxFreq))
			bodyStr = strconv.Itoa(freqInt)
		}
		s.logger.Info("requested to set current frequency", "frequency", bodyStr)
		success := true
		for _, core := range s.corePaths {
			err = os.WriteFile(core, []byte(bodyStr), 0)
			if err != nil {
				s.logger.Error("failed to set core frequency", "core", core, "error", err)
				success = false
			}
		}
		if success {
			w.WriteHeader(http.StatusOK)
		} else {
			w.WriteHeader(http.StatusInternalServerError)
		}
	default:
		s.logger.Error("unexpected method for current cpu", "method", r.Method)
		writeResponse("invalid method", http.StatusBadRequest, w, s.logger)
	}
}

func (s *cpuHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	requestedPath := getStrippedRequestPath(r.URL, r.Pattern)
	switch requestedPath {
	case "frequency/min":
		s.handleMin(w, r)
		return
	case "frequency/max":
		s.handleMax(w, r)
		return
	case "frequency/current":
		s.handleCurrent(w, r)
		return
	default:
		s.logger.Error("unknown resource", "path", requestedPath)
		writeResponse("not found", http.StatusNotFound, w, s.logger)
	}
}

func main() {
	appEnv := os.Getenv("AGENT_ENV")
	loggerOpts := &slog.HandlerOptions{
		Level: slog.LevelDebug,
	}
	var handler slog.Handler = slog.NewTextHandler(os.Stdout, loggerOpts)
	if appEnv == "production" {
		loggerOpts.Level = slog.LevelInfo
		handler = slog.NewJSONHandler(os.Stdout, loggerOpts)
	}
	logger := slog.New(handler)

	freqPaths, minF, maxF, err := enumerateCPUCores(logger)
	if err != nil {
		os.Exit(1)
	}

	cpuHandler := &cpuHandler{
		logger:    logger,
		maxFreq:   maxF,
		minFreq:   minF,
		corePaths: freqPaths,
	}
	http.Handle("/cpu/", logRequests(logger, cpuHandler))
	logger.Info("exit", "result", http.ListenAndServe(":8080", nil))
}