Feat support cpu affinity for eat in linux

.github/workflows/go.yml

@@ -19,7 +19,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v4
         with:
-          go-version: '1.20'
+          go-version: '1.22'
 
       - name: Build
         run: go build -v ./...

README.md

@@ -8,25 +8,53 @@ Developer will encounter the need to quickly occupy CPU and memory, I am also de
 - [x] Support `eat -c 35%` and `eat -m 35%`
 - [x] support gracefully exit: capture process signal SIGINT(2), SIGTERM(15)
 - [x] support deadline: `-t` specify the duration of eat progress. such as "300ms", "1.5h", "2h45m". (unit: "ns", "us" (or "µs"), "ms", "s", "m", "h")
-- [ ] CPU Affinity
+- [x] CPU Affinity
+  - [x] Linux
+  - [ ] macOs
+  - [ ] Windows
 - [x] Memory read/write periodically , prevent memory from being swapped out
 - [ ] Dynamic adjustment of CPU and memory usage
 - [ ] Eat GPU
 
 # Usage
 
 ```shell
-eat -c 4            # eating 4 CPU core
-eat -c 35%          # eating 35% CPU core (CPU count * 35%)
-eat -c 100%         # eating all CPU core
-eat -m 4g           # eating 4GB memory
-eat -m 20m          # eating 20MB memory
-eat -m 35%          # eating 35% memory (total memory * 35%)
-eat -m 100%         # eating all memory
-eat -c 2.5 -m 1.5g  # eating 2.5 CPU core and 1.5GB memory
-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.out --help
+A monster that eats cpu and memory 🦕
+
+Usage:
+    eat [flags]
+
+Flags:
+  --cpu-affinities ints                  Which cpu core(s) would you want to eat? multiple cores separate by ','
+  -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")
+  -m, --memory-usage string              How many memory would you want eat(GB) (default "0m")
+  -t, --time-deadline string             Deadline to quit eat process (default "0")
+```
+
+```shell
+eat -c 4										# eating 4 CPU core
+eat -c 35%										# eating 35% CPU core (CPU count * 35%)
+eat -c 100%										# eating all CPU core
+eat -m 4g										# eating 4GB memory
+eat -m 20m										# eating 20MB memory
+eat -m 35%										# eating 35% memory (total memory * 35%)
+eat -m 100%										# eating all memory
+eat -c 2.5 -m 1.5g								# eating 2.5 CPU core and 1.5GB memory
+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-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
+# Have 8C15G.
+# Error: failed to parse cpu affinities, reason: each request cpu cores need specify its affinity, aff 4 < req 8
+eat --cpu-affinities 0,1,2,3 -c 50%				# run eat in core #0,1,2,3 (first to fourth core)
+eat --cpu-affinities 0,1,2,3,4,5,6,7 -c 92% 	# run eat in all core(full of 7 cores, part of last core)
+
 ```
 
 > Tips:
@@ -35,11 +63,16 @@ eat -c 100% -t 1h # eating all CPU core and quit after 1hour
 # Build
 
 ```shell
-go build -o eat
+# Linux
+GOOS=linux GOARCH=amd64 go build -trimpath -ldflags "-s -w" -v -o eat
+# macOs
+GOOS=darwin GOARCH=amd64 go build -trimpath -ldflags "-s -w" -v -o eat_mac
+# Windows
+GOOS=windwos GOARCH=amd64 go build -trimpath -ldflags "-s -w" -v -o eat_win
 ```
 
 # 介绍
-<b>我是一个吃CPU和内存的怪兽🦕</b>
+<b>我是一只吃CPU和内存的怪兽🦕</b>
 
 开发者们经常会遇到需要快速占用 CPU 和内存的需求，我也是。所以我开发了一个名为 `eat` 的小工具来快速占用指定数量的 CPU 和内存。
 
@@ -48,25 +81,54 @@ go build -o eat
 - [x] 支持`eat -c 35%`和`eat -m 35%`
 - [x] 支持优雅退出: 捕捉进程 SIGINT, SIGTERM 信号实现有序退出
 - [x] 支持时限: `-t` 限制吃资源的时间，示例 "300ms", "1.5h", "2h45m". (单位: "ns", "us" (or "µs"), "ms", "s", "m", "h")
-- [ ] CPU亲和性
+- [x] CPU亲和性
+  - [X] Linux 
+  - [ ] macOS
+  - [ ] Windows
 - [x] 定期内存读写，防止内存被交换出去
 - [ ] 动态调整CPU和内存使用
 - [ ] 吃GPU
 
 # 使用
 
