feat: toggles in agentbakersvc (#4206)

Co-authored-by: Cameron Meissner <[email protected]>
Azure · Apr 5, 2024 · d70ff99 · d70ff99
1 parent cafd01c
commit d70ff99
Show file tree

Hide file tree

Showing 11 changed files with 445 additions and 51 deletions.
diff --git a/apiserver/apiserver.go b/apiserver/apiserver.go
@@ -6,15 +6,21 @@ import (
 	"log"
 	"net/http"
 	"time"
+
+	"github.com/Azure/agentbaker/pkg/agent/toggles"
 )
 
 const (
 	readHeaderTimeoutSeconds = 5
 )
 
+// OptionConfigurator is a function which can configure an Options object.
+type OptionConfigurator func(opts *Options)
+
 // Options holds the options for the api server.
 type Options struct {
-	Addr string
+	Addr    string
+	Toggles *toggles.Toggles
 }
 
 func (o *Options) validate() error {

diff --git a/apiserver/getdistrosigimageconfig.go b/apiserver/getdistrosigimageconfig.go
@@ -33,7 +33,15 @@ func (api *APIServer) GetDistroSigImageConfig(w http.ResponseWriter, r *http.Req
 		return
 	}
 
-	allDistros, err := agentBaker.GetDistroSigImageConfig(config.SIGConfig, config.Region)
+	if api.Options != nil && api.Options.Toggles != nil {
+		agentBaker = agentBaker.WithToggles(api.Options.Toggles)
+	}
+
+	allDistros, err := agentBaker.GetDistroSigImageConfig(config.SIGConfig, &datamodel.EnvironmentInfo{
+		SubscriptionID: config.SubscriptionID,
+		TenantID:       config.TenantID,
+		Region:         config.Region,
+	})
 	if err != nil {
 		log.Println(err.Error())
 		http.Error(w, err.Error(), http.StatusBadRequest)

diff --git a/apiserver/getlatestsigimageconfig.go b/apiserver/getlatestsigimageconfig.go
@@ -33,7 +33,15 @@ func (api *APIServer) GetLatestSigImageConfig(w http.ResponseWriter, r *http.Req
 		return
 	}
 
-	latestSigConfig, err := agentBaker.GetLatestSigImageConfig(config.SIGConfig, config.Region, config.Distro)
+	if api.Options != nil && api.Options.Toggles != nil {
+		agentBaker = agentBaker.WithToggles(api.Options.Toggles)
+	}
+
+	latestSigConfig, err := agentBaker.GetLatestSigImageConfig(config.SIGConfig, config.Distro, &datamodel.EnvironmentInfo{
+		SubscriptionID: config.SubscriptionID,
+		TenantID:       config.TenantID,
+		Region:         config.Region,
+	})
 	if err != nil {
 		log.Println(err.Error())
 		http.Error(w, err.Error(), http.StatusBadRequest)

diff --git a/apiserver/getnodebootstrapdata.go b/apiserver/getnodebootstrapdata.go
@@ -39,12 +39,18 @@ func (api *APIServer) GetNodeBootstrapData(w http.ResponseWriter, r *http.Reques
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
+
+	if api.Options != nil && api.Options.Toggles != nil {
+		agentBaker = agentBaker.WithToggles(api.Options.Toggles)
+	}
+
 	nodeBootStrapping, err := agentBaker.GetNodeBootstrapping(ctx, &config)
 	if err != nil {
 		log.Println(err.Error())
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
+
 	result, err := json.Marshal(nodeBootStrapping)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusBadRequest)

diff --git a/cmd/starter/start.go b/cmd/starter/start.go
@@ -13,10 +13,14 @@ import (
 )
 
 // Execute adds all child commands to the root command and sets flags appropriately.
-func Execute() {
+func Execute(configurators ...apiserver.OptionConfigurator) {
 	rootCmd.AddCommand(startCmd)
 	startCmd.Flags().StringVar(&options.Addr, "addr", ":8080", "the addr to serve the api on")
 
+	for _, configurator := range configurators {
+		configurator(options)
+	}
+
 	if err := rootCmd.Execute(); err != nil {
 		log.Println(err)
 		os.Exit(1)

diff --git a/pkg/agent/bakerapi.go b/pkg/agent/bakerapi.go
@@ -8,24 +8,36 @@ import (
 	"fmt"
 
 	"github.com/Azure/agentbaker/pkg/agent/datamodel"
+	"github.com/Azure/agentbaker/pkg/agent/toggles"
 )
 
 //nolint:revive // Name does not need to be modified to baker
 type AgentBaker interface {
 	GetNodeBootstrapping(ctx context.Context, config *datamodel.NodeBootstrappingConfiguration) (*datamodel.NodeBootstrapping, error)
-	GetLatestSigImageConfig(sigConfig datamodel.SIGConfig, region string, distro datamodel.Distro) (*datamodel.SigImageConfig, error)
-	GetDistroSigImageConfig(sigConfig datamodel.SIGConfig, region string) (map[datamodel.Distro]datamodel.SigImageConfig, error)
+	GetLatestSigImageConfig(sigConfig datamodel.SIGConfig, distro datamodel.Distro, envInfo *datamodel.EnvironmentInfo) (*datamodel.SigImageConfig, error)
+	GetDistroSigImageConfig(sigConfig datamodel.SIGConfig, envInfo *datamodel.EnvironmentInfo) (map[datamodel.Distro]datamodel.SigImageConfig, error)
 }
 
-func NewAgentBaker() (AgentBaker, error) {
-	return &agentBakerImpl{}, nil
+type agentBakerImpl struct {
+	toggles *toggles.Toggles
 }
 
-type agentBakerImpl struct{}
+var _ AgentBaker = (*agentBakerImpl)(nil)
+
+//nolint:revive // fine to return unexported type due to interface usage
+func NewAgentBaker() (*agentBakerImpl, error) {
+	return &agentBakerImpl{
+		toggles: toggles.New(),
+	}, nil
+}
+
+func (agentBaker *agentBakerImpl) WithToggles(toggles *toggles.Toggles) *agentBakerImpl {
+	agentBaker.toggles = toggles
+	return agentBaker
+}
 
 //nolint:revive, nolintlint // ctx is not used, but may be in the future
-func (agentBaker *agentBakerImpl) GetNodeBootstrapping(ctx context.Context,
-	config *datamodel.NodeBootstrappingConfiguration) (*datamodel.NodeBootstrapping, error) {
+func (agentBaker *agentBakerImpl) GetNodeBootstrapping(ctx context.Context, config *datamodel.NodeBootstrappingConfiguration) (*datamodel.NodeBootstrapping, error) {
 	// validate and fix input before passing config to the template generator.
 	if config.AgentPoolProfile.IsWindows() {
 		validateAndSetWindowsNodeBootstrappingConfiguration(config)
@@ -63,70 +75,102 @@ func (agentBaker *agentBakerImpl) GetNodeBootstrapping(ctx context.Context,
 		return nil, fmt.Errorf("can't find image for distro %s", distro)
 	}
 
-	return nodeBootstrapping, nil
-}
-
-func findSIGImageConfig(sigConfig datamodel.SIGAzureEnvironmentSpecConfig, distro datamodel.Distro) *datamodel.SigImageConfig {
-	if imageConfig, ok := sigConfig.SigUbuntuImageConfig[distro]; ok {
-		return &imageConfig
-	}
-	if imageConfig, ok := sigConfig.SigCBLMarinerImageConfig[distro]; ok {
-		return &imageConfig
-	}
-	if imageConfig, ok := sigConfig.SigAzureLinuxImageConfig[distro]; ok {
-		return &imageConfig
-	}
-	if imageConfig, ok := sigConfig.SigWindowsImageConfig[distro]; ok {
-		return &imageConfig
-	}
-	if imageConfig, ok := sigConfig.SigUbuntuEdgeZoneImageConfig[distro]; ok {
-		return &imageConfig
+	if !config.AgentPoolProfile.IsWindows() {
+		// handle node image version toggle/override
+		e := toggles.NewEntityFromNodeBootstrappingConfiguration(config)
+		imageVersionOverrides := agentBaker.toggles.GetLinuxNodeImageVersion(e)
+		if imageVersion, ok := imageVersionOverrides[string(distro)]; ok {
+			nodeBootstrapping.SigImageConfig.Version = imageVersion
+		}
 	}
 
-	return nil
+	return nodeBootstrapping, nil
 }
 
-func (agentBaker *agentBakerImpl) GetLatestSigImageConfig(
-	sigConfig datamodel.SIGConfig, region string, distro datamodel.Distro) (*datamodel.SigImageConfig, error) {
-	sigAzureEnvironmentSpecConfig, err := datamodel.GetSIGAzureCloudSpecConfig(sigConfig, region)
+func (agentBaker *agentBakerImpl) GetLatestSigImageConfig(sigConfig datamodel.SIGConfig,
+	distro datamodel.Distro, envInfo *datamodel.EnvironmentInfo) (*datamodel.SigImageConfig, error) {
+	sigAzureEnvironmentSpecConfig, err := datamodel.GetSIGAzureCloudSpecConfig(sigConfig, envInfo.Region)
 	if err != nil {
 		return nil, err
 	}
 
 	sigImageConfig := findSIGImageConfig(sigAzureEnvironmentSpecConfig, distro)
 	if sigImageConfig == nil {
-		return nil, fmt.Errorf("can't find SIG image config for distro %s in region %s", distro, region)
+		return nil, fmt.Errorf("can't find SIG image config for distro %s in region %s", distro, envInfo.Region)
+	}
+
+	if !distro.IsWindowsDistro() {
+		e := toggles.NewEntityFromEnvironmentInfo(envInfo)
+		imageVersionOverrides := agentBaker.toggles.GetLinuxNodeImageVersion(e)
+		if imageVersion, ok := imageVersionOverrides[string(distro)]; ok {
+			sigImageConfig.Version = imageVersion
+		}
 	}
 	return sigImageConfig, nil
 }
 
 func (agentBaker *agentBakerImpl) GetDistroSigImageConfig(
-	sigConfig datamodel.SIGConfig, region string) (map[datamodel.Distro]datamodel.SigImageConfig, error) {
-	allAzureSigConfig, err := datamodel.GetSIGAzureCloudSpecConfig(sigConfig, region)
+	sigConfig datamodel.SIGConfig, envInfo *datamodel.EnvironmentInfo) (map[datamodel.Distro]datamodel.SigImageConfig, error) {
+	allAzureSigConfig, err := datamodel.GetSIGAzureCloudSpecConfig(sigConfig, envInfo.Region)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get sig image config: %w", err)
 	}
 
+	e := toggles.NewEntityFromEnvironmentInfo(envInfo)
+	linuxImageVersionOverrides := agentBaker.toggles.GetLinuxNodeImageVersion(e)
+
 	allDistros := map[datamodel.Distro]datamodel.SigImageConfig{}
 	for distro, sigConfig := range allAzureSigConfig.SigWindowsImageConfig {
 		allDistros[distro] = sigConfig
 	}
 
 	for distro, sigConfig := range allAzureSigConfig.SigCBLMarinerImageConfig {
+		if version, ok := linuxImageVersionOverrides[string(distro)]; ok {
+			sigConfig.Version = version
+		}
 		allDistros[distro] = sigConfig
 	}
 
 	for distro, sigConfig := range allAzureSigConfig.SigAzureLinuxImageConfig {
+		if version, ok := linuxImageVersionOverrides[string(distro)]; ok {
+			sigConfig.Version = version
+		}
 		allDistros[distro] = sigConfig
 	}
 
 	for distro, sigConfig := range allAzureSigConfig.SigUbuntuImageConfig {
+		if version, ok := linuxImageVersionOverrides[string(distro)]; ok {
+			sigConfig.Version = version
+		}
 		allDistros[distro] = sigConfig
 	}
 
 	for distro, sigConfig := range allAzureSigConfig.SigUbuntuEdgeZoneImageConfig {
+		if version, ok := linuxImageVersionOverrides[string(distro)]; ok {
+			sigConfig.Version = version
+		}
 		allDistros[distro] = sigConfig
 	}
 
 	return allDistros, nil
 }
+
+func findSIGImageConfig(sigConfig datamodel.SIGAzureEnvironmentSpecConfig, distro datamodel.Distro) *datamodel.SigImageConfig {
+	if imageConfig, ok := sigConfig.SigUbuntuImageConfig[distro]; ok {
+		return &imageConfig
+	}
+	if imageConfig, ok := sigConfig.SigCBLMarinerImageConfig[distro]; ok {
+		return &imageConfig
+	}
+	if imageConfig, ok := sigConfig.SigAzureLinuxImageConfig[distro]; ok {
+		return &imageConfig
+	}
+	if imageConfig, ok := sigConfig.SigWindowsImageConfig[distro]; ok {
+		return &imageConfig
+	}
+	if imageConfig, ok := sigConfig.SigUbuntuEdgeZoneImageConfig[distro]; ok {
+		return &imageConfig
+	}
+
+	return nil
+}