Merge pull request #2 from Tomilla/feat_support_deadline_for_eat_process

Feat support deadline for eat process

README.md

@@ -6,6 +6,8 @@ Developer will encounter the need to quickly occupy CPU and memory, I am also de
 # Todo
 
 - [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 eat progress. such as "300ms", "1.5h", "2h45m". (unit: "ns", "us" (or "µs"), "ms", "s", "m", "h")
 - [] CPU Affinity
 - [] Memory read/write, prevent memory from being swapped out
 - [] Dynamic adjustment of CPU and memory usage
@@ -24,6 +26,7 @@ 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
 ```
 
 > Tips:
@@ -43,6 +46,8 @@ 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亲和性
 - [] 内存读写，防止内存被交换出去
 - [] 动态调整CPU和内存使用
@@ -61,6 +66,7 @@ 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核并在一小时后退出
 ```
 
 > 提示：

cmd/cpu.go

@@ -1,37 +1,75 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
+	"log"
+	"math"
 	"runtime"
 	"sync"
 	"time"
 )
 
-func busyWork() {
+const interval = 10000
+
+func busyWork(ctx context.Context) {
+	cnt := 0
 	for {
-		_ = 1 + 1
+		cnt += 1
+		if cnt%interval == 0 {
+			cnt = 0
+			select {
+			case <-ctx.Done():
+				log.Printf("busyWork: quit due to context be cancelled")
+				return
+			default:
+			}
+		}
 	}
 }
 
-func partialBusyWork(ratio float64) {
-	busyDuration := time.Duration(ratio*10) * time.Millisecond
-	idleDuration := time.Duration((1-ratio)*10) * time.Millisecond
-
+func partialBusyWork(ctx context.Context, ratio float64) {
+	const (
+		oneCycle  = 10 * time.Microsecond
+		precision = 1000
+	)
+	// round busy and idle percent
+	// case 1: ratio 0.8
+	//   busy 0.8                     idle 0.19999999999999996
+	//   busyRound 8ms                idleRound 2ms
+	//
+	// case 2: ratio 0.2
+	//   busy 0.16000000000000014     idle 0.8399999999999999
+	//   buseRound 1.6ms              idleRound 8.4ms
+	busyDuration := time.Duration(math.Floor(ratio*precision)) * oneCycle
+	idleDuration := time.Duration(math.Ceil((1-ratio)*precision)) * oneCycle
+	cnt := 0
 	for {
-		start := time.Now()
-		for time.Since(start) < busyDuration {
-			_ = 1 + 1
+		// Busy period
+		busyStart := time.Now()
+		for time.Since(busyStart) < busyDuration {
+			cnt += 1 // Simulate work
+			if cnt%interval == 0 {
+				cnt = 0
+				select {
+				case <-ctx.Done():
+					log.Printf("partialBusyWork: quit due to context being cancelled")
+					return
+				default:
+					//
+				}
+			}
 		}
+		// Idle period
 		time.Sleep(idleDuration)
 	}
 }
 
-func eatCPU(c float64) {
-	fmt.Printf("Eating CPU...          ")
+func eatCPU(ctx context.Context, wg *sync.WaitGroup, c float64) {
+	fmt.Printf("Eating %-12s", "CPU...")
 
 	runtime.GOMAXPROCS(runtime.NumCPU())
 
-	var wg sync.WaitGroup
 	fullCores := int(c)
 	partialCoreRatio := c - float64(fullCores)
 
@@ -40,7 +78,7 @@ func eatCPU(c float64) {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			busyWork()
+			busyWork(ctx)
 		}()
 	}
 
@@ -49,9 +87,9 @@ func eatCPU(c float64) {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			partialBusyWork(partialCoreRatio)
+			partialBusyWork(ctx, partialCoreRatio)
 		}()
 	}
 
-	fmt.Printf("Ate %2.1f CPU cores\n", c)
+	fmt.Printf("Ate %2.3f CPU cores\n", c)
 }

cmd/memory.go

@@ -1,14 +1,16 @@
 package cmd
 
-import "fmt"
+import (
+	"fmt"
+)
 
 func eatMemory(memoryBytes uint64) {
 	if memoryBytes == 0 {
 		return
 	}
 
 	memoryBlock := make([]byte, memoryBytes)
-	fmt.Printf("Eating memory...          ")
+	fmt.Printf("Eating %-12s", "memory...")
 	for i := range memoryBlock {
 		memoryBlock[i] = byte(i % 256)
 	}

cmd/parse.go

@@ -2,9 +2,11 @@ package cmd
 
 import (
 	"fmt"
-	"github.com/pbnjay/memory"
 	"runtime"
 	"strconv"
+	"time"
+
+	"github.com/pbnjay/memory"
 )
 
 func parseEatCPUCount(c string) float64 {
@@ -67,3 +69,14 @@ func parseEatMemoryBytes(m string) uint64 {
 	}
 	return 0
 }
+
+func parseEatDeadline(eta string) time.Duration {
+	duration, err := time.ParseDuration(eta)
+	if err != nil {
+		return time.Duration(0)
+	}
+	if duration <= 0 {
+		return time.Duration(0)
+	}
+	return duration
+}

cmd/root.go

@@ -1,15 +1,23 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
-	"github.com/pbnjay/memory"
-	"github.com/spf13/cobra"
+	"log"
 	"os"
 	"os/signal"
 	"runtime"
+	"strings"
+	"sync"
 	"syscall"
+	"time"
+
+	"github.com/pbnjay/memory"
+	"github.com/spf13/cobra"
 )
 
+const SleepDurationEachIteration = 100 * time.Millisecond
+
 func getCPUAndMemory() (uint64, uint64) {
 	cpuCount := uint64(runtime.NumCPU())
 	memoryBytes := memory.TotalMemory()
@@ -22,31 +30,62 @@ var RootCmd = &cobra.Command{
 	Run:   eatFunction,
 }
 
-func waitForever() {
-	sigs := make(chan os.Signal, 1)
+func getConsoleHelpTips(deadline time.Duration) string {
+	var helpTips = []string{"Press Ctrl + C to exit"}
+	if deadline > 0 {
+		eta := time.Now().Add(deadline)
+		helpTips = append(helpTips, fmt.Sprintf("or wait it util deadline %s", eta.Format(time.DateTime)))
+	}
+	helpTips = append(helpTips, "...")
+	return strings.Join(helpTips, " ")
+}
 
+func gracefulExit(ctx context.Context, ctxCancel context.CancelFunc) {
+	sigs := make(chan os.Signal, 1)
 	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case sig := <-sigs:
+			log.Printf("\nReceive exit signal: %v\n", sig)
+			ctxCancel()
+			return
+		default:
+			time.Sleep(SleepDurationEachIteration)
+		}
+	}
+}
 
-	done := make(chan bool, 1)
-
-	go func() {
-		sig := <-sigs
-		fmt.Println()
-		fmt.Println(sig)
-		done <- true
-	}()
+func waitUtil(ctx context.Context, wg *sync.WaitGroup, ctxCancel context.CancelFunc, deadline time.Duration) {
+	log.Println(getConsoleHelpTips(deadline))
+	gracefulExit(ctx, ctxCancel)
+	wg.Wait()
+}
 
-	fmt.Println("Press Ctrl + C to exit...")
-	<-done
+func getRootContext(dlEat time.Duration) (context.Context, context.CancelFunc) {
+	var (
+		cancel  context.CancelFunc
+		rootCtx context.Context
+	)
+	if dlEat > 0 {
+		deadline := time.Now().Add(dlEat)
+		rootCtx, cancel = context.WithDeadline(context.Background(), deadline)
+	} else {
+		rootCtx = context.Background()
+		cancel = func() {}
+	}
+	return rootCtx, cancel
 }
 
-func eatFunction(cmd *cobra.Command, args []string) {
+func eatFunction(cmd *cobra.Command, _ []string) {
 	cpuCount, memoryBytes := getCPUAndMemory()
 	fmt.Printf("Have %dC%dG.\n", cpuCount, memoryBytes/1024/1024/1024)
 
 	// Get the flags
 	c, _ := cmd.Flags().GetString("cpu_usage")
 	m, _ := cmd.Flags().GetString("memory_usage")
+	dl, _ := cmd.Flags().GetString("time_deadline")
 
 	if c == "0" && m == "0m" {
 		fmt.Println("Error: no cpu or memory usage specified")
@@ -55,9 +94,12 @@ func eatFunction(cmd *cobra.Command, args []string) {
 
 	cEat := parseEatCPUCount(c)
 	mEat := parseEatMemoryBytes(m)
+	dlEat := parseEatDeadline(dl)
 
-	fmt.Printf("Want to eat %2.1fCPU, %s Memory\n", cEat, m)
+	rootCtx, cancel := getRootContext(dlEat)
+	var wg sync.WaitGroup
+	fmt.Printf("Want to eat %2.3fCPU, %s Memory\n", cEat, m)
 	eatMemory(mEat)
-	eatCPU(cEat)
-	waitForever()
+	eatCPU(rootCtx, &wg, cEat)
+	waitUtil(rootCtx, &wg, cancel, dlEat)
 }

main.go

@@ -1,22 +1,20 @@
 package main
 
 import (
-	"eat/cmd"
 	"fmt"
 	"os"
-)
 
-var (
-	c string // how many cpu would you want eat
-	m string // how many memory would you want eat
+	"eat/cmd"
 )
 
 func main() {
 	rootCmd := cmd.RootCmd
 
 	// Add global flags
-	rootCmd.PersistentFlags().StringVarP(&c, "cpu_usage", "c", "0", "How many cpu would you want eat")
-	rootCmd.PersistentFlags().StringVarP(&m, "memory_usage", "m", "0m", "How many memory would you want eat(GB)")
+	rootCmd.PersistentFlags().StringP("cpu_usage", "c", "0", "How many cpu would you want eat")
+	rootCmd.PersistentFlags().StringP("memory_usage", "m", "0m", "How many memory would you want eat(GB)")
+	// such as "300ms", "1.5h", "2h45m". (unit: "ns", "us" (or "µs"), "ms", "s", "m", "h")
+	rootCmd.PersistentFlags().StringP("time_deadline", "t", "0", "deadline to quit eat process")
 
 	if err := rootCmd.Execute(); err != nil {
 		fmt.Println(err)