Move orchestrator config into a model (#1295)

Co-authored-by: Jakub Novák <[email protected]>

Author: djeebus

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

@@ -70,10 +70,12 @@ EOT
         TEMPLATE_BUCKET_NAME         = "${template_bucket_name}"
         OTEL_COLLECTOR_GRPC_ENDPOINT = "${otel_collector_grpc_endpoint}"
         ALLOW_SANDBOX_INTERNET       = "${allow_sandbox_internet}"
-        SHARED_CHUNK_CACHE_PATH    = "${shared_chunk_cache_path}"
+        SHARED_CHUNK_CACHE_PATH      = "${shared_chunk_cache_path}"
         CLICKHOUSE_CONNECTION_STRING = "${clickhouse_connection_string}"
         REDIS_URL                    = "${redis_url}"
         REDIS_CLUSTER_URL            = "${redis_cluster_url}"
+        GRPC_PORT                    = "${port}"
+        PROXY_PORT                   = "${proxy_port}"
 
 %{ if launch_darkly_api_key != "" }
         LAUNCH_DARKLY_API_KEY         = "${launch_darkly_api_key}"
@@ -82,7 +84,7 @@ EOT
 
       config {
         command = "/bin/bash"
-        args    = ["-c", " chmod +x local/orchestrator && local/orchestrator --port ${port} --proxy-port ${proxy_port}"]
+        args    = ["-c", " chmod +x local/orchestrator && local/orchestrator"]
       }
 
       artifact {

iac/provider-gcp/nomad/jobs/template-manager.hcl

@@ -77,6 +77,7 @@ job "template-manager-system" {
         SHARED_CHUNK_CACHE_PATH       = "${shared_chunk_cache_path}"
         CLICKHOUSE_CONNECTION_STRING  = "${clickhouse_connection_string}"
         DOCKERHUB_REMOTE_REPOSITORY_URL  = "${dockerhub_remote_repository_url}"
+        GRPC_PORT                     = "${port}"
 %{ if !update_stanza }
         FORCE_STOP                    = "true"
 %{ endif }
@@ -87,7 +88,7 @@ job "template-manager-system" {
 
       config {
         command = "/bin/bash"
-        args    = ["-c", " chmod +x local/template-manager && local/template-manager --port ${port}"]
+        args    = ["-c", " chmod +x local/template-manager && local/template-manager"]
       }
 
       artifact {

packages/client-proxy/internal/proxy/proxy.go

@@ -96,7 +96,13 @@ func catalogResolution(ctx context.Context, sandboxId string, c catalog.Sandboxe
 	return s.OrchestratorIP, nil
 }
 
-func NewClientProxy(meterProvider metric.MeterProvider, serviceName string, port uint, catalog catalog.SandboxesCatalog, useCatalogResolution bool, useDnsResolution bool) (*reverseproxy.Proxy, error) {
+func NewClientProxy(
+	meterProvider metric.MeterProvider,
+	serviceName string,
+	port uint16,
+	catalog catalog.SandboxesCatalog,
+	useCatalogResolution, useDnsResolution bool,
+) (*reverseproxy.Proxy, error) {
 	if !useCatalogResolution && !useDnsResolution {
 		return nil, errors.New("catalog resolution and DNS resolution are both disabled, at least one must be enabled")
 	}

packages/client-proxy/main.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"math"
 	"net"
 	"net/http"
 	"os"
@@ -104,6 +105,11 @@ func run() int {
 	zap.ReplaceGlobals(logger)
 
 	proxyPort := internal.GetProxyServicePort()
+	if proxyPort <= 0 || proxyPort > int(math.MaxUint16) {
+		logger.Error("Proxy port is outside the valid uint16 range", zap.Int("value", proxyPort))
+		return 1
+	}
+
 	edgePort := internal.GetEdgeServicePort()
 	edgeSecret := internal.GetEdgeServiceSecret()
 	orchestratorPort := internal.GetOrchestratorServicePort()
@@ -155,7 +161,7 @@ func run() int {
 	}
 
 	// Proxy sandbox http traffic to orchestrator nodes
-	trafficProxy, err := e2bproxy.NewClientProxy(tel.MeterProvider, serviceName, uint(proxyPort), catalog, useProxyCatalogResolution, useDnsResolution)
+	trafficProxy, err := e2bproxy.NewClientProxy(tel.MeterProvider, serviceName, uint16(proxyPort), catalog, useProxyCatalogResolution, useDnsResolution)
 	if err != nil {
 		logger.Error("Failed to create client proxy", zap.Error(err))
 		return 1

packages/orchestrator/benchmark_test.go

@@ -109,6 +109,11 @@ func BenchmarkBaseImageLaunch(b *testing.B) {
 	b.Setenv("SNAPSHOT_CACHE_DIR", abs(filepath.Join(tempDir, "snapshot-cache")))
 	b.Setenv("LOCAL_TEMPLATE_STORAGE_BASE_PATH", abs(filepath.Join(persistenceDir, "templates")))
 
+	networkConfig, err := network.ParseConfig()
+	if err != nil {
+		b.Fatalf("error parsing config: %v", err)
+	}
+
 	// prep directories
 	for _, subdir := range []string{"build", "build-templates" /*"fc-vm",*/, "sandbox", "snapshot-cache", "template"} {
 		fullDirName := filepath.Join(tempDir, subdir)
@@ -123,7 +128,7 @@ func BenchmarkBaseImageLaunch(b *testing.B) {
 	// sbxlogger.SetSandboxLoggerExternal(logger)
 
 	networkPool, err := network.NewPool(
-		b.Context(), noop.MeterProvider{}, 8, 8, clientID,
+		b.Context(), noop.MeterProvider{}, 8, 8, clientID, networkConfig,
 	)
 	require.NoError(b, err)
 	defer func() {
@@ -197,7 +202,7 @@ func BenchmarkBaseImageLaunch(b *testing.B) {
 	persistenceBuild, err := storage.GetBuildCacheStorageProvider(b.Context(), nil)
 	require.NoError(b, err)
 
-	var proxyPort uint = 5007
+	var proxyPort uint16 = 5007
 
 	sandboxes := smap.New[*sandbox.Sandbox]()
 

packages/orchestrator/cmd/build-template/main.go

@@ -38,18 +38,22 @@ const (
 )
 
 func main() {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
+	ctx := context.Background()
 
 	templateID := flag.String("template", "", "template id")
 	buildID := flag.String("build", "", "build id")
 	kernelVersion := flag.String("kernel", "", "kernel version")
 	fcVersion := flag.String("firecracker", "", "firecracker version")
 	flag.Parse()
 
-	err := buildTemplate(ctx, *kernelVersion, *fcVersion, *templateID, *buildID)
+	networkConfig, err := network.ParseConfig()
+	if err != nil {
+		log.Fatalf("error parsing config: %v", err)
+	}
+
+	err = buildTemplate(ctx, *kernelVersion, *fcVersion, *templateID, *buildID, networkConfig)
 	if err != nil {
-		log.Fatalf("error building template: %v", err) //nolint:gocritic // probably fine to bail if we're done?
+		log.Fatalf("error building template: %v", err)
 	}
 }
 
@@ -59,6 +63,7 @@ func buildTemplate(
 	fcVersion,
 	templateID,
 	buildID string,
+	networkConfig network.Config,
 ) error {
 	ctx, cancel := context.WithTimeout(parentCtx, time.Minute*5)
 	defer cancel()
@@ -121,7 +126,7 @@ func buildTemplate(
 		}
 	}()
 
-	networkPool, err := network.NewPool(ctx, noop.MeterProvider{}, 8, 8, clientID)
+	networkPool, err := network.NewPool(ctx, noop.MeterProvider{}, 8, 8, clientID, networkConfig)
 	if err != nil {
 		return fmt.Errorf("could not create network pool: %w", err)
 	}

packages/orchestrator/go.mod

@@ -18,6 +18,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/ecr v1.44.0
 	github.com/bits-and-blooms/bitset v1.22.0
 	github.com/bmatcuk/doublestar/v4 v4.9.1
+	github.com/caarlos0/env/v11 v11.3.1
 	github.com/containernetworking/plugins v1.6.0
 	github.com/containers/storage v1.58.0
 	github.com/coreos/go-iptables v0.8.0

packages/orchestrator/go.sum

@@ -0,0 +1,29 @@
+package cfg
+
+import (
+	"github.com/caarlos0/env/v11"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/network"
+)
+
+type Config struct {
+	AllowSandboxInternet       bool     `env:"ALLOW_SANDBOX_INTERNET"       envDefault:"true"`
+	ClickhouseConnectionString string   `env:"CLICKHOUSE_CONNECTION_STRING"`
+	ForceStop                  bool     `env:"FORCE_STOP"`
+	GRPCPort                   uint16   `env:"GRPC_PORT"                    envDefault:"5008"`
+	LaunchDarklyAPIKey         string   `env:"LAUNCH_DARKLY_API_KEY"`
+	OrchestratorBasePath       string   `env:"ORCHESTRATOR_BASE_PATH"       envDefault:"/orchestrator"`
+	OrchestratorLockPath       string   `env:"ORCHESTRATOR_LOCK_PATH"       envDefault:"/orchestrator.lock"`
+	ProxyPort                  uint16   `env:"PROXY_PORT"                   envDefault:"5007"`
+	RedisClusterURL            string   `env:"REDIS_CLUSTER_URL"`
+	RedisURL                   string   `env:"REDIS_URL"`
+	Services                   []string `env:"ORCHESTRATOR_SERVICES"        envDefault:"orchestrator"`
+
+	NetworkConfig network.Config
+}
+
+func Parse() (Config, error) {
+	var model Config
+	err := env.Parse(&model)
+	return model, err
+}

packages/orchestrator/internal/cfg/model.go

@@ -0,0 +1,35 @@
+package cfg
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestParse(t *testing.T) {
+	t.Run("network config local flag defaults to false", func(t *testing.T) {
+		config, err := Parse()
+		require.NoError(t, err)
+
+		assert.False(t, config.NetworkConfig.UseLocalNamespaceStorage)
+	})
+
+	t.Run("network config is parsed correctly", func(t *testing.T) {
+		t.Setenv("USE_LOCAL_NAMESPACE_STORAGE", "true")
+
+		config, err := Parse()
+		require.NoError(t, err)
+
+		assert.True(t, config.NetworkConfig.UseLocalNamespaceStorage)
+	})
+
+	t.Run("multiple services parses correctly", func(t *testing.T) {
+		t.Setenv("ORCHESTRATOR_SERVICES", "service1,service2")
+
+		config, err := Parse()
+		require.NoError(t, err)
+
+		assert.Equal(t, []string{"service1", "service2"}, config.Services)
+	})
+}