feat(api): discover orchestrators via nomad service

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/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/<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/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)
 	}
 

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")

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/<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/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
+}

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)
+}