Fix performance (#7)

* Update main.go

Signed-off-by: Md Imran <[email protected]>

* Update main.go

Signed-off-by: Md Imran <[email protected]>

* Update README.md

Signed-off-by: Md Imran <[email protected]>

---------

Signed-off-by: Md Imran <[email protected]>

Author: narmidm

README.md

@@ -34,7 +34,7 @@ Build the Docker image using the provided Dockerfile:
 
    ```shell
    docker build -t k8s-pod-cpu-stressor .
-  ```
+   ```
 
 Run the Docker container, specifying the desired CPU usage, stress duration, and optionally whether to run CPU stress indefinitely:
 
@@ -56,6 +56,24 @@ The `k8s-pod-cpu-stressor` allows you to specify the desired CPU usage and stres
 
 Adjust these parameters according to your requirements to simulate different CPU load scenarios.
 
+### Kubernetes Resource Requests and Limits
+
+It is recommended to specify Kubernetes resource requests and limits to control the amount of CPU resources consumed by the pod, and to prevent overloading your cluster. For example:
+
+- **Requests**: This defines the minimum amount of CPU that the pod is guaranteed to have.
+- **Limits**: This defines the maximum amount of CPU that the pod can use.
+
+Adding requests and limits helps Kubernetes manage resources efficiently and ensures that your cluster remains stable during stress testing.
+
+Example:
+
+```yaml
+resources:
+  requests:
+    cpu: "100m"
+  limits:
+    cpu: "200m"
+```
 
 ## Check the Public Docker Image
 
@@ -98,10 +116,42 @@ spec:
               cpu: "100m"
 ```
 
+## Sample Job Manifest
+
+If you want to run the CPU stressor for a fixed duration as a one-time job, you can use the following Kubernetes Job manifest:
+
+```yaml
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: cpu-stressor-job
+spec:
+  template:
+    metadata:
+      labels:
+        app: cpu-stressor
+    spec:
+      containers:
+        - name: cpu-stressor
+          image: narmidm/k8s-pod-cpu-stressor:latest
+          args:
+            - "-cpu=0.5"
+            - "-duration=5m"
+          resources:
+            limits:
+              cpu: "500m"
+            requests:
+              cpu: "250m"
+      restartPolicy: Never
+  backoffLimit: 3
+```
+
+This manifest runs the `k8s-pod-cpu-stressor` as a Kubernetes Job, which will execute the stress test once for 5 minutes and then stop. The `backoffLimit` specifies the number of retries if the job fails.
+
 ## Contributing
 
 Contributions are welcome! If you find a bug or have a suggestion, please open an issue or submit a pull request. For major changes, please discuss them first in the issue tracker.
 
 ## License
 
-This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
+This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.

main.go

@@ -3,57 +3,81 @@ package main
 import (
 	"flag"
 	"fmt"
+	"math/rand"
 	"os"
 	"os/signal"
 	"runtime"
+	"sync/atomic"
 	"time"
 )
 
 func main() {
-	cpuUsagePtr := flag.Float64("cpu", 0.2, "CPU usage as a fraction (e.g., 0.2 for 200m)")
+	cpuUsagePtr := flag.Float64("cpu", 0.2, "CPU usage as a fraction (e.g., 0.2 for 20% CPU usage)")
 	durationPtr := flag.Duration("duration", 10*time.Second, "Duration for the CPU stress (e.g., 10s)")
 	runForeverPtr := flag.Bool("forever", false, "Run CPU stress indefinitely")
 	flag.Parse()
 
 	numCPU := runtime.NumCPU()
 	runtime.GOMAXPROCS(numCPU)
 
+	// Number of goroutines to use for stressing CPU
 	numGoroutines := int(float64(numCPU) * (*cpuUsagePtr))
+	if numGoroutines < 1 {
+		numGoroutines = 1
+	}
 
-	fmt.Printf("Starting CPU stress with %d goroutines...\n", numGoroutines)
+	fmt.Printf("Starting CPU stress with %d goroutines targeting %.2f CPU usage...\n", numGoroutines, *cpuUsagePtr)
 
 	done := make(chan struct{})
 
 	// Capture termination signals
 	quit := make(chan os.Signal, 1)
 	signal.Notify(quit, os.Interrupt, os.Kill)
 
+	var stopFlag int32
+
+	// Improved workload generation
 	for i := 0; i < numGoroutines; i++ {
 		go func() {
+			workDuration := time.Duration(*cpuUsagePtr*1000) * time.Microsecond
+			idleDuration := time.Duration((1-*cpuUsagePtr)*1000) * time.Microsecond
+
 			for {
-				select {
-				case <-done:
+				if atomic.LoadInt32(&stopFlag) == 1 {
 					return
-				default:
 				}
+
+				// Busy loop for the specified work duration
+				endWork := time.Now().Add(workDuration)
+				for time.Now().Before(endWork) {
+					// Perform a small computation to keep the CPU active
+					_ = rand.Float64() * rand.Float64()
+				}
+
+				// Idle for the rest of the interval
+				time.Sleep(idleDuration)
 			}
 		}()
 	}
 
 	go func() {
 		// Wait for termination signal
 		<-quit
-		fmt.Println("Termination signal received. Stopping CPU stress...")
+		fmt.Println("\nTermination signal received. Stopping CPU stress...")
+		atomic.StoreInt32(&stopFlag, 1)
 		close(done)
 	}()
 
 	if !*runForeverPtr {
 		time.Sleep(*durationPtr)
-		fmt.Println("CPU stress completed.")
-		os.Exit(0)
+		fmt.Println("\nCPU stress completed.")
+		atomic.StoreInt32(&stopFlag, 1)
+		close(done)
+		// Keep the process running to prevent the pod from restarting
+		select {}
 	}
 
 	// Run stress indefinitely
 	fmt.Println("CPU stress will run indefinitely. Press Ctrl+C to stop.")
-	select {}
+	<-done
 }