refactor(api): retire the api's second discovery mechanism

GitOrigin-RevId: 136062996f2a78dd6cf552ddd817ef59c4e723e8

Author: dobrac

packages/api/go.mod

@@ -56,8 +56,6 @@ require (
 	golang.org/x/sync v0.22.0
 	google.golang.org/grpc v1.83.0
 	google.golang.org/protobuf v1.36.11
-	k8s.io/api v0.35.3
-	k8s.io/apimachinery v0.35.3
 	k8s.io/client-go v0.35.3
 )
 
@@ -417,6 +415,8 @@ require (
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
+	k8s.io/api v0.35.3 // indirect
+	k8s.io/apimachinery v0.35.3 // indirect
 	k8s.io/klog/v2 v2.140.0 // indirect
 	k8s.io/kube-openapi v0.0.0-20250910181357-589584f1c912 // indirect
 	k8s.io/utils v0.0.0-20251002143259-bc988d571ff4 // indirect

packages/api/internal/cfg/model.go

@@ -46,11 +46,15 @@ type Config struct {
 	LokiURL      string `env:"LOKI_URL,required"`
 	LokiUser     string `env:"LOKI_USER"`
 
-	// ServiceDiscoveryProvider selects how the API discovers orchestrator and template-manager instances.
+	// ServiceDiscoveryProvider selects how the API discovers orchestrator and
+	// template-manager instances. Left unset it is nomad, except in a local
+	// environment, which has no Nomad agent — see serviceDiscoveryProvider. It
+	// carries no envDefault precisely so that "unset" stays distinguishable
+	// from "set to nomad", which an operator running a local Nomad agent needs.
 	// Allowed values:
 	//   "nomad"      (default) - query the local Nomad agent's HTTP API.
 	//   "kubernetes"           - list pods via the in-cluster K8s API.
-	ServiceDiscoveryProvider string `env:"SERVICE_DISCOVERY_PROVIDER" envDefault:"nomad"`
+	ServiceDiscoveryProvider string `env:"SERVICE_DISCOVERY_PROVIDER"`
 
 	NomadAddress string `env:"NOMAD_ADDRESS" envDefault:"http://localhost:4646"`
 	NomadToken   string `env:"NOMAD_TOKEN"`
@@ -271,7 +275,7 @@ func Parse() (Config, error) {
 		config.AuthDBConnectionString = config.PostgresConnectionString
 	}
 
-	if !slices.Contains([]string{ServiceDiscoveryProviderNomad, ServiceDiscoveryProviderKubernetes, ServiceDiscoveryProviderLocal}, config.ServiceDiscoveryProvider) {
+	if !slices.Contains([]string{"", ServiceDiscoveryProviderNomad, ServiceDiscoveryProviderKubernetes, ServiceDiscoveryProviderLocal}, config.ServiceDiscoveryProvider) {
 		return config, newFailureError(
 			FailureConditionInvalidServiceDiscoveryProvider,
 			fmt.Sprintf("invalid service discovery provider: %s", config.ServiceDiscoveryProvider),

packages/api/internal/clusters/cluster.go

@@ -5,9 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"math/rand/v2"
-	"net"
 	"net/http"
-	"strconv"
 	"sync"
 	"time"
 
@@ -16,7 +14,6 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/e2b-dev/infra/packages/api/internal/cfg"
-	"github.com/e2b-dev/infra/packages/api/internal/clusters/discovery"
 	clickhouse "github.com/e2b-dev/infra/packages/clickhouse/pkg"
 	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
@@ -25,6 +22,7 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logs/loki"
 	"github.com/e2b-dev/infra/packages/shared/pkg/machineinfo"
+	"github.com/e2b-dev/infra/packages/shared/pkg/servicediscovery"
 	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
 	"github.com/e2b-dev/infra/packages/shared/pkg/synchronization"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
@@ -43,7 +41,7 @@ type Cluster struct {
 	AuthOrgID     string
 
 	instances       *smap.Map[*Instance]
-	synchronization *synchronization.Synchronize[discovery.Item, *Instance]
+	synchronization *synchronization.Synchronize[servicediscovery.Instance, *Instance]
 	resources       ClusterResource
 }
 
@@ -57,7 +55,7 @@ func NewCluster(
 	domain *string,
 	authOrgID string,
 	sandboxes *smap.Map[*Instance],
-	synchronization *synchronization.Synchronize[discovery.Item, *Instance],
+	synchronization *synchronization.Synchronize[servicediscovery.Instance, *Instance],
 	resources ClusterResource,
 ) *Cluster {
 	return &Cluster{
@@ -73,7 +71,7 @@ func NewCluster(
 func newLocalCluster(
 	ctx context.Context,
 	tel *telemetry.Client,
-	storeDiscovery discovery.Discovery,
+	storeDiscovery servicediscovery.Discoverer,
 	clickhouse clickhouse.Clickhouse,
 	queryLogsProvider *loki.LokiQueryProvider,
 	sandboxLogsReader ClickhouseLogsReader,
@@ -83,9 +81,9 @@ func newLocalCluster(
 	clusterID := consts.LocalClusterID
 
 	instances := smap.New[*Instance]()
-	instanceCreation := func(ctx context.Context, item discovery.Item) (*Instance, error) {
+	instanceCreation := func(ctx context.Context, item servicediscovery.Instance) (*Instance, error) {
 		// For local cluster we are doing direct connection to instance IP and API port and without additional cluster auth.
-		return newInstance(ctx, tel, nil, clusterID, item, net.JoinHostPort(item.LocalIPAddress, strconv.FormatUint(uint64(item.LocalInstanceApiPort), 10)), false)
+		return newInstance(ctx, tel, nil, clusterID, item, item.Address(), false)
 	}
 
 	store := instancesSyncStore{clusterID: clusterID, instances: instances, discovery: storeDiscovery, instanceCreation: instanceCreation}
@@ -105,6 +103,14 @@ func newLocalCluster(
 	return c
 }
 
+// remoteInstanceAuthorization routes a call through the remote proxy to one
+// instance. The discovery ID is that instance's service id: the edge API
+// reports it per run, which is what the proxy keys on — the machine would send
+// the call to whatever is on it now.
+func remoteInstanceAuthorization(secret string, tls bool, sd servicediscovery.Instance) *instanceAuthorization {
+	return &instanceAuthorization{secret: secret, tls: tls, serviceInstanceID: sd.WorkloadID}
+}
+
 func newRemoteCluster(
 	ctx context.Context,
 	tel *telemetry.Client,
@@ -142,14 +148,13 @@ func newRemoteCluster(
 	}
 
 	instances := smap.New[*Instance]()
-	instanceCreation := func(ctx context.Context, item discovery.Item) (*Instance, error) {
-		// For remote cluster we are doing connection to endpoint that works as gRPC proxy and handles auth and routing for us.
-		auth := &instanceAuthorization{secret: secret, tls: endpointTLS, serviceInstanceID: item.InstanceID}
-
-		return newInstance(ctx, tel, auth, clusterID, item, endpoint, endpointTLS)
+	instanceCreation := func(ctx context.Context, item servicediscovery.Instance) (*Instance, error) {
+		// A remote cluster is reached through an endpoint acting as a gRPC proxy,
+		// which handles auth and routing for us.
+		return newInstance(ctx, tel, remoteInstanceAuthorization(secret, endpointTLS, item), clusterID, item, endpoint, endpointTLS)
 	}
 
-	storeDiscovery := discovery.NewRemoteServiceDiscovery(clusterID, httpClient)
+	storeDiscovery := servicediscovery.NewRemote(httpClient)
 	store := instancesSyncStore{clusterID: clusterID, instances: instances, instanceCreation: instanceCreation, discovery: storeDiscovery}
 
 	c := NewCluster(
@@ -191,17 +196,61 @@ func (c *Cluster) Close(ctx context.Context) error {
 	return nil
 }
 
+// Scanned rather than looked up: the pool is keyed by the running process, and
+// a build persists the machine it was placed on (env_builds.cluster_node_id),
+// so this resolves the machine to whichever process is on it now.
 func (c *Cluster) GetTemplateBuilderByNodeID(nodeID string) (*Instance, error) {
-	instance, found := c.instances.Get(nodeID)
+	instance, found := builderOnNode(c.instances, nodeID)
 	if !found {
 		return nil, ErrTemplateBuilderNotFound
 	}
 
-	if info := instance.GetInfo(); info.Status == infogrpc.ServiceInfoStatus_Unhealthy || !info.IsBuilder {
-		return nil, ErrTemplateBuilderNotFound
+	return instance, nil
+}
+
+// instanceOnNode finds an instance running on a machine, preferring one that is
+// not unhealthy. Like builderOnNode this cannot be a map lookup, but it applies
+// no further filter: callers that only need to reach the machine should not
+// start rejecting instances the old lookup accepted.
+func instanceOnNode(instances *smap.Map[*Instance], nodeID string) (*Instance, bool) {
+	var unhealthy *Instance
+
+	for _, instance := range instances.Items() {
+		if instance.NodeID != nodeID {
+			continue
+		}
+
+		if instance.GetInfo().Status == infogrpc.ServiceInfoStatus_Unhealthy {
+			unhealthy = instance
+
+			continue
+		}
+
+		return instance, true
 	}
 
-	return instance, nil
+	return unhealthy, unhealthy != nil
+}
+
+// builderOnNode finds the usable template builder running on a machine. The
+// pool is keyed by process, so this cannot be a map lookup, and every entry on
+// the machine has to be considered rather than the first: a restart leaves the
+// dead run and the live one sharing a machine for a sync round, and stopping
+// early would answer with whichever the map happened to yield.
+func builderOnNode(instances *smap.Map[*Instance], nodeID string) (*Instance, bool) {
+	for _, instance := range instances.Items() {
+		if instance.NodeID != nodeID {
+			continue
+		}
+
+		if info := instance.GetInfo(); info.Status == infogrpc.ServiceInfoStatus_Unhealthy || !info.IsBuilder {
+			continue
+		}
+
+		return instance, true
+	}
+
+	return nil, false
 }
 
 func (c *Cluster) GetByServiceInstanceID(serviceInstanceID string) (*Instance, bool) {

packages/api/internal/clusters/cluster_builder_lookup_test.go

@@ -0,0 +1,134 @@
+package clusters
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	infogrpc "github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator-info"
+	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
+)
+
+// clusterWithInstances builds only the fields the builder lookup reads.
+func clusterWithInstances(t *testing.T, instances ...*Instance) *Cluster {
+	t.Helper()
+
+	pool := smap.New[*Instance]()
+	for _, i := range instances {
+		pool.Insert(i.workloadID, i)
+	}
+
+	return &Cluster{ID: uuid.New(), instances: pool}
+}
+
+func builderOn(nodeID, processID string) *Instance {
+	return &Instance{
+		workloadID: processID,
+		NodeID:     nodeID,
+		isBuilder:  true,
+		status:     infogrpc.ServiceInfoStatus_Healthy,
+	}
+}
+
+// A build persists the machine it was placed on (env_builds.cluster_node_id)
+// and re-resolves it later. The pool is keyed by process now, so this lookup
+// has to find the machine some other way or every in-flight build breaks.
+func TestGetTemplateBuilderByNodeID_ResolvesAPersistedMachine(t *testing.T) {
+	t.Parallel()
+
+	wanted := builderOn("node-b", "alloc-2")
+	c := clusterWithInstances(t, builderOn("node-a", "alloc-1"), wanted, builderOn("node-c", "alloc-3"))
+
+	got, err := c.GetTemplateBuilderByNodeID("node-b")
+	require.NoError(t, err)
+	assert.Same(t, wanted, got)
+}
+
+// After a restart the machine is the same and the process is not, which is the
+// case the re-keying exists for: the persisted node id must resolve to whatever
+// is running there now.
+func TestGetTemplateBuilderByNodeID_ResolvesToTheCurrentProcessAfterARestart(t *testing.T) {
+	t.Parallel()
+
+	replacement := builderOn("node-a", "alloc-2")
+	c := clusterWithInstances(t, replacement)
+
+	got, err := c.GetTemplateBuilderByNodeID("node-a")
+	require.NoError(t, err)
+	assert.Equal(t, "alloc-2", got.workloadID)
+}
+
+func TestGetTemplateBuilderByNodeID_RejectsAMachineItDoesNotHave(t *testing.T) {
+	t.Parallel()
+
+	c := clusterWithInstances(t, builderOn("node-a", "alloc-1"))
+
+	_, err := c.GetTemplateBuilderByNodeID("node-z")
+	require.ErrorIs(t, err, ErrTemplateBuilderNotFound)
+}
+
+func TestGetTemplateBuilderByNodeID_RejectsANonBuilderAndAnUnhealthyOne(t *testing.T) {
+	t.Parallel()
+
+	notABuilder := builderOn("node-a", "alloc-1")
+	notABuilder.isBuilder = false
+
+	unhealthy := builderOn("node-b", "alloc-2")
+	unhealthy.status = infogrpc.ServiceInfoStatus_Unhealthy
+
+	c := clusterWithInstances(t, notABuilder, unhealthy)
+
+	_, err := c.GetTemplateBuilderByNodeID("node-a")
+	require.ErrorIs(t, err, ErrTemplateBuilderNotFound)
+
+	_, err = c.GetTemplateBuilderByNodeID("node-b")
+	require.ErrorIs(t, err, ErrTemplateBuilderNotFound)
+}
+
+// A restart leaves the dead run and the live one on the same machine for a sync
+// round — the state the pool re-keying deliberately allows. Answering with
+// whichever entry the map yielded first sent builds to the dead one, or
+// reported the machine as having no builder at all.
+//
+// Repeated because the pool is a map: one pass only exercises whichever order
+// Go happened to randomise into, and the wrong answer is the order-dependent
+// one.
+func TestGetTemplateBuilderByNodeID_SkipsADeadRunSharingTheMachine(t *testing.T) {
+	t.Parallel()
+
+	dead := builderOn("node-a", "alloc-1")
+	dead.status = infogrpc.ServiceInfoStatus_Unhealthy
+	notABuilder := builderOn("node-a", "alloc-3")
+	notABuilder.isBuilder = false
+	live := builderOn("node-a", "alloc-2")
+
+	c := clusterWithInstances(t, dead, notABuilder, live)
+
+	for range 200 {
+		got, err := c.GetTemplateBuilderByNodeID("node-a")
+		require.NoError(t, err, "the live builder must be found whatever order the pool yields")
+		require.Equal(t, "alloc-2", got.workloadID)
+	}
+}
+
+// Build-log streaming resolves the builder from the machine the build was
+// placed on, which is what the database persists. The pool is keyed by process
+// now, so a map lookup on the machine finds nothing on Nomad and Kubernetes,
+// where an allocation ID or pod name is never a node name — and the failure is
+// silent, falling back to persistent logs instead of erroring.
+func TestBuildLogs_ResolveTheBuilderByMachineNotByPoolKey(t *testing.T) {
+	t.Parallel()
+
+	builder := builderOn("node-a", "alloc-1")
+	pool := smap.New[*Instance]()
+	pool.Insert(builder.workloadID, builder)
+
+	_, foundByPoolKey := pool.Get("node-a")
+	require.False(t, foundByPoolKey, "the machine is not the pool key; this is the lookup that broke")
+
+	got, found := instanceOnNode(pool, "node-a")
+	require.True(t, found, "the builder on the machine must still be reachable")
+	assert.Equal(t, "alloc-1", got.workloadID)
+}

packages/api/internal/clusters/clusters_sync.go

@@ -9,14 +9,14 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/e2b-dev/infra/packages/api/internal/cfg"
-	"github.com/e2b-dev/infra/packages/api/internal/clusters/discovery"
 	clickhouse "github.com/e2b-dev/infra/packages/clickhouse/pkg"
 	"github.com/e2b-dev/infra/packages/db/client"
 	"github.com/e2b-dev/infra/packages/db/queries"
 	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logs/loki"
+	"github.com/e2b-dev/infra/packages/shared/pkg/servicediscovery"
 	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
 	"github.com/e2b-dev/infra/packages/shared/pkg/synchronization"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
@@ -47,7 +47,7 @@ func NewPool(
 	ctx context.Context,
 	tel *telemetry.Client,
 	db *client.Client,
-	localDiscovery discovery.Discovery,
+	localDiscovery servicediscovery.Discoverer,
 	queryMetricsProvider clickhouse.Clickhouse,
 	queryLogsProvider *loki.LokiQueryProvider,
 	sandboxLogsReader ClickhouseLogsReader,
@@ -116,7 +116,7 @@ type clustersSyncStore struct {
 	tel                  *telemetry.Client
 	clusters             *smap.Map[*Cluster]
 	local                *queries.Cluster
-	localDiscovery       discovery.Discovery
+	localDiscovery       servicediscovery.Discoverer
 	queryMetricsProvider clickhouse.Clickhouse
 	queryLogsProvider    *loki.LokiQueryProvider
 	sandboxLogsReader    ClickhouseLogsReader