package main import ( "context" "encoding/json" "fmt" "log" "net/http" "os" "strconv" "sync" "time" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/client-go/kubernetes" "k8s.io/client-go/rest" "k8s.io/apimachinery/pkg/types" metricsv "k8s.io/metrics/pkg/client/clientset/versioned" ) var ( CPUOverloadThreshold int64 MemoryOverloadThreshold int64 CPUUnderloadThreshold int64 MemoryUnderloadThreshold int64 PodEvictionInterval time.Duration PodAnimationDuration time.Duration PodRespawnTimeout time.Duration ) func getEnvInt64(name string, defaultVal int64) int64 { valStr := os.Getenv(name) if valStr == "" { return defaultVal } val, err := strconv.ParseInt(valStr, 10, 64) if err != nil { log.Printf("Invalid %s: %s, using default %d", name, valStr, defaultVal) return defaultVal } return val } func getEnvSecondsDuration(name string, defaultSec int64) time.Duration { valStr := os.Getenv(name) if valStr == "" { return time.Duration(defaultSec) * time.Second } sec, err := strconv.ParseInt(valStr, 10, 64) if err != nil { log.Printf("Invalid %s: %s, using default %d seconds", name, valStr, defaultSec) return time.Duration(defaultSec) * time.Second } return time.Duration(sec) * time.Second } func init() { CPUOverloadThreshold = getEnvInt64("CPU_OVERLOAD_THRESHOLD", 80) MemoryOverloadThreshold = getEnvInt64("MEMORY_OVERLOAD_THRESHOLD", 80) CPUUnderloadThreshold = getEnvInt64("CPU_UNDERLOAD_THRESHOLD", 50) MemoryUnderloadThreshold = getEnvInt64("MEMORY_UNDERLOAD_THRESHOLD", 50) PodEvictionInterval = getEnvSecondsDuration("POD_EVICTION_INTERVAL", 30) // default 2 min PodAnimationDuration = getEnvSecondsDuration("POD_ANIMATION_DURATION", 2) // default 2 sec func init() { CPUOverloadThreshold = getEnvInt64("CPU_OVERLOAD_THRESHOLD", 80) MemoryOverloadThreshold = getEnvInt64("MEMORY_OVERLOAD_THRESHOLD", 80) CPUUnderloadThreshold = getEnvInt64("CPU_UNDERLOAD_THRESHOLD", 50) MemoryUnderloadThreshold = getEnvInt64("MEMORY_UNDERLOAD_THRESHOLD", 50) PodEvictionInterval = getEnvSecondsDuration("POD_EVICTION_INTERVAL", 30) PodAnimationDuration = getEnvSecondsDuration("POD_ANIMATION_DURATION", 2) PodRespawnTimeout = getEnvSecondsDuration("POD_RESPAWN_TIMEOUT", 60) } type PodMoveEvent struct { Time time.Time `json:"time"` Namespace string `json:"namespace"` Pod string `json:"pod"` FromNode string `json:"fromNode"` ToNode string `json:"toNode"` Reason string `json:"reason"` } type EventBus struct { mu sync.Mutex history []PodMoveEvent clients map[chan PodMoveEvent]struct{} } func NewEventBus() *EventBus { return &EventBus{ history: make([]PodMoveEvent, 0, 100), clients: make(map[chan PodMoveEvent]struct{}), } } func (b *EventBus) Emit(evt PodMoveEvent) { b.mu.Lock() defer b.mu.Unlock() if len(b.history) == 100 { b.history = b.history[1:] } b.history = append(b.history, evt) for ch := range b.clients { select { case ch <- evt: default: } } } func (b *EventBus) History() []PodMoveEvent { b.mu.Lock() defer b.mu.Unlock() out := make([]PodMoveEvent, len(b.history)) copy(out, b.history) return out } func (b *EventBus) Subscribe() chan PodMoveEvent { b.mu.Lock() defer b.mu.Unlock() ch := make(chan PodMoveEvent, 10) b.clients[ch] = struct{}{} return ch } func (b *EventBus) Unsubscribe(ch chan PodMoveEvent) { b.mu.Lock() defer b.mu.Unlock() delete(b.clients, ch) close(ch) } func main() { config, err := rest.InClusterConfig() if err != nil { log.Fatal(err) } clientset, err := kubernetes.NewForConfig(config) if err != nil { log.Fatal(err) } metricsClient, err := metricsv.NewForConfig(config) if err != nil { log.Fatal(err) } eventBus := NewEventBus() go startWebUI(eventBus, clientset) for { balanceNodes(clientset, metricsClient, eventBus) time.Sleep(PodEvictionInterval) } } func startWebUI(bus *EventBus, clientset *kubernetes.Clientset) { mux := http.NewServeMux() mux.Handle("/", http.FileServer(http.Dir("/web"))) mux.HandleFunc("/api/events", func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode(bus.History()) }) mux.HandleFunc("/api/stream", func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Cache-Control", "no-cache") flusher, ok := w.(http.Flusher) if !ok { http.Error(w, "streaming unsupported", 500) return } ch := bus.Subscribe() defer bus.Unsubscribe(ch) for { select { case evt := <-ch: data, _ := json.Marshal(evt) fmt.Fprintf(w, "data: %s\n\n", data) flusher.Flush() case <-r.Context().Done(): return } } }) log.Println("WebUI listening on :8080") mux.HandleFunc("/api/nodes", func(w http.ResponseWriter, r *http.Request) { nodes, err := clientset.CoreV1().Nodes().List(context.Background(), metav1.ListOptions{}) if err != nil { http.Error(w, "failed to list nodes", http.StatusInternalServerError) return } names := []string{} for _, n := range nodes.Items { names = append(names, n.Name) } w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode(names) }) mux.HandleFunc("/api/pods", func(w http.ResponseWriter, r *http.Request) { ctx := context.Background() pods, err := clientset.CoreV1().Pods("").List(ctx, metav1.ListOptions{}) if err != nil { http.Error(w, "failed to list pods", http.StatusInternalServerError) return } // group pods by node podsByNode := map[string][]string{} for _, pod := range pods.Items { if pod.Spec.NodeName == "" { continue } if _, ok := podsByNode[pod.Spec.NodeName]; !ok { podsByNode[pod.Spec.NodeName] = []string{} } podsByNode[pod.Spec.NodeName] = append(podsByNode[pod.Spec.NodeName], pod.Name) } w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode(podsByNode) }) mux.HandleFunc("/api/config", func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") cfg := map[string]interface{}{ "animationDurationMs": int64(2000), // number "version": "1.1", // string "clusterName": os.Getenv("CLUSTER_NAME"), // string } _ = json.NewEncoder(w).Encode(cfg) }) _ = http.ListenAndServe(":8080", mux) } // // ─── balance nodes ──────────────────────────────────────────────────────────────────── // func balanceNodes(clientset *kubernetes.Clientset) { log.Println("Starting balance cycle") nodes, err := clientset.CoreV1().Nodes().List(context.Background(), metav1.ListOptions{}) if err != nil { log.Printf("Error listing nodes: %v", err) return } podCounts := make(map[string]int) for _, node := range nodes.Items { if node.Spec.Unschedulable { continue } pods, err := clientset.CoreV1().Pods("").List(context.Background(), metav1.ListOptions{ FieldSelector: fmt.Sprintf("spec.nodeName=%s,status.phase=Running", node.Name), }) if err != nil { log.Printf("Error listing pods for node %s: %v", node.Name, err) continue } podCounts[node.Name] = len(pods.Items) } if len(podCounts) < 2 { return } var overloadedNode, underloadedNode string maxPods := -1 minPods := int(^uint(0) >> 1) for node, count := range podCounts { if count > maxPods { maxPods = count overloadedNode = node } if count < minPods { minPods = count underloadedNode = node } } if maxPods-minPods <= 1 { log.Println("Cluster balanced") return } pods, err := clientset.CoreV1().Pods("").List(context.Background(), metav1.ListOptions{ FieldSelector: fmt.Sprintf("spec.nodeName=%s,status.phase=Running", overloadedNode), }) if err != nil || len(pods.Items) == 0 { log.Println("No movable pods found") return } pod := pods.Items[0] log.Printf("Rescheduling pod %s/%s from %s", pod.Namespace, pod.Name, overloadedNode) // Cordon overloaded node nodeObj, err := clientset.CoreV1().Nodes().Get(context.Background(), overloadedNode, metav1.GetOptions{}) if err != nil { log.Printf("Failed to get node: %v", err) return } nodeObj.Spec.Unschedulable = true _, err = clientset.CoreV1().Nodes().Update(context.Background(), nodeObj, metav1.UpdateOptions{}) if err != nil { log.Printf("Failed to cordon node: %v", err) return } log.Printf("Successfully cordoned node %s", overloadedNode) // Delete pod err = clientset.CoreV1().Pods(pod.Namespace).Delete(context.Background(), pod.Name, metav1.DeleteOptions{}) if err != nil { log.Printf("Failed to delete pod: %v", err) return } log.Printf("Deleted pod %s", pod.Name) // -------- WAIT FOR REPLACEMENT (FIXED) -------- respawnCtx, cancel := context.WithTimeout(context.Background(), PodRespawnTimeout) log.Printf("PodRespawnTimeout = %v", PodRespawnTimeout) ticker := time.NewTicker(2 * time.Second) for { select { case <-respawnCtx.Done(): log.Println("Timeout waiting for replacement pod") ticker.Stop() cancel() // Uncordon node before exit nodeObj.Spec.Unschedulable = false _, _ = clientset.CoreV1().Nodes().Update(context.Background(), nodeObj, metav1.UpdateOptions{}) log.Printf("Uncordoned node %s after timeout", overloadedNode) return case <-ticker.C: podsList, err := clientset.CoreV1(). Pods(pod.Namespace). List(context.Background(), metav1.ListOptions{}) if err != nil { log.Printf("Error listing pods: %v", err) continue } for _, p := range podsList.Items { if p.Name != pod.Name && p.Spec.NodeName != overloadedNode && p.Status.Phase == v1.PodRunning { log.Printf("Replacement pod %s is Running on %s", p.Name, p.Spec.NodeName) ticker.Stop() cancel() // Uncordon original node nodeObj.Spec.Unschedulable = false _, _ = clientset.CoreV1().Nodes().Update(context.Background(), nodeObj, metav1.UpdateOptions{}) log.Printf("Uncordoned node %s", overloadedNode) return } } log.Println("Waiting for replacement pod to be Running...") } } } // // ─── HELPERS ──────────────────────────────────────────────────────────────────── // func mustEnv(key, def string) string { if val := os.Getenv(key); val != "" { return val } return def } func cordonNode(clientset *kubernetes.Clientset, nodeName string) error { ctx := context.Background() for i := 0; i < 5; i++ { // retry up to 5 times node, err := clientset.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{}) if err != nil { return err } if node.Spec.Unschedulable { return nil // already cordoned } node.Spec.Unschedulable = true _, err = clientset.CoreV1().Nodes().Update(ctx, node, metav1.UpdateOptions{}) if err == nil { return nil } log.Println("Cordoning failed, retrying:", err) time.Sleep(500 * time.Millisecond) } return fmt.Errorf("failed to cordon node %s after retries", nodeName) } func uncordonNode(clientset *kubernetes.Clientset, nodeName string) error { ctx := context.Background() for i := 0; i < 5; i++ { node, err := clientset.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{}) if err != nil { return err } if !node.Spec.Unschedulable { return nil // already uncordoned } node.Spec.Unschedulable = false _, err = clientset.CoreV1().Nodes().Update(ctx, node, metav1.UpdateOptions{}) if err == nil { return nil } log.Println("Uncordoning failed, retrying:", err) time.Sleep(500 * time.Millisecond) } return fmt.Errorf("failed to uncordon node %s after retries", nodeName) } func isDaemonSet(pod *corev1.Pod) bool { for _, owner := range pod.OwnerReferences { if owner.Kind == "DaemonSet" { return true } } return false } func pickTargetNode(nodes []string, usage map[string]map[string]int64, capacity map[string]map[string]int64) string { var bestNode string bestLoad := int64(1 << 62) for _, n := range nodes { cpuPercent := usage[n]["cpu"] * 100 / capacity[n]["cpu"] memPercent := usage[n]["mem"] * 100 / capacity[n]["mem"] load := cpuPercent + memPercent if load < bestLoad { bestLoad = load bestNode = n } } return bestNode } func estimatePodCPU(pod *corev1.Pod) int64 { var cpu int64 for _, c := range pod.Spec.Containers { if q, ok := c.Resources.Requests[corev1.ResourceCPU]; ok { cpu += q.MilliValue() } } return cpu } func estimatePodMem(pod *corev1.Pod) int64 { var mem int64 for _, c := range pod.Spec.Containers { if q, ok := c.Resources.Requests[corev1.ResourceMemory]; ok { mem += q.Value() / (1024 * 1024) } } return mem }