neonvm: apply code review fixes

Change cpuscaling logic to work with collection of IDs

Signed-off-by: Misha Sakhnov <[email protected]>

pkg/neonvm/cpuscaling/cpuscaler.go

@@ -1,12 +1,14 @@
 package cpuscaling
 
-import "fmt"
+import (
+	"errors"
+	"slices"
+)
 
 type CPUStater interface {
-	PossibleCPUs() (start int, end int, err error)
-	ActiveCPUsCount() (int, error)
-	SetState(cpuNum int, cpuState cpuState) error
-	GetState(cpuNum int) (cpuState, error)
+	OnlineCPUs() ([]int, error)
+	OfflineCPUs() ([]int, error)
+	SetState(cpuID int, cpuState cpuState) error
 }
 
 type cpuState string
@@ -27,72 +29,63 @@ func NewCPUScaler() *CPUScaler {
 }
 
 func (c *CPUScaler) ReconcileOnlineCPU(targetCount int) error {
-	onlineCount, err := c.cpuState.ActiveCPUsCount()
+	online, err := c.cpuState.OnlineCPUs()
 	if err != nil {
 		return err
 	}
 
-	targetCpuStateToSet := cpuOnline
-	if onlineCount < targetCount {
-		targetCpuStateToSet = cpuOnline
-	} else if onlineCount > targetCount {
-		targetCpuStateToSet = cpuOffline
+	if len(online) == targetCount {
+		return nil
 	}
-	return c.reconcileToState(int(onlineCount), targetCount, targetCpuStateToSet)
-}
 
-func (c *CPUScaler) reconcileToState(onlineCount int, targetCount int, targetState cpuState) error {
-	fistCPU, lastCPU, err := c.cpuState.PossibleCPUs()
-	if err != nil {
-		return err
-	}
+	if len(online) > targetCount {
+		diff := len(online) - targetCount
+		// offline 'diff' CPUs that are currently online
+		// reverse online slice so that we offline in the reverse order of onlining.
+		slices.Reverse(online)
+		return c.setStateTo(cpuOffline, diff, online)
 
-	if fistCPU == lastCPU {
-		// we can't scale only one CPU
-		// so we return early
-		return fmt.Errorf("failed to scale: only single CPU is available")
+	} else if len(online) < targetCount {
+		offline, err := c.cpuState.OfflineCPUs()
+		if err != nil {
+			return nil
+		}
+
+		diff := targetCount - len(online)
+		// online 'diff' CPUs that are currently offline
+		return c.setStateTo(cpuOnline, diff, offline)
 	}
 
-	for cpu := fistCPU; cpu <= lastCPU; cpu++ {
+	return nil
+}
 
-		// Skip CPU 0 as it is always online and can't be offed
-		if cpu == 0 && targetState == cpuOffline {
+func (c *CPUScaler) setStateTo(state cpuState, count int, candidateCPUs []int) error {
+	for _, cpuID := range candidateCPUs {
+		if cpuID == 0 {
+			// Not allowed to change the status of CPU 0
 			continue
 		}
 
-		cpuState, err := c.cpuState.GetState(cpu)
-		if err != nil {
+		if err := c.cpuState.SetState(cpuID, state); err != nil {
 			return err
 		}
 
-		if cpuState != targetState {
-			// mark cpu with targetState
-			err := c.cpuState.SetState(cpu, targetState)
-			if err != nil {
-				return err
-			}
-
-			// update counter
-			if targetState == cpuOnline {
-				onlineCount++
-			} else {
-				onlineCount--
-			}
-		}
-
-		// Stop when we reach the target count
-		if onlineCount == targetCount {
-			break
+		count -= 1
+		// nothing left to do
+		if count <= 0 {
+			return nil
 		}
 	}
 
-	if onlineCount != targetCount {
-		return fmt.Errorf("failed to ensure %d CPUs are online, current online CPUs: %d", targetCount, onlineCount)
-	}
-
-	return nil
+	// Got through the entire list but didn't change the state of enough CPUs
+	return errors.New("could not change the state of enough CPUs")
 }
 
+// ActiveCPUsCount() returns the count of online CPUs.
 func (c *CPUScaler) ActiveCPUsCount() (int, error) {
-	return c.cpuState.ActiveCPUsCount()
+	onlineCPUs, err := c.cpuState.OnlineCPUs()
+	if err != nil {
+		return 0, err
+	}
+	return len(onlineCPUs), nil
 }

pkg/neonvm/cpuscaling/cpuscaler_test.go

@@ -1,7 +1,6 @@
 package cpuscaling
 
 import (
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -24,33 +23,29 @@ func NewMockState(size int) *MockState {
 	return m
 }
 
-func (m *MockState) GetState(cpuNum int) (cpuState, error) {
-	if cpuNum < 0 || cpuNum >= len(m.state) {
-		return cpuOffline, os.ErrNotExist
+func (m *MockState) OnlineCPUs() ([]int, error) {
+	result := make([]int, 0, len(m.state))
+	for i, state := range m.state {
+		if state == cpuOnline {
+			result = append(result, i)
+		}
 	}
-	return m.state[cpuNum], nil
+	return result, nil
 }
 
-func (m *MockState) SetState(cpuNum int, state cpuState) error {
-	if cpuNum < 0 || cpuNum >= len(m.state) {
-		return os.ErrNotExist
+func (m *MockState) OfflineCPUs() ([]int, error) {
+	result := make([]int, 0, len(m.state))
+	for i, state := range m.state {
+		if state == cpuOffline {
+			result = append(result, i)
+		}
 	}
-	m.state[cpuNum] = state
-	return nil
-}
-
-func (cs *MockState) PossibleCPUs() (int, int, error) {
-	return 0, len(cs.state) - 1, nil
+	return result, nil
 }
 
-func (cs *MockState) ActiveCPUsCount() (int, error) {
-	count := 0
-	for _, state := range cs.state {
-		if state == cpuOnline {
-			count++
-		}
-	}
-	return count, nil
+func (m *MockState) SetState(cpuID int, cpuState cpuState) error {
+	m.state[cpuID] = cpuState
+	return nil
 }
 
 func TestReconcileCPU(t *testing.T) {
@@ -66,7 +61,7 @@ func TestReconcileCPU(t *testing.T) {
 		assert.NoError(t, scaler.ReconcileOnlineCPU(3))
 		assertActiveCPUsCount(t, scaler, 3)
 
-		// Scale down
+		// // Scale down
 		assert.NoError(t, scaler.ReconcileOnlineCPU(2))
 		assertActiveCPUsCount(t, scaler, 2)
 	})
@@ -116,7 +111,8 @@ func TestReconcileCPU(t *testing.T) {
 }
 
 func assertActiveCPUsCount(t *testing.T, scaler *CPUScaler, n int) {
-	activeCPUs, err := scaler.cpuState.ActiveCPUsCount()
+	t.Helper() // to tell the test suite that this is a helper method to correctly render the line number
+	onlineCPUs, err := scaler.cpuState.OnlineCPUs()
 	assert.NoError(t, err)
-	assert.Equal(t, n, activeCPUs)
+	assert.Equal(t, n, len(onlineCPUs))
 }

pkg/neonvm/cpuscaling/sysfs.go

@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
-	"slices"
 	"strconv"
 	"strings"
 )
@@ -31,43 +30,32 @@ func (cs *cpuSysfsState) SetState(cpuNum int, cpuState cpuState) error {
 	return nil
 }
 
-func (cs *cpuSysfsState) GetState(cpuNum int) (cpuState, error) {
-	onlineCPUs, err := cs.getAllOnlineCPUs()
-	if err != nil {
-		return cpuOffline, err
-	}
-	if slices.Contains(onlineCPUs, cpuNum) {
-		return cpuOnline, nil
-	}
-
-	return cpuOffline, nil
-}
-
-func (cs *cpuSysfsState) getAllOnlineCPUs() ([]int, error) {
+func (cs *cpuSysfsState) OnlineCPUs() ([]int, error) {
 	data, err := os.ReadFile(filepath.Join(cpuPath, "online"))
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to read online CPUs: %w", err)
 	}
-
-	onlineCPUs, err := cs.parseMultipleCPURange(string(data))
+	cpuIDs, err := cs.parseMultipleCPURange(string(data))
 	if err != nil {
-		return onlineCPUs, err
+		// log value of the file in case we can't parse to help debugging
+		return nil, fmt.Errorf("failed to parse online CPUs %q: %w", string(data), err)
 	}
-
-	return onlineCPUs, nil
+	return cpuIDs, nil
 }
 
-// PossibleCPUs returns the start and end indexes of all possible CPUs.
-func (cs *cpuSysfsState) PossibleCPUs() (int, int, error) {
-	data, err := os.ReadFile(filepath.Join(cpuPath, "possible"))
+func (cs *cpuSysfsState) OfflineCPUs() ([]int, error) {
+	data, err := os.ReadFile(filepath.Join(cpuPath, "offline"))
 	if err != nil {
-		return -1, -1, err
+		return nil, fmt.Errorf("failed to read offline CPUs: %w", err)
 	}
-
-	return cs.parseCPURange(string(data))
+	cpuIDs, err := cs.parseMultipleCPURange(string(data))
+	if err != nil {
+		// log value of the file in case we can't parse to help debugging
+		return nil, fmt.Errorf("failed to parse offline CPUs %q: %w", string(data), err)
+	}
+	return cpuIDs, nil
 }
 
-// parseCPURange parses the CPU range string (e.g., "0-3") and returns start and end indexes.
 func (cs *cpuSysfsState) parseCPURange(cpuRange string) (int, int, error) {
 	cpuRange = strings.TrimSpace(cpuRange)
 	parts := strings.Split(cpuRange, "-")
@@ -112,24 +100,3 @@ func (cs *cpuSysfsState) parseMultipleCPURange(cpuRanges string) ([]int, error)
 
 	return cpus, nil
 }
-
-// ActiveCPUsCount() returns the count of online CPUs.
-func (c *cpuSysfsState) ActiveCPUsCount() (int, error) {
-	start, end, err := c.PossibleCPUs()
-	if err != nil {
-		return 0, err
-	}
-
-	var onlineCount int
-	for cpu := start; cpu <= end; cpu++ {
-		state, err := c.GetState(cpu)
-		if err != nil {
-			return 0, err
-		}
-		if state == cpuOnline {
-			onlineCount++
-		}
-	}
-
-	return onlineCount, nil
-}