Don't assume an orchestrator port number, but pull it from nomad (#1368)

Author: djeebus

.editorconfig

@@ -0,0 +1,2 @@
+[*.hcl]
+indent_size = 2

iac/provider-gcp/nomad/jobs/orchestrator.hcl

@@ -5,23 +5,39 @@ job "orchestrator-${latest_orchestrator_job_id}" {
   priority = 90
 
   group "client-orchestrator" {
+    network {
+      mode = "host"
+
+      port "grpc" {
+        // todo: remove this once all API and client-proxy jobs
+        //  can pull the port number from nomad.
+        static = "${port}"
+      }
+
+      port "proxy" {
+        // todo: remove this once all API and client-proxy jobs
+        //  can pull the port number from nomad.
+        static = "${proxy_port}"
+      }
+    }
+
     service {
       name = "orchestrator"
-      port = "${port}"
+      port = "grpc"
 
       check {
         type         = "grpc"
         name         = "health"
         interval     = "20s"
         timeout      = "5s"
         grpc_use_tls = false
-        port         = "${port}"
+        port         = "grpc"
       }
     }
 
     service {
       name = "orchestrator-proxy"
-      port = "${proxy_port}"
+      port = "proxy"
     }
 
     task "check-placement" {
@@ -74,8 +90,8 @@ EOT
         CLICKHOUSE_CONNECTION_STRING = "${clickhouse_connection_string}"
         REDIS_URL                    = "${redis_url}"
         REDIS_CLUSTER_URL            = "${redis_cluster_url}"
-        GRPC_PORT                    = "${port}"
-        PROXY_PORT                   = "${proxy_port}"
+        GRPC_PORT                    = "$${NOMAD_PORT_grpc}"
+        PROXY_PORT                   = "$${NOMAD_PORT_proxy}"
         GIN_MODE                     = "release"
 
 %{ if launch_darkly_api_key != "" }

packages/api/internal/cfg/model.go

@@ -16,6 +16,13 @@ type Config struct {
 	NomadAddress string `env:"NOMAD_ADDRESS" envDefault:"http://localhost:4646"`
 	NomadToken   string `env:"NOMAD_TOKEN"`
 
+	// DefaultOrchestratorPort is the port that the Orchestrator listens on.
+	// Deprecated: Nomad knows which port the orchestrator is listening on. Keep this
+	// around temporarily until all nomad jobs have a port labeled "grpc", then this can be removed.
+	DefaultOrchestratorPort int `env:"ORCHESTRATOR_PORT" envDefault:"5008"`
+
+	OrchestratorPortLabel string `env:"ORCHESTRATOR_PORT_LABEL" envDefault:"grpc"`
+
 	PostgresConnectionString string `env:"POSTGRES_CONNECTION_STRING,required,notEmpty"`
 
 	PosthogAPIKey string `env:"POSTHOG_API_KEY"`

packages/api/internal/orchestrator/client.go

@@ -81,26 +81,121 @@ func (o *Orchestrator) listNomadNodes(ctx context.Context) ([]nodemanager.NomadS
 	defer listSpan.End()
 
 	options := &nomadapi.QueryOptions{
-		// TODO: Use variable for node pool name ("default")
-		Filter: "Status == \"ready\" and NodePool == \"default\"",
+		Filter: `ClientStatus == "running" and JobID contains "orchestrator-"`,
+		Params: map[string]string{"resources": "true"},
 	}
-	nomadNodes, _, err := o.nomadClient.Nodes().List(options.WithContext(ctx))
+	nomadAllocations, _, err := o.nomadClient.Allocations().List(options.WithContext(ctx))
 	if err != nil {
 		return nil, err
 	}
 
-	result := make([]nodemanager.NomadServiceDiscovery, 0, len(nomadNodes))
-	for _, n := range nomadNodes {
+	result := make([]nodemanager.NomadServiceDiscovery, 0, len(nomadAllocations))
+	for _, alloc := range nomadAllocations {
+		if !isHealthy(alloc) {
+			zap.L().Debug("Skipping unhealthy allocation", zap.String("allocation_id", alloc.ID))
+
+			continue
+		}
+
+		ip, port, ok := o.findPortInAllocation(alloc)
+		if !ok {
+			zap.L().Warn("Cannot find port in allocation",
+				zap.String("allocation_id", alloc.ID), zap.String("port_name", "grpc"))
+
+			continue
+		}
+
+		zap.L().Debug("Found port in allocation",
+			zap.String("allocation_id", alloc.ID),
+			zap.String("port_name", "grpc"),
+			zap.String("ip", ip),
+			zap.Int("port", port),
+		)
+
 		result = append(result, nodemanager.NomadServiceDiscovery{
-			NomadNodeShortID:    n.ID[:consts.NodeIDLength],
-			OrchestratorAddress: fmt.Sprintf("%s:%s", n.Address, consts.OrchestratorPort),
-			IPAddress:           n.Address,
+			NomadNodeShortID:    alloc.NodeID[:consts.NodeIDLength],
+			OrchestratorAddress: fmt.Sprintf("%s:%d", ip, port),
+			IPAddress:           ip,
 		})
 	}
 
 	return result, nil
 }
 
+func isHealthy(alloc *nomadapi.AllocationListStub) bool {
+	if alloc == nil {
+		zap.L().Warn("Allocation is nil")
+
+		return false
+	}
+
+	if alloc.DeploymentStatus == nil {
+		zap.L().Warn("Allocation deployment status is nil", zap.String("allocation_id", alloc.ID))
+
+		return false
+	}
+
+	if alloc.DeploymentStatus.Healthy == nil {
+		zap.L().Warn("Allocation deployment status healthy is nil", zap.String("allocation_id", alloc.ID))
+
+		return false
+	}
+
+	return *alloc.DeploymentStatus.Healthy
+}
+
+func (o *Orchestrator) findPortInAllocation(allocation *nomadapi.AllocationListStub) (string, int, bool) {
+	if allocation == nil {
+		return "", 0, false
+	}
+
+	if allocation.AllocatedResources == nil {
+		return "", 0, false
+	}
+
+	for _, task := range allocation.AllocatedResources.Tasks {
+		for _, network := range task.Networks {
+			host, port, ok := o.findPortInNetwork(network)
+			if ok {
+				return host, port, true
+			}
+		}
+	}
+
+	for _, net := range allocation.AllocatedResources.Shared.Networks {
+		host, port, ok := o.findPortInNetwork(net)
+		if ok {
+			return host, port, true
+		}
+	}
+
+	return "", 0, false
+}
+
+func (o *Orchestrator) findPortInNetwork(net *nomadapi.NetworkResource) (string, int, bool) {
+	for _, port := range net.ReservedPorts {
+		if port.Label == o.portLabel {
+			return net.IP, port.Value, true
+		}
+
+		if port.Value == o.defaultPort {
+			return net.IP, o.defaultPort, true
+		}
+	}
+
+	for _, port := range net.DynamicPorts {
+		if port.Label == o.portLabel {
+			return net.IP, port.Value, true
+		}
+
+		if port.Value == o.defaultPort {
+			return net.IP, o.defaultPort, true
+		}
+	}
+
+	return "", 0, false
+}
+
 func (o *Orchestrator) GetNode(clusterID uuid.UUID, nodeID string) *nodemanager.Node {
 	scopedKey := o.scopedNodeID(clusterID, nodeID)
 	n, _ := o.nodes.Get(scopedKey)