Merge pull request #9 from doraemonkeys/master

Add CPU Usage Maintenance Feature

Author: shawn-bluce

README.md

@@ -13,7 +13,7 @@ Developer will encounter the need to quickly occupy CPU and memory, I am also de
   - [ ] macOs
   - [ ] Windows
 - [x] Memory read/write periodically , prevent memory from being swapped out
-- [ ] Dynamic adjustment of CPU and memory usage
+- [ ] Dynamic adjustment of memory usage
 - [ ] Eat GPU
 
 # Usage
@@ -27,6 +27,7 @@ Usage:
 
 Flags:
   --cpu-affinities ints                  Which cpu core(s) would you want to eat? multiple cores separate by ','
+  --cpu-maintain string                  How many cpu would you want maintain(e.g. 50%)
   -c, --cpu-usage string                 How many cpu would you want eat (default "0")
   -h, --help                             help for eat
   -r, --memory-refresh-interval string   How often to trigger a refresh to prevent the ate memory from being swapped out (default "5m")
@@ -47,6 +48,9 @@ eat -c 3 -m 200m	# eating 3 CPU core and 200MB memory
 eat -c 100% -m 100%	# eating all CPU core and memory
 eat -c 100% -t 1h	# eating all CPU core and quit after 1hour
 
+eat --cpu-maintain 50%	# dynamic adjust to maintain minimum 50% CPU usage
+eat --cpu-maintain 50% -c 100%	# dynamic adjust to maintain minimum 50% CPU usage and use all CPU core
+
 eat --cpu-affinities 0 -c 1	# only run eat in core #0 (first core)
 eat --cpu-affinities 0,1 -c 2	# run eat in core #0,1 (first and second core)
 eat --cpu-affinities 0,1,2,3 -c 100% # error case: in-enough cpu affinities
@@ -101,6 +105,7 @@ $ ./eat.out --help
 
 Flags：
   --cpu-affinities 			整数	指定在几个核心上运行 Eat，多个核心索引之间用 ',' 分隔，索引从 0 开始。
+  --cpu-maintain 			  字符串	你想将CPU使用率维持在多少（e.g. 50%）
   -c, --cpu-usage 			字符串	你想吃掉多少个 CPU（默认为 '0'）？
   -h，--help				输出 eat 的帮助
   -r, --memory-refresh-interval 字符串	每隔多长时间触发一次刷新，以防止被吃掉的内存被交换出去（默认值为 '5m'）
@@ -121,6 +126,9 @@ eat -c 3 -m 200m	# 占用3个CPU核和200MB内存
 eat -c 100% -m 100%	# 占用所有CPU核和内存
 eat -c 100% -t 1h	# 占用所有CPU核并在一小时后退出
 
+eat --cpu-maintain 50%	# 动态调整维持50%的CPU使用率
+eat --cpu-maintain 50% -c 100%	# 使用所有CPU核心动态调整维持50%的CPU使用率
+
 eat --cpu-affinities 0 -c 1	# 只占用 #0 第一个核心
 eat --cpu-affinities 0,1 -c 2	# 占用 #0,1 前两个个核心
 eat --cpu-affinities 0,1,2,3 -c 100%	# 错误参数: 每个请求核都要指定对应的亲和性核心

cmd/cpu.go

