Merge pull request #1 from codeanker/feature/api-k8s-prefer-pod-ip

feat(api): dial Kubernetes-discovered orchestrators on their pod IP w…

Author: superbarne

docs/ARCHITECTURE.md

@@ -166,7 +166,9 @@ The control-plane entry point (Gin, OpenAPI-generated from `spec/openapi.yml`, p
   running-sandbox (Redis) and paused-snapshot (Postgres) state so they survive pause/resume and
   orchestrator re-sync; a fork starts a new workload and does not inherit them.
 - **Placement**: keeps a live map of orchestrator nodes (discovered via Nomad, Kubernetes, or a
-  static list). Chooses a node per sandbox with a **best-of-K** algorithm
+  static list; Kubernetes-discovered pods are dialed on their node's address, or on their pod
+  IP with `K8S_PREFER_POD_IP=true` for orchestrators that keep their own network namespace).
+  Chooses a node per sandbox with a **best-of-K** algorithm
   (`internal/orchestrator/placement/`): sample K ready nodes, score by the
   higher of CPU load and hugepage-pool load, pick the lowest; retry on
   exhausted nodes. A node that refuses a create (`ResourceExhausted`) stays

packages/api/internal/cfg/model.go

@@ -111,6 +111,14 @@ type Config struct {
 	K8sNamespace                       string `env:"K8S_NAMESPACE"                           envDefault:"default"`
 	K8sOrchestratorPodLabelSelector    string `env:"K8S_ORCHESTRATOR_POD_LABEL_SELECTOR"     envDefault:"app.kubernetes.io/name=orchestrator"`
 	K8sTemplateManagerPodLabelSelector string `env:"K8S_TEMPLATE_MANAGER_POD_LABEL_SELECTOR" envDefault:"app.kubernetes.io/name=template-manager"`
+	// K8sPreferPodIP dials the discovered orchestrator and template-manager
+	// pods on their pod IP instead of their node's. Off, the API dials
+	// status.hostIP, which is where a pod with host networking listens (the
+	// reference deploy). On, for a cluster whose orchestrator pods keep their
+	// own network namespace and are reachable only through the pod network,
+	// such as edge nodes behind an overlay that carries pod CIDRs but not node
+	// addresses.
+	K8sPreferPodIP bool `env:"K8S_PREFER_POD_IP"`
 
 	PostgresConnectionString string `env:"POSTGRES_CONNECTION_STRING,required,notEmpty"`
 	DBMaxOpenConnections     int32  `env:"DB_MAX_OPEN_CONNECTIONS"                      envDefault:"40"`

packages/api/internal/handlers/service_discovery_test.go

@@ -247,6 +247,50 @@ func TestNewServiceDiscovery_ExplicitProviderSurvivesALocalEnvironment(t *testin
 
 // The clusters registry owns orchestrator nodes only when the resolved provider
 // is the static local one — there both planes are the same single instance.
+// A pod with host networking reports the same address twice, so the fixture
+// pulls them apart: the node's address is where the reference deploy listens,
+// the pod's where an orchestrator with its own network namespace does.
+func TestNewServiceDiscovery_KubernetesDialsTheNodeUnlessPreferPodIP(t *testing.T) {
+	t.Parallel()
+
+	const (
+		nodeIP = "10.0.0.5"
+		podIP  = "10.42.1.3"
+	)
+
+	tests := map[string]struct {
+		preferPodIP bool
+		want        string
+	}{
+		"the default dials the node, where a host-network pod listens": {
+			preferPodIP: false, want: nodeIP,
+		},
+		"K8S_PREFER_POD_IP dials the pod network": {
+			preferPodIP: true, want: podIP,
+		},
+	}
+
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			orchestrator := discoveredPod("orchestrator-1", "orchestrator", nodeIP)
+			orchestrator.Status.PodIP = podIP
+			builder := discoveredPod("template-manager-1", "template-manager", nodeIP)
+			builder.Status.PodIP = podIP
+
+			config := discoveryConfig(cfg.ServiceDiscoveryProviderKubernetes, "")
+			config.K8sPreferPodIP = tt.preferPodIP
+
+			sd, err := newServiceDiscovery(t.Context(), config, fixedKubeClient(orchestrator, builder), cfg.ServiceDiscoveryProviderKubernetes)
+			require.NoError(t, err)
+
+			assert.Equal(t, []string{tt.want}, planeIPs(t, sd.nodes))
+			assert.Equal(t, []string{tt.want}, planeIPs(t, sd.templateBuilders))
+		})
+	}
+}
+
 // Deriving it from the environment instead builds a Nomad or Kubernetes node
 // plane that neither the periodic loop nor on-demand discovery ever consults,
 // so an orchestrator that is not also a template builder never registers.

packages/api/internal/handlers/store.go

@@ -44,6 +44,7 @@ import (
 	"github.com/e2b-dev/infra/packages/db/pkg/pool"
 	"github.com/e2b-dev/infra/packages/shared/pkg/apierrors"
 	sharedclusters "github.com/e2b-dev/infra/packages/shared/pkg/clusters/discovery"
+	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
 	"github.com/e2b-dev/infra/packages/shared/pkg/env"
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
 	managementv1 "github.com/e2b-dev/infra/packages/shared/pkg/grpc/contracts/secretsstore/management"
@@ -134,16 +135,23 @@ func newKubernetesServiceDiscovery(ctx context.Context, config cfg.Config, newKu
 		return serviceDiscovery{}, fmt.Errorf("kubernetes client: %w", err)
 	}
 
+	// Both planes dial the well-known orchestrator port. Which address they
+	// dial it on follows K8sPreferPodIP: the node's by default, the pod's when
+	// the orchestrator pods do not share their node's network namespace.
 	return serviceDiscovery{
-		nodes: kube.NewPods(
+		nodes: kube.NewPodsOnPort(
 			client,
 			config.K8sNamespace,
 			config.K8sOrchestratorPodLabelSelector,
+			consts.OrchestratorAPIPort,
+			!config.K8sPreferPodIP,
 		),
-		templateBuilders: kube.NewPods(
+		templateBuilders: kube.NewPodsOnPort(
 			client,
 			config.K8sNamespace,
 			config.K8sTemplateManagerPodLabelSelector,
+			consts.OrchestratorAPIPort,
+			!config.K8sPreferPodIP,
 		),
 	}, nil
 }

packages/shared/pkg/servicediscovery/kube/pods.go

@@ -18,9 +18,11 @@ var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/shared/pkg/servicedi
 // k8sDiscovery lists pods of the orchestrator DaemonSet via the K8s API
 // server.
 //
-// Those pods run with host networking, so each pod's status.HostIP equals
-// status.PodIP and is the address the orchestrator gRPC server listens on. A
-// pod is returned while it is Running and has an IP; the Ready condition is
+// In the reference deploy those pods run with host networking, so each pod's
+// status.HostIP equals status.PodIP and is the address the orchestrator gRPC
+// server listens on; a consumer whose orchestrators keep their own network
+// namespace asks for the pod IP instead (NewPodsOnPort with preferHostIP
+// false). A pod is returned while it is Running and has an IP; the Ready condition is
 // deliberately not consulted, because it is false for instances that still
 // have to be reachable — self-reported unhealthy, not yet probed, on a
 // NotReady node. Liveness is in-band and routing gates on that instead.