Add heartbeat discovery and provider-neutral fleet planning - #196
Add heartbeat discovery and provider-neutral fleet planning#196nikhilsimha wants to merge 1 commit into
Conversation
|
🔍 OpenCodeReview found 11 issue(s) in this PR.
|
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| if err := decodeFleetAdminJSON(r, &request); err != nil || strings.TrimSpace(request.ServiceInstanceID) == "" { | ||
| http.Error(w, "valid serviceInstanceId is required", http.StatusBadRequest) | ||
| return | ||
| } |
There was a problem hiding this comment.
TrimSpace is used only for the emptiness check, but the original value is forwarded to the scheduler. Because SetCordoned compares service-instance IDs byte-for-byte, a request such as {"serviceInstanceId":" service-a "} passes this gateway validation and then fails with a service-instance mismatch instead of acting on the canonical ID. Normalize once and send the trimmed value (or reject whitespace-containing IDs consistently).
Suggestion:
| if err := decodeFleetAdminJSON(r, &request); err != nil || strings.TrimSpace(request.ServiceInstanceID) == "" { | |
| http.Error(w, "valid serviceInstanceId is required", http.StatusBadRequest) | |
| return | |
| } | |
| if err := decodeFleetAdminJSON(r, &request); err != nil { | |
| http.Error(w, "valid serviceInstanceId is required", http.StatusBadRequest) | |
| return | |
| } | |
| request.ServiceInstanceID = strings.TrimSpace(request.ServiceInstanceID) | |
| if request.ServiceInstanceID == "" { | |
| http.Error(w, "valid serviceInstanceId is required", http.StatusBadRequest) | |
| return | |
| } |
| func decodeFleetAdminJSON(r *http.Request, target any) error { | ||
| decoder := json.NewDecoder(io.LimitReader(r.Body, maxFleetAdminBody+1)) | ||
| decoder.DisallowUnknownFields() |
There was a problem hiding this comment.
The LimitReader does not enforce the advertised 1 MiB maximum: a valid JSON document that fits within the limit followed by more than the allowed amount of whitespace is accepted, because the second decode sees EOF after the limited reader is exhausted. Track whether the limit was reached (or read one extra byte after decoding) and return 400 when the body exceeds maxFleetAdminBody.
Suggestion:
| func decodeFleetAdminJSON(r *http.Request, target any) error { | |
| decoder := json.NewDecoder(io.LimitReader(r.Body, maxFleetAdminBody+1)) | |
| decoder.DisallowUnknownFields() | |
| func decodeFleetAdminJSON(r *http.Request, target any) error { | |
| limited := io.LimitReader(r.Body, maxFleetAdminBody+1) | |
| decoder := json.NewDecoder(limited) | |
| decoder.DisallowUnknownFields() |
| if status.Code(err) != codes.Unavailable { | ||
| return response, err | ||
| } |
There was a problem hiding this comment.
This retries every Schedule RPC that returns codes.Unavailable, not just the scheduler's temporary no nodes available condition. Unavailable can also be produced by gRPC transport/server availability failures, and this loop has no retry cap or backoff beyond a fixed one-second sleep; during a scheduler outage, each gateway request will repeatedly issue RPCs until its request timeout, amplifying load and delaying the error response. Restrict retries to an explicitly identifiable capacity-exhaustion signal or add bounded/backoff retry behavior.
Suggestion:
| if status.Code(err) != codes.Unavailable { | |
| return response, err | |
| } | |
| if status.Code(err) != codes.Unavailable || status.Convert(err).Message() != "no nodes available" { | |
| return response, err | |
| } |
| observed := p.nodes.ListObserved("", now) | ||
| observedIDs := make(map[string]struct{}, len(observed)) |
There was a problem hiding this comment.
fleetNodeIDs is the caller-supplied membership boundary, but this query loads every observed node because ListObserved("") treats an empty cluster filter as global. The planner then counts unrelated nodes and can emit cordon, delete, or uncordon actions for nodes that are not in the requested fleet whenever the registry contains more than one fleet (or the request is a subset). Filter observed nodes and all candidate selection to the unique fleet IDs, or otherwise enforce ownership before calculating the plan.
Suggestion:
| observed := p.nodes.ListObserved("", now) | |
| observedIDs := make(map[string]struct{}, len(observed)) | |
| fleetNodeSet := make(map[string]struct{}, len(fleetNodeIDs)) | |
| for _, nodeID := range fleetNodeIDs { | |
| if nodeID != "" { | |
| fleetNodeSet[nodeID] = struct{}{} | |
| } | |
| } | |
| observed := p.nodes.ListObserved("", now) | |
| observedIDs := make(map[string]struct{}, len(observed)) |
| memoryReserved += snapshot.GetAllocatedMemoryBytes() + snapshot.GetPausedAllocatedMemoryBytes() | ||
| if snapshot.GetMemoryTotalBytes() > 0 { | ||
| memoryTotal += snapshot.GetMemoryTotalBytes() |
There was a problem hiding this comment.
These uint64 additions can wrap on large node-reported values before the planner computes capacity. The protobuf fields are uint64 and this path does not validate or cap them, so a node reporting values near MaxUint64 can make memoryReserved or memoryTotal artificially small; the subsequent averageTotal * MaxMemoryUsedPercent multiplication has the same overflow risk. Use checked/saturating arithmetic (or reject invalid telemetry) so overflow cannot cause under-provisioning.
Suggestion:
| memoryReserved += snapshot.GetAllocatedMemoryBytes() + snapshot.GetPausedAllocatedMemoryBytes() | |
| if snapshot.GetMemoryTotalBytes() > 0 { | |
| memoryTotal += snapshot.GetMemoryTotalBytes() | |
| memoryReserved = addUint64Saturating(memoryReserved, addUint64Saturating(snapshot.GetAllocatedMemoryBytes(), snapshot.GetPausedAllocatedMemoryBytes())) | |
| if snapshot.GetMemoryTotalBytes() > 0 { | |
| memoryTotal = addUint64Saturating(memoryTotal, snapshot.GetMemoryTotalBytes()) |
| func (s *Service) GetFleetPlan(_ context.Context, req *schedulerv1.GetFleetPlanRequest) (*schedulerv1.GetFleetPlanResponse, error) { | ||
| if s.fleetPlanner == nil { | ||
| return nil, status.Error(codes.FailedPrecondition, "scheduler fleet planning is disabled") | ||
| } | ||
| plan := s.fleetPlanner.Plan(req.GetFleetNodeIds(), time.Now()) |
There was a problem hiding this comment.
These administrative RPCs have no method-level authentication or authorization. The scheduler is started with only the metrics unary interceptor, so any caller that can reach the scheduler gRPC listen address can call CordonNode/UncordonNode to change scheduling state or call GetFleetPlan to inspect topology, bypassing the gateway's HTTP authentication. Enforce scheduler-side authentication/authorization for these RPCs (or otherwise ensure the gRPC listener is cryptographically and network-bound to trusted callers).
| if err := validateFleetNodeReference(req.GetNodeId(), req.GetServiceInstanceId()); err != nil { | ||
| return nil, err | ||
| } | ||
| if err := s.nodes.SetCordoned(req.GetNodeId(), req.GetServiceInstanceId(), true); err != nil { | ||
| return nil, fleetNodeStateError(err) | ||
| } | ||
| s.fleetPlanner.MarkCordoned(req.GetNodeId(), time.Now()) |
There was a problem hiding this comment.
Validation trims these identifiers only for the empty check, but the untrimmed request values are passed to the registry. A request such as node_id: " node-1 " passes validation and then fails lookup (or can create a whitespace-named discovered node), while the planner state is keyed by the canonical ID. Trim once and pass the canonical values to both SetCordoned and MarkCordoned/MarkUncordoned.
| if parsed.Fleet != nil { | ||
| s.Fleet = *parsed.Fleet | ||
| } |
There was a problem hiding this comment.
When loading a config file, cfg starts with defaultConfig and this unmarshaler is expected to preserve omitted fields. However, unmarshaling fleet into a new zero-valued SchedulerFleetConfig and assigning it here discards the existing defaults for omitted fields, notably min_nodes and warm_nodes (which applyDefaults does not restore). Thus a valid config such as { "scheduler": { "fleet": { "enabled": true } } } fails validation with scheduler.fleet requires positive min_nodes <= max_nodes instead of inheriting the defaults. Unmarshal into the existing s.Fleet value (or merge the parsed fields) before assigning.
Suggestion:
| if parsed.Fleet != nil { | |
| s.Fleet = *parsed.Fleet | |
| } | |
| if parsed.Fleet != nil { | |
| if err := json.Unmarshal(data, &struct { | |
| Fleet *SchedulerFleetConfig `json:"fleet"` | |
| }{Fleet: &s.Fleet}); err != nil { | |
| return err | |
| } | |
| } |
| /// Shared secret used to authenticate heartbeat-based registration. | ||
| #[config( | ||
| env = "AENV_OBSERVABILITY_REGISTRATION_TOKEN", | ||
| parse_env = parse_trimmed_string | ||
| )] | ||
| pub registration_token: Option<String>, |
There was a problem hiding this comment.
ObservabilitySchedulerReportConfig still derives Debug, so any debug formatting of AppConfig or ObservabilityConfig will include this shared registration secret in plaintext. This credential should be redacted in the type's Debug implementation (as is done for other secret-bearing config structs) or otherwise excluded from diagnostic output.
| address: endpoint.address.clone(), | ||
| }), | ||
| endpoint: node_endpoint.to_string(), | ||
| registration_token: registration_token.to_string(), |
There was a problem hiding this comment.
The registration secret is now included in every heartbeat, but scheduler_endpoint is accepted as an http:// URI by Endpoint::from_shared, so this can transmit the token in plaintext over the network. An observer could replay it to register nodes when heartbeat discovery is enabled. Require/use TLS for token-authenticated heartbeats (or reject non-HTTPS scheduler endpoints when a token is configured) before sending this field.
Suggestion:
| registration_token: registration_token.to_string(), | |
| registration_token: registration_token.to_string(), |
1 participant
Why this change
AgentENV knows which sandboxes exist, which nodes are full, and whether a node is safe to remove. A cloud autoscaler only sees VM-level signals, so it can remove a node that still owns a running or paused sandbox. A static node list also means every node addition or replacement requires a Scheduler config update and restart.
This change lets AgentENV make the sandbox-aware capacity and drain decisions while a small external executor creates and deletes VMs. We get automatic scale-out, safe scale-in, and node replacement without putting cloud credentials or cloud-specific code inside AgentENV.
High-level approach
AgentENV never receives cloud credentials, and the infrastructure executor never guesses whether a sandbox node is safe to remove.
Summary
Safety
Validation
make -C services testcargo fmt --all -- --checkcargo clippy -p agentenv --lib -- -D warningscargo test -p agentenv observability::reporter --lib