@@ -5,11 +5,13 @@ import (
 	"fmt"
 	"log"
 	"math"
+	"math/rand/v2"
 	"runtime"
 	"sync"
 	"time"
 
 	"eat/cmd/cpu_affinity"
+	"eat/cmd/sysinfo"
 )
 
 func busyWork(ctx context.Context) {
@@ -142,3 +144,153 @@ func eatCPU(ctx context.Context, wg *sync.WaitGroup, c float64, cpuAffinitiesEat
 
 	fmt.Printf("Ate %2.3f CPU cores\n", c)
 }
+
+func maintainCpuUsage(ctx context.Context, wg *sync.WaitGroup, coreNum float64, usagePercent float64, cpuAffinitiesEat []uint, cpuMonitor sysinfo.SystemCPUMonitor) {
+	if coreNum == 0 {
+		coreNum = float64(runtime.NumCPU())
+	}
+	fmt.Printf("CPU usage will be maintaining at minimum %.3f%%, eating %.3f cores, be patient...\n", usagePercent, coreNum)
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		MaintainCpuUsage(ctx, coreNum, usagePercent, cpuAffinitiesEat, cpuMonitor)
+	}()
+}
+
+func MaintainCpuUsage(ctx context.Context, coreNum float64, usagePercent float64, cpuAffinitiesEat []uint, cpuMonitor sysinfo.SystemCPUMonitor) {
+	runtime.GOMAXPROCS(runtime.NumCPU())
+
+	fullCores := int(coreNum)
+	partialCoreRatio := coreNum - float64(fullCores)
+
+	const maxIdleDuration = 1 * time.Second
+	const minIdleDuration = 1 * time.Millisecond
+	const initIdleDurationAdjustRatio float64 = 0.1
+	const minIdleDurationAdjustRatio = 0.002
+	var idleDuration = maxIdleDuration
+	var dynIdleDurationAdjustRatio float64 = initIdleDurationAdjustRatio
+	var stopWork = false
+	var cur float64 = 0
+	var ctxDone = false
+
+	var fixIdleDuration = func() {
+		var err error
+		cur, err = cpuMonitor.GetCPUUsage()
+		if err != nil {
+			log.Printf("MaintainCpuUsage: get cpu usage failed, reason: %s", err.Error())
+			return
+		}
+		// When the cpu usage fluctuates greatly, increase idleDurationAdjustRatio to stabilize the cpu usage
+		if dynIdleDurationAdjustRatio == minIdleDurationAdjustRatio {
+			if cur > usagePercent+20 || cur < usagePercent-20 {
+				dynIdleDurationAdjustRatio = initIdleDurationAdjustRatio
+			} else if cur > usagePercent+10 || cur < usagePercent-10 {
+				dynIdleDurationAdjustRatio = initIdleDurationAdjustRatio * 0.5
+			}
+		}
+		if cur > usagePercent {
+			idleDuration = time.Duration(float64(idleDuration) * (1 + dynIdleDurationAdjustRatio))
+		} else if cur < usagePercent {
+			idleDuration = time.Duration(float64(idleDuration) * (1 - dynIdleDurationAdjustRatio))
+		}
+		if idleDuration < minIdleDuration {
+			idleDuration = minIdleDuration
+		} else if idleDuration > maxIdleDuration {
+			idleDuration = maxIdleDuration
+			stopWork = true
+		} else {
+			stopWork = false
+		}
+		// gradually decrease the idle duration adjustment ratio, make the idle duration more stable
+		dynIdleDurationAdjustRatio -= 0.001
+		dynIdleDurationAdjustRatio = max(minIdleDurationAdjustRatio, dynIdleDurationAdjustRatio)
+	}
+	var worker = func(wg *sync.WaitGroup, idx int, workerName string, work func()) {
+		defer wg.Done()
+		cleanup, err := setCpuAffWrapper(idx, cpuAffinitiesEat)
+		if err != nil {
+			fmt.Printf("Error: %s failed to set cpu affinities, reason: %s\n", workerName, err.Error())
+			return
+		}
+		if cleanup != nil {
+			fmt.Printf("Worker %s: CPU affinities set to %d\n", workerName, cpuAffinitiesEat[idx])
+			defer cleanup()
+		}
+		for {
+			if ctxDone {
+				return
+			}
+			if !stopWork {
+				work()
+			}
+			// if idleDuration is less than 1ms, do not sleep, directly execute fixIdleDuration
+			if idleDuration > time.Millisecond*1 {
+				time.Sleep(idleDuration)
+			}
+		}
+	}
+	cpuIntensiveTask := GenerateCPUIntensiveTask(time.Microsecond * 2000) // 2ms is empirical data
+	wg := &sync.WaitGroup{}
+	for i := 0; i < fullCores; i++ {
+		wg.Add(1)
+		workerName := fmt.Sprintf("%d@fullCore", i)
+		go worker(wg, i, workerName, cpuIntensiveTask)
+	}
+	if partialCoreRatio > 0 {
+		wg.Add(1)
+		workerName := fmt.Sprintf("%d@partCore", fullCores)
+		go worker(wg, fullCores, workerName, cpuIntensiveTask)
+	}
+	fmt.Print("\033[?25l")       // hide cursor
+	defer fmt.Print("\033[?25h") // show cursor
+
+	ticker := time.NewTicker(time.Millisecond * 300)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			ctxDone = true
+			fmt.Println("MaintainCpuUsage: quit due to context being cancelled")
+			wg.Wait()
+			return
+		case <-ticker.C:
+			fixIdleDuration()
+			printCpuInfo(cur, idleDuration, dynIdleDurationAdjustRatio)
+		}
+	}
+}
+
+func printCpuInfo(usagePercent float64, idleDuration time.Duration, ratio float64) {
+	// clear current line
+	fmt.Print("\033[2K")
+	fmt.Printf("Idle: %s\n", idleDuration)
+	fmt.Printf("Ratio: %.3f%%\n", ratio)
+	fmt.Printf("CPU usage:\033[32m%6.2f%%\033[0m\n", usagePercent)
+	fmt.Print("\033[3A\033[G")
+}
+
+// GenerateCPUIntensiveTask returns a function that performs a duration CPU-intensive task
+func GenerateCPUIntensiveTask(duration time.Duration) func() {
+	const N = 1000
+	start := time.Now()
+	iteration := 0
+	var cnt int64 = 0
+	var i int
+	for time.Since(start) < duration {
+		for i = 0; i < N; i++ {
+			cnt = cnt * rand.Int64N(100)
+			iteration++
+		}
+	}
+	return func() {
+		var cnt2 int64 = cnt
+		for i = 0; i < iteration; i++ {
+			cnt2 = cnt2 * rand.Int64N(100)
+			cnt++
+		}
+		// KeepAlive ensures that the variable is not optimized away by the compiler
+		runtime.KeepAlive(cnt)
+		runtime.KeepAlive(cnt2)
+	}
+}

