feat: collect node sizes, instance type, and capacity type

cmd/agent/main.go

@@ -5,10 +5,12 @@ package main
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 	"os"
 	"os/signal"
+	"sync"
 	"syscall"
 
 	"k8s.io/client-go/kubernetes"
@@ -65,7 +67,7 @@ func run(ctx context.Context, logger *slog.Logger, clientset kubernetes.Interfac
 	logger.Info("agent starting", "version", version)
 	defer logger.Info("agent stopping")
 
-	watcher := collector.NewPodWatcher(clientset, func(p collector.PodInfo) {
+	podWatcher := collector.NewPodWatcher(clientset, func(p collector.PodInfo) {
 		logger.Info("pod observed",
 			"namespace", p.Namespace,
 			"pod", p.Name,
@@ -77,5 +79,40 @@ func run(ctx context.Context, logger *slog.Logger, clientset kubernetes.Interfac
 			"containers", p.Containers,
 		)
 	})
-	return watcher.Run(ctx)
+	nodeWatcher := collector.NewNodeWatcher(clientset, func(n collector.NodeInfo) {
+		logger.Info("node observed",
+			"node", n.Name,
+			"instance_type", n.InstanceType,
+			"capacity_type", n.CapacityType,
+			"allocatable_cpu_milli", n.AllocatableCPUMilli,
+			"allocatable_memory_bytes", n.AllocatableMemoryBytes,
+			"capacity_cpu_milli", n.CapacityCPUMilli,
+			"capacity_memory_bytes", n.CapacityMemoryBytes,
+		)
+	})
+
+	// Watchers run until ctx is canceled; a failing watcher must bring the
+	// agent down rather than leave it half-blind.
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	var errs []error
+	for name, run := range map[string]func(context.Context) error{
+		"pods":  podWatcher.Run,
+		"nodes": nodeWatcher.Run,
+	} {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if err := run(ctx); err != nil {
+				mu.Lock()
+				errs = append(errs, fmt.Errorf("%s watcher: %w", name, err))
+				mu.Unlock()
+			}
+			cancel()
+		}()
+	}
+	wg.Wait()
+	return errors.Join(errs...)
 }

internal/collector/nodes.go

@@ -0,0 +1,127 @@
+package collector
+
+import (
+	"context"
+	"fmt"
+
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/client-go/informers"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/tools/cache"
+)
+
+// NodeInfo is the collected view of one node: its size (allocatable and
+// capacity, normalized to millicores and bytes) and the two labels the cost
+// model needs — instance type and capacity type. Nothing else is read from
+// node objects (see docs/security.md §4).
+type NodeInfo struct {
+	Name                   string `json:"name"`
+	InstanceType           string `json:"instance_type,omitempty"`
+	CapacityType           string `json:"capacity_type,omitempty"` // "spot", "on-demand", or "" when undeterminable
+	AllocatableCPUMilli    int64  `json:"allocatable_cpu_milli"`
+	AllocatableMemoryBytes int64  `json:"allocatable_memory_bytes"`
+	CapacityCPUMilli       int64  `json:"capacity_cpu_milli"`
+	CapacityMemoryBytes    int64  `json:"capacity_memory_bytes"`
+}
+
+// NodeWatcher lists all nodes and keeps watching for new ones. Every observed
+// node is reported through the onNode callback.
+type NodeWatcher struct {
+	clientset kubernetes.Interface
+	onNode    func(NodeInfo)
+}
+
+// NewNodeWatcher returns a watcher that calls onNode for every node present
+// at start and for every node added afterwards. onNode is called from the
+// informer goroutine and must not block.
+func NewNodeWatcher(clientset kubernetes.Interface, onNode func(NodeInfo)) *NodeWatcher {
+	return &NodeWatcher{clientset: clientset, onNode: onNode}
+}
+
+// Run blocks until ctx is canceled.
+func (w *NodeWatcher) Run(ctx context.Context) error {
+	factory := informers.NewSharedInformerFactory(w.clientset, 0)
+
+	// Informers must be instantiated before Start, or the factory won't
+	// run them.
+	nodesInformer := factory.Core().V1().Nodes().Informer()
+
+	factory.Start(ctx.Done())
+	defer factory.Shutdown()
+
+	if !cache.WaitForCacheSync(ctx.Done(), nodesInformer.HasSynced) {
+		if ctx.Err() != nil {
+			return nil // canceled during sync — a normal shutdown, not a failure
+		}
+		return fmt.Errorf("node informer cache did not sync")
+	}
+
+	reg, err := nodesInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
+		AddFunc: func(obj any) {
+			if node, ok := obj.(*corev1.Node); ok {
+				w.onNode(describeNode(node))
+			}
+		},
+	})
+	if err != nil {
+		return fmt.Errorf("register node handler: %w", err)
+	}
+	defer func() { _ = nodesInformer.RemoveEventHandler(reg) }()
+
+	<-ctx.Done()
+	return nil
+}
+
+// describeNode reduces a node to the collected view.
+func describeNode(node *corev1.Node) NodeInfo {
+	return NodeInfo{
+		Name:                   node.Name,
+		InstanceType:           instanceType(node.Labels),
+		CapacityType:           capacityType(node.Labels),
+		AllocatableCPUMilli:    node.Status.Allocatable.Cpu().MilliValue(),
+		AllocatableMemoryBytes: node.Status.Allocatable.Memory().Value(),
+		CapacityCPUMilli:       node.Status.Capacity.Cpu().MilliValue(),
+		CapacityMemoryBytes:    node.Status.Capacity.Memory().Value(),
+	}
+}
+
+// instanceType reads the well-known instance-type label, falling back to the
+// deprecated beta name still set by older clusters.
+func instanceType(labels map[string]string) string {
+	if t := labels[corev1.LabelInstanceTypeStable]; t != "" {
+		return t
+	}
+	return labels[corev1.LabelInstanceType]
+}
+
+// capacityType normalizes the provider-specific spot/on-demand markers to
+// "spot" or "on-demand". Nodes without any known marker return "" — absence
+// of evidence is reported as unknown, not defaulted to on-demand.
+func capacityType(labels map[string]string) string {
+	// Karpenter sets the normalized value directly.
+	switch labels["karpenter.sh/capacity-type"] {
+	case "spot":
+		return "spot"
+	case "on-demand", "reserved":
+		return "on-demand"
+	}
+	// EKS managed node groups.
+	switch labels["eks.amazonaws.com/capacityType"] {
+	case "SPOT":
+		return "spot"
+	case "ON_DEMAND":
+		return "on-demand"
+	}
+	// GKE: spot VMs and legacy preemptible VMs.
+	if labels["cloud.google.com/gke-spot"] == "true" || labels["cloud.google.com/gke-preemptible"] == "true" {
+		return "spot"
+	}
+	// AKS scale set priority: "spot" or "regular".
+	switch labels["kubernetes.azure.com/scalesetpriority"] {
+	case "spot":
+		return "spot"
+	case "regular":
+		return "on-demand"
+	}
+	return ""
+}