+
 ```shell
-eat -c 4            # 占用4个CPU核
-eat -c 35%          # 占用35%CPU核（CPU核数 * 35%）
-eat -c 100%         # 占用所有CPU核
-eat -m 4g           # 占用4GB内存
-eat -m 20m          # 占用20MB内存
-eat -m 35%          # 占用35%内存（总内存 * 35%）
-eat -m 100%         # 占用所有内存
-eat -c 2.5 -m 1.5g  # 占用2.5个CPU核和1.5GB内存
-eat -c 3 -m 200m    # 占用3个CPU核和200MB内存
-eat -c 100% -m 100% # 占用所有CPU核和内存
-eat -c 100% -t 1h   # 占用所有CPU核并在一小时后退出
+$ ./eat.out --help
+我是一只吃CPU和内存的怪兽🦕
+
+使用方法
+    eat [flags］
+
+Flags：
+  --cpu-affinities 整数						指定在几个核心上运行 Eat，多个核心索引之间用 ',' 分隔，索引从 0 开始。
+  -c, --cpu-usage 字符串						你想吃掉多少个 CPU（默认为 '0'）？
+  -h，--help								输出 eat 的帮助
+  -r, --memory-refresh-interval 字符串		每隔多长时间触发一次刷新，以防止被吃掉的内存被交换出去（默认值为 '5m'）
+  -m, --memory-usage 字符串					你希望吃掉多少内存（GB）（默认值 '0m'）
+  -t，--time-deadline 字符串					退出 eat 进程的截止日期（默认为 "0'）。
+```
+
+```shell
+eat -c 4									# 占用4个CPU核
+eat -c 35%									# 占用35%CPU核（CPU核数 * 35%）
+eat -c 100%									# 占用所有CPU核
+eat -m 4g									# 占用4GB内存
+eat -m 20m									# 占用20MB内存
+eat -m 35%									# 占用35%内存（总内存 * 35%）
+eat -m 100%									# 占用所有内存
+eat -c 2.5 -m 1.5g							# 占用2.5个CPU核和1.5GB内存
+eat -c 3 -m 200m							# 占用3个CPU核和200MB内存
+eat -c 100% -m 100%							# 占用所有CPU核和内存
+eat -c 100% -t 1h							# 占用所有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% 		# 错误参数: 每个请求核都要指定对应的亲和性核心
+# Have 8C15G.
+# Error: failed to parse cpu affinities, reason: each request cpu cores need specify its affinity, aff 4 < phy 8
+# 出错: 无法解析 CPU 亲和性, 原因: 每个请求核都要指定对应的亲和性核心, 亲和核  4 < 请求核 8
+eat --cpu-affinities 0,1,2,3 -c 50%			# 占用前4个核心
+eat --cpu-affinities 0,1,2,3,4,5,6,7 -c 92%	# 占用前8个核心 (全部7个核心，部分的最后一个核心)
 ```
 
 > 提示：
@@ -75,5 +137,10 @@ eat -c 100% -t 1h   # 占用所有CPU核并在一小时后退出
 # 构建
 
 ```shell
-go build -o eat
-```
+# Linux
+GOOS=linux GOARCH=amd64 go build -trimpath -ldflags "-s -w" -v -o eat
+# macOs
+GOOS=darwin GOARCH=amd64 go build -trimpath -ldflags "-s -w" -v -o eat_mac
+# Windows
+GOOS=windwos GOARCH=amd64 go build -trimpath -ldflags "-s -w" -v -o eat_win
+```

cmd/constant.go

@@ -1,12 +1,28 @@
 package cmd
 
 import (
+	"fmt"
 	"time"
 )
 
+// contextKey is a value for use with context.WithValue.
+// It's used as a pointer. so it fits in an interface{} without allocation.
+type contextKey struct {
+	name      string
+	valueType string
+}
+
+func (k *contextKey) String() string {
+	return fmt.Sprintf("worker context value: name %s, type %s", k.name, k.valueType)
+}
+
 const (
 	intervalCpuWorkerCheckContextDone = 10000
 	durationMemoryWorkerDoRefresh     = 5 * time.Minute
 	durationEachSignCheck             = 100 * time.Millisecond
 	chunkSizeMemoryWorkerEachAllocate = 128 * 1024 * 1024 // 128MB
 )
+
+var (
+	cpuWorkerPartialCoreRatioContextKey = &contextKey{"partialCoreRatio", "float64"}
+)

cmd/cpu.go

@@ -8,6 +8,8 @@ import (
 	"runtime"
 	"sync"
 	"time"
+
+	"eat/cmd/cpu_affinity"
 )
 
 func busyWork(ctx context.Context) {
@@ -26,11 +28,16 @@ func busyWork(ctx context.Context) {
 	}
 }
 
-func partialBusyWork(ctx context.Context, ratio float64) {
+func partialBusyWork(ctx context.Context) {
 	const (
 		oneCycle  = 10 * time.Microsecond
 		precision = 1000
 	)
+	ratio, ok := ctx.Value(cpuWorkerPartialCoreRatioContextKey).(float64)
+	if !ok {
+		log.Printf("partialBusyWork: partial core ratio context key not set or type ")
+		return
+	}
 	// round busy and idle percent
 	// case 1: ratio 0.8
 	//   busy 0.8                     idle 0.19999999999999996
@@ -63,9 +70,55 @@ func partialBusyWork(ctx context.Context, ratio float64) {
 	}
 }
 
-func eatCPU(ctx context.Context, wg *sync.WaitGroup, c float64) {
-	fmt.Printf("Eating %-12s", "CPU...")
+func startEatCpuWorker(ctx context.Context, wg *sync.WaitGroup, workerName string, idx int, workerFunc func(ctx context.Context), cpuAffinitiesEat []uint) {
+	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()
+	}
+	workerFunc(ctx)
+}
 
+func setCpuAffWrapper(index int, cpuAffinitiesEat []uint) (func(), error) {
+	if len(cpuAffinitiesEat) == 0 { // user not set cpu affinities, skip...
+		return nil, nil
+	}
+	if len(cpuAffinitiesEat) <= index { // index error
+		return nil, fmt.Errorf("cpuAffinities: index out of range")
+	}
+	// LockOSThread wires the calling goroutine to its current operating system thread.
+	// The calling goroutine will **always execute** in that thread, and no other goroutine will execute in it,
+	// until the calling goroutine has made as many calls to [UnlockOSThread] as to LockOSThread.
+	// If the calling goroutine exits without unlocking the thread, the thread will be terminated.
+	//
+	// All init functions are run on the startup thread. Calling LockOSThread
+	// from an init function will cause the main function to be invoked on
+	// that thread.
+	//
+	// A goroutine should **call LockOSThread before** calling OS services or non-Go library functions
+	// that depend on per-thread state.
+	runtime.LockOSThread() // IMPORTANT!! Only limit the system thread affinity, not the whole go program process
+	var cpuAffDeputy = cpu_affinity.NewCpuAffinityDeputy()
+	if !cpuAffDeputy.IsImplemented() {
+		return nil, fmt.Errorf("SetCpuAffinities currently not support in this os: %s", runtime.GOOS)
+	}
+	tid := cpuAffDeputy.GetThreadId()
+	err := cpuAffDeputy.SetCpuAffinities(uint(tid), cpuAffinitiesEat[index])
+	if err != nil {
+		return nil, err
+	}
+	return func() {
+		runtime.UnlockOSThread()
+	}, nil
+}
+
+func eatCPU(ctx context.Context, wg *sync.WaitGroup, c float64, cpuAffinitiesEat []uint) {
+	fmt.Printf("Eating %-12s", "CPU...")
 	runtime.GOMAXPROCS(runtime.NumCPU())
 
 	fullCores := int(c)
@@ -74,19 +127,17 @@ func eatCPU(ctx context.Context, wg *sync.WaitGroup, c float64) {
 	// eat full cores
 	for i := 0; i < fullCores; i++ {
 		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			busyWork(ctx)
-		}()
+		workerName := fmt.Sprintf("%d@fullCore", i)
+		go startEatCpuWorker(ctx, wg, workerName, i, busyWork, cpuAffinitiesEat)
 	}
 
 	// eat partial core
 	if partialCoreRatio > 0 {
+		i := fullCores // the last core affinity
 		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			partialBusyWork(ctx, partialCoreRatio)
-		}()
+		workerName := fmt.Sprintf("%d@partCore", i)
+		childCtx := context.WithValue(ctx, cpuWorkerPartialCoreRatioContextKey, partialCoreRatio)
+		go startEatCpuWorker(childCtx, wg, workerName, i, partialBusyWork, cpuAffinitiesEat)
 	}
 
 	fmt.Printf("Ate %2.3f CPU cores\n", c)

cmd/cpu_affinity/cpu_affinity_linux.go

@@ -0,0 +1,51 @@
+//go:build linux
+// +build linux
+
+package cpu_affinity
+
+import (
+	"runtime"
+	"syscall"
+
+	"golang.org/x/sys/unix"
+)
+
+type CpuAffinityDeputy struct{}
+
+func (CpuAffinityDeputy) GetProcessId() uint {
+	return uint(syscall.Getpid())
+}
+
+func (CpuAffinityDeputy) GetThreadId() uint {
+	return uint(syscall.Gettid())
+}
+
+func (CpuAffinityDeputy) SetCpuAffinities(pid uint, cpus ...uint) error {
+	if len(cpus) == 0 {
+		return nil
+	}
+	mask := new(unix.CPUSet)
+	mask.Zero()
+	for _, c := range cpus {
+		mask.Set(int(c))
+	}
+	return unix.SchedSetaffinity(int(pid), mask)
+}
+
+func (CpuAffinityDeputy) GetCpuAffinities(pid uint) (map[uint]bool, error) {
+	mask := new(unix.CPUSet)
+	mask.Zero()
+	err := unix.SchedGetaffinity(int(pid), mask)
+	if err != nil {
+		return nil, err
+	}
+	var res = make(map[uint]bool)
+	for i := 0; i < runtime.NumCPU(); i++ {
+		res[uint(i)] = mask.IsSet(i)
+	}
+	return res, nil
+}
+
+func (CpuAffinityDeputy) IsImplemented() bool {
+	return true
+}