controller.go

package controller

import (
	"context"
	"fmt"
	"log"
	"node-controller/cgroup"
	"node-controller/config"
	"node-controller/metrics"
	k8sPods "node-controller/pods"
	rl "node-controller/rl-client"
	"node-controller/safety"
	"runtime/debug"
	"time"

	"k8s.io/client-go/kubernetes"
)

type Controller struct {
	clientset     *kubernetes.Clientset
	nodeName      string
	podWatcher    *k8sPods.PodWatcher
	cgroupManager *cgroup.Manager
	metricsCol    *metrics.Collector
	rlClient      *rl.ContextualBandit
	rewardCalc    *rl.RewardCalculator
	// rlClient      *rlclient.GRPCClient (not using for now as rl is implemented on golang itself)

	safety     *safety.Constraints
	rollback   *safety.RollbackManager
	config     *config.Config
	lastAction *int32 // for storing the last action

	// Observer for metrics collection and storage
	metricsObserver metrics.Observer
	iterationCount  int64     // Track iteration number for metrics
	iterStartTime   time.Time // Track start time of current iteration
}

// Public constructor for Controller
func NewController(
	clientset *kubernetes.Clientset,
	nodeName string,
	podWatcher *k8sPods.PodWatcher,
	cgroupManager *cgroup.Manager,
	metricsCol *metrics.Collector,
	rlClient *rl.ContextualBandit,
	rewardCalc *rl.RewardCalculator,
	safetyConstraints *safety.Constraints,
	rollbackMgr *safety.RollbackManager,
	cfg *config.Config,
) *Controller {
	return &Controller{
		clientset:      clientset,
		nodeName:       nodeName,
		podWatcher:     podWatcher,
		cgroupManager:  cgroupManager,
		metricsCol:     metricsCol,
		rlClient:       rlClient,
		rewardCalc:     rewardCalc,
		safety:         safetyConstraints,
		rollback:       rollbackMgr,
		config:         cfg,
		iterationCount: 0,
		iterStartTime:  time.Now(),
	}
}

// SetMetricsObserver sets the observer for metrics collection
func (c *Controller) SetMetricsObserver(observer metrics.Observer) {
	c.metricsObserver = observer
}

func (c *Controller) run_control_loop(ctx context.Context) error {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("[run_control_iteration] PANIC recovered: %v\n", r)
			debug.PrintStack()
		}
	}()

	// Initialize the components
	c.podWatcher.Start(ctx)

	// This sets up the metrics collector i.e sets up podClassifier
	if err := c.metricsCol.Start(ctx); err != nil {
		return fmt.Errorf("Error in starting metrics collector")
	}

	// Initialize cgroup paths for existing pods on the node
	initialPods := c.podWatcher.GetPods()
	c.cgroupManager.SetPodCgroupPath(initialPods)

	ticker := time.NewTicker(c.config.ControlInterval)
	defer ticker.Stop()

	//infinitely loop for error
	for {
		//switch statement for channels
		select {
		case <-ticker.C: //when ticker is triggered
			if err := c.run_control_iteration(ctx); err != nil {
				return fmt.Errorf("Control loop error: %v", err)
			}
		case <-ctx.Done(): //exit channel
			return ctx.Err()
		}
	}
}

// Run is a public wrapper to start the control loop
func (c *Controller) Run(ctx context.Context) error {
	return c.run_control_loop(ctx)
}

