diff --git a/packages/api/go.mod b/packages/api/go.mod index 2e29f2a951..e08c6fc4da 100644 --- a/packages/api/go.mod +++ b/packages/api/go.mod @@ -43,6 +43,7 @@ require ( github.com/oapi-codegen/runtime v1.4.0 github.com/posthog/posthog-go v0.0.0-20230801140217-d607812dee69 github.com/redis/go-redis/v9 v9.17.3 + github.com/samber/lo v1.53.0 github.com/stretchr/testify v1.11.1 go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.67.0 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.67.0 diff --git a/packages/api/go.sum b/packages/api/go.sum index db955eb386..26475cb39b 100644 --- a/packages/api/go.sum +++ b/packages/api/go.sum @@ -934,6 +934,8 @@ github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0= github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY= github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ= github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= +github.com/samber/lo v1.53.0 h1:t975lj2py4kJPQ6haz1QMgtId2gtmfktACxIXArw3HM= +github.com/samber/lo v1.53.0/go.mod h1:4+MXEGsJzbKGaUEQFKBq2xtfuznW9oz/WrgyzMzRoM0= github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 h1:KRzFb2m7YtdldCEkzs6KqmJw4nqEVZGK7IN2kJkjTuQ= github.com/santhosh-tekuri/jsonschema/v6 v6.0.2/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU= github.com/scaleway/scaleway-sdk-go v1.0.0-beta.36 h1:ObX9hZmK+VmijreZO/8x9pQ8/P/ToHD/bdSb4Eg4tUo= diff --git a/packages/api/internal/cfg/model.go b/packages/api/internal/cfg/model.go index f841ebddc0..ddd3a894b3 100644 --- a/packages/api/internal/cfg/model.go +++ b/packages/api/internal/cfg/model.go @@ -48,6 +48,24 @@ type Config struct { NomadAddress string `env:"NOMAD_ADDRESS" envDefault:"http://localhost:4646"` NomadToken string `env:"NOMAD_TOKEN"` + // NomadOrchestratorServiceNames is the comma-separated list of + // Nomad-native service names whose registrations enumerate orchestrator + // instances (GET /v1/service/ per name, results unioned). Every + // orchestrator jobspec registers one of these services, regardless of + // job type or node pool. Used when ServiceDiscoveryProvider=nomad. + NomadOrchestratorServiceNames []string `env:"NOMAD_ORCHESTRATOR_SERVICE_NAMES" envDefault:"orchestrator" envSeparator:","` + + // NomadOrchestratorLegacyDiscoveryEnabled enables a node-pool-based + // discovery FALLBACK unioned with the service-based discovery above: all + // ready Nomad nodes in the "default" pool are assumed to run an + // orchestrator on the well-known port. It covers orchestrator jobs + // deployed from jobspecs that register their service with an empty + // Address (pre-port-label-fix), which service discovery skips as + // unroutable, and removes any rollout ordering constraint between the + // API and orchestrator releases. Set to false once no legacy orchestrator + // jobs remain. Used when ServiceDiscoveryProvider=nomad. + NomadOrchestratorLegacyDiscoveryEnabled bool `env:"NOMAD_ORCHESTRATOR_LEGACY_DISCOVERY_ENABLED" envDefault:"true"` + // LocalOrchestratorAddress is the "host:port" address of a statically // configured orchestrator instance. Required when // ServiceDiscoveryProvider=local. Used for local dev against the darwin diff --git a/packages/api/internal/handlers/store.go b/packages/api/internal/handlers/store.go index 06a6293f5b..f154a1e923 100644 --- a/packages/api/internal/handlers/store.go +++ b/packages/api/internal/handlers/store.go @@ -171,7 +171,21 @@ func NewAPIStore(ctx context.Context, tel *telemetry.Client, redisClient redis.U if nomadErr != nil { logger.L().Fatal(ctx, "Initializing Nomad client", zap.Error(nomadErr)) } - nodeDiscovery = orchdiscovery.NewNomad(nomadClient, "default") + nodeDiscovery = orchdiscovery.NewNomad(nomadClient, config.NomadOrchestratorServiceNames) + // Migration fallback: orchestrator jobs deployed from jobspecs that + // predate the service port-label fix register their service with an + // empty Address, so service discovery alone would miss them until + // they are redeployed. Union in the legacy node-pool listing (service + // entries win on conflict) so the API flip has no rollout ordering + // constraint. Disable via NOMAD_ORCHESTRATOR_LEGACY_DISCOVERY_ENABLED + // once no legacy jobs remain. The pool is hardcoded: legacy jobs only + // ever ran on the "default" pool. + if config.NomadOrchestratorLegacyDiscoveryEnabled { + nodeDiscovery = orchdiscovery.NewMerged( + nodeDiscovery, + orchdiscovery.NewNomadNodePool(nomadClient, "default"), + ) + } templateBuilderDiscovery = clustersdiscovery.NewLocalDiscovery(consts.LocalClusterID, nomadClient) } diff --git a/packages/api/internal/orchestrator/client_test.go b/packages/api/internal/orchestrator/client_test.go index 40fb97a63d..9eee22ce24 100644 --- a/packages/api/internal/orchestrator/client_test.go +++ b/packages/api/internal/orchestrator/client_test.go @@ -42,7 +42,7 @@ func newTestOrchestrator(t *testing.T, nomad *nomadapi.Client) *Orchestrator { return &Orchestrator{ nodes: smap.New[*nodemanager.Node](), - nodeDiscovery: discovery.NewNomad(nomad, "default"), + nodeDiscovery: discovery.NewNomad(nomad, []string{"orchestrator"}), tel: telemetry.NewNoopClient(), } } @@ -139,18 +139,19 @@ func TestGetOrConnectNode_CacheMiss_TriggersNomadDiscovery(t *testing.T) { var discoveryAttempts atomic.Int32 nomadClient := newNomadMock(t, func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/v1/nodes" { + if r.URL.Path == "/v1/service/orchestrator" { discoveryAttempts.Add(1) - // Return a node stub. connectToNode will fail at the gRPC level - // (nodemanager.New dials the fake address), but the important thing - // is that discovery WAS attempted. - resp := []*nomadapi.NodeListStub{ + // Return a service registration. connectToNode will fail at the + // gRPC level (nodemanager.New dials the fake address), but the + // important thing is that discovery WAS attempted. + resp := []*nomadapi.ServiceRegistration{ { - ID: "abcdef1234567890abcdef1234567890abcdef12", - Address: "127.0.0.1", - Status: "ready", - NodePool: "default", + ID: "_nomad-task-abc-orchestrator", + ServiceName: "orchestrator", + NodeID: "abcdef1234567890abcdef1234567890abcdef12", + Address: "127.0.0.1", + Port: int(consts.OrchestratorAPIPort), }, } w.Header().Set("Content-Type", "application/json") @@ -184,13 +185,13 @@ func TestGetOrConnectNode_ConcurrentCacheMiss_SharesDiscovery(t *testing.T) { var discoveryAttempts atomic.Int32 nomadClient := newNomadMock(t, func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/v1/nodes" { + if r.URL.Path == "/v1/service/orchestrator" { discoveryAttempts.Add(1) // Slow response to ensure concurrent callers overlap. time.Sleep(100 * time.Millisecond) w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode([]*nomadapi.NodeListStub{}) + json.NewEncoder(w).Encode([]*nomadapi.ServiceRegistration{}) return } @@ -266,8 +267,8 @@ func TestConnectToNode_SingleflightDedup(t *testing.T) { // 3. This API instance has NOT yet synced (node is absent from o.nodes). // 4. A handler calls getOrConnectNode for a sandbox on that node. // -// The fake gRPC server listens on consts.OrchestratorAPIPort so that -// listNomadNodes builds the correct address (ip:OrchestratorAPIPort). +// The fake gRPC server listens on consts.OrchestratorAPIPort, matching the +// port carried by the mocked service registration. func TestGetOrConnectNode_CacheMiss_DiscoversAndConnects(t *testing.T) { t.Parallel() @@ -279,15 +280,17 @@ func TestGetOrConnectNode_CacheMiss_DiscoversAndConnects(t *testing.T) { listenAddr := fmt.Sprintf("127.0.0.1:%d", consts.OrchestratorAPIPort) startFakeOrchestratorGRPC(t, orchestratorNodeID, listenAddr) - // 2. Mock Nomad HTTP API returning a single ready node at 127.0.0.1. + // 2. Mock Nomad HTTP API returning a single service registration at + // 127.0.0.1. nomadClient := newNomadMock(t, func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/v1/nodes" { - resp := []*nomadapi.NodeListStub{ + if r.URL.Path == "/v1/service/orchestrator" { + resp := []*nomadapi.ServiceRegistration{ { - ID: nomadFullID, - Address: "127.0.0.1", - Status: "ready", - NodePool: "default", + ID: "_nomad-task-def-orchestrator", + ServiceName: "orchestrator", + NodeID: nomadFullID, + Address: "127.0.0.1", + Port: int(consts.OrchestratorAPIPort), }, } w.Header().Set("Content-Type", "application/json") diff --git a/packages/api/internal/orchestrator/discovery/discovery.go b/packages/api/internal/orchestrator/discovery/discovery.go index 600ea577b4..f7ea7a318c 100644 --- a/packages/api/internal/orchestrator/discovery/discovery.go +++ b/packages/api/internal/orchestrator/discovery/discovery.go @@ -1,13 +1,24 @@ // Package discovery enumerates running orchestrator (Firecracker host) instances // for the API to route sandbox calls to. // -// The Discovery interface has two implementations: +// The Discovery interface has several implementations: // -// - NomadDiscovery: queries the local Nomad agent's HTTP /v1/nodes endpoint. -// Used by the original Nomad-based deploy where every Nomad client node -// also runs an orchestrator process. +// - NewNomad: lists registrations of the orchestrator's Nomad-native +// services via the local Nomad agent's HTTP /v1/service/ endpoint +// (one query per configured name, unioned). Used by the Nomad-based +// deploy; node-pool- and job-agnostic. // -// - KubernetesDiscovery: lists pods of the orchestrator DaemonSet via the +// - NewNomadNodePool: the legacy implementation that lists ready Nomad +// NODES in a node pool and assumes each runs an orchestrator on the +// well-known port. Kept as a migration fallback for orchestrator jobs +// whose service registrations are broken (empty Address); see +// nomad_node_pool.go. +// +// - NewMerged: unions a primary and a fallback Discovery, deduplicated by +// ShortID with primary winning. Used to combine the two Nomad backends +// during the migration. +// +// - NewKubernetes: lists pods of the orchestrator DaemonSet via the // in-cluster K8s API. Used by the K8s deploy. // // The shape of the returned []Node mirrors what the existing diff --git a/packages/api/internal/orchestrator/discovery/merged.go b/packages/api/internal/orchestrator/discovery/merged.go new file mode 100644 index 0000000000..fdadd11ebb --- /dev/null +++ b/packages/api/internal/orchestrator/discovery/merged.go @@ -0,0 +1,53 @@ +package discovery + +import ( + "context" + "fmt" + "slices" + + "github.com/samber/lo" +) + +// mergedDiscovery unions a primary and a fallback Discovery, deduplicated by +// ShortID with the primary entry winning on conflict. It bridges the migration +// from node-pool-based to service-based discovery: the service registration +// carries the real bound port, so it wins when both backends report the same +// node. +// +// ListNodes fails if either backend fails: the caller treats a discovery +// error as "skip this cycle", which beats silently acting on a partial node +// list. +type mergedDiscovery struct { + primary Discovery + fallback Discovery +} + +// NewMerged creates a Discovery that unions primary's and fallback's nodes, +// deduplicated by ShortID with primary taking precedence. +func NewMerged(primary, fallback Discovery) Discovery { + return &mergedDiscovery{ + primary: primary, + fallback: fallback, + } +} + +func (d *mergedDiscovery) ListNodes(ctx context.Context) ([]Node, error) { + ctx, span := tracer.Start(ctx, "list-merged-nodes") + defer span.End() + + primaryNodes, err := d.primary.ListNodes(ctx) + if err != nil { + return nil, fmt.Errorf("primary discovery failed: %w", err) + } + + fallbackNodes, err := d.fallback.ListNodes(ctx) + if err != nil { + return nil, fmt.Errorf("fallback discovery failed: %w", err) + } + + // UniqBy keeps the first occurrence per key, so primary entries shadow + // fallback entries with the same ShortID. + return lo.UniqBy(slices.Concat(primaryNodes, fallbackNodes), func(n Node) string { + return n.ShortID + }), nil +} diff --git a/packages/api/internal/orchestrator/discovery/merged_test.go b/packages/api/internal/orchestrator/discovery/merged_test.go new file mode 100644 index 0000000000..4cfebc43ff --- /dev/null +++ b/packages/api/internal/orchestrator/discovery/merged_test.go @@ -0,0 +1,79 @@ +package discovery + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// staticDiscovery is a Discovery stub returning canned nodes or an error. +type staticDiscovery struct { + nodes []Node + err error +} + +func (d *staticDiscovery) ListNodes(context.Context) ([]Node, error) { + if d.err != nil { + return nil, d.err + } + + return d.nodes, nil +} + +// The primary (service-based) entry must win on conflict: it carries the real +// bound port. Dedupe itself is lo.UniqBy's job; this pins the concatenation +// order that makes primary win. +func TestMergedDiscovery_PrimaryWinsOnConflict(t *testing.T) { + t.Parallel() + + primary := &staticDiscovery{nodes: []Node{ + {ShortID: "aaaaaaaa", IPAddress: "10.0.0.1", OrchestratorAddress: "10.0.0.1:6123"}, + }} + fallback := &staticDiscovery{nodes: []Node{ + {ShortID: "aaaaaaaa", IPAddress: "10.0.0.1", OrchestratorAddress: "10.0.0.1:5008"}, + }} + + d := NewMerged(primary, fallback) + nodes, err := d.ListNodes(t.Context()) + require.NoError(t, err) + require.Len(t, nodes, 1) + + assert.Equal(t, "aaaaaaaa", nodes[0].ShortID) + assert.Equal(t, "10.0.0.1:6123", nodes[0].OrchestratorAddress, "primary (service-based) entry must win on conflict") +} + +// A primary failure fails the whole listing. +func TestMergedDiscovery_PrimaryError(t *testing.T) { + t.Parallel() + + primaryErr := errors.New("nomad agent unreachable") + primary := &staticDiscovery{err: primaryErr} + fallback := &staticDiscovery{nodes: []Node{ + {ShortID: "bbbbbbbb", IPAddress: "10.0.0.2", OrchestratorAddress: "10.0.0.2:5008"}, + }} + + d := NewMerged(primary, fallback) + nodes, err := d.ListNodes(t.Context()) + require.ErrorIs(t, err, primaryErr) + assert.Nil(t, nodes) +} + +// A fallback failure also fails the whole listing: no silent degradation to a +// partial node list. +func TestMergedDiscovery_FallbackError(t *testing.T) { + t.Parallel() + + fallbackErr := errors.New("nomad agent unreachable") + primary := &staticDiscovery{nodes: []Node{ + {ShortID: "aaaaaaaa", IPAddress: "10.0.0.1", OrchestratorAddress: "10.0.0.1:5008"}, + }} + fallback := &staticDiscovery{err: fallbackErr} + + d := NewMerged(primary, fallback) + nodes, err := d.ListNodes(t.Context()) + require.ErrorIs(t, err, fallbackErr) + assert.Nil(t, nodes) +} diff --git a/packages/api/internal/orchestrator/discovery/nomad.go b/packages/api/internal/orchestrator/discovery/nomad.go index 2073bd505d..fde71d7c10 100644 --- a/packages/api/internal/orchestrator/discovery/nomad.go +++ b/packages/api/internal/orchestrator/discovery/nomad.go @@ -13,45 +13,78 @@ import ( // nomadDiscovery implements Discovery against a local Nomad HTTP agent. // -// In the Nomad deploy each client node has the orchestrator binary running -// directly (raw_exec) and listening on consts.OrchestratorAPIPort, so every -// "ready" Nomad node in the configured pool is an orchestrator. +// It lists registrations of the Nomad-native services (provider = "nomad") +// that every orchestrator jobspec registers, via GET /v1/service/ per +// configured name, unioned. This is node-pool- and job-agnostic: any +// allocation that registers one of the services is discovered, which keeps +// discovery correct while orchestrators migrate between jobs, node pools, +// and service names. +// +// The endpoint returns registrations regardless of health-check status; that +// is acceptable because callers health-gate every discovered instance through +// the orchestrator's gRPC ServiceInfo probe after connecting. type nomadDiscovery struct { - client *nomadapi.Client - nodePool string + client *nomadapi.Client + serviceNames []string } -// NewNomad creates a Nomad-backed Discovery. -func NewNomad(client *nomadapi.Client, nodePool string) Discovery { +// NewNomad creates a Nomad-backed Discovery that enumerates registrations of +// the given Nomad-native service names (e.g. "orchestrator"), unioned and +// deduplicated by node. +func NewNomad(client *nomadapi.Client, serviceNames []string) Discovery { return &nomadDiscovery{ - client: client, - nodePool: nodePool, + client: client, + serviceNames: serviceNames, } } func (d *nomadDiscovery) ListNodes(ctx context.Context) ([]Node, error) { - ctx, span := tracer.Start(ctx, "list-nomad-nodes") + ctx, span := tracer.Start(ctx, "list-nomad-orchestrator-service") defer span.End() - options := &nomadapi.QueryOptions{ - Filter: fmt.Sprintf("Status == %q and NodePool == %q", "ready", d.nodePool), - } - nodes, _, err := d.client.Nodes().List(options.WithContext(ctx)) - if err != nil { - return nil, err - } + opts := (&nomadapi.QueryOptions{}).WithContext(ctx) + + var out []Node + seen := make(map[string]struct{}) + for _, serviceName := range d.serviceNames { + // Fail if ANY service listing fails, mirroring mergedDiscovery: + // silently degrading to a subset of services would make the sync + // loop treat the missing orchestrators as gone. + regs, _, err := d.client.Services().Get(serviceName, opts) + if err != nil { + return nil, fmt.Errorf("listing nomad service %q: %w", serviceName, err) + } + + for _, reg := range regs { + // Skip unroutable registrations so callers never dial an empty host. + if reg.Address == "" { + continue + } + + // At most one orchestrator runs per node (static port + host + // flock), so duplicate registrations on a node — within a service + // (e.g. a stopping allocation whose registration has not been + // reaped yet) or across services — are transient. Collapse them + // to keep ShortIDs unique; earlier service names win. + if _, ok := seen[reg.NodeID]; ok { + continue + } + seen[reg.NodeID] = struct{}{} + + // Truncated Nomad node ID, matching what the node-listing + // implementation produced, so node identity is stable across the + // discovery-backend switch. + shortID := reg.NodeID + if len(shortID) > consts.NodeIDLength { + shortID = shortID[:consts.NodeIDLength] + } - out := make([]Node, 0, len(nodes)) - for _, n := range nodes { - shortID := n.ID - if len(shortID) > consts.NodeIDLength { - shortID = shortID[:consts.NodeIDLength] + out = append(out, Node{ + ShortID: shortID, + IPAddress: reg.Address, + OrchestratorAddress: net.JoinHostPort(reg.Address, strconv.Itoa(reg.Port)), + }) } - out = append(out, Node{ - ShortID: shortID, - IPAddress: n.Address, - OrchestratorAddress: net.JoinHostPort(n.Address, strconv.FormatUint(uint64(consts.OrchestratorAPIPort), 10)), - }) } return out, nil diff --git a/packages/api/internal/orchestrator/discovery/nomad_node_pool.go b/packages/api/internal/orchestrator/discovery/nomad_node_pool.go new file mode 100644 index 0000000000..724b0c2d7c --- /dev/null +++ b/packages/api/internal/orchestrator/discovery/nomad_node_pool.go @@ -0,0 +1,72 @@ +package discovery + +import ( + "context" + "fmt" + "net" + "strconv" + + nomadapi "github.com/hashicorp/nomad/api" + + "github.com/e2b-dev/infra/packages/shared/pkg/consts" +) + +// nomadNodePoolDiscovery implements Discovery by listing Nomad NODES in a +// node pool via the local Nomad agent (GET /v1/nodes), assuming every "ready" +// node in the pool runs an orchestrator listening on the well-known +// consts.OrchestratorAPIPort. This is the pre-service-discovery +// implementation, re-added verbatim as a MIGRATION FALLBACK: orchestrator +// jobs deployed from jobspecs that predate the service port-label fix +// register their service with an empty Address, so the service-based +// discovery (NewNomad) skips them as unroutable. Unioning in this backend +// (see NewMerged) removes any rollout ordering constraint. Once no legacy +// jobs remain, disable via NOMAD_ORCHESTRATOR_LEGACY_DISCOVERY_ENABLED=false +// and delete this file. +// +// Note on draining: Nomad nodes keep Status == "ready" while they drain, so +// a draining orchestrator whose service registrations were already +// deregistered (PreKill) remains discoverable here for the drain window. +type nomadNodePoolDiscovery struct { + client *nomadapi.Client + nodePool string +} + +// NewNomadNodePool creates a Nomad-backed Discovery that lists ready nodes in +// the given node pool. See the nomadNodePoolDiscovery doc for why this exists. +func NewNomadNodePool(client *nomadapi.Client, nodePool string) Discovery { + return &nomadNodePoolDiscovery{ + client: client, + nodePool: nodePool, + } +} + +func (d *nomadNodePoolDiscovery) ListNodes(ctx context.Context) ([]Node, error) { + ctx, span := tracer.Start(ctx, "list-nomad-nodes") + defer span.End() + + options := &nomadapi.QueryOptions{ + Filter: fmt.Sprintf("Status == %q and NodePool == %q", "ready", d.nodePool), + } + nodes, _, err := d.client.Nodes().List(options.WithContext(ctx)) + if err != nil { + return nil, err + } + + out := make([]Node, 0, len(nodes)) + for _, n := range nodes { + // Truncated the same way the service-based backend truncates the + // registration's NodeID, so both backends yield the SAME ShortID for + // the same node and the union can never create duplicate identities. + shortID := n.ID + if len(shortID) > consts.NodeIDLength { + shortID = shortID[:consts.NodeIDLength] + } + out = append(out, Node{ + ShortID: shortID, + IPAddress: n.Address, + OrchestratorAddress: net.JoinHostPort(n.Address, strconv.FormatUint(uint64(consts.OrchestratorAPIPort), 10)), + }) + } + + return out, nil +} diff --git a/packages/api/internal/orchestrator/discovery/nomad_node_pool_test.go b/packages/api/internal/orchestrator/discovery/nomad_node_pool_test.go new file mode 100644 index 0000000000..eacd539f74 --- /dev/null +++ b/packages/api/internal/orchestrator/discovery/nomad_node_pool_test.go @@ -0,0 +1,86 @@ +package discovery + +import ( + "encoding/json" + "net" + "net/http" + "net/http/httptest" + "strconv" + "testing" + + nomadapi "github.com/hashicorp/nomad/api" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/shared/pkg/consts" +) + +// newNomadNodesMock returns a nomad API client backed by an httptest server +// that serves the given node stubs from GET /v1/nodes, mirroring the agent +// endpoint's behavior. The filter expression the client sends is captured +// into gotFilter (server-side filtering itself is not simulated). +func newNomadNodesMock(t *testing.T, stubs []*nomadapi.NodeListStub, gotFilter *string) *nomadapi.Client { + t.Helper() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/v1/nodes" { + if gotFilter != nil { + *gotFilter = r.URL.Query().Get("filter") + } + if stubs == nil { + stubs = []*nomadapi.NodeListStub{} + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(stubs) + + return + } + + http.NotFound(w, r) + })) + t.Cleanup(srv.Close) + + client, err := nomadapi.NewClient(&nomadapi.Config{Address: srv.URL}) + require.NoError(t, err) + + return client +} + +// TestNomadNodePoolDiscovery_MapsNodes verifies field mapping: ShortID is the +// truncated Nomad node ID (identical to what the service-based backend +// produces for the same node, so the merged union cannot create duplicate +// identities) and OrchestratorAddress uses the well-known orchestrator port. +func TestNomadNodePoolDiscovery_MapsNodes(t *testing.T) { + t.Parallel() + + fullNodeID := "aabbccdd11223344aabbccdd11223344aabbccdd" + client := newNomadNodesMock(t, []*nomadapi.NodeListStub{ + {ID: fullNodeID, Address: "10.0.0.1", Status: "ready", NodePool: "default"}, + }, nil) + + d := NewNomadNodePool(client, "default") + nodes, err := d.ListNodes(t.Context()) + require.NoError(t, err) + require.Len(t, nodes, 1) + + port := strconv.Itoa(int(consts.OrchestratorAPIPort)) + assert.Equal(t, fullNodeID[:consts.NodeIDLength], nodes[0].ShortID) + assert.Equal(t, "10.0.0.1", nodes[0].IPAddress) + assert.Equal(t, net.JoinHostPort("10.0.0.1", port), nodes[0].OrchestratorAddress) +} + +// TestNomadNodePoolDiscovery_FiltersByStatusAndPool verifies the server-side +// filter expression restricts the listing to ready nodes in the configured +// pool, matching the pre-service-discovery implementation. +func TestNomadNodePoolDiscovery_FiltersByStatusAndPool(t *testing.T) { + t.Parallel() + + var gotFilter string + client := newNomadNodesMock(t, nil, &gotFilter) + + d := NewNomadNodePool(client, "orchestrator-pool") + nodes, err := d.ListNodes(t.Context()) + require.NoError(t, err) + assert.Empty(t, nodes) + assert.Equal(t, `Status == "ready" and NodePool == "orchestrator-pool"`, gotFilter) +} diff --git a/packages/api/internal/orchestrator/discovery/nomad_test.go b/packages/api/internal/orchestrator/discovery/nomad_test.go new file mode 100644 index 0000000000..d87e15bfb3 --- /dev/null +++ b/packages/api/internal/orchestrator/discovery/nomad_test.go @@ -0,0 +1,162 @@ +package discovery + +import ( + "encoding/json" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + + nomadapi "github.com/hashicorp/nomad/api" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/shared/pkg/consts" +) + +const testServiceName = "orchestrator" + +// newNomadServicesMock returns a nomad API client backed by an httptest +// server that serves the given registrations per service name from +// GET /v1/service/, mirroring the agent endpoint's behavior +// (200 + JSON array, empty array when there are no registrations). Unlisted +// service names return 200 + [] as well, matching the real agent. +func newNomadServicesMock(t *testing.T, regsByService map[string][]*nomadapi.ServiceRegistration) *nomadapi.Client { + t.Helper() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + name, ok := strings.CutPrefix(r.URL.Path, "/v1/service/") + if !ok { + http.NotFound(w, r) + + return + } + + regs := regsByService[name] + if regs == nil { + regs = []*nomadapi.ServiceRegistration{} + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(regs) + })) + t.Cleanup(srv.Close) + + client, err := nomadapi.NewClient(&nomadapi.Config{Address: srv.URL}) + require.NoError(t, err) + + return client +} + +// newNomadServiceMock is a single-service convenience wrapper around +// newNomadServicesMock. +func newNomadServiceMock(t *testing.T, regs []*nomadapi.ServiceRegistration) *nomadapi.Client { + t.Helper() + + return newNomadServicesMock(t, map[string][]*nomadapi.ServiceRegistration{testServiceName: regs}) +} + +// TestNomadDiscovery_MapsRegistrations verifies field mapping: ShortID is the +// truncated Nomad node ID (stable across the node-listing -> service-listing +// backend switch) and OrchestratorAddress uses the registration's bound port. +func TestNomadDiscovery_MapsRegistrations(t *testing.T) { + t.Parallel() + + fullNodeID := "aabbccdd11223344aabbccdd11223344aabbccdd" + client := newNomadServiceMock(t, []*nomadapi.ServiceRegistration{ + { + ID: "_nomad-task-a-orchestrator", + ServiceName: testServiceName, + NodeID: fullNodeID, + Address: "10.0.0.1", + Port: 5008, + }, + }) + + d := NewNomad(client, []string{testServiceName}) + nodes, err := d.ListNodes(t.Context()) + require.NoError(t, err) + require.Len(t, nodes, 1) + + assert.Equal(t, fullNodeID[:consts.NodeIDLength], nodes[0].ShortID) + assert.Equal(t, "10.0.0.1", nodes[0].IPAddress) + assert.Equal(t, net.JoinHostPort("10.0.0.1", "5008"), nodes[0].OrchestratorAddress) +} + +// TestNomadDiscovery_DeduplicatesByNode ensures that two registrations on the +// same node (a transient state, e.g. a stopping allocation whose registration +// has not been reaped yet) collapse into a single discovered node, keeping +// ShortIDs unique for callers that key on them. +func TestNomadDiscovery_DeduplicatesByNode(t *testing.T) { + t.Parallel() + + nodeID := "aabbccdd11223344aabbccdd11223344aabbccdd" + client := newNomadServiceMock(t, []*nomadapi.ServiceRegistration{ + {ID: "_nomad-task-old-orchestrator", ServiceName: testServiceName, NodeID: nodeID, Address: "10.0.0.1", Port: 5008}, + {ID: "_nomad-task-new-orchestrator", ServiceName: testServiceName, NodeID: nodeID, Address: "10.0.0.1", Port: 5008}, + {ID: "_nomad-task-other-orchestrator", ServiceName: testServiceName, NodeID: "eeff00112233445566778899aabbccddeeff0011", Address: "10.0.0.2", Port: 5008}, + }) + + d := NewNomad(client, []string{testServiceName}) + nodes, err := d.ListNodes(t.Context()) + require.NoError(t, err) + require.Len(t, nodes, 2) + assert.NotEqual(t, nodes[0].ShortID, nodes[1].ShortID) +} + +// TestNomadDiscovery_UnionsServices ensures registrations from multiple +// configured service names are unioned, with nodes registered under several +// names collapsed to a single entry (earlier service names win). +func TestNomadDiscovery_UnionsServices(t *testing.T) { + t.Parallel() + + sharedNodeID := "aabbccdd11223344aabbccdd11223344aabbccdd" + client := newNomadServicesMock(t, map[string][]*nomadapi.ServiceRegistration{ + testServiceName: { + {ID: "_nomad-task-a-orchestrator", ServiceName: testServiceName, NodeID: sharedNodeID, Address: "10.0.0.1", Port: 5008}, + }, + "orchestrator-canary": { + // Same node also registered under the second service name: must + // collapse, with the first service's entry winning. + {ID: "_nomad-task-a-canary", ServiceName: "orchestrator-canary", NodeID: sharedNodeID, Address: "10.0.0.1", Port: 6008}, + {ID: "_nomad-task-b-canary", ServiceName: "orchestrator-canary", NodeID: "eeff00112233445566778899aabbccddeeff0011", Address: "10.0.0.2", Port: 5008}, + }, + }) + + d := NewNomad(client, []string{testServiceName, "orchestrator-canary"}) + nodes, err := d.ListNodes(t.Context()) + require.NoError(t, err) + require.Len(t, nodes, 2) + + assert.Equal(t, net.JoinHostPort("10.0.0.1", "5008"), nodes[0].OrchestratorAddress) + assert.Equal(t, net.JoinHostPort("10.0.0.2", "5008"), nodes[1].OrchestratorAddress) +} + +// TestNomadDiscovery_SkipsMissingAddress ensures registrations without an +// address are excluded so callers never dial an empty host. +func TestNomadDiscovery_SkipsMissingAddress(t *testing.T) { + t.Parallel() + + client := newNomadServiceMock(t, []*nomadapi.ServiceRegistration{ + {ID: "_nomad-task-a-orchestrator", ServiceName: testServiceName, NodeID: "aabbccdd11223344aabbccdd11223344aabbccdd", Port: 5008}, + }) + + d := NewNomad(client, []string{testServiceName}) + nodes, err := d.ListNodes(t.Context()) + require.NoError(t, err) + assert.Empty(t, nodes) +} + +// TestNomadDiscovery_EmptyService ensures an empty registration list (the +// agent returns 200 + [] when no instances are registered) yields an empty, +// non-error result. +func TestNomadDiscovery_EmptyService(t *testing.T) { + t.Parallel() + + client := newNomadServiceMock(t, nil) + + d := NewNomad(client, []string{testServiceName}) + nodes, err := d.ListNodes(t.Context()) + require.NoError(t, err) + assert.Empty(t, nodes) +}