cmd/cpu_test.go

@@ -0,0 +1,33 @@
+package cmd
+
+import (
+	"testing"
+	"time"
+)
+
+func TestGenerateCPUIntensiveTask(t *testing.T) {
+
+	tests := []struct {
+		name     string
+		duration time.Duration
+	}{
+		{"2ms", time.Millisecond * 2},
+		{"5ms", time.Millisecond * 5},
+		{"20ms", time.Millisecond * 20},
+		{"100ms", time.Millisecond * 100},
+		{"500ms", time.Millisecond * 500},
+		{"1s", time.Second},
+		{"2s", time.Second * 2},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := GenerateCPUIntensiveTask(tt.duration)
+			now := time.Now()
+			got()
+			after := time.Now()
+			if (after.Sub(now).Abs() - tt.duration) > max(tt.duration*8/10, time.Millisecond*2) {
+				t.Errorf("GenerateCPUIntensiveTask() = %v, want %v", after.Sub(now), max(tt.duration*8/10, time.Millisecond*2))
+			}
+		})
+	}
+}

cmd/parse.go

@@ -3,11 +3,14 @@ package cmd
 import (
 	"fmt"
 	"math"
+	"os"
 	"runtime"
 	"strconv"
+	"strings"
 	"time"
 
 	"eat/cmd/cpu_affinity"
+
 	"github.com/pbnjay/memory"
 )
 
@@ -35,6 +38,28 @@ func parseEatCPUCount(c string) float64 {
 	}
 }
 
+// parseCPUMaintainPercent parse cpu usage percent, return value is percent(0-100)
+func parseCPUMaintainPercent(c string) float64 {
+	if c == "" {
+		return 0
+	}
+	if !strings.HasSuffix(c, "%") {
+		fmt.Println("Error: invalid cpu maintain percent, must end with %")
+		os.Exit(1)
+	}
+	c = strings.TrimSuffix(c, "%")
+	cMaintain, err := strconv.ParseFloat(c, 32)
+	if err != nil {
+		fmt.Println("Error: invalid cpu maintain percent:", c)
+		os.Exit(1)
+	}
+	if cMaintain < 0 || cMaintain > 100 {
+		fmt.Println("Error: invalid cpu maintain percent:", cMaintain)
+		os.Exit(1)
+	}
+	return cMaintain
+}
+
 func parseEatMemoryBytes(m string) uint64 {
 	// allow g/G, m/M, k/K suffixes
 	// 1G = 1024M = 1048576K

cmd/root.go

@@ -13,7 +13,9 @@ import (
 	"syscall"
 	"time"
 
+	"eat/cmd/sysinfo"
 	"eat/cmd/version"
+
 	"github.com/pbnjay/memory"
 	"github.com/spf13/cobra"
 )
@@ -85,17 +87,22 @@ func eatFunction(cmd *cobra.Command, _ []string) {
 
 	// Get the flags
 	c, _ := cmd.Flags().GetString("cpu-usage")
+	cMaintain, _ := cmd.Flags().GetString("cpu-maintain")
 	cAff, _ := cmd.Flags().GetIntSlice("cpu-affinities")
 	m, _ := cmd.Flags().GetString("memory-usage")
 	dl, _ := cmd.Flags().GetString("time-deadline")
 	r, _ := cmd.Flags().GetString("memory-refresh-interval")
 
-	if c == "0" && m == "0m" {
+	if c == "0" && m == "0m" && cMaintain == "" {
 		fmt.Println("Error: no cpu or memory usage specified")
 		return
 	}
+	if c == "0" && cMaintain != "" {
+		c = cMaintain
+	}
 
 	cEat := parseEatCPUCount(c)
+	cMaintainPercent := parseCPUMaintainPercent(cMaintain)
 	phyCores := runtime.NumCPU()
 	if int(math.Ceil(cEat)) > phyCores {
 		fmt.Printf("Error: user specified cpu cores exceed system physical cores(%d)\n", phyCores)
@@ -115,7 +122,11 @@ func eatFunction(cmd *cobra.Command, _ []string) {
 	defer cancel()
 	fmt.Printf("Want to eat %2.3fCPU, %s Memory\n", cEat, m)
 	eatMemory(rootCtx, &wg, mEat, mAteRenew)
-	eatCPU(rootCtx, &wg, cEat, cpuAffinitiesEat)
+	if cMaintainPercent > 0 {
+		maintainCpuUsage(rootCtx, &wg, cEat, cMaintainPercent, cpuAffinitiesEat, sysinfo.Monitor)
+	} else {
+		eatCPU(rootCtx, &wg, cEat, cpuAffinitiesEat)
+	}
 	// in case that all sub goroutines are dead due to runtime error like memory not enough.
 	// so the main goroutine automatically quit as well, don't wait user ctrl+c or context deadline.
 	go func(wgp *sync.WaitGroup) {