func (c *Controller) executeAction(action rl.Action, state rl.State) error {
	log.Printf("[executeAction] Starting to execute action: %v\n", action.Type)

	// Read actual current weights from cgroups (not from state which may be stale/0)
	currentWeights := make(map[string]int)

	criticalWeight := 100
	batchWeight := 100
	defaultWeight := 1000

	// If state has non-zero weights, use them; otherwise use defaults
	if state.WeightCritical > 0 {
		w := int(state.WeightCritical * 10000)
		if w > 1 && w < 10000 {
			criticalWeight = w
		}
	}
	if state.WeightBatch > 0 {
		w := int(state.WeightBatch * 10000)
		if w > 1 && w < 10000 {
			batchWeight = w
		}
	}

	currentWeights["latency-critical"] = criticalWeight
	currentWeights["batch"] = batchWeight
	currentWeights["default"] = defaultWeight

	log.Printf("[executeAction] Current weights (from state or defaults): %v\n", currentWeights)

	// Apply the action to get new weights using the Action.Apply() method
	newWeights := action.Apply(currentWeights)
	if newWeights == nil {
		log.Printf("[executeAction] ERROR: action.Apply() returned nil map for action type %v\n", action.Type)
		return fmt.Errorf("action.Apply() returned nil map")
	}
	log.Printf("[executeAction] Applying action %v: %v -> %v\n", action.Type, currentWeights, newWeights)

	// Get only latency-critical pods using the classifier
	criticalPods := c.metricsCol.GetClassifier().GetPodsInClass(k8sPods.LatencyCritical)
	log.Printf("[executeAction] Found %d latency-critical pods to update\n", len(criticalPods))

	// Apply the new CPU weights to latency-critical pods only
	for _, pod := range criticalPods {
		cgroupPath, err := c.cgroupManager.GetPodCgroupPath(string(pod.UID))
		if err != nil {
			log.Printf("[executeAction] WARNING: Could not get cgroup path for pod %s: %v\n", pod.Name, err)
			continue
		}

		// Set the new weight for this pod (latency-critical class gets the adjusted weight)
		newWeight := newWeights["latency-critical"]
		if err := c.cgroupManager.SetCgroupWeight(cgroupPath, newWeight); err != nil {
			log.Printf("[executeAction] ERROR: Failed to set weight for pod %s: %v\n", pod.Name, err)
			continue
		}

		log.Printf("[executeAction] Successfully set CPU weight %d for pod %s\n", newWeight, pod.Name)
	}

	// Store the new weights in the action for reference
	action.NewWeights = newWeights
	log.Printf("[executeAction] Action %v executed successfully\n", action.Type)

	return nil
}

// Reward calculation is done by controller
func (c *Controller) calculateReward(state rl.State, newState rl.State, action rl.Action) (float64, float64, float64) {
	return c.rewardCalc.Calculate(state, newState, action.Type)
}

