diff --git a/packages/api/internal/cfg/model.go b/packages/api/internal/cfg/model.go index dae172b46d..8d3fee86e5 100644 --- a/packages/api/internal/cfg/model.go +++ b/packages/api/internal/cfg/model.go @@ -112,6 +112,13 @@ type Config struct { DefaultPersistentVolumeType string `env:"DEFAULT_PERSISTENT_VOLUME_TYPE"` + // DefaultPersistentVolumeTypeByRegion is the per-region default volume type, + // e.g. "us-west3:zonalfilestore-us-west3". A team's region is resolved from + // the region= node labels of the nodes its scheduling labels select; Terraform + // derives this map from the volume types themselves and fails the plan when a + // region with several types lacks an explicit default. + DefaultPersistentVolumeTypeByRegion map[string]string `env:"DEFAULT_PERSISTENT_VOLUME_TYPE_BY_REGION"` + DomainName string `env:"DOMAIN_NAME" envDefault:""` } diff --git a/packages/api/internal/cfg/model_test.go b/packages/api/internal/cfg/model_test.go index f9d0821941..dad4c0fdeb 100644 --- a/packages/api/internal/cfg/model_test.go +++ b/packages/api/internal/cfg/model_test.go @@ -43,6 +43,17 @@ func TestParse(t *testing.T) { assert.Equal(t, content, result.VolumesToken.SigningKey) }) + t.Run("default persistent volume type by region is parsed as a map", func(t *testing.T) { + t.Setenv("DEFAULT_PERSISTENT_VOLUME_TYPE_BY_REGION", "us-west3:zonalfilestore-us-west3,other:other-type") + + result, err := Parse() + require.NoError(t, err) + assert.Equal(t, map[string]string{ + "us-west3": "zonalfilestore-us-west3", + "other": "other-type", + }, result.DefaultPersistentVolumeTypeByRegion) + }) + t.Run("invalid service discovery provider exposes failure condition", func(t *testing.T) { t.Setenv("SERVICE_DISCOVERY_PROVIDER", "invalid") diff --git a/packages/api/internal/handlers/volume_create.go b/packages/api/internal/handlers/volume_create.go index 5b3f3087b7..a156de8726 100644 --- a/packages/api/internal/handlers/volume_create.go +++ b/packages/api/internal/handlers/volume_create.go @@ -6,12 +6,14 @@ import ( "fmt" "net/http" "regexp" + "strings" "github.com/gin-gonic/gin" "github.com/google/uuid" "github.com/e2b-dev/infra/packages/api/internal/api" "github.com/e2b-dev/infra/packages/api/internal/clusters" + "github.com/e2b-dev/infra/packages/auth/pkg/types" "github.com/e2b-dev/infra/packages/db/pkg/dberrors" "github.com/e2b-dev/infra/packages/db/queries" clustershared "github.com/e2b-dev/infra/packages/shared/pkg/clusters" @@ -72,9 +74,12 @@ func (a *APIStore) PostVolumes(c *gin.Context) { return } - ctx = featureflags.AddToContext(ctx, featureflags.VolumeContext(body.Name)) + ctx = featureflags.AddToContext(ctx, + featureflags.VolumeContext(body.Name), + featureflags.TeamContext(team.ID.String()), + ) - volumeType := a.getVolumeType(ctx) + volumeType := a.getVolumeType(ctx, team) if volumeType == "" { a.sendAPIStoreError(c, http.StatusInternalServerError, "No persistent volume type is configured") telemetry.ReportCriticalError(ctx, "default persistent volume type is not configured", nil) @@ -174,13 +179,86 @@ func (a *APIStore) PostVolumes(c *gin.Context) { c.JSON(http.StatusCreated, result) } -func (a *APIStore) getVolumeType(ctx context.Context) string { - volumeType := a.featureFlags.StringFlag(ctx, featureflags.DefaultPersistentVolumeType) - if volumeType == "" { - volumeType = a.config.DefaultPersistentVolumeType +const ( + // regionNodeLabelPrefix marks the node label naming the region a node runs + // in, e.g. "region=us-west3". Regions live on nodes, never on teams: + // Terraform appends the label to every client cluster automatically. + regionNodeLabelPrefix = "region=" + + // defaultSchedulingLabel is the pool a team without scheduling labels of + // its own lands on. + defaultSchedulingLabel = "default" +) + +// getVolumeType resolves the volume type a new volume of the given team gets, +// in order of precedence: the LaunchDarkly override, the per-region default of +// the region the team schedules into, and finally the deployment-wide default. +// +// Node labels only answer *where* the team runs; *what* a new volume there +// should be is policy and comes exclusively from the region map. A region +// mounting several volume types therefore never needs runtime guessing - the +// map names its default explicitly. +func (a *APIStore) getVolumeType(ctx context.Context, team *types.Team) string { + if volumeType := a.featureFlags.StringFlag(ctx, featureflags.DefaultPersistentVolumeType); volumeType != "" { + return volumeType + } + + // Regional defaulting is opt-in: without a map there is nothing to look + // up, so don't walk the cluster. + if len(a.config.DefaultPersistentVolumeTypeByRegion) == 0 || a.orchestrator == nil { + return a.config.DefaultPersistentVolumeType + } + + // Mirrors generateRequiredNodeLabels: a team without labels of its own + // runs on the "default" pool, so we resolve the region over the same set + // of nodes that placement would choose from. + requiredLabels := team.SandboxSchedulingLabels + if len(requiredLabels) == 0 { + requiredLabels = []string{defaultSchedulingLabel} + } + + // Collect the distinct regions advertised by the ready nodes carrying all + // of requiredLabels; the same subset semantics as sandbox placement. + regions := make(map[string]struct{}) + for _, node := range a.orchestrator.GetClusterNodes(clustershared.WithClusterFallback(team.ClusterID)) { + if node.Status() != api.NodeStatusReady { + continue + } + + labels := node.Labels() + if !hasAllLabels(labels, requiredLabels) { + continue + } + + for label := range labels { + if region, ok := strings.CutPrefix(label, regionNodeLabelPrefix); ok && region != "" { + regions[region] = struct{}{} + } + } + } + + // Exactly one region with a mapped default is the only affirmative + // answer. Zero regions (labels match nothing yet) or several (labels do + // not pin a region) fall open to the deployment-wide default. + if len(regions) == 1 { + for region := range regions { + if volumeType, ok := a.config.DefaultPersistentVolumeTypeByRegion[region]; ok { + return volumeType + } + } + } + + return a.config.DefaultPersistentVolumeType +} + +func hasAllLabels(labels map[string]struct{}, requiredLabels []string) bool { + for _, required := range requiredLabels { + if _, ok := labels[required]; !ok { + return false + } } - return volumeType + return true } var validVolumeNameRegex = regexp.MustCompile(`^[a-zA-Z0-9_-]+$`) diff --git a/packages/api/internal/handlers/volume_test.go b/packages/api/internal/handlers/volume_test.go index 95d2e76ec2..2fc87f673e 100644 --- a/packages/api/internal/handlers/volume_test.go +++ b/packages/api/internal/handlers/volume_test.go @@ -3,7 +3,16 @@ package handlers import ( "testing" + "github.com/google/uuid" + "github.com/launchdarkly/go-sdk-common/v3/ldvalue" + "github.com/launchdarkly/go-server-sdk/v7/testhelpers/ldtestdata" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/api/internal/cfg" + "github.com/e2b-dev/infra/packages/auth/pkg/types" + authqueries "github.com/e2b-dev/infra/packages/db/pkg/auth/queries" + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" ) func TestIsValidVolumeName(t *testing.T) { @@ -135,3 +144,88 @@ func TestIsValidVolumeName(t *testing.T) { }) } } + +// TestGetVolumeType covers the precedence paths that don't need a cluster; +// APIStore.orchestrator is a concrete *orchestrator.Orchestrator with no test +// constructor, so the node-derived resolution is not covered here. +func TestGetVolumeType(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + // flagValue is the LaunchDarkly default-persistent-volume-type + // override; empty leaves the flag unset. + flagValue string + regionVolumeType map[string]string + defaultType string + expected string + }{ + { + name: "no region map falls back to global default", + defaultType: "global-type", + expected: "global-type", + }, + { + // The orchestrator is only nil in tests, but a missing cluster + // view must not be what decides whether a volume can be created. + name: "no orchestrator falls back to global default", + regionVolumeType: map[string]string{"us-west3": "zonalfilestore-us-west3"}, + defaultType: "global-type", + expected: "global-type", + }, + { + name: "feature flag wins over region map and global default", + flagValue: "flag-type", + regionVolumeType: map[string]string{"us-west3": "zonalfilestore-us-west3"}, + defaultType: "global-type", + expected: "flag-type", + }, + { + name: "no default configured at all", + expected: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // Each subtest gets its own datasource/client so parallel runs + // don't race on the shared flag value. + td := ldtestdata.DataSource() + ff, err := featureflags.NewClientWithDatasource(td) + require.NoError(t, err) + t.Cleanup(func() { + assert.NoError(t, ff.Close(t.Context())) + }) + + if tt.flagValue != "" { + td.Update(td.Flag(featureflags.DefaultPersistentVolumeType.Key()). + ValueForAll(ldvalue.String(tt.flagValue))) + } + + store := &APIStore{ + featureFlags: ff, + config: cfg.Config{ + DefaultPersistentVolumeType: tt.defaultType, + DefaultPersistentVolumeTypeByRegion: tt.regionVolumeType, + }, + } + team := &types.Team{Team: &authqueries.Team{ID: uuid.New()}} + + assert.Equal(t, tt.expected, store.getVolumeType(t.Context(), team)) + }) + } +} + +func TestHasAllLabels(t *testing.T) { + t.Parallel() + + labels := map[string]struct{}{"gpu": {}, "highmem": {}, "region=us-west3": {}} + + assert.True(t, hasAllLabels(labels, nil)) + assert.True(t, hasAllLabels(labels, []string{"gpu"})) + assert.True(t, hasAllLabels(labels, []string{"gpu", "highmem"})) + assert.False(t, hasAllLabels(labels, []string{"gpu", "default"})) + assert.False(t, hasAllLabels(map[string]struct{}{}, []string{"default"})) +}