feat(api): discover orchestrators via nomad service (#3176)

Switch discovery to use nomad's service discovery instead of blindly
checking every node in the default nodepool. Lets us use multiple
nodepools for cutting over orchestrator from system jobs to service
jobs.

Existing orchestrator-ee system job already registers to the service.

Author: wj-e2b

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

packages/api/go.sum

@@ -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/<name> 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

packages/api/internal/cfg/model.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)
 	}
 

packages/api/internal/handlers/store.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")

packages/api/internal/orchestrator/client_test.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/<name> 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

packages/api/internal/orchestrator/discovery/discovery.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
+}

packages/api/internal/orchestrator/discovery/merged.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)
+}