diff --git a/api/v1alpha1/agentdeployment_types.go b/api/v1alpha1/agentdeployment_types.go index 3cbae92..f93a2c9 100644 --- a/api/v1alpha1/agentdeployment_types.go +++ b/api/v1alpha1/agentdeployment_types.go @@ -234,6 +234,7 @@ type AgentDeploymentList struct { Items []AgentDeployment `json:"items"` } +// init registers AgentDeployment and AgentDeploymentList types with the SchemeBuilder. func init() { SchemeBuilder.Register(&AgentDeployment{}, &AgentDeploymentList{}) } diff --git a/api/v1alpha1/error_rate_test.go b/api/v1alpha1/error_rate_test.go index 933b58a..bcc3793 100644 --- a/api/v1alpha1/error_rate_test.go +++ b/api/v1alpha1/error_rate_test.go @@ -22,6 +22,7 @@ import ( agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" ) +// TestParseErrorRate verifies percentage string parsing across valid, invalid, and boundary inputs. func TestParseErrorRate(t *testing.T) { t.Parallel() tests := []struct { @@ -62,6 +63,7 @@ func TestParseErrorRate(t *testing.T) { } } +// abs returns the absolute value of a float64. func abs(f float64) float64 { if f < 0 { return -f diff --git a/api/v1alpha1/tenantquota_types.go b/api/v1alpha1/tenantquota_types.go index 011e9eb..203d738 100644 --- a/api/v1alpha1/tenantquota_types.go +++ b/api/v1alpha1/tenantquota_types.go @@ -86,6 +86,7 @@ type TenantQuotaList struct { Items []TenantQuota `json:"items"` } +// init registers TenantQuota and TenantQuotaList types with the SchemeBuilder. func init() { SchemeBuilder.Register(&TenantQuota{}, &TenantQuotaList{}) } diff --git a/cmd/main.go b/cmd/main.go index fbadaa4..f8afa3b 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -26,6 +26,8 @@ import ( // to ensure that exec-entrypoint and run can make use of them. _ "k8s.io/client-go/plugin/pkg/client/auth" + autoscalingv2 "k8s.io/api/autoscaling/v2" + apiextensionsv1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1" "k8s.io/apimachinery/pkg/runtime" utilruntime "k8s.io/apimachinery/pkg/util/runtime" clientgoscheme "k8s.io/client-go/kubernetes/scheme" @@ -36,6 +38,8 @@ import ( metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server" "sigs.k8s.io/controller-runtime/pkg/webhook" + monitoringv1 "github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring/v1" + agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" "github.com/gitcommitankit/agentrax/internal/controller" "github.com/gitcommitankit/agentrax/internal/quota" @@ -48,13 +52,18 @@ var ( setupLog = ctrl.Log.WithName("setup") ) +// init registers all Kubernetes core, CRD, and monitoring schemes. func init() { utilruntime.Must(clientgoscheme.AddToScheme(scheme)) + utilruntime.Must(autoscalingv2.AddToScheme(scheme)) + utilruntime.Must(apiextensionsv1.AddToScheme(scheme)) + utilruntime.Must(monitoringv1.AddToScheme(scheme)) utilruntime.Must(agentraxv1alpha1.AddToScheme(scheme)) // +kubebuilder:scaffold:scheme } +// main is the entrypoint for the Agentrax controller manager binary. func main() { var metricsAddr string var enableLeaderElection bool @@ -98,9 +107,18 @@ func main() { tlsOpts = append(tlsOpts, disableHTTP2) } - webhookServer := webhook.NewServer(webhook.Options{ - TLSOpts: tlsOpts, - }) + // Resolve the webhook-enabled flag once so both the server creation and + // handler registration use the same value. Log it explicitly so operators + // can confirm the resolved state at startup. + enableWebhooks := os.Getenv("ENABLE_WEBHOOKS") != "false" + setupLog.Info("webhook state resolved", "enabled", enableWebhooks) + + var webhookServer webhook.Server + if enableWebhooks { + webhookServer = webhook.NewServer(webhook.Options{ + TLSOpts: tlsOpts, + }) + } // Metrics endpoint is enabled in 'config/default/kustomization.yaml'. The Metrics options configure the server. // More info: @@ -154,8 +172,9 @@ func main() { quotaEnforcer := quota.NewEnforcer(gpuResourceName) if err = (&controller.AgentDeploymentReconciler{ - Client: mgr.GetClient(), - Scheme: mgr.GetScheme(), + Client: mgr.GetClient(), + Scheme: mgr.GetScheme(), + GPUResourceName: gpuResourceName, }).SetupWithManager(mgr); err != nil { setupLog.Error(err, "unable to create controller", "controller", "AgentDeployment") os.Exit(1) @@ -168,9 +187,11 @@ func main() { setupLog.Error(err, "unable to create controller", "controller", "TenantQuota") os.Exit(1) } - if err = agentraxwebhook.SetupAgentDeploymentWebhookWithManager(mgr, quotaEnforcer); err != nil { - setupLog.Error(err, "unable to register webhook", "webhook", "AgentDeployment") - os.Exit(1) + if enableWebhooks { + if err = agentraxwebhook.SetupAgentDeploymentWebhookWithManager(mgr, quotaEnforcer); err != nil { + setupLog.Error(err, "unable to register webhook", "webhook", "AgentDeployment") + os.Exit(1) + } } // +kubebuilder:scaffold:builder diff --git a/config/prometheus-adapter/custom-metrics-config.yaml b/config/prometheus-adapter/custom-metrics-config.yaml new file mode 100644 index 0000000..63e3152 --- /dev/null +++ b/config/prometheus-adapter/custom-metrics-config.yaml @@ -0,0 +1,51 @@ +# Prometheus Adapter custom metrics configuration for Agentrax. +# +# This ConfigMap is consumed by the prometheus-adapter deployment (typically in +# the monitoring namespace). It maps PromQL queries to named custom metrics that +# the HorizontalPodAutoscaler can target via the custom.metrics.k8s.io API. +# +# Deploy with: +# kubectl apply -f config/prometheus-adapter/custom-metrics-config.yaml +# kubectl rollout restart deployment/prometheus-adapter -n monitoring +# +# Verify metrics are registered: +# kubectl get --raw /apis/external.metrics.k8s.io/v1beta1 | jq . +# +apiVersion: v1 +kind: ConfigMap +metadata: + # Use a distinct name so this ConfigMap does not collide with or overwrite + # the upstream prometheus-adapter ConfigMap (commonly named adapter-config). + # Reference this name in your prometheus-adapter Deployment via + # --config=/etc/adapter/config.yaml mounted from this ConfigMap. + name: agentrax-custom-metrics + namespace: monitoring + labels: + app.kubernetes.io/name: prometheus-adapter + app.kubernetes.io/managed-by: agentrax +data: + config.yaml: | + externalRules: + # ── queueDepth ──────────────────────────────────────────────────────────── + # Exposes agentrax_queue_depth via external.metrics.k8s.io, which is the + # API group queried by HPAs using ExternalMetricSourceType (the type + # BuildHPA produces). The <<.LabelMatchers>> template is populated by the + # Prometheus Adapter from the HPA metric selector. Prometheus sanitizes + # label names (app.kubernetes.io/name → app_kubernetes_io_name), so the + # query references the sanitized forms that actually exist in storage. + - seriesQuery: 'agentrax_queue_depth{namespace!=""}' + name: + matches: "^agentrax_queue_depth$" + as: "agentrax_queue_depth" + metricsQuery: 'sum(<<.Series>>{<<.LabelMatchers>>}) by (<<.GroupBy>>)' + + # ── gpuUtilization ──────────────────────────────────────────────────────── + # Exposes agentrax_gpu_utilization via external.metrics.k8s.io. + # If your GPU device plugin exposes a different metric name (e.g., from DCGM), + # update the seriesQuery and the as: name here; the HPA target in the + # AgentDeployment spec references the as: name, which stays stable. + - seriesQuery: 'agentrax_gpu_utilization{namespace!=""}' + name: + matches: "^agentrax_gpu_utilization$" + as: "agentrax_gpu_utilization" + metricsQuery: 'sum(<<.Series>>{<<.LabelMatchers>>}) by (<<.GroupBy>>)' diff --git a/config/prometheus-adapter/kustomization.yaml b/config/prometheus-adapter/kustomization.yaml new file mode 100644 index 0000000..96d6b35 --- /dev/null +++ b/config/prometheus-adapter/kustomization.yaml @@ -0,0 +1,41 @@ +apiVersion: kustomize.config.k8s.io/v1beta1 +kind: Kustomization + +# Prometheus Adapter custom metrics configuration for Agentrax. +# +# Apply this overlay after the Prometheus Adapter base is installed: +# kubectl apply -k config/prometheus-adapter/ +# +# Note: This kustomization targets the monitoring namespace where the +# prometheus-adapter deployment is expected to run. If your cluster uses +# a different namespace, update the namespace field below. +namespace: monitoring + +resources: + - custom-metrics-config.yaml + +# To apply this configuration, either: +# 1. Include the prometheus-adapter Deployment base in resources above, then +# uncomment the patches section below, OR +# 2. Manually configure your existing prometheus-adapter Deployment to mount +# this ConfigMap at /etc/adapter/config.yaml and pass --config=/etc/adapter/config.yaml +# +# patches: +# - target: +# kind: Deployment +# name: prometheus-adapter +# patch: |- +# - op: add +# path: /spec/template/spec/volumes/- +# value: +# name: adapter-config +# configMap: +# name: agentrax-custom-metrics +# - op: add +# path: /spec/template/spec/containers/0/volumeMounts/- +# value: +# name: adapter-config +# mountPath: /etc/adapter +# - op: add +# path: /spec/template/spec/containers/0/args/- +# value: --config=/etc/adapter/config.yaml diff --git a/config/rbac/role.yaml b/config/rbac/role.yaml index bc3d84c..25b554a 100644 --- a/config/rbac/role.yaml +++ b/config/rbac/role.yaml @@ -59,6 +59,18 @@ rules: - patch - update - watch +- apiGroups: + - autoscaling + resources: + - horizontalpodautoscalers + verbs: + - create + - delete + - get + - list + - patch + - update + - watch - apiGroups: - "" resources: diff --git a/config/samples/agentrax_v1alpha1_agentdeployment.yaml b/config/samples/agentrax_v1alpha1_agentdeployment.yaml index 50ee7a2..ccf91be 100644 --- a/config/samples/agentrax_v1alpha1_agentdeployment.yaml +++ b/config/samples/agentrax_v1alpha1_agentdeployment.yaml @@ -31,6 +31,10 @@ spec: maxTotalReplicas: 10 # Maximum number of distinct AgentDeployments allowed. maxAgents: 5 + # Maximum GPU count allowed across all agents in this namespace. + maxGPUs: 2 + # Maximum replica count allowed for any single AgentDeployment. + maxReplicasPerAgent: 5 --- # 3. AgentDeployment — the main resource. diff --git a/config/webhook/manifests.yaml b/config/webhook/manifests.yaml index fa32a1a..baa6fd1 100644 --- a/config/webhook/manifests.yaml +++ b/config/webhook/manifests.yaml @@ -50,3 +50,4 @@ webhooks: resources: - agentdeployments sideEffects: None + timeoutSeconds: 10 diff --git a/internal/controller/agentdeployment_builder_test.go b/internal/controller/agentdeployment_builder_test.go index 9162660..848f43c 100644 --- a/internal/controller/agentdeployment_builder_test.go +++ b/internal/controller/agentdeployment_builder_test.go @@ -62,6 +62,7 @@ func makeAD(name, image string, port int32, minReplicas int32) *agentraxv1alpha1 // ── desiredDeployment ───────────────────────────────────────────────────────── +// TestDesiredDeployment_Image verifies that the desired Deployment container image matches spec.image. func TestDesiredDeployment_Image(t *testing.T) { r := &AgentDeploymentReconciler{} ad := makeAD("my-agent", "registry.io/agent:v1", 8080, 1) @@ -72,6 +73,7 @@ func TestDesiredDeployment_Image(t *testing.T) { } } +// TestDesiredDeployment_Port verifies container port configuration and default fallback. func TestDesiredDeployment_Port(t *testing.T) { tests := []struct { name string @@ -94,6 +96,7 @@ func TestDesiredDeployment_Port(t *testing.T) { } } +// TestDesiredDeployment_Replicas verifies that desired Deployment replicas match spec.replicas.min. func TestDesiredDeployment_Replicas(t *testing.T) { r := &AgentDeploymentReconciler{} ad := makeAD(testDefaultAgent, testDefaultImage, 8080, 3) @@ -104,6 +107,7 @@ func TestDesiredDeployment_Replicas(t *testing.T) { } } +// TestDesiredDeployment_EnvAndArgs verifies propagation of environment variables and container arguments. func TestDesiredDeployment_EnvAndArgs(t *testing.T) { r := &AgentDeploymentReconciler{} ad := makeAD(testDefaultAgent, testDefaultImage, 8080, 1) @@ -121,6 +125,7 @@ func TestDesiredDeployment_EnvAndArgs(t *testing.T) { } } +// TestDesiredDeployment_Resources verifies container CPU and Memory resource requests/limits propagation. func TestDesiredDeployment_Resources(t *testing.T) { r := &AgentDeploymentReconciler{} ad := makeAD(testDefaultAgent, testDefaultImage, 8080, 1) @@ -153,6 +158,7 @@ func TestDesiredDeployment_Resources(t *testing.T) { } } +// TestDesiredDeployment_Labels verifies required standard labels on Deployment and Pod template. func TestDesiredDeployment_Labels(t *testing.T) { r := &AgentDeploymentReconciler{} ad := makeAD(testAgentName, testDefaultImage, 8080, 1) @@ -173,6 +179,7 @@ func TestDesiredDeployment_Labels(t *testing.T) { } } +// TestDesiredDeployment_SelectorMatchesPodLabels verifies Deployment selector matches Pod template labels. func TestDesiredDeployment_SelectorMatchesPodLabels(t *testing.T) { r := &AgentDeploymentReconciler{} ad := makeAD(testDefaultAgent, testDefaultImage, 8080, 1) @@ -187,6 +194,7 @@ func TestDesiredDeployment_SelectorMatchesPodLabels(t *testing.T) { // ── desiredService ──────────────────────────────────────────────────────────── +// TestDesiredService_Port verifies Service port configuration and default fallback. func TestDesiredService_Port(t *testing.T) { tests := []struct { name string @@ -208,6 +216,7 @@ func TestDesiredService_Port(t *testing.T) { } } +// TestDesiredService_Selector verifies Service selector targets the agent pod label. func TestDesiredService_Selector(t *testing.T) { r := &AgentDeploymentReconciler{} ad := makeAD(testAgentName, testDefaultImage, 8080, 1) @@ -218,6 +227,7 @@ func TestDesiredService_Selector(t *testing.T) { } } +// TestDesiredService_ClusterIPType verifies that the created Service is of type ClusterIP. func TestDesiredService_ClusterIPType(t *testing.T) { r := &AgentDeploymentReconciler{} ad := makeAD(testDefaultAgent, testDefaultImage, 8080, 1) @@ -230,6 +240,7 @@ func TestDesiredService_ClusterIPType(t *testing.T) { // ── agentLabels ─────────────────────────────────────────────────────────────── +// TestAgentLabels verifies standard label generation for an AgentDeployment. func TestAgentLabels(t *testing.T) { ad := &agentraxv1alpha1.AgentDeployment{ ObjectMeta: metav1.ObjectMeta{Name: "foo"}, @@ -251,6 +262,7 @@ func TestAgentLabels(t *testing.T) { // ── condition helpers ───────────────────────────────────────────────────────── +// TestSetAndGetCondition verifies setting and reading status conditions on AgentDeployment. func TestSetAndGetCondition(t *testing.T) { ad := &agentraxv1alpha1.AgentDeployment{} @@ -268,6 +280,7 @@ func TestSetAndGetCondition(t *testing.T) { } } +// TestSetCondition_Overwrite verifies that updating an existing condition updates in-place without duplicates. func TestSetCondition_Overwrite(t *testing.T) { ad := &agentraxv1alpha1.AgentDeployment{} @@ -284,6 +297,7 @@ func TestSetCondition_Overwrite(t *testing.T) { } } +// TestRemoveCondition verifies condition removal from the status condition slice. func TestRemoveCondition(t *testing.T) { ad := &agentraxv1alpha1.AgentDeployment{} @@ -300,6 +314,7 @@ func TestRemoveCondition(t *testing.T) { } } +// TestRemoveCondition_NonExistent verifies that removing a non-existent condition is a safe no-op. func TestRemoveCondition_NonExistent(t *testing.T) { ad := &agentraxv1alpha1.AgentDeployment{} // Should be a no-op, not panic. @@ -309,6 +324,7 @@ func TestRemoveCondition_NonExistent(t *testing.T) { } } +// TestGetCondition_Absent verifies that querying an un-set condition returns nil. func TestGetCondition_Absent(t *testing.T) { ad := &agentraxv1alpha1.AgentDeployment{} c := GetCondition(ad, agentraxv1alpha1.ConditionReady) @@ -319,6 +335,7 @@ func TestGetCondition_Absent(t *testing.T) { // ── desiredServiceMonitor ───────────────────────────────────────────────────── +// TestDesiredServiceMonitor_Endpoint verifies ServiceMonitor metrics endpoint configuration. func TestDesiredServiceMonitor_Endpoint(t *testing.T) { r := &AgentDeploymentReconciler{} ad := makeAD(testDefaultAgent, testDefaultImage, 8080, 1) @@ -336,6 +353,7 @@ func TestDesiredServiceMonitor_Endpoint(t *testing.T) { } } +// TestDesiredServiceMonitor_SelectorMatchesLabels verifies ServiceMonitor selector matches agent labels. func TestDesiredServiceMonitor_SelectorMatchesLabels(t *testing.T) { r := &AgentDeploymentReconciler{} ad := makeAD(testAgentName, testDefaultImage, 8080, 1) diff --git a/internal/controller/agentdeployment_controller.go b/internal/controller/agentdeployment_controller.go index 8190afc..a3f517d 100644 --- a/internal/controller/agentdeployment_controller.go +++ b/internal/controller/agentdeployment_controller.go @@ -24,11 +24,13 @@ import ( "time" appsv1 "k8s.io/api/apps/v1" + autoscalingv2 "k8s.io/api/autoscaling/v2" corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/equality" apierrors "k8s.io/apimachinery/pkg/api/errors" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" "k8s.io/apimachinery/pkg/util/intstr" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/client" @@ -39,6 +41,28 @@ import ( monitoringv1 "github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring/v1" agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" + "github.com/gitcommitankit/agentrax/internal/scaling" +) + +// quotaState is the outcome of reconcileHPA's quota evaluation, used by +// updateStatus to apply the correct QuotaLimited condition without relying on +// a bool that cannot distinguish between "not capped" and "not evaluated". +type quotaState int + +const ( + // quotaStateUncapped means HPA was reconciled and quota headroom was sufficient. + // updateStatus will clear the QuotaLimited condition. + quotaStateUncapped quotaState = iota + // quotaStateCapped means HPA was reconciled but maxReplicas was capped by quota. + // updateStatus will set QuotaLimited=True with reason HPAMaxReplicasCapped. + quotaStateCapped + // quotaStateUnknown means quota could not be evaluated (e.g., TenantQuota not found). + // updateStatus will set QuotaLimited=True with reason TenantQuotaNotFound. + quotaStateUnknown + // quotaStateSkipped means HPA reconciliation was skipped (canary rollout in progress). + // updateStatus must NOT touch the QuotaLimited condition so a pre-canary + // capped state is neither erroneously cleared nor re-set. + quotaStateSkipped ) // AgentDeploymentReconciler reconciles an AgentDeployment object. @@ -46,6 +70,11 @@ type AgentDeploymentReconciler struct { client.Client Scheme *runtime.Scheme + // GPUResourceName is the Kubernetes resource name used to count GPU units + // (e.g. "nvidia.com/gpu"). Injected from the --gpu-resource-name operator flag. + // Used when computing quota headroom for HPA max-replicas capping. + GPUResourceName string + // hasServiceMonitorCRD is set once during SetupWithManager and determines // whether ServiceMonitor reconciliation is attempted at all. hasServiceMonitorCRD bool @@ -73,11 +102,13 @@ func (r *AgentDeploymentReconciler) SetDeregister(fn func(ctx context.Context, a // +kubebuilder:rbac:groups=agentrax.io,resources=agentdeployments/status,verbs=get;update;patch // +kubebuilder:rbac:groups=agentrax.io,resources=agentdeployments/finalizers,verbs=update // +kubebuilder:rbac:groups=apps,resources=deployments,verbs=get;list;watch;create;update;patch;delete +// +kubebuilder:rbac:groups=autoscaling,resources=horizontalpodautoscalers,verbs=get;list;watch;create;update;patch;delete // +kubebuilder:rbac:groups=core,resources=services,verbs=get;list;watch;create;update;patch;delete // +kubebuilder:rbac:groups=core,resources=pods,verbs=get;list;watch // +kubebuilder:rbac:groups=core,resources=events,verbs=create;patch // +kubebuilder:rbac:groups=monitoring.coreos.com,resources=servicemonitors,verbs=get;list;watch;create;update;patch;delete // +kubebuilder:rbac:groups=apiextensions.k8s.io,resources=customresourcedefinitions,verbs=get;list;watch +// +kubebuilder:rbac:groups=agentrax.io,resources=tenantquotas,verbs=get;list;watch // Reconcile drives the AgentDeployment's observed state toward its declared spec. // It creates and self-heals a Deployment, Service, and (when Prometheus Operator is present) @@ -136,9 +167,29 @@ func (r *AgentDeploymentReconciler) Reconcile(ctx context.Context, req ctrl.Requ return ctrl.Result{}, fmt.Errorf("reconciling servicemonitor: %w", err) } - // 6. Derive status from the live Deployment and update it — always last. - if result, err := r.updateStatus(ctx, ad, logger); err != nil || result.RequeueAfter > 0 { - return result, err + // 6. Reconcile the managed HPA (skip during active canary — Phase 4 owns it). + // reconcileHPA also returns the quota evaluation state so updateStatus can + // write the correct QuotaLimited condition onto the freshly re-fetched object. + hpaResult, qs, err := r.reconcileHPA(ctx, ad) + if err != nil { + return ctrl.Result{}, fmt.Errorf("reconciling hpa: %w", err) + } + + // 7. Derive status from the live Deployment and update it — always last. + // We continue into updateStatus even when hpaResult requests a requeue so + // that the QuotaLimited condition is written in the same reconcile cycle. + // Return the shorter of the two requeue intervals. + statusResult, err := r.updateStatus(ctx, ad, logger, qs) + if err != nil { + return statusResult, fmt.Errorf("updating status: %w", err) + } + if hpaResult.RequeueAfter > 0 { + if statusResult.RequeueAfter == 0 || hpaResult.RequeueAfter < statusResult.RequeueAfter { + return hpaResult, nil + } + } + if statusResult.RequeueAfter > 0 { + return statusResult, nil } return ctrl.Result{}, nil @@ -256,9 +307,112 @@ func (r *AgentDeploymentReconciler) reconcileServiceMonitor(ctx context.Context, return err } +// reconcileHPA creates or updates the HorizontalPodAutoscaler owned by this +// AgentDeployment. It is skipped when a canary rollout is in progress because +// Phase 4 owns the HPA lifecycle during rollout. The HPA's maxReplicas is +// capped at the tenant quota headroom; the returned quotaState tells the caller +// which QuotaLimited condition (if any) to apply to status. +func (r *AgentDeploymentReconciler) reconcileHPA(ctx context.Context, ad *agentraxv1alpha1.AgentDeployment) (ctrl.Result, quotaState, error) { + // Phase 4 owns the HPA when a canary rollout is in progress. Signal + // quotaStateSkipped so updateStatus does not touch the QuotaLimited condition. + if ad.Status.Phase == agentraxv1alpha1.PhaseRolloutInProgress { + return ctrl.Result{}, quotaStateSkipped, nil + } + + // Fetch the TenantQuota to compute quota headroom. + tq := &agentraxv1alpha1.TenantQuota{} + if err := r.Get(ctx, types.NamespacedName{Name: ad.Spec.TenantRef, Namespace: ad.Namespace}, tq); err != nil { + if apierrors.IsNotFound(err) { + // TenantQuota missing — the webhook prevents this on create, but it + // can happen if the TenantQuota is deleted while agents exist. + // Signal quotaStateUnknown so updateStatus sets QuotaLimited with a + // TenantQuotaNotFound reason (not the "HPA capped" reason). + // Do NOT return an error so updateStatus still runs this cycle. + return ctrl.Result{RequeueAfter: 10 * time.Second}, quotaStateUnknown, nil + } + return ctrl.Result{}, quotaStateUncapped, fmt.Errorf("fetching TenantQuota %q: %w", ad.Spec.TenantRef, err) + } + + // Compute how many replicas the other agents in this tenant already consume + // so QuotaHeadroom accounts for the full tenant budget. + usedByOthers, err := r.replicasUsedByOtherAgents(ctx, ad) + if err != nil { + return ctrl.Result{}, quotaStateUncapped, fmt.Errorf("computing replica usage for quota headroom: %w", err) + } + + headroom := scaling.QuotaHeadroom(tq.Spec, ad.Spec, usedByOthers) + desiredHPA := scaling.BuildHPA(ad, headroom) + + // Set controller owner reference before CreateOrUpdate so garbage collection + // removes the HPA when the AgentDeployment is deleted. + if err := controllerutil.SetControllerReference(ad, desiredHPA, r.Scheme); err != nil { + return ctrl.Result{}, quotaStateUncapped, fmt.Errorf("setting HPA owner reference: %w", err) + } + + existing := &autoscalingv2.HorizontalPodAutoscaler{} + existing.Name = desiredHPA.Name + existing.Namespace = desiredHPA.Namespace + + _, err = controllerutil.CreateOrUpdate(ctx, r.Client, existing, func() error { + existing.Labels = desiredHPA.Labels + existing.Spec.ScaleTargetRef = desiredHPA.Spec.ScaleTargetRef + existing.Spec.MinReplicas = desiredHPA.Spec.MinReplicas + existing.Spec.MaxReplicas = desiredHPA.Spec.MaxReplicas + existing.Spec.Metrics = desiredHPA.Spec.Metrics + existing.Spec.Behavior = desiredHPA.Spec.Behavior + // Re-apply owner reference in case it was cleared out-of-band. + if err := controllerutil.SetControllerReference(ad, existing, r.Scheme); err != nil { + return fmt.Errorf("setting HPA owner reference: %w", err) + } + return nil + }) + if err != nil { + return ctrl.Result{}, quotaStateUncapped, fmt.Errorf("creating/updating HPA: %w", err) + } + + if scaling.IsQuotaCapped(ad, headroom) { + return ctrl.Result{}, quotaStateCapped, nil + } + return ctrl.Result{}, quotaStateUncapped, nil +} + +// replicasUsedByOtherAgents returns the sum of spec.replicas.max across all +// AgentDeployments in the same namespace that reference the same TenantQuota, +// excluding the AgentDeployment being reconciled and any that are terminating +// (DeletionTimestamp set). Terminating ADs will be removed shortly, so +// including their replicas would inflate the quota headroom calculation and +// cause premature rejection of new creates. +func (r *AgentDeploymentReconciler) replicasUsedByOtherAgents(ctx context.Context, ad *agentraxv1alpha1.AgentDeployment) (int32, error) { + list := &agentraxv1alpha1.AgentDeploymentList{} + if err := r.List(ctx, list, client.InNamespace(ad.Namespace)); err != nil { + return 0, fmt.Errorf("listing AgentDeployments for quota headroom: %w", err) + } + + var total int32 + for i := range list.Items { + other := &list.Items[i] + if other.Spec.TenantRef != ad.Spec.TenantRef { + continue + } + if other.Name == ad.Name && other.Namespace == ad.Namespace { + continue // exclude self + } + // Exclude terminating ADs: they will be deleted soon and their replicas + // should not count against the remaining budget. + if !other.DeletionTimestamp.IsZero() { + continue + } + total += other.Spec.Replicas.Max + } + return total, nil +} + // updateStatus derives the AgentDeployment status from the live Deployment and writes it. +// qs is the quota evaluation outcome from reconcileHPA; the QuotaLimited condition +// is applied to the freshly re-fetched object here so it is never silently discarded. +// quotaStateSkipped means the condition must not be modified (canary in progress). // This is always the last step in the reconcile loop. -func (r *AgentDeploymentReconciler) updateStatus(ctx context.Context, ad *agentraxv1alpha1.AgentDeployment, logger logr.Logger) (ctrl.Result, error) { +func (r *AgentDeploymentReconciler) updateStatus(ctx context.Context, ad *agentraxv1alpha1.AgentDeployment, logger logr.Logger, qs quotaState) (ctrl.Result, error) { // Re-fetch the live Deployment to get accurate replica counts. dep := &appsv1.Deployment{} depKey := client.ObjectKey{Name: ad.Name, Namespace: ad.Namespace} @@ -287,8 +441,29 @@ func (r *AgentDeploymentReconciler) updateStatus(ctx context.Context, ad *agentr // condition count) that a scalar/len check would silently miss. prevStatus := latest.Status.DeepCopy() - latest.Status.CurrentReplicas = dep.Status.ReadyReplicas + // Apply the QuotaLimited condition from reconcileHPA onto the freshly + // re-fetched object. This must happen before the DeepEqual check so the + // condition is written in the same API call as the rest of the status. + // quotaStateSkipped means HPA reconciliation was bypassed (canary in progress) + // — the existing condition must be left untouched. + switch qs { + case quotaStateCapped: + SetCondition(latest, agentraxv1alpha1.ConditionQuotaLimited, metav1.ConditionTrue, + "HPAMaxReplicasCapped", + fmt.Sprintf("spec.replicas.max (%d) exceeds quota headroom; HPA capped", + latest.Spec.Replicas.Max)) + case quotaStateUnknown: + SetCondition(latest, agentraxv1alpha1.ConditionQuotaLimited, metav1.ConditionTrue, + "TenantQuotaNotFound", + fmt.Sprintf("TenantQuota %q not found in namespace %s; requeuing", + latest.Spec.TenantRef, latest.Namespace)) + case quotaStateUncapped: + RemoveCondition(latest, agentraxv1alpha1.ConditionQuotaLimited) + default: + // quotaStateSkipped (canary in progress): leave the existing condition as-is. + } + latest.Status.CurrentReplicas = dep.Status.ReadyReplicas // Detect ImagePullBackOff by inspecting pod list; requeue if listing fails. imagePullFailed, failMsg, err := r.detectImagePullFailure(ctx, ad) if err != nil { @@ -463,6 +638,11 @@ func (r *AgentDeploymentReconciler) desiredService(ad *agentraxv1alpha1.AgentDep } // desiredServiceMonitor builds the ServiceMonitor spec that scrapes /metrics on the agent pods. +// TargetLabels propagates app.kubernetes.io/name and app.kubernetes.io/managed-by from the +// Service object into every Prometheus sample. Prometheus sanitizes these label names +// (dots become underscores: app_kubernetes_io_name, app_kubernetes_io_managed_by) before +// storage, and the HPA metric selector and Prometheus Adapter query reference the sanitized +// forms to match the stored labels. func (r *AgentDeploymentReconciler) desiredServiceMonitor(ad *agentraxv1alpha1.AgentDeployment) *monitoringv1.ServiceMonitor { labels := agentLabels(ad) @@ -476,6 +656,13 @@ func (r *AgentDeploymentReconciler) desiredServiceMonitor(ad *agentraxv1alpha1.A Selector: metav1.LabelSelector{ MatchLabels: labels, }, + // Carry the two labels used by the HPA ExternalMetric selector into + // every scraped sample. Prometheus will sanitize the dots to underscores + // during ingestion (app.kubernetes.io/name → app_kubernetes_io_name). + TargetLabels: []string{ + "app.kubernetes.io/name", + "app.kubernetes.io/managed-by", + }, Endpoints: []monitoringv1.Endpoint{ { Port: "agent", @@ -502,7 +689,12 @@ func (r *AgentDeploymentReconciler) SetupWithManager(mgr ctrl.Manager) error { bldr := ctrl.NewControllerManagedBy(mgr). For(&agentraxv1alpha1.AgentDeployment{}). Owns(&appsv1.Deployment{}). - Owns(&corev1.Service{}) + Owns(&corev1.Service{}). + Owns(&autoscalingv2.HorizontalPodAutoscaler{}). + Watches( + &agentraxv1alpha1.TenantQuota{}, + enqueueAgentDeploymentsForTenantQuota(mgr.GetClient()), + ) if r.hasServiceMonitorCRD { bldr = bldr.Owns(&monitoringv1.ServiceMonitor{}) diff --git a/internal/controller/agentdeployment_controller_test.go b/internal/controller/agentdeployment_controller_test.go index c63f493..3b6577b 100644 --- a/internal/controller/agentdeployment_controller_test.go +++ b/internal/controller/agentdeployment_controller_test.go @@ -24,8 +24,10 @@ import ( . "github.com/onsi/gomega" appsv1 "k8s.io/api/apps/v1" + autoscalingv2 "k8s.io/api/autoscaling/v2" corev1 "k8s.io/api/core/v1" apierrors "k8s.io/apimachinery/pkg/api/errors" + apimeta "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/client" @@ -135,6 +137,18 @@ func deleteChildResources(key types.NamespacedName) { return apierrors.IsNotFound(k8sClient.Get(ctx, key, &monitoringv1.ServiceMonitor{})) }, testTimeout, testInterval).Should(BeTrue(), "child ServiceMonitor should be deleted") } + + hpa := &autoscalingv2.HorizontalPodAutoscaler{} + err := k8sClient.Get(ctx, key, hpa) + if err != nil && !apierrors.IsNotFound(err) { + Expect(err).NotTo(HaveOccurred(), "unexpected error reading HPA during cleanup") + } + if err == nil { + Expect(k8sClient.Delete(ctx, hpa)).To(Succeed(), "HPA cleanup delete should succeed") + Eventually(func() bool { + return apierrors.IsNotFound(k8sClient.Get(ctx, key, &autoscalingv2.HorizontalPodAutoscaler{})) + }, testTimeout, testInterval).Should(BeTrue(), "child HPA should be deleted") + } } var _ = Describe("AgentDeployment Controller", func() { @@ -336,6 +350,14 @@ var _ = Describe("AgentDeployment Controller", func() { Expect(sm.OwnerReferences[0].Kind).To(Equal("AgentDeployment")) Expect(sm.OwnerReferences[0].Controller).NotTo(BeNil()) Expect(*sm.OwnerReferences[0].Controller).To(BeTrue(), "ServiceMonitor owner reference must have Controller=true") + + // TargetLabels must include both labels used by the HPA ExternalMetric selector + // so that Prometheus carries them into scraped samples and the Adapter + // query's <<.LabelMatchers>> can filter by agent name. + Expect(sm.Spec.TargetLabels).To(ContainElements( + "app.kubernetes.io/name", + "app.kubernetes.io/managed-by", + ), "TargetLabels must propagate HPA selector labels into Prometheus samples") }) It("sets status.phase to Pending initially (no running pods in envtest)", func() { @@ -704,3 +726,384 @@ var _ = Describe("AgentDeployment Controller", func() { }) }) }) + +// ── Phase 3: HPA lifecycle integration tests ────────────────────────────────── + +var _ = Describe("AgentDeployment HPA lifecycle", func() { + Describe("HPA creation and spec", func() { + var key types.NamespacedName + + BeforeEach(func() { + ns := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: "test-hpa"}} + err := k8sClient.Create(ctx, ns) + Expect(err == nil || apierrors.IsAlreadyExists(err)).To(BeTrue()) + key = createAgentDeployment("ad-hpa", "test-hpa", testNginxImage, 8080, 1) + }) + + AfterEach(func() { + deleteAgentDeployment(key) + deleteChildResources(key) + }) + + It("creates a managed HPA after AgentDeployment is reconciled", func() { + hpa := &autoscalingv2.HorizontalPodAutoscaler{} + Eventually(func() error { + return k8sClient.Get(ctx, key, hpa) + }, testTimeout, testInterval).Should(Succeed(), "HPA should be created") + + Expect(hpa.Spec.ScaleTargetRef.Kind).To(Equal("Deployment")) + Expect(hpa.Spec.ScaleTargetRef.Name).To(Equal(key.Name)) + Expect(hpa.Spec.MinReplicas).NotTo(BeNil()) + Expect(*hpa.Spec.MinReplicas).To(Equal(int32(1))) + Expect(hpa.Spec.MaxReplicas).To(Equal(int32(3))) + }) + + It("sets the HPA owner reference to the AgentDeployment", func() { + parent := &agentraxv1alpha1.AgentDeployment{} + Eventually(func() error { + return k8sClient.Get(ctx, key, parent) + }, testTimeout, testInterval).Should(Succeed()) + + hpa := &autoscalingv2.HorizontalPodAutoscaler{} + Eventually(func() error { + return k8sClient.Get(ctx, key, hpa) + }, testTimeout, testInterval).Should(Succeed()) + + Expect(hpa.OwnerReferences).To(HaveLen(1)) + Expect(hpa.OwnerReferences[0].Kind).To(Equal("AgentDeployment")) + Expect(hpa.OwnerReferences[0].Name).To(Equal(key.Name)) + Expect(hpa.OwnerReferences[0].UID).To(Equal(parent.UID)) + Expect(hpa.OwnerReferences[0].Controller).NotTo(BeNil()) + Expect(*hpa.OwnerReferences[0].Controller).To(BeTrue()) + }) + + It("configures an External metric source", func() { + hpa := &autoscalingv2.HorizontalPodAutoscaler{} + Eventually(func() error { + return k8sClient.Get(ctx, key, hpa) + }, testTimeout, testInterval).Should(Succeed()) + + Expect(hpa.Spec.Metrics).To(HaveLen(1)) + Expect(hpa.Spec.Metrics[0].Type).To(Equal(autoscalingv2.ExternalMetricSourceType)) + Expect(hpa.Spec.Metrics[0].External).NotTo(BeNil()) + // createAgentDeployment uses metric: queueDepth + Expect(hpa.Spec.Metrics[0].External.Metric.Name).To(Equal("agentrax_queue_depth")) + }) + + It("sets stabilization windows on the HPA behavior", func() { + hpa := &autoscalingv2.HorizontalPodAutoscaler{} + Eventually(func() error { + return k8sClient.Get(ctx, key, hpa) + }, testTimeout, testInterval).Should(Succeed()) + + Expect(hpa.Spec.Behavior).NotTo(BeNil()) + Expect(hpa.Spec.Behavior.ScaleUp).NotTo(BeNil()) + Expect(hpa.Spec.Behavior.ScaleUp.StabilizationWindowSeconds).NotTo(BeNil()) + Expect(*hpa.Spec.Behavior.ScaleUp.StabilizationWindowSeconds).To(Equal(int32(60))) + Expect(hpa.Spec.Behavior.ScaleDown).NotTo(BeNil()) + Expect(hpa.Spec.Behavior.ScaleDown.StabilizationWindowSeconds).NotTo(BeNil()) + Expect(*hpa.Spec.Behavior.ScaleDown.StabilizationWindowSeconds).To(Equal(int32(300))) + }) + }) + + Describe("HPA self-healing", func() { + var key types.NamespacedName + + BeforeEach(func() { + ns := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: "test-hpa-selfheal"}} + err := k8sClient.Create(ctx, ns) + Expect(err == nil || apierrors.IsAlreadyExists(err)).To(BeTrue()) + key = createAgentDeployment("ad-hpa-sh", "test-hpa-selfheal", testNginxImage, 8080, 1) + }) + + AfterEach(func() { + deleteAgentDeployment(key) + deleteChildResources(key) + }) + + It("recreates the HPA when it is deleted out-of-band", func() { + // Wait for the initial HPA to be created. + hpa := &autoscalingv2.HorizontalPodAutoscaler{} + Eventually(func() error { + return k8sClient.Get(ctx, key, hpa) + }, testTimeout, testInterval).Should(Succeed(), "HPA should appear initially") + origUID := hpa.UID + + // Delete the HPA out-of-band. + Expect(k8sClient.Delete(ctx, hpa)).To(Succeed()) + + // The reconciler should restore it within one reconcile interval with a new UID. + recreated := &autoscalingv2.HorizontalPodAutoscaler{} + Eventually(func(g Gomega) { + g.Expect(k8sClient.Get(ctx, key, recreated)).To(Succeed()) + g.Expect(recreated.UID).NotTo(Equal(origUID)) + }, testTimeout, testInterval).Should(Succeed(), "HPA should be self-healed with a new UID") + + ad := &agentraxv1alpha1.AgentDeployment{} + Expect(k8sClient.Get(ctx, key, ad)).To(Succeed()) + + Expect(recreated.OwnerReferences).To(HaveLen(1)) + Expect(recreated.OwnerReferences[0].Kind).To(Equal("AgentDeployment")) + Expect(recreated.OwnerReferences[0].Name).To(Equal(key.Name)) + Expect(recreated.OwnerReferences[0].UID).To(Equal(ad.UID)) + Expect(recreated.OwnerReferences[0].Controller).NotTo(BeNil()) + Expect(*recreated.OwnerReferences[0].Controller).To(BeTrue()) + }) + }) + + Describe("HPA created with correct spec after AgentDeployment creation", func() { + // Phase 3 DoD: "managed HPA exists after AgentDeployment creation; changes to + // replicas spec update HPA". This block verifies the initial HPA spec correctness — + // metric source, scaleTargetRef, min/max replicas, and stabilization windows. + var key types.NamespacedName + + BeforeEach(func() { + ns := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: "test-hpa-spec"}} + err := k8sClient.Create(ctx, ns) + Expect(err == nil || apierrors.IsAlreadyExists(err)).To(BeTrue()) + // Create with min=2, max=4, metric=queueDepth, target=75. + ensureTenantQuota("test-hpa-spec") + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "ad-hpa-spec", Namespace: "test-hpa-spec"}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: testNginxImage, + Port: 8080, + TenantRef: "team-test", + Replicas: agentraxv1alpha1.ScalingPolicy{ + Min: 2, Max: 4, Metric: "queueDepth", Target: 75, + }, + }, + } + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) + key = types.NamespacedName{Name: "ad-hpa-spec", Namespace: "test-hpa-spec"} + }) + + AfterEach(func() { + deleteAgentDeployment(key) + deleteChildResources(key) + }) + + It("creates an HPA with correct scaleTargetRef pointing to the Deployment", func() { + hpa := &autoscalingv2.HorizontalPodAutoscaler{} + Eventually(func() error { + return k8sClient.Get(ctx, key, hpa) + }, testTimeout, testInterval).Should(Succeed()) + + Expect(hpa.Spec.ScaleTargetRef.APIVersion).To(Equal("apps/v1")) + Expect(hpa.Spec.ScaleTargetRef.Kind).To(Equal("Deployment")) + Expect(hpa.Spec.ScaleTargetRef.Name).To(Equal(key.Name)) + }) + + It("creates an HPA with min and max replicas matching spec", func() { + hpa := &autoscalingv2.HorizontalPodAutoscaler{} + Eventually(func() error { + return k8sClient.Get(ctx, key, hpa) + }, testTimeout, testInterval).Should(Succeed()) + + Expect(hpa.Spec.MinReplicas).NotTo(BeNil()) + Expect(*hpa.Spec.MinReplicas).To(Equal(int32(2)), "minReplicas should match spec.replicas.min") + Expect(hpa.Spec.MaxReplicas).To(Equal(int32(4)), "maxReplicas should match spec.replicas.max") + }) + + It("creates an HPA wired to the agentrax_queue_depth external metric", func() { + hpa := &autoscalingv2.HorizontalPodAutoscaler{} + Eventually(func() error { + return k8sClient.Get(ctx, key, hpa) + }, testTimeout, testInterval).Should(Succeed()) + + Expect(hpa.Spec.Metrics).To(HaveLen(1)) + m := hpa.Spec.Metrics[0] + Expect(m.Type).To(Equal(autoscalingv2.ExternalMetricSourceType)) + Expect(m.External).NotTo(BeNil()) + Expect(m.External.Metric.Name).To(Equal("agentrax_queue_depth")) + Expect(m.External.Target.Type).To(Equal(autoscalingv2.AverageValueMetricType)) + }) + + It("creates an HPA with the correct stabilization windows (60s up, 300s down)", func() { + hpa := &autoscalingv2.HorizontalPodAutoscaler{} + Eventually(func() error { + return k8sClient.Get(ctx, key, hpa) + }, testTimeout, testInterval).Should(Succeed()) + + Expect(hpa.Spec.Behavior).NotTo(BeNil()) + Expect(hpa.Spec.Behavior.ScaleUp).NotTo(BeNil()) + Expect(hpa.Spec.Behavior.ScaleUp.StabilizationWindowSeconds).NotTo(BeNil()) + Expect(*hpa.Spec.Behavior.ScaleUp.StabilizationWindowSeconds).To(Equal(int32(60)), + "scale-up stabilization should be 60s") + Expect(hpa.Spec.Behavior.ScaleDown).NotTo(BeNil()) + Expect(hpa.Spec.Behavior.ScaleDown.StabilizationWindowSeconds).NotTo(BeNil()) + Expect(*hpa.Spec.Behavior.ScaleDown.StabilizationWindowSeconds).To(Equal(int32(300)), + "scale-down stabilization should be 300s") + }) + }) + + Describe("HPA spec update on replicas change", func() { + var key types.NamespacedName + + BeforeEach(func() { + ns := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: "test-hpa-update"}} + err := k8sClient.Create(ctx, ns) + Expect(err == nil || apierrors.IsAlreadyExists(err)).To(BeTrue()) + key = createAgentDeployment("ad-hpa-upd", "test-hpa-update", testNginxImage, 8080, 1) + }) + + AfterEach(func() { + deleteAgentDeployment(key) + deleteChildResources(key) + }) + + It("updates HPA maxReplicas when spec.replicas.max changes", func() { + // Wait for initial HPA. + Eventually(func() error { + return k8sClient.Get(ctx, key, &autoscalingv2.HorizontalPodAutoscaler{}) + }, testTimeout, testInterval).Should(Succeed()) + + // Patch spec.replicas.max to 5. + ad := &agentraxv1alpha1.AgentDeployment{} + Expect(k8sClient.Get(ctx, key, ad)).To(Succeed()) + patch := client.MergeFrom(ad.DeepCopy()) + ad.Spec.Replicas.Max = 5 + Expect(k8sClient.Patch(ctx, ad, patch)).To(Succeed()) + + // Expect the HPA to be updated to maxReplicas=5 + // (quota headroom in tests is maxReplicasPerAgent=10, so no capping). + Eventually(func(g Gomega) { + hpa := &autoscalingv2.HorizontalPodAutoscaler{} + g.Expect(k8sClient.Get(ctx, key, hpa)).To(Succeed()) + g.Expect(hpa.Spec.MaxReplicas).To(Equal(int32(5))) + }, testTimeout, testInterval).Should(Succeed()) + + // Verify QuotaLimited condition is NOT True — max=5 is within headroom of 10. + // Use Consistently so a transient True that later settles does not go undetected. + Consistently(func(g Gomega) { + latest := &agentraxv1alpha1.AgentDeployment{} + g.Expect(k8sClient.Get(ctx, key, latest)).To(Succeed()) + c := apimeta.FindStatusCondition(latest.Status.Conditions, agentraxv1alpha1.ConditionQuotaLimited) + g.Expect(c != nil && c.Status == metav1.ConditionTrue).To(BeFalse(), + "QuotaLimited should never be True when max (5) <= headroom (10)") + }, 3*time.Second, testInterval).Should(Succeed()) + }) + }) + + Describe("QuotaLimited condition when HPA is capped", func() { + var key types.NamespacedName + const nsName = "test-hpa-quota" + + BeforeEach(func() { + ns := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: nsName}} + err := k8sClient.Create(ctx, ns) + Expect(err == nil || apierrors.IsAlreadyExists(err)).To(BeTrue()) + + // Start with a generous quota (maxReplicasPerAgent=5) so the initial + // create passes webhook validation. + tq := &agentraxv1alpha1.TenantQuota{ + ObjectMeta: metav1.ObjectMeta{Name: "team-quota-test", Namespace: nsName}, + Spec: agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 10, + MaxGPUs: 0, + MaxTotalReplicas: 5, + MaxReplicasPerAgent: 5, + }, + } + err = k8sClient.Create(ctx, tq) + if apierrors.IsAlreadyExists(err) { + existing := &agentraxv1alpha1.TenantQuota{} + Expect(k8sClient.Get(ctx, namespacedName("team-quota-test", nsName), existing)).To(Succeed()) + existing.Spec = tq.Spec + Expect(k8sClient.Update(ctx, existing)).To(Succeed()) + } else { + Expect(err).NotTo(HaveOccurred()) + } + + // Create an AgentDeployment with max=5, which fits the initial quota. + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "ad-quota-capped", Namespace: nsName}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: testNginxImage, + Port: 8080, + TenantRef: "team-quota-test", + Replicas: agentraxv1alpha1.ScalingPolicy{ + Min: 1, + Max: 5, + Metric: "queueDepth", + Target: 50, + }, + }, + } + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) + key = types.NamespacedName{Name: ad.Name, Namespace: ad.Namespace} + }) + + AfterEach(func() { + deleteAgentDeployment(key) + deleteChildResources(key) + // Delete the TenantQuota so each It block starts from a clean slate. + // Without this, the first It (which lowers the ceiling) bleeds state + // into the next BeforeEach, causing webhook rejection on the AD create. + tqKey := namespacedName("team-quota-test", nsName) + tq := &agentraxv1alpha1.TenantQuota{} + if err := k8sClient.Get(ctx, tqKey, tq); err != nil { + if !apierrors.IsNotFound(err) { + Expect(err).NotTo(HaveOccurred()) + } + } else { + Expect(client.IgnoreNotFound(k8sClient.Delete(ctx, tq))).To(Succeed()) + Eventually(func() bool { + return apierrors.IsNotFound(k8sClient.Get(ctx, tqKey, &agentraxv1alpha1.TenantQuota{})) + }, testTimeout, testInterval).Should(BeTrue(), "TenantQuota should be deleted") + } + }) + + It("caps HPA maxReplicas and sets QuotaLimited condition when quota is lowered", func() { + // Wait for initial HPA with maxReplicas=5 (uncapped). + Eventually(func(g Gomega) { + hpa := &autoscalingv2.HorizontalPodAutoscaler{} + g.Expect(k8sClient.Get(ctx, key, hpa)).To(Succeed()) + g.Expect(hpa.Spec.MaxReplicas).To(Equal(int32(5))) + }, testTimeout, testInterval).Should(Succeed()) + + // Lower the TenantQuota ceiling to maxReplicasPerAgent=2, maxTotalReplicas=2. + // This causes the reconciler to cap HPA.maxReplicas at 2 and set QuotaLimited. + tq := &agentraxv1alpha1.TenantQuota{} + Expect(k8sClient.Get(ctx, namespacedName("team-quota-test", nsName), tq)).To(Succeed()) + patch := tq.DeepCopy() + patch.Spec.MaxReplicasPerAgent = 2 + patch.Spec.MaxTotalReplicas = 2 + Expect(k8sClient.Patch(ctx, patch, client.MergeFrom(tq))).To(Succeed()) + + // HPA maxReplicas must be reduced to the new quota ceiling (2). + Eventually(func(g Gomega) { + hpa := &autoscalingv2.HorizontalPodAutoscaler{} + g.Expect(k8sClient.Get(ctx, key, hpa)).To(Succeed()) + g.Expect(hpa.Spec.MaxReplicas).To(Equal(int32(2))) + }, testTimeout, testInterval).Should(Succeed(), + "HPA maxReplicas should be capped at the new quota ceiling") + + // QuotaLimited condition must be True because spec.max (5) > headroom (2). + Eventually(func(g Gomega) { + latest := &agentraxv1alpha1.AgentDeployment{} + g.Expect(k8sClient.Get(ctx, key, latest)).To(Succeed()) + c := apimeta.FindStatusCondition(latest.Status.Conditions, agentraxv1alpha1.ConditionQuotaLimited) + g.Expect(c).NotTo(BeNil()) + g.Expect(c.Status).To(Equal(metav1.ConditionTrue)) + }, testTimeout, testInterval).Should(Succeed(), + "QuotaLimited condition should be True when HPA is capped by quota") + }) + + It("QuotaLimited condition is absent when max equals headroom", func() { + // With the generous initial quota (max=5, ceiling=5), the condition + // must never be True. Use Consistently to guard against transient flap. + Eventually(func() error { + return k8sClient.Get(ctx, key, &autoscalingv2.HorizontalPodAutoscaler{}) + }, testTimeout, testInterval).Should(Succeed(), "wait for first reconcile") + + Consistently(func(g Gomega) { + latest := &agentraxv1alpha1.AgentDeployment{} + g.Expect(k8sClient.Get(ctx, key, latest)).To(Succeed()) + c := apimeta.FindStatusCondition(latest.Status.Conditions, agentraxv1alpha1.ConditionQuotaLimited) + g.Expect(c != nil && c.Status == metav1.ConditionTrue).To(BeFalse(), + "QuotaLimited should never be True when max == headroom") + }, 3*time.Second, testInterval).Should(Succeed()) + }) + }) +}) diff --git a/internal/controller/enqueue_handlers.go b/internal/controller/enqueue_handlers.go index 87e64f3..d403b1b 100644 --- a/internal/controller/enqueue_handlers.go +++ b/internal/controller/enqueue_handlers.go @@ -22,6 +22,7 @@ import ( "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/handler" + "sigs.k8s.io/controller-runtime/pkg/log" "sigs.k8s.io/controller-runtime/pkg/reconcile" agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" @@ -47,3 +48,35 @@ func enqueueTenantQuota() handler.EventHandler { } }) } + +// enqueueAgentDeploymentsForTenantQuota returns an EventHandler that maps every +// TenantQuota event to reconcile requests for all AgentDeployments in the same +// namespace that reference that TenantQuota. This ensures that lowering or +// raising quota ceilings updates HPA maxReplicas and QuotaLimited conditions +// across all affected agents without requiring out-of-band edits to each AD. +func enqueueAgentDeploymentsForTenantQuota(c client.Client) handler.EventHandler { + return handler.EnqueueRequestsFromMapFunc(func(ctx context.Context, obj client.Object) []reconcile.Request { + tq, ok := obj.(*agentraxv1alpha1.TenantQuota) + if !ok { + return nil + } + list := &agentraxv1alpha1.AgentDeploymentList{} + if err := c.List(ctx, list, client.InNamespace(tq.Namespace)); err != nil { + log.FromContext(ctx).Error(err, "failed to list AgentDeployments for TenantQuota watch", + "tenantQuota", tq.Name, "namespace", tq.Namespace) + return nil + } + var reqs []reconcile.Request + for _, ad := range list.Items { + if ad.Spec.TenantRef == tq.Name { + reqs = append(reqs, reconcile.Request{ + NamespacedName: types.NamespacedName{ + Namespace: ad.Namespace, + Name: ad.Name, + }, + }) + } + } + return reqs + }) +} diff --git a/internal/controller/suite_test.go b/internal/controller/suite_test.go index 41eca0b..f58d907 100644 --- a/internal/controller/suite_test.go +++ b/internal/controller/suite_test.go @@ -27,6 +27,7 @@ import ( . "github.com/onsi/gomega" appsv1 "k8s.io/api/apps/v1" + autoscalingv2 "k8s.io/api/autoscaling/v2" corev1 "k8s.io/api/core/v1" apiextensionsv1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1" "k8s.io/client-go/kubernetes/scheme" @@ -66,6 +67,7 @@ var testReconciler *AgentDeploymentReconciler // and the validating webhook in integration tests. var testEnforcer *quota.Enforcer +// TestControllers is the Ginkgo test suite runner for controller integration tests. func TestControllers(t *testing.T) { RegisterFailHandler(Fail) RunSpecs(t, "Controller Suite") @@ -110,6 +112,8 @@ var _ = BeforeSuite(func() { Expect(agentraxv1alpha1.AddToScheme(scheme.Scheme)).To(Succeed()) Expect(appsv1.AddToScheme(scheme.Scheme)).To(Succeed()) Expect(corev1.AddToScheme(scheme.Scheme)).To(Succeed()) + // Register autoscaling/v2 so HPA objects can be created/read in integration tests. + Expect(autoscalingv2.AddToScheme(scheme.Scheme)).To(Succeed()) // Register prometheus-operator types so the reconciler can handle ServiceMonitor objects. Expect(monitoringv1.AddToScheme(scheme.Scheme)).To(Succeed()) // Register apiextensions types so serviceMonitorCRDExists can decode CRD objects @@ -140,8 +144,9 @@ var _ = BeforeSuite(func() { testEnforcer = quota.NewEnforcer(quota.DefaultGPUResourceName) testReconciler = &AgentDeploymentReconciler{ - Client: mgr.GetClient(), - Scheme: mgr.GetScheme(), + Client: mgr.GetClient(), + Scheme: mgr.GetScheme(), + GPUResourceName: quota.DefaultGPUResourceName, } Expect(testReconciler.SetupWithManager(mgr)).To(Succeed()) diff --git a/internal/controller/tenantquota_controller_test.go b/internal/controller/tenantquota_controller_test.go index 08ce49d..808c4ff 100644 --- a/internal/controller/tenantquota_controller_test.go +++ b/internal/controller/tenantquota_controller_test.go @@ -350,6 +350,7 @@ var _ = Describe("TenantQuota Controller", func() { // ── Test helpers ────────────────────────────────────────────────────────────── +// makeTQ creates a TenantQuota object with the provided limits for test fixtures. func makeTQ(name, ns string, maxAgents, maxGPUs, maxTotalReplicas, maxReplicasPerAgent int32) *agentraxv1alpha1.TenantQuota { //nolint:unparam return &agentraxv1alpha1.TenantQuota{ ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: ns}, diff --git a/internal/controller/webhook_integration_test.go b/internal/controller/webhook_integration_test.go new file mode 100644 index 0000000..33187d2 --- /dev/null +++ b/internal/controller/webhook_integration_test.go @@ -0,0 +1,507 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package controller + +// webhook_integration_test.go verifies Phase-2 admission-webhook behaviour +// end-to-end through the real envtest webhook server (not a fake client). +// The envtest manager started in suite_test.go already registers the webhook, +// so every k8sClient.Create / k8sClient.Update call here goes through the +// full validation and defaulting path. + +import ( + "context" + "sync" + "sync/atomic" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + "sigs.k8s.io/controller-runtime/pkg/client" + + agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" +) + +// ── helpers local to this file ──────────────────────────────────────────────── + +// whCreateNS ensures the namespace exists; idempotent. +func whCreateNS(ctx context.Context, name string) { + ns := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: name}} + err := k8sClient.Create(ctx, ns) + Expect(err == nil || apierrors.IsAlreadyExists(err)).To(BeTrue(), + "ensuring namespace %s: %v", name, err) +} + +// whCreateTQ creates a TenantQuota in the given namespace with the provided spec. +func whCreateTQ(ctx context.Context, name, namespace string, spec agentraxv1alpha1.TenantQuotaSpec) { + tq := &agentraxv1alpha1.TenantQuota{ + ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: namespace}, + Spec: spec, + } + err := k8sClient.Create(ctx, tq) + Expect(err == nil || apierrors.IsAlreadyExists(err)).To(BeTrue(), + "creating TenantQuota %s/%s: %v", namespace, name, err) +} + +// whMinimalAD returns an AgentDeployment with a valid spec for the given TQ. +func whMinimalAD(name, namespace, tenantRef string, maxReplicas int32) *agentraxv1alpha1.AgentDeployment { + return &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: namespace}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: testNginxImage, + Port: 8080, + TenantRef: tenantRef, + Replicas: agentraxv1alpha1.ScalingPolicy{ + Min: 1, + Max: maxReplicas, + Metric: "queueDepth", + Target: 50, + }, + }, + } +} + +// whGPUAD returns an AgentDeployment that requests gpus GPU units per replica. +func whGPUAD(name, namespace, tenantRef string, gpus int64) *agentraxv1alpha1.AgentDeployment { + ad := whMinimalAD(name, namespace, tenantRef, 1) + ad.Spec.Resources = corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + corev1.ResourceName("nvidia.com/gpu"): *resource.NewQuantity(gpus, resource.DecimalSI), + }, + } + return ad +} + +// whCleanupAD deletes an AD (ignoring not-found) and strips finalizers if needed. +// Envtest does not run the GC controller, so the caller is responsible for +// cleaning up child resources separately if required. +func whCleanupAD(ctx context.Context, key types.NamespacedName) { + ad := &agentraxv1alpha1.AgentDeployment{} + if err := k8sClient.Get(ctx, key, ad); err != nil { + return + } + // Strip finalizer so deletion is not blocked by the controller. + if len(ad.Finalizers) > 0 { + patch := ad.DeepCopy() + patch.Finalizers = nil + _ = k8sClient.Patch(ctx, patch, client.MergeFrom(ad)) + } + _ = k8sClient.Delete(ctx, ad) + Eventually(func() bool { + return apierrors.IsNotFound(k8sClient.Get(ctx, key, &agentraxv1alpha1.AgentDeployment{})) + }, testTimeout, testInterval).Should(BeTrue(), "AD %s/%s should be fully gone", key.Namespace, key.Name) +} + +// ── Webhook integration tests ───────────────────────────────────────────────── + +var _ = Describe("Admission Webhook (integration)", func() { + ctx := context.Background() + + // ── 1. Over-maxAgents rejection ────────────────────────────────────────── + + Describe("rejects AgentDeployment exceeding maxAgents", func() { + const ( + ns = "wh-agents" + tqName = "tq-agents" + ) + + BeforeEach(func() { + whCreateNS(ctx, ns) + // Allow exactly 1 agent. + whCreateTQ(ctx, tqName, ns, agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 1, MaxGPUs: 0, MaxTotalReplicas: 10, MaxReplicasPerAgent: 5, + }) + }) + + AfterEach(func() { + whCleanupAD(ctx, namespacedName("ad-first", ns)) + deleteChildResources(namespacedName("ad-first", ns)) + }) + + It("rejects a second create that would exceed maxAgents", func() { + // First create must succeed (1 agent, within limit). + ad1 := whMinimalAD("ad-first", ns, tqName, 1) + Expect(k8sClient.Create(ctx, ad1)).To(Succeed()) + + // Wait for the TenantQuota reconciler to reflect usedAgents=1. + Eventually(func() int32 { + tq := &agentraxv1alpha1.TenantQuota{} + if err := k8sClient.Get(ctx, namespacedName(tqName, ns), tq); err != nil { + return -1 + } + return tq.Status.UsedAgents + }, testTimeout, testInterval).Should(Equal(int32(1))) + + // Second create should be rejected by the webhook. + ad2 := whMinimalAD("ad-second", ns, tqName, 1) + err := k8sClient.Create(ctx, ad2) + Expect(err).To(HaveOccurred(), "second create should be rejected (maxAgents=1 already used)") + Expect(apierrors.IsInvalid(err) || apierrors.IsForbidden(err)).To(BeTrue(), + "expected 422 Invalid or 403 Forbidden, got: %v", err) + }) + }) + + // ── 2. Over-maxGPUs rejection ───────────────────────────────────────────── + + Describe("rejects AgentDeployment exceeding maxGPUs", func() { + const ( + ns = "wh-gpus" + tqName = "tq-gpus" + ) + + BeforeEach(func() { + whCreateNS(ctx, ns) + // Allow exactly 1 GPU total. + whCreateTQ(ctx, tqName, ns, agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 5, MaxGPUs: 1, MaxTotalReplicas: 10, MaxReplicasPerAgent: 5, + }) + }) + + AfterEach(func() { + whCleanupAD(ctx, namespacedName("ad-gpu-ok", ns)) + deleteChildResources(namespacedName("ad-gpu-ok", ns)) + }) + + It("rejects an AD that would exceed maxGPUs", func() { + // First AD requests 1 GPU — within limit. + ad1 := whGPUAD("ad-gpu-ok", ns, tqName, 1) + Expect(k8sClient.Create(ctx, ad1)).To(Succeed()) + + // Wait for TQ status to reflect 1 GPU used. + Eventually(func() int32 { + tq := &agentraxv1alpha1.TenantQuota{} + if err := k8sClient.Get(ctx, namespacedName(tqName, ns), tq); err != nil { + return -1 + } + return tq.Status.UsedGPUs + }, testTimeout, testInterval).Should(Equal(int32(1))) + + // Second AD requests 1 more GPU — must be rejected (total=2 > maxGPUs=1). + ad2 := whGPUAD("ad-gpu-over", ns, tqName, 1) + err := k8sClient.Create(ctx, ad2) + Expect(err).To(HaveOccurred(), "should be rejected: over maxGPUs") + Expect(apierrors.IsInvalid(err) || apierrors.IsForbidden(err)).To(BeTrue()) + }) + }) + + // ── 3. Over-maxTotalReplicas rejection ─────────────────────────────────── + + Describe("rejects AgentDeployment exceeding maxTotalReplicas", func() { + const ( + ns = "wh-total-rep" + tqName = "tq-total-rep" + ) + + BeforeEach(func() { + whCreateNS(ctx, ns) + // Allow max 3 total replicas. + whCreateTQ(ctx, tqName, ns, agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 5, MaxGPUs: 0, MaxTotalReplicas: 3, MaxReplicasPerAgent: 3, + }) + }) + + AfterEach(func() { + whCleanupAD(ctx, namespacedName("ad-rep-ok", ns)) + deleteChildResources(namespacedName("ad-rep-ok", ns)) + }) + + It("rejects an AD that would push total replicas past the ceiling", func() { + // First AD uses max=3 — fills the budget exactly. + ad1 := whMinimalAD("ad-rep-ok", ns, tqName, 3) + Expect(k8sClient.Create(ctx, ad1)).To(Succeed()) + + Eventually(func() int32 { + tq := &agentraxv1alpha1.TenantQuota{} + if err := k8sClient.Get(ctx, namespacedName(tqName, ns), tq); err != nil { + return -1 + } + return tq.Status.UsedTotalReplicas + }, testTimeout, testInterval).Should(Equal(int32(3))) + + // Any new AD (even max=1) should be rejected. + ad2 := whMinimalAD("ad-rep-over", ns, tqName, 1) + err := k8sClient.Create(ctx, ad2) + Expect(err).To(HaveOccurred(), "should be rejected: over maxTotalReplicas") + Expect(apierrors.IsInvalid(err) || apierrors.IsForbidden(err)).To(BeTrue()) + }) + }) + + // ── 4. Over-maxReplicasPerAgent rejection ──────────────────────────────── + + Describe("rejects spec.replicas.max > maxReplicasPerAgent", func() { + const ( + ns = "wh-per-agent" + tqName = "tq-per-agent" + ) + + BeforeEach(func() { + whCreateNS(ctx, ns) + whCreateTQ(ctx, tqName, ns, agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 5, MaxGPUs: 0, MaxTotalReplicas: 50, MaxReplicasPerAgent: 2, + }) + }) + + It("rejects an AD with max replicas exceeding the per-agent ceiling", func() { + // max=5 > maxReplicasPerAgent=2 — must be rejected at admission. + ad := whMinimalAD("ad-per-agent-over", ns, tqName, 5) + err := k8sClient.Create(ctx, ad) + Expect(err).To(HaveOccurred(), "max > maxReplicasPerAgent should be rejected") + Expect(apierrors.IsInvalid(err) || apierrors.IsForbidden(err)).To(BeTrue()) + }) + }) + + // ── 5. Concurrent near-limit creates ───────────────────────────────────── + // Two goroutines create simultaneously when only one slot remains. + // Exactly one must succeed and one must be rejected via the in-flight + // reservation mechanism in internal/quota.Enforcer. + + Describe("concurrent near-limit creates (race prevention)", func() { + const ( + ns = "wh-concurrent" + tqName = "tq-concurrent" + ) + + BeforeEach(func() { + whCreateNS(ctx, ns) + // Allow exactly 1 agent — the race window. + whCreateTQ(ctx, tqName, ns, agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 1, MaxGPUs: 0, MaxTotalReplicas: 10, MaxReplicasPerAgent: 5, + }) + }) + + AfterEach(func() { + for _, name := range []string{"ad-race-a", "ad-race-b"} { + whCleanupAD(ctx, namespacedName(name, ns)) + deleteChildResources(namespacedName(name, ns)) + } + }) + + It("allows exactly one of two simultaneous creates when only one slot remains", func() { + var ( + successCount int64 + failCount int64 + wg sync.WaitGroup + ) + + // Start barrier: ensure both goroutines are scheduled before either + // calls k8sClient.Create, maximising the chance of genuine concurrency + // at the webhook admission layer. + start := make(chan struct{}) + errs := make([]error, 2) + for i, name := range []string{"ad-race-a", "ad-race-b"} { + i, name := i, name + wg.Add(1) + go func() { + defer wg.Done() + <-start // wait until both goroutines are ready + ad := whMinimalAD(name, ns, tqName, 1) + err := k8sClient.Create(ctx, ad) + errs[i] = err + if err == nil { + atomic.AddInt64(&successCount, 1) + } else { + atomic.AddInt64(&failCount, 1) + } + }() + } + close(start) // release both goroutines simultaneously + wg.Wait() + + Expect(successCount).To(Equal(int64(1)), + "exactly one concurrent create should succeed (in-flight reservation protects the quota)") + Expect(failCount).To(Equal(int64(1)), + "exactly one concurrent create should be rejected by the webhook") + + // The rejection must come from quota admission, not an unrelated failure. + for _, err := range errs { + if err == nil { + continue + } + Expect(apierrors.IsInvalid(err) || apierrors.IsForbidden(err)).To(BeTrue(), + "expected 422 Invalid or 403 Forbidden from quota admission, got: %v", err) + } + }) + }) + + // ── 6. TenantQuota status accuracy ─────────────────────────────────────── + + Describe("TenantQuota status accurately reflects live usage", func() { + const ( + ns = "wh-tq-status" + tqName = "tq-status" + ) + + BeforeEach(func() { + whCreateNS(ctx, ns) + whCreateTQ(ctx, tqName, ns, agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 5, MaxGPUs: 0, MaxTotalReplicas: 20, MaxReplicasPerAgent: 5, + }) + }) + + It("increments usedAgents on create and decrements after deletion", func() { + key := types.NamespacedName{Name: "ad-tq-track", Namespace: ns} + ad := whMinimalAD(key.Name, ns, tqName, 2) + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) + + // UsedAgents must reach 1 and UsedTotalReplicas must reach 2. + Eventually(func() agentraxv1alpha1.TenantQuotaStatus { + tq := &agentraxv1alpha1.TenantQuota{} + _ = k8sClient.Get(ctx, namespacedName(tqName, ns), tq) + return tq.Status + }, testTimeout, testInterval).Should(And( + WithTransform(func(s agentraxv1alpha1.TenantQuotaStatus) int32 { return s.UsedAgents }, + Equal(int32(1))), + WithTransform(func(s agentraxv1alpha1.TenantQuotaStatus) int32 { return s.UsedTotalReplicas }, + Equal(int32(2))), + ), "TQ status should reflect the newly created AD") + + // Delete the AD; status should drop back to 0. + whCleanupAD(ctx, key) + deleteChildResources(key) + + Eventually(func() int32 { + tq := &agentraxv1alpha1.TenantQuota{} + _ = k8sClient.Get(ctx, namespacedName(tqName, ns), tq) + return tq.Status.UsedAgents + }, testTimeout, testInterval).Should(Equal(int32(0)), + "usedAgents should drop to 0 after the AD is deleted") + }) + }) + + // ── 7. Mutating webhook defaults ───────────────────────────────────────── + + Describe("mutating webhook applies defaults", func() { + const ( + ns = "wh-defaults" + tqName = "tq-defaults" + ) + + BeforeEach(func() { + whCreateNS(ctx, ns) + whCreateTQ(ctx, tqName, ns, agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 5, MaxGPUs: 0, MaxTotalReplicas: 20, MaxReplicasPerAgent: 5, + }) + }) + + AfterEach(func() { + whCleanupAD(ctx, namespacedName("ad-defaults", ns)) + deleteChildResources(namespacedName("ad-defaults", ns)) + }) + + It("defaults port=8080, strategy=Recreate, and resources when omitted", func() { + // Create with only the required fields — no port, no rollout, no resources. + ad := &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "ad-defaults", Namespace: ns}, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: testNginxImage, + TenantRef: tqName, + Replicas: agentraxv1alpha1.ScalingPolicy{ + Min: 1, Max: 2, Metric: "queueDepth", Target: 50, + }, + }, + } + Expect(k8sClient.Create(ctx, ad)).To(Succeed()) + + // Re-read to see server-side-applied defaults. + got := &agentraxv1alpha1.AgentDeployment{} + Expect(k8sClient.Get(ctx, namespacedName("ad-defaults", ns), got)).To(Succeed()) + + Expect(got.Spec.Port).To(Equal(int32(8080)), "port should be defaulted to 8080") + Expect(got.Spec.Rollout.Strategy).To(Equal("Recreate"), + "strategy should be defaulted to Recreate") + Expect(got.Spec.Resources.Requests).NotTo(BeEmpty(), + "resources.requests should be defaulted") + Expect(got.Spec.Resources.Limits).NotTo(BeEmpty(), + "resources.limits should be defaulted") + + // Verify the specific CPU default (100m request). + cpuReq := got.Spec.Resources.Requests[corev1.ResourceCPU] + Expect(cpuReq.Cmp(resource.MustParse("100m"))).To(Equal(0), + "default CPU request should be 100m, got %s", cpuReq.String()) + }) + }) + + // ── 8. OverQuota condition set when quota is lowered below usage ────────── + + Describe("OverQuota condition when quota is lowered below existing usage", func() { + const ( + ns = "wh-overquota" + tqName = "tq-overquota" + ) + + BeforeEach(func() { + whCreateNS(ctx, ns) + // Start with generous quota. + whCreateTQ(ctx, tqName, ns, agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 3, MaxGPUs: 0, MaxTotalReplicas: 15, MaxReplicasPerAgent: 5, + }) + }) + + AfterEach(func() { + for _, name := range []string{"ad-oq-1", "ad-oq-2"} { + whCleanupAD(ctx, namespacedName(name, ns)) + deleteChildResources(namespacedName(name, ns)) + } + }) + + It("sets OverQuota condition without deleting existing ADs", func() { + // Create two agents (both succeed under the generous quota). + Expect(k8sClient.Create(ctx, whMinimalAD("ad-oq-1", ns, tqName, 1))).To(Succeed()) + Expect(k8sClient.Create(ctx, whMinimalAD("ad-oq-2", ns, tqName, 1))).To(Succeed()) + + // Wait for both to be reflected in TQ status. + Eventually(func() int32 { + tq := &agentraxv1alpha1.TenantQuota{} + _ = k8sClient.Get(ctx, namespacedName(tqName, ns), tq) + return tq.Status.UsedAgents + }, testTimeout, testInterval).Should(Equal(int32(2))) + + // Lower maxAgents to 1 (below current usage of 2). + tq := &agentraxv1alpha1.TenantQuota{} + Expect(k8sClient.Get(ctx, namespacedName(tqName, ns), tq)).To(Succeed()) + patched := tq.DeepCopy() + patched.Spec.MaxAgents = 1 + Expect(k8sClient.Patch(ctx, patched, client.MergeFrom(tq))).To(Succeed()) + + // Expect the OverQuota condition to be set. + Eventually(func() bool { + latest := &agentraxv1alpha1.TenantQuota{} + if err := k8sClient.Get(ctx, namespacedName(tqName, ns), latest); err != nil { + return false + } + for _, c := range latest.Status.Conditions { + if c.Type == agentraxv1alpha1.ConditionOverQuota && + c.Status == metav1.ConditionTrue { + return true + } + } + return false + }, testTimeout, testInterval).Should(BeTrue(), + "OverQuota condition should be set after quota is lowered below usage") + + // Both ADs must still exist — no forced deletions. + Expect(k8sClient.Get(ctx, namespacedName("ad-oq-1", ns), + &agentraxv1alpha1.AgentDeployment{})).To(Succeed()) + Expect(k8sClient.Get(ctx, namespacedName("ad-oq-2", ns), + &agentraxv1alpha1.AgentDeployment{})).To(Succeed()) + }) + }) +}) diff --git a/internal/metrics/prometheus.go b/internal/metrics/prometheus.go new file mode 100644 index 0000000..2a9e794 --- /dev/null +++ b/internal/metrics/prometheus.go @@ -0,0 +1,260 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package metrics provides a lightweight Prometheus HTTP query client shared +// by the scaling and rollout packages. It wraps the raw Prometheus API so +// that callers can perform PromQL instant and range queries without importing +// the full Prometheus client library. +package metrics + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strconv" + "strings" + "time" +) + +// defaultTimeout is the per-request HTTP timeout used when no custom timeout is set. +const defaultTimeout = 10 * time.Second + +// Client is a lightweight HTTP client for the Prometheus query API. +// Create one with NewClient; the zero value is not usable. +type Client struct { + baseURL string + httpClient *http.Client +} + +// Option is a functional option for NewClient. +type Option func(*Client) + +// WithTimeout overrides the default 10-second per-request timeout. +func WithTimeout(d time.Duration) Option { + return func(c *Client) { + c.httpClient.Timeout = d + } +} + +// NewClient returns a Client pointed at the given Prometheus base URL +// (e.g. "http://prometheus.monitoring.svc:9090"). +// Trailing slashes on baseURL are stripped so that URL concatenation always +// produces a single-slash boundary (e.g. baseURL + "/api/v1/query"). +func NewClient(baseURL string, opts ...Option) *Client { + c := &Client{ + baseURL: strings.TrimRight(baseURL, "/"), + httpClient: &http.Client{Timeout: defaultTimeout}, + } + for _, o := range opts { + o(c) + } + return c +} + +// QueryScalar executes a PromQL instant query and returns the scalar result. +// It returns an error if the query fails, if the result is empty, or if the +// returned value is not a scalar or single-element vector. +// +// Phase 4 (canary rollout) will call this to evaluate threshold queries during +// pause windows. +func (c *Client) QueryScalar(ctx context.Context, query string) (float64, error) { + u, err := url.Parse(c.baseURL + "/api/v1/query") + if err != nil { + return 0, fmt.Errorf("parsing Prometheus URL: %w", err) + } + + q := u.Query() + q.Set("query", query) + u.RawQuery = q.Encode() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil) + if err != nil { + return 0, fmt.Errorf("building Prometheus request: %w", err) + } + + resp, err := c.httpClient.Do(req) + if err != nil { + return 0, fmt.Errorf("querying Prometheus: %w", err) + } + defer resp.Body.Close() //nolint:errcheck + + body, err := io.ReadAll(resp.Body) + if err != nil { + return 0, fmt.Errorf("reading Prometheus response: %w", err) + } + if resp.StatusCode != http.StatusOK { + return 0, fmt.Errorf("Prometheus returned HTTP %d: %s", resp.StatusCode, body) + } + + return parseScalarFromQueryResponse(body) +} + +// QueryRange executes a PromQL range query and returns the last value +// for the first returned series. start/end/step follow Prometheus range +// query semantics. +// +// This is used by the canary rollout controller (Phase 4) to evaluate +// error rate and p99 latency over a look-back window. +func (c *Client) QueryRange(ctx context.Context, query string, start, end time.Time, step time.Duration) (float64, error) { + u, err := url.Parse(c.baseURL + "/api/v1/query_range") + if err != nil { + return 0, fmt.Errorf("parsing Prometheus URL: %w", err) + } + + q := u.Query() + q.Set("query", query) + q.Set("start", formatTime(start)) + q.Set("end", formatTime(end)) + q.Set("step", formatDuration(step)) + u.RawQuery = q.Encode() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil) + if err != nil { + return 0, fmt.Errorf("building Prometheus range request: %w", err) + } + + resp, err := c.httpClient.Do(req) + if err != nil { + return 0, fmt.Errorf("querying Prometheus range: %w", err) + } + defer resp.Body.Close() //nolint:errcheck + + body, err := io.ReadAll(resp.Body) + if err != nil { + return 0, fmt.Errorf("reading Prometheus range response: %w", err) + } + if resp.StatusCode != http.StatusOK { + return 0, fmt.Errorf("Prometheus returned HTTP %d: %s", resp.StatusCode, body) + } + + return parseLastValueFromRangeResponse(body) +} + +// ── Internal response parsing ───────────────────────────────────────────────── + +// prometheusResponse is a partial deserialization of the Prometheus HTTP API +// query response. Only the fields needed for scalar extraction are decoded. +type prometheusResponse struct { + Status string `json:"status"` + Data struct { + ResultType string `json:"resultType"` + Result []json.RawMessage `json:"result"` + } `json:"data"` +} + +// parseScalarFromQueryResponse extracts a single float64 from a Prometheus +// instant query response body. It handles both "vector" and "scalar" result +// types. +func parseScalarFromQueryResponse(body []byte) (float64, error) { + var r prometheusResponse + if err := json.Unmarshal(body, &r); err != nil { + return 0, fmt.Errorf("decoding Prometheus response: %w", err) + } + if r.Status != "success" { + return 0, fmt.Errorf("Prometheus query status: %q", r.Status) + } + + switch r.Data.ResultType { + case "scalar": + // Scalar result: Data.Result is a [timestamp, "value"] pair. + // Require exactly two elements to guard against a malformed payload; + // we read the value string from index 1 (index 0 is the Unix timestamp). + if len(r.Data.Result) != 2 { + return 0, fmt.Errorf("Prometheus scalar result must contain exactly 2 elements, got %d", len(r.Data.Result)) + } + var valStr string + if err := json.Unmarshal(r.Data.Result[1], &valStr); err != nil { + return 0, fmt.Errorf("decoding scalar value string: %w", err) + } + return strconv.ParseFloat(valStr, 64) + case "vector": + if len(r.Data.Result) != 1 { + return 0, fmt.Errorf("Prometheus vector result must contain exactly 1 element, got %d (metric may not exist yet or query returned multiple series)", len(r.Data.Result)) + } + // Vector result: each element is {"metric":{},"value":[timestamp,"value"]}. + return extractValueFromVectorElement(r.Data.Result[0]) + default: + return 0, fmt.Errorf("unsupported Prometheus result type: %q", r.Data.ResultType) + } +} + +// parseLastValueFromRangeResponse extracts the last value from a Prometheus +// range query response — suitable for looking at the most recent datapoint +// in a look-back window. +func parseLastValueFromRangeResponse(body []byte) (float64, error) { + var r prometheusResponse + if err := json.Unmarshal(body, &r); err != nil { + return 0, fmt.Errorf("decoding Prometheus range response: %w", err) + } + if r.Status != "success" { + return 0, fmt.Errorf("Prometheus range query status: %q", r.Status) + } + if len(r.Data.Result) == 0 { + return 0, fmt.Errorf("Prometheus range result is empty") + } + + // Decode the matrix element to get the values array. + var elem struct { + Values []json.RawMessage `json:"values"` + } + if err := json.Unmarshal(r.Data.Result[0], &elem); err != nil { + return 0, fmt.Errorf("decoding range matrix element: %w", err) + } + if len(elem.Values) == 0 { + return 0, fmt.Errorf("range result has no values") + } + + return extractValueFromPair(elem.Values[len(elem.Values)-1]) +} + +// extractValueFromPair decodes a JSON [timestamp, "value"] pair and returns +// the float64 value. The timestamp is discarded. +func extractValueFromPair(raw json.RawMessage) (float64, error) { + var pair [2]json.RawMessage + if err := json.Unmarshal(raw, &pair); err != nil { + return 0, fmt.Errorf("decoding value pair: %w", err) + } + var valStr string + if err := json.Unmarshal(pair[1], &valStr); err != nil { + return 0, fmt.Errorf("decoding value string: %w", err) + } + return strconv.ParseFloat(valStr, 64) +} + +// extractValueFromVectorElement decodes a Prometheus vector element +// {"metric":{},"value":[ts,"val"]} and returns the numeric value. +func extractValueFromVectorElement(raw json.RawMessage) (float64, error) { + var elem struct { + Value json.RawMessage `json:"value"` + } + if err := json.Unmarshal(raw, &elem); err != nil { + return 0, fmt.Errorf("decoding vector element: %w", err) + } + return extractValueFromPair(elem.Value) +} + +// formatTime formats a time.Time as a Unix timestamp string for the Prometheus API. +func formatTime(t time.Time) string { + return strconv.FormatFloat(float64(t.UnixNano())/1e9, 'f', 3, 64) +} + +// formatDuration formats a time.Duration as a seconds string for the Prometheus API. +func formatDuration(d time.Duration) string { + return strconv.FormatFloat(d.Seconds(), 'f', 3, 64) + "s" +} diff --git a/internal/metrics/prometheus_test.go b/internal/metrics/prometheus_test.go new file mode 100644 index 0000000..aaec1a2 --- /dev/null +++ b/internal/metrics/prometheus_test.go @@ -0,0 +1,236 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package metrics + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +func TestParseScalarFromQueryResponse(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + body string + wantVal float64 + wantErr bool + }{ + { + name: "valid scalar response", + body: `{"status":"success","data":{"resultType":"scalar","result":[1435781451.781,"42.5"]}}`, + wantVal: 42.5, + wantErr: false, + }, + { + name: "scalar with fewer than 2 elements", + body: `{"status":"success","data":{"resultType":"scalar","result":[1435781451.781]}}`, + wantErr: true, + }, + { + name: "scalar with more than 2 elements", + body: `{"status":"success","data":{"resultType":"scalar","result":[1435781451.781,"42.5","extra"]}}`, + wantErr: true, + }, + { + name: "scalar with non-string value", + body: `{"status":"success","data":{"resultType":"scalar","result":[1435781451.781,42.5]}}`, + wantErr: true, + }, + { + name: "scalar with non-float string value", + body: `{"status":"success","data":{"resultType":"scalar","result":[1435781451.781,"invalid"]}}`, + wantErr: true, + }, + { + name: "valid vector response", + body: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"http_requests_total"},"value":[1435781451.781,"100.5"]}]}}`, + wantVal: 100.5, + wantErr: false, + }, + { + name: "empty vector response", + body: `{"status":"success","data":{"resultType":"vector","result":[]}}`, + wantErr: true, + }, + { + name: "multiple series in vector response", + body: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"http_requests_total"},"value":[1435781451.781,"100.5"]},{"metric":{"__name__":"http_requests_total"},"value":[1435781451.781,"200.5"]}]}}`, + wantErr: true, + }, + { + name: "unsupported resultType", + body: `{"status":"success","data":{"resultType":"matrix","result":[]}}`, + wantErr: true, + }, + { + name: "status error", + body: `{"status":"error","error":"bad query"}`, + wantErr: true, + }, + { + name: "malformed JSON", + body: `{"status":`, + wantErr: true, + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got, err := parseScalarFromQueryResponse([]byte(tc.body)) + if tc.wantErr { + if err == nil { + t.Errorf("expected error for %s, got nil", tc.name) + } + return + } + if err != nil { + t.Fatalf("unexpected error for %s: %v", tc.name, err) + } + if got != tc.wantVal { + t.Errorf("got %v, want %v", got, tc.wantVal) + } + }) + } +} + +func TestParseLastValueFromRangeResponse(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + body string + wantVal float64 + wantErr bool + }{ + { + name: "valid range response", + body: `{"status":"success","data":{"resultType":"matrix","result":[{"metric":{},"values":[[1435781430.781,"10"],[1435781451.781,"20.5"]]}]}}`, + wantVal: 20.5, + wantErr: false, + }, + { + name: "empty result", + body: `{"status":"success","data":{"resultType":"matrix","result":[]}}`, + wantErr: true, + }, + { + name: "empty values in series", + body: `{"status":"success","data":{"resultType":"matrix","result":[{"metric":{},"values":[]}]}}`, + wantErr: true, + }, + { + name: "status error", + body: `{"status":"error","error":"bad query"}`, + wantErr: true, + }, + { + name: "malformed json", + body: `invalid json`, + wantErr: true, + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got, err := parseLastValueFromRangeResponse([]byte(tc.body)) + if tc.wantErr { + if err == nil { + t.Errorf("expected error for %s, got nil", tc.name) + } + return + } + if err != nil { + t.Fatalf("unexpected error for %s: %v", tc.name, err) + } + if got != tc.wantVal { + t.Errorf("got %v, want %v", got, tc.wantVal) + } + }) + } +} + +func TestClient_QueryScalar(t *testing.T) { + t.Parallel() + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/v1/query" { + http.NotFound(w, r) + return + } + q := r.URL.Query().Get("query") + if q == "scalar_metric" { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"status":"success","data":{"resultType":"scalar","result":[1435781451.781,"12.34"]}}`)) + return + } + if q == "error_metric" { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`bad query`)) + return + } + http.Error(w, "unknown query", http.StatusInternalServerError) + })) + defer ts.Close() + + c := NewClient(ts.URL, WithTimeout(2*time.Second)) + val, err := c.QueryScalar(context.Background(), "scalar_metric") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if val != 12.34 { + t.Errorf("expected 12.34, got %v", val) + } + + _, err = c.QueryScalar(context.Background(), "error_metric") + if err == nil { + t.Error("expected error for error_metric, got nil") + } +} + +func TestClient_QueryRange(t *testing.T) { + t.Parallel() + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/v1/query_range" { + http.NotFound(w, r) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"status":"success","data":{"resultType":"matrix","result":[{"metric":{},"values":[[1000,"5.5"],[2000,"8.5"]]}]}}`)) + })) + defer ts.Close() + + c := NewClient(ts.URL) + now := time.Now() + val, err := c.QueryRange(context.Background(), "range_query", now.Add(-10*time.Minute), now, time.Minute) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if val != 8.5 { + t.Errorf("expected 8.5, got %v", val) + } +} diff --git a/internal/quota/enforcer.go b/internal/quota/enforcer.go index c3e3601..d3a438b 100644 --- a/internal/quota/enforcer.go +++ b/internal/quota/enforcer.go @@ -171,11 +171,9 @@ func (e *Enforcer) ComputeUsage(ads []agentraxv1alpha1.AgentDeploymentSpec) agen // by concurrent admission requests. // // admissionKey must be unique per AgentDeployment — use "namespace/adName". -// It identifies the reservation slot for this specific AD so that a retried -// webhook call replaces (not duplicates) any prior reservation, and so that -// sibling ADs in the same TenantQuota each hold independent slots. -// oldSpec may be nil for CREATE requests; for UPDATE requests it is the -// previous spec so we can compute the delta rather than a full new addition. +// CanAdmit is exported for use in dry-run validation where no reservation +// should be created. For persisted requests, use AdmitAndReserve which +// eliminates the TOCTOU window between checking and reserving. // // Returns (true, "") if admission is allowed, or (false, reason) if not. func (e *Enforcer) CanAdmit( @@ -185,82 +183,178 @@ func (e *Enforcer) CanAdmit( requested agentraxv1alpha1.AgentDeploymentSpec, oldSpec *agentraxv1alpha1.AgentDeploymentSpec, ) (bool, string) { - // Compute the delta this request adds on top of committed usage. - // For CREATE: delta = full requested resources. - // For UPDATE: delta = requested - old (can be negative if scaling down). - var deltaAgents, deltaGPUs, deltaReplicas int32 + delta := e.computeDelta(requested, oldSpec) + inFlight := e.sumInflight(admissionKey) + return evalQuotaRules(quota, committedUsage, inFlight, delta, + requested.Replicas.Max, sameTenantUpdate(requested, oldSpec), oldMax(oldSpec)) +} + +// sameTenantUpdate returns true for UPDATE operations where the tenant has not changed. +// Returns false for CREATE (oldSpec == nil) or cross-tenant moves. +func sameTenantUpdate(requested agentraxv1alpha1.AgentDeploymentSpec, oldSpec *agentraxv1alpha1.AgentDeploymentSpec) bool { if oldSpec == nil { - // CREATE - deltaAgents = 1 - deltaGPUs = e.gpusForAD(requested) - deltaReplicas = requested.Replicas.Max - } else { - // UPDATE — agent count doesn't change, only resources may shift. - deltaAgents = 0 - deltaGPUs = e.gpusForAD(requested) - e.gpusForAD(*oldSpec) - deltaReplicas = requested.Replicas.Max - oldSpec.Replicas.Max + return false } + return requested.TenantRef == oldSpec.TenantRef +} - // Add in-flight reservations (excluding this AD's own slot, which will - // be replaced when we Reserve below). - inFlight := e.sumInflight(admissionKey) +// computeDelta returns the resource delta for this admission request. +// For CREATE (oldSpec == nil): the full resources of one new agent. +// For UPDATE with same TenantRef: only the incremental change relative to the previous spec. +// For UPDATE with different TenantRef (tenant move): the FULL requested resources, +// because the old allocation was never charged against the new tenant's quota. +func (e *Enforcer) computeDelta( + requested agentraxv1alpha1.AgentDeploymentSpec, + oldSpec *agentraxv1alpha1.AgentDeploymentSpec, +) reservationEntry { + if oldSpec == nil { + // CREATE: charge full amount + return reservationEntry{ + agents: 1, + gpus: e.gpusForAD(requested), + replicas: requested.Replicas.Max, + } + } + // UPDATE: check if tenant changed + if requested.TenantRef != oldSpec.TenantRef { + // Cross-tenant move: charge full amount to the new tenant + return reservationEntry{ + agents: 1, + gpus: e.gpusForAD(requested), + replicas: requested.Replicas.Max, + } + } + // Same tenant: charge only the delta + return reservationEntry{ + agents: 0, + gpus: e.gpusForAD(requested) - e.gpusForAD(*oldSpec), + replicas: requested.Replicas.Max - oldSpec.Replicas.Max, + } +} + +// oldMax returns oldSpec.Replicas.Max or 0 for CREATE requests. +func oldMax(oldSpec *agentraxv1alpha1.AgentDeploymentSpec) int32 { + if oldSpec == nil { + return 0 + } + return oldSpec.Replicas.Max +} - projectedAgents := committedUsage.UsedAgents + inFlight.agents + deltaAgents - projectedGPUs := committedUsage.UsedGPUs + inFlight.gpus + deltaGPUs - projectedReplicas := committedUsage.UsedTotalReplicas + inFlight.replicas + deltaReplicas +// evalQuotaRules checks all four quota limits given committed usage, in-flight +// totals, and the delta contributed by this request. It is a pure, lock-free +// helper shared by canAdmit and AdmitAndReserve; callers are responsible for +// holding any locks before reading in-flight state. +// Returns (false, reason) for the first violated limit, or (true, "") otherwise. +func evalQuotaRules( + quota agentraxv1alpha1.TenantQuotaSpec, + committedUsage agentraxv1alpha1.TenantQuotaStatus, + inFlight reservationEntry, + delta reservationEntry, + requestedMaxReplicas int32, + isUpdate bool, + prevMaxReplicas int32, // 0 for creates; used by per-agent ceiling check +) (bool, string) { + // Reject the math.MaxInt32 sentinel used by gpusForAD to signal overflow. + // If delta.gpus is MaxInt32, the GPU calculation overflowed and we must + // fail closed to prevent undercount-based admission. + if delta.gpus == math.MaxInt32 { + return false, "GPU resource calculation overflowed; request denied" + } + + // Perform projection calculations in int64 to prevent overflow, then + // compare against int64-converted quota limits. + projAgents := int64(committedUsage.UsedAgents) + int64(inFlight.agents) + int64(delta.agents) + projGPUs := int64(committedUsage.UsedGPUs) + int64(inFlight.gpus) + int64(delta.gpus) + projReplicas := int64(committedUsage.UsedTotalReplicas) + int64(inFlight.replicas) + int64(delta.replicas) // For UPDATE requests, only reject when the delta increases a dimension that // is already at or over quota. If quota was lowered below current usage, the // existing ADs are already OverQuota (indicated by the TQ condition) — we // must not block updates that don't make things worse, otherwise finalizer // removal and spec corrections are deadlocked. - isUpdate := oldSpec != nil - - if projectedAgents > quota.MaxAgents && (!isUpdate || deltaAgents > 0) { + if projAgents > int64(quota.MaxAgents) && (!isUpdate || delta.agents > 0) { return false, fmt.Sprintf( "would exceed maxAgents (%d): current=%d in-flight=%d delta=%d", - quota.MaxAgents, committedUsage.UsedAgents, inFlight.agents, deltaAgents, + quota.MaxAgents, committedUsage.UsedAgents, inFlight.agents, delta.agents, ) } - if quota.MaxGPUs > 0 && projectedGPUs > quota.MaxGPUs && (!isUpdate || deltaGPUs > 0) { + if projGPUs > int64(quota.MaxGPUs) && (!isUpdate || delta.gpus > 0) { return false, fmt.Sprintf( "would exceed maxGPUs (%d): current=%d in-flight=%d delta=%d", - quota.MaxGPUs, committedUsage.UsedGPUs, inFlight.gpus, deltaGPUs, + quota.MaxGPUs, committedUsage.UsedGPUs, inFlight.gpus, delta.gpus, ) } - if projectedReplicas > quota.MaxTotalReplicas && (!isUpdate || deltaReplicas > 0) { + if projReplicas > int64(quota.MaxTotalReplicas) && (!isUpdate || delta.replicas > 0) { return false, fmt.Sprintf( "would exceed maxTotalReplicas (%d): current=%d in-flight=%d delta=%d", - quota.MaxTotalReplicas, committedUsage.UsedTotalReplicas, inFlight.replicas, deltaReplicas, + quota.MaxTotalReplicas, committedUsage.UsedTotalReplicas, inFlight.replicas, delta.replicas, ) } // MaxReplicasPerAgent is only enforced when the request would increase the // per-agent replica ceiling. If the quota ceiling was lowered below an // existing AD's replicas.max, updates that don't raise replicas.max further // must still be allowed — blocking them would deadlock spec corrections. - if requested.Replicas.Max > quota.MaxReplicasPerAgent && (!isUpdate || requested.Replicas.Max > oldSpec.Replicas.Max) { + perAgentIncreases := !isUpdate || requestedMaxReplicas > prevMaxReplicas + if requestedMaxReplicas > quota.MaxReplicasPerAgent && perAgentIncreases { return false, fmt.Sprintf( "spec.replicas.max (%d) exceeds maxReplicasPerAgent (%d)", - requested.Replicas.Max, quota.MaxReplicasPerAgent, + requestedMaxReplicas, quota.MaxReplicasPerAgent, ) } - return true, "" } -// sumInflight returns the total in-flight resource counts excluding the entry -// for excludeKey (so this AD's existing slot is not double-counted when the -// same AD retries admission). -func (e *Enforcer) sumInflight(excludeKey string) reservationEntry { +// AdmitAndReserve is an atomic version of canAdmit followed by reserve. +// It holds the in-flight mutex continuously from the quota check through the +// reservation write, eliminating the TOCTOU window that exists when the two +// operations are called separately: a concurrent near-limit create cannot slip +// through because no other goroutine can observe stale in-flight counts between +// the check and the write. +// +// Returns (true, "") and writes the reservation when admission is allowed. +// Returns (false, reason) without touching the map when admission is denied. +// The reservation expires after ttl and is cleaned up by the sweep goroutine. +func (e *Enforcer) AdmitAndReserve( + admissionKey string, + quota agentraxv1alpha1.TenantQuotaSpec, + committedUsage agentraxv1alpha1.TenantQuotaStatus, + requested agentraxv1alpha1.AgentDeploymentSpec, + oldSpec *agentraxv1alpha1.AgentDeploymentSpec, + ttl time.Duration, +) (bool, string) { + delta := e.computeDelta(requested, oldSpec) + + // Hold the mutex for the entire check-then-reserve operation so no + // concurrent admission can observe an inconsistent in-flight snapshot. e.mu.Lock() defer e.mu.Unlock() + now := e.nowFn() + inFlight := e.sumInflightLocked(admissionKey, now) + + ok, reason := evalQuotaRules(quota, committedUsage, inFlight, delta, + requested.Replicas.Max, sameTenantUpdate(requested, oldSpec), oldMax(oldSpec)) + if !ok { + return false, reason + } + + // Admission passed — write the reservation under the same lock so the + // slot is visible to any concurrent AdmitAndReserve caller immediately. + e.reservations[admissionKey] = &reservationEntry{ + agents: delta.agents, + gpus: delta.gpus, + replicas: delta.replicas, + expiry: now.Add(ttl), + } + return true, "" +} + +// sumInflightLocked returns the total in-flight resource counts excluding the +// entry for excludeKey. The caller must hold e.mu before invoking this method. +func (e *Enforcer) sumInflightLocked(excludeKey string, now time.Time) reservationEntry { var total reservationEntry for k, v := range e.reservations { - if k == excludeKey { - continue - } - if now.After(v.expiry) { + if k == excludeKey || now.After(v.expiry) { continue } total.agents += v.agents @@ -270,28 +364,27 @@ func (e *Enforcer) sumInflight(excludeKey string) reservationEntry { return total } -// Reserve creates or replaces an in-flight reservation for admissionKey (a -// per-request unique string, e.g. UID of the AdmissionRequest) lasting ttl. -// The reservation is automatically swept when it expires. Call Release when -// the webhook handler returns (succeeded or failed), or let it expire on its -// own if the process crashes. -func (e *Enforcer) Reserve(admissionKey string, spec agentraxv1alpha1.AgentDeploymentSpec, oldSpec *agentraxv1alpha1.AgentDeploymentSpec, ttl time.Duration) { - var deltaAgents, deltaGPUs, deltaReplicas int32 - if oldSpec == nil { - deltaAgents = 1 - deltaGPUs = e.gpusForAD(spec) - deltaReplicas = spec.Replicas.Max - } else { - deltaGPUs = e.gpusForAD(spec) - e.gpusForAD(*oldSpec) - deltaReplicas = spec.Replicas.Max - oldSpec.Replicas.Max - } +// sumInflight returns the total in-flight resource counts excluding the entry +// for excludeKey (so this AD's existing slot is not double-counted when the +// same AD retries admission). +func (e *Enforcer) sumInflight(excludeKey string) reservationEntry { + e.mu.Lock() + defer e.mu.Unlock() + return e.sumInflightLocked(excludeKey, e.nowFn()) +} +// reserve creates or replaces an in-flight reservation for admissionKey lasting +// ttl. It is intentionally unexported: production code must use AdmitAndReserve +// to avoid the TOCTOU race between checking and reserving. Tests that need to +// pre-seed the in-flight map in isolation may call this directly. +func (e *Enforcer) reserve(admissionKey string, spec agentraxv1alpha1.AgentDeploymentSpec, oldSpec *agentraxv1alpha1.AgentDeploymentSpec, ttl time.Duration) { + delta := e.computeDelta(spec, oldSpec) e.mu.Lock() defer e.mu.Unlock() e.reservations[admissionKey] = &reservationEntry{ - agents: deltaAgents, - gpus: deltaGPUs, - replicas: deltaReplicas, + agents: delta.agents, + gpus: delta.gpus, + replicas: delta.replicas, expiry: e.nowFn().Add(ttl), } } @@ -314,7 +407,7 @@ func (e *Enforcer) IsOverQuota( if usage.UsedAgents > quota.MaxAgents { return true, fmt.Sprintf("usedAgents (%d) exceeds maxAgents (%d)", usage.UsedAgents, quota.MaxAgents) } - if quota.MaxGPUs > 0 && usage.UsedGPUs > quota.MaxGPUs { + if usage.UsedGPUs > quota.MaxGPUs { return true, fmt.Sprintf("usedGPUs (%d) exceeds maxGPUs (%d)", usage.UsedGPUs, quota.MaxGPUs) } if usage.UsedTotalReplicas > quota.MaxTotalReplicas { diff --git a/internal/quota/enforcer_test.go b/internal/quota/enforcer_test.go index 018c0b9..fa72fb3 100644 --- a/internal/quota/enforcer_test.go +++ b/internal/quota/enforcer_test.go @@ -14,10 +14,16 @@ See the License for the specific language governing permissions and limitations under the License. */ -package quota_test +// White-box test: package quota (not quota_test) so unexported methods +// like reserve are accessible for isolated unit testing. +// AdmitAndReserve remains the production-facing atomic API. +package quota import ( + "fmt" "strings" + "sync" + "sync/atomic" "testing" "time" @@ -25,11 +31,11 @@ import ( "k8s.io/apimachinery/pkg/api/resource" agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" - "github.com/gitcommitankit/agentrax/internal/quota" ) // ── helpers ────────────────────────────────────────────────────────────────── +// makeSpec creates an AgentDeploymentSpec test fixture with given maxReplicas and GPU limit. func makeSpec(maxReplicas int32, gpuLimit string) agentraxv1alpha1.AgentDeploymentSpec { spec := agentraxv1alpha1.AgentDeploymentSpec{ Image: "test-image:v1", @@ -51,6 +57,7 @@ func makeSpec(maxReplicas int32, gpuLimit string) agentraxv1alpha1.AgentDeployme return spec } +// makeQuota creates a TenantQuotaSpec test fixture with given quota limits. func makeQuota(maxAgents, maxGPUs, maxTotalReplicas, maxReplicasPerAgent int32) agentraxv1alpha1.TenantQuotaSpec { return agentraxv1alpha1.TenantQuotaSpec{ MaxAgents: maxAgents, @@ -60,6 +67,7 @@ func makeQuota(maxAgents, maxGPUs, maxTotalReplicas, maxReplicasPerAgent int32) } } +// makeUsage creates a TenantQuotaStatus test fixture with given observed usage. func makeUsage(agents, gpus, replicas int32) agentraxv1alpha1.TenantQuotaStatus { return agentraxv1alpha1.TenantQuotaStatus{ UsedAgents: agents, @@ -70,15 +78,16 @@ func makeUsage(agents, gpus, replicas int32) agentraxv1alpha1.TenantQuotaStatus // newTestEnforcer creates an Enforcer for tests and registers Stop() as a cleanup // function so the background sweep goroutine is terminated when the test ends. -func newTestEnforcer(t *testing.T) *quota.Enforcer { +func newTestEnforcer(t *testing.T) *Enforcer { t.Helper() - e := quota.NewEnforcer(quota.DefaultGPUResourceName) + e := NewEnforcer(DefaultGPUResourceName) t.Cleanup(e.Stop) return e } // ── CanAdmit tests ──────────────────────────────────────────────────────────── +// TestCanAdmit_Create verifies admission decisions for new AgentDeployment creates against various quota scenarios. func TestCanAdmit_Create(t *testing.T) { t.Parallel() tests := []struct { @@ -143,6 +152,21 @@ func TestCanAdmit_Create(t *testing.T) { spec: makeSpec(2, ""), // no GPU requested → no increase wantAdmit: true, }, + { + name: "zero GPU quota rejects positive GPU request", + quota: makeQuota(6, 0, 12, 6), + usage: makeUsage(1, 0, 2), + spec: makeSpec(2, "1"), // 1 GPU × 2 replicas = 2 > 0 + wantAdmit: false, + wantContain: "maxGPUs", + }, + { + name: "zero GPU quota admits non-GPU request", + quota: makeQuota(6, 0, 12, 6), + usage: makeUsage(1, 0, 2), + spec: makeSpec(2, ""), // 0 GPUs requested + wantAdmit: true, + }, } for _, tc := range tests { @@ -165,6 +189,7 @@ func TestCanAdmit_Create(t *testing.T) { } } +// TestCanAdmit_Update verifies admission decisions when updating existing AgentDeployment specs. func TestCanAdmit_Update(t *testing.T) { t.Parallel() e := newTestEnforcer(t) @@ -192,6 +217,7 @@ func TestCanAdmit_Update(t *testing.T) { } } +// TestCanAdmit_Update_MaxReplicasPerAgent_Downgrade verifies that lowering maxReplicasPerAgent does not deadlock updates that do not increase replicas. func TestCanAdmit_Update_MaxReplicasPerAgent_Downgrade(t *testing.T) { // When maxReplicasPerAgent is lowered below an existing AD's replicas.max, // updates that do NOT further increase replicas.max must still be allowed. @@ -228,8 +254,193 @@ func TestCanAdmit_Update_MaxReplicasPerAgent_Downgrade(t *testing.T) { } } +// TestCanAdmit_Update_GPUCeiling verifies admission behavior when updating GPU requests or when GPU quota is reduced. +func TestCanAdmit_Update_GPUCeiling(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + quota agentraxv1alpha1.TenantQuotaSpec + usage agentraxv1alpha1.TenantQuotaStatus + oldSpec agentraxv1alpha1.AgentDeploymentSpec + newSpec agentraxv1alpha1.AgentDeploymentSpec + wantAdmit bool + wantContain string + }{ + { + name: "zero quota rejects increasing GPU allocation", + quota: makeQuota(6, 0, 10, 6), + usage: makeUsage(1, 0, 2), + oldSpec: makeSpec(2, ""), + newSpec: makeSpec(2, "1"), // requests 2 GPUs when quota is 0 + wantAdmit: false, + wantContain: "maxGPUs", + }, + { + name: "non-increasing GPU update when over-quota is admitted", + quota: makeQuota(6, 2, 10, 6), // quota lowered to 2 GPUs + usage: makeUsage(1, 4, 2), // current usage is 4 GPUs (already over) + oldSpec: makeSpec(2, "2"), // 2 GPU × 2 = 4 GPUs + newSpec: makeSpec(2, "2"), // same GPUs + wantAdmit: true, + wantContain: "", + }, + { + name: "increasing GPU update when over-quota is rejected", + quota: makeQuota(6, 2, 10, 6), // quota lowered to 2 GPUs + usage: makeUsage(1, 4, 2), // current usage is 4 GPUs (already over) + oldSpec: makeSpec(2, "2"), // 2 GPU × 2 = 4 GPUs + newSpec: makeSpec(3, "2"), // 2 GPU × 3 = 6 GPUs (delta +2) + wantAdmit: false, + wantContain: "maxGPUs", + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + e := newTestEnforcer(t) + got, reason := e.CanAdmit("ns/ad-A", tc.quota, tc.usage, tc.newSpec, &tc.oldSpec) + if got != tc.wantAdmit { + t.Errorf("CanAdmit() = %v, want %v; reason: %q", got, tc.wantAdmit, reason) + } + if !tc.wantAdmit && tc.wantContain != "" { + if !strings.Contains(reason, tc.wantContain) { + t.Errorf("expected reason to contain %q, got %q", tc.wantContain, reason) + } + } + }) + } +} + +// TestCanAdmit_Update_CrossTenantMove verifies that when an AgentDeployment +// moves from one tenant to another (TenantRef changes), the full requested +// resources are charged to the target tenant's quota, not a delta. +func TestCanAdmit_Update_CrossTenantMove(t *testing.T) { + t.Parallel() + const ( + oldTenant = "old-tenant" + targetTenant = "target-tenant" + ) + + tests := []struct { + name string + quota agentraxv1alpha1.TenantQuotaSpec + usage agentraxv1alpha1.TenantQuotaStatus + oldSpec agentraxv1alpha1.AgentDeploymentSpec + newSpec agentraxv1alpha1.AgentDeploymentSpec + wantAdmit bool + wantContain string + }{ + { + name: "cross-tenant move into full quota is rejected", + quota: makeQuota(2, 4, 6, 4), // maxAgents=2, maxGPUs=4, maxTotalReplicas=6 + usage: makeUsage(2, 4, 6), // already at full capacity + oldSpec: func() agentraxv1alpha1.AgentDeploymentSpec { + s := makeSpec(3, "1") // 3 replicas, 1 GPU per replica = 3 GPUs total + s.TenantRef = oldTenant + return s + }(), + newSpec: func() agentraxv1alpha1.AgentDeploymentSpec { + s := makeSpec(3, "1") // same resources: 3 replicas, 3 GPUs + s.TenantRef = targetTenant + return s + }(), + wantAdmit: false, + wantContain: "maxAgents", + }, + { + name: "cross-tenant move into quota with capacity is admitted", + quota: makeQuota(3, 8, 10, 4), // enough room: maxAgents=3, maxGPUs=8, maxTotalReplicas=10 + usage: makeUsage(1, 2, 4), // current: 1 agent, 2 GPUs, 4 replicas + oldSpec: func() agentraxv1alpha1.AgentDeploymentSpec { + s := makeSpec(3, "1") // 3 replicas, 3 GPUs + s.TenantRef = oldTenant + return s + }(), + newSpec: func() agentraxv1alpha1.AgentDeploymentSpec { + s := makeSpec(3, "1") // same resources + s.TenantRef = targetTenant + return s + }(), + wantAdmit: true, + wantContain: "", + }, + { + name: "cross-tenant move with resource increase when quota allows", + quota: makeQuota(3, 5, 8, 5), + usage: makeUsage(1, 1, 2), // 1 agent, 1 GPU, 2 replicas + oldSpec: func() agentraxv1alpha1.AgentDeploymentSpec { + s := makeSpec(2, "1") // 2 replicas, 2 GPUs + s.TenantRef = oldTenant + return s + }(), + newSpec: func() agentraxv1alpha1.AgentDeploymentSpec { + s := makeSpec(4, "1") // 4 replicas, 4 GPUs + s.TenantRef = targetTenant + return s + }(), + wantAdmit: true, + wantContain: "", + }, + { + name: "cross-tenant move into insufficient GPU quota is rejected", + quota: makeQuota(3, 4, 8, 5), // maxGPUs=4 (too low for +4) + usage: makeUsage(1, 1, 2), + oldSpec: func() agentraxv1alpha1.AgentDeploymentSpec { + s := makeSpec(2, "1") + s.TenantRef = oldTenant + return s + }(), + newSpec: func() agentraxv1alpha1.AgentDeploymentSpec { + s := makeSpec(4, "1") // 4 GPUs needed + s.TenantRef = targetTenant + return s + }(), + wantAdmit: false, + wantContain: "maxGPUs", + }, + { + name: "cross-tenant move of 8-replica workload into tenant with maxReplicasPerAgent=2 is rejected", + quota: makeQuota(5, 10, 20, 2), // maxReplicasPerAgent=2 + usage: makeUsage(1, 0, 5), + oldSpec: func() agentraxv1alpha1.AgentDeploymentSpec { + s := makeSpec(8, "") // 8 replicas + s.TenantRef = oldTenant + return s + }(), + newSpec: func() agentraxv1alpha1.AgentDeploymentSpec { + s := makeSpec(8, "") // 8 replicas + s.TenantRef = targetTenant + return s + }(), + wantAdmit: false, + wantContain: "maxReplicasPerAgent", + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + e := newTestEnforcer(t) + got, reason := e.CanAdmit("ns/ad-move", tc.quota, tc.usage, tc.newSpec, &tc.oldSpec) + if got != tc.wantAdmit { + t.Errorf("CanAdmit() = %v, want %v; reason: %q", got, tc.wantAdmit, reason) + } + if !tc.wantAdmit && tc.wantContain != "" { + if !strings.Contains(reason, tc.wantContain) { + t.Errorf("expected reason to contain %q, got %q", tc.wantContain, reason) + } + } + }) + } +} + // ── In-flight reservation tests ─────────────────────────────────────────────── +// TestReservation_BlocksConcurrentCreate verifies that an in-flight reservation blocks concurrent creation of the same remaining slot. func TestReservation_BlocksConcurrentCreate(t *testing.T) { t.Parallel() e := newTestEnforcer(t) @@ -243,7 +454,7 @@ func TestReservation_BlocksConcurrentCreate(t *testing.T) { if !ok1 { t.Fatal("first CanAdmit should have passed") } - e.Reserve("ns/ad-A", spec, nil, 5*time.Second) + e.reserve("ns/ad-A", spec, nil, 5*time.Second) // Second concurrent request for the same remaining slot should now be blocked // because ad-A's reservation already claimed it. @@ -260,6 +471,7 @@ func TestReservation_BlocksConcurrentCreate(t *testing.T) { } } +// TestReservation_DoesNotDoubleCount verifies that re-admission for the same AD key excludes its own prior reservation. func TestReservation_DoesNotDoubleCount(t *testing.T) { // A re-admission for the same AD key should exclude its own prior reservation // so it isn't double-counted. @@ -269,19 +481,161 @@ func TestReservation_DoesNotDoubleCount(t *testing.T) { usage := makeUsage(1, 0, 2) spec := makeSpec(2, "") - e.Reserve("ns/ad-X", spec, nil, 5*time.Second) + e.reserve("ns/ad-X", spec, nil, 5*time.Second) - // Calling CanAdmit with the same admissionKey should exclude its own + // Calling canAdmit with the same admissionKey should exclude its own // reservation from the in-flight sum (no double-count). ok, _ := e.CanAdmit("ns/ad-X", q, usage, spec, nil) // usage.agents=1, in-flight from ad-X is excluded, delta=1 → projected=2 ≤ 3 → ok if !ok { - t.Error("CanAdmit for the same AD key should not be blocked by its own reservation") + t.Error("canAdmit for the same AD key should not be blocked by its own reservation") + } +} + +// TestRelease_Concurrent verifies that concurrent Release calls on distinct keys +// are race-free and that all reservations are removed. Table-driven so we cover +// different concurrency fan-outs. +func TestRelease_Concurrent(t *testing.T) { + tests := []struct { + name string + numKeys int + }{ + {"2 concurrent releases", 2}, + {"5 concurrent releases", 5}, + {"10 concurrent releases", 10}, + } + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + e := newTestEnforcer(t) + // Quota large enough to hold all initial reservations at once. + q := makeQuota(int32(tc.numKeys*2), 0, int32(tc.numKeys*10), int32(tc.numKeys+1)) + usage := makeUsage(0, 0, 0) + spec := makeSpec(1, "") + + // Create one reservation per key sequentially (no concurrency yet). + keys := make([]string, tc.numKeys) + for i := range keys { + keys[i] = fmt.Sprintf("ns/ad-%d", i) + ok, _ := e.AdmitAndReserve(keys[i], q, usage, spec, nil, 30*time.Second) + if !ok { + t.Fatalf("initial AdmitAndReserve(%q) failed unexpectedly", keys[i]) + } + } + + // Release all reservations concurrently from a start barrier so + // goroutines are likely to overlap rather than run sequentially. + start := make(chan struct{}) + var wg sync.WaitGroup + for _, k := range keys { + k := k + wg.Add(1) + go func() { + defer wg.Done() + <-start // wait until all goroutines are ready + e.Release(k) + }() + } + close(start) // release all goroutines simultaneously + wg.Wait() + + // After all releases, each key's slot should be free: a fresh + // canAdmit (zero usage, zero in-flight) must succeed for every key. + for _, k := range keys { + ok, reason := e.CanAdmit(k, q, usage, spec, nil) + if !ok { + t.Errorf("after Release, CanAdmit(%q) = false; reason: %q", k, reason) + } + } + + // Every reservation must be gone, not merely within headroom. + e.mu.Lock() + remaining := len(e.reservations) + e.mu.Unlock() + if remaining != 0 { + t.Errorf("after concurrent Release, %d reservations remain; want 0", remaining) + } + }) + } +} + +// TestAdmitAndReserve_AtomicRaceProtection verifies that concurrent calls to +// AdmitAndReserve for the same final quota slot allow exactly one through. +// This is the property that the separate CanAdmit+Reserve two-call pattern +// could not guarantee (TOCTOU race). +func TestAdmitAndReserve_AtomicRaceProtection(t *testing.T) { + t.Parallel() + e := newTestEnforcer(t) + // Exactly 1 agent slot remaining. + q := makeQuota(1, 0, 5, 5) + usage := makeUsage(0, 0, 0) + spec := makeSpec(1, "") + + var ( + successCount int64 + failCount int64 + wg sync.WaitGroup + ) + // Start barrier: ensure both goroutines are scheduled before either calls + // AdmitAndReserve, maximising the chance of a real concurrent execution. + start := make(chan struct{}) + for _, key := range []string{"ns/ad-A", "ns/ad-B"} { + key := key + wg.Add(1) + go func() { + defer wg.Done() + <-start // wait until both goroutines are running + ok, _ := e.AdmitAndReserve(key, q, usage, spec, nil, 5*time.Second) + if ok { + atomic.AddInt64(&successCount, 1) + } else { + atomic.AddInt64(&failCount, 1) + } + }() + } + close(start) // release both goroutines simultaneously + wg.Wait() + + if successCount != 1 { + t.Errorf("expected exactly 1 AdmitAndReserve to succeed; got successCount=%d failCount=%d", + successCount, failCount) + } +} + +// TestAdmitAndReserve_DenialLeavesNoReservation is a white-box test that +// verifies a denied admission does not pollute the in-flight map. +func TestAdmitAndReserve_DenialLeavesNoReservation(t *testing.T) { + t.Parallel() + e := newTestEnforcer(t) + // Quota already exhausted: no remaining agents. + q := makeQuota(1, 10, 10, 10) + usage := makeUsage(1, 0, 0) + spec := makeSpec(1, "") + + ok, reason := e.AdmitAndReserve("ns/ad-denied", q, usage, spec, nil, 5*time.Second) + if ok { + t.Errorf("AdmitAndReserve should have denied admission with exhausted quota, but it succeeded") + } + if reason == "" { + t.Errorf("AdmitAndReserve denial should include a reason, got empty string") + } + + // White-box check: under the mutex, verify no reservation was created. + e.mu.Lock() + defer e.mu.Unlock() + if len(e.reservations) != 0 { + t.Errorf("expected 0 reservations after denial; got %d entries: %v (denial reason: %s)", + len(e.reservations), e.reservations, reason) + } + if _, exists := e.reservations["ns/ad-denied"]; exists { + t.Errorf("denied key ns/ad-denied should not exist in reservations (denial reason: %s)", reason) } } // ── ComputeUsage tests ──────────────────────────────────────────────────────── +// TestComputeUsage verifies aggregation of agents, GPUs, and replicas across multiple AgentDeployment specs. func TestComputeUsage(t *testing.T) { t.Parallel() e := newTestEnforcer(t) @@ -302,6 +656,7 @@ func TestComputeUsage(t *testing.T) { } } +// TestComputeUsage_Empty verifies that empty input returns all-zero usage. func TestComputeUsage_Empty(t *testing.T) { t.Parallel() e := newTestEnforcer(t) @@ -313,27 +668,29 @@ func TestComputeUsage_Empty(t *testing.T) { // ── IsOverQuota tests ───────────────────────────────────────────────────────── +// TestIsOverQuota verifies over-quota condition detection across agents, GPUs, and replicas dimensions. func TestIsOverQuota(t *testing.T) { t.Parallel() e := newTestEnforcer(t) - q := makeQuota(3, 4, 10, 3) - tests := []struct { name string + quota agentraxv1alpha1.TenantQuotaSpec usage agentraxv1alpha1.TenantQuotaStatus wantOQ bool wantMsg string }{ - {"within limits", makeUsage(2, 3, 8), false, ""}, - {"agents over", makeUsage(4, 3, 8), true, "maxAgents"}, - {"GPUs over", makeUsage(2, 5, 8), true, "maxGPUs"}, - {"replicas over", makeUsage(2, 3, 11), true, "maxTotalReplicas"}, + {"within limits", makeQuota(3, 4, 10, 3), makeUsage(2, 3, 8), false, ""}, + {"agents over", makeQuota(3, 4, 10, 3), makeUsage(4, 3, 8), true, "maxAgents"}, + {"GPUs over", makeQuota(3, 4, 10, 3), makeUsage(2, 5, 8), true, "maxGPUs"}, + {"replicas over", makeQuota(3, 4, 10, 3), makeUsage(2, 3, 11), true, "maxTotalReplicas"}, + {"zero GPU quota with used GPUs over", makeQuota(3, 0, 10, 3), makeUsage(2, 1, 8), true, "maxGPUs"}, + {"zero GPU quota with 0 used GPUs ok", makeQuota(3, 0, 10, 3), makeUsage(2, 0, 8), false, ""}, } for _, tc := range tests { tc := tc t.Run(tc.name, func(t *testing.T) { t.Parallel() - over, msg := e.IsOverQuota(q, tc.usage) + over, msg := e.IsOverQuota(tc.quota, tc.usage) if over != tc.wantOQ { t.Errorf("IsOverQuota() = %v, want %v; msg=%q", over, tc.wantOQ, msg) } @@ -348,6 +705,7 @@ func TestIsOverQuota(t *testing.T) { // ParseErrorRate lives in api/v1alpha1 to avoid an import cycle. // Its tests are in api/v1alpha1/webhook_test.go. +// TestParseErrorRate_ViaV1alpha1 verifies ParseErrorRate is accessible and correct from outside api/v1alpha1. func TestParseErrorRate_ViaV1alpha1(t *testing.T) { // Smoke-test that ParseErrorRate is accessible from outside api/v1alpha1. t.Parallel() @@ -362,6 +720,7 @@ func TestParseErrorRate_ViaV1alpha1(t *testing.T) { // ── helpers ─────────────────────────────────────────────────────────────────── +// absFloat returns the absolute value of a float64. func absFloat(f float64) float64 { if f < 0 { return -f diff --git a/internal/scaling/autoscaler.go b/internal/scaling/autoscaler.go new file mode 100644 index 0000000..91d4868 --- /dev/null +++ b/internal/scaling/autoscaler.go @@ -0,0 +1,212 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package scaling translates an AgentDeployment's replica policy into a managed +// HorizontalPodAutoscaler wired to the custom metrics exposed by Prometheus Adapter. +package scaling + +import ( + autoscalingv2 "k8s.io/api/autoscaling/v2" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" +) + +const ( + // MetricQueueDepth is the value of spec.replicas.metric that selects the + // request-queue depth custom metric. + MetricQueueDepth = "queueDepth" + + // MetricGPUUtilization is the value of spec.replicas.metric that selects + // the GPU utilization custom metric. + MetricGPUUtilization = "gpuUtilization" + + // customMetricQueueDepth is the name registered with the Prometheus Adapter + // for queue depth (External metric, namespace-scoped via label selector). + customMetricQueueDepth = "agentrax_queue_depth" + + // customMetricGPUUtilization is the name registered with the Prometheus Adapter + // for GPU utilization. + customMetricGPUUtilization = "agentrax_gpu_utilization" + + // scaleUpStabilizationSec is the HPA stabilization window for scale-up events + // (seconds). A 60-second window prevents thrashing on transient spikes. + scaleUpStabilizationSec int32 = 60 + + // scaleDownStabilizationSec is the HPA stabilization window for scale-down events + // (seconds). A 5-minute window prevents premature scale-down while load subsides. + scaleDownStabilizationSec int32 = 300 +) + +// BuildHPA constructs the desired HorizontalPodAutoscaler for the given AgentDeployment. +// quotaHeadroom is the maximum number of replicas the tenant quota currently allows; +// the HPA maxReplicas is capped at min(spec.replicas.max, quotaHeadroom). +// +// BuildHPA is a pure function — it does not make any Kubernetes API calls. +// The caller (reconciler) is responsible for CreateOrUpdate and owner reference. +func BuildHPA(ad *agentraxv1alpha1.AgentDeployment, quotaHeadroom int32) *autoscalingv2.HorizontalPodAutoscaler { + maxReplicas := ad.Spec.Replicas.Max + if quotaHeadroom < maxReplicas { + maxReplicas = quotaHeadroom + } + // Never let maxReplicas drop below minReplicas — a degenerate quota that + // leaves zero headroom is surfaced as QuotaLimited in the reconciler, but + // we still need a valid HPA spec. + minReplicas := ad.Spec.Replicas.Min + if maxReplicas < minReplicas { + maxReplicas = minReplicas + } + + metricName := customMetricNameFor(ad.Spec.Replicas.Metric) + + // Use an AverageValue target so the HPA scales to keep the per-replica + // metric value near the declared target (e.g., queue depth of 50 per pod). + targetValue := resource.NewMilliQuantity(int64(ad.Spec.Replicas.Target)*1000, resource.DecimalSI) + + scaleUpWindow := scaleUpStabilizationSec + scaleDownWindow := scaleDownStabilizationSec + + return &autoscalingv2.HorizontalPodAutoscaler{ + ObjectMeta: metav1.ObjectMeta{ + Name: ad.Name, + Namespace: ad.Namespace, + Labels: hpaLabels(ad), + }, + Spec: autoscalingv2.HorizontalPodAutoscalerSpec{ + ScaleTargetRef: autoscalingv2.CrossVersionObjectReference{ + APIVersion: "apps/v1", + Kind: "Deployment", + Name: ad.Name, + }, + MinReplicas: &minReplicas, + MaxReplicas: maxReplicas, + Metrics: []autoscalingv2.MetricSpec{ + { + Type: autoscalingv2.ExternalMetricSourceType, + External: &autoscalingv2.ExternalMetricSource{ + Metric: autoscalingv2.MetricIdentifier{ + Name: metricName, + // Scope the metric to this specific AgentDeployment. + // Prometheus sanitizes label names with dots to underscores, + // so use the sanitized forms that match what's in storage. + Selector: &metav1.LabelSelector{ + MatchLabels: map[string]string{ + "app_kubernetes_io_name": ad.Name, + "app_kubernetes_io_managed_by": "agentrax", + }, + }, + }, + Target: autoscalingv2.MetricTarget{ + Type: autoscalingv2.AverageValueMetricType, + AverageValue: targetValue, + }, + }, + }, + }, + Behavior: &autoscalingv2.HorizontalPodAutoscalerBehavior{ + ScaleUp: &autoscalingv2.HPAScalingRules{ + StabilizationWindowSeconds: &scaleUpWindow, + SelectPolicy: policyPtr(autoscalingv2.MaxChangePolicySelect), + Policies: []autoscalingv2.HPAScalingPolicy{ + { + Type: autoscalingv2.PodsScalingPolicy, + Value: 4, + PeriodSeconds: 60, + }, + }, + }, + ScaleDown: &autoscalingv2.HPAScalingRules{ + StabilizationWindowSeconds: &scaleDownWindow, + SelectPolicy: policyPtr(autoscalingv2.MinChangePolicySelect), + Policies: []autoscalingv2.HPAScalingPolicy{ + { + Type: autoscalingv2.PodsScalingPolicy, + Value: 1, + PeriodSeconds: 60, + }, + }, + }, + }, + }, + } +} + +// IsQuotaCapped returns true when spec.replicas.max exceeds the available +// quota headroom, indicating the HPA's maxReplicas was limited by quota. +// The caller should surface a QuotaLimited condition when this is true. +func IsQuotaCapped(ad *agentraxv1alpha1.AgentDeployment, quotaHeadroom int32) bool { + return quotaHeadroom < ad.Spec.Replicas.Max +} + +// QuotaHeadroom returns the raw quota ceiling for this AgentDeployment: the +// smaller of MaxReplicasPerAgent and the total replica budget remaining after +// accounting for other agents in the same tenant. The return value may be 0 +// when the total budget is exhausted. +// +// Callers must not confuse a 0 return with an error — it simply means no +// additional replicas are budgeted. BuildHPA applies a floor of minReplicas +// solely to keep the HPA spec valid; that does not change quota accounting. +// +// usedReplicasByOthers is the sum of spec.replicas.max for all OTHER +// AgentDeployments in the same tenant (excluding this one). +// The caller is responsible for computing this from the live list. +func QuotaHeadroom(tqSpec agentraxv1alpha1.TenantQuotaSpec, adSpec agentraxv1alpha1.AgentDeploymentSpec, usedReplicasByOthers int32) int32 { + // Per-agent ceiling from the TenantQuota. + perAgentCeiling := tqSpec.MaxReplicasPerAgent + + // Total replica budget remaining after accounting for other agents. + totalBudgetRemaining := tqSpec.MaxTotalReplicas - usedReplicasByOthers + if totalBudgetRemaining < 0 { + totalBudgetRemaining = 0 + } + + // The effective headroom is the smaller of the two constraints so neither + // the per-agent ceiling nor the total budget is violated. + headroom := perAgentCeiling + if totalBudgetRemaining < headroom { + headroom = totalBudgetRemaining + } + + return headroom +} + +// hpaLabels returns the label set applied to the managed HPA. +func hpaLabels(ad *agentraxv1alpha1.AgentDeployment) map[string]string { + return map[string]string{ + "app.kubernetes.io/name": ad.Name, + "app.kubernetes.io/managed-by": "agentrax", + "agentrax.io/tenant": ad.Spec.TenantRef, + } +} + +// customMetricNameFor maps a spec.replicas.metric value to the Prometheus +// Adapter custom metric name registered in the cluster. +func customMetricNameFor(metric string) string { + switch metric { + case MetricGPUUtilization: + return customMetricGPUUtilization + default: + // queueDepth and any unrecognized value default to queue depth. + return customMetricQueueDepth + } +} + +// policyPtr returns a pointer to an HPAScalingPolicySelect value. +func policyPtr(p autoscalingv2.ScalingPolicySelect) *autoscalingv2.ScalingPolicySelect { + v := p + return &v +} diff --git a/internal/scaling/autoscaler_test.go b/internal/scaling/autoscaler_test.go new file mode 100644 index 0000000..ca16f42 --- /dev/null +++ b/internal/scaling/autoscaler_test.go @@ -0,0 +1,357 @@ +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package scaling + +import ( + "testing" + + autoscalingv2 "k8s.io/api/autoscaling/v2" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + agentraxv1alpha1 "github.com/gitcommitankit/agentrax/api/v1alpha1" +) + +// makeAD is a test helper that builds a minimal AgentDeployment with the +// given replica policy. +func makeAD(name string, minR, maxR, target int32, metric string) *agentraxv1alpha1.AgentDeployment { + return &agentraxv1alpha1.AgentDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: name, + Namespace: "tenant-test", + }, + Spec: agentraxv1alpha1.AgentDeploymentSpec{ + Image: "nginx:latest", + TenantRef: "team-test", + Replicas: agentraxv1alpha1.ScalingPolicy{ + Min: minR, + Max: maxR, + Metric: metric, + Target: target, + }, + }, + } +} + +// makeTQSpec is a test helper that builds a TenantQuotaSpec with the two +// limits that vary across HPA/headroom tests. MaxAgents and MaxGPUs are fixed +// at 10 and 0 respectively — they do not affect HPA or headroom calculations. +func makeTQSpec(maxTotalReplicas, maxReplicasPerAgent int32) agentraxv1alpha1.TenantQuotaSpec { + return agentraxv1alpha1.TenantQuotaSpec{ + MaxAgents: 10, + MaxGPUs: 0, + MaxTotalReplicas: maxTotalReplicas, + MaxReplicasPerAgent: maxReplicasPerAgent, + } +} + +// ── BuildHPA tests ──────────────────────────────────────────────────────────── + +// TestBuildHPA_MinMaxReplicas verifies that HPA MinReplicas and MaxReplicas are set directly when headroom exceeds spec.replicas.max. +func TestBuildHPA_MinMaxReplicas(t *testing.T) { + t.Parallel() + + ad := makeAD("query-agent", 2, 8, 50, MetricQueueDepth) + hpa := BuildHPA(ad, 10) // headroom > max; no capping + + if hpa.Spec.MinReplicas == nil || *hpa.Spec.MinReplicas != 2 { + t.Errorf("expected minReplicas=2, got %v", hpa.Spec.MinReplicas) + } + if hpa.Spec.MaxReplicas != 8 { + t.Errorf("expected maxReplicas=8, got %d", hpa.Spec.MaxReplicas) + } +} + +// TestBuildHPA_QuotaCapApplied verifies that HPA MaxReplicas is capped at quota headroom when headroom is less than spec.replicas.max. +func TestBuildHPA_QuotaCapApplied(t *testing.T) { + t.Parallel() + + ad := makeAD("query-agent", 1, 10, 50, MetricQueueDepth) + hpa := BuildHPA(ad, 5) // headroom < spec.replicas.max → capped at 5 + + if hpa.Spec.MaxReplicas != 5 { + t.Errorf("expected maxReplicas capped at 5, got %d", hpa.Spec.MaxReplicas) + } +} + +// TestBuildHPA_QuotaCapNotApplied verifies that HPA MaxReplicas is not capped when headroom equals spec.replicas.max. +func TestBuildHPA_QuotaCapNotApplied(t *testing.T) { + t.Parallel() + + ad := makeAD("query-agent", 1, 6, 50, MetricQueueDepth) + hpa := BuildHPA(ad, 6) // headroom == max; no capping + + if hpa.Spec.MaxReplicas != 6 { + t.Errorf("expected maxReplicas=6 (not capped), got %d", hpa.Spec.MaxReplicas) + } +} + +// TestBuildHPA_QuotaHeadroomBelowMin verifies that MaxReplicas is clamped to MinReplicas when headroom is zero. +func TestBuildHPA_QuotaHeadroomBelowMin(t *testing.T) { + t.Parallel() + + // Edge: quota headroom is 0 but minReplicas is 1. + // BuildHPA must clamp maxReplicas to minReplicas (1) so the HPA spec stays valid. + ad := makeAD("query-agent", 1, 5, 50, MetricQueueDepth) + hpa := BuildHPA(ad, 0) + + if hpa.Spec.MinReplicas == nil || *hpa.Spec.MinReplicas != 1 { + t.Errorf("expected minReplicas=1, got %v", hpa.Spec.MinReplicas) + } + if hpa.Spec.MaxReplicas != 1 { + t.Errorf("expected maxReplicas=1 (clamped to minReplicas when headroom=0), got %d", + hpa.Spec.MaxReplicas) + } +} + +// TestBuildHPA_ScaleTargetRef verifies that ScaleTargetRef points to the agent Deployment apps/v1. +func TestBuildHPA_ScaleTargetRef(t *testing.T) { + t.Parallel() + + ad := makeAD("my-agent", 1, 4, 100, MetricQueueDepth) + hpa := BuildHPA(ad, 10) + + ref := hpa.Spec.ScaleTargetRef + if ref.Kind != "Deployment" { + t.Errorf("expected ScaleTargetRef.Kind=Deployment, got %s", ref.Kind) + } + if ref.Name != "my-agent" { + t.Errorf("expected ScaleTargetRef.Name=my-agent, got %s", ref.Name) + } + if ref.APIVersion != "apps/v1" { + t.Errorf("expected ScaleTargetRef.APIVersion=apps/v1, got %s", ref.APIVersion) + } +} + +// TestBuildHPA_MetricQueueDepth verifies External metric configuration for queueDepth metric. +func TestBuildHPA_MetricQueueDepth(t *testing.T) { + t.Parallel() + + ad := makeAD("query-agent", 1, 5, 75, MetricQueueDepth) + hpa := BuildHPA(ad, 10) + + requireExternalMetricName(t, hpa, customMetricQueueDepth) +} + +// TestBuildHPA_MetricGPUUtilization verifies External metric configuration for gpuUtilization metric. +func TestBuildHPA_MetricGPUUtilization(t *testing.T) { + t.Parallel() + + ad := makeAD("gpu-agent", 1, 4, 80, MetricGPUUtilization) + hpa := BuildHPA(ad, 10) + + requireExternalMetricName(t, hpa, customMetricGPUUtilization) +} + +// TestBuildHPA_MetricUnknownDefaultsToQueueDepth verifies fallback to queueDepth for unrecognized metric names. +func TestBuildHPA_MetricUnknownDefaultsToQueueDepth(t *testing.T) { + t.Parallel() + + // Unknown metric values currently fall back to queue depth. + // This is intentional for forward-compatibility until the CRD enum + // validation makes unknown values impossible at admission time. + ad := makeAD("agent", 1, 3, 50, "unknownMetric") + hpa := BuildHPA(ad, 10) + + requireExternalMetricName(t, hpa, customMetricQueueDepth) +} + +// TestBuildHPA_TargetAverageValue verifies external metric target AverageValue quantity configuration. +func TestBuildHPA_TargetAverageValue(t *testing.T) { + t.Parallel() + + ad := makeAD("agent", 1, 5, 50, MetricQueueDepth) + hpa := BuildHPA(ad, 10) + + if len(hpa.Spec.Metrics) != 1 { + t.Fatalf("expected 1 metric, got %d", len(hpa.Spec.Metrics)) + } + m := hpa.Spec.Metrics[0] + if m.External == nil { + t.Fatal("expected External metric source, got nil") + } + if m.External.Target.Type != autoscalingv2.AverageValueMetricType { + t.Errorf("expected AverageValue target type, got %s", m.External.Target.Type) + } + + want := resource.NewMilliQuantity(50*1000, resource.DecimalSI) + if m.External.Target.AverageValue == nil || !m.External.Target.AverageValue.Equal(*want) { + t.Errorf("expected AverageValue=%s, got %v", want, m.External.Target.AverageValue) + } +} + +// TestBuildHPA_StabilizationWindows verifies scale-up and scale-down stabilization window behavior. +func TestBuildHPA_StabilizationWindows(t *testing.T) { + t.Parallel() + + ad := makeAD("agent", 1, 5, 50, MetricQueueDepth) + hpa := BuildHPA(ad, 10) + + b := hpa.Spec.Behavior + if b == nil { + t.Fatal("expected HPA behavior to be set") + } + if b.ScaleUp == nil || b.ScaleUp.StabilizationWindowSeconds == nil { + t.Fatal("expected ScaleUp stabilization window") + } + if *b.ScaleUp.StabilizationWindowSeconds != scaleUpStabilizationSec { + t.Errorf("ScaleUp window: want %d, got %d", scaleUpStabilizationSec, *b.ScaleUp.StabilizationWindowSeconds) + } + if b.ScaleDown == nil || b.ScaleDown.StabilizationWindowSeconds == nil { + t.Fatal("expected ScaleDown stabilization window") + } + if *b.ScaleDown.StabilizationWindowSeconds != scaleDownStabilizationSec { + t.Errorf("ScaleDown window: want %d, got %d", scaleDownStabilizationSec, *b.ScaleDown.StabilizationWindowSeconds) + } +} + +// TestBuildHPA_LabelsAndNamespace verifies HPA metadata names, namespaces, and tenant labels. +func TestBuildHPA_LabelsAndNamespace(t *testing.T) { + t.Parallel() + + ad := makeAD("agent", 1, 5, 50, MetricQueueDepth) + hpa := BuildHPA(ad, 10) + + if hpa.Name != "agent" { + t.Errorf("expected HPA name=agent, got %s", hpa.Name) + } + if hpa.Namespace != "tenant-test" { + t.Errorf("expected HPA namespace=tenant-test, got %s", hpa.Namespace) + } + if hpa.Labels["agentrax.io/tenant"] != "team-test" { + t.Errorf("expected tenant label, got %v", hpa.Labels) + } +} + +// ── IsQuotaCapped tests ─────────────────────────────────────────────────────── + +// TestIsQuotaCapped_WhenCapped verifies IsQuotaCapped returns true when headroom is less than spec.replicas.max. +func TestIsQuotaCapped_WhenCapped(t *testing.T) { + t.Parallel() + + ad := makeAD("agent", 1, 10, 50, MetricQueueDepth) + if !IsQuotaCapped(ad, 5) { + t.Error("expected IsQuotaCapped=true when headroom < spec.replicas.max") + } +} + +// TestIsQuotaCapped_WhenExact verifies IsQuotaCapped returns false when headroom equals spec.replicas.max. +func TestIsQuotaCapped_WhenExact(t *testing.T) { + t.Parallel() + + ad := makeAD("agent", 1, 5, 50, MetricQueueDepth) + if IsQuotaCapped(ad, 5) { + t.Error("expected IsQuotaCapped=false when headroom == spec.replicas.max") + } +} + +// TestIsQuotaCapped_WhenNotCapped verifies IsQuotaCapped returns false when headroom exceeds spec.replicas.max. +func TestIsQuotaCapped_WhenNotCapped(t *testing.T) { + t.Parallel() + + ad := makeAD("agent", 1, 5, 50, MetricQueueDepth) + if IsQuotaCapped(ad, 10) { + t.Error("expected IsQuotaCapped=false when headroom > spec.replicas.max") + } +} + +// ── QuotaHeadroom tests ─────────────────────────────────────────────────────── + +// TestQuotaHeadroom_PerAgentCeilingLimits verifies QuotaHeadroom respects maxReplicasPerAgent. +func TestQuotaHeadroom_PerAgentCeilingLimits(t *testing.T) { + t.Parallel() + + // maxReplicasPerAgent=4 is smaller than the total budget remaining. + tqSpec := makeTQSpec(20, 4) + adSpec := agentraxv1alpha1.AgentDeploymentSpec{Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 10}} + + h := QuotaHeadroom(tqSpec, adSpec, 0) + if h != 4 { + t.Errorf("expected headroom=4 (per-agent ceiling), got %d", h) + } +} + +// TestQuotaHeadroom_TotalBudgetLimits verifies QuotaHeadroom respects remaining total replica budget. +func TestQuotaHeadroom_TotalBudgetLimits(t *testing.T) { + t.Parallel() + + // Total budget remaining = 20 - 18 = 2; per-agent ceiling = 6. + // Headroom should be min(6, 2) = 2. + tqSpec := makeTQSpec(20, 6) + adSpec := agentraxv1alpha1.AgentDeploymentSpec{Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 6}} + + h := QuotaHeadroom(tqSpec, adSpec, 18) + if h != 2 { + t.Errorf("expected headroom=2 (total budget limited), got %d", h) + } +} + +// TestQuotaHeadroom_ZeroBudgetReturnsZero verifies QuotaHeadroom returns 0 when total budget is exhausted. +func TestQuotaHeadroom_ZeroBudgetReturnsZero(t *testing.T) { + t.Parallel() + + // Total already consumed: QuotaHeadroom should return 0 (raw budget = 0). + // BuildHPA then raises maxReplicas to minReplicas solely for HPA API validity, + // and the reconciler sets QuotaLimited=True because headroom (0) < spec.max (6). + tqSpec := makeTQSpec(10, 6) + adSpec := agentraxv1alpha1.AgentDeploymentSpec{Replicas: agentraxv1alpha1.ScalingPolicy{Min: 2, Max: 6}} + + h := QuotaHeadroom(tqSpec, adSpec, 10) // total budget remaining = 0 + if h != 0 { + t.Errorf("expected headroom=0 when budget exhausted, got %d", h) + } +} + +// TestQuotaHeadroom_NegativeUsedByOthers verifies QuotaHeadroom clamps over-consumed budget to 0. +func TestQuotaHeadroom_NegativeUsedByOthers(t *testing.T) { + t.Parallel() + + // usedReplicasByOthers > MaxTotalReplicas: the negative totalBudgetRemaining + // must be clamped to 0 and headroom returns 0 (raw budget exhausted). + // The reconciler will set QuotaLimited=True and BuildHPA will floor maxReplicas + // at minReplicas=1 solely to keep the HPA spec valid. + tqSpec := makeTQSpec(6, 4) + adSpec := agentraxv1alpha1.AgentDeploymentSpec{Replicas: agentraxv1alpha1.ScalingPolicy{Min: 1, Max: 4}} + + // usedByOthers=10 > MaxTotalReplicas=6 → totalBudgetRemaining clamps to 0; + // headroom = min(perAgentCeiling=4, 0) = 0. + h := QuotaHeadroom(tqSpec, adSpec, 10) + if h != 0 { + t.Errorf("expected headroom=0 when budget is negative (over-consumed), got %d", h) + } +} + +// ── helpers ─────────────────────────────────────────────────────────────────── + +// requireExternalMetricName asserts that the HPA has exactly one External +// metric source with the given metric name. +func requireExternalMetricName(t *testing.T, hpa *autoscalingv2.HorizontalPodAutoscaler, want string) { + t.Helper() + if len(hpa.Spec.Metrics) != 1 { + t.Fatalf("expected 1 metric spec, got %d", len(hpa.Spec.Metrics)) + } + m := hpa.Spec.Metrics[0] + if m.Type != autoscalingv2.ExternalMetricSourceType { + t.Errorf("expected External metric type, got %s", m.Type) + } + if m.External == nil { + t.Fatal("External metric source is nil") + } + if m.External.Metric.Name != want { + t.Errorf("metric name: want %q, got %q", want, m.External.Metric.Name) + } +} diff --git a/internal/webhook/agentdeployment_validator_test.go b/internal/webhook/agentdeployment_validator_test.go index 8f177e6..ec301e7 100644 --- a/internal/webhook/agentdeployment_validator_test.go +++ b/internal/webhook/agentdeployment_validator_test.go @@ -104,6 +104,7 @@ func validCanaryAD() *agentraxv1alpha1.AgentDeployment { // ── ValidateCreate tests ────────────────────────────────────────────────────── +// TestValidateCreate_HappyPath verifies successful admission validation of a valid AgentDeployment. func TestValidateCreate_HappyPath(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -121,6 +122,7 @@ func TestValidateCreate_HappyPath(t *testing.T) { } } +// TestValidateCreate_TenantRefNotFound verifies admission rejection when the referenced TenantQuota does not exist. func TestValidateCreate_TenantRefNotFound(t *testing.T) { t.Parallel() v := newValidatorNoTQ(t) @@ -137,6 +139,7 @@ func TestValidateCreate_TenantRefNotFound(t *testing.T) { } } +// TestValidateCreate_ReplicasMinGtMax verifies admission rejection when replicas.min exceeds replicas.max. func TestValidateCreate_ReplicasMinGtMax(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -153,6 +156,7 @@ func TestValidateCreate_ReplicasMinGtMax(t *testing.T) { } } +// TestValidateCreate_OverMaxAgents verifies admission rejection when tenant maxAgents quota is reached. func TestValidateCreate_OverMaxAgents(t *testing.T) { t.Parallel() tqSpec := agentraxv1alpha1.TenantQuotaSpec{ @@ -174,6 +178,7 @@ func TestValidateCreate_OverMaxAgents(t *testing.T) { } } +// TestValidateCreate_OverMaxReplicasPerAgent verifies admission rejection when replicas.max exceeds maxReplicasPerAgent. func TestValidateCreate_OverMaxReplicasPerAgent(t *testing.T) { t.Parallel() tqSpec := agentraxv1alpha1.TenantQuotaSpec{ @@ -194,6 +199,7 @@ func TestValidateCreate_OverMaxReplicasPerAgent(t *testing.T) { } } +// TestValidateCreate_WrongType verifies that passing an unexpected runtime object type returns an error. func TestValidateCreate_WrongType(t *testing.T) { t.Parallel() v := newValidatorNoTQ(t) @@ -205,6 +211,7 @@ func TestValidateCreate_WrongType(t *testing.T) { // ── validateCanarySpec tests ────────────────────────────────────────────────── +// TestValidateCreate_Canary_HappyPath verifies admission validation of a well-formed canary rollout spec. func TestValidateCreate_Canary_HappyPath(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -214,6 +221,7 @@ func TestValidateCreate_Canary_HappyPath(t *testing.T) { } } +// TestValidateCreate_Canary_NoSteps verifies rejection when canary rollout has no steps. func TestValidateCreate_Canary_NoSteps(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -225,6 +233,7 @@ func TestValidateCreate_Canary_NoSteps(t *testing.T) { } } +// TestValidateCreate_Canary_StepBothFields verifies rejection when a canary step sets both setWeight and pause. func TestValidateCreate_Canary_StepBothFields(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -239,6 +248,7 @@ func TestValidateCreate_Canary_StepBothFields(t *testing.T) { } } +// TestValidateCreate_Canary_NoFullWeight verifies rejection when no canary step promotes to 100% weight. func TestValidateCreate_Canary_NoFullWeight(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -252,6 +262,7 @@ func TestValidateCreate_Canary_NoFullWeight(t *testing.T) { } } +// TestValidateCreate_Canary_MissingMaxErrorRate verifies rejection when canary rollback criteria omits maxErrorRate. func TestValidateCreate_Canary_MissingMaxErrorRate(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -263,6 +274,7 @@ func TestValidateCreate_Canary_MissingMaxErrorRate(t *testing.T) { } } +// TestValidateCreate_Canary_InvalidMaxErrorRate verifies rejection when maxErrorRate is not a valid percentage string. func TestValidateCreate_Canary_InvalidMaxErrorRate(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -274,6 +286,7 @@ func TestValidateCreate_Canary_InvalidMaxErrorRate(t *testing.T) { } } +// TestValidateCreate_Canary_MissingMaxP99 verifies rejection when canary rollback criteria omits maxP99LatencyMs. func TestValidateCreate_Canary_MissingMaxP99(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -285,6 +298,7 @@ func TestValidateCreate_Canary_MissingMaxP99(t *testing.T) { } } +// TestValidateCreate_Canary_MissingMinRequestSample verifies rejection when minRequestSample is missing or zero. func TestValidateCreate_Canary_MissingMinRequestSample(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -298,6 +312,7 @@ func TestValidateCreate_Canary_MissingMinRequestSample(t *testing.T) { // ── validateMCPTools tests ──────────────────────────────────────────────────── +// TestValidateCreate_MCPTools_EmptyName verifies rejection of empty MCP tool names. func TestValidateCreate_MCPTools_EmptyName(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -315,6 +330,7 @@ func TestValidateCreate_MCPTools_EmptyName(t *testing.T) { } } +// TestValidateCreate_MCPTools_Duplicate verifies rejection of duplicate MCP tool names. func TestValidateCreate_MCPTools_Duplicate(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -334,6 +350,7 @@ func TestValidateCreate_MCPTools_Duplicate(t *testing.T) { // ── ValidateUpdate tests ────────────────────────────────────────────────────── +// TestValidateUpdate_DeletionTimestampBypass verifies that objects being deleted bypass quota checks to avoid deadlocking finalizers. func TestValidateUpdate_DeletionTimestampBypass(t *testing.T) { t.Parallel() // Even with a TQ that would reject quota, updates with DeletionTimestamp @@ -358,6 +375,7 @@ func TestValidateUpdate_DeletionTimestampBypass(t *testing.T) { } } +// TestValidateUpdate_RolloutInProgress_BlocksImageChange verifies that image changes are blocked during an active rollout. func TestValidateUpdate_RolloutInProgress_BlocksImageChange(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -377,6 +395,7 @@ func TestValidateUpdate_RolloutInProgress_BlocksImageChange(t *testing.T) { } } +// TestValidateUpdate_RolloutInProgress_AllowsNonImageChange verifies that non-image updates (e.g. labels) are permitted during rollout. func TestValidateUpdate_RolloutInProgress_AllowsNonImageChange(t *testing.T) { t.Parallel() v := newValidatorWithTQ(t, permissiveTQSpec(), agentraxv1alpha1.TenantQuotaStatus{}) @@ -397,6 +416,7 @@ func TestValidateUpdate_RolloutInProgress_AllowsNonImageChange(t *testing.T) { } } +// TestValidateUpdate_WrongType verifies that passing wrong object types to ValidateUpdate returns an error. func TestValidateUpdate_WrongType(t *testing.T) { t.Parallel() v := newValidatorNoTQ(t) @@ -411,6 +431,7 @@ func TestValidateUpdate_WrongType(t *testing.T) { // ── ValidateDelete test ─────────────────────────────────────────────────────── +// TestValidateDelete_AlwaysNil verifies that ValidateDelete always allows deletion without error. func TestValidateDelete_AlwaysNil(t *testing.T) { t.Parallel() v := newValidatorNoTQ(t) diff --git a/internal/webhook/agentdeployment_webhook.go b/internal/webhook/agentdeployment_webhook.go index 8913faf..d11a225 100644 --- a/internal/webhook/agentdeployment_webhook.go +++ b/internal/webhook/agentdeployment_webhook.go @@ -57,7 +57,7 @@ func SetupAgentDeploymentWebhookWithManager(mgr ctrl.Manager, enforcer *quota.En For(&agentraxv1alpha1.AgentDeployment{}). WithDefaulter(&AgentDeploymentCustomDefaulter{}). WithValidator(&AgentDeploymentCustomValidator{ - Client: mgr.GetClient(), + Client: mgr.GetAPIReader(), Enforcer: enforcer, }). Complete() @@ -118,7 +118,7 @@ func isZeroResources(r corev1.ResourceRequirements) bool { // AgentDeploymentCustomValidator validates AgentDeployment specs at admission // time. It implements admission.CustomValidator. type AgentDeploymentCustomValidator struct { - Client client.Client + Client client.Reader Enforcer *quota.Enforcer } @@ -126,6 +126,7 @@ var _ webhook.CustomValidator = &AgentDeploymentCustomValidator{} // ValidateCreate validates a new AgentDeployment against quota and spec rules. func (v *AgentDeploymentCustomValidator) ValidateCreate(ctx context.Context, obj runtime.Object) (admission.Warnings, error) { + start := time.Now() ad, ok := obj.(*agentraxv1alpha1.AgentDeployment) if !ok { return nil, fmt.Errorf("expected AgentDeployment, got %T", obj) @@ -133,6 +134,22 @@ func (v *AgentDeploymentCustomValidator) ValidateCreate(ctx context.Context, obj webhookLog.Info("validating create", "name", ad.Name, "namespace", ad.Namespace) allErrs := v.validateSpec(ctx, ad, nil) + + latency := time.Since(start) + webhookLog.V(1).Info("admission complete", + "operation", "create", + "name", ad.Name, + "namespace", ad.Namespace, + "latency_ms", latency.Milliseconds(), + "admitted", len(allErrs) == 0) + if latency > 2*time.Second { + webhookLog.Info("slow admission request detected", + "operation", "create", + "name", ad.Name, + "namespace", ad.Namespace, + "latency_ms", latency.Milliseconds()) + } + if len(allErrs) > 0 { return nil, apierrors.NewInvalid( agentraxv1alpha1.GroupVersion.WithKind("AgentDeployment").GroupKind(), @@ -143,6 +160,7 @@ func (v *AgentDeploymentCustomValidator) ValidateCreate(ctx context.Context, obj // ValidateUpdate validates an updated AgentDeployment against quota and spec rules. func (v *AgentDeploymentCustomValidator) ValidateUpdate(ctx context.Context, oldObj, newObj runtime.Object) (admission.Warnings, error) { + start := time.Now() ad, ok := newObj.(*agentraxv1alpha1.AgentDeployment) if !ok { return nil, fmt.Errorf("expected AgentDeployment, got %T", newObj) @@ -172,6 +190,22 @@ func (v *AgentDeploymentCustomValidator) ValidateUpdate(ctx context.Context, old } allErrs := v.validateSpec(ctx, ad, &oldAD.Spec) + + latency := time.Since(start) + webhookLog.V(1).Info("admission complete", + "operation", "update", + "name", ad.Name, + "namespace", ad.Namespace, + "latency_ms", latency.Milliseconds(), + "admitted", len(allErrs) == 0) + if latency > 2*time.Second { + webhookLog.Info("slow admission request detected", + "operation", "update", + "name", ad.Name, + "namespace", ad.Namespace, + "latency_ms", latency.Milliseconds()) + } + if len(allErrs) > 0 { return nil, apierrors.NewInvalid( agentraxv1alpha1.GroupVersion.WithKind("AgentDeployment").GroupKind(), @@ -230,18 +264,36 @@ func (v *AgentDeploymentCustomValidator) validateSpec( allErrs = append(allErrs, validateMCPTools(specPath.Child("mcp", "tools"), ad.Spec.MCP.Tools)...) } - // ── 5. Quota admission check ── - // Compute current usage from the TenantQuota status. The status is kept - // accurate by the TenantQuota reconciler; we add in-flight reservations to - // handle concurrent near-limit creates. + // ── 5. Quota admission check (atomic check-and-reserve) ── + // AdmitAndReserve holds the in-flight mutex across both the quota check + // and the reservation write, preventing concurrent near-limit creates from + // both slipping through the quota ceiling. + // Only run when all earlier checks pass: a spec that is already invalid + // must not create a reservation that would transiently block valid admits. admissionKey := fmt.Sprintf("%s/%s", ad.Namespace, ad.Name) - ok, reason := v.Enforcer.CanAdmit(admissionKey, tq.Spec, tq.Status, ad.Spec, oldSpec) - if !ok { - allErrs = append(allErrs, field.Forbidden(specPath, fmt.Sprintf("quota exceeded: %s", reason))) - } else { - // Reserve in-flight slot for the duration the webhook response is in transit. - // The reservation is automatically swept after reservationTTL. - v.Enforcer.Reserve(admissionKey, ad.Spec, oldSpec, reservationTTL) + if len(allErrs) == 0 { + // Detect server-side dry-run to avoid creating reservations for + // requests that will never be persisted. + isDryRun := false + if req, err := admission.RequestFromContext(ctx); err == nil && req.DryRun != nil && *req.DryRun { + isDryRun = true + } + + if isDryRun { + // For dry-run requests, perform the quota check without writing + // to the reservations map. This still acquires the mutex to read + // consistent in-flight state, but does not create a reservation. + ok, reason := v.Enforcer.CanAdmit(admissionKey, tq.Spec, tq.Status, ad.Spec, oldSpec) + if !ok { + allErrs = append(allErrs, field.Forbidden(specPath, fmt.Sprintf("quota exceeded: %s", reason))) + } + } else { + // For persisted requests, use the atomic check-and-reserve. + ok, reason := v.Enforcer.AdmitAndReserve(admissionKey, tq.Spec, tq.Status, ad.Spec, oldSpec, reservationTTL) + if !ok { + allErrs = append(allErrs, field.Forbidden(specPath, fmt.Sprintf("quota exceeded: %s", reason))) + } + } } return allErrs diff --git a/internal/webhook/agentdeployment_webhook_test.go b/internal/webhook/agentdeployment_webhook_test.go index 910ce05..06279ec 100644 --- a/internal/webhook/agentdeployment_webhook_test.go +++ b/internal/webhook/agentdeployment_webhook_test.go @@ -33,6 +33,7 @@ func newDefaulter() *agentraxwebhook.AgentDeploymentCustomDefaulter { return &agentraxwebhook.AgentDeploymentCustomDefaulter{} } +// baseAD returns a minimal AgentDeployment fixture for defaulter testing. func baseAD() *agentraxv1alpha1.AgentDeployment { return &agentraxv1alpha1.AgentDeployment{ ObjectMeta: metav1.ObjectMeta{Name: "test-ad", Namespace: "test-ns"}, @@ -44,6 +45,7 @@ func baseAD() *agentraxv1alpha1.AgentDeployment { } } +// TestDefaulter_PortDefaulted verifies that port defaults to 8080 when unset. func TestDefaulter_PortDefaulted(t *testing.T) { t.Parallel() ad := baseAD() @@ -55,6 +57,7 @@ func TestDefaulter_PortDefaulted(t *testing.T) { } } +// TestDefaulter_PortNotOverwritten verifies that an explicit port is preserved. func TestDefaulter_PortNotOverwritten(t *testing.T) { t.Parallel() ad := baseAD() @@ -67,6 +70,7 @@ func TestDefaulter_PortNotOverwritten(t *testing.T) { } } +// TestDefaulter_StrategyDefaulted verifies that rollout strategy defaults to Recreate when unset. func TestDefaulter_StrategyDefaulted(t *testing.T) { t.Parallel() ad := baseAD() @@ -78,6 +82,7 @@ func TestDefaulter_StrategyDefaulted(t *testing.T) { } } +// TestDefaulter_StrategyNotOverwritten verifies that explicit rollout strategy is preserved. func TestDefaulter_StrategyNotOverwritten(t *testing.T) { t.Parallel() ad := baseAD() @@ -90,6 +95,7 @@ func TestDefaulter_StrategyNotOverwritten(t *testing.T) { } } +// TestDefaulter_ResourcesDefaulted verifies that default CPU/memory requests and limits are populated when unset. func TestDefaulter_ResourcesDefaulted(t *testing.T) { t.Parallel() ad := baseAD() @@ -109,6 +115,7 @@ func TestDefaulter_ResourcesDefaulted(t *testing.T) { } } +// TestDefaulter_ResourcesNotOverwritten verifies that explicitly specified resource requirements are preserved. func TestDefaulter_ResourcesNotOverwritten(t *testing.T) { t.Parallel() ad := baseAD() @@ -126,6 +133,7 @@ func TestDefaulter_ResourcesNotOverwritten(t *testing.T) { } } +// TestDefaulter_WrongType verifies that passing an unexpected runtime object type returns an error. func TestDefaulter_WrongType(t *testing.T) { t.Parallel() d := newDefaulter() diff --git a/test/e2e/scaling_test.go b/test/e2e/scaling_test.go new file mode 100644 index 0000000..186ae7d --- /dev/null +++ b/test/e2e/scaling_test.go @@ -0,0 +1,293 @@ +//go:build e2e + +/* +Copyright 2026. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package e2e contains end-to-end tests that run against a real kind cluster +// with Prometheus, Prometheus Adapter, and cert-manager installed. +// Build with -tags e2e to include this file; it is excluded from normal CI. +// Run with: make e2e-soak +package e2e + +// scaling_test.go — Phase 3 soak test for metrics-driven autoscaling. +// +// This file verifies the full autoscaling lifecycle: +// - An AgentDeployment with metric:queueDepth produces a correctly configured HPA. +// - Injecting synthetic queue-depth load causes replica count to increase +// within one HPA polling cycle (default 15s). +// - Removing load causes scale-down no faster than the 300s stabilization +// window (verifiable: no flapping observed in a 10-minute soak). +// - Scale-up capped by tenant quota: QuotaLimited condition is set when +// the HPA would exceed the tenant's remaining replica headroom. +// - spec.replicas.min is never violated (starts at min, not zero). +// +// Prerequisites (installed by `make deploy-deps`): +// - Prometheus Operator + Prometheus scraping agentrax pods +// - Prometheus Adapter with agentrax custom-metrics-config.yaml applied +// - cert-manager (for webhook TLS) +// - A conformant Gateway API implementation (for Phase 4 — optional here) +// +// The soak load generator uses a Kubernetes Job that pushes synthetic +// agentrax_queue_depth metrics via a pushgateway stub, or alternatively +// a direct metric exporter sidecar injected by the test. + +import ( + "context" + "fmt" + "io" + "os/exec" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/gitcommitankit/agentrax/test/utils" +) + +const ( + // soakNamespace is the namespace used for all soak test resources. + soakNamespace = "agentrax-soak" + + // soakTenantQuota is the name of the TenantQuota used in soak tests. + soakTenantQuota = "soak-quota" + + // soakAgentName is the name of the AgentDeployment under test. + soakAgentName = "soak-agent" + + // hpaPollingInterval is the time to wait for the HPA to react to a metric change. + // Kubernetes HPA default polling is 15s; we allow 3× headroom. + hpaPollingInterval = 60 * time.Second + + // soakDuration is how long the no-flapping soak runs (10 minutes per DoD). + soakDuration = 10 * time.Minute + + // scaleDownStabilizationWindow is the expected HPA scale-down stabilization. + // The soak verifies no scale-down occurs faster than this after load is removed. + scaleDownStabilizationWindow = 5 * time.Minute +) + +var _ = Describe("Phase 3 — Metrics-Driven Autoscaling (soak)", Ordered, func() { + ctx := context.Background() + _ = ctx // used in BeforeAll/AfterAll below + + BeforeAll(func() { + By("creating soak namespace") + cmd := exec.Command("kubectl", "create", "ns", soakNamespace, "--dry-run=client", "-o", "yaml") + out, err := utils.Run(cmd) + Expect(err).NotTo(HaveOccurred()) + applyCmd := exec.Command("kubectl", "apply", "-f", "-") + applyCmd.Stdin = bytesReader(out) + _, err = utils.Run(applyCmd) + Expect(err).NotTo(HaveOccurred()) + + By("applying TenantQuota for soak tests") + tqYAML := fmt.Sprintf(` +apiVersion: agentrax.io/v1alpha1 +kind: TenantQuota +metadata: + name: %s + namespace: %s +spec: + maxAgents: 3 + maxGPUs: 0 + maxTotalReplicas: 12 + maxReplicasPerAgent: 6 +`, soakTenantQuota, soakNamespace) + Expect(kubectlApplyStdin(tqYAML)).To(Succeed()) + }) + + AfterAll(func() { + By("cleaning up soak namespace") + cmd := exec.Command("kubectl", "delete", "ns", soakNamespace, "--ignore-not-found") + _, _ = utils.Run(cmd) + }) + + // ── Scenario 1: HPA created with correct spec ───────────────────────────── + + It("creates a managed HPA when an AgentDeployment is applied", func() { + PendingNow("requires Prometheus Adapter installed via make deploy-deps") + + adYAML := fmt.Sprintf(` +apiVersion: agentrax.io/v1alpha1 +kind: AgentDeployment +metadata: + name: %s + namespace: %s +spec: + image: nginx:latest + tenantRef: %s + replicas: + min: 1 + max: 6 + metric: queueDepth + target: 50 +`, soakAgentName, soakNamespace, soakTenantQuota) + Expect(kubectlApplyStdin(adYAML)).To(Succeed()) + + By("waiting for the HPA to be created") + Eventually(func() error { + cmd := exec.Command("kubectl", "get", "hpa", soakAgentName, "-n", soakNamespace) + _, err := utils.Run(cmd) + return err + }, 30*time.Second, time.Second).Should(Succeed()) + + By("verifying HPA scaleTargetRef points to the Deployment") + cmd := exec.Command("kubectl", "get", "hpa", soakAgentName, "-n", soakNamespace, + "-o", "jsonpath={.spec.scaleTargetRef.name}") + out, err := utils.Run(cmd) + Expect(err).NotTo(HaveOccurred()) + Expect(string(out)).To(Equal(soakAgentName)) + + By("verifying HPA metric is agentrax_queue_depth (External type)") + metricCmd := exec.Command("kubectl", "get", "hpa", soakAgentName, "-n", soakNamespace, + "-o", "jsonpath={.spec.metrics[0].external.metric.name}") + metricOut, err := utils.Run(metricCmd) + Expect(err).NotTo(HaveOccurred()) + Expect(string(metricOut)).To(Equal("agentrax_queue_depth")) + }) + + // ── Scenario 2: Scale-out under synthetic load ──────────────────────────── + + It("scales replicas up when queue depth exceeds the target threshold", func() { + PendingNow("requires synthetic load generator + Prometheus Adapter + kind cluster") + + By("injecting synthetic load: push agentrax_queue_depth > 50 per replica") + // TODO: deploy the load-generator Job from test/fixtures/load-gen.yaml. + // The load generator writes time-series to a Pushgateway stub that + // Prometheus scrapes per the ServiceMonitor. + + By("waiting for replica count to exceed the initial minimum") + Eventually(func() int { + cmd := exec.Command("kubectl", "get", "deploy", soakAgentName, + "-n", soakNamespace, "-o", "jsonpath={.status.readyReplicas}") + out, err := utils.Run(cmd) + if err != nil { + return -1 + } + count := 0 + _, _ = fmt.Sscanf(string(out), "%d", &count) + return count + }, hpaPollingInterval, 5*time.Second).Should(BeNumerically(">", 1), + "replica count should increase above the minimum under load") + }) + + // ── Scenario 3: No flapping during 10-minute soak ──────────────────────── + + It("does not flap replicas during a 10-minute steady-state soak", func() { + PendingNow("requires synthetic load generator + kind cluster running for 10+ minutes") + + By(fmt.Sprintf("observing replica count for %s under steady load", soakDuration)) + // Track replica count samples; assert coefficient of variation is < 20%. + // Flapping definition: more than 2 scale events in a 5-minute window. + + // TODO: collect replica samples every 30s for soakDuration. + // Fail if abs(sample[i+1] - sample[i]) > 1 more than twice in any 5-minute window. + }) + + // ── Scenario 4: Scale-down respects stabilization window ───────────────── + + It("does not scale down faster than the 300s stabilization window after load removal", func() { + PendingNow("requires load generator + timing control in kind cluster") + + By("removing synthetic load") + // TODO: delete the load-generator Job. + + By(fmt.Sprintf("asserting replica count stays elevated for at least %s", scaleDownStabilizationWindow)) + // Sample replica count every 30s for scaleDownStabilizationWindow. + // Fail if replicas drop before the window elapses. + }) + + // ── Scenario 5: QuotaLimited condition when scale-up exceeds quota ──────── + + It("sets QuotaLimited condition and caps HPA when quota headroom is exhausted", func() { + PendingNow("requires a second AgentDeployment consuming remaining quota headroom") + + By("creating a second AD that consumes the remaining replica budget") + // TODO: apply a second AD with max=high to exhaust the TenantQuota headroom. + + By("verifying the first AD's QuotaLimited condition is True") + Eventually(func() bool { + cmd := exec.Command("kubectl", "get", "agentdeployment", soakAgentName, + "-n", soakNamespace, + "-o", `jsonpath={.status.conditions[?(@.type=="QuotaLimited")].status}`) + out, err := utils.Run(cmd) + if err != nil { + return false + } + return string(out) == "True" + }, 30*time.Second, time.Second).Should(BeTrue(), + "QuotaLimited condition should be True when HPA is capped by quota") + }) + + // ── Scenario 6: spec.replicas.min never violated ────────────────────────── + + It("never scales below spec.replicas.min even at zero traffic", func() { + PendingNow("requires zero-traffic steady state in kind cluster") + + By("removing all load and waiting for the stabilization window") + // TODO: ensure no load for scaleDownStabilizationWindow + buffer. + + By("asserting ready replicas >= spec.replicas.min") + Consistently(func() int { + cmd := exec.Command("kubectl", "get", "deploy", soakAgentName, + "-n", soakNamespace, "-o", "jsonpath={.status.readyReplicas}") + out, err := utils.Run(cmd) + if err != nil { + return -1 + } + count := 0 + _, _ = fmt.Sscanf(string(out), "%d", &count) + return count + }, 2*time.Minute, 10*time.Second).Should(BeNumerically(">=", 1), + "ready replicas must never drop below spec.replicas.min=1") + }) +}) + +// ── helpers ────────────────────────────────────────────────────────────────── + +// kubectlApplyStdin pipes a YAML string to kubectl apply -f -. +func kubectlApplyStdin(yaml string) error { + cmd := exec.Command("kubectl", "apply", "-f", "-") + cmd.Stdin = stringReader(yaml) + _, err := utils.Run(cmd) + return err +} + +// stringReader wraps a string as an io.Reader for test inputs. +func stringReader(s string) *bytesReaderWrapper { + return &bytesReaderWrapper{data: []byte(s), pos: 0} +} + +// bytesReader wraps a byte slice as an io.Reader for test inputs. +func bytesReader(b []byte) *bytesReaderWrapper { + return &bytesReaderWrapper{data: b, pos: 0} +} + +// bytesReaderWrapper is a minimal io.Reader over a byte slice. +type bytesReaderWrapper struct { + data []byte + pos int +} + +// Read implements io.Reader for bytesReaderWrapper. +func (r *bytesReaderWrapper) Read(p []byte) (n int, err error) { + if r.pos >= len(r.data) { + return 0, io.EOF + } + n = copy(p, r.data[r.pos:]) + r.pos += n + return n, nil +} diff --git a/test/utils/utils.go b/test/utils/utils.go index 9c1691e..d49f72e 100644 --- a/test/utils/utils.go +++ b/test/utils/utils.go @@ -35,6 +35,7 @@ const ( certmanagerURLTmpl = "https://github.com/jetstack/cert-manager/releases/download/%s/cert-manager.yaml" ) +// warnError logs a non-fatal warning message to the GinkgoWriter. func warnError(err error) { _, _ = fmt.Fprintf(GinkgoWriter, "warning: %v\n", err) }