func (c *Controller) run_control_iteration(ctx context.Context) error {
	iterStartTime := time.Now()

	// 1. Collect current state
	state, err := c.collectState(ctx)
	if err != nil {
		//handle collecting state error
		return fmt.Errorf("Error in collecting state: %w", err)
	}

	// 2. Query RL for action (Assuming rl-client is set up)
	// returns the ActionType
	actionType := c.rlClient.SelectAction(ctx)

	action := rl.Action{
		Type: actionType,
	}

	//3. Validate action, return error if action is not safe
	actionConstraintViolated := false
	if err := c.safety.ValidateAction(action); err != nil {
		fmt.Printf("Action rejected: %v", err)
		actionConstraintViolated = true
		// Continue with the iteration to record the constraint violation
	}

	// cgroup handling required
	//4. Execute action with current state weights
	if err := c.executeAction(action, state); err != nil {
		return fmt.Errorf("Error in executing action: %w", err)
	}

	// Save this action for next iteration's state
	actionValue := int32(action.Type)
	c.lastAction = &actionValue

	//5. Wait and collect new state feedback
	log.Printf("[run_control_iteration] Waiting %v before collecting new state...\n", c.config.WaitTime)
	select {
	case <-time.After(c.config.WaitTime):
		// Continue after waiting or timeout
		log.Println("[run_control_iteration] Wait completed, now collecting new state...")
	case <-ctx.Done():
		return ctx.Err()
	}
	log.Println("[run_control_iteration] Starting new state collection...")
	newState, err := c.collectState(ctx)
	if err != nil {
		log.Printf("[run_control_iteration] ERROR collecting new state: %v\n", err)
		return err
	}
	log.Println("[run_control_iteration] New state collected successfully")

	//6. Calculate reward
	log.Printf("[run_control_iteration] Calculating reward with state transition...\n")
	if c.rewardCalc == nil {
		log.Printf("[run_control_iteration] ERROR: rewardCalc is nil\n")
		return fmt.Errorf("rewardCalc is nil")
	}
	reward, oscillationPenalty, violation := c.calculateReward(state, newState, action)
	log.Printf("[run_control_iteration] Reward calculated: %f\n", reward)

	//6.1 check if state is good, save snapshot (for rollback)
	if newState.P99LatencyCritical < 100 && reward > 0 {
		log.Printf("[run_control_iteration] Saving snapshot for rollback (latency: %f, reward: %f)\n", newState.P99LatencyCritical, reward)
		c.rollback.SaveSnapshot(newState)
	}

	//7. Send feedback to RL
	log.Printf("[run_control_iteration] Updating RL client with action type %v and reward %f\n", action.Type, reward)
	if c.rlClient == nil {
		log.Printf("[run_control_iteration] ERROR: rlClient is nil\n")
		return fmt.Errorf("rlClient is nil")
	}
	c.rlClient.Update(ctx, action.Type, reward)
	log.Printf("[run_control_iteration] RL client updated successfully\n")

	// 8. Check for degradation
	log.Printf("[run_control_iteration] Checking if rollback is needed\n")
	if c.rollback.ShouldRollBack(newState) {
		log.Printf("[run_control_iteration] Rollback triggered!\n")
		// get most recent snapshot for recent good state
		snapshot, err := c.rollback.GetLastSnapshot()
		if err != nil {
			log.Printf("[run_control_iteration] ERROR getting snapshot: %v\n", err)
			return err
		}
		if snapshot == nil {
			log.Printf("[run_control_iteration] WARNING: snapshot is nil, cannot rollback\n")
			return nil
		}

		// Restore weights from the snapshot
		log.Printf("[run_control_iteration] Restoring weights from snapshot\n")
		for class, weight := range snapshot.Weights {
			pods := c.metricsCol.GetClassifier().GetPodsInClass(class)
			log.Printf("[run_control_iteration] Setting weight %d for class %v (%d pods)\n", weight, class, len(pods))
			c.cgroupManager.SetClassWeight(pods, weight)
		}

		if err := c.rollback.Execute(); err != nil {
			//Handle rollback
			log.Printf("[run_control_iteration] ERROR: Rollback execution failed: %v\n", err)
			return fmt.Errorf("Rollback failed: %v", err)
		}
		log.Printf("[run_control_iteration] Rollback executed successfully\n")
	}

	// 9. Notify metrics observer (if configured)
	if c.metricsObserver != nil {
		c.iterationCount++
		iterDurationNs := time.Since(iterStartTime).Nanoseconds()

		// Populate IterationMetric with all collected data
		iterMetric := &metrics.IterationMetric{
			Timestamp:                iterStartTime,
			IterationNum:             c.iterationCount,
			P99LatencyCritical:       state.P99LatencyCritical,
			P99LatencyBatch:          state.P99LatencyBatch,
			P99LatencyDefault:        0, // Not tracked separately in state
			CPUUsageCritical:         state.CPUUsageCritical,
			CPUUsageBatch:            state.CPUUsageBatch,
			WeightCritical:           int(state.WeightCritical * 10000),
			WeightBatch:              int(state.WeightBatch * 10000),
			WeightDefault:            1000,
			RunqueueLength:           state.RunqueueLength,
			ActionType:               action.Type,
			ActionDescription:        metrics.ActionTypeToString(action.Type),
			ActionConstraintViolated: actionConstraintViolated,
			TotalReward:              reward,
			IterationDurationNs:      iterDurationNs,
			LoadAverage:              state.RunqueueLength, // Approximation from runqueue
			ActivePodsCount: len(c.metricsCol.GetClassifier().GetPodsInClass(k8sPods.LatencyCritical)) +
				len(c.metricsCol.GetClassifier().GetPodsInClass(k8sPods.Batch)) +
				len(c.metricsCol.GetClassifier().GetPodsInClass(k8sPods.Default)),
			ErrorOccurred: false,
			ErrorMessage:  "",
		}

		// Populate CgroupChanges if action was executed
		if len(action.NewWeights) > 0 {
			iterMetric.CgroupChanges = &metrics.CgroupChanges{
				CriticalOldWeight: int(state.WeightCritical * 10000),
				CriticalNewWeight: action.NewWeights["latency-critical"],
				CriticalChanged:   action.NewWeights["latency-critical"] != int(state.WeightCritical*10000),
				BatchOldWeight:    int(state.WeightBatch * 10000),
				BatchNewWeight:    action.NewWeights["batch"],
				BatchChanged:      action.NewWeights["batch"] != int(state.WeightBatch*10000),
			}
		}

		// Populate RewardComponents from reward calculator
		iterMetric.RewardComponents = &metrics.RewardComponents{
			Reward:             reward,
			ConstraintPenalty:  oscillationPenalty,
			OscillationPenalty: violation,
			PreviousP99Latency: state.P99LatencyCritical,
			LatencyImprovement: state.P99LatencyCritical - newState.P99LatencyCritical,
		}

		// Notify observer (non-blocking)
		c.metricsObserver.OnIterationComplete(iterMetric)
	}

	return nil
}