internal/collector/nodes_test.go

@@ -0,0 +1,206 @@
+package collector
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/client-go/kubernetes/fake"
+)
+
+func node(name string, labels map[string]string) *corev1.Node {
+	return &corev1.Node{
+		ObjectMeta: metav1.ObjectMeta{Name: name, Labels: labels},
+		Status: corev1.NodeStatus{
+			Capacity: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("4"),
+				corev1.ResourceMemory: resource.MustParse("16Gi"),
+			},
+			Allocatable: corev1.ResourceList{
+				corev1.ResourceCPU:    resource.MustParse("3900m"),
+				corev1.ResourceMemory: resource.MustParse("14Gi"),
+			},
+		},
+	}
+}
+
+// collectNodes runs a NodeWatcher over the given objects and returns
+// everything reported until at least want nodes are seen or the timeout hits.
+func collectNodes(t *testing.T, want int, objects ...runtime.Object) map[string]NodeInfo {
+	t.Helper()
+	clientset := fake.NewClientset(objects...)
+
+	var mu sync.Mutex
+	seen := make(map[string]NodeInfo)
+	got := make(chan struct{}, 64)
+	watcher := NewNodeWatcher(clientset, func(n NodeInfo) {
+		mu.Lock()
+		seen[n.Name] = n
+		mu.Unlock()
+		got <- struct{}{}
+	})
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	done := make(chan error, 1)
+	go func() { done <- watcher.Run(ctx) }()
+
+	deadline := time.After(5 * time.Second)
+	for n := 0; n < want; {
+		select {
+		case <-got:
+			n++
+		case <-deadline:
+			t.Fatalf("saw %d nodes before timeout, want %d", n, want)
+		}
+	}
+	cancel()
+	if err := <-done; err != nil {
+		t.Fatalf("watcher returned error: %v", err)
+	}
+	mu.Lock()
+	defer mu.Unlock()
+	return seen
+}
+
+func TestCollectsNodeSizes(t *testing.T) {
+	seen := collectNodes(t, 1, node("node-1", nil))
+
+	info := seen["node-1"]
+	want := NodeInfo{
+		Name:                   "node-1",
+		AllocatableCPUMilli:    3900,
+		AllocatableMemoryBytes: 14 << 30,
+		CapacityCPUMilli:       4000,
+		CapacityMemoryBytes:    16 << 30,
+	}
+	if info != want {
+		t.Errorf("node = %+v, want %+v", info, want)
+	}
+}
+
+func TestNodeInstanceAndCapacityTypes(t *testing.T) {
+	cases := []struct {
+		name             string
+		labels           map[string]string
+		wantInstanceType string
+		wantCapacityType string
+	}{
+		{
+			name: "karpenter-spot",
+			labels: map[string]string{
+				"node.kubernetes.io/instance-type": "m6i.large",
+				"karpenter.sh/capacity-type":       "spot",
+			},
+			wantInstanceType: "m6i.large",
+			wantCapacityType: "spot",
+		},
+		{
+			name: "eks-on-demand",
+			labels: map[string]string{
+				"node.kubernetes.io/instance-type": "m5.xlarge",
+				"eks.amazonaws.com/capacityType":   "ON_DEMAND",
+			},
+			wantInstanceType: "m5.xlarge",
+			wantCapacityType: "on-demand",
+		},
+		{
+			name: "gke-spot",
+			labels: map[string]string{
+				"node.kubernetes.io/instance-type": "e2-standard-4",
+				"cloud.google.com/gke-spot":        "true",
+			},
+			wantInstanceType: "e2-standard-4",
+			wantCapacityType: "spot",
+		},
+		{
+			name: "gke-preemptible",
+			labels: map[string]string{
+				"cloud.google.com/gke-preemptible": "true",
+			},
+			wantCapacityType: "spot",
+		},
+		{
+			name: "aks-regular",
+			labels: map[string]string{
+				"kubernetes.azure.com/scalesetpriority": "regular",
+			},
+			wantCapacityType: "on-demand",
+		},
+		{
+			name: "beta-instance-type-fallback",
+			labels: map[string]string{
+				"beta.kubernetes.io/instance-type": "t3.medium",
+			},
+			wantInstanceType: "t3.medium",
+		},
+		{
+			name:   "no-markers",
+			labels: map[string]string{"kubernetes.io/hostname": "bare-metal-7"},
+		},
+	}
+
+	objects := make([]runtime.Object, 0, len(cases))
+	for _, c := range cases {
+		objects = append(objects, node(c.name, c.labels))
+	}
+	seen := collectNodes(t, len(cases), objects...)
+
+	for _, c := range cases {
+		info, ok := seen[c.name]
+		if !ok {
+			t.Errorf("node %s was not reported", c.name)
+			continue
+		}
+		if info.InstanceType != c.wantInstanceType {
+			t.Errorf("%s: instance type = %q, want %q", c.name, info.InstanceType, c.wantInstanceType)
+		}
+		if info.CapacityType != c.wantCapacityType {
+			t.Errorf("%s: capacity type = %q, want %q", c.name, info.CapacityType, c.wantCapacityType)
+		}
+	}
+}
+
+func TestReportsNodesAddedAfterStart(t *testing.T) {
+	clientset := fake.NewClientset(node("existing", nil))
+
+	events := make(chan NodeInfo, 16)
+	watcher := NewNodeWatcher(clientset, func(n NodeInfo) { events <- n })
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	done := make(chan error, 1)
+	go func() { done <- watcher.Run(ctx) }()
+
+	waitFor := func(name string) {
+		t.Helper()
+		deadline := time.After(5 * time.Second)
+		for {
+			select {
+			case n := <-events:
+				if n.Name == name {
+					return
+				}
+			case <-deadline:
+				t.Fatalf("node %s was not reported before timeout", name)
+			}
+		}
+	}
+
+	waitFor("existing")
+
+	if _, err := clientset.CoreV1().Nodes().Create(ctx, node("scaled-up", nil), metav1.CreateOptions{}); err != nil {
+		t.Fatalf("creating node: %v", err)
+	}
+	waitFor("scaled-up")
+
+	cancel()
+	if err := <-done; err != nil {
+		t.Fatalf("watcher returned error: %v", err)
+	}
+}