renamemain

4 files changed, 75 insertions, 21 deletions

diff --git a/internal/cpu/helpers.go b/internal/cpu/helpers.go
index 9fcc6b9..138cf8c 100644
--- a/internal/cpu/helpers.go
+++ b/internal/cpu/helpers.go
@@ -11,6 +11,16 @@ func (s *Stresser) randomIntensity() int {
 	return s.rand.Intn(s.config.IntensityMax-s.config.IntensityMin+1) + s.config.IntensityMin
 }
 
+func (s *Stresser) determineCPUCount() int {
+	if s.config.ChaosEnabled {
+		return s.rand.Intn(s.config.MaxCPUs-s.config.MinCPUs+1) + s.config.MinCPUs
+	}
+	if s.config.Workers > 0 {
+		return s.config.Workers
+	}
+	return s.config.MaxCPUs
+}
+
 func (s *Stresser) getRandomMode() Mode {
 	return allModes[s.rand.Intn(len(allModes))]
 }
diff --git a/internal/cpu/stresser.go b/internal/cpu/stresser.go
index 906467a..62ae3b7 100644
--- a/internal/cpu/stresser.go
+++ b/internal/cpu/stresser.go
@@ -13,12 +13,13 @@ import (
 // Stresser uses pointer receivers because it contains a sync.Mutex
 // and shared mutable state that must not be copied.
 type Stresser struct {
-	config  Config
-	cancel  context.CancelFunc
-	wg      sync.WaitGroup
-	mu      sync.Mutex
-	current Mode
-	rand    *rand.Rand
+	config      Config
+	cancel      context.CancelFunc
+	wg          sync.WaitGroup
+	mu          sync.Mutex
+	current     Mode
+	currentCPUs int
+	rand        *rand.Rand
 }
 
 func New(cfg Config) *Stresser {
@@ -34,6 +35,11 @@ func New(cfg Config) *Stresser {
 	if cfg.Mode == "" {
 		cfg.Mode = ModePrime
 	}
+	if !cfg.ChaosEnabled && cfg.MinCPUs <= 0 {
+		cfg.MinCPUs = cfg.Workers
+		cfg.MaxCPUs = cfg.Workers
+	}
+	cfg.normalize()
 	return &Stresser{
 		config: cfg,
 		rand:   rand.New(rand.NewSource(time.Now().UnixNano())),
@@ -44,8 +50,12 @@ func (s *Stresser) Start() {
 	ctx, cancel := context.WithCancel(context.Background())
 	s.cancel = cancel
 
-	s.wg.Add(s.config.Workers)
-	for i := 0; i < s.config.Workers; i++ {
+	s.mu.Lock()
+	s.currentCPUs = s.determineCPUCount()
+	s.mu.Unlock()
+
+	s.wg.Add(s.currentCPUs)
+	for i := 0; i < s.currentCPUs; i++ {
 		go s.worker(ctx, i)
 	}
 
diff --git a/internal/cpu/types.go b/internal/cpu/types.go
index 58da0a9..45f6ff5 100644
--- a/internal/cpu/types.go
+++ b/internal/cpu/types.go
@@ -1,7 +1,10 @@
 // Package cpu provides CPU stress testing with multiple algorithms and chaos modes.
 package cpu
 
-import "time"
+import (
+	"runtime"
+	"time"
+)
 
 // types.go defines the configuration types and mode constants for CPU stress testing.
 
@@ -18,7 +21,6 @@ const (
 	ModeEncrypt  Mode = "encrypt"
 	ModeSort     Mode = "sort"
 	ModeBusy     Mode = "busy"
-	ModeRandom   Mode = "random"
 )
 
 var allModes = []Mode{
@@ -42,4 +44,24 @@ type Config struct {
 	ChaosInterval time.Duration
 	IntensityMin  int
 	IntensityMax  int
+	MinCPUs       int
+	MaxCPUs       int
+}
+
+func (c *Config) normalize() {
+	if c.MinCPUs <= 0 {
+		c.MinCPUs = 1
+	}
+	if c.MaxCPUs <= 0 {
+		c.MaxCPUs = runtime.NumCPU()
+	}
+	if c.MinCPUs > runtime.NumCPU() {
+		c.MinCPUs = runtime.NumCPU()
+	}
+	if c.MaxCPUs > runtime.NumCPU() {
+		c.MaxCPUs = runtime.NumCPU()
+	}
+	if c.MinCPUs > c.MaxCPUs {
+		c.MinCPUs, c.MaxCPUs = c.MaxCPUs, c.MinCPUs
+	}
 }
diff --git a/internal/cpu/worker.go b/internal/cpu/worker.go
index 06f6812..b11dc9a 100644
--- a/internal/cpu/worker.go
+++ b/internal/cpu/worker.go
@@ -12,10 +12,6 @@ func (s *Stresser) worker(ctx context.Context, id int) {
 	defer s.wg.Done()
 
 	mode := s.config.Mode
-	if s.config.ChaosEnabled {
-		mode = s.getRandomMode()
-	}
-
 	intensity := s.randomIntensity()
 
 	ticker := time.NewTicker(10 * time.Millisecond)
@@ -31,13 +27,10 @@ func (s *Stresser) worker(ctx context.Context, id int) {
 				mode = s.current
 				s.mu.Unlock()
 				intensity = s.randomIntensity()
+				s.sleepWithJitter(1, 50)
 			}
 
 			s.executeStress(ctx, mode, intensity)
-
-			if s.config.ChaosEnabled {
-				s.sleepWithJitter(1, 50)
-			}
 		}
 	}
 }
@@ -53,11 +46,33 @@ func (s *Stresser) chaosController(ctx context.Context) {
 		case <-ticker.C:
 			s.mu.Lock()
 			s.current = s.getRandomMode()
+			newCPUCount := s.determineCPUCount()
+			if newCPUCount != s.currentCPUs {
+				s.adjustWorkers(ctx, newCPUCount)
+			}
 			s.mu.Unlock()
 		}
 	}
 }
 
+func (s *Stresser) adjustWorkers(ctx context.Context, newCount int) {
+	if newCount > s.currentCPUs {
+		for i := s.currentCPUs; i < newCount; i++ {
+			s.wg.Add(1)
+			go s.worker(ctx, i)
+		}
+	} else if newCount < s.currentCPUs {
+		s.cancel()
+		s.wg.Wait()
+		ctx, s.cancel = context.WithCancel(context.Background())
+		for i := 0; i < newCount; i++ {
+			s.wg.Add(1)
+			go s.worker(ctx, i)
+		}
+	}
+	s.currentCPUs = newCount
+}
+
 func (s *Stresser) executeStress(ctx context.Context, mode Mode, intensity int) {
 	switch mode {
 	case ModePrime:
@@ -80,8 +95,5 @@ func (s *Stresser) executeStress(ctx context.Context, mode Mode, intensity int)
 		s.stressSort(ctx, intensity)
 	case ModeBusy:
 		s.stressBusy(ctx, intensity)
-	case ModeRandom:
-		randomMode := allModes[s.rand.Intn(len(allModes))]
-		s.executeStress(ctx, randomMode, intensity)
 	}
 }