build source

internal/k8s/metrics.go

@@ -0,0 +1,135 @@
+package k8s
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"log/slog"
+	"strings"
+	"time"
+
+	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+// UsageRecorder records resource usage samples into the store.
+type UsageRecorder interface {
+	RecordResourceSample(ctx context.Context, userID, podName string, cpuMillicores, memBytes float64) error
+}
+
+// ResourceSampler periodically samples pod resource usage and records it.
+type ResourceSampler struct {
+	client   *Client
+	recorder UsageRecorder
+	logger   *slog.Logger
+	interval time.Duration
+}
+
+// NewResourceSampler creates a new resource sampler with a 60-second interval.
+func NewResourceSampler(client *Client, recorder UsageRecorder, logger *slog.Logger) *ResourceSampler {
+	return &ResourceSampler{
+		client:   client,
+		recorder: recorder,
+		logger:   logger,
+		interval: 60 * time.Second,
+	}
+}
+
+// Start runs the sampling loop until the context is cancelled.
+func (s *ResourceSampler) Start(ctx context.Context) {
+	ticker := time.NewTicker(s.interval)
+	defer ticker.Stop()
+
+	s.logger.Info("resource sampler started", "interval", s.interval)
+
+	for {
+		select {
+		case <-ctx.Done():
+			s.logger.Info("resource sampler stopped")
+			return
+		case <-ticker.C:
+			s.sample(ctx)
+		}
+	}
+}
+
+// podMetricsList mirrors the k8s metrics API response.
+type podMetricsList struct {
+	Items []podMetrics `json:"items"`
+}
+
+type podMetrics struct {
+	Metadata   podMetricsMetadata `json:"metadata"`
+	Containers []containerMetrics `json:"containers"`
+}
+
+type podMetricsMetadata struct {
+	Name      string            `json:"name"`
+	Namespace string            `json:"namespace"`
+	Labels    map[string]string `json:"labels"`
+}
+
+type containerMetrics struct {
+	Name  string            `json:"name"`
+	Usage map[string]string `json:"usage"`
+}
+
+func (s *ResourceSampler) sample(ctx context.Context) {
+	// List dev-pod namespaces
+	nsList, err := s.client.Clientset.CoreV1().Namespaces().List(ctx, metav1.ListOptions{
+		LabelSelector: "managed-by=dev-pod-api",
+	})
+	if err != nil {
+		s.logger.Error("failed to list namespaces for sampling", "error", err)
+		return
+	}
+
+	for _, ns := range nsList.Items {
+		user := ns.Labels["user"]
+		if user == "" {
+			continue
+		}
+
+		// Query metrics API for this namespace
+		data, err := s.client.Clientset.Discovery().RESTClient().Get().
+			AbsPath(fmt.Sprintf("/apis/metrics.k8s.io/v1beta1/namespaces/%s/pods", ns.Name)).
+			DoRaw(ctx)
+		if err != nil {
+			s.logger.Warn("failed to get pod metrics", "namespace", ns.Name, "error", err)
+			continue
+		}
+
+		var metrics podMetricsList
+		if err := json.Unmarshal(data, &metrics); err != nil {
+			s.logger.Warn("failed to parse pod metrics", "namespace", ns.Name, "error", err)
+			continue
+		}
+
+		for _, pm := range metrics.Items {
+			if !strings.HasPrefix(pm.Metadata.Name, "dev-pod-") {
+				continue
+			}
+			podName := strings.TrimPrefix(pm.Metadata.Name, "dev-pod-")
+
+			var totalCPU, totalMem float64
+			for _, cm := range pm.Containers {
+				if cpuStr, ok := cm.Usage["cpu"]; ok {
+					q, err := resource.ParseQuantity(cpuStr)
+					if err == nil {
+						totalCPU += float64(q.MilliValue())
+					}
+				}
+				if memStr, ok := cm.Usage["memory"]; ok {
+					q, err := resource.ParseQuantity(memStr)
+					if err == nil {
+						totalMem += float64(q.Value())
+					}
+				}
+			}
+
+			if err := s.recorder.RecordResourceSample(ctx, user, podName, totalCPU, totalMem); err != nil {
+				s.logger.Warn("failed to record sample", "user", user, "pod", podName, "error", err)
+			}
+		}
+	}
+}