feat: feature flag guard provision version to allow rolling update (#1399)

dobrac · web-flow · commit fd02c957984e · 2025-10-27T10:41:44.000Z
diff --git a/packages/api/internal/middleware/launchdarkly.go b/packages/api/internal/middleware/launchdarkly.go
@@ -46,7 +46,7 @@ func createLaunchDarklyTeamContext(c *gin.Context) (ldcontext.Context, bool) {
 
 	var contexts []ldcontext.Context
 	if team != nil {
-		contexts = append(contexts, featureflags.TeamContext(team.ID.String(), team.Name))
+		contexts = append(contexts, featureflags.TeamContextWithName(team.ID.String(), team.Name))
 		if clusterID := team.ClusterID; clusterID != nil {
 			contexts = append(contexts, featureflags.ClusterContext(clusterID.String()))
 		}
diff --git a/packages/orchestrator/benchmark_test.go b/packages/orchestrator/benchmark_test.go
@@ -232,6 +232,7 @@ func BenchmarkBaseImageLaunch(b *testing.B) {
 
 	builder := build.NewBuilder(
 		logger,
+		featureFlags,
 		sandboxFactory,
 		persistenceTemplate,
 		persistenceBuild,
diff --git a/packages/orchestrator/cmd/build-template/main.go b/packages/orchestrator/cmd/build-template/main.go
@@ -190,6 +190,7 @@ func buildTemplate(
 
 	builder := build.NewBuilder(
 		logger,
+		featureFlags,
 		sandboxFactory,
 		persistenceTemplate,
 		persistenceBuild,
diff --git a/packages/orchestrator/internal/template/build/builder.go b/packages/orchestrator/internal/template/build/builder.go
@@ -32,6 +32,7 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/constants"
 	artifactsregistry "github.com/e2b-dev/infra/packages/shared/pkg/artifacts-registry"
 	"github.com/e2b-dev/infra/packages/shared/pkg/dockerhub"
+	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
@@ -55,10 +56,12 @@ type Builder struct {
 	sandboxes           *sandbox.Map
 	templateCache       *sbxtemplate.Cache
 	metrics             *metrics.BuildMetrics
+	featureFlags        *featureflags.Client
 }
 
 func NewBuilder(
 	logger *zap.Logger,
+	featureFlags *featureflags.Client,
 	sandboxFactory *sandbox.Factory,
 	templateStorage storage.StorageProvider,
 	buildStorage storage.StorageProvider,
@@ -71,6 +74,7 @@ func NewBuilder(
 ) *Builder {
 	return &Builder{
 		logger:              logger,
+		featureFlags:        featureFlags,
 		sandboxFactory:      sandboxFactory,
 		templateStorage:     templateStorage,
 		buildStorage:        buildStorage,
@@ -217,6 +221,7 @@ func runBuild(
 
 	baseBuilder := base.New(
 		bc,
+		builder.featureFlags,
 		builder.logger,
 		builder.proxy,
 		builder.templateStorage,
diff --git a/packages/orchestrator/internal/template/build/phases/base/builder.go b/packages/orchestrator/internal/template/build/phases/base/builder.go
@@ -30,6 +30,7 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg"
 	artifactsregistry "github.com/e2b-dev/infra/packages/shared/pkg/artifacts-registry"
 	"github.com/e2b-dev/infra/packages/shared/pkg/dockerhub"
+	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/id"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
@@ -61,6 +62,7 @@ type BaseBuilder struct {
 	templateStorage     storage.StorageProvider
 	artifactRegistry    artifactsregistry.ArtifactsRegistry
 	dockerhubRepository dockerhub.RemoteRepository
+	featureFlags        *featureflags.Client
 
 	layerExecutor *layer.LayerExecutor
 	index         cache.Index
@@ -69,6 +71,7 @@ type BaseBuilder struct {
 
 func New(
 	buildContext buildcontext.BuildContext,
+	featureFlags *featureflags.Client,
 	logger *zap.Logger,
 	proxy *proxy.SandboxProxy,
 	templateStorage storage.StorageProvider,
@@ -89,6 +92,7 @@ func New(
 		artifactRegistry:    artifactRegistry,
 		dockerhubRepository: dockerhubRepository,
 		sandboxFactory:      sandboxFactory,
+		featureFlags:        featureFlags,
 
 		layerExecutor: layerExecutor,
 		index:         index,
diff --git a/packages/orchestrator/internal/template/build/phases/base/hash.go b/packages/orchestrator/internal/template/build/phases/base/hash.go
@@ -1,14 +1,16 @@
 package base
 
 import (
+	"context"
 	"fmt"
 	"strconv"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/phases"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/storage/cache"
+	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
 )
 
-func (bb *BaseBuilder) Hash(_ phases.LayerResult) (string, error) {
+func (bb *BaseBuilder) Hash(ctx context.Context, _ phases.LayerResult) (string, error) {
 	var baseSource string
 	if bb.Config.FromTemplate != nil {
 		// When building from template, use the base template metadata
@@ -23,9 +25,20 @@ func (bb *BaseBuilder) Hash(_ phases.LayerResult) (string, error) {
 		baseSource = bb.Config.FromImage
 	}
 
+	provisionVersion := provisionScriptFile
+	if val, err := bb.featureFlags.IntFlag(
+		ctx,
+		featureflags.BuildProvisionVersion,
+		featureflags.TemplateContext(bb.Config.TemplateID),
+		featureflags.TeamContext(bb.Config.TeamID),
+	// for dev environments (fallback value), use the provision script hash
+	); val != featureflags.BuildProvisionVersion.Fallback() && err == nil {
+		provisionVersion = strconv.FormatInt(int64(val), 10)
+	}
+
 	return cache.HashKeys(
 		bb.index.Version(),
-		provisionScriptFile,
+		provisionVersion,
 		strconv.FormatInt(bb.Config.DiskSizeMB, 10),
 		baseSource,
 	), nil
diff --git a/packages/orchestrator/internal/template/build/phases/finalize/builder.go b/packages/orchestrator/internal/template/build/phases/finalize/builder.go
@@ -75,7 +75,7 @@ func (ppb *PostProcessingBuilder) Metadata() phases.PhaseMeta {
 	}
 }
 
-func (ppb *PostProcessingBuilder) Hash(sourceLayer phases.LayerResult) (string, error) {
+func (ppb *PostProcessingBuilder) Hash(_ context.Context, sourceLayer phases.LayerResult) (string, error) {
 	return cache.HashKeys(sourceLayer.Hash, "config-run-cmd"), nil
 }
 
diff --git a/packages/orchestrator/internal/template/build/phases/phase.go b/packages/orchestrator/internal/template/build/phases/phase.go
@@ -31,7 +31,7 @@ type BuilderPhase interface {
 	String(ctx context.Context) (string, error)
 	Metadata() PhaseMeta
 
-	Hash(sourceLayer LayerResult) (string, error)
+	Hash(ctx context.Context, sourceLayer LayerResult) (string, error)
 	Layer(ctx context.Context, sourceLayer LayerResult, hash string) (LayerResult, error)
 	Build(ctx context.Context, userLogger *zap.Logger, prefix string, sourceLayer LayerResult, currentLayer LayerResult) (LayerResult, error)
 }
@@ -83,7 +83,7 @@ func Run(
 		)
 
 		phaseStartTime := time.Now()
-		hash, err := builder.Hash(sourceLayer)
+		hash, err := builder.Hash(ctx, sourceLayer)
 		if err != nil {
 			return LayerResult{}, fmt.Errorf("getting hash: %w", err)
 		}
diff --git a/packages/orchestrator/internal/template/build/phases/steps/hash.go b/packages/orchestrator/internal/template/build/phases/steps/hash.go
@@ -1,14 +1,15 @@
 package steps
 
 import (
+	"context"
 	"strings"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/phases"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/storage/cache"
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
-func (sb *StepBuilder) Hash(sourceLayer phases.LayerResult) (string, error) {
+func (sb *StepBuilder) Hash(_ context.Context, sourceLayer phases.LayerResult) (string, error) {
 	return cache.HashKeys(
 		sourceLayer.Hash,
 		sb.step.GetType(),
diff --git a/packages/orchestrator/internal/template/build/phases/user/hash.go b/packages/orchestrator/internal/template/build/phases/user/hash.go
@@ -1,13 +1,15 @@
 package user
 
 import (
+	"context"
+
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/phases"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/storage/cache"
 )
 
 const prefix = "DEFAULT USER"
 
-func (ub *UserBuilder) Hash(sourceLayer phases.LayerResult) (string, error) {
+func (ub *UserBuilder) Hash(_ context.Context, sourceLayer phases.LayerResult) (string, error) {
 	return cache.HashKeys(
 		sourceLayer.Hash,
 		prefix,
diff --git a/packages/orchestrator/internal/template/server/main.go b/packages/orchestrator/internal/template/server/main.go
@@ -20,6 +20,7 @@ import (
 	artifactsregistry "github.com/e2b-dev/infra/packages/shared/pkg/artifacts-registry"
 	"github.com/e2b-dev/infra/packages/shared/pkg/dockerhub"
 	"github.com/e2b-dev/infra/packages/shared/pkg/env"
+	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
 	templatemanager "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
 	"github.com/e2b-dev/infra/packages/shared/pkg/limit"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
@@ -48,6 +49,7 @@ type ServerStore struct {
 
 func New(
 	ctx context.Context,
+	featureFlags *featureflags.Client,
 	meterProvider metric.MeterProvider,
 	logger *zap.Logger,
 	buildLogger *zap.Logger,
@@ -98,6 +100,7 @@ func New(
 
 	builder := build.NewBuilder(
 		logger,
+		featureFlags,
 		sandboxFactory,
 		templatePersistence,
 		buildPersistance,
diff --git a/packages/orchestrator/main.go b/packages/orchestrator/main.go
@@ -405,6 +405,7 @@ func run(config cfg.Config) (success bool) {
 	if slices.Contains(services, cfg.TemplateManager) {
 		tmpl, err = tmplserver.New(
 			ctx,
+			featureFlags,
 			tel.MeterProvider,
 			globalLogger,
 			tmplSbxLoggerExternal,
diff --git a/packages/shared/pkg/feature-flags/context.go b/packages/shared/pkg/feature-flags/context.go
@@ -130,13 +130,19 @@ func SandboxContext(sandboxID string) ldcontext.Context {
 	return ldcontext.NewWithKind(SandboxKind, sandboxID)
 }
 
-func TeamContext(teamID, teamName string) ldcontext.Context {
+func TeamContextWithName(teamID, teamName string) ldcontext.Context {
 	return ldcontext.NewBuilder(teamID).
 		Kind(TeamKind).
 		Name(teamName).
 		Build()
 }
 
+func TeamContext(teamID string) ldcontext.Context {
+	return ldcontext.NewBuilder(teamID).
+		Kind(TeamKind).
+		Build()
+}
+
 func TierContext(tierID, tierName string) ldcontext.Context {
 	return ldcontext.NewBuilder(tierID).
 		Kind(TierKind).
@@ -151,3 +157,7 @@ func UserContext(userID string) ldcontext.Context {
 func ServiceContext(serviceName string) ldcontext.Context {
 	return ldcontext.NewWithKind(ServiceKind, serviceName)
 }
+
+func TemplateContext(templateID string) ldcontext.Context {
+	return ldcontext.NewWithKind(TemplateKind, templateID)
+}
diff --git a/packages/shared/pkg/feature-flags/flags.go b/packages/shared/pkg/feature-flags/flags.go
@@ -23,6 +23,8 @@ const (
 	TierKind ldcontext.Kind = "tier"
 
 	ServiceKind ldcontext.Kind = "service"
+
+	TemplateKind ldcontext.Kind = "template"
 )
 
 // All flags must be defined here: https://app.launchdarkly.com/projects/default/flags/
@@ -93,4 +95,5 @@ var (
 	// BuildCacheMaxUsagePercentage the maximum percentage of the cache disk storage
 	// that can be used before the cache starts evicting items.
 	BuildCacheMaxUsagePercentage = newIntFlag("build-cache-max-usage-percentage", 85)
+	BuildProvisionVersion        = newIntFlag("build-provision-version", 0)
 )

Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,7 @@ func createLaunchDarklyTeamContext(c *gin.Context) (ldcontext.Context, bool) {`
`46`	`46`
`47`	`47`	`var contexts []ldcontext.Context`
`48`	`48`	`if team != nil {`
`49`		`- contexts = append(contexts, featureflags.TeamContext(team.ID.String(), team.Name))`
	`49`	`+ contexts = append(contexts, featureflags.TeamContextWithName(team.ID.String(), team.Name))`
`50`	`50`	`if clusterID := team.ClusterID; clusterID != nil {`
`51`	`51`	`contexts = append(contexts, featureflags.ClusterContext(clusterID.String()))`
`52`	`52`	`}`
Original file line number	Diff line number	Diff line change
`@@ -75,7 +75,7 @@ func (ppb *PostProcessingBuilder) Metadata() phases.PhaseMeta {`
`75`	`75`	`}`
`76`	`76`	`}`
`77`	`77`
`78`		`-func (ppb *PostProcessingBuilder) Hash(sourceLayer phases.LayerResult) (string, error) {`
	`78`	`+func (ppb *PostProcessingBuilder) Hash(_ context.Context, sourceLayer phases.LayerResult) (string, error) {`
`79`	`79`	`return cache.HashKeys(sourceLayer.Hash, "config-run-cmd"), nil`
`80`	`80`	`}`
`81`	`81`
Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,7 @@ type BuilderPhase interface {`
`31`	`31`	`String(ctx context.Context) (string, error)`
`32`	`32`	`Metadata() PhaseMeta`
`33`	`33`
`34`		`- Hash(sourceLayer LayerResult) (string, error)`
	`34`	`+ Hash(ctx context.Context, sourceLayer LayerResult) (string, error)`
`35`	`35`	`Layer(ctx context.Context, sourceLayer LayerResult, hash string) (LayerResult, error)`
`36`	`36`	`Build(ctx context.Context, userLogger *zap.Logger, prefix string, sourceLayer LayerResult, currentLayer LayerResult) (LayerResult, error)`
`37`	`37`	`}`
`@@ -83,7 +83,7 @@ func Run(`
`83`	`83`	`)`
`84`	`84`
`85`	`85`	`phaseStartTime := time.Now()`
`86`		`- hash, err := builder.Hash(sourceLayer)`
	`86`	`+ hash, err := builder.Hash(ctx, sourceLayer)`
`87`	`87`	`if err != nil {`
`88`	`88`	`return LayerResult{}, fmt.Errorf("getting hash: %w", err)`
`89`	`89`	`}`