From 92b2089ee24f2cdfdc660666d2e8b88fbe9e81e8 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 13:05:20 +0800 Subject: [PATCH 01/26] feat(scheduler): add leader-election config surface Introduce scheduler.leader_election config (enabled, lease_name, lease_namespace, lease_duration, renew_deadline, retry_period) for Kubernetes Lease-based scheduler HA (#259), plus: - scheduler.node_admin_api_key for sync-node-snapshots pulls against node admin APIs; it is a credential, so SCHEDULER_NODE_ADMIN_API_KEY env (mounted from the agentenv-auth Secret) overrides it and is the recommended path on Kubernetes. - leader_election.snapshot_pull_concurrency bounding the post-acquisition snapshot pull fan-out; zero means the default of 4. Scope: services/shared/config only. Disabled by default; when disabled the scheduler behaves exactly as a single-writer process. Startup validation (fail fast): mutually exclusive with --query-only; client-go timing rule lease_duration > renew_deadline > 2*retry_period. Redis stays optional (without it, failover loses pre-failover sandbox routing; documented degraded mode). Tests: validation cases for no-redis pass, query-only conflict, invalid timing triples, disabled transparency, and the snapshot_pull_concurrency default. Relates to #259, helps #191. Co-Authored-By: Claude Code --- services/shared/config/config.go | 141 ++++++++++++++++-- .../config/scheduler_leader_election_test.go | 120 +++++++++++++++ 2 files changed, 250 insertions(+), 11 deletions(-) create mode 100644 services/shared/config/scheduler_leader_election_test.go diff --git a/services/shared/config/config.go b/services/shared/config/config.go index 7977084e2..7a731b5f0 100644 --- a/services/shared/config/config.go +++ b/services/shared/config/config.go @@ -57,18 +57,42 @@ type NodeResourceLimit struct { MaxAllocatedMemoryBytesIncludingPaused *uint64 `json:"max_allocated_memory_bytes_including_paused"` } +// SchedulerLeaderElectionConfig configures Kubernetes Lease-based leader +// election for the scheduler (#259). Disabled by default; when disabled the +// scheduler behaves exactly as a single-writer process (today's behaviour). +// A Redis binding store is recommended but not required: without it, failover +// loses routing for pre-failover sandboxes (degraded mode). +type SchedulerLeaderElectionConfig struct { + Enabled bool `json:"enabled"` + LeaseName string `json:"lease_name"` + LeaseNamespace string `json:"lease_namespace"` + LeaseDuration time.Duration `json:"lease_duration"` + RenewDeadline time.Duration `json:"renew_deadline"` + RetryPeriod time.Duration `json:"retry_period"` + // SnapshotPullConcurrency bounds how many nodes the post-acquisition + // sync-node-snapshots pull fans out to at once. Zero means the default. + SnapshotPullConcurrency int `json:"snapshot_pull_concurrency"` +} + type SchedulerConfig struct { - GRPCListenAddr string `json:"grpc_listen_addr"` - MetricsListenAddr string `json:"metrics_listen_addr"` - Strategy string `json:"strategy"` - ReportTTL time.Duration `json:"report_ttl"` - BindingTTL time.Duration `json:"binding_ttl"` - RedisAddr string `json:"redis_addr"` - ArtifactStoreCapacity int `json:"artifact_store_capacity"` - ArtifactLookupNodeLimit int `json:"artifact_lookup_node_limit"` - Nodes []Node `json:"nodes"` - Discovery SchedulerDiscoveryConfig `json:"discovery"` - NodeResourceLimit *NodeResourceLimit `json:"node_resource_limit"` + GRPCListenAddr string `json:"grpc_listen_addr"` + MetricsListenAddr string `json:"metrics_listen_addr"` + Strategy string `json:"strategy"` + ReportTTL time.Duration `json:"report_ttl"` + BindingTTL time.Duration `json:"binding_ttl"` + RedisAddr string `json:"redis_addr"` + // NodeAdminAPIKey authenticates the scheduler against each node's admin + // API (x-api-key header) for sync-node-snapshots pulls (#259). It is a + // credential: prefer the SCHEDULER_NODE_ADMIN_API_KEY env (mounted from + // the agentenv-auth Secret) over this JSON field. Optional; without it, + // pulls fail auth and nodes stay unobserved until heartbeats. + NodeAdminAPIKey string `json:"node_admin_api_key"` + ArtifactStoreCapacity int `json:"artifact_store_capacity"` + ArtifactLookupNodeLimit int `json:"artifact_lookup_node_limit"` + Nodes []Node `json:"nodes"` + Discovery SchedulerDiscoveryConfig `json:"discovery"` + NodeResourceLimit *NodeResourceLimit `json:"node_resource_limit"` + LeaderElection SchedulerLeaderElectionConfig `json:"leader_election"` } func (s *SchedulerConfig) UnmarshalJSON(data []byte) error { @@ -79,11 +103,13 @@ func (s *SchedulerConfig) UnmarshalJSON(data []byte) error { ReportTTL json.RawMessage `json:"report_ttl"` BindingTTL json.RawMessage `json:"binding_ttl"` RedisAddr *string `json:"redis_addr"` + NodeAdminAPIKey *string `json:"node_admin_api_key"` ArtifactStoreCapacity *int `json:"artifact_store_capacity"` ArtifactLookupNodeLimit *int `json:"artifact_lookup_node_limit"` Nodes *[]Node `json:"nodes"` Discovery *SchedulerDiscoveryConfig `json:"discovery"` NodeResourceLimit *NodeResourceLimit `json:"node_resource_limit"` + LeaderElection *leaderElectionWire `json:"leader_election"` } parsed := wire{} @@ -112,6 +138,9 @@ func (s *SchedulerConfig) UnmarshalJSON(data []byte) error { if parsed.RedisAddr != nil { s.RedisAddr = *parsed.RedisAddr } + if parsed.NodeAdminAPIKey != nil { + s.NodeAdminAPIKey = *parsed.NodeAdminAPIKey + } if parsed.ArtifactStoreCapacity != nil { s.ArtifactStoreCapacity = *parsed.ArtifactStoreCapacity } @@ -134,9 +163,56 @@ func (s *SchedulerConfig) UnmarshalJSON(data []byte) error { s.BindingTTL = d } + if parsed.LeaderElection != nil { + le := parsed.LeaderElection + if le.Enabled != nil { + s.LeaderElection.Enabled = *le.Enabled + } + if le.LeaseName != nil { + s.LeaderElection.LeaseName = *le.LeaseName + } + if le.LeaseNamespace != nil { + s.LeaderElection.LeaseNamespace = *le.LeaseNamespace + } + if le.SnapshotPullConcurrency != nil { + s.LeaderElection.SnapshotPullConcurrency = *le.SnapshotPullConcurrency + } + durationFields := []struct { + raw json.RawMessage + field string + dst *time.Duration + }{ + {le.LeaseDuration, "scheduler.leader_election.lease_duration", &s.LeaderElection.LeaseDuration}, + {le.RenewDeadline, "scheduler.leader_election.renew_deadline", &s.LeaderElection.RenewDeadline}, + {le.RetryPeriod, "scheduler.leader_election.retry_period", &s.LeaderElection.RetryPeriod}, + } + for _, f := range durationFields { + if len(bytes.TrimSpace(f.raw)) == 0 { + continue + } + d, err := parseSchedulerDuration(f.raw, f.field) + if err != nil { + return err + } + *f.dst = d + } + } + return nil } +// leaderElectionWire is the JSON wire form of SchedulerLeaderElectionConfig; +// durations arrive as strings and are parsed via parseSchedulerDuration. +type leaderElectionWire struct { + Enabled *bool `json:"enabled"` + LeaseName *string `json:"lease_name"` + LeaseNamespace *string `json:"lease_namespace"` + LeaseDuration json.RawMessage `json:"lease_duration"` + RenewDeadline json.RawMessage `json:"renew_deadline"` + RetryPeriod json.RawMessage `json:"retry_period"` + SnapshotPullConcurrency *int `json:"snapshot_pull_concurrency"` +} + func parseSchedulerDuration(raw json.RawMessage, field string) (time.Duration, error) { var asString string if err := json.Unmarshal(raw, &asString); err == nil { @@ -320,6 +396,9 @@ func overrideWithEnv(cfg *Config) error { set("SCHEDULER_METRICS_LISTEN_ADDR", &cfg.Scheduler.MetricsListenAddr) set("SCHEDULER_STRATEGY", &cfg.Scheduler.Strategy) set("SCHEDULER_REDIS_ADDR", &cfg.Scheduler.RedisAddr) + // The node admin API key is a credential: pass it via env/Secret (same + // agentenv-auth secret the nodes use), never via config files. + set("SCHEDULER_NODE_ADMIN_API_KEY", &cfg.Scheduler.NodeAdminAPIKey) set("GATEWAY_HTTP_LISTEN_ADDR", &cfg.Gateway.HTTPListenAddr) set("GATEWAY_METRICS_LISTEN_ADDR", &cfg.Gateway.MetricsListenAddr) set("GATEWAY_SCHEDULER_ADDR", &cfg.Gateway.SchedulerAddr) @@ -403,6 +482,27 @@ func (c *Config) applyDefaults() { if strings.TrimSpace(c.Gateway.MetricsListenAddr) == "" { c.Gateway.MetricsListenAddr = ":9102" } + if c.Scheduler.LeaderElection.Enabled { + le := &c.Scheduler.LeaderElection + if strings.TrimSpace(le.LeaseName) == "" { + le.LeaseName = "agentenv-scheduler" + } + if strings.TrimSpace(le.LeaseNamespace) == "" { + le.LeaseNamespace = "agentenv-system" + } + if le.LeaseDuration <= 0 { + le.LeaseDuration = 15 * time.Second + } + if le.RenewDeadline <= 0 { + le.RenewDeadline = 10 * time.Second + } + if le.RetryPeriod <= 0 { + le.RetryPeriod = 2 * time.Second + } + if le.SnapshotPullConcurrency <= 0 { + le.SnapshotPullConcurrency = 4 + } + } } func (c Config) Validate() error { @@ -437,6 +537,9 @@ func (c Config) validate(schedulerQueryOnly bool) error { if c.Scheduler.BindingTTL <= 0 { return errors.New("scheduler.binding_ttl must be greater than zero") } + if c.Scheduler.LeaderElection.Enabled && schedulerQueryOnly { + return errors.New("scheduler.leader_election is mutually exclusive with --query-only") + } if schedulerQueryOnly { if strings.TrimSpace(c.Scheduler.RedisAddr) == "" { return errors.New("scheduler --query-only requires scheduler.redis_addr") @@ -473,6 +576,22 @@ func (c Config) validate(schedulerQueryOnly bool) error { default: return errors.New("scheduler.discovery.mode must be one of static, kubernetes") } + if c.Scheduler.LeaderElection.Enabled { + le := c.Scheduler.LeaderElection + if strings.TrimSpace(le.LeaseName) == "" { + return errors.New("scheduler.leader_election.lease_name is required") + } + if strings.TrimSpace(le.LeaseNamespace) == "" { + return errors.New("scheduler.leader_election.lease_namespace is required") + } + // client-go leaderelection timing constraints. + if le.LeaseDuration <= le.RenewDeadline { + return errors.New("scheduler.leader_election requires lease_duration > renew_deadline") + } + if le.RenewDeadline <= 2*le.RetryPeriod { + return errors.New("scheduler.leader_election requires renew_deadline > 2 * retry_period") + } + } } if c.Service == "gateway" { if c.Gateway.HTTPListenAddr == "" { diff --git a/services/shared/config/scheduler_leader_election_test.go b/services/shared/config/scheduler_leader_election_test.go new file mode 100644 index 000000000..ec34d5b30 --- /dev/null +++ b/services/shared/config/scheduler_leader_election_test.go @@ -0,0 +1,120 @@ +package config + +import ( + "strings" + "testing" + "time" +) + +// Leader-election config surface tests (#259, test list group A). + +func schedulerConfigWithLeaderElection(le SchedulerLeaderElectionConfig) Config { + return Config{ + Service: "scheduler", + LogLevel: "info", + LogFormat: "json", + Scheduler: SchedulerConfig{ + GRPCListenAddr: ":9090", + MetricsListenAddr: ":9091", + ReportTTL: 30 * time.Second, + BindingTTL: 30 * time.Second, + ArtifactStoreCapacity: 1, + Nodes: []Node{{ID: "n1", Endpoint: "http://n1:8080"}}, + Discovery: SchedulerDiscoveryConfig{Mode: "static"}, + LeaderElection: le, + }, + } +} + +func defaultLeaderElection() SchedulerLeaderElectionConfig { + return SchedulerLeaderElectionConfig{ + Enabled: true, + LeaseName: "agentenv-scheduler", + LeaseNamespace: "agentenv-system", + LeaseDuration: 15 * time.Second, + RenewDeadline: 10 * time.Second, + RetryPeriod: 2 * time.Second, + } +} + +// Case A1: leader election on without redis_addr still passes validation. +// Redis is recommended for failover routing, but not a hard requirement; +// the degraded mode (lost routing for pre-failover sandboxes) is documented. +func TestLeaderElectionWithoutRedisPassesValidation(t *testing.T) { + c := schedulerConfigWithLeaderElection(defaultLeaderElection()) + // Deliberately no RedisAddr. + if err := c.validate(false); err != nil { + t.Fatalf("leader election without redis_addr must pass validation, got %v", err) + } +} + +// Case A2: leader election is mutually exclusive with --query-only. +func TestLeaderElectionRejectsQueryOnly(t *testing.T) { + c := schedulerConfigWithLeaderElection(defaultLeaderElection()) + err := c.validate(true) + if err == nil { + t.Fatal("leader election with --query-only must fail validation") + } + if !strings.Contains(err.Error(), "query-only") { + t.Fatalf("error should mention the query-only conflict, got %v", err) + } +} + +// Case A3: invalid client-go timing triples are rejected. +// Rules: lease_duration > renew_deadline > 2 * retry_period. +func TestLeaderElectionRejectsInvalidTiming(t *testing.T) { + cases := map[string]SchedulerLeaderElectionConfig{ + "renew_deadline_not_below_lease_duration": { + Enabled: true, LeaseName: "l", LeaseNamespace: "ns", + LeaseDuration: 10 * time.Second, RenewDeadline: 15 * time.Second, RetryPeriod: 2 * time.Second, + }, + "retry_period_too_large": { + Enabled: true, LeaseName: "l", LeaseNamespace: "ns", + LeaseDuration: 15 * time.Second, RenewDeadline: 10 * time.Second, RetryPeriod: 6 * time.Second, + }, + "missing_lease_name": { + Enabled: true, LeaseNamespace: "ns", + LeaseDuration: 15 * time.Second, RenewDeadline: 10 * time.Second, RetryPeriod: 2 * time.Second, + }, + } + for name, le := range cases { + t.Run(name, func(t *testing.T) { + if err := schedulerConfigWithLeaderElection(le).validate(false); err == nil { + t.Fatalf("expected validation failure for %v", le) + } + }) + } +} + +// Case A4: leader election off (default) behaves exactly like today — +// validation ignores leader-election settings entirely. +func TestLeaderElectionDisabledIsTransparent(t *testing.T) { + c := schedulerConfigWithLeaderElection(SchedulerLeaderElectionConfig{ + Enabled: false, + LeaseDuration: 1 * time.Second, // invalid triple, must be ignored while disabled + RenewDeadline: 10 * time.Second, + }) + if err := c.validate(false); err != nil { + t.Fatalf("disabled leader election must not affect validation, got %v", err) + } + if err := c.validate(true); err == nil || !strings.Contains(err.Error(), "redis_addr") { + t.Fatalf("disabled leader election must keep existing query-only rules, got %v", err) + } +} + +// snapshot_pull_concurrency: defaults to 4 when enabled and unset; +// an explicit value is preserved. +func TestSnapshotPullConcurrencyDefault(t *testing.T) { + c := schedulerConfigWithLeaderElection(defaultLeaderElection()) + c.applyDefaults() + if got := c.Scheduler.LeaderElection.SnapshotPullConcurrency; got != 4 { + t.Fatalf("unset snapshot_pull_concurrency must default to 4, got %d", got) + } + + c = schedulerConfigWithLeaderElection(defaultLeaderElection()) + c.Scheduler.LeaderElection.SnapshotPullConcurrency = 16 + c.applyDefaults() + if got := c.Scheduler.LeaderElection.SnapshotPullConcurrency; got != 16 { + t.Fatalf("explicit snapshot_pull_concurrency must be preserved, got %d", got) + } +} From f92133c75efdd12e0f763e0fe374a6d2c6a55cf5 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 13:05:56 +0800 Subject: [PATCH 02/26] feat(scheduler): leadership domain with Lease election and recovery window MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the leadership domain for scheduler HA (#259), behind a single exported contract: - Leadership interface + NewLeadership(cfg) factory are the only exported surface. Two implementations, both package-private: leadershipManager (election enabled) and a null object (election disabled), so callers wire unconditionally and 'feature off' is an implementation difference, not a wiring difference. - leadershipSnapshot: immutable leadership state (leader?, acquired since, recovery TTL). leadershipManager owns transitions via atomic.Pointer whole-instance swaps; readers load the current snapshot per call — stable facade, swappable payload. - Read/write gate interceptor: health checks never gated; leader serves everything; standbys with a shared (Redis) binding store serve LookupNode/GetNode and reject the rest with Unavailable("not the leader"); without a shared store reads retry onto the leader (degraded mode). The RPC table classifies every method explicitly and panics on unclassified ones, so proto upgrades fail loudly instead of silently gating. - Leader readiness health service scheduler.v1.Scheduler/leader: SERVING only while holding the lease, so Service endpoints contain only the leader; liveness is unchanged. - client-go Lease elector with fencing: OnStoppedLeading flips readiness, then stops the process via the composition root's shutdown hook; a partitioned ex-leader stops serving within renew_deadline. - Recovery window in the Service: while open, Schedule only considers nodes observed at or after acquisition (fresh-observations-only, via the new NodeRegistry.LastReportAt) and returns Unavailable with zero fresh observations — the #191 fix, no snapshot means "unknown", not "unlimited"; LookupNode/GetNode return Unavailable (not NotFound) for missing entries so "not yet rebuilt" is not misreported as "does not exist". With election disabled all of this short-circuits to today's behaviour. - Heartbeat ingest is extracted into a shared path so pulled snapshots take the exact same code as heartbeats. Scope: services/scheduler/internal (leadership.go, service.go, node_registry.go) with unit and fake-clientset election tests (exactly-one-leader, failover after cancel, fencing on renew failure, callback ordering, recovery-window boundaries, gate matrix). Relates to #259, #191. Co-Authored-By: Claude Code --- services/scheduler/internal/leadership.go | 454 +++++++++++++++++ .../scheduler/internal/leadership_test.go | 466 ++++++++++++++++++ services/scheduler/internal/node_registry.go | 19 + .../internal/recovery_window_test.go | 226 +++++++++ services/scheduler/internal/service.go | 97 +++- 5 files changed, 1255 insertions(+), 7 deletions(-) create mode 100644 services/scheduler/internal/leadership.go create mode 100644 services/scheduler/internal/leadership_test.go create mode 100644 services/scheduler/internal/recovery_window_test.go diff --git a/services/scheduler/internal/leadership.go b/services/scheduler/internal/leadership.go new file mode 100644 index 000000000..9258e6503 --- /dev/null +++ b/services/scheduler/internal/leadership.go @@ -0,0 +1,454 @@ +package scheduler + +import ( + "context" + "fmt" + "os" + "strings" + "sync/atomic" + "time" + + schedulerv1 "agentenv/services/api/proto" + "agentenv/services/shared/config" + + "go.uber.org/zap" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/health" + "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/status" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/client-go/kubernetes" + "k8s.io/client-go/rest" + "k8s.io/client-go/tools/leaderelection" + "k8s.io/client-go/tools/leaderelection/resourcelock" +) + +// leadershipView is the narrow read-only seam the Service uses for +// recovery-window semantics (#259). It is intentionally unexported: outside +// this package, everything goes through the Leadership facade. +type leadershipView interface { + InRecoveryWindow(now time.Time) bool + LeaderSince() (time.Time, bool) +} + +// Leadership is the single external contract for the scheduler's +// leader-election behavior (#259). Callers (main, the Service) depend only +// on this interface and never on a concrete implementation. Two +// implementations exist: LeadershipManager (election enabled) and +// nonLeadership (a no-op for election disabled), chosen by NewLeadership. +type Leadership interface { + leadershipView + + // ServiceOption attaches the recovery-window view to a Service under + // construction. + ServiceOption() ServiceOption + // GateInterceptor returns the read/write gate for the gRPC chain; + // a pass-through when election is disabled. + GateInterceptor() grpc.UnaryServerInterceptor + // RegisterHealth registers the leader readiness service; no-op when + // election is disabled. + RegisterHealth(hs *health.Server) + // BindRuntime attaches the post-construction dependencies for + // sync-node-snapshots; no-op when election is disabled. + BindRuntime(svc *Service, registry NodeRegistry) + // Run drives leader election until ctx is done; returns immediately + // when election is disabled. + Run(ctx context.Context, onStop func()) error + // MarkNotServing flips leader readiness off; no-op when disabled. + MarkNotServing() +} + +// NewLeadership builds the Leadership implementation for the given config: +// LeadershipManager when leader election is enabled, a no-op otherwise. +// Callers never branch on the config themselves. +func NewLeadership(logger *zap.Logger, cfg config.SchedulerConfig) Leadership { + if !cfg.LeaderElection.Enabled { + return nonLeadership{} + } + return newLeadershipManager(logger, cfg) +} + +// nonLeadership is the null-object Leadership: election disabled. Every +// method is a no-op or a conservative default, so callers behave exactly as +// a single-writer scheduler without any conditional wiring. +type nonLeadership struct{} + +func (nonLeadership) InRecoveryWindow(time.Time) bool { return false } +func (nonLeadership) LeaderSince() (time.Time, bool) { return time.Time{}, false } +func (nonLeadership) ServiceOption() ServiceOption { return func(*Service) {} } +func (nonLeadership) RegisterHealth(*health.Server) {} +func (nonLeadership) BindRuntime(*Service, NodeRegistry) {} +func (nonLeadership) MarkNotServing() {} +func (nonLeadership) GateInterceptor() grpc.UnaryServerInterceptor { + return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { + return handler(ctx, req) + } +} +func (nonLeadership) Run(ctx context.Context, _ func()) error { + <-ctx.Done() + return nil +} + +// ============================================================ + +// leadershipSnapshot is an immutable snapshot of this replica's leadership state: +// whether it currently holds the leader lease and, if so, when it acquired +// it. The acquisition time drives the recovery window: right after winning, +// observations and bindings are still rebuilding, so scheduling must only +// consider nodes observed after that point (#259). +// +// The value object carries no lock and no mutators. Transitions are owned by +// LeadershipManager, which atomically swaps the current instance; readers +// load the current snapshot and compute against it. +type leadershipSnapshot struct { + leader bool + since time.Time + recoveryTTL time.Duration +} + +// newLeadership is the not-leading state. recoveryTTL bounds the recovery +// window after each acquisition; it is the scheduler report TTL: after it, +// every live node has either reported or been pulled, so a missing binding +// means "does not exist", not "not yet rebuilt". +func newLeadershipSnapshot(recoveryTTL time.Duration) *leadershipSnapshot { + return &leadershipSnapshot{recoveryTTL: recoveryTTL} +} + +// acquiredLeadership is the leading state, acquired at now. +func acquiredLeadershipSnapshot(recoveryTTL time.Duration, now time.Time) *leadershipSnapshot { + return &leadershipSnapshot{leader: true, since: now, recoveryTTL: recoveryTTL} +} + +// InRecoveryWindow reports whether the recovery window following leadership +// acquisition is still open at now. While open, LookupNode/GetNode must +// return Unavailable for missing entries instead of NotFound ("not yet +// rebuilt" is not "does not exist"). A replica that is not the leader is +// never in the window. +func (l *leadershipSnapshot) InRecoveryWindow(now time.Time) bool { + if !l.leader { + return false + } + return now.Before(l.since.Add(l.recoveryTTL)) +} + +func (l *leadershipSnapshot) IsLeader() bool { + return l.leader +} + +// LeaderSince reports when leadership was acquired. The second return value +// is false while this replica is not the leader. +func (l *leadershipSnapshot) LeaderSince() (time.Time, bool) { + return l.since, l.leader +} + +// ============================================================ + +// LeaderHealthService is the readiness health service name for leader +// election (#259): SERVING only while this replica holds the lease, so the +// K8S Service's endpoints contain only the leader. The deployment's +// readiness probe points here; liveness keeps reporting the overall "" +// status. Single definition point: deploy manifests reference the same +// string, and it is derived from the generated service descriptor so a proto +// package rename cannot drift. +var LeaderHealthService = schedulerv1.Scheduler_ServiceDesc.ServiceName + "/leader" + +// leadershipManager owns every leader-election policy the scheduler process +// applies at runtime: the shared Leadership state (the single source for +// "who leads and since when"), the leader-aware gate, the leader readiness +// health service, and the elector callbacks (readiness flip + +// sync-node-snapshots on acquire; readiness flip + stop hook on loss). +// +// Construction is two-phase to cut a construction cycle: the Service needs +// the leadership state early (ServiceOption), while the post-acquisition +// retrieval needs the Service and registry late (BindRuntime). Call order: +// +// NewLeadershipManager → ServiceOption / GateInterceptor → (build Service) +// → BindRuntime → RegisterHealth → Run +type leadershipManager struct { + logger *zap.Logger + cfg config.SchedulerConfig + state atomic.Pointer[leadershipSnapshot] + sharedBindingStore bool + health *health.Server + svc *Service + registry NodeRegistry +} + +// NewLeadershipManager builds the manager and its Leadership state. It must +// be created before the Service so the recovery window can be attached via +// ServiceOption. +func newLeadershipManager(logger *zap.Logger, cfg config.SchedulerConfig) *leadershipManager { + if logger == nil { + logger = zap.NewNop() + } + l := &leadershipManager{ + logger: logger, + cfg: cfg, + // Whether standbys may serve reads depends on the binding store: + // shared (Redis) → reads served; in-memory → reads rejected too. + sharedBindingStore: strings.TrimSpace(cfg.RedisAddr) != "", + } + l.state.Store(newLeadershipSnapshot(cfg.ReportTTL)) + logger.Info("scheduler leader election enabled", + zap.String("lease", cfg.LeaderElection.LeaseNamespace+"/"+cfg.LeaderElection.LeaseName), + ) + if strings.TrimSpace(cfg.NodeAdminAPIKey) == "" { + logger.Warn("sync-node-snapshots pulls will fail auth without scheduler.node_admin_api_key; nodes stay unobserved until heartbeats") + } + return l +} + +// ServiceOption attaches the recovery-window view to the Service under +// construction. The Service receives the manager itself as a narrow +// leadershipView — the internal leadership object is never handed out. +func (l *leadershipManager) ServiceOption() ServiceOption { + return WithLeadership(l) +} + +// InRecoveryWindow delegates to the internal leadership state: the only +// recovery-window read the Service needs (#259). +func (l *leadershipManager) InRecoveryWindow(now time.Time) bool { + return l.state.Load().InRecoveryWindow(now) +} + +// LeaderSince delegates to the internal leadership state. +func (l *leadershipManager) LeaderSince() (time.Time, bool) { + return l.state.Load().LeaderSince() +} + +// standbyReadableMethods classifies every Scheduler RPC: true = a non-leader +// replica sharing the binding store (Redis) may serve it; false = gated +// behind leadership. The table lists all methods explicitly so a proto +// change adding an RPC cannot silently fall into "not a read, reject": +// init panics on any unclassified method (fail fast at startup, #259). +// +// Notes on the reads: +// - LookupNode reads bindings from the shared store: correct on standbys. +// - GetNode reads replica-local observations, which are empty on standbys +// until heartbeats/pulls arrive; classified readable per the design +// matrix anyway (a miss degrades to NotFound, not to wrong data). +// - ListNodes reads the informer-backed registry (warm on every replica) +// and is semantically standby-serveable, but stays gated for now — +// deferred to a PR decision (#259 review). +var standbyReadableMethods = map[string]bool{ + "/scheduler.v1.Scheduler/Schedule": false, + "/scheduler.v1.Scheduler/ListNodes": false, + "/scheduler.v1.Scheduler/LookupNode": true, + "/scheduler.v1.Scheduler/RecordAssignment": false, + "/scheduler.v1.Scheduler/Heartbeat": false, + "/scheduler.v1.Scheduler/ReportSandboxEvent": false, + "/scheduler.v1.Scheduler/ListObservedNodes": false, + "/scheduler.v1.Scheduler/ListP2pPeers": false, + "/scheduler.v1.Scheduler/RecordP2pArtifact": false, + "/scheduler.v1.Scheduler/ForgetP2pArtifact": false, + "/scheduler.v1.Scheduler/LookupP2pArtifact": false, + "/scheduler.v1.Scheduler/GetNode": true, + "/scheduler.v1.Scheduler/UnregisterNode": false, +} + +func init() { + serviceName := schedulerv1.Scheduler_ServiceDesc.ServiceName + for _, method := range schedulerv1.Scheduler_ServiceDesc.Methods { + fullMethod := "/" + serviceName + "/" + method.MethodName + if _, ok := standbyReadableMethods[fullMethod]; !ok { + panic("scheduler: unclassified RPC in standbyReadableMethods: " + fullMethod) + } + } +} + +// GateInterceptor returns the read/write gate to chain into the gRPC server +// after the metrics interceptor. gRPC registration must happen before +// Serve, so gating is an interceptor, not late registration. Behaviour: +// +// - health-check RPCs are never gated (liveness must work on standbys); +// - the leader serves everything; +// - a standby with a shared binding store serves the read RPCs above from +// that store and rejects the rest with Unavailable("not the leader"); +// - a standby without a shared store (in-memory bindings) rejects reads +// too, and lookups reach the leader via client retry (degraded mode). +// +// When leader election is disabled the lifecycle (and this interceptor) is +// not installed at all. +func (l *leadershipManager) GateInterceptor() grpc.UnaryServerInterceptor { + return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { + if strings.HasPrefix(info.FullMethod, "/grpc.health.v1.") { + return handler(ctx, req) + } + if l.state.Load().IsLeader() { + return handler(ctx, req) + } + if l.sharedBindingStore && standbyReadableMethods[info.FullMethod] { + return handler(ctx, req) + } + return nil, status.Error(codes.Unavailable, "not the leader") + } +} + +// RegisterHealth registers the leader readiness service (NOT_SERVING until +// the lease is won) and keeps a handle for the acquire/loss/shutdown +// transitions. +func (l *leadershipManager) RegisterHealth(hs *health.Server) { + l.health = hs + hs.SetServingStatus(LeaderHealthService, grpc_health_v1.HealthCheckResponse_NOT_SERVING) +} + +// MarkNotServing flips the leader readiness service off; called on +// leadership loss (before process shutdown) and during graceful shutdown. +func (l *leadershipManager) MarkNotServing() { + if l.health != nil { + l.health.SetServingStatus(LeaderHealthService, grpc_health_v1.HealthCheckResponse_NOT_SERVING) + } +} + +// BindRuntime attaches the service-side dependencies needed by the +// sync-node-snapshots retrieval on acquire. Must be called after Service +// construction and before Run. The registry is the NodeRegistry interface: +// the manager only lists nodes through it, never the concrete type. +func (l *leadershipManager) BindRuntime(svc *Service, registry NodeRegistry) { + l.svc = svc + l.registry = registry +} + +// Run drives leader election until ctx is done or the lease is lost; it +// blocks. onStop is the composition root's shutdown hook (e.g. canceling the +// process root context); it runs after the readiness flip so a partitioned +// ex-leader stops serving before exiting (fencing, #259). +func (l *leadershipManager) Run(ctx context.Context, onStop func()) error { + restCfg, err := rest.InClusterConfig() + if err != nil { + return fmt.Errorf("leader election requires in-cluster kubernetes access: %w", err) + } + clientset, err := kubernetes.NewForConfig(restCfg) + if err != nil { + return fmt.Errorf("leader election kubernetes client failed: %w", err) + } + hostname, err := os.Hostname() + if err != nil { + return fmt.Errorf("leader election identity failed: %w", err) + } + + elector := newLeaderElector(l.logger, l.cfg.LeaderElection, + clientset, hostname, + l.onStartedLeading, + func() { + l.state.Store(newLeadershipSnapshot(l.cfg.ReportTTL)) + l.MarkNotServing() + if onStop != nil { + onStop() + } + }, + ) + return elector.Run(ctx) +} + +// onStartedLeading is the elector's OnStartedLeading hook: readiness +// SERVING, then kick off sync-node-snapshots — refresh observations and +// bindings by pulling every node's admin snapshot instead of waiting for the +// next reporter backoff (#259). +func (l *leadershipManager) onStartedLeading(ctx context.Context) { + l.state.Store(acquiredLeadershipSnapshot(l.cfg.ReportTTL, time.Now())) + if l.health != nil { + l.health.SetServingStatus(LeaderHealthService, grpc_health_v1.HealthCheckResponse_SERVING) + } + if l.svc == nil || l.registry == nil { + l.logger.Error("leader lifecycle run before BindRuntime; skipping sync-node-snapshots") + return + } + ingestor := func(report *nodeReport, now time.Time) error { + _, err := l.svc.ingestNodeReport(report, now) + return err + } + var refresher nodeSnapshotRefresher = NewConcurrentNodeSnapshotRefresher(l.logger, ingestor, + NewAdminSnapshotFetcher(l.cfg.NodeAdminAPIKey), + l.cfg.LeaderElection.SnapshotPullConcurrency) + go refresher.Refresh(ctx, l.registry.Snapshot(false)) +} + +// ============================================================ + +// LeaderElector runs client-go leader election over a +// coordination.k8s.io/Lease (#259). Exactly one replica holds the lease at a +// time; the others wait. Election itself never touches Redis. +// +// The elector holds no leadership state; it forwards client-go callbacks to +// the caller (LeadershipManager), which owns the state transitions. +type leaderElector struct { + logger *zap.Logger + cfg config.SchedulerLeaderElectionConfig + clientset kubernetes.Interface + identity string + onStarted func(ctx context.Context) + onStopped func() +} + +// NewLeaderElector builds a runner. The clientset and identity are +// injected: production passes the in-cluster client and the pod hostname; +// tests pass a fake clientset and a stable identity. +func newLeaderElector( + logger *zap.Logger, + cfg config.SchedulerLeaderElectionConfig, + clientset kubernetes.Interface, + identity string, + onStarted func(ctx context.Context), + onStopped func(), +) *leaderElector { + if logger == nil { + logger = zap.NewNop() + } + return &leaderElector{ + logger: logger, + cfg: cfg, + clientset: clientset, + identity: identity, + onStarted: onStarted, + onStopped: onStopped, + } +} + +func (r *leaderElector) Run(ctx context.Context) error { + lock := &resourcelock.LeaseLock{ + LeaseMeta: metav1.ObjectMeta{ + Name: r.cfg.LeaseName, + Namespace: r.cfg.LeaseNamespace, + }, + Client: r.clientset.CoordinationV1(), + LockConfig: resourcelock.ResourceLockConfig{ + Identity: r.identity, + }, + } + + elector, err := leaderelection.NewLeaderElector(leaderelection.LeaderElectionConfig{ + Lock: lock, + LeaseDuration: r.cfg.LeaseDuration, + RenewDeadline: r.cfg.RenewDeadline, + RetryPeriod: r.cfg.RetryPeriod, + Callbacks: leaderelection.LeaderCallbacks{ + OnStartedLeading: func(leaderCtx context.Context) { + r.logger.Info("scheduler acquired leadership", + zap.String("lease", r.cfg.LeaseNamespace+"/"+r.cfg.LeaseName), + ) + if r.onStarted != nil { + r.onStarted(leaderCtx) + } + }, + OnStoppedLeading: func() { + r.logger.Warn("scheduler lost leadership; shutting down for fencing") + if r.onStopped != nil { + r.onStopped() + } + }, + OnNewLeader: func(identity string) { + r.logger.Info("scheduler leader changed", zap.String("leader", identity)) + }, + }, + // Release the lease on graceful shutdown so failover is fast. + ReleaseOnCancel: true, + }) + if err != nil { + return fmt.Errorf("leader election config: %w", err) + } + + elector.Run(ctx) + return nil +} diff --git a/services/scheduler/internal/leadership_test.go b/services/scheduler/internal/leadership_test.go new file mode 100644 index 000000000..f14858b96 --- /dev/null +++ b/services/scheduler/internal/leadership_test.go @@ -0,0 +1,466 @@ +package scheduler + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "agentenv/services/shared/config" + + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/health" + "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/status" + coordinationv1 "k8s.io/api/coordination/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/client-go/kubernetes/fake" + ktesting "k8s.io/client-go/testing" +) + +// ============================================================ + +// Leadership contract (#259): the immutable state values and the +// recovery-window boundaries that lookup/schedule semantics read. +// Transitions are owned by LeadershipManager and covered in its tests. + +func TestLeadershipStates(t *testing.T) { + standby := newLeadershipSnapshot(30 * time.Second) + if standby.IsLeader() { + t.Fatal("a fresh replica must not be the leader") + } + if _, ok := standby.LeaderSince(); ok { + t.Fatal("LeaderSince must be false before acquisition") + } + + t0 := time.Date(2026, 9, 29, 12, 0, 0, 0, time.UTC) + leader := acquiredLeadershipSnapshot(30*time.Second, t0) + if !leader.IsLeader() { + t.Fatal("acquiredLeadership must be the leader") + } + since, ok := leader.LeaderSince() + if !ok || !since.Equal(t0) { + t.Fatalf("LeaderSince after acquisition = (%v, %v); want (%v, true)", since, ok, t0) + } +} + +func TestLeadershipRecoveryWindowBoundaries(t *testing.T) { + const ttl = 30 * time.Second + + if newLeadershipSnapshot(ttl).InRecoveryWindow(time.Now()) { + t.Fatal("a replica that never acquired leadership is never in the recovery window") + } + + t0 := time.Date(2026, 9, 29, 12, 0, 0, 0, time.UTC) + l := acquiredLeadershipSnapshot(ttl, t0) + + if !l.InRecoveryWindow(t0.Add(time.Second)) { + t.Fatal("the window must be open right after acquisition") + } + if !l.InRecoveryWindow(t0.Add(ttl - time.Nanosecond)) { + t.Fatal("the window must stay open for the whole TTL") + } + if l.InRecoveryWindow(t0.Add(ttl)) { + t.Fatal("the window must close exactly at the TTL boundary") + } + if l.InRecoveryWindow(t0.Add(ttl + time.Minute)) { + t.Fatal("the window must stay closed after the TTL") + } + + if newLeadershipSnapshot(ttl).InRecoveryWindow(t0.Add(time.Second)) { + t.Fatal("a released replica is never in the recovery window") + } +} + +// ============================================================ + +func newTestManager() *leadershipManager { + return newLeadershipManager(nil, config.SchedulerConfig{ + ReportTTL: 30 * time.Second, + RedisAddr: "localhost:6379", // shared store → standbys serve reads + LeaderElection: config.SchedulerLeaderElectionConfig{Enabled: true}, + }) +} + +// The leader readiness service is NOT_SERVING from registration, flips +// SERVING on acquisition, and flips back on loss/shutdown (#259). +func TestLeadershipManagerHealthTransitions(t *testing.T) { + lifecycle := newTestManager() + hs := health.NewServer() + lifecycle.RegisterHealth(hs) + + check := func() grpc_health_v1.HealthCheckResponse_ServingStatus { + resp, err := hs.Check(context.Background(), &grpc_health_v1.HealthCheckRequest{Service: LeaderHealthService}) + if err != nil { + t.Fatalf("leader health check failed: %v", err) + } + return resp.GetStatus() + } + + if got := check(); got != grpc_health_v1.HealthCheckResponse_NOT_SERVING { + t.Fatalf("initial leader health = %v; want NOT_SERVING until the lease is won", got) + } + + lifecycle.onStartedLeading(context.Background()) + if got := check(); got != grpc_health_v1.HealthCheckResponse_SERVING { + t.Fatalf("leader health after acquisition = %v; want SERVING", got) + } + + lifecycle.MarkNotServing() + if got := check(); got != grpc_health_v1.HealthCheckResponse_NOT_SERVING { + t.Fatalf("leader health after MarkNotServing = %v; want NOT_SERVING", got) + } +} + +// Construction-order guard: onStartedLeading before BindRuntime must not +// panic or dereference the unbound service/registry; it logs and skips the +// sync-node-snapshots retrieval. +func TestLeadershipManagerAcquireBeforeBindRuntimeIsSafe(t *testing.T) { + lifecycle := newTestManager() + lifecycle.onStartedLeading(context.Background()) +} + +// ServiceOption attaches the lifecycle's Leadership to the Service, enabling +// the recovery-window rules only when election is on. +func TestLeadershipManagerServiceOptionAttachesLeadership(t *testing.T) { + lifecycle := newTestManager() + svc := NewService(nil, nil, NewRandomStrategy(), NewInMemoryBindingStore(0), + lifecycle.ServiceOption(), + ) + if svc.leadership == nil { + t.Fatal("ServiceOption must attach the lifecycle's Leadership to the Service") + } +} + +// Leader gate tests (#259, test list group B, unit cases B1-B4). + +func invokeGate(t *testing.T, isLeader, serveReads bool, method string) (bool, error) { + t.Helper() + l := &leadershipManager{ + sharedBindingStore: serveReads, + } + l.state.Store(newLeadershipSnapshot(30 * time.Second)) + if isLeader { + l.state.Store(acquiredLeadershipSnapshot(30*time.Second, time.Now())) + } + interceptor := l.GateInterceptor() + handlerCalled := false + handler := func(ctx context.Context, req any) (any, error) { + handlerCalled = true + return "served", nil + } + _, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{FullMethod: method}, handler) + return handlerCalled, err +} + +// Case B1: a standby without a shared store rejects writes. +func TestGateStandbyWithoutStoreRejectsWrites(t *testing.T) { + called, err := invokeGate(t, false, false, "/scheduler.v1.Scheduler/Schedule") + if called { + t.Fatal("write handler must not run on a standby") + } + if status.Code(err) != codes.Unavailable { + t.Fatalf("expected Unavailable, got %v", err) + } + if status.Convert(err).Message() != "not the leader" { + t.Fatalf("unexpected message: %v", err) + } +} + +// Case B2: a standby without a shared store rejects reads too +// (degraded mode: reads retry until they land on the leader). +func TestGateStandbyWithoutStoreRejectsReads(t *testing.T) { + called, err := invokeGate(t, false, false, "/scheduler.v1.Scheduler/LookupNode") + if called { + t.Fatal("read handler must not run on a standby without a shared store") + } + if status.Code(err) != codes.Unavailable { + t.Fatalf("expected Unavailable, got %v", err) + } +} + +// Case B3: a standby with a shared Redis store serves reads from it, +// while writes stay gated. +func TestGateStandbyWithStoreServesReads(t *testing.T) { + called, err := invokeGate(t, false, true, "/scheduler.v1.Scheduler/LookupNode") + if !called || err != nil { + t.Fatalf("standby with a shared store must serve LookupNode: called=%v err=%v", called, err) + } + + called, err = invokeGate(t, false, true, "/scheduler.v1.Scheduler/GetNode") + if !called || err != nil { + t.Fatalf("standby with a shared store must serve GetNode: called=%v err=%v", called, err) + } + + called, err = invokeGate(t, false, true, "/scheduler.v1.Scheduler/RecordAssignment") + if called || status.Code(err) != codes.Unavailable { + t.Fatalf("writes must stay gated even with a shared store: called=%v err=%v", called, err) + } + + called, err = invokeGate(t, false, true, "/scheduler.v1.Scheduler/Heartbeat") + if called || status.Code(err) != codes.Unavailable { + t.Fatalf("heartbeats must stay gated even with a shared store: called=%v err=%v", called, err) + } +} + +// Case B4: health-check RPCs bypass the gate (else probes deadlock). +func TestGateHealthChecksBypass(t *testing.T) { + for _, method := range []string{ + "/grpc.health.v1.Health/Check", + "/grpc.health.v1.Health/Watch", + } { + called, err := invokeGate(t, false, false, method) + if !called || err != nil { + t.Fatalf("health check %s must bypass the gate: called=%v err=%v", method, called, err) + } + } +} + +// The leader serves everything, including with no shared store. +func TestGateLeaderServesEverything(t *testing.T) { + for _, method := range []string{ + "/scheduler.v1.Scheduler/Schedule", + "/scheduler.v1.Scheduler/LookupNode", + "/scheduler.v1.Scheduler/Heartbeat", + } { + called, err := invokeGate(t, true, false, method) + if !called || err != nil { + t.Fatalf("leader must serve %s: called=%v err=%v", method, called, err) + } + } +} + +// ============================================================ + +// Leader-election behavior tests against a fake clientset (#259, group C). +// The fake API server handles Lease CRUD and renewal failures, so these run +// offline in seconds without envtest/Kind. A real-apiserver partition test +// stays in Kind/CI. + +var electionTiming = config.SchedulerLeaderElectionConfig{ + Enabled: true, + LeaseName: "agentenv-scheduler", + LeaseNamespace: "agentenv-system", + LeaseDuration: 2 * time.Second, + RenewDeadline: 1500 * time.Millisecond, + RetryPeriod: 300 * time.Millisecond, +} + +type runnerHandle struct { + state atomic.Pointer[leadershipSnapshot] + started atomic.Int32 + stopped atomic.Int32 + cancel context.CancelFunc +} + +func startRunner(t *testing.T, clientset *fake.Clientset, identity string) *runnerHandle { + t.Helper() + h := &runnerHandle{} + h.state.Store(newLeadershipSnapshot(30 * time.Second)) + ctx, cancel := context.WithCancel(context.Background()) + h.cancel = cancel + runner := newLeaderElector(nil, electionTiming, clientset, identity, + func(context.Context) { + h.state.Store(acquiredLeadershipSnapshot(30*time.Second, time.Now())) + h.started.Add(1) + }, + func() { + h.state.Store(newLeadershipSnapshot(30 * time.Second)) + h.stopped.Add(1) + }, + ) + go func() { + if err := runner.Run(ctx); err != nil { + t.Errorf("runner %s failed: %v", identity, err) + } + }() + t.Cleanup(cancel) + return h +} + +func eventually(t *testing.T, timeout time.Duration, what string, cond func() bool) { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if cond() { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Fatalf("timed out waiting for %s", what) +} + +// Case C1: exactly one leader among two replicas competing for one Lease. +func TestElectionExactlyOneLeader(t *testing.T) { + clientset := fake.NewSimpleClientset() + a := startRunner(t, clientset, "replica-a") + b := startRunner(t, clientset, "replica-b") + + eventually(t, 5*time.Second, "one leader", func() bool { + return a.state.Load().IsLeader() != b.state.Load().IsLeader() + }) + eventually(t, 2*time.Second, "leader remains unique", func() bool { + return a.state.Load().IsLeader() != b.state.Load().IsLeader() + }) +} + +// Case C2: graceful leader loss (context cancel) — the standby takes over +// within the lease + retry budget, and the ex-leader's fencing callback runs. +func TestElectionFailoverAfterCancel(t *testing.T) { + clientset := fake.NewSimpleClientset() + a := startRunner(t, clientset, "replica-a") + b := startRunner(t, clientset, "replica-b") + + eventually(t, 5*time.Second, "one leader", func() bool { + return a.state.Load().IsLeader() != b.state.Load().IsLeader() + }) + + leader, standby := a, b + if b.state.Load().IsLeader() { + leader, standby = b, a + } + + leader.cancel() + + eventually(t, 5*time.Second, "ex-leader fencing callback", func() bool { + return leader.stopped.Load() >= 1 && !leader.state.Load().IsLeader() + }) + eventually(t, 5*time.Second, "standby takeover", func() bool { + return standby.state.Load().IsLeader() && standby.started.Load() >= 1 + }) +} + +// Case C3: fencing under renewal failure — when the leader can no longer +// renew the Lease (partition equivalent), it must call OnStoppedLeading +// within renew_deadline and the standby must take over. No dual-primary +// window: the standby only starts after the Lease actually expires. +func TestElectionFencingOnRenewFailure(t *testing.T) { + clientset := fake.NewSimpleClientset() + a := startRunner(t, clientset, "replica-a") + b := startRunner(t, clientset, "replica-b") + + eventually(t, 5*time.Second, "one leader", func() bool { + return a.state.Load().IsLeader() != b.state.Load().IsLeader() + }) + + leader, standby := a, b + leaderID := "replica-a" + if b.state.Load().IsLeader() { + leader, standby = b, a + leaderID = "replica-b" + } + + // Simulate a partition: Lease updates by the current holder now fail. + clientset.PrependReactor("update", "leases", func(action ktesting.Action) (bool, runtime.Object, error) { + lease, ok := action.(ktesting.UpdateAction).GetObject().(*coordinationv1.Lease) + if !ok || lease.Spec.HolderIdentity == nil || *lease.Spec.HolderIdentity != leaderID { + return false, nil, nil + } + return true, nil, errors.New("apiserver unreachable (simulated partition)") + }) + + eventually(t, 5*time.Second, "ex-leader stops serving within renew_deadline", func() bool { + return leader.stopped.Load() >= 1 && !leader.state.Load().IsLeader() + }) + eventually(t, 5*time.Second, "standby takes over after expiry", func() bool { + return standby.state.Load().IsLeader() + }) +} + +// Case C3-bis: callback ordering — the ex-leader's OnStoppedLeading must fire +// before the standby's OnStartedLeading. This is the callback-level proof of +// the no-dual-primary window. +func TestElectionCallbackOrderNoDualPrimary(t *testing.T) { + type event struct { + kind string + id string + } + var mu sync.Mutex + var seq []event + record := func(kind, id string) { + mu.Lock() + seq = append(seq, event{kind: kind, id: id}) + mu.Unlock() + } + + clientset := fake.NewSimpleClientset() + start := func(id string) *atomic.Pointer[leadershipSnapshot] { + state := &atomic.Pointer[leadershipSnapshot]{} + state.Store(newLeadershipSnapshot(30 * time.Second)) + ctx, cancel := context.WithCancel(context.Background()) + runner := newLeaderElector(nil, electionTiming, clientset, id, + func(context.Context) { + state.Store(acquiredLeadershipSnapshot(30*time.Second, time.Now())) + record("start", id) + }, + func() { + state.Store(newLeadershipSnapshot(30 * time.Second)) + record("stop", id) + }, + ) + go func() { _ = runner.Run(ctx) }() + t.Cleanup(cancel) + return state + } + + a := start("replica-a") + b := start("replica-b") + + eventually(t, 5*time.Second, "one leader", func() bool { + return a.Load().IsLeader() != b.Load().IsLeader() + }) + + mu.Lock() + var firstStart, lastStop, takeOver int = -1, -1, -1 + for i, e := range seq { + switch { + case e.kind == "start" && firstStart == -1: + firstStart = i + case e.kind == "stop": + lastStop = i + case e.kind == "start" && i > firstStart: + takeOver = i + } + } + mu.Unlock() + + // Kill the leader (cancel via cleanup happens later; simulate loss by + // renewing failure is C3's job — here we only assert ordering of the + // events seen so far and after a forced loss). + // Force a loss: prepend the same renew-failure reactor as C3. + leaderID := "replica-a" + if b.Load().IsLeader() { + leaderID = "replica-b" + } + clientset.PrependReactor("update", "leases", func(action ktesting.Action) (bool, runtime.Object, error) { + lease, ok := action.(ktesting.UpdateAction).GetObject().(*coordinationv1.Lease) + if !ok || lease.Spec.HolderIdentity == nil || *lease.Spec.HolderIdentity != leaderID { + return false, nil, nil + } + return true, nil, errors.New("apiserver unreachable (simulated partition)") + }) + + eventually(t, 5*time.Second, "stop then takeover ordering", func() bool { + mu.Lock() + defer mu.Unlock() + lastStop, takeOver = -1, -1 + for i, e := range seq { + if e.kind == "stop" { + lastStop = i + } + if e.kind == "start" && i != firstStart { + takeOver = i + } + } + return lastStop != -1 && takeOver != -1 + }) + + mu.Lock() + defer mu.Unlock() + if lastStop >= takeOver { + t.Fatalf("dual-primary window: OnStoppedLeading (idx %d) must precede takeover OnStartedLeading (idx %d); seq=%v", lastStop, takeOver, seq) + } +} diff --git a/services/scheduler/internal/node_registry.go b/services/scheduler/internal/node_registry.go index 5596d2980..6054349d5 100644 --- a/services/scheduler/internal/node_registry.go +++ b/services/scheduler/internal/node_registry.go @@ -26,6 +26,11 @@ type NodeRegistry interface { // and returns only the raw snapshot suitable for scheduling decisions. // Returns nil if the node has never sent a heartbeat. PeekObserved(nodeID string) *schedulerv1.NodeSnapshot + // LastReportAt returns when the given node last reported (heartbeat or + // pulled snapshot). The second return value is false if the node has + // never reported. Used by the leader-election recovery window to tell + // fresh observations from pre-acquisition ones (#259). + LastReportAt(nodeID string) (time.Time, bool) UnregisterObserved(nodeID string, serviceInstanceID string) error } @@ -345,6 +350,20 @@ func (r *AtomicNodeRegistry) PeekObserved(nodeID string) *schedulerv1.NodeSnapsh return cloneSnapshot(snapshot) } +func (r *AtomicNodeRegistry) LastReportAt(nodeID string) (time.Time, bool) { + r.mu.RLock() + defer r.mu.RUnlock() + record, ok := r.observed[nodeID] + if !ok || record.node == nil { + return time.Time{}, false + } + ms := record.node.GetLastSeenUnixMs() + if ms <= 0 { + return time.Time{}, false + } + return time.UnixMilli(ms), true +} + func (r *AtomicNodeRegistry) UnregisterObserved(nodeID string, serviceInstanceID string) error { r.mu.Lock() defer r.mu.Unlock() diff --git a/services/scheduler/internal/recovery_window_test.go b/services/scheduler/internal/recovery_window_test.go new file mode 100644 index 000000000..f8371f0ba --- /dev/null +++ b/services/scheduler/internal/recovery_window_test.go @@ -0,0 +1,226 @@ +package scheduler + +import ( + "context" + "errors" + "testing" + "time" + + schedulerv1 "agentenv/services/api/proto" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// Recovery-window and sync-node-snapshots tests (#259, test list group D). +// +// These tests pin the contract ahead of the Phase 3 business code; the ones +// covering still-unimplemented rules are expected to fail (red) until it +// lands. + +func newTestService(t *testing.T, nodeIDs []string) (*Service, *AtomicNodeRegistry, *InMemoryBindingStore) { + t.Helper() + nodes := make([]Node, 0, len(nodeIDs)) + for _, id := range nodeIDs { + nodes = append(nodes, Node{ID: id, Endpoint: "http://" + id + ":8080"}) + } + registry := NewAtomicNodeRegistry(nodes, 30*time.Second) + store := NewInMemoryBindingStore(30 * time.Second) + svc := NewService(nil, registry, NewRandomStrategy(), store) + return svc, registry, store +} + +func heartbeatFor(nodeID, instanceID string, sandboxIDs ...string) *schedulerv1.HeartbeatRequest { + return &schedulerv1.HeartbeatRequest{ + NodeId: nodeID, + ServiceInstanceId: instanceID, + SandboxIds: sandboxIDs, + } +} + +func reportFor(nodeID, instanceID string, sandboxIDs ...string) *nodeReport { + return nodeReportFromProto(heartbeatFor(nodeID, instanceID, sandboxIDs...)) +} + +// fixedClock is a controllable time source for window tests. +type fixedClock struct{ now time.Time } + +func (c *fixedClock) Now() time.Time { return c.now } + +// newLeaderTestService builds a service with leadership attached (recovery +// window = recoveryTTL) and a fixed clock the caller can advance. +func newLeaderTestService(t *testing.T, nodeIDs []string, recoveryTTL time.Duration, at, acquireAt time.Time) (*Service, *AtomicNodeRegistry, *fixedClock) { + t.Helper() + nodes := make([]Node, 0, len(nodeIDs)) + for _, id := range nodeIDs { + nodes = append(nodes, Node{ID: id, Endpoint: "http://" + id + ":8080"}) + } + registry := NewAtomicNodeRegistry(nodes, 30*time.Second) + store := NewInMemoryBindingStore(30 * time.Second) + clock := &fixedClock{now: at} + svc := NewService(nil, registry, NewStrategy("round_robin"), store, + WithLeadership(acquiredLeadershipSnapshot(recoveryTTL, acquireAt)), + WithClock(clock.Now), + ) + return svc, registry, clock +} + +// Case D3: the pulled-snapshot ingest path is the same code path as Heartbeat. +// Sync-node-snapshots feeds admin /nodes responses through ingestNodeReport; +// this test pins the contract: both entry points update observations and +// reconcile bindings identically. +func TestIngestNodeReportSharedWithHeartbeat(t *testing.T) { + svc, registry, store := newTestService(t, []string{"node-a"}) + + // Entry point 1: the Heartbeat RPC. + if _, err := svc.Heartbeat(context.Background(), heartbeatFor("node-a", "inst-1", "sbx-1")); err != nil { + t.Fatalf("heartbeat failed: %v", err) + } + if snap := registry.PeekObserved("node-a"); snap == nil { + t.Fatal("heartbeat must record an observation") + } + node, ok, err := store.Get("sbx-1", time.Now()) + if err != nil || !ok || node.ID != "node-a" { + t.Fatalf("heartbeat must reconcile bindings: ok=%v node=%v err=%v", ok, node, err) + } + + // Entry point 2: direct ingest (what sync-node-snapshots calls with a + // pulled admin /nodes snapshot). Same observable effects. + if _, err := svc.ingestNodeReport(reportFor("node-a", "inst-1", "sbx-2"), time.Now()); err != nil { + t.Fatalf("direct ingest failed: %v", err) + } + if snap := registry.PeekObserved("node-a"); snap == nil { + t.Fatal("direct ingest must record an observation") + } + node, ok, err = store.Get("sbx-2", time.Now()) + if err != nil || !ok || node.ID != "node-a" { + t.Fatalf("direct ingest must reconcile bindings: ok=%v node=%v err=%v", ok, node, err) + } + + // Both entry points also stamp the same freshness marker the recovery + // window reads. + if _, ok := registry.LastReportAt("node-a"); !ok { + t.Fatal("ingest must stamp LastReportAt for the recovery window") + } +} + +// Case D1: fresh-observations-only scheduling inside the recovery window. +// +// node-b's only observation predates acquisition (stale); node-a reported +// after (fresh). Round-robin over both candidates would alternate +// node-a → node-b on consecutive calls; the fresh-observations-only filter must pin every +// pick to node-a. +func TestScheduleFreshOnlyDuringRecoveryWindow(t *testing.T) { + t0 := time.Date(2026, 9, 29, 12, 0, 0, 0, time.UTC) + svc, _, _ := newLeaderTestService(t, []string{"node-a", "node-b"}, 30*time.Second, t0.Add(2*time.Second), t0) + + // Stale observation: reported before acquisition. + if _, err := svc.ingestNodeReport(reportFor("node-b", "inst-b"), t0.Add(-10*time.Second)); err != nil { + t.Fatalf("stale ingest failed: %v", err) + } + // Fresh observation: reported after acquisition. + if _, err := svc.ingestNodeReport(reportFor("node-a", "inst-a"), t0.Add(time.Second)); err != nil { + t.Fatalf("fresh ingest failed: %v", err) + } + + for i := 0; i < 2; i++ { + resp, err := svc.Schedule(context.Background(), &schedulerv1.ScheduleRequest{}) + if err != nil { + t.Fatalf("schedule %d failed: %v", i, err) + } + if got := resp.GetNode().GetNodeId(); got != "node-a" { + t.Fatalf("schedule %d picked %q; want node-a — node-b's observation predates acquisition and must not be a candidate", i, got) + } + } +} + +// Case D2: zero fresh observations — Schedule must return Unavailable and +// never fall back to unobserved nodes (#191: no snapshot means "unknown", +// not "unlimited"). +func TestScheduleZeroFreshObservationsReturnsUnavailable(t *testing.T) { + t0 := time.Date(2026, 9, 29, 12, 0, 0, 0, time.UTC) + svc, _, _ := newLeaderTestService(t, []string{"node-a"}, 30*time.Second, t0.Add(time.Second), t0) + + // The only observation predates acquisition: nothing is fresh. + if _, err := svc.ingestNodeReport(reportFor("node-a", "inst-a"), t0.Add(-10*time.Second)); err != nil { + t.Fatalf("stale ingest failed: %v", err) + } + + _, err := svc.Schedule(context.Background(), &schedulerv1.ScheduleRequest{}) + if status.Code(err) != codes.Unavailable { + t.Fatalf("want Unavailable with zero fresh observations, got %v", err) + } +} + +// Case D4: a node whose rebuild pull fails stays unobserved (and is kept out +// of scheduling by the fresh-observations-only rule) until its next report; successful +// pulls ingest through the shared path. +func TestPullFailureKeepsNodeUnobserved(t *testing.T) { + t0 := time.Date(2026, 9, 29, 12, 0, 0, 0, time.UTC) + svc, registry, _ := newLeaderTestService(t, []string{"node-a", "node-b"}, 30*time.Second, t0.Add(2*time.Second), t0) + + fetchCalls := 0 + fetch := func(_ context.Context, node Node) (*nodeReport, error) { + fetchCalls++ + if node.ID == "node-b" { + return nil, errors.New("admin endpoint unreachable") + } + return reportFor("node-a", "inst-a"), nil + } + ingest := func(report *nodeReport, now time.Time) error { + _, err := svc.ingestNodeReport(report, now) + return err + } + refresher := NewConcurrentNodeSnapshotRefresher(nil, ingest, fetch, 2) + refresher.Refresh(context.Background(), registry.Snapshot(false)) + + if fetchCalls != 2 { + t.Fatalf("rebuild must pull every node exactly once; got %d fetches for 2 nodes", fetchCalls) + } + if _, ok := registry.LastReportAt("node-a"); !ok { + t.Fatal("a successful pull must observe node-a") + } + if _, ok := registry.LastReportAt("node-b"); ok { + t.Fatal("a failed pull must leave node-b unobserved") + } + + // Fresh-only scheduling: only node-a is a candidate. + resp, err := svc.Schedule(context.Background(), &schedulerv1.ScheduleRequest{}) + if err != nil { + t.Fatalf("schedule failed: %v", err) + } + if got := resp.GetNode().GetNodeId(); got != "node-a" { + t.Fatalf("schedule picked %q; want node-a — node-b's pull failed and it must stay unobserved", got) + } +} + +// Case D5: LookupNode/GetNode semantics during vs after the recovery window. +// A missing binding/observation is "not yet rebuilt" (Unavailable) while the +// window is open, and "does not exist" (NotFound) once it closes. +func TestLookupNodeWindowSemantics(t *testing.T) { + t0 := time.Date(2026, 9, 29, 12, 0, 0, 0, time.UTC) + svc, _, clock := newLeaderTestService(t, []string{"node-a"}, 30*time.Second, t0.Add(time.Second), t0) + + // Window open: missing entries are "not yet rebuilt". + clock.now = t0.Add(time.Second) + _, err := svc.LookupNode(context.Background(), &schedulerv1.LookupNodeRequest{SandboxId: "sbx-ghost"}) + if status.Code(err) != codes.Unavailable { + t.Fatalf("window open: LookupNode want Unavailable, got %v", err) + } + _, err = svc.GetNode(context.Background(), &schedulerv1.GetNodeRequest{NodeId: "node-ghost"}) + if status.Code(err) != codes.Unavailable { + t.Fatalf("window open: GetNode want Unavailable, got %v", err) + } + + // Window closed (acquired more than the recovery TTL ago): missing + // entries are "does not exist". + clock.now = t0.Add(31 * time.Second) + _, err = svc.LookupNode(context.Background(), &schedulerv1.LookupNodeRequest{SandboxId: "sbx-ghost"}) + if status.Code(err) != codes.NotFound { + t.Fatalf("window closed: LookupNode want NotFound, got %v", err) + } + _, err = svc.GetNode(context.Background(), &schedulerv1.GetNodeRequest{NodeId: "node-ghost"}) + if status.Code(err) != codes.NotFound { + t.Fatalf("window closed: GetNode want NotFound, got %v", err) + } +} diff --git a/services/scheduler/internal/service.go b/services/scheduler/internal/service.go index a284dfa84..050c6c4d1 100644 --- a/services/scheduler/internal/service.go +++ b/services/scheduler/internal/service.go @@ -15,6 +15,10 @@ import ( "google.golang.org/grpc/status" ) +// The Service sees leadership state only through the narrow leadershipView +// seam (defined in leadership.go); LeadershipManager is the only production +// implementation, and the internal snapshot object is never handed out. + type Service struct { schedulerv1.UnimplementedSchedulerServer logger *zap.Logger @@ -23,6 +27,9 @@ type Service struct { store BindingStore artifacts ArtifactStore resourceLimit *config.NodeResourceLimit + leadership leadershipView + // for test cases, inject for time-related function checking on `inRecoveryWindow` + now func() time.Time } func NewService(logger *zap.Logger, nodes NodeRegistry, strategy Strategy, store BindingStore, opts ...ServiceOption) *Service { @@ -38,6 +45,7 @@ func NewService(logger *zap.Logger, nodes NodeRegistry, strategy Strategy, store strategy: strategy, store: store, artifacts: NewInMemoryArtifactStore(defaultArtifactStoreCapacity, 0), + now: time.Now, } for _, opt := range opts { opt(s) @@ -61,6 +69,22 @@ func WithArtifactStore(store ArtifactStore) ServiceOption { } } +// WithLeadership attaches the read-only leadership view so scheduling and +// lookup can apply recovery-window rules right after a leadership +// transition (#259). +func WithLeadership(leadership leadershipView) ServiceOption { + return func(s *Service) { + s.leadership = leadership + } +} + +// WithClock overrides the time source (tests only). +func WithClock(now func() time.Time) ServiceOption { + return func(s *Service) { + s.now = now + } +} + type QueryOnlyService struct { schedulerv1.UnimplementedSchedulerServer logger *zap.Logger @@ -93,6 +117,8 @@ func (s *Service) Schedule(_ context.Context, req *schedulerv1.ScheduleRequest) }) } + rich = s.filterFreshObservations(rich) + eligible := FilterByResourceLimit(rich, s.resourceLimit) node, selectErr := s.strategy.Select(eligible, req.GetHint()) @@ -122,6 +148,31 @@ func (s *Service) Schedule(_ context.Context, req *schedulerv1.ScheduleRequest) return &schedulerv1.ScheduleResponse{Node: node.Node.ToProto()}, nil } +// filterFreshObservations applies the fresh-observations-only rule (#259): +// while the recovery window is open, only nodes observed at or after the +// leadership acquisition are scheduling candidates. Outside the window (or +// with leader election disabled) all nodes pass, preserving today's +// fail-open behaviour for new nodes. +func (s *Service) filterFreshObservations(rich []RichNode) []RichNode { + if !s.inRecoveryWindow() { + return rich + } + since, _ := s.leadership.LeaderSince() + fresh := make([]RichNode, 0, len(rich)) + for _, n := range rich { + if at, ok := s.nodes.LastReportAt(n.Node.ID); ok && !at.Before(since) { + fresh = append(fresh, n) + } + } + return fresh +} + +// inRecoveryWindow reports whether the recovery window following leadership +// acquisition is open right now. Without leader election it is always closed. +func (s *Service) inRecoveryWindow() bool { + return s.leadership != nil && s.leadership.InRecoveryWindow(s.now()) +} + // summarizeScheduleHint renders a compact, log-friendly description of a // scheduling hint. func summarizeScheduleHint(hint *schedulerv1.ScheduleRequestHint) string { @@ -149,7 +200,13 @@ func (s *Service) ListNodes(_ context.Context, _ *schedulerv1.ListNodesRequest) } func (s *Service) LookupNode(_ context.Context, req *schedulerv1.LookupNodeRequest) (*schedulerv1.LookupNodeResponse, error) { - return lookupNode(s.logger, s.store, req) + resp, err := lookupNode(s.logger, s.store, req) + if err != nil && status.Code(err) == codes.NotFound && s.inRecoveryWindow() { + // Recovery window open: a missing binding means "not yet rebuilt", + // not "does not exist" (#259). + return nil, status.Error(codes.Unavailable, "sandbox assignment not rebuilt yet") + } + return resp, err } func lookupNode(logger *zap.Logger, store BindingStore, req *schedulerv1.LookupNodeRequest) (*schedulerv1.LookupNodeResponse, error) { @@ -212,20 +269,42 @@ func (s *Service) Heartbeat(_ context.Context, req *schedulerv1.HeartbeatRequest return nil, status.Error(codes.InvalidArgument, "node_id and service_instance_id are required") } - now := time.Now() - node, cpuConfigJSON, err := s.nodes.Heartbeat(req, now) + return s.ingestNodeReport(nodeReportFromProto(req), s.now()) +} + +// nodeReportFromProto adapts a Heartbeat RPC request to the neutral ingest +// currency. Identity validation already happened in Heartbeat. +func nodeReportFromProto(req *schedulerv1.HeartbeatRequest) *nodeReport { + return &nodeReport{ + nodeID: req.GetNodeId(), + clusterID: req.GetClusterId(), + serviceInstanceID: req.GetServiceInstanceId(), + version: req.GetVersion(), + commit: req.GetCommit(), + machineInfo: req.GetMachineInfo(), + snapshot: req.GetSnapshot(), + sandboxIDs: req.GetSandboxIds(), + p2pEndpoint: req.GetP2PEndpoint(), + } +} + +// ingestNodeReport is the shared ingest path for heartbeat RPCs and pulled +// node snapshots (sync-node-snapshots on leadership acquisition, #259): +// update registry observations, then reconcile sandbox bindings. +func (s *Service) ingestNodeReport(report *nodeReport, now time.Time) (*schedulerv1.HeartbeatResponse, error) { + node, cpuConfigJSON, err := s.nodes.Heartbeat(report.toHeartbeatRequest(), now) if err != nil { if errors.Is(err, ErrNodeNotInRegistry) { s.logger.Warn("scheduler rejected observed registration for unknown node", - zap.String("node_id", nodeID), + zap.String("node_id", report.nodeID), ) return nil, status.Error(codes.InvalidArgument, "node is not in scheduler node list") } return nil, status.Error(codes.Internal, "node registry heartbeat failed") } - if err := s.store.ReconcileNode(node, req.GetSandboxIds(), now); err != nil { + if err := s.store.ReconcileNode(node, report.sandboxIDs, now); err != nil { s.logger.Warn("scheduler heartbeat binding reconcile failed", - zap.String("node_id", nodeID), + zap.String("node_id", report.nodeID), zap.Error(err), ) return nil, status.Error(codes.Unavailable, "binding store unavailable") @@ -336,8 +415,12 @@ func (s *Service) GetNode(_ context.Context, req *schedulerv1.GetNodeRequest) (* return nil, status.Error(codes.InvalidArgument, "node_id is required") } - node, ok := s.nodes.GetObserved(nodeID, req.GetClusterId(), time.Now()) + node, ok := s.nodes.GetObserved(nodeID, req.GetClusterId(), s.now()) if !ok { + if s.inRecoveryWindow() { + // Same recovery-window rule as LookupNode (#259). + return nil, status.Error(codes.Unavailable, "observed node not rebuilt yet") + } return nil, status.Error(codes.NotFound, "observed node not found") } From f1a67c37725acb94287754552b5fa15be23094ab Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 13:06:25 +0800 Subject: [PATCH 03/26] feat(scheduler): refresh node snapshots after leadership acquisition Add sync-node-snapshots (#259): instead of waiting out a reporter backoff after winning the lease, the new leader refreshes observations and bindings in one RTT per node. - nodeSnapshotRefresher contract with a single production implementation, ConcurrentNodeSnapshotRefresher: bounded fan-out (leader_election.snapshot_pull_concurrency) pulling every registry node exactly once; failed pulls are logged and leave the node unobserved, where the recovery window's fresh-observations-only rule keeps it out of scheduling until its next report. - nodeReport is the neutral ingest currency: both the Heartbeat RPC handler and the pull path convert into it, and it feeds the shared ingest path (registry observation update + ReconcileNode binding refresh), so pulled and heartbeated state are byte-identical in effect. - NodeSnapshotFetcher seam with a production implementation talking to each node's admin HTTP API (GET /nodes + GET /sandboxes, x-api-key from SCHEDULER_NODE_ADMIN_API_KEY); the HTTP client is an internal default with a bounded timeout, consistent with the gateway's default-client pattern. Scope: services/scheduler/internal (node_snapshot.go) with fetcher mapping, auth-failure, ingest-compatibility, and pull-failure tests. Relates to #259. Co-Authored-By: Claude Code --- services/scheduler/internal/node_snapshot.go | 282 ++++++++++++++++++ .../scheduler/internal/node_snapshot_test.go | 124 ++++++++ 2 files changed, 406 insertions(+) create mode 100644 services/scheduler/internal/node_snapshot.go create mode 100644 services/scheduler/internal/node_snapshot_test.go diff --git a/services/scheduler/internal/node_snapshot.go b/services/scheduler/internal/node_snapshot.go new file mode 100644 index 000000000..35111ba56 --- /dev/null +++ b/services/scheduler/internal/node_snapshot.go @@ -0,0 +1,282 @@ +package scheduler + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "sync" + "time" + + schedulerv1 "agentenv/services/api/proto" + + "go.uber.org/zap" +) + +// ============================================================ + +// nodeSnapshotRefresher is this file's contract: refresh the leader's +// observations and bindings for a set of nodes by pulling their admin +// snapshots (sync-node-snapshots, #259). The only production implementation +// is ConcurrentNodeSnapshotRefresher; callers (the leadership manager) depend on this +// interface, not on the concrete type. +type nodeSnapshotRefresher interface { + Refresh(ctx context.Context, nodes []Node) +} + +// nodeReport is the neutral currency of node-state ingestion (#259): what a +// fetcher returns after successfully pulling a node, and what the shared +// ingest path consumes. Both the Heartbeat RPC handler and +// sync-node-snapshots convert into it; a fetcher never fabricates an RPC +// request. +type nodeReport struct { + nodeID string + serviceInstanceID string + clusterID string + version string + commit string + machineInfo *schedulerv1.MachineInfo + snapshot *schedulerv1.NodeSnapshot + sandboxIDs []string + p2pEndpoint *schedulerv1.P2PEndpoint +} + +// toHeartbeatRequest converts the report to the registry's proto shape; the +// conversion is contained here so the registry interface stays unchanged. +func (r *nodeReport) toHeartbeatRequest() *schedulerv1.HeartbeatRequest { + return &schedulerv1.HeartbeatRequest{ + NodeId: r.nodeID, + ClusterId: r.clusterID, + ServiceInstanceId: r.serviceInstanceID, + Version: r.version, + Commit: r.commit, + MachineInfo: r.machineInfo, + Snapshot: r.snapshot, + SandboxIds: r.sandboxIDs, + P2PEndpoint: r.p2pEndpoint, + } +} + +// NodeSnapshotFetcher retrieves a single node's current snapshot from its +// admin endpoints for sync-node-snapshots on leadership acquisition (#259). +// It returns the fetched data (a nodeReport), not an RPC request. It is an +// injected seam so the retrieval logic is testable without HTTP or node +// credentials. +type NodeSnapshotFetcher func(ctx context.Context, node Node) (*nodeReport, error) + +// ConcurrentNodeSnapshotRefresher drives the post-acquisition state rebuild: fan out a +// retrieval of every registry node's snapshot with bounded concurrency and +// feed each response through the shared ingest path (ingestNodeReport — the +// same code a heartbeat takes). A failed retrieval leaves that node +// unobserved; the recovery window's fresh-observations-only rule keeps it out of +// scheduling until its next report. +// nodeReportIngester is the retriever's only dependency on the Service: the +// shared ingest path. The retriever never holds the Service itself. +type nodeReportIngester func(report *nodeReport, now time.Time) error + +type ConcurrentNodeSnapshotRefresher struct { + logger *zap.Logger + ingest nodeReportIngester + fetch NodeSnapshotFetcher + concurrency int +} + +func NewConcurrentNodeSnapshotRefresher(logger *zap.Logger, ingest nodeReportIngester, fetch NodeSnapshotFetcher, concurrency int) *ConcurrentNodeSnapshotRefresher { + if logger == nil { + logger = zap.NewNop() + } + if concurrency <= 0 { + concurrency = 4 + } + return &ConcurrentNodeSnapshotRefresher{logger: logger, ingest: ingest, fetch: fetch, concurrency: concurrency} +} + +// Refresh pulls a snapshot from every node in nodes and ingests the results. +// Every node is fetched exactly once, successes are ingested, failures are +// logged and leave the node unobserved (the recovery window's +// fresh-observations-only rule keeps it out of scheduling until its next +// report). +func (r *ConcurrentNodeSnapshotRefresher) Refresh(ctx context.Context, nodes []Node) { + sem := make(chan struct{}, r.concurrency) + var wg sync.WaitGroup + for _, node := range nodes { + node := node + wg.Add(1) + go func() { + defer wg.Done() + sem <- struct{}{} + defer func() { <-sem }() + report, err := r.fetch(ctx, node) + if err != nil { + r.logger.Warn("sync-node-snapshots: node snapshot pull failed", + zap.String("node_id", node.ID), zap.Error(err)) + return + } + if report == nil { + return + } + if err := r.ingest(report, time.Now()); err != nil { + r.logger.Warn("sync-node-snapshots: ingesting pulled snapshot failed", + zap.String("node_id", node.ID), zap.Error(err)) + } + }() + } + wg.Wait() +} + +// ============================================================ + +// adminNodeResponse mirrors the agentenv server admin GET /nodes entry +// (openapi: Node). Field names follow the API's camelCase JSON. +type adminNodeResponse struct { + Version string `json:"version"` + Commit string `json:"commit"` + ID string `json:"id"` + ServiceInstanceID string `json:"serviceInstanceID"` + ClusterID string `json:"clusterID"` + SandboxCount uint32 `json:"sandboxCount"` + CreateSuccesses uint64 `json:"createSuccesses"` + CreateFails uint64 `json:"createFails"` + SandboxStartingCnt uint32 `json:"sandboxStartingCount"` + SandboxPausedCount uint32 `json:"sandboxPausedCount"` + MachineInfo struct { + CPUFamily string `json:"cpuFamily"` + CPUModel string `json:"cpuModel"` + CPUModelName string `json:"cpuModelName"` + CPUArchitecture string `json:"cpuArchitecture"` + CPUConfigJSON string `json:"cpuConfigJSON"` + } `json:"machineInfo"` + Metrics struct { + AllocatedCPU uint32 `json:"allocatedCPU"` + AllocatedMemoryBytes uint64 `json:"allocatedMemoryBytes"` + CPUPercent uint32 `json:"cpuPercent"` + CPUCount uint32 `json:"cpuCount"` + MemoryUsedBytes uint64 `json:"memoryUsedBytes"` + MemoryTotalBytes uint64 `json:"memoryTotalBytes"` + PausedAllocatedCPU uint32 `json:"pausedAllocatedCPU"` + PausedAllocatedMemoryBytes uint64 `json:"pausedAllocatedMemoryBytes"` + Disks []struct { + MountPoint string `json:"mountPoint"` + Device string `json:"device"` + FilesystemType string `json:"filesystemType"` + UsedBytes uint64 `json:"usedBytes"` + TotalBytes uint64 `json:"totalBytes"` + } `json:"disks"` + } `json:"metrics"` +} + +// adminSandboxEntry mirrors a ListedSandbox entry from GET /sandboxes; only +// the id is needed for binding reconciliation. +type adminSandboxEntry struct { + SandboxID string `json:"sandboxID"` +} + +// NewAdminSnapshotFetcher builds the production NodeSnapshotFetcher (#259, +// sync-node-snapshots): it pulls GET {endpoint}/nodes for observations and +// GET {endpoint}/sandboxes for the sandbox id list, and returns the fetched +// data as a nodeReport for the shared ingest path. +// The node's admin API requires the x-api-key header. The HTTP client is an +// implementation detail with a bounded per-request timeout; tests inject +// through the NodeSnapshotFetcher seam, not this constructor. +func NewAdminSnapshotFetcher(apiKey string) NodeSnapshotFetcher { + client := &http.Client{Timeout: 5 * time.Second} + return func(ctx context.Context, node Node) (*nodeReport, error) { + base := strings.TrimRight(node.Endpoint, "/") + + var nodes []adminNodeResponse + if err := adminGetJSON(ctx, client, apiKey, base+"/nodes", &nodes); err != nil { + return nil, fmt.Errorf("pull %s /nodes: %w", node.ID, err) + } + var entry *adminNodeResponse + for i := range nodes { + if nodes[i].ID == node.ID { + entry = &nodes[i] + break + } + } + if entry == nil && len(nodes) == 1 { + entry = &nodes[0] + } + if entry == nil { + return nil, fmt.Errorf("pull %s /nodes: node not in admin response", node.ID) + } + + var sandboxes []adminSandboxEntry + if err := adminGetJSON(ctx, client, apiKey, base+"/sandboxes", &sandboxes); err != nil { + return nil, fmt.Errorf("pull %s /sandboxes: %w", node.ID, err) + } + + return reportFromAdmin(entry, sandboxes), nil + } +} + +func adminGetJSON(ctx context.Context, client *http.Client, apiKey, url string, out any) error { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return err + } + if strings.TrimSpace(apiKey) != "" { + req.Header.Set("x-api-key", apiKey) + } + resp, err := client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("unexpected status %s", resp.Status) + } + return json.NewDecoder(resp.Body).Decode(out) +} + +func reportFromAdmin(n *adminNodeResponse, sandboxes []adminSandboxEntry) *nodeReport { + ids := make([]string, 0, len(sandboxes)) + for _, s := range sandboxes { + if strings.TrimSpace(s.SandboxID) != "" { + ids = append(ids, s.SandboxID) + } + } + disks := make([]*schedulerv1.DiskMetric, 0, len(n.Metrics.Disks)) + for _, d := range n.Metrics.Disks { + disks = append(disks, &schedulerv1.DiskMetric{ + MountPoint: d.MountPoint, + Device: d.Device, + FilesystemType: d.FilesystemType, + UsedBytes: d.UsedBytes, + TotalBytes: d.TotalBytes, + }) + } + return &nodeReport{ + nodeID: n.ID, + clusterID: n.ClusterID, + serviceInstanceID: n.ServiceInstanceID, + version: n.Version, + commit: n.Commit, + machineInfo: &schedulerv1.MachineInfo{ + CpuFamily: n.MachineInfo.CPUFamily, + CpuModel: n.MachineInfo.CPUModel, + CpuModelName: n.MachineInfo.CPUModelName, + CpuArchitecture: n.MachineInfo.CPUArchitecture, + CpuConfigJson: n.MachineInfo.CPUConfigJSON, + }, + snapshot: &schedulerv1.NodeSnapshot{ + AllocatedCpu: n.Metrics.AllocatedCPU, + AllocatedMemoryBytes: n.Metrics.AllocatedMemoryBytes, + CpuPercent: n.Metrics.CPUPercent, + CpuCount: n.Metrics.CPUCount, + MemoryUsedBytes: n.Metrics.MemoryUsedBytes, + MemoryTotalBytes: n.Metrics.MemoryTotalBytes, + Disks: disks, + SandboxCount: n.SandboxCount, + SandboxStartingCount: n.SandboxStartingCnt, + CreateSuccesses: n.CreateSuccesses, + CreateFails: n.CreateFails, + ReportedAtUnixMs: time.Now().UnixMilli(), + PausedSandboxCount: n.SandboxPausedCount, + PausedAllocatedCpu: n.Metrics.PausedAllocatedCPU, + PausedAllocatedMemoryBytes: n.Metrics.PausedAllocatedMemoryBytes, + }, + sandboxIDs: ids, + } +} diff --git a/services/scheduler/internal/node_snapshot_test.go b/services/scheduler/internal/node_snapshot_test.go new file mode 100644 index 000000000..a2c183918 --- /dev/null +++ b/services/scheduler/internal/node_snapshot_test.go @@ -0,0 +1,124 @@ +package scheduler + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +// Sync-node-snapshots fetcher tests (#259): the production pull assembles a +// heartbeat-shaped request from GET /nodes + GET /sandboxes, honoring the +// x-api-key header. + +func adminTestServer(t *testing.T, wantKey string, nodeJSON, sandboxesJSON string) *httptest.Server { + t.Helper() + mux := http.NewServeMux() + mux.HandleFunc("/nodes", func(w http.ResponseWriter, r *http.Request) { + if wantKey != "" && r.Header.Get("x-api-key") != wantKey { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(nodeJSON)) + }) + mux.HandleFunc("/sandboxes", func(w http.ResponseWriter, r *http.Request) { + if wantKey != "" && r.Header.Get("x-api-key") != wantKey { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(sandboxesJSON)) + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + return srv +} + +const adminNodeFixture = `[{ + "version": "0.2.0", + "commit": "abc123", + "id": "node-a", + "serviceInstanceID": "inst-1", + "clusterID": "cluster-1", + "sandboxCount": 2, + "createSuccesses": 10, + "createFails": 1, + "sandboxStartingCount": 1, + "sandboxPausedCount": 3, + "machineInfo": {"cpuFamily": "6", "cpuModel": "85", "cpuModelName": "Xeon", "cpuArchitecture": "x86_64", "cpuConfigJSON": "{}"}, + "metrics": { + "allocatedCPU": 4, + "allocatedMemoryBytes": 8589934592, + "cpuPercent": 55, + "cpuCount": 16, + "memoryUsedBytes": 17179869184, + "memoryTotalBytes": 34359738368, + "pausedAllocatedCPU": 2, + "pausedAllocatedMemoryBytes": 4294967296, + "disks": [{"mountPoint": "/", "device": "/dev/ublkb0", "filesystemType": "ext4", "usedBytes": 1024, "totalBytes": 4096}] + } +}]` + +const adminSandboxesFixture = `[{"sandboxID": "sbx-1"}, {"sandboxID": "sbx-2"}]` + +func TestAdminSnapshotFetcherAssemblesHeartbeatShape(t *testing.T) { + srv := adminTestServer(t, "secret", adminNodeFixture, adminSandboxesFixture) + fetch := NewAdminSnapshotFetcher("secret") + + req, err := fetch(context.Background(), Node{ID: "node-a", Endpoint: srv.URL}) + if err != nil { + t.Fatalf("fetch failed: %v", err) + } + + if req.nodeID != "node-a" || req.serviceInstanceID != "inst-1" || req.clusterID != "cluster-1" { + t.Fatalf("identity fields wrong: %v", req) + } + snap := req.snapshot + if snap == nil { + t.Fatal("snapshot must be populated") + } + if snap.GetAllocatedCpu() != 4 || snap.GetCpuPercent() != 55 || snap.GetSandboxCount() != 2 { + t.Fatalf("metrics mapping wrong: %+v", snap) + } + if snap.GetPausedSandboxCount() != 3 || snap.GetPausedAllocatedCpu() != 2 { + t.Fatalf("paused mapping wrong: %+v", snap) + } + if len(snap.GetDisks()) != 1 || snap.GetDisks()[0].GetDevice() != "/dev/ublkb0" { + t.Fatalf("disk mapping wrong: %+v", snap.GetDisks()) + } + ids := req.sandboxIDs + if len(ids) != 2 || ids[0] != "sbx-1" || ids[1] != "sbx-2" { + t.Fatalf("sandbox ids wrong: %v", ids) + } +} + +func TestAdminSnapshotFetcherAuthFailure(t *testing.T) { + srv := adminTestServer(t, "secret", adminNodeFixture, adminSandboxesFixture) + fetch := NewAdminSnapshotFetcher("wrong-key") + + if _, err := fetch(context.Background(), Node{ID: "node-a", Endpoint: srv.URL}); err == nil { + t.Fatal("fetch with a wrong key must fail") + } +} + +func TestAdminSnapshotFetcherHeartbeatIngestCompatibility(t *testing.T) { + srv := adminTestServer(t, "", adminNodeFixture, adminSandboxesFixture) + fetch := NewAdminSnapshotFetcher("") + + svc, registry, store := newTestService(t, []string{"node-a"}) + req, err := fetch(context.Background(), Node{ID: "node-a", Endpoint: srv.URL}) + if err != nil { + t.Fatalf("fetch failed: %v", err) + } + if _, err := svc.ingestNodeReport(req, time.Now()); err != nil { + t.Fatalf("pulled snapshot must ingest through the shared path: %v", err) + } + if registry.PeekObserved("node-a") == nil { + t.Fatal("pulled snapshot must record an observation") + } + if _, ok, err := store.Get("sbx-1", time.Now()); err != nil || !ok { + t.Fatalf("pulled snapshot must reconcile bindings: ok=%v err=%v", ok, err) + } +} From aac26dd1eda34272e330d64ceefb5c0364914e85 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 13:08:01 +0800 Subject: [PATCH 04/26] feat(scheduler): wire leader election in main and floor binding TTL Wire the leadership facade into the scheduler process (#259): - main builds Leadership via the factory and wires it unconditionally: gate interceptor chained after metrics, Service option attached, leader health registered, BindRuntime + Run driven in a goroutine. A leadership loss cancels the process root context, driving the same graceful shutdown as SIGTERM so the ex-leader exits and restarts as a standby (fencing). - Binding TTL floor under election: effective TTL = max(binding_ttl, lease_duration + 90s), so bindings outlive the worst-case failover budget (lease detection + endpoints propagation + one reporter reconnect backoff) and live sandboxes are not misreported NotFound mid-failover. Authoritative cleanup stays with ReconcileNode; the TTL only guards crashed nodes, so raising it is harmless. Logged at startup when raised. Scope: services/scheduler/cmd, plus go.mod/go.sum entries for client-go's testing/fixture packages used by the fake-clientset election tests. Relates to #259. Co-Authored-By: Claude Code --- services/go.mod | 8 ++- services/go.sum | 22 ++++--- services/scheduler/cmd/binding_ttl_test.go | 49 ++++++++++++++ services/scheduler/cmd/main.go | 75 ++++++++++++++++++---- 4 files changed, 132 insertions(+), 22 deletions(-) create mode 100644 services/scheduler/cmd/binding_ttl_test.go diff --git a/services/go.mod b/services/go.mod index 42f4e915b..2d108aea4 100644 --- a/services/go.mod +++ b/services/go.mod @@ -20,6 +20,7 @@ require ( github.com/davecgh/go-spew v1.1.1 // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/emicklei/go-restful/v3 v3.11.0 // indirect + github.com/evanphx/json-patch v4.12.0+incompatible // indirect github.com/go-logr/logr v1.4.3 // indirect github.com/go-openapi/jsonpointer v0.19.6 // indirect github.com/go-openapi/jsonreference v0.20.2 // indirect @@ -37,10 +38,13 @@ require ( github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect + github.com/onsi/ginkgo/v2 v2.14.0 // indirect + github.com/onsi/gomega v1.30.0 // indirect + github.com/pkg/errors v0.9.1 // indirect github.com/prometheus/client_model v0.6.1 // indirect github.com/prometheus/common v0.55.0 // indirect github.com/prometheus/procfs v0.15.1 // indirect - go.uber.org/multierr v1.10.0 // indirect + go.uber.org/multierr v1.11.0 // indirect golang.org/x/net v0.55.0 // indirect golang.org/x/oauth2 v0.36.0 // indirect golang.org/x/sys v0.45.0 // indirect @@ -56,5 +60,5 @@ require ( k8s.io/utils v0.0.0-20230726121419-3b25d923346b // indirect sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect sigs.k8s.io/structured-merge-diff/v4 v4.4.1 // indirect - sigs.k8s.io/yaml v1.3.0 // indirect + sigs.k8s.io/yaml v1.4.0 // indirect ) diff --git a/services/go.sum b/services/go.sum index 8fa5d2de0..4eb1a629e 100644 --- a/services/go.sum +++ b/services/go.sum @@ -14,6 +14,8 @@ github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/r github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/emicklei/go-restful/v3 v3.11.0 h1:rAQeMHw1c7zTmncogyy8VvRZwtkmkZ4FxERmMY4rD+g= github.com/emicklei/go-restful/v3 v3.11.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc= +github.com/evanphx/json-patch v4.12.0+incompatible h1:4onqiflcdA9EOZ4RxV643DvftH5pOlLGNtQ5lPWQu84= +github.com/evanphx/json-patch v4.12.0+incompatible/go.mod h1:50XU6AFN0ol/bzJsmQLiYLvXMP4fmwYFNcr97nuDLSk= github.com/go-logr/logr v1.3.0/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= @@ -29,6 +31,8 @@ github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEe github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls= github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= +github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da h1:oI5xCqsCo564l8iNU+DwB5epxmsaqB+rhGL0m5jtYqE= +github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= github.com/google/gnostic-models v0.6.8 h1:yo/ABAfM5IMRsS1VnXjTBvUb61tFIHozhlYvRgGre9I= @@ -71,10 +75,12 @@ github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9G github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= -github.com/onsi/ginkgo/v2 v2.13.0 h1:0jY9lJquiL8fcf3M4LAXN5aMlS/b2BV86HFFPCPMgE4= -github.com/onsi/ginkgo/v2 v2.13.0/go.mod h1:TE309ZR8s5FsKKpuB1YAQYBzCaAfUgatB/xlT/ETL/o= -github.com/onsi/gomega v1.29.0 h1:KIA/t2t5UBzoirT4H9tsML45GEbo3ouUnBHsCfD2tVg= -github.com/onsi/gomega v1.29.0/go.mod h1:9sxs+SwGrKI0+PWe4Fxa9tFQQBG5xSsSbMXOI8PPpoQ= +github.com/onsi/ginkgo/v2 v2.14.0 h1:vSmGj2Z5YPb9JwCWT6z6ihcUvDhuXLc3sJiqd3jMKAY= +github.com/onsi/ginkgo/v2 v2.14.0/go.mod h1:JkUdW7JkN0V6rFvsHcJ478egV3XH9NxpD27Hal/PhZw= +github.com/onsi/gomega v1.30.0 h1:hvMK7xYz4D3HapigLTeGdId/NcfQx1VHMJc60ew99+8= +github.com/onsi/gomega v1.30.0/go.mod h1:9sxs+SwGrKI0+PWe4Fxa9tFQQBG5xSsSbMXOI8PPpoQ= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prometheus/client_golang v1.20.5 h1:cxppBPuYhUnsO6yo/aoRol4L7q7UFfdm+bR9r+8l63Y= @@ -116,8 +122,8 @@ go.opentelemetry.io/otel/trace v1.44.0 h1:jxF5CsGYCe74MCRx2X4g7WsY/VBKRqqpNvXlX/ go.opentelemetry.io/otel/trace v1.44.0/go.mod h1:oLl1jrMQAVo6v3GAggN+1VH9VIz9iUSvW53sW1Q8PIE= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= -go.uber.org/multierr v1.10.0 h1:S0h4aNzvfcFsC3dRF1jLoaov7oRaKqRGC/pUEJ2yvPQ= -go.uber.org/multierr v1.10.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= +go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0= +go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8= go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= @@ -194,5 +200,5 @@ sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd h1:EDPBXCAspyGV4jQlpZSudPeMm sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd/go.mod h1:B8JuhiUyNFVKdsE8h686QcCxMaH6HrOAZj4vswFpcB0= sigs.k8s.io/structured-merge-diff/v4 v4.4.1 h1:150L+0vs/8DA78h1u02ooW1/fFq/Lwr+sGiqlzvrtq4= sigs.k8s.io/structured-merge-diff/v4 v4.4.1/go.mod h1:N8hJocpFajUSSeSJ9bOZ77VzejKZaXsTtZo4/u7Io08= -sigs.k8s.io/yaml v1.3.0 h1:a2VclLzOGrwOHDiV8EfBGhvjHvP46CtW5j6POvhYGGo= -sigs.k8s.io/yaml v1.3.0/go.mod h1:GeOyir5tyXNByN85N/dRIT9es5UQNerPYEKK56eTBm8= +sigs.k8s.io/yaml v1.4.0 h1:Mk1wCc2gy/F0THH0TAp1QYyJNzRm2KCLy3o5ASXVI5E= +sigs.k8s.io/yaml v1.4.0/go.mod h1:Ejl7/uTz7PSA4eKMyQCUTnhZYNmLIl+5c2lQPGR2BPY= diff --git a/services/scheduler/cmd/binding_ttl_test.go b/services/scheduler/cmd/binding_ttl_test.go new file mode 100644 index 000000000..7479eac8c --- /dev/null +++ b/services/scheduler/cmd/binding_ttl_test.go @@ -0,0 +1,49 @@ +package main + +import ( + "testing" + "time" + + "agentenv/services/shared/config" +) + +// Binding TTL floor tests (#259): under leader election the effective TTL +// must cover the failover budget (lease_duration + 90s); without election the +// configured value passes through unchanged. + +func baseCfg(bindingTTL time.Duration) config.Config { + return config.Config{ + Scheduler: config.SchedulerConfig{BindingTTL: bindingTTL}, + } +} + +func TestEffectiveBindingTTLWithoutElection(t *testing.T) { + cfg := baseCfg(30 * time.Second) + if got := effectiveBindingTTL(cfg); got != 30*time.Second { + t.Fatalf("election off: want configured TTL 30s, got %v", got) + } +} + +func TestEffectiveBindingTTLWithElectionFloors(t *testing.T) { + cfg := baseCfg(30 * time.Second) + cfg.Scheduler.LeaderElection = config.SchedulerLeaderElectionConfig{ + Enabled: true, + LeaseDuration: 15 * time.Second, + } + // floor = 15s + 90s = 105s > 30s + if got := effectiveBindingTTL(cfg); got != 105*time.Second { + t.Fatalf("election on: want floor 105s, got %v", got) + } +} + +func TestEffectiveBindingTTLWithElectionKeepsLargerConfigured(t *testing.T) { + cfg := baseCfg(5 * time.Minute) + cfg.Scheduler.LeaderElection = config.SchedulerLeaderElectionConfig{ + Enabled: true, + LeaseDuration: 15 * time.Second, + } + // floor = 105s < 5m: the larger configured value wins. + if got := effectiveBindingTTL(cfg); got != 5*time.Minute { + t.Fatalf("election on: want configured 5m (larger than floor), got %v", got) + } +} diff --git a/services/scheduler/cmd/main.go b/services/scheduler/cmd/main.go index 3e8574a5e..1fea11a84 100644 --- a/services/scheduler/cmd/main.go +++ b/services/scheduler/cmd/main.go @@ -42,19 +42,33 @@ func main() { } defer logger.Sync() - sigCtx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + // rootCancel lets a leadership loss drive the same graceful shutdown as a + // signal: the ex-leader exits and restarts as a standby (fencing, #259). + rootCtx, rootCancel := context.WithCancel(context.Background()) + defer rootCancel() + sigCtx, stop := signal.NotifyContext(rootCtx, os.Interrupt, syscall.SIGTERM) defer stop() store, closeStore := createBindingStore(logger, cfg) defer closeStore() - g := grpc.NewServer(grpc.UnaryInterceptor(scheduler.MetricsUnaryInterceptor())) + // Leadership facade (#259): the factory picks the election manager (election + // on) or a no-op (election off); all wiring below is unconditional. + leadership := scheduler.NewLeadership(logger, cfg.Scheduler) + + interceptors := []grpc.UnaryServerInterceptor{ + scheduler.MetricsUnaryInterceptor(), + leadership.GateInterceptor(), + } + g := grpc.NewServer(grpc.ChainUnaryInterceptor(interceptors...)) + var registry *scheduler.AtomicNodeRegistry + var svc *scheduler.Service if *queryOnly { - svc := scheduler.NewQueryOnlyService(logger, store) - schedulerv1.RegisterSchedulerServer(g, svc) + qo := scheduler.NewQueryOnlyService(logger, store) + schedulerv1.RegisterSchedulerServer(g, qo) logger.Info("scheduler query-only service enabled", zap.String("redis_addr", cfg.Scheduler.RedisAddr)) } else { - registry := scheduler.NewAtomicNodeRegistry(nil, cfg.Scheduler.ReportTTL) + registry = scheduler.NewAtomicNodeRegistry(nil, cfg.Scheduler.ReportTTL) switch strings.ToLower(strings.TrimSpace(cfg.Scheduler.Discovery.Mode)) { case "kubernetes": go runKubernetesDiscoveryWithRetry(sigCtx, logger, cfg.Scheduler.Discovery.Kubernetes, registry) @@ -66,16 +80,20 @@ func main() { registry.Set(nodes, nil) } - svc := scheduler.NewService( - logger, - registry, - scheduler.NewStrategy(cfg.Scheduler.Strategy), - store, + svcOpts := []scheduler.ServiceOption{ scheduler.WithArtifactStore(scheduler.NewInMemoryArtifactStore( cfg.Scheduler.ArtifactStoreCapacity, cfg.Scheduler.ArtifactLookupNodeLimit, )), scheduler.WithNodeResourceLimit(cfg.Scheduler.NodeResourceLimit), + } + svcOpts = append(svcOpts, leadership.ServiceOption()) + svc = scheduler.NewService( + logger, + registry, + scheduler.NewStrategy(cfg.Scheduler.Strategy), + store, + svcOpts..., ) go svc.RunObservedNodesMetrics(sigCtx, 15*time.Second) schedulerv1.RegisterSchedulerServer(g, svc) @@ -84,8 +102,16 @@ func main() { hs := health.NewServer() hs.SetServingStatus("", grpc_health_v1.HealthCheckResponse_SERVING) hs.SetServingStatus(schedulerv1.Scheduler_ServiceDesc.ServiceName, grpc_health_v1.HealthCheckResponse_SERVING) + leadership.RegisterHealth(hs) grpc_health_v1.RegisterHealthServer(g, hs) + leadership.BindRuntime(svc, registry) + go func() { + if err := leadership.Run(sigCtx, rootCancel); err != nil { + logger.Fatal("leader election failed", zap.Error(err)) + } + }() + lis, err := net.Listen("tcp", cfg.Scheduler.GRPCListenAddr) if err != nil { logger.Fatal("listen failed", zap.Error(err), zap.String("addr", cfg.Scheduler.GRPCListenAddr)) @@ -130,6 +156,7 @@ func main() { logger.Info("scheduler shutdown signal received") hs.SetServingStatus("", grpc_health_v1.HealthCheckResponse_NOT_SERVING) hs.SetServingStatus(schedulerv1.Scheduler_ServiceDesc.ServiceName, grpc_health_v1.HealthCheckResponse_NOT_SERVING) + leadership.MarkNotServing() gracefulStopDone := make(chan struct{}) go func() { @@ -160,12 +187,36 @@ func main() { } } +// effectiveBindingTTL floors the binding TTL under leader election (#259, +// failover safeguard): bindings must outlive the worst-case failover budget +// (lease detection + endpoints propagation + one reporter reconnect backoff), +// or live sandboxes would be misreported as NotFound mid-failover. The +// authoritative cleanup is ReconcileNode; the TTL only guards against a node +// that crashed for good, so raising it is harmless. +func effectiveBindingTTL(cfg config.Config) time.Duration { + ttl := cfg.Scheduler.BindingTTL + if cfg.Scheduler.LeaderElection.Enabled { + floor := cfg.Scheduler.LeaderElection.LeaseDuration + 90*time.Second + if floor > ttl { + ttl = floor + } + } + return ttl +} + func createBindingStore(logger *zap.Logger, cfg config.Config) (scheduler.BindingStore, func()) { + ttl := effectiveBindingTTL(cfg) + if ttl != cfg.Scheduler.BindingTTL { + logger.Info("binding TTL raised to cover the failover budget", + zap.Duration("configured", cfg.Scheduler.BindingTTL), + zap.Duration("effective", ttl), + ) + } if strings.TrimSpace(cfg.Scheduler.RedisAddr) == "" { - return scheduler.NewInMemoryBindingStore(cfg.Scheduler.BindingTTL), func() {} + return scheduler.NewInMemoryBindingStore(ttl), func() {} } - store, err := scheduler.NewRedisBindingStore(cfg.Scheduler.RedisAddr, cfg.Scheduler.BindingTTL) + store, err := scheduler.NewRedisBindingStore(cfg.Scheduler.RedisAddr, ttl) if err != nil { logger.Fatal("create redis binding store failed", zap.Error(err), zap.String("addr", cfg.Scheduler.RedisAddr)) } From 8a36b6c055fd7140d29484468b16b1de58857ec4 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 13:08:26 +0800 Subject: [PATCH 05/26] test(scheduler): spec Kind integration cases for HA failover Add build-tagged integration specs (#259) for the scenarios that need a real cluster and stay in CI (the repo has no Kind harness yet, and the dev machine is macOS): - C: exactly one leader; failover within the lease budget with seconds-level pull-driven rebuild; fencing with no dual-primary window (renewal-failure partition is covered offline by the fake-clientset tests). - E: the #191 resource-limit scenario passes with 3 replicas; observation coverage equals all nodes. - F: routing continuity across failover via standbys + Redis; degraded no-Redis reads retrying onto the leader; seconds-level recovery; binding TTL floor follow-up; rollback to single replica. Scope: services/scheduler/kind only. Relates to #259, #191. Co-Authored-By: Claude Code --- services/scheduler/kind/kind_test.go | 93 ++++++++++++++++++++++++++++ 1 file changed, 93 insertions(+) create mode 100644 services/scheduler/kind/kind_test.go diff --git a/services/scheduler/kind/kind_test.go b/services/scheduler/kind/kind_test.go new file mode 100644 index 000000000..8a70c02c3 --- /dev/null +++ b/services/scheduler/kind/kind_test.go @@ -0,0 +1,93 @@ +//go:build integration + +package kind + +import "testing" + +// Kind integration specs for scheduler HA (#259, test list groups C, E, F). +// +// These run in CI against a Kind cluster with: scheduler Deployment x3 +// (leader election on, redis_addr set), Redis, gateway, and at least 2 +// runtime nodes. They are spec'd here so each case maps to a reviewable +// assertion; bodies are filled when the CI harness lands. + +// Case C1: exactly one leader among N replicas; OnNewLeader is logged. +// +// setup: 3 replicas, empty Lease +// expect: exactly one pod readiness=SERVING; Service endpoints = 1; +// other pods liveness healthy +func TestC1ExactlyOneLeader(t *testing.T) { + t.Skip("requires Kind harness") +} + +// Case C2: failover within budget, rebuild in seconds (pull-driven). +// +// setup: C1 steady state, nodes heartbeating +// action: kill -9 the leader pod +// expect: new leader readiness=SERVING within lease_duration; +// sync-node-snapshots completes within seconds (not a reporter backoff); +// Service endpoints converge on the new leader +func TestC2FailoverWithinBudget(t *testing.T) { + t.Skip("requires Kind harness") +} + +// Case C3: fencing — a partitioned old leader stops serving within +// renew_deadline; no dual-primary window. +// +// action: network-partition the leader from the API server +// expect: old leader exits within renew_deadline; new leader takes over; +// at no point do two pods serve Schedule +func TestC3FencingNoDualPrimary(t *testing.T) { + t.Skip("requires Kind harness") +} + +// Case E1: the #191 resource-limit scenario passes with 3 replicas. +// +// setup: 3 replicas, max_sandbox_count=10 per node +// action: create sandboxes at full speed until well past nodes*10 +// expect: every node <= 10; Schedule returns Unavailable at the ceiling; +// zero overshoot (vs. the measured 3.0x in #191) +func TestE1ResourceLimitWithReplicas(t *testing.T) { + t.Skip("requires Kind harness") +} + +// Case E2: the leader's observed set equals all nodes (heartbeats converge +// on the single leader). +func TestE2ObservationCoverage(t *testing.T) { + t.Skip("requires Kind harness") +} + +// Case F1: routing continuity across failover (the issue's core assertion). +// +// setup: N sandboxes created, Redis-backed bindings +// action: kill the leader, keep calling LookupNode for old sandboxes +// (standbys serve from Redis) +// expect: correct node every time; zero NotFound, zero gap +func TestF1RoutingContinuity(t *testing.T) { + t.Skip("requires Kind harness") +} + +// Case F2: degraded mode without Redis — reads get Unavailable on standbys +// and client retries land on the leader; no error surfaces to the caller. +func TestF2NoRedisDegradedReads(t *testing.T) { + t.Skip("requires Kind harness") +} + +// Case F3: post-failover first successful read/schedule lands in seconds +// (pull-driven rebuild), not after a heartbeat backoff (~60s). +func TestF3SecondsLevelRecovery(t *testing.T) { + t.Skip("requires Kind harness") +} + +// Case F4 (follow-up): binding TTL floor — with bindings silent for 80s +// across a failover, bindings survive and F1 still holds. Enable once the +// floor (max(binding_ttl, lease_duration + 90s) under election) lands. +func TestF4BindingTTLFloor(t *testing.T) { + t.Skip("follow-up: TTL floor") +} + +// Case F5: rollback — disable election, scale back to 1 replica, behaviour +// returns to the single-replica shape. +func TestF5Rollback(t *testing.T) { + t.Skip("requires Kind harness") +} From 2ac0f4caec53488b036030c7c5f0bb1e71504f52 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 13:09:12 +0800 Subject: [PATCH 06/26] feat(deploy): add HA scheduler overlay and leader-election docs Kubernetes manifests and docs for scheduler HA (#259): - base role.yaml gains coordination.k8s.io/leases get/create/update; inert when election is off. - New deploy/k8s/overlays/ha: 3 leader-elected scheduler replicas, readiness probe targeting the leader health service (-service=scheduler.v1.Scheduler/leader) so Service endpoints hold only the leader, leader_election enabled in the scheduler config, and the SCHEDULER_NODE_ADMIN_API_KEY env mounted from the shared agentenv-auth Secret for sync-node-snapshots pulls. The overlay README covers topology, traffic rules, prerequisites, and the binary-first upgrade ordering with rollback. Base manifests are untouched; existing overlays are unaffected. - services/README.md gains a "Leader election (HA)" section: switch semantics, probe split, traffic rules, degraded mode without Redis, the binding TTL floor, and credential passing via env/Secret. Scope: deploy/k8s and services/README.md. Verified with kubectl kustomize (replicas, probe flag, leader_election config, leases RBAC, secret env all render). Relates to #259. Co-Authored-By: Claude Code --- deploy/k8s/base/role.yaml | 10 ++++ deploy/k8s/overlays/ha/README.md | 39 ++++++++++++++++ deploy/k8s/overlays/ha/config/scheduler.json | 25 ++++++++++ deploy/k8s/overlays/ha/kustomization.yaml | 49 ++++++++++++++++++++ services/README.md | 11 +++++ 5 files changed, 134 insertions(+) create mode 100644 deploy/k8s/overlays/ha/README.md create mode 100644 deploy/k8s/overlays/ha/config/scheduler.json create mode 100644 deploy/k8s/overlays/ha/kustomization.yaml diff --git a/deploy/k8s/base/role.yaml b/deploy/k8s/base/role.yaml index 4b6484f15..13d288568 100644 --- a/deploy/k8s/base/role.yaml +++ b/deploy/k8s/base/role.yaml @@ -19,3 +19,13 @@ rules: - get - list - watch + # Leader election (#259): the scheduler competes for a Lease when + # leader_election.enabled is set. Inert otherwise. + - apiGroups: + - coordination.k8s.io + resources: + - leases + verbs: + - get + - create + - update diff --git a/deploy/k8s/overlays/ha/README.md b/deploy/k8s/overlays/ha/README.md new file mode 100644 index 000000000..c8d81b899 --- /dev/null +++ b/deploy/k8s/overlays/ha/README.md @@ -0,0 +1,39 @@ +# HA scheduler overlay (#259) + +Three leader-elected scheduler replicas behind one Service, competing for a +`coordination.k8s.io/Lease` (`agentenv-system/agentenv-scheduler`). Exactly +one pod schedules and processes node reports; standbys stay liveness-healthy +and serve `LookupNode`/`GetNode` from the shared Redis bindings. + +Traffic rules: + +1. Node heartbeats → `agentenv-scheduler` Service (leader only). +2. Gateway writes (`Schedule`, `RecordAssignment`) → same Service. +3. Gateway reads (`LookupNode`) → same Service; served by standbys when + `scheduler.redis_addr` is set, otherwise retried onto the leader. + +Deploy: + +```bash +kubectl apply -k deploy/k8s/overlays/ha +``` + +Prerequisites: + +- RBAC: the scheduler Role needs `coordination.k8s.io/leases` + get/create/update (already in `base/role.yaml`). +- Redis (recommended): set `scheduler.redis_addr` in + `config/scheduler.json`. Without it the scheduler runs the documented + degraded mode: reads retry to the leader, and failover loses routing for + pre-failover sandboxes. +- `scheduler.node_admin_api_key`: the x-api-key for node admin APIs, used by + the leader to pull node snapshots right after winning (sync-node-snapshots). + +Upgrade ordering (rolling this out over a single-replica deployment): + +1. Upgrade the scheduler binary first — with election off it behaves exactly + like today, and it must know the leader health service before the probe + points at it. +2. Then apply this overlay (RBAC, probe, replicas, config). + +Rollback: set `leader_election.enabled` to false and scale back to 1 replica. diff --git a/deploy/k8s/overlays/ha/config/scheduler.json b/deploy/k8s/overlays/ha/config/scheduler.json new file mode 100644 index 000000000..45156e050 --- /dev/null +++ b/deploy/k8s/overlays/ha/config/scheduler.json @@ -0,0 +1,25 @@ +{ + "log_level": "info", + "log_format": "json", + "scheduler": { + "grpc_listen_addr": ":9090", + "strategy": "round_robin", + "discovery": { + "mode": "kubernetes", + "kubernetes": { + "namespace": "agentenv-system", + "service_name": "agentenv-nodes", + "port": 8000, + "scheme": "http" + } + }, + "leader_election": { + "enabled": true, + "lease_name": "agentenv-scheduler", + "lease_namespace": "agentenv-system", + "lease_duration": "15s", + "renew_deadline": "10s", + "retry_period": "2s" + } + } +} diff --git a/deploy/k8s/overlays/ha/kustomization.yaml b/deploy/k8s/overlays/ha/kustomization.yaml new file mode 100644 index 000000000..f7cdb9059 --- /dev/null +++ b/deploy/k8s/overlays/ha/kustomization.yaml @@ -0,0 +1,49 @@ +apiVersion: kustomize.config.k8s.io/v1beta1 +kind: Kustomization + +# HA scheduler topology (#259): 3 leader-elected scheduler replicas behind +# one Service. Only the leader is ready (readiness probes the leader-specific +# health service), so Service endpoints contain only the leader; standbys +# stay liveness-healthy and serve reads from the shared Redis bindings. +# +# Point scheduler.redis_addr at your Redis to enable standby reads and +# failover-preserved routing; without it the scheduler runs the documented +# degraded mode (reads retry to the leader; pre-failover routing is lost). + +resources: + - ../../base + +generatorOptions: + disableNameSuffixHash: true + +configMapGenerator: + - name: scheduler-k8s-config + behavior: replace + files: + - config/scheduler.json + +patches: + - target: + kind: Deployment + name: agentenv-scheduler + patch: |- + - op: replace + path: /spec/replicas + value: 3 + - op: replace + path: /spec/template/spec/containers/0/readinessProbe/exec/command + value: + - /grpc_health_probe + - -addr=127.0.0.1:9090 + - -service=scheduler.v1.Scheduler/leader + # sync-node-snapshots pulls authenticate against node admin APIs with + # the cluster-wide admin key; mount it from the shared agentenv-auth + # Secret, never from a ConfigMap. + - op: add + path: /spec/template/spec/containers/0/env + value: + - name: SCHEDULER_NODE_ADMIN_API_KEY + valueFrom: + secretKeyRef: + name: agentenv-auth + key: AENV_API_KEY diff --git a/services/README.md b/services/README.md index 59fe1194e..c5187df73 100644 --- a/services/README.md +++ b/services/README.md @@ -108,6 +108,17 @@ General config notes: - `SCHEDULER_ARTIFACT_STORE_CAPACITY=` overrides `scheduler.artifact_store_capacity` from the environment. - `SCHEDULER_ARTIFACT_LOOKUP_NODE_LIMIT=` overrides `scheduler.artifact_lookup_node_limit` from the environment. +### Leader election (HA) + +`scheduler.leader_election` enables Kubernetes Lease-based leader election (#259). Off by default; with it off, the scheduler behaves exactly as a single-writer process. + +- With it on, N replicas compete for a `coordination.k8s.io/Lease`. Exactly one leader schedules and processes heartbeats; standbys stay liveness-healthy, reject writes with `Unavailable`, and serve `LookupNode`/`GetNode` from the shared Redis bindings (reads are rejected too when no `redis_addr` is set, and clients retry onto the leader). +- Readiness probes should target the leader-specific health service `scheduler.v1.Scheduler/leader` (`grpc_health_probe -service=scheduler.v1.Scheduler/leader`) so Service endpoints contain only the leader. Liveness keeps probing the overall health status. +- Traffic rules: node heartbeats → scheduler Service (leader only); gateway writes (`Schedule`, `RecordAssignment`) → same Service; gateway reads (`LookupNode`) → same Service, answered by standbys when Redis is configured. +- On failover, the new leader pulls each node's admin `/nodes` snapshot (sync-node-snapshots) instead of waiting for the next heartbeat; pulls authenticate with the `x-api-key` from `SCHEDULER_NODE_ADMIN_API_KEY` — pass it via env/Secret (the HA overlay mounts the shared `agentenv-auth` Secret), not via config files. +- `redis_addr` is optional but recommended: without it, failover loses routing for pre-failover sandboxes (documented degraded mode). Under election the binding TTL is floored to `lease_duration + 90s` so bindings outlive the failover budget. +- Mutually exclusive with `--query-only`. Ready-made manifests: `deploy/k8s/overlays/ha` (see its README for upgrade ordering). + ### Scheduling strategy `scheduler.strategy` selects the algorithm used to pick a node from the eligible candidate list. Built-in strategies: From d089509e045e547ed42421e1d365454682a193c1 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 14:49:55 +0800 Subject: [PATCH 07/26] fix(scheduler): address Copilot review on HA failover correctness MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four correctness fixes from the PR #341 review: - Gate GetNode on standbys (F2): it reads replica-local observations, which are always empty on standbys (heartbeats and pulls only reach the leader), so serving it there returned a terminal NotFound for real nodes. LookupNode stays standby-readable; it reads the shared binding store. - Hold the lease until process exit (F3): ReleaseOnCancel=true let a standby acquire while the ex-leader was still draining in-flight RPCs through GracefulStop, opening the dual-primary window the fencing design exists to prevent. Failover after a graceful stop is now bounded by the remaining lease duration instead. - Never reconcile bindings from a snapshot pull (F4): the roster is now three-state — non-nil (heartbeat) is authoritative and reconciles, nil (pull) leaves bindings untouched. A pull cannot see paused sandboxes or in-flight template builds (list_sandbox_ids covers both; no admin list endpoint does), so reconciling from it deleted live bindings exactly on failover. The admin fetcher drops the GET /sandboxes call; binding refresh stays with heartbeats, and the binding TTL floor covers the gap. Includes a regression test that a pre-existing binding survives a roster-less pull ingest. - Reject explicitly invalid election config values (F5): defaults now apply only to omitted fields — an explicit non-positive duration (e.g. "retry_period": "-1s") or snapshot_pull_concurrency fails at config parse instead of being silently replaced by the default, keeping the advertised fail-fast behavior. Scope: services/scheduler/internal, services/shared/config. The F1 finding (standby reads unreachable behind leader-only readiness) is a design fork rooted in two conflicting requirements and is answered in the PR discussion rather than changed here. Relates to #259. Co-Authored-By: Claude Code --- services/scheduler/internal/leadership.go | 15 ++++--- .../scheduler/internal/leadership_test.go | 6 ++- services/scheduler/internal/node_snapshot.go | 42 ++++++++----------- .../scheduler/internal/node_snapshot_test.go | 13 +++--- .../internal/recovery_window_test.go | 26 ++++++++++++ services/scheduler/internal/service.go | 26 ++++++++---- services/shared/config/config.go | 8 ++++ 7 files changed, 92 insertions(+), 44 deletions(-) diff --git a/services/scheduler/internal/leadership.go b/services/scheduler/internal/leadership.go index 9258e6503..fe0987356 100644 --- a/services/scheduler/internal/leadership.go +++ b/services/scheduler/internal/leadership.go @@ -225,9 +225,9 @@ func (l *leadershipManager) LeaderSince() (time.Time, bool) { // // Notes on the reads: // - LookupNode reads bindings from the shared store: correct on standbys. -// - GetNode reads replica-local observations, which are empty on standbys -// until heartbeats/pulls arrive; classified readable per the design -// matrix anyway (a miss degrades to NotFound, not to wrong data). +// - GetNode reads replica-local observations, which are always empty on +// standbys (heartbeats/pulls only reach the leader), so serving it there +// would return a terminal NotFound for real nodes. Gated (#341 review). // - ListNodes reads the informer-backed registry (warm on every replica) // and is semantically standby-serveable, but stays gated for now — // deferred to a PR decision (#259 review). @@ -243,7 +243,7 @@ var standbyReadableMethods = map[string]bool{ "/scheduler.v1.Scheduler/RecordP2pArtifact": false, "/scheduler.v1.Scheduler/ForgetP2pArtifact": false, "/scheduler.v1.Scheduler/LookupP2pArtifact": false, - "/scheduler.v1.Scheduler/GetNode": true, + "/scheduler.v1.Scheduler/GetNode": false, "/scheduler.v1.Scheduler/UnregisterNode": false, } @@ -442,8 +442,11 @@ func (r *leaderElector) Run(ctx context.Context) error { r.logger.Info("scheduler leader changed", zap.String("leader", identity)) }, }, - // Release the lease on graceful shutdown so failover is fast. - ReleaseOnCancel: true, + // Hold the lease until process exit: releasing on cancel would let + // a standby acquire while this ex-leader is still draining in-flight + // RPCs, opening a dual-primary window (#341 review). Failover after a + // graceful stop is bounded by the remaining lease duration. + ReleaseOnCancel: false, }) if err != nil { return fmt.Errorf("leader election config: %w", err) diff --git a/services/scheduler/internal/leadership_test.go b/services/scheduler/internal/leadership_test.go index f14858b96..4bfcd20cc 100644 --- a/services/scheduler/internal/leadership_test.go +++ b/services/scheduler/internal/leadership_test.go @@ -190,9 +190,11 @@ func TestGateStandbyWithStoreServesReads(t *testing.T) { t.Fatalf("standby with a shared store must serve LookupNode: called=%v err=%v", called, err) } + // GetNode reads replica-local observations, which are always empty on a + // standby — it stays gated even with a shared store (#341 review, F2). called, err = invokeGate(t, false, true, "/scheduler.v1.Scheduler/GetNode") - if !called || err != nil { - t.Fatalf("standby with a shared store must serve GetNode: called=%v err=%v", called, err) + if called || status.Code(err) != codes.Unavailable { + t.Fatalf("GetNode must stay gated on standbys (observations are leader-local): called=%v err=%v", called, err) } called, err = invokeGate(t, false, true, "/scheduler.v1.Scheduler/RecordAssignment") diff --git a/services/scheduler/internal/node_snapshot.go b/services/scheduler/internal/node_snapshot.go index 35111ba56..d3560939d 100644 --- a/services/scheduler/internal/node_snapshot.go +++ b/services/scheduler/internal/node_snapshot.go @@ -38,8 +38,17 @@ type nodeReport struct { commit string machineInfo *schedulerv1.MachineInfo snapshot *schedulerv1.NodeSnapshot - sandboxIDs []string - p2pEndpoint *schedulerv1.P2PEndpoint + // sandboxIDs carries the node's authoritative sandbox roster. Three + // states matter: + // non-nil (including empty) — authoritative (heartbeat): reconcile + // bindings to exactly this set; + // nil — unknown (snapshot pull): bindings are left untouched. A pull + // cannot see paused sandboxes or in-flight template builds, so + // reconciling from it would delete live bindings on failover + // (#341 review). Binding refresh stays with heartbeats, and the + // binding TTL floor covers the gap. + sandboxIDs []string + p2pEndpoint *schedulerv1.P2PEndpoint } // toHeartbeatRequest converts the report to the registry's proto shape; the @@ -166,16 +175,10 @@ type adminNodeResponse struct { } `json:"metrics"` } -// adminSandboxEntry mirrors a ListedSandbox entry from GET /sandboxes; only -// the id is needed for binding reconciliation. -type adminSandboxEntry struct { - SandboxID string `json:"sandboxID"` -} - // NewAdminSnapshotFetcher builds the production NodeSnapshotFetcher (#259, // sync-node-snapshots): it pulls GET {endpoint}/nodes for observations and -// GET {endpoint}/sandboxes for the sandbox id list, and returns the fetched -// data as a nodeReport for the shared ingest path. +// returns the fetched data as a nodeReport for the shared ingest path. +// Bindings are deliberately not pulled (see nodeReport.sandboxIDs). // The node's admin API requires the x-api-key header. The HTTP client is an // implementation detail with a bounded per-request timeout; tests inject // through the NodeSnapshotFetcher seam, not this constructor. @@ -202,12 +205,10 @@ func NewAdminSnapshotFetcher(apiKey string) NodeSnapshotFetcher { return nil, fmt.Errorf("pull %s /nodes: node not in admin response", node.ID) } - var sandboxes []adminSandboxEntry - if err := adminGetJSON(ctx, client, apiKey, base+"/sandboxes", &sandboxes); err != nil { - return nil, fmt.Errorf("pull %s /sandboxes: %w", node.ID, err) - } - - return reportFromAdmin(entry, sandboxes), nil + // Only /nodes is pulled: observations, not the sandbox roster. See + // nodeReport.sandboxIDs for why the roster is never reconciled from + // a pull. + return reportFromAdmin(entry), nil } } @@ -230,13 +231,7 @@ func adminGetJSON(ctx context.Context, client *http.Client, apiKey, url string, return json.NewDecoder(resp.Body).Decode(out) } -func reportFromAdmin(n *adminNodeResponse, sandboxes []adminSandboxEntry) *nodeReport { - ids := make([]string, 0, len(sandboxes)) - for _, s := range sandboxes { - if strings.TrimSpace(s.SandboxID) != "" { - ids = append(ids, s.SandboxID) - } - } +func reportFromAdmin(n *adminNodeResponse) *nodeReport { disks := make([]*schedulerv1.DiskMetric, 0, len(n.Metrics.Disks)) for _, d := range n.Metrics.Disks { disks = append(disks, &schedulerv1.DiskMetric{ @@ -277,6 +272,5 @@ func reportFromAdmin(n *adminNodeResponse, sandboxes []adminSandboxEntry) *nodeR PausedAllocatedCpu: n.Metrics.PausedAllocatedCPU, PausedAllocatedMemoryBytes: n.Metrics.PausedAllocatedMemoryBytes, }, - sandboxIDs: ids, } } diff --git a/services/scheduler/internal/node_snapshot_test.go b/services/scheduler/internal/node_snapshot_test.go index a2c183918..72233a8c0 100644 --- a/services/scheduler/internal/node_snapshot_test.go +++ b/services/scheduler/internal/node_snapshot_test.go @@ -88,9 +88,10 @@ func TestAdminSnapshotFetcherAssemblesHeartbeatShape(t *testing.T) { if len(snap.GetDisks()) != 1 || snap.GetDisks()[0].GetDevice() != "/dev/ublkb0" { t.Fatalf("disk mapping wrong: %+v", snap.GetDisks()) } - ids := req.sandboxIDs - if len(ids) != 2 || ids[0] != "sbx-1" || ids[1] != "sbx-2" { - t.Fatalf("sandbox ids wrong: %v", ids) + // A pull must not carry a sandbox roster: reconciling from it would + // delete live bindings (paused sandboxes, template builds) on failover. + if req.sandboxIDs != nil { + t.Fatalf("pull must leave sandboxIDs nil, got %v", req.sandboxIDs) } } @@ -118,7 +119,9 @@ func TestAdminSnapshotFetcherHeartbeatIngestCompatibility(t *testing.T) { if registry.PeekObserved("node-a") == nil { t.Fatal("pulled snapshot must record an observation") } - if _, ok, err := store.Get("sbx-1", time.Now()); err != nil || !ok { - t.Fatalf("pulled snapshot must reconcile bindings: ok=%v err=%v", ok, err) + // The pull refreshes observations only; bindings are refreshed by + // heartbeats, so nothing may be written to the store here. + if _, ok, err := store.Get("sbx-1", time.Now()); err != nil || ok { + t.Fatalf("pull-ingest must not touch bindings: ok=%v err=%v", ok, err) } } diff --git a/services/scheduler/internal/recovery_window_test.go b/services/scheduler/internal/recovery_window_test.go index f8371f0ba..a64c63463 100644 --- a/services/scheduler/internal/recovery_window_test.go +++ b/services/scheduler/internal/recovery_window_test.go @@ -224,3 +224,29 @@ func TestLookupNodeWindowSemantics(t *testing.T) { t.Fatalf("window closed: GetNode want NotFound, got %v", err) } } + +// Case F4 regression (#341 review): a snapshot pull carries no roster, so +// pulling must never delete existing bindings — including bindings for +// paused sandboxes that no list endpoint would report. +func TestPullIngestPreservesBindingsWithoutRoster(t *testing.T) { + svc, registry, store := newTestService(t, []string{"node-a"}) + + // A binding exists, e.g. for a paused sandbox. + if err := store.Record("sbx-paused", Node{ID: "node-a", Endpoint: "http://node-a:8080"}, time.Now()); err != nil { + t.Fatalf("seed binding failed: %v", err) + } + + // A pull-shaped report: observations present, roster unknown (nil). + pulled := &nodeReport{nodeID: "node-a", serviceInstanceID: "inst-a"} + if _, err := svc.ingestNodeReport(pulled, time.Now()); err != nil { + t.Fatalf("pull ingest failed: %v", err) + } + + if registry.PeekObserved("node-a") == nil { + t.Fatal("pull must still record the observation") + } + node, ok, err := store.Get("sbx-paused", time.Now()) + if err != nil || !ok || node.ID != "node-a" { + t.Fatalf("pull without a roster must preserve bindings: ok=%v node=%v err=%v", ok, node, err) + } +} diff --git a/services/scheduler/internal/service.go b/services/scheduler/internal/service.go index 050c6c4d1..165ace6fe 100644 --- a/services/scheduler/internal/service.go +++ b/services/scheduler/internal/service.go @@ -275,6 +275,13 @@ func (s *Service) Heartbeat(_ context.Context, req *schedulerv1.HeartbeatRequest // nodeReportFromProto adapts a Heartbeat RPC request to the neutral ingest // currency. Identity validation already happened in Heartbeat. func nodeReportFromProto(req *schedulerv1.HeartbeatRequest) *nodeReport { + // A heartbeat roster is always authoritative, including when empty: + // proto3 decodes absent and empty repeated fields identically, so force + // a non-nil slice to keep heartbeats reconciling (see nodeReport). + sandboxIDs := req.GetSandboxIds() + if sandboxIDs == nil { + sandboxIDs = []string{} + } return &nodeReport{ nodeID: req.GetNodeId(), clusterID: req.GetClusterId(), @@ -283,7 +290,7 @@ func nodeReportFromProto(req *schedulerv1.HeartbeatRequest) *nodeReport { commit: req.GetCommit(), machineInfo: req.GetMachineInfo(), snapshot: req.GetSnapshot(), - sandboxIDs: req.GetSandboxIds(), + sandboxIDs: sandboxIDs, p2pEndpoint: req.GetP2PEndpoint(), } } @@ -302,12 +309,17 @@ func (s *Service) ingestNodeReport(report *nodeReport, now time.Time) (*schedule } return nil, status.Error(codes.Internal, "node registry heartbeat failed") } - if err := s.store.ReconcileNode(node, report.sandboxIDs, now); err != nil { - s.logger.Warn("scheduler heartbeat binding reconcile failed", - zap.String("node_id", report.nodeID), - zap.Error(err), - ) - return nil, status.Error(codes.Unavailable, "binding store unavailable") + // Reconcile only from an authoritative roster (heartbeat). A snapshot + // pull leaves sandboxIDs nil — bindings are refreshed by heartbeats and + // guarded meanwhile by the binding TTL floor (#341 review). + if report.sandboxIDs != nil { + if err := s.store.ReconcileNode(node, report.sandboxIDs, now); err != nil { + s.logger.Warn("scheduler heartbeat binding reconcile failed", + zap.String("node_id", report.nodeID), + zap.Error(err), + ) + return nil, status.Error(codes.Unavailable, "binding store unavailable") + } } return &schedulerv1.HeartbeatResponse{CpuConfigJson: cpuConfigJSON}, nil } diff --git a/services/shared/config/config.go b/services/shared/config/config.go index 7a731b5f0..cb4f2163e 100644 --- a/services/shared/config/config.go +++ b/services/shared/config/config.go @@ -175,6 +175,9 @@ func (s *SchedulerConfig) UnmarshalJSON(data []byte) error { s.LeaderElection.LeaseNamespace = *le.LeaseNamespace } if le.SnapshotPullConcurrency != nil { + if *le.SnapshotPullConcurrency <= 0 { + return fmt.Errorf("scheduler.leader_election.snapshot_pull_concurrency must be greater than zero, got %d", *le.SnapshotPullConcurrency) + } s.LeaderElection.SnapshotPullConcurrency = *le.SnapshotPullConcurrency } durationFields := []struct { @@ -194,6 +197,11 @@ func (s *SchedulerConfig) UnmarshalJSON(data []byte) error { if err != nil { return err } + // Defaults apply to omitted fields only; an explicitly provided + // non-positive duration is invalid, not a request for the default. + if d <= 0 { + return fmt.Errorf("%s must be greater than zero, got %q", f.field, d.String()) + } *f.dst = d } } From 2b15c2cfab484a4225b24b4aac41a2716a45f736 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 15:06:08 +0800 Subject: [PATCH 08/26] fix(scheduler): keep slow snapshot pulls from overwriting newer heartbeats MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a freshness guard to the post-acquisition refresh (#341 review, OpenCodeReview): a nodeReport from a pull now carries the capture time (fetchedAt), and ingest skips a pulled report when the node has already reported something newer — e.g. a heartbeat that landed while the HTTP request was in flight. Heartbeat-shaped reports carry a zero fetchedAt and are never skipped. Without the guard, a slow pull could regress observations by up to one request latency during the rebuild window. The remaining pull/heartbeat parity gap (admin /nodes lacks sandboxIDs, p2pEndpoint, and cpuConfigJSON) is bounded and self-healing via the next heartbeat; it is documented in the PR discussion as follow-up rather than changed here. Scope: services/scheduler/internal. Includes guard-semantics tests. Relates to #259. Co-Authored-By: Claude Code --- services/scheduler/internal/leadership.go | 5 ++++ services/scheduler/internal/node_snapshot.go | 19 ++++++++++++ .../scheduler/internal/node_snapshot_test.go | 30 +++++++++++++++++++ 3 files changed, 54 insertions(+) diff --git a/services/scheduler/internal/leadership.go b/services/scheduler/internal/leadership.go index fe0987356..0be33a5cf 100644 --- a/services/scheduler/internal/leadership.go +++ b/services/scheduler/internal/leadership.go @@ -356,6 +356,11 @@ func (l *leadershipManager) onStartedLeading(ctx context.Context) { return } ingestor := func(report *nodeReport, now time.Time) error { + // Freshness guard (#341 review): a slow pull must not overwrite a + // heartbeat that arrived while the request was in flight. + if skipStalePull(l.registry, report) { + return nil + } _, err := l.svc.ingestNodeReport(report, now) return err } diff --git a/services/scheduler/internal/node_snapshot.go b/services/scheduler/internal/node_snapshot.go index d3560939d..298a92bf1 100644 --- a/services/scheduler/internal/node_snapshot.go +++ b/services/scheduler/internal/node_snapshot.go @@ -49,6 +49,12 @@ type nodeReport struct { // binding TTL floor covers the gap. sandboxIDs []string p2pEndpoint *schedulerv1.P2PEndpoint + // fetchedAt records when a pulled report's data was captured. A report + // built from a heartbeat leaves it zero (always fresh). The ingest guard + // uses it to keep a slow pull from overwriting a newer heartbeat + // (#341 review): skip the pull when the node has already reported + // something captured after this data. + fetchedAt time.Time } // toHeartbeatRequest converts the report to the registry's proto shape; the @@ -84,6 +90,18 @@ type NodeSnapshotFetcher func(ctx context.Context, node Node) (*nodeReport, erro // shared ingest path. The retriever never holds the Service itself. type nodeReportIngester func(report *nodeReport, now time.Time) error +// skipStalePull reports whether a pulled report should be dropped because +// the node has already reported something newer — a heartbeat that arrived +// while the pull was in flight (#341 review). Heartbeat-shaped reports +// (zero fetchedAt) are never skipped. +func skipStalePull(registry NodeRegistry, report *nodeReport) bool { + if report.fetchedAt.IsZero() { + return false + } + at, ok := registry.LastReportAt(report.nodeID) + return ok && at.After(report.fetchedAt) +} + type ConcurrentNodeSnapshotRefresher struct { logger *zap.Logger ingest nodeReportIngester @@ -248,6 +266,7 @@ func reportFromAdmin(n *adminNodeResponse) *nodeReport { serviceInstanceID: n.ServiceInstanceID, version: n.Version, commit: n.Commit, + fetchedAt: time.Now(), machineInfo: &schedulerv1.MachineInfo{ CpuFamily: n.MachineInfo.CPUFamily, CpuModel: n.MachineInfo.CPUModel, diff --git a/services/scheduler/internal/node_snapshot_test.go b/services/scheduler/internal/node_snapshot_test.go index 72233a8c0..a9c1bd0a8 100644 --- a/services/scheduler/internal/node_snapshot_test.go +++ b/services/scheduler/internal/node_snapshot_test.go @@ -125,3 +125,33 @@ func TestAdminSnapshotFetcherHeartbeatIngestCompatibility(t *testing.T) { t.Fatalf("pull-ingest must not touch bindings: ok=%v err=%v", ok, err) } } + +// Freshness guard (#341 review, N1): a pull captured before a heartbeat +// must not overwrite it; a pull captured after is applied; heartbeat-shaped +// reports are never skipped. +func TestSkipStalePull(t *testing.T) { + _, registry, _ := newTestService(t, []string{"node-a"}) + t0 := time.Date(2026, 10, 5, 12, 0, 0, 0, time.UTC) + + // A heartbeat lands at t0+1s. + if _, _, err := registry.Heartbeat(heartbeatFor("node-a", "inst-1"), t0.Add(time.Second)); err != nil { + t.Fatalf("heartbeat failed: %v", err) + } + + stale := &nodeReport{nodeID: "node-a", fetchedAt: t0} + if !skipStalePull(registry, stale) { + t.Fatal("a pull captured before the latest heartbeat must be skipped") + } + fresh := &nodeReport{nodeID: "node-a", fetchedAt: t0.Add(2 * time.Second)} + if skipStalePull(registry, fresh) { + t.Fatal("a pull captured after the latest heartbeat must be applied") + } + heartbeatShaped := &nodeReport{nodeID: "node-a"} + if skipStalePull(registry, heartbeatShaped) { + t.Fatal("heartbeat-shaped reports (zero fetchedAt) must never be skipped") + } + unknown := &nodeReport{nodeID: "node-b", fetchedAt: t0} + if skipStalePull(registry, unknown) { + t.Fatal("an unobserved node must not be skipped") + } +} From d4c18afb5af270f3d52738c0462c39ca01e266b9 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 15:19:08 +0800 Subject: [PATCH 09/26] fix(scheduler): keep known-but-unobserved nodes excluded past the recovery TTL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address the remaining Copilot finding (#341 review, service.go): the recovery window ended on a timer (one report_ttl), after which a node whose pull failed and whose reporter is still backing off (up to 60s) was readmitted with a nil snapshot — FilterByResourceLimit's fail-open then kept it, reintroducing the #191 capacity bug during failover. The acquisition snapshot now carries the node set known at that moment. Scheduling excludes a known node until it produces a report at or after the acquisition — no timer readmission — while nodes that joined later are outside the recovery-pending set and keep the steady-state fail-open behavior. The recovery TTL still bounds only the lookup semantics (Unavailable vs NotFound), unchanged. leadershipView gains RecoveryPending; the manager delegates; the null object returns false. Includes a regression test: past-TTL exclusion of an unreported known node, fail-open admission of a later-joined node, and readmission after the late report. Relates to #259, #191. Co-Authored-By: Claude Code --- services/scheduler/internal/leadership.go | 59 +++++++++++++++---- .../internal/recovery_window_test.go | 45 +++++++++++++- services/scheduler/internal/service.go | 24 +++++--- 3 files changed, 110 insertions(+), 18 deletions(-) diff --git a/services/scheduler/internal/leadership.go b/services/scheduler/internal/leadership.go index 0be33a5cf..4e931747e 100644 --- a/services/scheduler/internal/leadership.go +++ b/services/scheduler/internal/leadership.go @@ -30,6 +30,9 @@ import ( type leadershipView interface { InRecoveryWindow(now time.Time) bool LeaderSince() (time.Time, bool) + // RecoveryPending reports whether nodeID was known at acquisition and + // has not yet produced a post-acquisition report. + RecoveryPending(nodeID string, lastReportAt time.Time, ok bool) bool } // Leadership is the single external contract for the scheduler's @@ -74,12 +77,13 @@ func NewLeadership(logger *zap.Logger, cfg config.SchedulerConfig) Leadership { // a single-writer scheduler without any conditional wiring. type nonLeadership struct{} -func (nonLeadership) InRecoveryWindow(time.Time) bool { return false } -func (nonLeadership) LeaderSince() (time.Time, bool) { return time.Time{}, false } -func (nonLeadership) ServiceOption() ServiceOption { return func(*Service) {} } -func (nonLeadership) RegisterHealth(*health.Server) {} -func (nonLeadership) BindRuntime(*Service, NodeRegistry) {} -func (nonLeadership) MarkNotServing() {} +func (nonLeadership) InRecoveryWindow(time.Time) bool { return false } +func (nonLeadership) LeaderSince() (time.Time, bool) { return time.Time{}, false } +func (nonLeadership) RecoveryPending(string, time.Time, bool) bool { return false } +func (nonLeadership) ServiceOption() ServiceOption { return func(*Service) {} } +func (nonLeadership) RegisterHealth(*health.Server) {} +func (nonLeadership) BindRuntime(*Service, NodeRegistry) {} +func (nonLeadership) MarkNotServing() {} func (nonLeadership) GateInterceptor() grpc.UnaryServerInterceptor { return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { return handler(ctx, req) @@ -105,6 +109,12 @@ type leadershipSnapshot struct { leader bool since time.Time recoveryTTL time.Duration + // knownNodes is the node set known at acquisition. A known node with no + // post-acquisition report stays excluded from scheduling until it + // reports — a timer alone must not readmit it while its reporter is + // still backing off (#341 review). Nodes that join later are not in the + // set and keep the steady-state fail-open behavior. + knownNodes map[string]bool } // newLeadership is the not-leading state. recoveryTTL bounds the recovery @@ -115,9 +125,24 @@ func newLeadershipSnapshot(recoveryTTL time.Duration) *leadershipSnapshot { return &leadershipSnapshot{recoveryTTL: recoveryTTL} } -// acquiredLeadership is the leading state, acquired at now. -func acquiredLeadershipSnapshot(recoveryTTL time.Duration, now time.Time) *leadershipSnapshot { - return &leadershipSnapshot{leader: true, since: now, recoveryTTL: recoveryTTL} +// acquiredLeadershipSnapshot is the leading state, acquired at now, with the +// node set known at that moment (the recovery-pending set). +func acquiredLeadershipSnapshot(recoveryTTL time.Duration, now time.Time, knownNodes ...string) *leadershipSnapshot { + known := make(map[string]bool, len(knownNodes)) + for _, id := range knownNodes { + known[id] = true + } + return &leadershipSnapshot{leader: true, since: now, recoveryTTL: recoveryTTL, knownNodes: known} +} + +// recoveryPending reports whether nodeID was known at acquisition and has +// not yet produced a post-acquisition report (lastReportAt is the node's +// latest report time, ok=false when it never reported). +func (l *leadershipSnapshot) RecoveryPending(nodeID string, lastReportAt time.Time, ok bool) bool { + if !l.leader || !l.knownNodes[nodeID] { + return false + } + return !ok || lastReportAt.Before(l.since) } // InRecoveryWindow reports whether the recovery window following leadership @@ -217,6 +242,11 @@ func (l *leadershipManager) LeaderSince() (time.Time, bool) { return l.state.Load().LeaderSince() } +// RecoveryPending delegates to the internal leadership state. +func (l *leadershipManager) RecoveryPending(nodeID string, lastReportAt time.Time, ok bool) bool { + return l.state.Load().RecoveryPending(nodeID, lastReportAt, ok) +} + // standbyReadableMethods classifies every Scheduler RPC: true = a non-leader // replica sharing the binding store (Redis) may serve it; false = gated // behind leadership. The table lists all methods explicitly so a proto @@ -347,7 +377,16 @@ func (l *leadershipManager) Run(ctx context.Context, onStop func()) error { // bindings by pulling every node's admin snapshot instead of waiting for the // next reporter backoff (#259). func (l *leadershipManager) onStartedLeading(ctx context.Context) { - l.state.Store(acquiredLeadershipSnapshot(l.cfg.ReportTTL, time.Now())) + // Capture the acquisition-time node set for recovery-pending semantics; + // tolerate being called before BindRuntime (tests, construction-order + // guard). + known := make([]string, 0) + if l.registry != nil { + for _, n := range l.registry.Snapshot(false) { + known = append(known, n.ID) + } + } + l.state.Store(acquiredLeadershipSnapshot(l.cfg.ReportTTL, time.Now(), known...)) if l.health != nil { l.health.SetServingStatus(LeaderHealthService, grpc_health_v1.HealthCheckResponse_SERVING) } diff --git a/services/scheduler/internal/recovery_window_test.go b/services/scheduler/internal/recovery_window_test.go index a64c63463..2e1b659c2 100644 --- a/services/scheduler/internal/recovery_window_test.go +++ b/services/scheduler/internal/recovery_window_test.go @@ -59,7 +59,7 @@ func newLeaderTestService(t *testing.T, nodeIDs []string, recoveryTTL time.Durat store := NewInMemoryBindingStore(30 * time.Second) clock := &fixedClock{now: at} svc := NewService(nil, registry, NewStrategy("round_robin"), store, - WithLeadership(acquiredLeadershipSnapshot(recoveryTTL, acquireAt)), + WithLeadership(acquiredLeadershipSnapshot(recoveryTTL, acquireAt, nodeIDs...)), WithClock(clock.Now), ) return svc, registry, clock @@ -250,3 +250,46 @@ func TestPullIngestPreservesBindingsWithoutRoster(t *testing.T) { t.Fatalf("pull without a roster must preserve bindings: ok=%v node=%v err=%v", ok, node, err) } } + +// Case: #341 review (Copilot, service.go) — a node known at acquisition +// whose pull failed and whose reporter is still backing off must stay +// excluded after the recovery TTL expires; a node that joined after the +// acquisition keeps the steady-state fail-open admission. +func TestKnownUnobservedNodeStaysExcludedPastRecoveryTTL(t *testing.T) { + t0 := time.Date(2026, 10, 5, 12, 0, 0, 0, time.UTC) + svc, registry, clock := newLeaderTestService(t, []string{"node-old"}, 30*time.Second, t0.Add(2*time.Second), t0) + + // node-old was known at acquisition but never reports (pull failed, + // reporter backing off up to 60s). node-new is registered later, so it + // is not in the acquisition-time recovery-pending set. + registry.Set([]Node{ + {ID: "node-old", Endpoint: "http://node-old:8080"}, + {ID: "node-new", Endpoint: "http://node-new:8080"}, + }, nil) + + // Well past the recovery TTL: the timer alone must not readmit node-old. + clock.now = t0.Add(90 * time.Second) + resp, err := svc.Schedule(context.Background(), &schedulerv1.ScheduleRequest{}) + if err != nil { + t.Fatalf("schedule failed: %v", err) + } + if got := resp.GetNode().GetNodeId(); got != "node-new" { + t.Fatalf("past the TTL, node-old must still be excluded; got %q", got) + } + + // Once node-old finally reports, it becomes a candidate again. + if _, err := svc.Heartbeat(context.Background(), heartbeatFor("node-old", "inst-o")); err != nil { + t.Fatalf("late heartbeat failed: %v", err) + } + seen := map[string]bool{} + for i := 0; i < 2; i++ { + resp, err := svc.Schedule(context.Background(), &schedulerv1.ScheduleRequest{}) + if err != nil { + t.Fatalf("schedule %d failed: %v", i, err) + } + seen[resp.GetNode().GetNodeId()] = true + } + if !seen["node-old"] { + t.Fatalf("after its late report, node-old must be schedulable; seen %v", seen) + } +} diff --git a/services/scheduler/internal/service.go b/services/scheduler/internal/service.go index 165ace6fe..fbcf7c6ef 100644 --- a/services/scheduler/internal/service.go +++ b/services/scheduler/internal/service.go @@ -149,18 +149,28 @@ func (s *Service) Schedule(_ context.Context, req *schedulerv1.ScheduleRequest) } // filterFreshObservations applies the fresh-observations-only rule (#259): -// while the recovery window is open, only nodes observed at or after the -// leadership acquisition are scheduling candidates. Outside the window (or -// with leader election disabled) all nodes pass, preserving today's -// fail-open behaviour for new nodes. +// a node that was known at leadership acquisition is a scheduling candidate +// only once it has reported at or after the acquisition — a timer alone +// must not readmit a node whose reporter is still backing off (#341 +// review). Nodes that joined after the acquisition are not in the +// recovery-pending set and keep the steady-state fail-open behaviour; with +// leader election disabled everything passes. func (s *Service) filterFreshObservations(rich []RichNode) []RichNode { - if !s.inRecoveryWindow() { + if s.leadership == nil { + return rich + } + since, ok := s.leadership.LeaderSince() + if !ok { return rich } - since, _ := s.leadership.LeaderSince() fresh := make([]RichNode, 0, len(rich)) for _, n := range rich { - if at, ok := s.nodes.LastReportAt(n.Node.ID); ok && !at.Before(since) { + at, reported := s.nodes.LastReportAt(n.Node.ID) + if reported && !at.Before(since) { + fresh = append(fresh, n) + continue + } + if !s.leadership.RecoveryPending(n.Node.ID, at, reported) { fresh = append(fresh, n) } } From c9ea0fc419a394c3e28dd06781d20bda57d52c8a Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 20:17:36 +0800 Subject: [PATCH 10/26] fix(scheduler): address the second OpenCodeReview round on snapshot pulls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Seven correctness and hardening fixes from the follow-up review (#341): - Freshness is now enforced atomically inside the registry: the new HeartbeatUnlessStale checks and writes under the same lock, closing the TOCTOU window between a pull's freshness check and its ingest. The manager's pre-check is removed; the registry is the single point of truth. Includes an atomic-path test. - fetchedAt is stamped before the /nodes request is sent, not after the response is decoded, so a heartbeat landing mid-flight correctly counts as newer than the pulled data. - Pulled reports map the admin status field into NodeSnapshot.Status instead of leaving UNSPECIFIED (which the registry derives to CONNECTING), and the fetcher no longer accepts a single-entry response whose ID differs from the requested node — a stale or misconfigured endpoint cannot overwrite another node's observation. - Admin responses are decoded with a 1 MiB body limit, so a faulty or compromised node cannot make the leader allocate from an arbitrarily large document during the acquisition fan-out. - The registry preserves a previously known P2P endpoint when a report arrives without one (the admin API does not serve it), instead of erasing it on every pull. - Freshness comparisons run at millisecond precision in both filterFreshObservations and RecoveryPending: LastReportAt is reconstructed from a millisecond timestamp, so a report in the same millisecond as the acquisition counts as fresh instead of stale. The remaining heartbeat-parity gaps in admin /nodes (sandboxIDs, p2pEndpoint, cpuConfigJSON) stay on the documented follow-up track; the cpuConfigJSON finding was already dispositioned in the previous round. Scope: services/scheduler/internal. Relates to #259. Co-Authored-By: Claude Code --- services/scheduler/internal/leadership.go | 13 ++--- services/scheduler/internal/node_registry.go | 41 +++++++++++++- services/scheduler/internal/node_snapshot.go | 54 ++++++++++++------- .../scheduler/internal/node_snapshot_test.go | 44 +++++++++------ services/scheduler/internal/service.go | 45 ++++++++++++---- 5 files changed, 145 insertions(+), 52 deletions(-) diff --git a/services/scheduler/internal/leadership.go b/services/scheduler/internal/leadership.go index 4e931747e..176783dd4 100644 --- a/services/scheduler/internal/leadership.go +++ b/services/scheduler/internal/leadership.go @@ -142,7 +142,10 @@ func (l *leadershipSnapshot) RecoveryPending(nodeID string, lastReportAt time.Ti if !l.leader || !l.knownNodes[nodeID] { return false } - return !ok || lastReportAt.Before(l.since) + // Millisecond-precision comparison, matching the scheduler's freshness + // rule (#341 review): a report in the same millisecond as the + // acquisition counts as post-acquisition. + return !ok || lastReportAt.UnixMilli() < l.since.UnixMilli() } // InRecoveryWindow reports whether the recovery window following leadership @@ -395,11 +398,9 @@ func (l *leadershipManager) onStartedLeading(ctx context.Context) { return } ingestor := func(report *nodeReport, now time.Time) error { - // Freshness guard (#341 review): a slow pull must not overwrite a - // heartbeat that arrived while the request was in flight. - if skipStalePull(l.registry, report) { - return nil - } + // Freshness is enforced atomically inside the registry by + // HeartbeatUnlessStale (#341 review): a slow pull cannot overwrite a + // heartbeat that landed while the request was in flight. _, err := l.svc.ingestNodeReport(report, now) return err } diff --git a/services/scheduler/internal/node_registry.go b/services/scheduler/internal/node_registry.go index 6054349d5..dffc9cd50 100644 --- a/services/scheduler/internal/node_registry.go +++ b/services/scheduler/internal/node_registry.go @@ -31,6 +31,9 @@ type NodeRegistry interface { // never reported. Used by the leader-election recovery window to tell // fresh observations from pre-acquisition ones (#259). LastReportAt(nodeID string) (time.Time, bool) + // HeartbeatUnlessStale applies a report only when the node has no newer + // report than notOlderThan, atomically (#341 review). + HeartbeatUnlessStale(req *schedulerv1.HeartbeatRequest, now, notOlderThan time.Time) (applied bool, node Node, cpuConfigJSON string, err error) UnregisterObserved(nodeID string, serviceInstanceID string) error } @@ -143,6 +146,29 @@ func (r *AtomicNodeRegistry) Set(active []Node, lingering []Node) { } func (r *AtomicNodeRegistry) Heartbeat(req *schedulerv1.HeartbeatRequest, now time.Time) (Node, string, error) { + node, cpu, err := r.heartbeat(req, now, time.Time{}) + return node, cpu, err +} + +// HeartbeatUnlessStale applies a report only when the node has no newer +// report than notOlderThan. The check and the write hold the same lock, so +// a heartbeat landing between a pull's freshness check and its ingest +// cannot be overwritten by that pull (#341 review). applied is false when +// a newer report already exists; node and cpuConfigJSON are zero then. +func (r *AtomicNodeRegistry) HeartbeatUnlessStale(req *schedulerv1.HeartbeatRequest, now, notOlderThan time.Time) (applied bool, node Node, cpuConfigJSON string, err error) { + node, cpu, err := r.heartbeat(req, now, notOlderThan) + if err != nil { + if errors.Is(err, errStaleReport) { + return false, Node{}, "", nil + } + return false, Node{}, "", err + } + return true, node, cpu, nil +} + +var errStaleReport = errors.New("a newer report already exists") + +func (r *AtomicNodeRegistry) heartbeat(req *schedulerv1.HeartbeatRequest, now time.Time, notOlderThan time.Time) (Node, string, error) { nowMs := now.UTC().UnixMilli() machineInfo := cloneMachineInfo(req.GetMachineInfo()) @@ -153,15 +179,28 @@ func (r *AtomicNodeRegistry) Heartbeat(req *schedulerv1.HeartbeatRequest, now ti if !ok { return Node{}, "", ErrNodeNotInRegistry } + if !notOlderThan.IsZero() { + if prev, ok := r.observed[req.GetNodeId()]; ok && prev.node.GetLastSeenUnixMs() > notOlderThan.UTC().UnixMilli() { + return Node{}, "", errStaleReport + } + } prevCPU, existed := "", false + prevP2P := (*schedulerv1.P2PEndpoint)(nil) if prev, ok := r.observed[req.GetNodeId()]; ok { existed = true prevCPU = prev.node.GetMachineInfo().GetCpuConfigJson() + prevP2P = prev.p2pEndpoint if machineInfo != nil && machineInfo.CpuConfigJson == "" { machineInfo.CpuConfigJson = prevCPU } } + // A report without a P2P endpoint (e.g. a snapshot pull, which the admin + // API does not serve) must not erase a previously known one (#341 review). + p2pEndpoint := cloneP2PEndpoint(req.GetP2PEndpoint()) + if p2pEndpoint == nil { + p2pEndpoint = prevP2P + } record := observedNodeRecord{ node: &schedulerv1.ObservedNode{ @@ -175,7 +214,7 @@ func (r *AtomicNodeRegistry) Heartbeat(req *schedulerv1.HeartbeatRequest, now ti LastSeenUnixMs: nowMs, Snapshot: cloneSnapshot(req.GetSnapshot()), }, - p2pEndpoint: cloneP2PEndpoint(req.GetP2PEndpoint()), + p2pEndpoint: p2pEndpoint, reportTTL: r.observedTTL, } if record.node.Snapshot.GetReportedAtUnixMs() == 0 { diff --git a/services/scheduler/internal/node_snapshot.go b/services/scheduler/internal/node_snapshot.go index 298a92bf1..dd125143c 100644 --- a/services/scheduler/internal/node_snapshot.go +++ b/services/scheduler/internal/node_snapshot.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "io" "net/http" "strings" "sync" @@ -90,18 +91,6 @@ type NodeSnapshotFetcher func(ctx context.Context, node Node) (*nodeReport, erro // shared ingest path. The retriever never holds the Service itself. type nodeReportIngester func(report *nodeReport, now time.Time) error -// skipStalePull reports whether a pulled report should be dropped because -// the node has already reported something newer — a heartbeat that arrived -// while the pull was in flight (#341 review). Heartbeat-shaped reports -// (zero fetchedAt) are never skipped. -func skipStalePull(registry NodeRegistry, report *nodeReport) bool { - if report.fetchedAt.IsZero() { - return false - } - at, ok := registry.LastReportAt(report.nodeID) - return ok && at.After(report.fetchedAt) -} - type ConcurrentNodeSnapshotRefresher struct { logger *zap.Logger ingest nodeReportIngester @@ -167,6 +156,7 @@ type adminNodeResponse struct { CreateFails uint64 `json:"createFails"` SandboxStartingCnt uint32 `json:"sandboxStartingCount"` SandboxPausedCount uint32 `json:"sandboxPausedCount"` + Status string `json:"status"` MachineInfo struct { CPUFamily string `json:"cpuFamily"` CPUModel string `json:"cpuModel"` @@ -193,6 +183,24 @@ type adminNodeResponse struct { } `json:"metrics"` } +// adminStatusToProto maps the admin /nodes status string to the proto +// NodeStatus. Unknown values stay UNSPECIFIED (the registry then derives +// CONNECTING, same as a first heartbeat). +func adminStatusToProto(status string) schedulerv1.NodeStatus { + switch strings.ToLower(strings.TrimSpace(status)) { + case "ready": + return schedulerv1.NodeStatus_NODE_STATUS_READY + case "connecting": + return schedulerv1.NodeStatus_NODE_STATUS_CONNECTING + case "unhealthy": + return schedulerv1.NodeStatus_NODE_STATUS_UNHEALTHY + case "lingering": + return schedulerv1.NodeStatus_NODE_STATUS_LINGERING + default: + return schedulerv1.NodeStatus_NODE_STATUS_UNSPECIFIED + } +} + // NewAdminSnapshotFetcher builds the production NodeSnapshotFetcher (#259, // sync-node-snapshots): it pulls GET {endpoint}/nodes for observations and // returns the fetched data as a nodeReport for the shared ingest path. @@ -204,6 +212,10 @@ func NewAdminSnapshotFetcher(apiKey string) NodeSnapshotFetcher { client := &http.Client{Timeout: 5 * time.Second} return func(ctx context.Context, node Node) (*nodeReport, error) { base := strings.TrimRight(node.Endpoint, "/") + // Stamp the capture time before the request: a heartbeat landing + // while this request is in flight must count as newer than whatever + // the response contains (#341 review). + fetchedAt := time.Now() var nodes []adminNodeResponse if err := adminGetJSON(ctx, client, apiKey, base+"/nodes", &nodes); err != nil { @@ -216,9 +228,9 @@ func NewAdminSnapshotFetcher(apiKey string) NodeSnapshotFetcher { break } } - if entry == nil && len(nodes) == 1 { - entry = &nodes[0] - } + // Never accept a response for a different node ID: a stale or + // misconfigured endpoint must not overwrite another node's + // observation (#341 review). if entry == nil { return nil, fmt.Errorf("pull %s /nodes: node not in admin response", node.ID) } @@ -226,7 +238,7 @@ func NewAdminSnapshotFetcher(apiKey string) NodeSnapshotFetcher { // Only /nodes is pulled: observations, not the sandbox roster. See // nodeReport.sandboxIDs for why the roster is never reconciled from // a pull. - return reportFromAdmin(entry), nil + return reportFromAdmin(entry, fetchedAt), nil } } @@ -246,10 +258,13 @@ func adminGetJSON(ctx context.Context, client *http.Client, apiKey, url string, if resp.StatusCode != http.StatusOK { return fmt.Errorf("unexpected status %s", resp.Status) } - return json.NewDecoder(resp.Body).Decode(out) + // Bound the decoded body: pulls fan out across nodes during acquisition, + // and a faulty or compromised node must not make the leader allocate + // from an arbitrarily large response (#341 review). + return json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(out) } -func reportFromAdmin(n *adminNodeResponse) *nodeReport { +func reportFromAdmin(n *adminNodeResponse, fetchedAt time.Time) *nodeReport { disks := make([]*schedulerv1.DiskMetric, 0, len(n.Metrics.Disks)) for _, d := range n.Metrics.Disks { disks = append(disks, &schedulerv1.DiskMetric{ @@ -266,7 +281,7 @@ func reportFromAdmin(n *adminNodeResponse) *nodeReport { serviceInstanceID: n.ServiceInstanceID, version: n.Version, commit: n.Commit, - fetchedAt: time.Now(), + fetchedAt: fetchedAt, machineInfo: &schedulerv1.MachineInfo{ CpuFamily: n.MachineInfo.CPUFamily, CpuModel: n.MachineInfo.CPUModel, @@ -275,6 +290,7 @@ func reportFromAdmin(n *adminNodeResponse) *nodeReport { CpuConfigJson: n.MachineInfo.CPUConfigJSON, }, snapshot: &schedulerv1.NodeSnapshot{ + Status: adminStatusToProto(n.Status), AllocatedCpu: n.Metrics.AllocatedCPU, AllocatedMemoryBytes: n.Metrics.AllocatedMemoryBytes, CpuPercent: n.Metrics.CPUPercent, diff --git a/services/scheduler/internal/node_snapshot_test.go b/services/scheduler/internal/node_snapshot_test.go index a9c1bd0a8..0aa498dc1 100644 --- a/services/scheduler/internal/node_snapshot_test.go +++ b/services/scheduler/internal/node_snapshot_test.go @@ -126,32 +126,42 @@ func TestAdminSnapshotFetcherHeartbeatIngestCompatibility(t *testing.T) { } } -// Freshness guard (#341 review, N1): a pull captured before a heartbeat -// must not overwrite it; a pull captured after is applied; heartbeat-shaped -// reports are never skipped. -func TestSkipStalePull(t *testing.T) { - _, registry, _ := newTestService(t, []string{"node-a"}) +// Freshness guard (#341 review): the registry applies a pulled report +// atomically only when nothing newer exists — a pull captured before a +// heartbeat must not overwrite it, a pull captured after is applied, and +// heartbeat-shaped reports are never skipped. +func TestPulledIngestIsAtomicWithNewerReports(t *testing.T) { + svc, registry, _ := newTestService(t, []string{"node-a"}) t0 := time.Date(2026, 10, 5, 12, 0, 0, 0, time.UTC) // A heartbeat lands at t0+1s. if _, _, err := registry.Heartbeat(heartbeatFor("node-a", "inst-1"), t0.Add(time.Second)); err != nil { t.Fatalf("heartbeat failed: %v", err) } + before, _ := registry.LastReportAt("node-a") - stale := &nodeReport{nodeID: "node-a", fetchedAt: t0} - if !skipStalePull(registry, stale) { - t.Fatal("a pull captured before the latest heartbeat must be skipped") + // A pull captured before that heartbeat must not overwrite it. + stale := &nodeReport{nodeID: "node-a", serviceInstanceID: "inst-1", fetchedAt: t0} + if _, err := svc.ingestNodeReport(stale, t0.Add(2*time.Second)); err != nil { + t.Fatalf("stale pull ingest failed: %v", err) } - fresh := &nodeReport{nodeID: "node-a", fetchedAt: t0.Add(2 * time.Second)} - if skipStalePull(registry, fresh) { - t.Fatal("a pull captured after the latest heartbeat must be applied") + after, _ := registry.LastReportAt("node-a") + if !after.Equal(before) { + t.Fatalf("stale pull must not overwrite the newer heartbeat: before=%v after=%v", before, after) } - heartbeatShaped := &nodeReport{nodeID: "node-a"} - if skipStalePull(registry, heartbeatShaped) { - t.Fatal("heartbeat-shaped reports (zero fetchedAt) must never be skipped") + + // A pull captured after the heartbeat is applied. + fresh := &nodeReport{nodeID: "node-a", serviceInstanceID: "inst-1", fetchedAt: t0.Add(3 * time.Second)} + if _, err := svc.ingestNodeReport(fresh, t0.Add(3*time.Second)); err != nil { + t.Fatalf("fresh pull ingest failed: %v", err) + } + after2, _ := registry.LastReportAt("node-a") + if !after2.After(after) { + t.Fatalf("fresh pull must be applied: after=%v after2=%v", after, after2) } - unknown := &nodeReport{nodeID: "node-b", fetchedAt: t0} - if skipStalePull(registry, unknown) { - t.Fatal("an unobserved node must not be skipped") + + // Heartbeat-shaped reports (zero fetchedAt) always apply. + if _, err := svc.ingestNodeReport(reportFor("node-a", "inst-1"), t0.Add(4*time.Second)); err != nil { + t.Fatalf("heartbeat-shaped ingest failed: %v", err) } } diff --git a/services/scheduler/internal/service.go b/services/scheduler/internal/service.go index fbcf7c6ef..f6b096586 100644 --- a/services/scheduler/internal/service.go +++ b/services/scheduler/internal/service.go @@ -163,10 +163,15 @@ func (s *Service) filterFreshObservations(rich []RichNode) []RichNode { if !ok { return rich } + // Compare at millisecond precision: LastReportAt is reconstructed from a + // millisecond timestamp while the acquisition time has nanoseconds, so a + // report in the same millisecond as the acquisition must count as fresh + // (#341 review). + sinceMs := since.UnixMilli() fresh := make([]RichNode, 0, len(rich)) for _, n := range rich { at, reported := s.nodes.LastReportAt(n.Node.ID) - if reported && !at.Before(since) { + if reported && at.UnixMilli() >= sinceMs { fresh = append(fresh, n) continue } @@ -282,6 +287,17 @@ func (s *Service) Heartbeat(_ context.Context, req *schedulerv1.HeartbeatRequest return s.ingestNodeReport(nodeReportFromProto(req), s.now()) } +// ingestNodeError maps registry ingest failures to gRPC statuses. +func ingestNodeError(logger *zap.Logger, nodeID string, err error) error { + if errors.Is(err, ErrNodeNotInRegistry) { + logger.Warn("scheduler rejected observed registration for unknown node", + zap.String("node_id", nodeID), + ) + return status.Error(codes.InvalidArgument, "node is not in scheduler node list") + } + return status.Error(codes.Internal, "node registry heartbeat failed") +} + // nodeReportFromProto adapts a Heartbeat RPC request to the neutral ingest // currency. Identity validation already happened in Heartbeat. func nodeReportFromProto(req *schedulerv1.HeartbeatRequest) *nodeReport { @@ -309,15 +325,26 @@ func nodeReportFromProto(req *schedulerv1.HeartbeatRequest) *nodeReport { // node snapshots (sync-node-snapshots on leadership acquisition, #259): // update registry observations, then reconcile sandbox bindings. func (s *Service) ingestNodeReport(report *nodeReport, now time.Time) (*schedulerv1.HeartbeatResponse, error) { - node, cpuConfigJSON, err := s.nodes.Heartbeat(report.toHeartbeatRequest(), now) - if err != nil { - if errors.Is(err, ErrNodeNotInRegistry) { - s.logger.Warn("scheduler rejected observed registration for unknown node", - zap.String("node_id", report.nodeID), - ) - return nil, status.Error(codes.InvalidArgument, "node is not in scheduler node list") + req := report.toHeartbeatRequest() + var node Node + var cpuConfigJSON string + if !report.fetchedAt.IsZero() { + // Pulled report: apply only if nothing newer landed meanwhile; the + // check is atomic with the write inside the registry (#341 review). + applied, n, cpu, err := s.nodes.HeartbeatUnlessStale(req, now, report.fetchedAt) + if err != nil { + return nil, ingestNodeError(s.logger, report.nodeID, err) + } + if !applied { + return &schedulerv1.HeartbeatResponse{}, nil + } + node, cpuConfigJSON = n, cpu + } else { + n, cpu, err := s.nodes.Heartbeat(req, now) + if err != nil { + return nil, ingestNodeError(s.logger, report.nodeID, err) } - return nil, status.Error(codes.Internal, "node registry heartbeat failed") + node, cpuConfigJSON = n, cpu } // Reconcile only from an authoritative roster (heartbeat). A snapshot // pull leaves sandboxIDs nil — bindings are refreshed by heartbeats and From 8b8b03bd16348ab1794ee921ce238189c586be58 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 21:57:47 +0800 Subject: [PATCH 11/26] refactor(scheduler): restore registry Heartbeat, contain pull semantics Address the review concern that the previous commit reshaped the shared AtomicNodeRegistry.Heartbeat for a problem that only exists on the new snapshot-pull path. Heartbeat is now byte-identical to its pre-HA form; everything pull-specific lives in the pull path: - The atomic HeartbeatUnlessStale and the in-registry P2P merge are removed. Freshness is guarded in Service.ingestNodeReport: a pulled report is skipped when LastReportAt is newer than its fetchedAt. The residual check-then-act window is microseconds and self-heals on the next heartbeat (5s). - P2P endpoint preservation moves to the pull path too: a read-only P2PEndpointFor getter (purely additive to the registry interface) lets ingest fill the endpoint into the pulled request before calling the untouched Heartbeat. Scope: services/scheduler/internal (node_registry.go, service.go, node_snapshot_test.go). Relates to #259. Co-Authored-By: Claude Code --- services/scheduler/internal/node_registry.go | 56 ++++++------------- .../scheduler/internal/node_snapshot_test.go | 10 ++-- services/scheduler/internal/service.go | 33 +++++------ 3 files changed, 38 insertions(+), 61 deletions(-) diff --git a/services/scheduler/internal/node_registry.go b/services/scheduler/internal/node_registry.go index dffc9cd50..a5c50af03 100644 --- a/services/scheduler/internal/node_registry.go +++ b/services/scheduler/internal/node_registry.go @@ -31,9 +31,10 @@ type NodeRegistry interface { // never reported. Used by the leader-election recovery window to tell // fresh observations from pre-acquisition ones (#259). LastReportAt(nodeID string) (time.Time, bool) - // HeartbeatUnlessStale applies a report only when the node has no newer - // report than notOlderThan, atomically (#341 review). - HeartbeatUnlessStale(req *schedulerv1.HeartbeatRequest, now, notOlderThan time.Time) (applied bool, node Node, cpuConfigJSON string, err error) + // P2PEndpointFor returns the node's last reported P2P endpoint, or nil. + // Read-only lookup used by the snapshot-pull path, whose admin source + // cannot provide the endpoint (#341 review). + P2PEndpointFor(nodeID string) *schedulerv1.P2PEndpoint UnregisterObserved(nodeID string, serviceInstanceID string) error } @@ -146,29 +147,6 @@ func (r *AtomicNodeRegistry) Set(active []Node, lingering []Node) { } func (r *AtomicNodeRegistry) Heartbeat(req *schedulerv1.HeartbeatRequest, now time.Time) (Node, string, error) { - node, cpu, err := r.heartbeat(req, now, time.Time{}) - return node, cpu, err -} - -// HeartbeatUnlessStale applies a report only when the node has no newer -// report than notOlderThan. The check and the write hold the same lock, so -// a heartbeat landing between a pull's freshness check and its ingest -// cannot be overwritten by that pull (#341 review). applied is false when -// a newer report already exists; node and cpuConfigJSON are zero then. -func (r *AtomicNodeRegistry) HeartbeatUnlessStale(req *schedulerv1.HeartbeatRequest, now, notOlderThan time.Time) (applied bool, node Node, cpuConfigJSON string, err error) { - node, cpu, err := r.heartbeat(req, now, notOlderThan) - if err != nil { - if errors.Is(err, errStaleReport) { - return false, Node{}, "", nil - } - return false, Node{}, "", err - } - return true, node, cpu, nil -} - -var errStaleReport = errors.New("a newer report already exists") - -func (r *AtomicNodeRegistry) heartbeat(req *schedulerv1.HeartbeatRequest, now time.Time, notOlderThan time.Time) (Node, string, error) { nowMs := now.UTC().UnixMilli() machineInfo := cloneMachineInfo(req.GetMachineInfo()) @@ -179,28 +157,15 @@ func (r *AtomicNodeRegistry) heartbeat(req *schedulerv1.HeartbeatRequest, now ti if !ok { return Node{}, "", ErrNodeNotInRegistry } - if !notOlderThan.IsZero() { - if prev, ok := r.observed[req.GetNodeId()]; ok && prev.node.GetLastSeenUnixMs() > notOlderThan.UTC().UnixMilli() { - return Node{}, "", errStaleReport - } - } prevCPU, existed := "", false - prevP2P := (*schedulerv1.P2PEndpoint)(nil) if prev, ok := r.observed[req.GetNodeId()]; ok { existed = true prevCPU = prev.node.GetMachineInfo().GetCpuConfigJson() - prevP2P = prev.p2pEndpoint if machineInfo != nil && machineInfo.CpuConfigJson == "" { machineInfo.CpuConfigJson = prevCPU } } - // A report without a P2P endpoint (e.g. a snapshot pull, which the admin - // API does not serve) must not erase a previously known one (#341 review). - p2pEndpoint := cloneP2PEndpoint(req.GetP2PEndpoint()) - if p2pEndpoint == nil { - p2pEndpoint = prevP2P - } record := observedNodeRecord{ node: &schedulerv1.ObservedNode{ @@ -214,7 +179,7 @@ func (r *AtomicNodeRegistry) heartbeat(req *schedulerv1.HeartbeatRequest, now ti LastSeenUnixMs: nowMs, Snapshot: cloneSnapshot(req.GetSnapshot()), }, - p2pEndpoint: p2pEndpoint, + p2pEndpoint: cloneP2PEndpoint(req.GetP2PEndpoint()), reportTTL: r.observedTTL, } if record.node.Snapshot.GetReportedAtUnixMs() == 0 { @@ -389,6 +354,17 @@ func (r *AtomicNodeRegistry) PeekObserved(nodeID string) *schedulerv1.NodeSnapsh return cloneSnapshot(snapshot) } +// P2PEndpointFor returns the node's last reported P2P endpoint, or nil. +func (r *AtomicNodeRegistry) P2PEndpointFor(nodeID string) *schedulerv1.P2PEndpoint { + r.mu.RLock() + defer r.mu.RUnlock() + record, ok := r.observed[nodeID] + if !ok { + return nil + } + return record.p2pEndpoint +} + func (r *AtomicNodeRegistry) LastReportAt(nodeID string) (time.Time, bool) { r.mu.RLock() defer r.mu.RUnlock() diff --git a/services/scheduler/internal/node_snapshot_test.go b/services/scheduler/internal/node_snapshot_test.go index 0aa498dc1..b9cad219b 100644 --- a/services/scheduler/internal/node_snapshot_test.go +++ b/services/scheduler/internal/node_snapshot_test.go @@ -126,11 +126,11 @@ func TestAdminSnapshotFetcherHeartbeatIngestCompatibility(t *testing.T) { } } -// Freshness guard (#341 review): the registry applies a pulled report -// atomically only when nothing newer exists — a pull captured before a -// heartbeat must not overwrite it, a pull captured after is applied, and -// heartbeat-shaped reports are never skipped. -func TestPulledIngestIsAtomicWithNewerReports(t *testing.T) { +// Freshness guard (#341 review): the pull path skips reports older than +// the node's latest observation — a pull captured before a heartbeat must +// not overwrite it, a pull captured after is applied, and heartbeat-shaped +// reports are never skipped. +func TestPulledIngestSkipsStaleReports(t *testing.T) { svc, registry, _ := newTestService(t, []string{"node-a"}) t0 := time.Date(2026, 10, 5, 12, 0, 0, 0, time.UTC) diff --git a/services/scheduler/internal/service.go b/services/scheduler/internal/service.go index f6b096586..667fcef31 100644 --- a/services/scheduler/internal/service.go +++ b/services/scheduler/internal/service.go @@ -175,7 +175,11 @@ func (s *Service) filterFreshObservations(rich []RichNode) []RichNode { fresh = append(fresh, n) continue } - if !s.leadership.RecoveryPending(n.Node.ID, at, reported) { + var lastReport *time.Time + if reported { + lastReport = &at + } + if !s.leadership.RecoveryPending(n.Node.ID, lastReport) { fresh = append(fresh, n) } } @@ -326,25 +330,22 @@ func nodeReportFromProto(req *schedulerv1.HeartbeatRequest) *nodeReport { // update registry observations, then reconcile sandbox bindings. func (s *Service) ingestNodeReport(report *nodeReport, now time.Time) (*schedulerv1.HeartbeatResponse, error) { req := report.toHeartbeatRequest() - var node Node - var cpuConfigJSON string if !report.fetchedAt.IsZero() { - // Pulled report: apply only if nothing newer landed meanwhile; the - // check is atomic with the write inside the registry (#341 review). - applied, n, cpu, err := s.nodes.HeartbeatUnlessStale(req, now, report.fetchedAt) - if err != nil { - return nil, ingestNodeError(s.logger, report.nodeID, err) - } - if !applied { + // Pulled report: skip when the node already reported something + // newer (#341 review). The residual check-then-act window is + // microseconds and self-heals on the next heartbeat. + if at, ok := s.nodes.LastReportAt(report.nodeID); ok && at.After(report.fetchedAt) { return &schedulerv1.HeartbeatResponse{}, nil } - node, cpuConfigJSON = n, cpu - } else { - n, cpu, err := s.nodes.Heartbeat(req, now) - if err != nil { - return nil, ingestNodeError(s.logger, report.nodeID, err) + // The admin API does not serve the P2P endpoint; carry over the + // previously known one so a pull does not erase it (#341 review). + if req.P2PEndpoint == nil { + req.P2PEndpoint = s.nodes.P2PEndpointFor(report.nodeID) } - node, cpuConfigJSON = n, cpu + } + node, cpuConfigJSON, err := s.nodes.Heartbeat(req, now) + if err != nil { + return nil, ingestNodeError(s.logger, report.nodeID, err) } // Reconcile only from an authoritative roster (heartbeat). A snapshot // pull leaves sandboxIDs nil — bindings are refreshed by heartbeats and From 148c98360b4108633c22e5fe480edf63a90e7800 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 21:58:11 +0800 Subject: [PATCH 12/26] refactor(scheduler): pointer-typed lastReportAt in RecoveryPending Replace the (time.Time, bool) pair with *time.Time in the RecoveryPending signature: nil means "never reported", which is self-explanatory at the call site instead of a free-floating ok flag. Applies uniformly to the leadershipView seam, the snapshot implementation, the manager delegate, and the null object. Scope: services/scheduler/internal (leadership.go, service.go). Relates to #259. Co-Authored-By: Claude Code --- services/scheduler/internal/leadership.go | 32 ++++++++++++++--------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/services/scheduler/internal/leadership.go b/services/scheduler/internal/leadership.go index 176783dd4..fe7b8085d 100644 --- a/services/scheduler/internal/leadership.go +++ b/services/scheduler/internal/leadership.go @@ -32,7 +32,10 @@ type leadershipView interface { LeaderSince() (time.Time, bool) // RecoveryPending reports whether nodeID was known at acquisition and // has not yet produced a post-acquisition report. - RecoveryPending(nodeID string, lastReportAt time.Time, ok bool) bool + // RecoveryPending reports whether nodeID was known at acquisition and + // has not yet produced a post-acquisition report. A nil lastReportAt + // means the node never reported. + RecoveryPending(nodeID string, lastReportAt *time.Time) bool } // Leadership is the single external contract for the scheduler's @@ -77,13 +80,13 @@ func NewLeadership(logger *zap.Logger, cfg config.SchedulerConfig) Leadership { // a single-writer scheduler without any conditional wiring. type nonLeadership struct{} -func (nonLeadership) InRecoveryWindow(time.Time) bool { return false } -func (nonLeadership) LeaderSince() (time.Time, bool) { return time.Time{}, false } -func (nonLeadership) RecoveryPending(string, time.Time, bool) bool { return false } -func (nonLeadership) ServiceOption() ServiceOption { return func(*Service) {} } -func (nonLeadership) RegisterHealth(*health.Server) {} -func (nonLeadership) BindRuntime(*Service, NodeRegistry) {} -func (nonLeadership) MarkNotServing() {} +func (nonLeadership) InRecoveryWindow(time.Time) bool { return false } +func (nonLeadership) LeaderSince() (time.Time, bool) { return time.Time{}, false } +func (nonLeadership) RecoveryPending(string, *time.Time) bool { return false } +func (nonLeadership) ServiceOption() ServiceOption { return func(*Service) {} } +func (nonLeadership) RegisterHealth(*health.Server) {} +func (nonLeadership) BindRuntime(*Service, NodeRegistry) {} +func (nonLeadership) MarkNotServing() {} func (nonLeadership) GateInterceptor() grpc.UnaryServerInterceptor { return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { return handler(ctx, req) @@ -135,17 +138,20 @@ func acquiredLeadershipSnapshot(recoveryTTL time.Duration, now time.Time, knownN return &leadershipSnapshot{leader: true, since: now, recoveryTTL: recoveryTTL, knownNodes: known} } -// recoveryPending reports whether nodeID was known at acquisition and has +// RecoveryPending reports whether nodeID was known at acquisition and has // not yet produced a post-acquisition report (lastReportAt is the node's // latest report time, ok=false when it never reported). -func (l *leadershipSnapshot) RecoveryPending(nodeID string, lastReportAt time.Time, ok bool) bool { +func (l *leadershipSnapshot) RecoveryPending(nodeID string, lastReportAt *time.Time) bool { if !l.leader || !l.knownNodes[nodeID] { return false } + if lastReportAt == nil { + return true + } // Millisecond-precision comparison, matching the scheduler's freshness // rule (#341 review): a report in the same millisecond as the // acquisition counts as post-acquisition. - return !ok || lastReportAt.UnixMilli() < l.since.UnixMilli() + return lastReportAt.UnixMilli() < l.since.UnixMilli() } // InRecoveryWindow reports whether the recovery window following leadership @@ -246,8 +252,8 @@ func (l *leadershipManager) LeaderSince() (time.Time, bool) { } // RecoveryPending delegates to the internal leadership state. -func (l *leadershipManager) RecoveryPending(nodeID string, lastReportAt time.Time, ok bool) bool { - return l.state.Load().RecoveryPending(nodeID, lastReportAt, ok) +func (l *leadershipManager) RecoveryPending(nodeID string, lastReportAt *time.Time) bool { + return l.state.Load().RecoveryPending(nodeID, lastReportAt) } // standbyReadableMethods classifies every Scheduler RPC: true = a non-leader From 04b371aa3cb8bdade41c6dea5ad1e04152f6d10a Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 21:58:58 +0800 Subject: [PATCH 13/26] feat(scheduler,gateway): enable standby reads via round_robin gateway retry Resolve the F1 review finding: with leader-only readiness endpoints, standbys never received read traffic, so the standby Redis-read path was unreachable. This implements the "standbys stay Ready" option: - HA overlay readiness probes overall health again, so standbys are in Service endpoints and gateway reads reach them. The leader health service (scheduler.v1.Scheduler/leader) stays registered for observability but no longer gates endpoints. - The gateway's scheduler connection switches from the default pick_first (which pinned one pod) to round_robin plus a retry policy on Unavailable, so writes that land on a leader-gated standby are retried onto the leader (at most two retries with 3 replicas). Trade-offs, called out for review: endpoints now include standbys (deviating from the issue's "readiness exposes only the leader"); ~2/3 of write RPCs take one retry under 3 replicas; the retry policy also retries semantic Unavailable errors (harmless, same result). If the "drop standby reads" option is preferred, this commit reverts cleanly on its own. Scope: services/gateway/cmd, deploy/k8s/overlays/ha. Relates to #259. Co-Authored-By: Claude Code --- deploy/k8s/overlays/ha/kustomization.yaml | 7 ++++++- services/gateway/cmd/main.go | 19 +++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/deploy/k8s/overlays/ha/kustomization.yaml b/deploy/k8s/overlays/ha/kustomization.yaml index f7cdb9059..ba6d23bc1 100644 --- a/deploy/k8s/overlays/ha/kustomization.yaml +++ b/deploy/k8s/overlays/ha/kustomization.yaml @@ -30,12 +30,17 @@ patches: - op: replace path: /spec/replicas value: 3 + # Standbys are Ready (overall health) so they are in Service endpoints + # and can serve reads from the shared Redis bindings. Writes that land + # on a standby are rejected by the leader gate and retried onto the + # leader by the gateway's round_robin + retry policy. The leader health + # service (scheduler.v1.Scheduler/leader) stays registered for + # observability but does not gate endpoints in this topology. - op: replace path: /spec/template/spec/containers/0/readinessProbe/exec/command value: - /grpc_health_probe - -addr=127.0.0.1:9090 - - -service=scheduler.v1.Scheduler/leader # sync-node-snapshots pulls authenticate against node admin APIs with # the cluster-wide admin key; mount it from the shared agentenv-auth # Secret, never from a ConfigMap. diff --git a/services/gateway/cmd/main.go b/services/gateway/cmd/main.go index e7bbef260..8faafc381 100644 --- a/services/gateway/cmd/main.go +++ b/services/gateway/cmd/main.go @@ -32,10 +32,29 @@ const ( maxAPIKeyFileLen = maxAPIKeyLen + 2 ) +// schedulerServiceConfig enables round_robin load balancing across scheduler +// replicas and retries Unavailable once or twice, so write RPCs that land on +// a leader-gated standby (Unavailable "not the leader") are retried onto the +// leader (#341, HA topology with standbys in Service endpoints). +const schedulerServiceConfig = `{ + "loadBalancingConfig": [{"round_robin": {}}], + "methodConfig": [{ + "name": [{"service": "scheduler.v1.Scheduler"}], + "retryPolicy": { + "maxAttempts": 3, + "initialBackoff": "0.2s", + "maxBackoff": "1s", + "backoffMultiplier": 2, + "retryableStatusCodes": ["UNAVAILABLE"] + } + }] +}` + func newSchedulerConn(addr string) (*grpc.ClientConn, error) { return grpc.NewClient( addr, grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithDefaultServiceConfig(schedulerServiceConfig), ) } From 2f29b5906be8d427d67ef08fc6353cce022efc70 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Mon, 5 Oct 2026 22:33:41 +0800 Subject: [PATCH 14/26] Revert "feat(scheduler,gateway): enable standby reads via round_robin gateway retry" This reverts commit 04b371aa3cb8bdade41c6dea5ad1e04152f6d10a. --- deploy/k8s/overlays/ha/kustomization.yaml | 7 +------ services/gateway/cmd/main.go | 19 ------------------- 2 files changed, 1 insertion(+), 25 deletions(-) diff --git a/deploy/k8s/overlays/ha/kustomization.yaml b/deploy/k8s/overlays/ha/kustomization.yaml index ba6d23bc1..f7cdb9059 100644 --- a/deploy/k8s/overlays/ha/kustomization.yaml +++ b/deploy/k8s/overlays/ha/kustomization.yaml @@ -30,17 +30,12 @@ patches: - op: replace path: /spec/replicas value: 3 - # Standbys are Ready (overall health) so they are in Service endpoints - # and can serve reads from the shared Redis bindings. Writes that land - # on a standby are rejected by the leader gate and retried onto the - # leader by the gateway's round_robin + retry policy. The leader health - # service (scheduler.v1.Scheduler/leader) stays registered for - # observability but does not gate endpoints in this topology. - op: replace path: /spec/template/spec/containers/0/readinessProbe/exec/command value: - /grpc_health_probe - -addr=127.0.0.1:9090 + - -service=scheduler.v1.Scheduler/leader # sync-node-snapshots pulls authenticate against node admin APIs with # the cluster-wide admin key; mount it from the shared agentenv-auth # Secret, never from a ConfigMap. diff --git a/services/gateway/cmd/main.go b/services/gateway/cmd/main.go index 8faafc381..e7bbef260 100644 --- a/services/gateway/cmd/main.go +++ b/services/gateway/cmd/main.go @@ -32,29 +32,10 @@ const ( maxAPIKeyFileLen = maxAPIKeyLen + 2 ) -// schedulerServiceConfig enables round_robin load balancing across scheduler -// replicas and retries Unavailable once or twice, so write RPCs that land on -// a leader-gated standby (Unavailable "not the leader") are retried onto the -// leader (#341, HA topology with standbys in Service endpoints). -const schedulerServiceConfig = `{ - "loadBalancingConfig": [{"round_robin": {}}], - "methodConfig": [{ - "name": [{"service": "scheduler.v1.Scheduler"}], - "retryPolicy": { - "maxAttempts": 3, - "initialBackoff": "0.2s", - "maxBackoff": "1s", - "backoffMultiplier": 2, - "retryableStatusCodes": ["UNAVAILABLE"] - } - }] -}` - func newSchedulerConn(addr string) (*grpc.ClientConn, error) { return grpc.NewClient( addr, grpc.WithTransportCredentials(insecure.NewCredentials()), - grpc.WithDefaultServiceConfig(schedulerServiceConfig), ) } From c4f4d576c31f351db169b8722a40578f4a5af35a Mon Sep 17 00:00:00 2001 From: NickNYU Date: Fri, 9 Oct 2026 17:17:37 +0800 Subject: [PATCH 15/26] fix(scheduler): close takeover races and plug the admin-key redirect leak MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address the maintainer's required takeover fixes on PR #341: - Discovery race (OCR high): on kubernetes discovery the informer runs asynchronously and needs up to ~30s for its first sync, while leader election started immediately. An early acquisition would capture an empty registry (empty recovery-pending set, refresh pulls nothing) and schedule fail-open with zero observations. Discovery now signals first-sync completion, and leader election waits for it (bounded at 30s) before starting; static discovery and non-election deployments are unaffected. - In-flight handover (OCR security-high): OnStoppedLeading only cancelled the root context, and GracefulStop then drained in-flight RPCs for up to 10s — a window where the demoted leader keeps writing while the next leader takes over. The leadership-loss path now force-stops the gRPC server instead: in-flight RPCs fail fast and clients retry onto the new leader, eliminating the dual-writer overlap. Signal-triggered shutdown keeps the graceful drain. - Admin-key redirect leak: the snapshot-pull HTTP client followed redirects, re-sending the x-api-key header to arbitrary targets. It now returns the redirect response as-is instead of following. Scope: services/scheduler (cmd, internal). Relates to #259. Co-Authored-By: Claude Code --- services/scheduler/cmd/main.go | 76 ++++++++++++++----- .../internal/kubernetes_discovery.go | 10 +++ services/scheduler/internal/node_snapshot.go | 10 ++- 3 files changed, 75 insertions(+), 21 deletions(-) diff --git a/services/scheduler/cmd/main.go b/services/scheduler/cmd/main.go index 1fea11a84..4871062ef 100644 --- a/services/scheduler/cmd/main.go +++ b/services/scheduler/cmd/main.go @@ -10,6 +10,7 @@ import ( "os" "os/signal" "strings" + "sync/atomic" "syscall" "time" @@ -42,10 +43,18 @@ func main() { } defer logger.Sync() - // rootCancel lets a leadership loss drive the same graceful shutdown as a - // signal: the ex-leader exits and restarts as a standby (fencing, #259). + // rootCancel lets a leadership loss drive the shutdown path as a signal, + // but with an important difference (#341 review): on leadership loss the + // process force-stops instead of draining, so in-flight writes cannot + // overlap the next leader (no dual-writer window). leadershipLost tells + // the shutdown path which case it is in. + var leadershipLost atomic.Bool rootCtx, rootCancel := context.WithCancel(context.Background()) defer rootCancel() + leadershipLossStop := func() { + leadershipLost.Store(true) + rootCancel() + } sigCtx, stop := signal.NotifyContext(rootCtx, os.Interrupt, syscall.SIGTERM) defer stop() @@ -63,6 +72,16 @@ func main() { g := grpc.NewServer(grpc.ChainUnaryInterceptor(interceptors...)) var registry *scheduler.AtomicNodeRegistry var svc *scheduler.Service + // Closed once discovery has produced its first sync. Closed immediately + // unless leader election runs on kubernetes discovery — there, election + // must not start before the first informer sync, or an early acquisition + // would capture an empty registry and recover nothing (#341 review). + waitForDiscoverySync := cfg.Scheduler.LeaderElection.Enabled && !*queryOnly && + strings.EqualFold(strings.TrimSpace(cfg.Scheduler.Discovery.Mode), "kubernetes") + discoveryReady := make(chan struct{}) + if !waitForDiscoverySync { + close(discoveryReady) + } if *queryOnly { qo := scheduler.NewQueryOnlyService(logger, store) schedulerv1.RegisterSchedulerServer(g, qo) @@ -71,7 +90,7 @@ func main() { registry = scheduler.NewAtomicNodeRegistry(nil, cfg.Scheduler.ReportTTL) switch strings.ToLower(strings.TrimSpace(cfg.Scheduler.Discovery.Mode)) { case "kubernetes": - go runKubernetesDiscoveryWithRetry(sigCtx, logger, cfg.Scheduler.Discovery.Kubernetes, registry) + go runKubernetesDiscoveryWithRetry(sigCtx, logger, cfg.Scheduler.Discovery.Kubernetes, registry, discoveryReady) default: nodes := make([]scheduler.Node, 0, len(cfg.Scheduler.Nodes)) for _, n := range cfg.Scheduler.Nodes { @@ -106,8 +125,16 @@ func main() { grpc_health_v1.RegisterHealthServer(g, hs) leadership.BindRuntime(svc, registry) + if waitForDiscoverySync { + select { + case <-discoveryReady: + case <-time.After(30 * time.Second): + logger.Fatal("kubernetes discovery initial sync timed out before leader election") + case <-sigCtx.Done(): + } + } go func() { - if err := leadership.Run(sigCtx, rootCancel); err != nil { + if err := leadership.Run(sigCtx, leadershipLossStop); err != nil { logger.Fatal("leader election failed", zap.Error(err)) } }() @@ -158,22 +185,30 @@ func main() { hs.SetServingStatus(schedulerv1.Scheduler_ServiceDesc.ServiceName, grpc_health_v1.HealthCheckResponse_NOT_SERVING) leadership.MarkNotServing() - gracefulStopDone := make(chan struct{}) - go func() { - g.GracefulStop() - close(gracefulStopDone) - }() - - timer := time.NewTimer(10 * time.Second) - defer timer.Stop() - - select { - case <-gracefulStopDone: - logger.Info("scheduler stopped gracefully") - case <-timer.C: - logger.Warn("scheduler graceful shutdown timed out; forcing stop") + if leadershipLost.Load() { + // Leadership loss: stop immediately. In-flight RPCs fail and clients + // retry onto the new leader — better than a dual-writer overlap + // (#341 review). + logger.Warn("leadership lost; forcing immediate stop to avoid dual-writer overlap") g.Stop() - <-gracefulStopDone + } else { + gracefulStopDone := make(chan struct{}) + go func() { + g.GracefulStop() + close(gracefulStopDone) + }() + + timer := time.NewTimer(10 * time.Second) + defer timer.Stop() + + select { + case <-gracefulStopDone: + logger.Info("scheduler stopped gracefully") + case <-timer.C: + logger.Warn("scheduler graceful shutdown timed out; forcing stop") + g.Stop() + <-gracefulStopDone + } } metricsShutdownCtx, cancelMetricsShutdown := context.WithTimeout(context.Background(), 5*time.Second) @@ -239,6 +274,7 @@ func runKubernetesDiscoveryWithRetry( logger *zap.Logger, cfg config.SchedulerDiscoveryKubernetesConfig, registry *scheduler.AtomicNodeRegistry, + ready ...chan<- struct{}, ) { const ( initialBackoff = 1 * time.Second @@ -254,7 +290,7 @@ func runKubernetesDiscoveryWithRetry( } attempt++ - discovery, err := scheduler.NewKubernetesDiscovery(logger, cfg, registry) + discovery, err := scheduler.NewKubernetesDiscovery(logger, cfg, registry, ready...) if err != nil { if errors.Is(err, rest.ErrNotInCluster) { logger.Error("kubernetes discovery initialization failed with non-retryable error; stopping discovery loop", diff --git a/services/scheduler/internal/kubernetes_discovery.go b/services/scheduler/internal/kubernetes_discovery.go index c7e36f537..f184ce05a 100644 --- a/services/scheduler/internal/kubernetes_discovery.go +++ b/services/scheduler/internal/kubernetes_discovery.go @@ -25,6 +25,7 @@ import ( const kubernetesDiscoveryCacheSyncTimeout = 30 * time.Second type KubernetesDiscovery struct { + ready []chan<- struct{} logger *zap.Logger config config.SchedulerDiscoveryKubernetesConfig registry *AtomicNodeRegistry @@ -33,10 +34,15 @@ type KubernetesDiscovery struct { noSchedulePodInformer cache.SharedIndexInformer } +// NewKubernetesDiscovery builds discovery. When ready is non-nil, it is +// closed once the initial informer cache sync completes (the first time the +// registry actually reflects the cluster); leader election waits on it so an +// early acquisition cannot capture an empty registry (#341 review). func NewKubernetesDiscovery( logger *zap.Logger, cfg config.SchedulerDiscoveryKubernetesConfig, registry *AtomicNodeRegistry, + ready ...chan<- struct{}, ) (*KubernetesDiscovery, error) { if logger == nil { logger = zap.NewNop() @@ -81,6 +87,7 @@ func NewKubernetesDiscovery( } discovery := &KubernetesDiscovery{ + ready: ready, logger: logger, config: cfg, registry: registry, @@ -150,6 +157,9 @@ func (d *KubernetesDiscovery) Run(ctx context.Context) error { syncCtx, cancelSync := context.WithTimeout(ctx, kubernetesDiscoveryCacheSyncTimeout) defer cancelSync() + if len(d.ready) > 0 { + defer close(d.ready[0]) + } if !cache.WaitForCacheSync(syncCtx.Done(), cacheSyncs...) { if err := ctx.Err(); err != nil { return err diff --git a/services/scheduler/internal/node_snapshot.go b/services/scheduler/internal/node_snapshot.go index dd125143c..7509ef5f9 100644 --- a/services/scheduler/internal/node_snapshot.go +++ b/services/scheduler/internal/node_snapshot.go @@ -209,7 +209,15 @@ func adminStatusToProto(status string) schedulerv1.NodeStatus { // implementation detail with a bounded per-request timeout; tests inject // through the NodeSnapshotFetcher seam, not this constructor. func NewAdminSnapshotFetcher(apiKey string) NodeSnapshotFetcher { - client := &http.Client{Timeout: 5 * time.Second} + client := &http.Client{ + Timeout: 5 * time.Second, + // Never follow redirects: Go re-sends headers (including x-api-key) + // to the redirect target, which would leak the admin key (#341 + // review). A redirect from a node admin API is an error, not a hint. + CheckRedirect: func(_ *http.Request, _ []*http.Request) error { + return http.ErrUseLastResponse + }, + } return func(ctx context.Context, node Node) (*nodeReport, error) { base := strings.TrimRight(node.Endpoint, "/") // Stamp the capture time before the request: a heartbeat landing From 8b333a080f1a9875d749df081bf853064f7a87e3 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Fri, 9 Oct 2026 17:41:55 +0800 Subject: [PATCH 16/26] feat(api,scheduler): expose heartbeat-parity fields on admin /nodes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolve the CPU recovery limitation the maintainer flagged (#341 review): after failover, nodes had already cleared their pending cpu_config_json (reporter.rs clears it after the first successful heartbeat) and the new leader could not learn it anywhere — admin /nodes did not expose it — so CPU config distribution never recovered until node restarts. The same gap made pulled sandbox rosters and P2P endpoints permanently absent. The admin /nodes response now carries the same fields a heartbeat has: - MachineInfo.cpuConfigJSON, populated from the node's machine info; - Node.sandboxIDs, the full tracked roster from orchestrator.list_sandbox_ids (paused sandboxes and template builds included — the authoritative source heartbeats use); - Node.p2pEndpoint, plumbed from the server's P2P transport into ObservabilityService and the snapshot model. On the scheduler side, the snapshot fetcher decodes all three. A roster present in a pull is now authoritative, so ingest reconciles bindings from it (parity with heartbeats); an absent roster (older nodes) still skips reconciliation, preserving the F4 guard semantics. Scope: src/api (openapi + regenerated models + admin conversion), src/observability, src/bin/server.rs, services/scheduler/internal. openapi.yml edited first and models regenerated with `make agentenv-server` (JDK 17; system Java 8 is too old for the generator). Relates to #259. Co-Authored-By: Claude Code --- services/scheduler/internal/node_snapshot.go | 48 +++- .../scheduler/internal/node_snapshot_test.go | 66 +++++- src/api/generated/src/models.rs | 213 +++++++++++++++++- src/api/impls/admin.rs | 15 +- src/api/openapi.yml | 22 ++ src/bin/server.rs | 1 + src/observability/model.rs | 3 + src/observability/service.rs | 4 + 8 files changed, 350 insertions(+), 22 deletions(-) diff --git a/services/scheduler/internal/node_snapshot.go b/services/scheduler/internal/node_snapshot.go index 7509ef5f9..c6fc1e44d 100644 --- a/services/scheduler/internal/node_snapshot.go +++ b/services/scheduler/internal/node_snapshot.go @@ -146,18 +146,23 @@ func (r *ConcurrentNodeSnapshotRefresher) Refresh(ctx context.Context, nodes []N // adminNodeResponse mirrors the agentenv server admin GET /nodes entry // (openapi: Node). Field names follow the API's camelCase JSON. type adminNodeResponse struct { - Version string `json:"version"` - Commit string `json:"commit"` - ID string `json:"id"` - ServiceInstanceID string `json:"serviceInstanceID"` - ClusterID string `json:"clusterID"` - SandboxCount uint32 `json:"sandboxCount"` - CreateSuccesses uint64 `json:"createSuccesses"` - CreateFails uint64 `json:"createFails"` - SandboxStartingCnt uint32 `json:"sandboxStartingCount"` - SandboxPausedCount uint32 `json:"sandboxPausedCount"` - Status string `json:"status"` - MachineInfo struct { + Version string `json:"version"` + Commit string `json:"commit"` + ID string `json:"id"` + ServiceInstanceID string `json:"serviceInstanceID"` + ClusterID string `json:"clusterID"` + SandboxCount uint32 `json:"sandboxCount"` + CreateSuccesses uint64 `json:"createSuccesses"` + CreateFails uint64 `json:"createFails"` + SandboxStartingCnt uint32 `json:"sandboxStartingCount"` + SandboxPausedCount uint32 `json:"sandboxPausedCount"` + Status string `json:"status"` + SandboxIDs []string `json:"sandboxIDs"` + P2pEndpoint *struct { + Backend string `json:"backend"` + Address string `json:"address"` + } `json:"p2pEndpoint"` + MachineInfo struct { CPUFamily string `json:"cpuFamily"` CPUModel string `json:"cpuModel"` CPUModelName string `json:"cpuModelName"` @@ -273,6 +278,23 @@ func adminGetJSON(ctx context.Context, client *http.Client, apiKey, url string, } func reportFromAdmin(n *adminNodeResponse, fetchedAt time.Time) *nodeReport { + // The admin roster is authoritative when present (same source as + // heartbeat sandbox_ids: orchestrator.list_sandbox_ids), so a parity + // pull may reconcile bindings; an absent roster stays nil and skips + // reconciliation (#341 review). + var sandboxIDs []string + if n.SandboxIDs != nil { + sandboxIDs = []string{} + for _, id := range n.SandboxIDs { + if strings.TrimSpace(id) != "" { + sandboxIDs = append(sandboxIDs, id) + } + } + } + var p2p *schedulerv1.P2PEndpoint + if n.P2pEndpoint != nil { + p2p = &schedulerv1.P2PEndpoint{Backend: n.P2pEndpoint.Backend, Address: n.P2pEndpoint.Address} + } disks := make([]*schedulerv1.DiskMetric, 0, len(n.Metrics.Disks)) for _, d := range n.Metrics.Disks { disks = append(disks, &schedulerv1.DiskMetric{ @@ -285,6 +307,8 @@ func reportFromAdmin(n *adminNodeResponse, fetchedAt time.Time) *nodeReport { } return &nodeReport{ nodeID: n.ID, + sandboxIDs: sandboxIDs, + p2pEndpoint: p2p, clusterID: n.ClusterID, serviceInstanceID: n.ServiceInstanceID, version: n.Version, diff --git a/services/scheduler/internal/node_snapshot_test.go b/services/scheduler/internal/node_snapshot_test.go index b9cad219b..082d085cf 100644 --- a/services/scheduler/internal/node_snapshot_test.go +++ b/services/scheduler/internal/node_snapshot_test.go @@ -88,10 +88,70 @@ func TestAdminSnapshotFetcherAssemblesHeartbeatShape(t *testing.T) { if len(snap.GetDisks()) != 1 || snap.GetDisks()[0].GetDevice() != "/dev/ublkb0" { t.Fatalf("disk mapping wrong: %+v", snap.GetDisks()) } - // A pull must not carry a sandbox roster: reconciling from it would - // delete live bindings (paused sandboxes, template builds) on failover. + // Nodes without the parity fields (older admin API) yield no roster; + // ingest must skip reconciliation for them. if req.sandboxIDs != nil { - t.Fatalf("pull must leave sandboxIDs nil, got %v", req.sandboxIDs) + t.Fatalf("degraded pull must leave sandboxIDs nil, got %v", req.sandboxIDs) + } +} + +const adminParityNodeFixture = `[{ + "version": "0.2.0", + "commit": "abc123", + "id": "node-a", + "serviceInstanceID": "inst-1", + "clusterID": "cluster-1", + "sandboxCount": 2, + "createSuccesses": 10, + "createFails": 1, + "sandboxStartingCount": 1, + "sandboxPausedCount": 3, + "sandboxIDs": ["sbx-1", "sbx-paused"], + "p2pEndpoint": {"backend": "iroh", "address": "node-a.p2p.local"}, + "machineInfo": {"cpuFamily": "6", "cpuModel": "85", "cpuModelName": "Xeon", "cpuArchitecture": "x86_64", "cpuConfigJSON": "{\"ht\":true}"}, + "metrics": { + "allocatedCPU": 4, + "allocatedMemoryBytes": 8589934592, + "cpuPercent": 55, + "cpuCount": 16, + "memoryUsedBytes": 17179869184, + "memoryTotalBytes": 34359738368, + "pausedAllocatedCPU": 2, + "pausedAllocatedMemoryBytes": 4294967296, + "disks": [{"mountPoint": "/", "device": "/dev/ublkb0", "filesystemType": "ext4", "usedBytes": 1024, "totalBytes": 4096}] + } +}]` + +// Parity path (#341 review): when the admin API exposes the heartbeat fields +// (sandboxIDs, p2pEndpoint, cpuConfigJSON), the pull maps them through and +// ingest reconciles bindings from the authoritative roster. +func TestAdminSnapshotFetcherParityFields(t *testing.T) { + srv := adminTestServer(t, "", adminParityNodeFixture, "[]") + fetch := NewAdminSnapshotFetcher("") + + req, err := fetch(context.Background(), Node{ID: "node-a", Endpoint: srv.URL}) + if err != nil { + t.Fatalf("fetch failed: %v", err) + } + if len(req.sandboxIDs) != 2 || req.sandboxIDs[0] != "sbx-1" || req.sandboxIDs[1] != "sbx-paused" { + t.Fatalf("parity roster mapping wrong: %v", req.sandboxIDs) + } + if req.p2pEndpoint == nil || req.p2pEndpoint.GetBackend() != "iroh" || req.p2pEndpoint.GetAddress() != "node-a.p2p.local" { + t.Fatalf("parity p2p mapping wrong: %v", req.p2pEndpoint) + } + if req.machineInfo.GetCpuConfigJson() != `{"ht":true}` { + t.Fatalf("parity cpuConfigJSON mapping wrong: %q", req.machineInfo.GetCpuConfigJson()) + } + + svc, _, store := newTestService(t, []string{"node-a"}) + if _, err := svc.ingestNodeReport(req, time.Now()); err != nil { + t.Fatalf("parity pull ingest failed: %v", err) + } + // The authoritative roster reconciles bindings, including the paused one. + for _, id := range []string{"sbx-1", "sbx-paused"} { + if _, ok, err := store.Get(id, time.Now()); err != nil || !ok { + t.Fatalf("parity pull must reconcile binding %s: ok=%v err=%v", id, ok, err) + } } } diff --git a/src/api/generated/src/models.rs b/src/api/generated/src/models.rs index 4c31c0270..5d6389a47 100644 --- a/src/api/generated/src/models.rs +++ b/src/api/generated/src/models.rs @@ -2576,6 +2576,12 @@ pub struct MachineInfo { #[serde(rename = "cpuArchitecture")] #[validate(custom(function = "check_xss_string"))] pub cpu_architecture: String, + + /// Cluster-level CPU configuration assigned by the scheduler; empty until the node reports one + #[serde(rename = "cpuConfigJSON")] + #[validate(custom(function = "check_xss_string"))] + #[serde(skip_serializing_if = "Option::is_none")] + pub cpu_config_json: Option, } impl MachineInfo { @@ -2591,6 +2597,7 @@ impl MachineInfo { cpu_model, cpu_model_name, cpu_architecture, + cpu_config_json: None, } } } @@ -2609,6 +2616,9 @@ impl std::fmt::Display for MachineInfo { Some(self.cpu_model_name.to_string()), Some("cpuArchitecture".to_string()), Some(self.cpu_architecture.to_string()), + self.cpu_config_json.as_ref().map(|cpu_config_json| { + ["cpuConfigJSON".to_string(), cpu_config_json.to_string()].join(",") + }), ]; write!( @@ -2634,6 +2644,7 @@ impl std::str::FromStr for MachineInfo { pub cpu_model: Vec, pub cpu_model_name: Vec, pub cpu_architecture: Vec, + pub cpu_config_json: Vec, } let mut intermediate_rep = IntermediateRep::default(); @@ -2671,6 +2682,10 @@ impl std::str::FromStr for MachineInfo { "cpuArchitecture" => intermediate_rep.cpu_architecture.push( ::from_str(val).map_err(|x| x.to_string())?, ), + #[allow(clippy::redundant_clone)] + "cpuConfigJSON" => intermediate_rep.cpu_config_json.push( + ::from_str(val).map_err(|x| x.to_string())?, + ), _ => { return std::result::Result::Err( "Unexpected key while parsing MachineInfo".to_string(), @@ -2705,6 +2720,7 @@ impl std::str::FromStr for MachineInfo { .into_iter() .next() .ok_or_else(|| "cpuArchitecture missing in MachineInfo".to_string())?, + cpu_config_json: intermediate_rep.cpu_config_json.into_iter().next(), }) } } @@ -4117,6 +4133,17 @@ pub struct Node { #[validate(nested)] pub machine_info: models::MachineInfo, + /// IDs of all tracked sandboxes on the node, including paused sandboxes and template builds + #[serde(rename = "sandboxIDs")] + #[validate(custom(function = "check_xss_vec_string"))] + #[serde(skip_serializing_if = "Option::is_none")] + pub sandbox_ids: Option>, + + #[serde(rename = "p2pEndpoint")] + #[validate(nested)] + #[serde(skip_serializing_if = "Option::is_none")] + pub p2p_endpoint: Option, + #[serde(rename = "status")] #[validate(nested)] pub status: models::NodeStatus, @@ -4170,6 +4197,8 @@ impl Node { service_instance_id, cluster_id, machine_info, + sandbox_ids: None, + p2p_endpoint: None, status, sandbox_count, metrics, @@ -4198,6 +4227,18 @@ impl std::fmt::Display for Node { Some("clusterID".to_string()), Some(self.cluster_id.to_string()), // Skipping machineInfo in query parameter serialization + self.sandbox_ids.as_ref().map(|sandbox_ids| { + [ + "sandboxIDs".to_string(), + sandbox_ids + .iter() + .map(|x| x.to_string()) + .collect::>() + .join(","), + ] + .join(",") + }), + // Skipping p2pEndpoint in query parameter serialization // Skipping status in query parameter serialization Some("sandboxCount".to_string()), @@ -4238,6 +4279,8 @@ impl std::str::FromStr for Node { pub service_instance_id: Vec, pub cluster_id: Vec, pub machine_info: Vec, + pub sandbox_ids: Vec>, + pub p2p_endpoint: Vec, pub status: Vec, pub sandbox_count: Vec, pub metrics: Vec, @@ -4257,9 +4300,7 @@ impl std::str::FromStr for Node { let val = match string_iter.next() { Some(x) => x, None => { - return std::result::Result::Err( - "Missing value while parsing Node".to_string(), - ); + return std::result::Result::Err("Missing value while parsing Node".to_string()); } }; @@ -4291,6 +4332,17 @@ impl std::str::FromStr for Node { ::from_str(val) .map_err(|x| x.to_string())?, ), + "sandboxIDs" => { + return std::result::Result::Err( + "Parsing a container in this style is not supported in Node" + .to_string(), + ); + } + #[allow(clippy::redundant_clone)] + "p2pEndpoint" => intermediate_rep.p2p_endpoint.push( + ::from_str(val) + .map_err(|x| x.to_string())?, + ), #[allow(clippy::redundant_clone)] "status" => intermediate_rep.status.push( ::from_str(val) @@ -4365,6 +4417,8 @@ impl std::str::FromStr for Node { .into_iter() .next() .ok_or_else(|| "machineInfo missing in Node".to_string())?, + sandbox_ids: intermediate_rep.sandbox_ids.into_iter().next(), + p2p_endpoint: intermediate_rep.p2p_endpoint.into_iter().next(), status: intermediate_rep .status .into_iter() @@ -5171,6 +5225,159 @@ impl std::str::FromStr for OrderDirection { } } +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, validator::Validate)] +#[cfg_attr(feature = "conversion", derive(frunk::LabelledGeneric))] +pub struct P2pEndpoint { + /// Transport backend that understands the address (e.g. iroh) + #[serde(rename = "backend")] + #[validate(custom(function = "check_xss_string"))] + pub backend: String, + + /// Backend-specific serialized endpoint address + #[serde(rename = "address")] + #[validate(custom(function = "check_xss_string"))] + pub address: String, +} + +impl P2pEndpoint { + #[allow(clippy::new_without_default, clippy::too_many_arguments)] + pub fn new(backend: String, address: String) -> P2pEndpoint { + P2pEndpoint { backend, address } + } +} + +/// Converts the P2pEndpoint value to the Query Parameters representation (style=form, explode=false) +/// specified in https://swagger.io/docs/specification/serialization/ +/// Should be implemented in a serde serializer +impl std::fmt::Display for P2pEndpoint { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let params: Vec> = vec![ + Some("backend".to_string()), + Some(self.backend.to_string()), + Some("address".to_string()), + Some(self.address.to_string()), + ]; + + write!( + f, + "{}", + params.into_iter().flatten().collect::>().join(",") + ) + } +} + +/// Converts Query Parameters representation (style=form, explode=false) to a P2pEndpoint value +/// as specified in https://swagger.io/docs/specification/serialization/ +/// Should be implemented in a serde deserializer +impl std::str::FromStr for P2pEndpoint { + type Err = String; + + fn from_str(s: &str) -> std::result::Result { + /// An intermediate representation of the struct to use for parsing. + #[derive(Default)] + #[allow(dead_code)] + struct IntermediateRep { + pub backend: Vec, + pub address: Vec, + } + + let mut intermediate_rep = IntermediateRep::default(); + + // Parse into intermediate representation + let mut string_iter = s.split(','); + let mut key_result = string_iter.next(); + + while key_result.is_some() { + let val = match string_iter.next() { + Some(x) => x, + None => { + return std::result::Result::Err( + "Missing value while parsing P2pEndpoint".to_string(), + ); + } + }; + + if let Some(key) = key_result { + #[allow(clippy::match_single_binding)] + match key { + #[allow(clippy::redundant_clone)] + "backend" => intermediate_rep.backend.push( + ::from_str(val).map_err(|x| x.to_string())?, + ), + #[allow(clippy::redundant_clone)] + "address" => intermediate_rep.address.push( + ::from_str(val).map_err(|x| x.to_string())?, + ), + _ => { + return std::result::Result::Err( + "Unexpected key while parsing P2pEndpoint".to_string(), + ); + } + } + } + + // Get the next key + key_result = string_iter.next(); + } + + // Use the intermediate representation to return the struct + std::result::Result::Ok(P2pEndpoint { + backend: intermediate_rep + .backend + .into_iter() + .next() + .ok_or_else(|| "backend missing in P2pEndpoint".to_string())?, + address: intermediate_rep + .address + .into_iter() + .next() + .ok_or_else(|| "address missing in P2pEndpoint".to_string())?, + }) + } +} + +// Methods for converting between header::IntoHeaderValue and HeaderValue + +#[cfg(feature = "server")] +impl std::convert::TryFrom> for HeaderValue { + type Error = String; + + fn try_from( + hdr_value: header::IntoHeaderValue, + ) -> std::result::Result { + let hdr_value = hdr_value.to_string(); + match HeaderValue::from_str(&hdr_value) { + std::result::Result::Ok(value) => std::result::Result::Ok(value), + std::result::Result::Err(e) => std::result::Result::Err(format!( + r#"Invalid header value for P2pEndpoint - value: {hdr_value} is invalid {e}"# + )), + } + } +} + +#[cfg(feature = "server")] +impl std::convert::TryFrom for header::IntoHeaderValue { + type Error = String; + + fn try_from(hdr_value: HeaderValue) -> std::result::Result { + match hdr_value.to_str() { + std::result::Result::Ok(value) => { + match ::from_str(value) { + std::result::Result::Ok(value) => { + std::result::Result::Ok(header::IntoHeaderValue(value)) + } + std::result::Result::Err(err) => std::result::Result::Err(format!( + r#"Unable to convert header value '{value}' into P2pEndpoint - {err}"# + )), + } + } + std::result::Result::Err(e) => std::result::Result::Err(format!( + r#"Unable to convert header: {hdr_value:?} to string: {e}"# + )), + } + } +} + #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, validator::Validate)] #[cfg_attr(feature = "conversion", derive(frunk::LabelledGeneric))] pub struct ResumedSandbox { diff --git a/src/api/impls/admin.rs b/src/api/impls/admin.rs index 6ea7034ca..524be8953 100644 --- a/src/api/impls/admin.rs +++ b/src/api/impls/admin.rs @@ -10,12 +10,14 @@ use super::ApiImpl; impl From for models::MachineInfo { fn from(machine_info: MachineInfo) -> Self { - models::MachineInfo::new( + let mut out = models::MachineInfo::new( machine_info.cpu_family, machine_info.cpu_model, machine_info.cpu_model_name, machine_info.cpu_architecture, - ) + ); + out.cpu_config_json = machine_info.cpu_config_json; + out } } @@ -53,7 +55,7 @@ impl From for models::NodeMetrics { impl From for models::Node { fn from(node: NodeSnapshot) -> Self { - models::Node::new( + let mut out = models::Node::new( node.version, node.commit, node.node_id, @@ -67,7 +69,12 @@ impl From for models::Node { node.create_fails, node.sandbox_starting_count, node.paused_sandbox_count, - ) + ); + out.sandbox_ids = Some(node.sandbox_ids.iter().map(|id| id.to_string()).collect()); + out.p2p_endpoint = node + .p2p_endpoint + .map(|ep| models::P2pEndpoint::new(ep.backend, ep.address)); + out } } diff --git a/src/api/openapi.yml b/src/api/openapi.yml index ec391ed49..d46fcae09 100644 --- a/src/api/openapi.yml +++ b/src/api/openapi.yml @@ -1341,6 +1341,18 @@ components: type: integer format: uint64 description: Sum of memory reservations (bytes) across all sandboxes currently in the Paused state + P2pEndpoint: + required: + - backend + - address + properties: + backend: + type: string + description: Transport backend that understands the address (e.g. iroh) + address: + type: string + description: Backend-specific serialized endpoint address + MachineInfo: required: - cpuFamily @@ -1360,6 +1372,9 @@ components: cpuArchitecture: type: string description: CPU architecture of the node + cpuConfigJSON: + type: string + description: Cluster-level CPU configuration assigned by the scheduler; empty until the node reports one Node: required: @@ -1394,6 +1409,13 @@ components: description: Identifier of the cluster machineInfo: $ref: "#/components/schemas/MachineInfo" + sandboxIDs: + type: array + items: + type: string + description: IDs of all tracked sandboxes on the node, including paused sandboxes and template builds + p2pEndpoint: + $ref: "#/components/schemas/P2pEndpoint" status: $ref: "#/components/schemas/NodeStatus" sandboxCount: diff --git a/src/bin/server.rs b/src/bin/server.rs index 81b544d3b..bd6a48057 100644 --- a/src/bin/server.rs +++ b/src/bin/server.rs @@ -139,6 +139,7 @@ async fn main() -> anyhow::Result<()> { Arc::clone(&orchestrator), config.resolved_cpu_template_helper(), cluster_cpu_arc, + p2p_local_endpoint.clone(), ) .await, )) diff --git a/src/observability/model.rs b/src/observability/model.rs index 3d4684e3e..aff11d7c4 100644 --- a/src/observability/model.rs +++ b/src/observability/model.rs @@ -1,4 +1,5 @@ use super::DiskMetric; +use crate::p2p::P2pEndpoint; use crate::types::SandboxId; /// Static machine descriptors reported as part of node observability. @@ -48,4 +49,6 @@ pub struct NodeSnapshot { pub sandbox_starting_count: u32, /// Number of sandboxes currently in the Paused state on this node. pub paused_sandbox_count: u32, + /// The node's P2P artifact transport endpoint, when P2P is enabled. + pub p2p_endpoint: Option, } diff --git a/src/observability/service.rs b/src/observability/service.rs index eeca27590..7c6935760 100644 --- a/src/observability/service.rs +++ b/src/observability/service.rs @@ -28,6 +28,7 @@ pub struct ObservabilityService { host_metrics: HostMetricsCollector, pending_cpu_config: Arc>>, cluster_cpu_config: Arc>>, + p2p_endpoint: Option, } impl ObservabilityService { @@ -36,6 +37,7 @@ impl ObservabilityService { orchestrator: Arc, cpu_template_helper: Option, cluster_cpu_arc: Arc>>, + p2p_endpoint: Option, ) -> Self { let machine_info = detect_machine_info(); let host_metrics = HostMetricsCollector::new(); @@ -51,6 +53,7 @@ impl ObservabilityService { host_metrics, pending_cpu_config: Arc::new(Mutex::new(cpu_config_json)), cluster_cpu_config: cluster_cpu_arc, + p2p_endpoint, } } @@ -95,6 +98,7 @@ impl ObservabilityService { create_fails: runtime.create_fails, sandbox_starting_count: runtime.starting_sandbox_count, paused_sandbox_count: runtime.paused_sandbox_count, + p2p_endpoint: self.p2p_endpoint.clone(), }) } From d7c21ae53efd275b65ee2b49bc2551c8da5a0f17 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Fri, 9 Oct 2026 23:25:56 +0800 Subject: [PATCH 17/26] test(scheduler): Kind-based HA failover e2e harness with four scenarios MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address the maintainer's requirement that failover tests run for real instead of staying skipped (#341 review, R3). Adds a Kind-based harness plus a GitHub Actions workflow that runs it on push/PR: - stubnode: a minimal agentenv-node double serving the admin /nodes API with heartbeat-parity fields (sandboxIDs, p2pEndpoint, cpuConfigJSON), sending periodic gRPC heartbeats, and exposing a /control plane for metrics, rosters, and heartbeat pauses. Distroless image built from services/scheduler/e2e/Dockerfile.stubnode. - k8s manifests: redis, scheduler x3 with leader election (6s/4s/1s timings, leader-readiness probe, Lease RBAC), and two stub nodes behind the agentenv-nodes headless Service the discovery watches. - harness_test.go (build tag e2e), driving kubectl and the scheduler gRPC API: T1 kill the leader pod — takeover within the lease budget, pre-seeded bindings keep resolving (no NotFound mid-failover), scheduling resumes on freshly observed nodes; T2 partition the leader from the API with a pod-targeted NetworkPolicy egress block (distroless-safe, no exec) — standby takes over, and after restoring egress the ex-leader has force-stopped and its pod restarted, never a second primary; T3 with node heartbeats paused, scheduling right after takeover returns Unavailable (never guesses capacity), and recovers once heartbeats resume; T4 a write hitting the demoted leader mid-takeover never commits; the client retries and succeeds on the new leader. The workflow builds both images, loads them into Kind, deploys the harness, waits for a leader, then runs the tests with a port-forward and dumps diagnostics on failure. Relates to #259. Co-Authored-By: Claude Code --- .github/workflows/scheduler-ha-e2e.yml | 81 ++++++ services/scheduler/e2e/Dockerfile.stubnode | 12 + services/scheduler/e2e/harness_test.go | 281 +++++++++++++++++++++ services/scheduler/e2e/k8s/redis.yaml | 30 +++ services/scheduler/e2e/k8s/scheduler.yaml | 127 ++++++++++ services/scheduler/e2e/k8s/stubnodes.yaml | 50 ++++ services/scheduler/e2e/stubnode/main.go | 237 +++++++++++++++++ 7 files changed, 818 insertions(+) create mode 100644 .github/workflows/scheduler-ha-e2e.yml create mode 100644 services/scheduler/e2e/Dockerfile.stubnode create mode 100644 services/scheduler/e2e/harness_test.go create mode 100644 services/scheduler/e2e/k8s/redis.yaml create mode 100644 services/scheduler/e2e/k8s/scheduler.yaml create mode 100644 services/scheduler/e2e/k8s/stubnodes.yaml create mode 100644 services/scheduler/e2e/stubnode/main.go diff --git a/.github/workflows/scheduler-ha-e2e.yml b/.github/workflows/scheduler-ha-e2e.yml new file mode 100644 index 000000000..66c962d60 --- /dev/null +++ b/.github/workflows/scheduler-ha-e2e.yml @@ -0,0 +1,81 @@ +name: scheduler-ha-e2e + +on: + push: + paths: + - "services/scheduler/**" + - "services/shared/**" + - "services/api/**" + - "deploy/k8s/**" + - ".github/workflows/scheduler-ha-e2e.yml" + pull_request: + paths: + - "services/scheduler/**" + - "services/shared/**" + - "services/api/**" + - "deploy/k8s/**" + - ".github/workflows/scheduler-ha-e2e.yml" + +jobs: + failover: + runs-on: ubuntu-latest + timeout-minutes: 25 + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-go@v5 + with: + go-version: "1.25" + + - name: Create Kind cluster + uses: helm/kind-action@v1 + with: + cluster_name: scheduler-ha-e2e + + - name: Build scheduler image + run: docker build -f deploy/docker/Dockerfile.scheduler -t agentenv-scheduler:e2e . + + - name: Build stubnode image + run: docker build -f services/scheduler/e2e/Dockerfile.stubnode -t agentenv-stubnode:e2e . + + - name: Load images into Kind + run: | + kind load docker-image agentenv-scheduler:e2e --name scheduler-ha-e2e + kind load docker-image agentenv-stubnode:e2e --name scheduler-ha-e2e + + - name: Deploy harness + run: | + kubectl apply -f services/scheduler/e2e/k8s/redis.yaml + kubectl apply -f services/scheduler/e2e/k8s/scheduler.yaml + kubectl apply -f services/scheduler/e2e/k8s/stubnodes.yaml + kubectl rollout status deployment/redis --timeout=120s + kubectl rollout status deployment/stubnode --timeout=120s + kubectl rollout status deployment/agentenv-scheduler --timeout=180s + + - name: Wait for a leader + run: | + for i in $(seq 1 60); do + HOLDER=$(kubectl get lease agentenv-scheduler -o jsonpath='{.spec.holderIdentity}' 2>/dev/null || true) + if [ -n "$HOLDER" ]; then echo "leader: $HOLDER"; exit 0; fi + sleep 2 + done + echo "no leader elected in time" >&2 + kubectl get pods + exit 1 + + - name: Run failover tests + run: | + kubectl port-forward svc/agentenv-scheduler 19090:9090 & + PF_PID=$! + trap "kill $PF_PID 2>/dev/null || true" EXIT + sleep 3 + cd services + go test -tags e2e -count=1 -timeout 10m -v ./scheduler/e2e/... + + - name: Dump diagnostics on failure + if: failure() + run: | + kubectl get pods -o wide || true + kubectl describe lease agentenv-scheduler || true + kubectl logs -l app=agentenv-scheduler --tail=100 --prefix=true || true + kubectl logs -l app=stubnode --tail=50 --prefix=true || true diff --git a/services/scheduler/e2e/Dockerfile.stubnode b/services/scheduler/e2e/Dockerfile.stubnode new file mode 100644 index 000000000..626847531 --- /dev/null +++ b/services/scheduler/e2e/Dockerfile.stubnode @@ -0,0 +1,12 @@ +FROM golang:1.25-bookworm AS build +ARG TARGETOS=linux +ARG TARGETARCH +WORKDIR /src +COPY services/go.mod services/go.sum ./ +RUN go mod download +COPY services/ . +RUN CGO_ENABLED=0 GOOS=${TARGETOS} GOARCH=${TARGETARCH} go build -o /out/stubnode ./scheduler/e2e/stubnode + +FROM gcr.io/distroless/static-debian12 +COPY --from=build /out/stubnode /stubnode +ENTRYPOINT ["/stubnode"] diff --git a/services/scheduler/e2e/harness_test.go b/services/scheduler/e2e/harness_test.go new file mode 100644 index 000000000..f85de7900 --- /dev/null +++ b/services/scheduler/e2e/harness_test.go @@ -0,0 +1,281 @@ +//go:build e2e + +// Scheduler HA failover e2e tests (#259, review requirement R3). +// Runs against a Kind cluster with: scheduler x3 (leader election on, +// Redis bindings) and stub nodes (admin /nodes + heartbeats). The workflow +// builds and loads the images and applies services/scheduler/e2e/k8s. +package e2e + +import ( + "context" + "fmt" + "os" + "os/exec" + "strings" + "testing" + "time" + + schedulerv1 "agentenv/services/api/proto" + + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/status" +) + +var schedAddr = envOr("E2E_SCHED_ADDR", "127.0.0.1:19090") + +func writeFile(t *testing.T, path, content string) { + t.Helper() + if err := os.WriteFile(path, []byte(content), 0o644); err != nil { + t.Fatalf("write %s: %v", path, err) + } +} + +func envOr(key, fallback string) string { + if v := os.Getenv(key); v != "" { + return v + } + return fallback +} + +func kube(t *testing.T, args ...string) string { + t.Helper() + out, err := exec.Command("kubectl", args...).CombinedOutput() + if err != nil { + t.Fatalf("kubectl %s failed: %v\n%s", strings.Join(args, " "), err, out) + } + return strings.TrimSpace(string(out)) +} + +func dial(t *testing.T) schedulerv1.SchedulerClient { + t.Helper() + conn, err := grpc.NewClient(schedAddr, grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("dial %s: %v", schedAddr, err) + } + t.Cleanup(func() { conn.Close() }) + return schedulerv1.NewSchedulerClient(conn) +} + +func eventually(t *testing.T, timeout time.Duration, what string, cond func() (bool, error)) { + t.Helper() + deadline := time.Now().Add(timeout) + var lastErr error + for time.Now().Before(deadline) { + ok, err := cond() + if err == nil && ok { + return + } + lastErr = err + time.Sleep(300 * time.Millisecond) + } + t.Fatalf("timed out waiting for %s (last error: %v)", what, lastErr) +} + +func leaderIdentity(t *testing.T) string { + t.Helper() + out := kube(t, "get", "lease", "agentenv-scheduler", "-o", "jsonpath={.spec.holderIdentity}") + return strings.TrimSpace(out) +} + +func schedulerEndpoints(t *testing.T) []string { + t.Helper() + out := kube(t, "get", "endpoints", "agentenv-scheduler", "-o", "jsonpath={.subsets[*].addresses[*].targetRef.name}") + return strings.Fields(out) +} + +// T1: kill the leader pod — a standby must take over within the lease +// budget, pre-existing bindings keep resolving through Redis, and +// scheduling resumes on freshly observed nodes. +func TestFailoverLeaderKill(t *testing.T) { + client := dial(t) + ctx := context.Background() + + if _, err := client.RecordAssignment(ctx, &schedulerv1.RecordAssignmentRequest{ + SandboxId: "sbx-e2e-1", + Node: firstDiscoveredNode(t, client), + }); err != nil { + t.Fatalf("seed binding failed: %v", err) + } + + // Continuous lookup during the failover: must never fail with NotFound. + lookupErrs := make(chan error, 1) + stop := make(chan struct{}) + go func() { + for { + select { + case <-stop: + return + default: + } + _, err := client.LookupNode(ctx, &schedulerv1.LookupNodeRequest{SandboxId: "sbx-e2e-1"}) + if err != nil { + if s, ok := status.FromError(err); ok && strings.Contains(s.Message(), "not found") { + lookupErrs <- fmt.Errorf("lookup returned not found mid-failover: %v", err) + return + } + } + time.Sleep(100 * time.Millisecond) + } + }() + defer close(stop) + + victim := leaderIdentity(t) + t.Logf("killing leader pod %s", victim) + kube(t, "delete", "pod", victim, "--force", "--grace-period=0") + + eventually(t, 30*time.Second, "a new leader", func() (bool, error) { + cur := leaderIdentity(t) + return cur != "" && cur != victim, nil + }) + eventually(t, 30*time.Second, "endpoints converge on one leader", func() (bool, error) { + return len(schedulerEndpoints(t)) == 1, nil + }) + + select { + case err := <-lookupErrs: + t.Fatal(err) + case <-time.After(2 * time.Second): + } + + eventually(t, 30*time.Second, "scheduling works again on the new leader", func() (bool, error) { + _, err := client.Schedule(ctx, &schedulerv1.ScheduleRequest{}) + return err == nil, nil + }) +} + +// T2: freeze the leader (SIGSTOP = API partition equivalent) — the standby +// must take over, and the frozen ex-leader must exit on resume instead of +// serving as a second primary. +func TestPartitionFrozenLeader(t *testing.T) { + victim := leaderIdentity(t) + t.Logf("partitioning leader pod %s from the API (egress block)", victim) + kube(t, "label", "pod", victim, "e2e-partition=true", "--overwrite") + writeFile(t, "/tmp/e2e-netpol.yaml", `apiVersion: networking.k8s.io/v1 +kind: NetworkPolicy +metadata: + name: e2e-partition +spec: + podSelector: + matchLabels: + e2e-partition: "true" + policyTypes: ["Egress"] + egress: [] +`) + kube(t, "apply", "-f", "/tmp/e2e-netpol.yaml") + t.Cleanup(func() { kube(t, "delete", "networkpolicy", "e2e-partition", "--ignore-not-found") }) + + eventually(t, 45*time.Second, "standby takes over while old leader partitioned", func() (bool, error) { + cur := leaderIdentity(t) + return cur != "" && cur != victim, nil + }) + + // While partitioned, there must be exactly one serving endpoint. + eps := schedulerEndpoints(t) + if len(eps) != 1 { + t.Fatalf("expected exactly one endpoint during partition, got %v", eps) + } + + kube(t, "delete", "networkpolicy", "e2e-partition") + // With egress restored, the ex-leader's renew failure has already fired + // OnStoppedLeading: it force-stops and its pod restarts as standby. + eventually(t, 60*time.Second, "frozen ex-leader exits (pod restarts)", func() (bool, error) { + out := kube(t, "get", "pod", victim, "-o", "jsonpath={.status.containerStatuses[0].restartCount}") + return out != "0", nil + }) + eventually(t, 30*time.Second, "still exactly one leader after resume", func() (bool, error) { + return len(schedulerEndpoints(t)) == 1, nil + }) +} + +// T3: right after takeover, with node heartbeats paused, scheduling must +// return Unavailable (never fall back to unobserved nodes); once heartbeats +// resume, scheduling recovers. +func TestTakeoverSchedulingSemantics(t *testing.T) { + client := dial(t) + ctx := context.Background() + + stubPost(t, "control/pause", "") + victim := leaderIdentity(t) + kube(t, "delete", "pod", victim, "--force", "--grace-period=0") + + eventually(t, 30*time.Second, "new leader", func() (bool, error) { + cur := leaderIdentity(t) + return cur != "" && cur != victim, nil + }) + + // No node has freshly reported to the new leader: Unavailable, not a guess. + if _, err := client.Schedule(ctx, &schedulerv1.ScheduleRequest{}); err == nil { + t.Fatal("schedule must fail with Unavailable before fresh observations, not guess capacity") + } + + stubPost(t, "control/resume", "") + eventually(t, 30*time.Second, "scheduling recovers after heartbeats resume", func() (bool, error) { + _, err := client.Schedule(ctx, &schedulerv1.ScheduleRequest{}) + return err == nil, nil + }) +} + +// T4: a write that lands on the demoted leader mid-takeover must not +// commit after the takeover; the client retries and succeeds on the new +// leader. +func TestDemotedLeaderDelayedWrite(t *testing.T) { + client := dial(t) + ctx := context.Background() + + victim := leaderIdentity(t) + done := make(chan error, 1) + go func() { + _, err := client.RecordAssignment(ctx, &schedulerv1.RecordAssignmentRequest{ + SandboxId: "sbx-e2e-t4", + Node: firstDiscoveredNode(t, client), + }) + done <- err + }() + + kube(t, "delete", "pod", victim, "--force", "--grace-period=0") + writeErr := <-done + t.Logf("write during takeover returned: %v", writeErr) + + eventually(t, 30*time.Second, "new leader", func() (bool, error) { + cur := leaderIdentity(t) + return cur != "" && cur != victim, nil + }) + + // The client-side retry path: writing again must succeed on the new leader. + if _, err := client.RecordAssignment(ctx, &schedulerv1.RecordAssignmentRequest{ + SandboxId: "sbx-e2e-t4", + Node: firstDiscoveredNode(t, client), + }); err != nil { + t.Fatalf("write retry on the new leader failed: %v", err) + } + if _, err := client.LookupNode(ctx, &schedulerv1.LookupNodeRequest{SandboxId: "sbx-e2e-t4"}); err != nil { + t.Fatalf("binding must exist after retry, got %v", err) + } +} + +// firstDiscoveredNode resolves a node through the scheduler's own discovery, +// so RecordAssignment passes the service's known-node validation. +func firstDiscoveredNode(t *testing.T, client schedulerv1.SchedulerClient) *schedulerv1.Node { + t.Helper() + resp, err := client.ListNodes(context.Background(), &schedulerv1.ListNodesRequest{}) + if err != nil || len(resp.GetNodes()) == 0 { + t.Fatalf("ListNodes must return discovered stub nodes: %v", err) + } + return resp.GetNodes()[0] +} + +// stubPost hits every stub's control plane through an ephemeral busybox +// debug container (the stub image is distroless, so kubectl exec into the +// container itself is not possible). +func stubPost(t *testing.T, path, body string) { + t.Helper() + pods := kube(t, "get", "pods", "-l", "app=stubnode", "-o", "jsonpath={.items[*].metadata.name}") + for _, pod := range strings.Fields(pods) { + out, err := exec.Command("kubectl", "debug", "-q", pod, "--image=busybox:1.36", + "--", "wget", "-q", "-O", "-", "--post-data", body, "http://127.0.0.1:8000/"+path).CombinedOutput() + if err != nil { + t.Fatalf("stub control %s on %s failed: %v\n%s", path, pod, err, out) + } + } +} diff --git a/services/scheduler/e2e/k8s/redis.yaml b/services/scheduler/e2e/k8s/redis.yaml new file mode 100644 index 000000000..a4c529ef3 --- /dev/null +++ b/services/scheduler/e2e/k8s/redis.yaml @@ -0,0 +1,30 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: redis +spec: + replicas: 1 + selector: + matchLabels: + app: redis + template: + metadata: + labels: + app: redis + spec: + containers: + - name: redis + image: redis:7-alpine + ports: + - containerPort: 6379 +--- +apiVersion: v1 +kind: Service +metadata: + name: redis +spec: + selector: + app: redis + ports: + - port: 6379 + targetPort: 6379 diff --git a/services/scheduler/e2e/k8s/scheduler.yaml b/services/scheduler/e2e/k8s/scheduler.yaml new file mode 100644 index 000000000..80fa62e51 --- /dev/null +++ b/services/scheduler/e2e/k8s/scheduler.yaml @@ -0,0 +1,127 @@ +apiVersion: v1 +kind: ServiceAccount +metadata: + name: agentenv-scheduler +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: agentenv-scheduler +rules: + - apiGroups: ["coordination.k8s.io"] + resources: ["leases"] + verbs: ["get", "create", "update"] + - apiGroups: ["discovery.k8s.io"] + resources: ["endpointslices"] + verbs: ["get", "list", "watch"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: agentenv-scheduler +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: Role + name: agentenv-scheduler +subjects: + - kind: ServiceAccount + name: agentenv-scheduler +--- +apiVersion: v1 +kind: ConfigMap +metadata: + name: scheduler-config +data: + scheduler.json: | + { + "log_level": "info", + "log_format": "json", + "scheduler": { + "grpc_listen_addr": ":9090", + "metrics_listen_addr": ":9091", + "strategy": "round_robin", + "report_ttl": "10s", + "binding_ttl": "60s", + "discovery": { + "mode": "kubernetes", + "kubernetes": { + "namespace": "default", + "service_name": "agentenv-nodes", + "port": 8000, + "scheme": "http" + } + }, + "leader_election": { + "enabled": true, + "lease_name": "agentenv-scheduler", + "lease_namespace": "default", + "lease_duration": "6s", + "renew_deadline": "4s", + "retry_period": "1s" + } + } + } +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: agentenv-scheduler +spec: + replicas: 3 + selector: + matchLabels: + app: agentenv-scheduler + template: + metadata: + labels: + app: agentenv-scheduler + spec: + serviceAccountName: agentenv-scheduler + containers: + - name: scheduler + image: agentenv-scheduler:e2e + imagePullPolicy: IfNotPresent + args: ["-config", "/config/scheduler.json"] + env: + - name: SCHEDULER_REDIS_ADDR + value: redis:6379 + - name: SCHEDULER_NODE_ADMIN_API_KEY + value: e2e-admin-key + ports: + - name: grpc + containerPort: 9090 + readinessProbe: + exec: + command: + - /grpc_health_probe + - -addr=127.0.0.1:9090 + - -service=scheduler.v1.Scheduler/leader + initialDelaySeconds: 2 + periodSeconds: 2 + livenessProbe: + exec: + command: + - /grpc_health_probe + - -addr=127.0.0.1:9090 + initialDelaySeconds: 5 + periodSeconds: 10 + volumeMounts: + - name: config + mountPath: /config/scheduler.json + subPath: scheduler.json + readOnly: true + volumes: + - name: config + configMap: + name: scheduler-config +--- +apiVersion: v1 +kind: Service +metadata: + name: agentenv-scheduler +spec: + selector: + app: agentenv-scheduler + ports: + - port: 9090 + targetPort: 9090 diff --git a/services/scheduler/e2e/k8s/stubnodes.yaml b/services/scheduler/e2e/k8s/stubnodes.yaml new file mode 100644 index 000000000..23ec6ebfe --- /dev/null +++ b/services/scheduler/e2e/k8s/stubnodes.yaml @@ -0,0 +1,50 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: stubnode +spec: + replicas: 2 + selector: + matchLabels: + app: stubnode + template: + metadata: + labels: + app: stubnode + spec: + containers: + - name: stubnode + image: agentenv-stubnode:e2e + imagePullPolicy: IfNotPresent + env: + - name: AENV_API_KEY + value: e2e-admin-key + - name: NODE_ID + valueFrom: + fieldRef: + fieldPath: metadata.name + - name: SCHEDULER_ADDR + value: agentenv-scheduler:9090 + args: ["-interval", "2s"] + ports: + - containerPort: 8000 + readinessProbe: + httpGet: + path: /healthz + port: 8000 + initialDelaySeconds: 1 + periodSeconds: 2 +--- +apiVersion: v1 +kind: Service +metadata: + name: agentenv-nodes + labels: + kubernetes.io/service-name: agentenv-nodes +spec: + clusterIP: None + selector: + app: stubnode + ports: + - port: 8000 + targetPort: 8000 diff --git a/services/scheduler/e2e/stubnode/main.go b/services/scheduler/e2e/stubnode/main.go new file mode 100644 index 000000000..0c6854417 --- /dev/null +++ b/services/scheduler/e2e/stubnode/main.go @@ -0,0 +1,237 @@ +// stubnode is a minimal agentenv-node double for scheduler HA e2e tests. +// It serves the admin HTTP API the leader pulls (with heartbeat-parity +// fields), sends periodic gRPC heartbeats, and exposes a /control plane +// the test driver uses to inject metrics, sandbox rosters, and heartbeat +// pauses. +package main + +import ( + "context" + "encoding/json" + "flag" + "fmt" + "log" + "net/http" + "os" + "sync" + "time" + + schedulerv1 "agentenv/services/api/proto" + + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" +) + +type stubState struct { + mu sync.Mutex + sandboxIDs []string + sandboxCount uint32 + cpuPercent uint32 + pauseHeartbeats bool +} + +type stub struct { + id string + clusterID string + apiKey string + state *stubState + client schedulerv1.SchedulerClient +} + +func main() { + id := flag.String("id", mustEnv("NODE_ID", "stub-1"), "node id") + listen := flag.String("listen", ":8000", "admin HTTP listen address") + schedulerAddr := flag.String("scheduler", mustEnv("SCHEDULER_ADDR", "agentenv-scheduler:9090"), "scheduler gRPC address") + clusterID := flag.String("cluster-id", "e2e", "cluster id") + apiKey := flag.String("api-key", os.Getenv("AENV_API_KEY"), "admin API key for /nodes and /sandboxes") + interval := flag.Duration("interval", 5*time.Second, "heartbeat interval") + flag.Parse() + + conn, err := grpc.NewClient(*schedulerAddr, grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + log.Fatalf("dial scheduler: %v", err) + } + defer conn.Close() + + s := &stub{ + id: *id, + clusterID: *clusterID, + apiKey: *apiKey, + state: &stubState{sandboxIDs: []string{}, cpuPercent: 5}, + client: schedulerv1.NewSchedulerClient(conn), + } + + go s.heartbeatLoop(*interval) + + mux := http.NewServeMux() + mux.HandleFunc("GET /nodes", s.withAuth(s.handleNodes)) + mux.HandleFunc("GET /sandboxes", s.withAuth(s.handleSandboxes)) + mux.HandleFunc("POST /control/metrics", s.handleSetMetrics) + mux.HandleFunc("POST /control/sandboxes", s.handleSetSandboxes) + mux.HandleFunc("POST /control/pause", s.handlePause) + mux.HandleFunc("POST /control/resume", s.handleResume) + mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) }) + + log.Printf("stubnode %s listening on %s, heartbeating to %s", *id, *listen, *schedulerAddr) + log.Fatal(http.ListenAndServe(*listen, mux)) +} + +func mustEnv(key, fallback string) string { + if v := os.Getenv(key); v != "" { + return v + } + return fallback +} + +func (s *stub) withAuth(next http.HandlerFunc) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + if s.apiKey != "" && r.Header.Get("x-api-key") != s.apiKey { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + next(w, r) + } +} + +func (s *stub) heartbeatLoop(interval time.Duration) { + for { + s.state.mu.Lock() + paused := s.state.pauseHeartbeats + s.state.mu.Unlock() + if !paused { + if _, err := s.client.Heartbeat(context.Background(), s.heartbeatRequest()); err != nil { + log.Printf("heartbeat failed: %v", err) + } + } + time.Sleep(interval) + } +} + +func (s *stub) heartbeatRequest() *schedulerv1.HeartbeatRequest { + s.state.mu.Lock() + defer s.state.mu.Unlock() + return &schedulerv1.HeartbeatRequest{ + NodeId: s.id, + ClusterId: s.clusterID, + ServiceInstanceId: s.id + "-instance", + Version: "e2e", + Commit: "e2e", + MachineInfo: &schedulerv1.MachineInfo{ + CpuFamily: "e2e", + CpuModel: "e2e", + CpuModelName: "stub-cpu", + CpuArchitecture: "x86_64", + CpuConfigJson: `{"e2e":true}`, + }, + Snapshot: &schedulerv1.NodeSnapshot{ + Status: schedulerv1.NodeStatus_NODE_STATUS_READY, + AllocatedCpu: 1, + CpuPercent: s.state.cpuPercent, + CpuCount: 8, + MemoryUsedBytes: 1 << 30, + MemoryTotalBytes: 8 << 30, + SandboxCount: s.state.sandboxCount, + }, + SandboxIds: append([]string(nil), s.state.sandboxIDs...), + P2PEndpoint: &schedulerv1.P2PEndpoint{Backend: "iroh", Address: s.id + ".p2p.e2e"}, + } +} + +func (s *stub) handleNodes(w http.ResponseWriter, _ *http.Request) { + s.state.mu.Lock() + defer s.state.mu.Unlock() + writeJSON(w, []map[string]any{{ + "version": "e2e", + "commit": "e2e", + "id": s.id, + "serviceInstanceID": s.id + "-instance", + "clusterID": s.clusterID, + "status": "ready", + "sandboxCount": s.state.sandboxCount, + "sandboxIDs": s.state.sandboxIDs, + "createSuccesses": 0, + "createFails": 0, + "sandboxStartingCount": 0, + "sandboxPausedCount": 0, + "p2pEndpoint": map[string]string{"backend": "iroh", "address": s.id + ".p2p.e2e"}, + "machineInfo": map[string]string{ + "cpuFamily": "e2e", + "cpuModel": "e2e", + "cpuModelName": "stub-cpu", + "cpuArchitecture": "x86_64", + "cpuConfigJSON": `{"e2e":true}`, + }, + "metrics": map[string]any{ + "allocatedCPU": 1, + "allocatedMemoryBytes": 1 << 30, + "cpuPercent": s.state.cpuPercent, + "cpuCount": 8, + "memoryUsedBytes": 1 << 30, + "memoryTotalBytes": 8 << 30, + "pausedAllocatedCPU": 0, + "pausedAllocatedMemoryBytes": 0, + "disks": []any{}, + }, + }}) +} + +func (s *stub) handleSandboxes(w http.ResponseWriter, _ *http.Request) { + s.state.mu.Lock() + defer s.state.mu.Unlock() + out := make([]map[string]string, 0, len(s.state.sandboxIDs)) + for _, id := range s.state.sandboxIDs { + out = append(out, map[string]string{"sandboxID": id}) + } + writeJSON(w, out) +} + +func (s *stub) handleSetMetrics(w http.ResponseWriter, r *http.Request) { + var req struct { + SandboxCount uint32 `json:"sandboxCount"` + CpuPercent uint32 `json:"cpuPercent"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + s.state.mu.Lock() + s.state.sandboxCount = req.SandboxCount + s.state.cpuPercent = req.CpuPercent + s.state.mu.Unlock() + w.WriteHeader(http.StatusNoContent) +} + +func (s *stub) handleSetSandboxes(w http.ResponseWriter, r *http.Request) { + var req struct { + SandboxIDs []string `json:"sandboxIDs"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + s.state.mu.Lock() + s.state.sandboxIDs = req.SandboxIDs + s.state.mu.Unlock() + w.WriteHeader(http.StatusNoContent) +} + +func (s *stub) handlePause(w http.ResponseWriter, _ *http.Request) { + s.state.mu.Lock() + s.state.pauseHeartbeats = true + s.state.mu.Unlock() + w.WriteHeader(http.StatusNoContent) +} + +func (s *stub) handleResume(w http.ResponseWriter, _ *http.Request) { + s.state.mu.Lock() + s.state.pauseHeartbeats = false + s.state.mu.Unlock() + w.WriteHeader(http.StatusNoContent) +} + +func writeJSON(w http.ResponseWriter, v any) { + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(v); err != nil { + fmt.Fprintf(w, `{"error":%q}`, err.Error()) + } +} From 3fe2c056bc846b936cdf9b48f50c489b4dd99eb0 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Sat, 10 Oct 2026 10:15:18 +0800 Subject: [PATCH 18/26] fix(scheduler): close discovery ready channel right after first sync MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ready signal added for the leader-election discovery gate used a deferred close, which only fires when Run returns — never during normal operation. Every scheduler pod then hit the 30s waitForDiscoverySync timeout and crashlooped on 'kubernetes discovery initial sync timed out before leader election' (first e2e CI run, #341). Close the channel immediately after WaitForCacheSync succeeds instead. Co-Authored-By: Claude Code --- services/scheduler/internal/kubernetes_discovery.go | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/services/scheduler/internal/kubernetes_discovery.go b/services/scheduler/internal/kubernetes_discovery.go index f184ce05a..d16762942 100644 --- a/services/scheduler/internal/kubernetes_discovery.go +++ b/services/scheduler/internal/kubernetes_discovery.go @@ -157,9 +157,6 @@ func (d *KubernetesDiscovery) Run(ctx context.Context) error { syncCtx, cancelSync := context.WithTimeout(ctx, kubernetesDiscoveryCacheSyncTimeout) defer cancelSync() - if len(d.ready) > 0 { - defer close(d.ready[0]) - } if !cache.WaitForCacheSync(syncCtx.Done(), cacheSyncs...) { if err := ctx.Err(); err != nil { return err @@ -167,7 +164,11 @@ func (d *KubernetesDiscovery) Run(ctx context.Context) error { if err := syncCtx.Err(); err != nil { return fmt.Errorf("kubernetes discovery cache sync timed out after %s", kubernetesDiscoveryCacheSyncTimeout) } - return fmt.Errorf("kubernetes discovery cache sync failed") + } + // Signal first sync right after it completes — deferring this to Run's + // return would never fire during normal operation (#341 e2e finding). + if len(d.ready) > 0 { + close(d.ready[0]) } d.syncFromStore() From 6eb48064319021a11c6077f6d3c663d6d3a90c66 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Sat, 10 Oct 2026 12:10:15 +0800 Subject: [PATCH 19/26] ci(scheduler-e2e): don't rollout-status the scheduler deployment By HA design only the elected leader is Ready (standbys stay NOT_SERVING on the leader health service), so rollout status can never complete for all 3 replicas and the step timed out at 180s. Redis and stubnode keep their rollout waits; the scheduler gate is the existing Wait-for-a-leader step. Co-Authored-By: Claude Code --- .github/workflows/scheduler-ha-e2e.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/scheduler-ha-e2e.yml b/.github/workflows/scheduler-ha-e2e.yml index 66c962d60..0dbc34a8f 100644 --- a/.github/workflows/scheduler-ha-e2e.yml +++ b/.github/workflows/scheduler-ha-e2e.yml @@ -50,7 +50,9 @@ jobs: kubectl apply -f services/scheduler/e2e/k8s/stubnodes.yaml kubectl rollout status deployment/redis --timeout=120s kubectl rollout status deployment/stubnode --timeout=120s - kubectl rollout status deployment/agentenv-scheduler --timeout=180s + # NOTE: do not rollout-status the scheduler Deployment — by HA + # design only the elected leader is Ready, so a full rollout can + # never complete. The "Wait for a leader" step below is the gate. - name: Wait for a leader run: | From ea788a1f8b6d118dda08bf0e3f9922d5a48825d6 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Sat, 10 Oct 2026 14:37:09 +0800 Subject: [PATCH 20/26] test(scheduler-e2e): forward to the leader pod and freeze via debug container MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two root causes from the first real test run: - kubectl port-forward to a Service ignores readiness and can land on a gated standby, so RPCs fail with "not the leader" and every test cascaded from that. Each test now resolves the current Lease holder and port-forwards to the leader pod directly, re-resolving after failovers. - The NetworkPolicy-based API partition did not block renewal in the Kind CNI, so no takeover happened. T2 now freezes the leader process with SIGSTOP through an ephemeral busybox debug container (scheduler pods run shareProcessNamespace in the e2e manifest) and resumes with SIGCONT — a CNI-independent partition. The workflow no longer port-forwards the Service; tests manage their own forwards. Co-Authored-By: Claude Code --- .github/workflows/scheduler-ha-e2e.yml | 4 -- services/scheduler/e2e/harness_test.go | 70 ++++++++++++++--------- services/scheduler/e2e/k8s/scheduler.yaml | 3 + 3 files changed, 47 insertions(+), 30 deletions(-) diff --git a/.github/workflows/scheduler-ha-e2e.yml b/.github/workflows/scheduler-ha-e2e.yml index 0dbc34a8f..0bc1b077a 100644 --- a/.github/workflows/scheduler-ha-e2e.yml +++ b/.github/workflows/scheduler-ha-e2e.yml @@ -67,10 +67,6 @@ jobs: - name: Run failover tests run: | - kubectl port-forward svc/agentenv-scheduler 19090:9090 & - PF_PID=$! - trap "kill $PF_PID 2>/dev/null || true" EXIT - sleep 3 cd services go test -tags e2e -count=1 -timeout 10m -v ./scheduler/e2e/... diff --git a/services/scheduler/e2e/harness_test.go b/services/scheduler/e2e/harness_test.go index f85de7900..9838ee8d2 100644 --- a/services/scheduler/e2e/harness_test.go +++ b/services/scheduler/e2e/harness_test.go @@ -24,13 +24,6 @@ import ( var schedAddr = envOr("E2E_SCHED_ADDR", "127.0.0.1:19090") -func writeFile(t *testing.T, path, content string) { - t.Helper() - if err := os.WriteFile(path, []byte(content), 0o644); err != nil { - t.Fatalf("write %s: %v", path, err) - } -} - func envOr(key, fallback string) string { if v := os.Getenv(key); v != "" { return v @@ -78,6 +71,33 @@ func leaderIdentity(t *testing.T) string { return strings.TrimSpace(out) } +// forwardLeader port-forwards to the current leader POD. kubectl port-forward +// to a Service ignores readiness and can land on a gated standby (every RPC +// rejected with "not the leader"), so tests must target the leader pod +// directly and re-resolve it after every failover. +func forwardLeader(t *testing.T) { + t.Helper() + leader := leaderIdentity(t) + if leader == "" { + t.Fatal("no leader elected yet") + } + cmd := exec.Command("kubectl", "port-forward", "pod/"+leader, "19090:9090") + if err := cmd.Start(); err != nil { + t.Fatalf("port-forward to leader %s: %v", leader, err) + } + t.Cleanup(func() { _ = cmd.Process.Kill() }) + deadline := time.Now().Add(10 * time.Second) + for time.Now().Before(deadline) { + conn, err := grpc.NewClient(schedAddr, grpc.WithTransportCredentials(insecure.NewCredentials())) + if err == nil { + conn.Close() + return + } + time.Sleep(300 * time.Millisecond) + } + t.Fatalf("port-forward to leader %s never came up", leader) +} + func schedulerEndpoints(t *testing.T) []string { t.Helper() out := kube(t, "get", "endpoints", "agentenv-scheduler", "-o", "jsonpath={.subsets[*].addresses[*].targetRef.name}") @@ -88,6 +108,7 @@ func schedulerEndpoints(t *testing.T) []string { // budget, pre-existing bindings keep resolving through Redis, and // scheduling resumes on freshly observed nodes. func TestFailoverLeaderKill(t *testing.T) { + forwardLeader(t) client := dial(t) ctx := context.Background() @@ -149,34 +170,29 @@ func TestFailoverLeaderKill(t *testing.T) { // serving as a second primary. func TestPartitionFrozenLeader(t *testing.T) { victim := leaderIdentity(t) - t.Logf("partitioning leader pod %s from the API (egress block)", victim) - kube(t, "label", "pod", victim, "e2e-partition=true", "--overwrite") - writeFile(t, "/tmp/e2e-netpol.yaml", `apiVersion: networking.k8s.io/v1 -kind: NetworkPolicy -metadata: - name: e2e-partition -spec: - podSelector: - matchLabels: - e2e-partition: "true" - policyTypes: ["Egress"] - egress: [] -`) - kube(t, "apply", "-f", "/tmp/e2e-netpol.yaml") - t.Cleanup(func() { kube(t, "delete", "networkpolicy", "e2e-partition", "--ignore-not-found") }) - - eventually(t, 45*time.Second, "standby takes over while old leader partitioned", func() (bool, error) { + t.Logf("freezing leader process in pod %s (SIGSTOP = cannot renew)", victim) + out, err := exec.Command("kubectl", "debug", "-q", victim, "--image=busybox:1.36", + "--", "kill", "-STOP", "1").CombinedOutput() + if err != nil { + t.Fatalf("freeze leader via debug container failed: %v\n%s", err, out) + } + + eventually(t, 45*time.Second, "standby takes over while old leader frozen", func() (bool, error) { cur := leaderIdentity(t) return cur != "" && cur != victim, nil }) - // While partitioned, there must be exactly one serving endpoint. + // While frozen, there must be exactly one serving endpoint. eps := schedulerEndpoints(t) if len(eps) != 1 { t.Fatalf("expected exactly one endpoint during partition, got %v", eps) } - kube(t, "delete", "networkpolicy", "e2e-partition") + out, err = exec.Command("kubectl", "debug", "-q", victim, "--image=busybox:1.36", + "--", "kill", "-CONT", "1").CombinedOutput() + if err != nil { + t.Fatalf("resume leader via debug container failed: %v\n%s", err, out) + } // With egress restored, the ex-leader's renew failure has already fired // OnStoppedLeading: it force-stops and its pod restarts as standby. eventually(t, 60*time.Second, "frozen ex-leader exits (pod restarts)", func() (bool, error) { @@ -192,6 +208,7 @@ spec: // return Unavailable (never fall back to unobserved nodes); once heartbeats // resume, scheduling recovers. func TestTakeoverSchedulingSemantics(t *testing.T) { + forwardLeader(t) client := dial(t) ctx := context.Background() @@ -220,6 +237,7 @@ func TestTakeoverSchedulingSemantics(t *testing.T) { // commit after the takeover; the client retries and succeeds on the new // leader. func TestDemotedLeaderDelayedWrite(t *testing.T) { + forwardLeader(t) client := dial(t) ctx := context.Background() diff --git a/services/scheduler/e2e/k8s/scheduler.yaml b/services/scheduler/e2e/k8s/scheduler.yaml index 80fa62e51..54e2cc941 100644 --- a/services/scheduler/e2e/k8s/scheduler.yaml +++ b/services/scheduler/e2e/k8s/scheduler.yaml @@ -77,6 +77,9 @@ spec: app: agentenv-scheduler spec: serviceAccountName: agentenv-scheduler + # Lets debug containers signal the scheduler process (T2 freezes the + # leader with SIGSTOP via an ephemeral busybox container). + shareProcessNamespace: true containers: - name: scheduler image: agentenv-scheduler:e2e From 2b31f00bae95234e97b49896afbf8ffbe52e26c5 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Sat, 10 Oct 2026 15:00:26 +0800 Subject: [PATCH 21/26] test(scheduler-e2e): re-forward and re-dial after every takeover MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous fixes were necessary but not sufficient, as the latest run showed: - grpc.NewClient is lazy, so the forward readiness probe passed before kubectl had the tunnel up; probe with a real TCP dial instead. - Tests kill the very pod they are forwarded to. After each takeover, re-resolve the Lease holder, kill the stale forward, and re-dial — otherwise every post-kill RPC goes to a dead tunnel. - SIGSTOP hit PID 1, which under shareProcessNamespace is the sandbox pause container, not the scheduler; target it with pidof. Co-Authored-By: Claude Code --- services/scheduler/e2e/harness_test.go | 39 ++++++++++++++++++++++---- 1 file changed, 33 insertions(+), 6 deletions(-) diff --git a/services/scheduler/e2e/harness_test.go b/services/scheduler/e2e/harness_test.go index 9838ee8d2..c898b6ed1 100644 --- a/services/scheduler/e2e/harness_test.go +++ b/services/scheduler/e2e/harness_test.go @@ -9,6 +9,7 @@ package e2e import ( "context" "fmt" + "net" "os" "os/exec" "strings" @@ -71,12 +72,19 @@ func leaderIdentity(t *testing.T) string { return strings.TrimSpace(out) } +var currentForward *exec.Cmd + // forwardLeader port-forwards to the current leader POD. kubectl port-forward // to a Service ignores readiness and can land on a gated standby (every RPC // rejected with "not the leader"), so tests must target the leader pod -// directly and re-resolve it after every failover. +// directly and re-resolve it after every failover. Any previous forward is +// killed first, since the pod it targeted may have been deleted mid-test. func forwardLeader(t *testing.T) { t.Helper() + if currentForward != nil { + _ = currentForward.Process.Kill() + currentForward = nil + } leader := leaderIdentity(t) if leader == "" { t.Fatal("no leader elected yet") @@ -85,10 +93,18 @@ func forwardLeader(t *testing.T) { if err := cmd.Start(); err != nil { t.Fatalf("port-forward to leader %s: %v", leader, err) } - t.Cleanup(func() { _ = cmd.Process.Kill() }) - deadline := time.Now().Add(10 * time.Second) + currentForward = cmd + t.Cleanup(func() { + if currentForward == cmd { + _ = cmd.Process.Kill() + currentForward = nil + } + }) + // Probe the port with a real TCP dial: grpc.NewClient is lazy and would + // report ready long before kubectl has the tunnel up. + deadline := time.Now().Add(15 * time.Second) for time.Now().Before(deadline) { - conn, err := grpc.NewClient(schedAddr, grpc.WithTransportCredentials(insecure.NewCredentials())) + conn, err := net.DialTimeout("tcp", "127.0.0.1:19090", 500*time.Millisecond) if err == nil { conn.Close() return @@ -159,6 +175,9 @@ func TestFailoverLeaderKill(t *testing.T) { case <-time.After(2 * time.Second): } + // The forward targeted the killed pod; re-forward and re-dial. + forwardLeader(t) + client = dial(t) eventually(t, 30*time.Second, "scheduling works again on the new leader", func() (bool, error) { _, err := client.Schedule(ctx, &schedulerv1.ScheduleRequest{}) return err == nil, nil @@ -172,7 +191,7 @@ func TestPartitionFrozenLeader(t *testing.T) { victim := leaderIdentity(t) t.Logf("freezing leader process in pod %s (SIGSTOP = cannot renew)", victim) out, err := exec.Command("kubectl", "debug", "-q", victim, "--image=busybox:1.36", - "--", "kill", "-STOP", "1").CombinedOutput() + "--", "sh", "-c", "kill -STOP $(pidof scheduler)").CombinedOutput() if err != nil { t.Fatalf("freeze leader via debug container failed: %v\n%s", err, out) } @@ -189,7 +208,7 @@ func TestPartitionFrozenLeader(t *testing.T) { } out, err = exec.Command("kubectl", "debug", "-q", victim, "--image=busybox:1.36", - "--", "kill", "-CONT", "1").CombinedOutput() + "--", "sh", "-c", "kill -CONT $(pidof scheduler)").CombinedOutput() if err != nil { t.Fatalf("resume leader via debug container failed: %v\n%s", err, out) } @@ -221,6 +240,10 @@ func TestTakeoverSchedulingSemantics(t *testing.T) { return cur != "" && cur != victim, nil }) + // The forward targeted the killed pod; re-forward and re-dial. + forwardLeader(t) + client = dial(t) + // No node has freshly reported to the new leader: Unavailable, not a guess. if _, err := client.Schedule(ctx, &schedulerv1.ScheduleRequest{}); err == nil { t.Fatal("schedule must fail with Unavailable before fresh observations, not guess capacity") @@ -260,6 +283,10 @@ func TestDemotedLeaderDelayedWrite(t *testing.T) { return cur != "" && cur != victim, nil }) + // The forward targeted the killed pod; re-forward and re-dial. + forwardLeader(t) + client = dial(t) + // The client-side retry path: writing again must succeed on the new leader. if _, err := client.RecordAssignment(ctx, &schedulerv1.RecordAssignmentRequest{ SandboxId: "sbx-e2e-t4", From 858235ec932517f136dfd59f69ae7be108cffa1a Mon Sep 17 00:00:00 2001 From: NickNYU Date: Sat, 10 Oct 2026 15:33:25 +0800 Subject: [PATCH 22/26] test(scheduler-e2e): query EndpointSlice and make T3's zero-observation real MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two findings from the latest run (where T4 now passes): - kubectl get endpoints prints a deprecation warning on v1.33 that polluted the name list, so the converge checks always failed. Query EndpointSlice directly and filter by pod-name prefix. - T3's premise predated sync-node-snapshots: pausing heartbeats does not stop the leader's admin pull, so scheduling recovered instantly and the Unavailable assertion was wrong. Zero fresh observations now requires both channels down — the stub gains an admin pause toggle (/nodes and /sandboxes return 503), and T3 pauses heartbeats and the admin API before the takeover, asserting Unavailable, then resumes both and asserts the pull-driven rebuild works. Co-Authored-By: Claude Code --- services/scheduler/e2e/harness_test.go | 26 ++++++++++++++++++------- services/scheduler/e2e/stubnode/main.go | 24 +++++++++++++++++++++++ 2 files changed, 43 insertions(+), 7 deletions(-) diff --git a/services/scheduler/e2e/harness_test.go b/services/scheduler/e2e/harness_test.go index c898b6ed1..3e05c6a52 100644 --- a/services/scheduler/e2e/harness_test.go +++ b/services/scheduler/e2e/harness_test.go @@ -116,8 +116,15 @@ func forwardLeader(t *testing.T) { func schedulerEndpoints(t *testing.T) []string { t.Helper() - out := kube(t, "get", "endpoints", "agentenv-scheduler", "-o", "jsonpath={.subsets[*].addresses[*].targetRef.name}") - return strings.Fields(out) + out := kube(t, "get", "endpointslice", "-l", "kubernetes.io/service-name=agentenv-scheduler", + "-o", "jsonpath={.items[*].endpoints[*].targetRef.name}") + var names []string + for _, f := range strings.Fields(out) { + if strings.HasPrefix(f, "agentenv-scheduler-") { + names = append(names, f) + } + } + return names } // T1: kill the leader pod — a standby must take over within the lease @@ -223,15 +230,19 @@ func TestPartitionFrozenLeader(t *testing.T) { }) } -// T3: right after takeover, with node heartbeats paused, scheduling must -// return Unavailable (never fall back to unobserved nodes); once heartbeats -// resume, scheduling recovers. +// T3: with node heartbeats paused AND the admin pull failing, scheduling +// right after takeover must return Unavailable (never fall back to +// unobserved nodes); once both resume, the pull-driven rebuild makes +// scheduling work again. func TestTakeoverSchedulingSemantics(t *testing.T) { forwardLeader(t) client := dial(t) ctx := context.Background() + // Zero fresh observations requires both channels down: heartbeats alone + // are not enough, since sync-node-snapshots rebuilds via the admin pull. stubPost(t, "control/pause", "") + stubPost(t, "control/admin-pause", "") victim := leaderIdentity(t) kube(t, "delete", "pod", victim, "--force", "--grace-period=0") @@ -244,13 +255,14 @@ func TestTakeoverSchedulingSemantics(t *testing.T) { forwardLeader(t) client = dial(t) - // No node has freshly reported to the new leader: Unavailable, not a guess. + // No fresh observations from either channel: Unavailable, not a guess. if _, err := client.Schedule(ctx, &schedulerv1.ScheduleRequest{}); err == nil { t.Fatal("schedule must fail with Unavailable before fresh observations, not guess capacity") } stubPost(t, "control/resume", "") - eventually(t, 30*time.Second, "scheduling recovers after heartbeats resume", func() (bool, error) { + stubPost(t, "control/admin-resume", "") + eventually(t, 30*time.Second, "scheduling recovers once pull and heartbeats resume", func() (bool, error) { _, err := client.Schedule(ctx, &schedulerv1.ScheduleRequest{}) return err == nil, nil }) diff --git a/services/scheduler/e2e/stubnode/main.go b/services/scheduler/e2e/stubnode/main.go index 0c6854417..d0da0cb03 100644 --- a/services/scheduler/e2e/stubnode/main.go +++ b/services/scheduler/e2e/stubnode/main.go @@ -28,6 +28,7 @@ type stubState struct { sandboxCount uint32 cpuPercent uint32 pauseHeartbeats bool + pauseAdmin bool } type stub struct { @@ -70,6 +71,8 @@ func main() { mux.HandleFunc("POST /control/sandboxes", s.handleSetSandboxes) mux.HandleFunc("POST /control/pause", s.handlePause) mux.HandleFunc("POST /control/resume", s.handleResume) + mux.HandleFunc("POST /control/admin-pause", s.handleAdminPause) + mux.HandleFunc("POST /control/admin-resume", s.handleAdminResume) mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) }) log.Printf("stubnode %s listening on %s, heartbeating to %s", *id, *listen, *schedulerAddr) @@ -89,6 +92,13 @@ func (s *stub) withAuth(next http.HandlerFunc) http.HandlerFunc { http.Error(w, "unauthorized", http.StatusUnauthorized) return } + s.state.mu.Lock() + paused := s.state.pauseAdmin + s.state.mu.Unlock() + if paused { + http.Error(w, "admin paused", http.StatusServiceUnavailable) + return + } next(w, r) } } @@ -229,6 +239,20 @@ func (s *stub) handleResume(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusNoContent) } +func (s *stub) handleAdminPause(w http.ResponseWriter, _ *http.Request) { + s.state.mu.Lock() + s.state.pauseAdmin = true + s.state.mu.Unlock() + w.WriteHeader(http.StatusNoContent) +} + +func (s *stub) handleAdminResume(w http.ResponseWriter, _ *http.Request) { + s.state.mu.Lock() + s.state.pauseAdmin = false + s.state.mu.Unlock() + w.WriteHeader(http.StatusNoContent) +} + func writeJSON(w http.ResponseWriter, v any) { w.Header().Set("Content-Type", "application/json") if err := json.NewEncoder(w).Encode(v); err != nil { From 11e8307fc3fee2f6f07c1d82b838fbbeb5d4676e Mon Sep 17 00:00:00 2001 From: NickNYU Date: Sat, 10 Oct 2026 16:13:06 +0800 Subject: [PATCH 23/26] test(scheduler-e2e): use production lease timings and wider budgets The 6s/4s/1s lease timings were too aggressive for a loaded CI Kind: renewals intermittently exceeded the 4s renew deadline, so candidates kept grabbing the lease from each other (two acquisitions in the same second), and every endpoint/forward assertion drowned in the churn. Use the production 15s/10s/2s timings and widen takeover/converge budgets to 60-90s to match. Co-Authored-By: Claude Code --- services/scheduler/e2e/harness_test.go | 18 +++++++++--------- services/scheduler/e2e/k8s/scheduler.yaml | 6 +++--- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/services/scheduler/e2e/harness_test.go b/services/scheduler/e2e/harness_test.go index 3e05c6a52..fc59d3875 100644 --- a/services/scheduler/e2e/harness_test.go +++ b/services/scheduler/e2e/harness_test.go @@ -168,11 +168,11 @@ func TestFailoverLeaderKill(t *testing.T) { t.Logf("killing leader pod %s", victim) kube(t, "delete", "pod", victim, "--force", "--grace-period=0") - eventually(t, 30*time.Second, "a new leader", func() (bool, error) { + eventually(t, 60*time.Second, "a new leader", func() (bool, error) { cur := leaderIdentity(t) return cur != "" && cur != victim, nil }) - eventually(t, 30*time.Second, "endpoints converge on one leader", func() (bool, error) { + eventually(t, 60*time.Second, "endpoints converge on one leader", func() (bool, error) { return len(schedulerEndpoints(t)) == 1, nil }) @@ -185,7 +185,7 @@ func TestFailoverLeaderKill(t *testing.T) { // The forward targeted the killed pod; re-forward and re-dial. forwardLeader(t) client = dial(t) - eventually(t, 30*time.Second, "scheduling works again on the new leader", func() (bool, error) { + eventually(t, 60*time.Second, "scheduling works again on the new leader", func() (bool, error) { _, err := client.Schedule(ctx, &schedulerv1.ScheduleRequest{}) return err == nil, nil }) @@ -203,7 +203,7 @@ func TestPartitionFrozenLeader(t *testing.T) { t.Fatalf("freeze leader via debug container failed: %v\n%s", err, out) } - eventually(t, 45*time.Second, "standby takes over while old leader frozen", func() (bool, error) { + eventually(t, 60*time.Second, "standby takes over while old leader frozen", func() (bool, error) { cur := leaderIdentity(t) return cur != "" && cur != victim, nil }) @@ -221,11 +221,11 @@ func TestPartitionFrozenLeader(t *testing.T) { } // With egress restored, the ex-leader's renew failure has already fired // OnStoppedLeading: it force-stops and its pod restarts as standby. - eventually(t, 60*time.Second, "frozen ex-leader exits (pod restarts)", func() (bool, error) { + eventually(t, 90*time.Second, "frozen ex-leader exits (pod restarts)", func() (bool, error) { out := kube(t, "get", "pod", victim, "-o", "jsonpath={.status.containerStatuses[0].restartCount}") return out != "0", nil }) - eventually(t, 30*time.Second, "still exactly one leader after resume", func() (bool, error) { + eventually(t, 60*time.Second, "still exactly one leader after resume", func() (bool, error) { return len(schedulerEndpoints(t)) == 1, nil }) } @@ -246,7 +246,7 @@ func TestTakeoverSchedulingSemantics(t *testing.T) { victim := leaderIdentity(t) kube(t, "delete", "pod", victim, "--force", "--grace-period=0") - eventually(t, 30*time.Second, "new leader", func() (bool, error) { + eventually(t, 60*time.Second, "new leader", func() (bool, error) { cur := leaderIdentity(t) return cur != "" && cur != victim, nil }) @@ -262,7 +262,7 @@ func TestTakeoverSchedulingSemantics(t *testing.T) { stubPost(t, "control/resume", "") stubPost(t, "control/admin-resume", "") - eventually(t, 30*time.Second, "scheduling recovers once pull and heartbeats resume", func() (bool, error) { + eventually(t, 60*time.Second, "scheduling recovers once pull and heartbeats resume", func() (bool, error) { _, err := client.Schedule(ctx, &schedulerv1.ScheduleRequest{}) return err == nil, nil }) @@ -290,7 +290,7 @@ func TestDemotedLeaderDelayedWrite(t *testing.T) { writeErr := <-done t.Logf("write during takeover returned: %v", writeErr) - eventually(t, 30*time.Second, "new leader", func() (bool, error) { + eventually(t, 60*time.Second, "new leader", func() (bool, error) { cur := leaderIdentity(t) return cur != "" && cur != victim, nil }) diff --git a/services/scheduler/e2e/k8s/scheduler.yaml b/services/scheduler/e2e/k8s/scheduler.yaml index 54e2cc941..b81b83eed 100644 --- a/services/scheduler/e2e/k8s/scheduler.yaml +++ b/services/scheduler/e2e/k8s/scheduler.yaml @@ -55,9 +55,9 @@ data: "enabled": true, "lease_name": "agentenv-scheduler", "lease_namespace": "default", - "lease_duration": "6s", - "renew_deadline": "4s", - "retry_period": "1s" + "lease_duration": "15s", + "renew_deadline": "10s", + "retry_period": "2s" } } } From 3845688133e1339687db23597a69120fb4c04e79 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Sat, 10 Oct 2026 17:09:04 +0800 Subject: [PATCH 24/26] test(scheduler-e2e): filter EndpointSlice by ready condition EndpointSlice's endpoints array includes not-ready addresses, so the converge checks always counted all three pods even though the leader-only readiness gate worked correctly. Filter with ?(@.conditions.ready==true). Co-Authored-By: Claude Code --- services/scheduler/e2e/harness_test.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/services/scheduler/e2e/harness_test.go b/services/scheduler/e2e/harness_test.go index fc59d3875..1a2a4e135 100644 --- a/services/scheduler/e2e/harness_test.go +++ b/services/scheduler/e2e/harness_test.go @@ -116,8 +116,10 @@ func forwardLeader(t *testing.T) { func schedulerEndpoints(t *testing.T) []string { t.Helper() + // EndpointSlice lists not-ready addresses too; filter by the ready + // condition, or standbys always show up in the count. out := kube(t, "get", "endpointslice", "-l", "kubernetes.io/service-name=agentenv-scheduler", - "-o", "jsonpath={.items[*].endpoints[*].targetRef.name}") + "-o", "jsonpath={.items[*].endpoints[?(@.conditions.ready==true)].targetRef.name}") var names []string for _, f := range strings.Fields(out) { if strings.HasPrefix(f, "agentenv-scheduler-") { From a4c52281067e6957cb4ecd51fe38bfa4ecb7668a Mon Sep 17 00:00:00 2001 From: NickNYU Date: Sat, 10 Oct 2026 19:06:30 +0800 Subject: [PATCH 25/26] test(scheduler-e2e): wait for readiness flip in T2's endpoint assertion The Lease flips before readiness does, so zero ready endpoints is the correct transient right after takeover. T2 asserted exactly-one instantly and raced the probe. Wrap the assertion in eventually (and fail fast if the count ever exceeds one). Co-Authored-By: Claude Code --- services/scheduler/e2e/harness_test.go | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/services/scheduler/e2e/harness_test.go b/services/scheduler/e2e/harness_test.go index 1a2a4e135..de432ee7c 100644 --- a/services/scheduler/e2e/harness_test.go +++ b/services/scheduler/e2e/harness_test.go @@ -210,11 +210,15 @@ func TestPartitionFrozenLeader(t *testing.T) { return cur != "" && cur != victim, nil }) - // While frozen, there must be exactly one serving endpoint. - eps := schedulerEndpoints(t) - if len(eps) != 1 { - t.Fatalf("expected exactly one endpoint during partition, got %v", eps) - } + // The Lease flips before readiness does, so zero endpoints is the + // correct transient — wait for the new leader's probe to flip SERVING. + eventually(t, 60*time.Second, "exactly one ready endpoint while old leader frozen", func() (bool, error) { + eps := schedulerEndpoints(t) + if len(eps) > 1 { + return false, fmt.Errorf("more than one ready endpoint: %v", eps) + } + return len(eps) == 1, nil + }) out, err = exec.Command("kubectl", "debug", "-q", victim, "--image=busybox:1.36", "--", "sh", "-c", "kill -CONT $(pidof scheduler)").CombinedOutput() From 9c6b4228a84b4024b8b3802f33232c829313eba0 Mon Sep 17 00:00:00 2001 From: NickNYU Date: Sat, 10 Oct 2026 19:59:00 +0800 Subject: [PATCH 26/26] docs(scheduler): document the failover read gap (R4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With leader-only endpoints, LookupNode is served only by the leader, so failover has a read gap from the leader's death until the new leader's readiness flips (~15-25s with default timings). Document it in both READMEs, including the two binding-store outcomes: with Redis the floored TTL (lease_duration + 90s) carries lookups across the gap with no false NotFound; without it, pre-failover bindings are lost (degraded mode). Also correct the earlier bullets that implied standbys actively serve reads — the gate allows it, but endpoints do not route there; closing that fork is under review in #341. Addresses the maintainer's R4 requirement on PR #341. Relates to #259. Co-Authored-By: Claude Code --- deploy/k8s/overlays/ha/README.md | 11 +++++++++-- services/README.md | 12 +++++++++--- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/deploy/k8s/overlays/ha/README.md b/deploy/k8s/overlays/ha/README.md index c8d81b899..267cb1c61 100644 --- a/deploy/k8s/overlays/ha/README.md +++ b/deploy/k8s/overlays/ha/README.md @@ -9,8 +9,15 @@ Traffic rules: 1. Node heartbeats → `agentenv-scheduler` Service (leader only). 2. Gateway writes (`Schedule`, `RecordAssignment`) → same Service. -3. Gateway reads (`LookupNode`) → same Service; served by standbys when - `scheduler.redis_addr` is set, otherwise retried onto the leader. +3. Gateway reads (`LookupNode`) → same Service, answered by the leader. + +Read gap on failover: because endpoints contain only the leader, lookups +fail fast from the leader's death until the new leader is ready +(~lease_duration + election + probe, roughly 15–25s with defaults). With +`scheduler.redis_addr` bindings survive (TTL floored to +`lease_duration + 90s`) and lookups resume automatically; without it, +pre-failover bindings are lost (degraded mode). See "Failover read gap" +in services/README.md. Deploy: diff --git a/services/README.md b/services/README.md index c5187df73..0aa5109cd 100644 --- a/services/README.md +++ b/services/README.md @@ -112,14 +112,20 @@ General config notes: `scheduler.leader_election` enables Kubernetes Lease-based leader election (#259). Off by default; with it off, the scheduler behaves exactly as a single-writer process. -- With it on, N replicas compete for a `coordination.k8s.io/Lease`. Exactly one leader schedules and processes heartbeats; standbys stay liveness-healthy, reject writes with `Unavailable`, and serve `LookupNode`/`GetNode` from the shared Redis bindings (reads are rejected too when no `redis_addr` is set, and clients retry onto the leader). +- With it on, N replicas compete for a `coordination.k8s.io/Lease`. Exactly one leader schedules and processes heartbeats; standbys stay liveness-healthy and reject writes with `Unavailable` (standbys may serve reads from the shared Redis bindings at the gate level, but with leader-only endpoints they receive no traffic — see the read gap below). - Readiness probes should target the leader-specific health service `scheduler.v1.Scheduler/leader` (`grpc_health_probe -service=scheduler.v1.Scheduler/leader`) so Service endpoints contain only the leader. Liveness keeps probing the overall health status. -- Traffic rules: node heartbeats → scheduler Service (leader only); gateway writes (`Schedule`, `RecordAssignment`) → same Service; gateway reads (`LookupNode`) → same Service, answered by standbys when Redis is configured. +- Traffic rules: node heartbeats → scheduler Service (leader only); gateway writes (`Schedule`, `RecordAssignment`) → same Service; gateway reads (`LookupNode`) → same Service, answered by the leader. - On failover, the new leader pulls each node's admin `/nodes` snapshot (sync-node-snapshots) instead of waiting for the next heartbeat; pulls authenticate with the `x-api-key` from `SCHEDULER_NODE_ADMIN_API_KEY` — pass it via env/Secret (the HA overlay mounts the shared `agentenv-auth` Secret), not via config files. - `redis_addr` is optional but recommended: without it, failover loses routing for pre-failover sandboxes (documented degraded mode). Under election the binding TTL is floored to `lease_duration + 90s` so bindings outlive the failover budget. - Mutually exclusive with `--query-only`. Ready-made manifests: `deploy/k8s/overlays/ha` (see its README for upgrade ordering). -### Scheduling strategy +#### Failover read gap (lookup continuity) + +Because Service endpoints contain only the leader, `LookupNode` is served only by the leader. During failover there is a read gap: from the leader's death until the new leader's readiness flips, the Service has zero ready endpoints and lookups fail fast (connection refused). With default timings (`lease_duration=15s`) the window is roughly 15–25s (detection + election + readiness probe). + +- With `redis_addr`: bindings survive — under election the binding TTL is floored to `lease_duration + 90s`, deliberately covering the worst-case failover budget (detection + endpoint propagation + one reporter reconnect backoff). Lookups resume automatically and correctly once the new leader is ready; live sandboxes are never misreported as NotFound. +- Without `redis_addr` (degraded mode): bindings die with the old leader's process, so lookups for pre-failover sandboxes keep failing after the gap until those sandboxes are recreated. New scheduling is unaffected. +- Closing the gap entirely (standbys serving reads) is an open design fork under review in #341; this section describes the shipped behavior. `scheduler.strategy` selects the algorithm used to pick a node from the eligible candidate list. Built-in strategies: