Merge pull request #125052 from p0lyn0mial/upstream-client-go-env-var-set

client-go/features: add Set method to the envvar impl

Kubernetes-commit: 8991e825d5f2affce16a05666c11d7a2a773a628

features/envvar.go

@@ -47,6 +47,10 @@ var _ Gates = &envVarFeatureGates{}
 //
 // Please note that environmental variables can only be set to the boolean value.
 // Incorrect values will be ignored and logged.
+//
+// Features can also be set directly via the Set method.
+// In that case, these features take precedence over
+// features set via environmental variables.
 func newEnvVarFeatureGates(features map[Feature]FeatureSpec) *envVarFeatureGates {
 	known := map[Feature]FeatureSpec{}
 	for name, spec := range features {
@@ -57,7 +61,8 @@ func newEnvVarFeatureGates(features map[Feature]FeatureSpec) *envVarFeatureGates
 		callSiteName: naming.GetNameFromCallsite(internalPackages...),
 		known:        known,
 	}
-	fg.enabled.Store(map[Feature]bool{})
+	fg.enabledViaEnvVar.Store(map[Feature]bool{})
+	fg.enabledViaSetMethod = map[Feature]bool{}
 
 	return fg
 }
@@ -74,17 +79,34 @@ type envVarFeatureGates struct {
 	// known holds known feature gates
 	known map[Feature]FeatureSpec
 
-	// enabled holds a map[Feature]bool
+	// enabledViaEnvVar holds a map[Feature]bool
 	// with values explicitly set via env var
-	enabled atomic.Value
+	enabledViaEnvVar atomic.Value
+
+	// lockEnabledViaSetMethod protects enabledViaSetMethod
+	lockEnabledViaSetMethod sync.RWMutex
+
+	// enabledViaSetMethod holds values explicitly set
+	// via Set method, features stored in this map take
+	// precedence over features stored in enabledViaEnvVar
+	enabledViaSetMethod map[Feature]bool
 
 	// readEnvVars holds the boolean value which
 	// indicates whether readEnvVarsOnce has been called.
 	readEnvVars atomic.Bool
 }
 
-// Enabled returns true if the key is enabled.  If the key is not known, this call will panic.
+// Enabled returns true if the key is enabled. If the key is not known, this call will panic.
 func (f *envVarFeatureGates) Enabled(key Feature) bool {
+	if v, ok := f.wasFeatureEnabledViaSetMethod(key); ok {
+		// ensue that the state of all known features
+		// is loaded from environment variables
+		// on the first call to Enabled method.
+		if !f.hasAlreadyReadEnvVar() {
+			_ = f.getEnabledMapFromEnvVar()
+		}
+		return v
+	}
 	if v, ok := f.getEnabledMapFromEnvVar()[key]; ok {
 		return v
 	}
@@ -94,6 +116,26 @@ func (f *envVarFeatureGates) Enabled(key Feature) bool {
 	panic(fmt.Errorf("feature %q is not registered in FeatureGates %q", key, f.callSiteName))
 }
 
+// Set sets the given feature to the given value.
+//
+// Features set via this method take precedence over
+// the features set via environment variables.
+func (f *envVarFeatureGates) Set(featureName Feature, featureValue bool) error {
+	feature, ok := f.known[featureName]
+	if !ok {
+		return fmt.Errorf("feature %q is not registered in FeatureGates %q", featureName, f.callSiteName)
+	}
+	if feature.LockToDefault && feature.Default != featureValue {
+		return fmt.Errorf("cannot set feature gate %q to %v, feature is locked to %v", featureName, featureValue, feature.Default)
+	}
+
+	f.lockEnabledViaSetMethod.Lock()
+	defer f.lockEnabledViaSetMethod.Unlock()
+	f.enabledViaSetMethod[featureName] = featureValue
+
+	return nil
+}
+
 // getEnabledMapFromEnvVar will fill the enabled map on the first call.
 // This is the only time a known feature can be set to a value
 // read from the corresponding environmental variable.
@@ -119,7 +161,7 @@ func (f *envVarFeatureGates) getEnabledMapFromEnvVar() map[Feature]bool {
 				featureGatesState[feature] = boolVal
 			}
 		}
-		f.enabled.Store(featureGatesState)
+		f.enabledViaEnvVar.Store(featureGatesState)
 		f.readEnvVars.Store(true)
 
 		for feature, featureSpec := range f.known {
@@ -130,7 +172,15 @@ func (f *envVarFeatureGates) getEnabledMapFromEnvVar() map[Feature]bool {
 			klog.V(1).InfoS("Feature gate default state", "feature", feature, "enabled", featureSpec.Default)
 		}
 	})
-	return f.enabled.Load().(map[Feature]bool)
+	return f.enabledViaEnvVar.Load().(map[Feature]bool)
+}
+
+func (f *envVarFeatureGates) wasFeatureEnabledViaSetMethod(key Feature) (bool, bool) {
+	f.lockEnabledViaSetMethod.RLock()
+	defer f.lockEnabledViaSetMethod.RUnlock()
+
+	value, found := f.enabledViaSetMethod[key]
+	return value, found
 }
 
 func (f *envVarFeatureGates) hasAlreadyReadEnvVar() bool {

features/envvar_test.go

@@ -23,21 +23,21 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestEnvVarFeatureGates(t *testing.T) {
+var defaultTestFeatures = map[Feature]FeatureSpec{
-	defaultTestFeatures := map[Feature]FeatureSpec{
+	"TestAlpha": {
-		"TestAlpha": {
+		Default:       false,
-			Default:       false,
+		LockToDefault: false,
-			LockToDefault: false,
+		PreRelease:    "Alpha",
-			PreRelease:    "Alpha",
+	},
-		},
+	"TestBeta": {
-		"TestBeta": {
+		Default:       true,
-			Default:       true,
+		LockToDefault: false,
-			LockToDefault: false,
+		PreRelease:    "Beta",
-			PreRelease:    "Beta",
+	},
-		},
+}
-	}
-	expectedDefaultFeaturesState := map[Feature]bool{"TestAlpha": false, "TestBeta": true}
 
+func TestEnvVarFeatureGates(t *testing.T) {
+	expectedDefaultFeaturesState := map[Feature]bool{"TestAlpha": false, "TestBeta": true}
 	copyExpectedStateMap := func(toCopy map[Feature]bool) map[Feature]bool {
 		m := map[Feature]bool{}
 		for k, v := range toCopy {
@@ -47,11 +47,14 @@ func TestEnvVarFeatureGates(t *testing.T) {
 	}
 
 	scenarios := []struct {
-		name                                string
+		name              string
-		features                            map[Feature]FeatureSpec
+		features          map[Feature]FeatureSpec
-		envVariables                        map[string]string
+		envVariables      map[string]string
-		expectedFeaturesState               map[Feature]bool
+		setMethodFeatures map[Feature]bool
-		expectedInternalEnabledFeatureState map[Feature]bool
+
+		expectedFeaturesState                           map[Feature]bool
+		expectedInternalEnabledViaEnvVarFeatureState    map[Feature]bool
+		expectedInternalEnabledViaSetMethodFeatureState map[Feature]bool
 	}{
 		{
 			name: "can add empty features",
@@ -76,7 +79,7 @@ func TestEnvVarFeatureGates(t *testing.T) {
 				expectedDefaultFeaturesStateCopy["TestAlpha"] = true
 				return expectedDefaultFeaturesStateCopy
 			}(),
-			expectedInternalEnabledFeatureState: map[Feature]bool{"TestAlpha": true},
+			expectedInternalEnabledViaEnvVarFeatureState: map[Feature]bool{"TestAlpha": true},
 		},
 		{
 			name:                  "incorrect env var value gets ignored",
@@ -111,9 +114,25 @@ func TestEnvVarFeatureGates(t *testing.T) {
 					PreRelease:    "Alpha",
 				},
 			},
-			envVariables:                        map[string]string{"KUBE_FEATURE_TestAlpha": "True"},
+			envVariables:          map[string]string{"KUBE_FEATURE_TestAlpha": "True"},
-			expectedFeaturesState:               map[Feature]bool{"TestAlpha": true},
+			expectedFeaturesState: map[Feature]bool{"TestAlpha": true},
-			expectedInternalEnabledFeatureState: map[Feature]bool{"TestAlpha": true},
+			expectedInternalEnabledViaEnvVarFeatureState: map[Feature]bool{"TestAlpha": true},
+		},
+		{
+			name:                  "setting a feature via the Set method works",
+			features:              defaultTestFeatures,
+			setMethodFeatures:     map[Feature]bool{"TestAlpha": true},
+			expectedFeaturesState: map[Feature]bool{"TestAlpha": true},
+			expectedInternalEnabledViaSetMethodFeatureState: map[Feature]bool{"TestAlpha": true},
+		},
+		{
+			name:                  "setting a feature via the Set method wins",
+			features:              defaultTestFeatures,
+			setMethodFeatures:     map[Feature]bool{"TestAlpha": false},
+			envVariables:          map[string]string{"KUBE_FEATURE_TestAlpha": "True"},
+			expectedFeaturesState: map[Feature]bool{"TestAlpha": false},
+			expectedInternalEnabledViaEnvVarFeatureState:    map[Feature]bool{"TestAlpha": true},
+			expectedInternalEnabledViaSetMethodFeatureState: map[Feature]bool{"TestAlpha": false},
 		},
 	}
 	for _, scenario := range scenarios {
@@ -123,20 +142,33 @@ func TestEnvVarFeatureGates(t *testing.T) {
 			}
 			target := newEnvVarFeatureGates(scenario.features)
 
+			for k, v := range scenario.setMethodFeatures {
+				err := target.Set(k, v)
+				require.NoError(t, err)
+			}
 			for expectedFeature, expectedValue := range scenario.expectedFeaturesState {
 				actualValue := target.Enabled(expectedFeature)
 				require.Equal(t, actualValue, expectedValue, "expected feature=%v, to be=%v, not=%v", expectedFeature, expectedValue, actualValue)
 			}
 
-			enabledInternalMap := target.enabled.Load().(map[Feature]bool)
+			enabledViaEnvVarInternalMap := target.enabledViaEnvVar.Load().(map[Feature]bool)
-			require.Len(t, enabledInternalMap, len(scenario.expectedInternalEnabledFeatureState))
+			require.Len(t, enabledViaEnvVarInternalMap, len(scenario.expectedInternalEnabledViaEnvVarFeatureState))
+			for expectedFeatureName, expectedFeatureValue := range scenario.expectedInternalEnabledViaEnvVarFeatureState {
+				actualFeatureValue, wasExpectedFeatureFound := enabledViaEnvVarInternalMap[expectedFeatureName]
+				if !wasExpectedFeatureFound {
+					t.Errorf("feature %v has not been found in enabledViaEnvVarInternalMap", expectedFeatureName)
+				}
+				require.Equal(t, expectedFeatureValue, actualFeatureValue, "feature %v has incorrect value = %v, expected = %v", expectedFeatureName, actualFeatureValue, expectedFeatureValue)
+			}
 
-			for expectedFeature, expectedInternalPresence := range scenario.expectedInternalEnabledFeatureState {
+			enabledViaSetMethodInternalMap := target.enabledViaSetMethod
-				featureInternalValue, featureSet := enabledInternalMap[expectedFeature]
+			require.Len(t, enabledViaSetMethodInternalMap, len(scenario.expectedInternalEnabledViaSetMethodFeatureState))
-				require.Equal(t, expectedInternalPresence, featureSet, "feature %v present = %v, expected = %v", expectedFeature, featureSet, expectedInternalPresence)
+			for expectedFeatureName, expectedFeatureValue := range scenario.expectedInternalEnabledViaSetMethodFeatureState {
-
+				actualFeatureValue, wasExpectedFeatureFound := enabledViaSetMethodInternalMap[expectedFeatureName]
-				expectedFeatureInternalValue := scenario.expectedFeaturesState[expectedFeature]
+				if !wasExpectedFeatureFound {
-				require.Equal(t, expectedFeatureInternalValue, featureInternalValue)
+					t.Errorf("feature %v has not been found in enabledViaSetMethod", expectedFeatureName)
+				}
+				require.Equal(t, expectedFeatureValue, actualFeatureValue, "feature %v has incorrect value = %v, expected = %v", expectedFeatureName, actualFeatureValue, expectedFeatureValue)
 			}
 		})
 	}
@@ -154,3 +186,48 @@ func TestHasAlreadyReadEnvVar(t *testing.T) {
 	_ = target.getEnabledMapFromEnvVar()
 	require.True(t, target.hasAlreadyReadEnvVar())
 }
+
+func TestEnvVarFeatureGatesSetNegative(t *testing.T) {
+	scenarios := []struct {
+		name         string
+		features     map[Feature]FeatureSpec
+		featureName  Feature
+		featureValue bool
+
+		expectedErr func(string) error
+	}{
+		{
+			name:     "empty feature name returns an error",
+			features: defaultTestFeatures,
+			expectedErr: func(callSiteName string) error {
+				return fmt.Errorf("feature %q is not registered in FeatureGates %q", "", callSiteName)
+			},
+		},
+		{
+			name:        "setting unknown feature returns an error",
+			features:    defaultTestFeatures,
+			featureName: "Unknown",
+			expectedErr: func(callSiteName string) error {
+				return fmt.Errorf("feature %q is not registered in FeatureGates %q", "Unknown", callSiteName)
+			},
+		},
+		{
+			name:         "setting locked feature returns an error",
+			features:     map[Feature]FeatureSpec{"LockedFeature": {LockToDefault: true, Default: true}},
+			featureName:  "LockedFeature",
+			featureValue: false,
+			expectedErr: func(_ string) error {
+				return fmt.Errorf("cannot set feature gate %q to %v, feature is locked to %v", "LockedFeature", false, true)
+			},
+		},
+	}
+
+	for _, scenario := range scenarios {
+		t.Run(scenario.name, func(t *testing.T) {
+			target := newEnvVarFeatureGates(scenario.features)
+
+			err := target.Set(scenario.featureName, scenario.featureValue)
+			require.Equal(t, scenario.expectedErr(target.callSiteName), err)
+		})
+	}
+}

features/features_test.go

@@ -30,6 +30,15 @@ func TestAddFeaturesToExistingFeatureGates(t *testing.T) {
 	require.Equal(t, defaultKubernetesFeatureGates, fakeFeatureGates.specs)
 }
 
+func TestReplaceFeatureGatesWithWarningIndicator(t *testing.T) {
+	defaultFeatureGates := FeatureGates()
+	require.Panics(t, func() { defaultFeatureGates.Enabled("Foo") }, "reading an unregistered feature gate Foo should panic")
+
+	if !replaceFeatureGatesWithWarningIndicator(defaultFeatureGates) {
+		t.Error("replacing the default feature gates after reading a value hasn't produced a warning")
+	}
+}
+
 type fakeRegistry struct {
 	specs map[Feature]FeatureSpec
 }