Fix feature setsjc-feature-set-improvement

13 files changed, 92 insertions, 62 deletions

diff --git a/changelogs/unreleased/jc-feature-set-improvement.yml b/changelogs/unreleased/jc-feature-set-improvement.yml
new file mode 100644
index 000000000..230f1eaea
--- /dev/null
+++ b/changelogs/unreleased/jc-feature-set-improvement.yml
@@ -0,0 +1,5 @@
+---
+title: Fix feature sets
+merge_request: 2411
+author:
+type: fixed
diff --git a/cmd/gitaly-ssh/main.go b/cmd/gitaly-ssh/main.go
index d611da5c7..88721ea46 100644
--- a/cmd/gitaly-ssh/main.go
+++ b/cmd/gitaly-ssh/main.go
@@ -90,7 +90,7 @@ func (cmd gitalySSHCommand) run() (int, error) {
 			if len(flagPairSplit) != 2 {
 				continue
 			}
-			ctx = featureflag.OutgoingCtxWithFeatureFlagValue(ctx, featureflag.FeatureFlag{Name: flagPairSplit[0]}, flagPairSplit[1])
+			ctx = featureflag.OutgoingCtxWithFeatureFlagValue(ctx, featureflag.NewGoFeatureFlag(flagPairSplit[0], false), flagPairSplit[1])
 		}
 	}
 
diff --git a/internal/cache/cachedb_test.go b/internal/cache/cachedb_test.go
index 9743a0c0c..3513af37e 100644
--- a/internal/cache/cachedb_test.go
+++ b/internal/cache/cachedb_test.go
@@ -94,7 +94,7 @@ func TestStreamDBNaiveKeyer(t *testing.T) {
 
 	// enabled feature flags affect caching
 	oldCtx := ctx
-	ctx = featureflag.IncomingCtxWithFeatureFlag(ctx, featureflag.FeatureFlag{"meow", false})
+	ctx = featureflag.IncomingCtxWithFeatureFlag(ctx, featureflag.NewGoFeatureFlag("meow", false))
 	expectGetMiss(req1)
 	ctx = oldCtx
 	expectGetHit(expectStream2, req1)
diff --git a/internal/metadata/featureflag/context.go b/internal/metadata/featureflag/context.go
index 6522b4018..04d5d6e7a 100644
--- a/internal/metadata/featureflag/context.go
+++ b/internal/metadata/featureflag/context.go
@@ -18,7 +18,7 @@ func OutgoingCtxWithFeatureFlags(ctx context.Context, flags ...FeatureFlag) cont
 	}
 
 	for _, flag := range flags {
-		md.Set(HeaderKey(flag.Name), "true")
+		md.Set(HeaderKey(flag.GetName()), "true")
 	}
 
 	return metadata.NewOutgoingContext(ctx, md)
@@ -35,7 +35,7 @@ func OutgoingCtxWithDisabledFeatureFlags(ctx context.Context, flags ...FeatureFl
 	}
 
 	for _, flag := range flags {
-		md.Set(HeaderKey(flag.Name), "false")
+		md.Set(HeaderKey(flag.GetName()), "false")
 	}
 
 	return metadata.NewOutgoingContext(ctx, md)
@@ -53,7 +53,7 @@ func OutgoingCtxWithFeatureFlagValue(ctx context.Context, flag FeatureFlag, val
 		md = metadata.New(map[string]string{})
 	}
 
-	md.Set(HeaderKey(flag.Name), val)
+	md.Set(HeaderKey(flag.GetName()), val)
 
 	return metadata.NewOutgoingContext(ctx, md)
 }
@@ -66,7 +66,7 @@ func IncomingCtxWithFeatureFlag(ctx context.Context, flag FeatureFlag) context.C
 	if !ok {
 		md = metadata.New(map[string]string{})
 	}
-	md.Set(HeaderKey(flag.Name), "true")
+	md.Set(HeaderKey(flag.GetName()), "true")
 	return metadata.NewIncomingContext(ctx, md)
 }
 
@@ -77,7 +77,7 @@ func OutgoingCtxWithRubyFeatureFlags(ctx context.Context, flags ...FeatureFlag)
 	}
 
 	for _, flag := range flags {
-		md.Set(rubyHeaderKey(flag.Name), "true")
+		md.Set(rubyHeaderKey(flag.GetName()), "true")
 	}
 
 	return metadata.NewOutgoingContext(ctx, md)
diff --git a/internal/metadata/featureflag/context_test.go b/internal/metadata/featureflag/context_test.go
index e184340f4..4c02dd628 100644
--- a/internal/metadata/featureflag/context_test.go
+++ b/internal/metadata/featureflag/context_test.go
@@ -8,7 +8,7 @@ import (
 	"google.golang.org/grpc/metadata"
 )
 
-var mockFeatureFlag = FeatureFlag{"turn meow on", false}
+var mockFeatureFlag = GoFeatureFlag{"turn meow on", false}
 
 func TestIncomingCtxWithFeatureFlag(t *testing.T) {
 	ctx := context.Background()
diff --git a/internal/metadata/featureflag/feature_flags.go b/internal/metadata/featureflag/feature_flags.go
index 19f703759..ee8c60682 100644
--- a/internal/metadata/featureflag/feature_flags.go
+++ b/internal/metadata/featureflag/feature_flags.go
@@ -1,31 +1,55 @@
 package featureflag
 
-type FeatureFlag struct {
-	Name        string `json:"name"`
+type FeatureFlag interface {
+	GetName() string
+	IsOnByDefault() bool
+}
+
+type RubyFeatureFlag struct {
+	GoFeatureFlag
+}
+
+type GoFeatureFlag struct {
+	Name        string `json:"Name"`
 	OnByDefault bool   `json:"on_by_default"`
 }
 
+func (g GoFeatureFlag) GetName() string {
+	return g.Name
+}
+
+func (g GoFeatureFlag) IsOnByDefault() bool {
+	return g.OnByDefault
+}
+
+func NewGoFeatureFlag(Name string, OnByDefault bool) *GoFeatureFlag {
+	return &GoFeatureFlag{
+		Name:        Name,
+		OnByDefault: OnByDefault,
+	}
+}
+
 var (
 	// GoUpdateHook will bypass the ruby update hook and use the go implementation of custom hooks
-	GoUpdateHook = FeatureFlag{Name: "go_update_hook", OnByDefault: true}
+	GoUpdateHook = GoFeatureFlag{Name: "go_update_hook", OnByDefault: true}
 	// RemoteBranchesLsRemote will use `ls-remote` for remote branches
-	RemoteBranchesLsRemote = FeatureFlag{Name: "ruby_remote_branches_ls_remote", OnByDefault: true}
+	RemoteBranchesLsRemote = GoFeatureFlag{Name: "ruby_remote_branches_ls_remote", OnByDefault: true}
 	// GoFetchSourceBranch enables a go implementation of FetchSourceBranch
-	GoFetchSourceBranch = FeatureFlag{Name: "go_fetch_source_branch", OnByDefault: false}
+	GoFetchSourceBranch = GoFeatureFlag{Name: "go_fetch_source_branch", OnByDefault: false}
 	// DistributedReads allows praefect to redirect accessor operations to up-to-date secondaries
-	DistributedReads = FeatureFlag{Name: "distributed_reads", OnByDefault: false}
+	DistributedReads = GoFeatureFlag{Name: "distributed_reads", OnByDefault: false}
 	// GoPreReceiveHook will bypass the ruby pre-receive hook and use the go implementation
-	GoPreReceiveHook = FeatureFlag{Name: "go_prereceive_hook", OnByDefault: true}
+	GoPreReceiveHook = GoFeatureFlag{Name: "go_prereceive_hook", OnByDefault: true}
 	// GoPostReceiveHook will bypass the ruby post-receive hook and use the go implementation
-	GoPostReceiveHook = FeatureFlag{Name: "go_postreceive_hook", OnByDefault: false}
+	GoPostReceiveHook = GoFeatureFlag{Name: "go_postreceive_hook", OnByDefault: false}
 	// ReferenceTransactions will handle Git reference updates via the transaction service for strong consistency
-	ReferenceTransactions = FeatureFlag{Name: "reference_transactions", OnByDefault: false}
+	ReferenceTransactions = GoFeatureFlag{Name: "reference_transactions", OnByDefault: false}
 	// ReferenceTransactionsOperationService will enable reference transactions for the OperationService
-	ReferenceTransactionsOperationService = FeatureFlag{Name: "reference_transactions_operation_service", OnByDefault: true}
+	ReferenceTransactionsOperationService = GoFeatureFlag{Name: "reference_transactions_operation_service", OnByDefault: true}
 	// ReferenceTransactionsSmartHTTPService will enable reference transactions for the SmartHTTPService
-	ReferenceTransactionsSmartHTTPService = FeatureFlag{Name: "reference_transactions_smarthttp_service", OnByDefault: true}
+	ReferenceTransactionsSmartHTTPService = GoFeatureFlag{Name: "reference_transactions_smarthttp_service", OnByDefault: true}
 	// ReferenceTransactionsSSHService will enable reference transactions for the SSHService
-	ReferenceTransactionsSSHService = FeatureFlag{Name: "reference_transactions_ssh_service", OnByDefault: true}
+	ReferenceTransactionsSSHService = GoFeatureFlag{Name: "reference_transactions_ssh_service", OnByDefault: true}
 )
 
 const (
diff --git a/internal/metadata/featureflag/grpc_header.go b/internal/metadata/featureflag/grpc_header.go
index 23e144d86..9e037f7b2 100644
--- a/internal/metadata/featureflag/grpc_header.go
+++ b/internal/metadata/featureflag/grpc_header.go
@@ -27,14 +27,14 @@ func init() {
 // IsEnabled checks if the feature flag is enabled for the passed context.
 // Only returns true if the metadata for the feature flag is set to "true"
 func IsEnabled(ctx context.Context, flag FeatureFlag) bool {
-	val, ok := getFlagVal(ctx, flag.Name)
+	val, ok := getFlagVal(ctx, flag.GetName())
 	if !ok {
-		return flag.OnByDefault
+		return flag.IsOnByDefault()
 	}
 
 	enabled := val == "true"
 
-	flagChecks.WithLabelValues(flag.Name, strconv.FormatBool(enabled)).Inc()
+	flagChecks.WithLabelValues(flag.GetName(), strconv.FormatBool(enabled)).Inc()
 
 	return enabled
 }
diff --git a/internal/metadata/featureflag/grpc_header_test.go b/internal/metadata/featureflag/grpc_header_test.go
index e3810900a..a4ba26087 100644
--- a/internal/metadata/featureflag/grpc_header_test.go
+++ b/internal/metadata/featureflag/grpc_header_test.go
@@ -31,7 +31,7 @@ func TestGRPCMetadataFeatureFlag(t *testing.T) {
 			md := metadata.New(tc.headers)
 			ctx := metadata.NewIncomingContext(context.Background(), md)
 
-			assert.Equal(t, tc.enabled, IsEnabled(ctx, FeatureFlag{tc.flag, tc.onByDefault}))
+			assert.Equal(t, tc.enabled, IsEnabled(ctx, GoFeatureFlag{tc.flag, tc.onByDefault}))
 		})
 	}
 }
diff --git a/internal/praefect/datastore/queue_test.go b/internal/praefect/datastore/queue_test.go
index 3ea4de70d..4cfcf2200 100644
--- a/internal/praefect/datastore/queue_test.go
+++ b/internal/praefect/datastore/queue_test.go
@@ -4,8 +4,8 @@ package datastore
 
 import (
 	"context"
-	"testing"
 	"sync"
+	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
diff --git a/internal/service/operations/rebase_test.go b/internal/service/operations/rebase_test.go
index 2c6cecbdb..0b8bd8951 100644
--- a/internal/service/operations/rebase_test.go
+++ b/internal/service/operations/rebase_test.go
@@ -24,7 +24,7 @@ var (
 )
 
 func TestSuccessfulUserRebaseConfirmableRequest(t *testing.T) {
-	featureSets, err := testhelper.NewFeatureSets(nil, featureflag.GoPostReceiveHook)
+	featureSets, err := testhelper.NewFeatureSets([]featureflag.FeatureFlag{featureflag.RubyFeatureFlag{featureflag.GoPostReceiveHook}})
 	require.NoError(t, err)
 
 	ctx, cancel := testhelper.Context()
diff --git a/internal/service/operations/tags_test.go b/internal/service/operations/tags_test.go
index bc06a693c..a4a356fd8 100644
--- a/internal/service/operations/tags_test.go
+++ b/internal/service/operations/tags_test.go
@@ -62,7 +62,7 @@ func TestSuccessfulGitHooksForUserDeleteTagRequest(t *testing.T) {
 	ctx, cancel := testhelper.Context()
 	defer cancel()
 
-	featureSets, err := testhelper.NewFeatureSets(nil, featureflag.GoPostReceiveHook)
+	featureSets, err := testhelper.NewFeatureSets([]featureflag.FeatureFlag{featureflag.RubyFeatureFlag{featureflag.GoPostReceiveHook}})
 	require.NoError(t, err)
 
 	for _, featureSet := range featureSets {
diff --git a/internal/service/operations/update_branches_test.go b/internal/service/operations/update_branches_test.go
index 542d6cff2..52ca772d1 100644
--- a/internal/service/operations/update_branches_test.go
+++ b/internal/service/operations/update_branches_test.go
@@ -102,12 +102,12 @@ func testSuccessfulGitHooksForUserUpdateBranchRequest(t *testing.T, ctx context.
 }
 
 func TestFailedUserUpdateBranchDueToHooks(t *testing.T) {
-	featureSet, err := testhelper.NewFeatureSets(nil, featureflag.GoUpdateHook)
+	featureSets, err := testhelper.NewFeatureSets([]featureflag.FeatureFlag{featureflag.RubyFeatureFlag{featureflag.GoUpdateHook}})
 	require.NoError(t, err)
 	ctx, cancel := testhelper.Context()
 	defer cancel()
 
-	for _, features := range featureSet {
+	for _, features := range featureSets {
 		t.Run(features.String(), func(t *testing.T) {
 			ctx = features.WithParent(ctx)
 			testFailedUserUpdateBranchDueToHooks(t, ctx)
diff --git a/internal/testhelper/testhelper.go b/internal/testhelper/testhelper.go
index 72d803176..cb13ca76b 100644
--- a/internal/testhelper/testhelper.go
+++ b/internal/testhelper/testhelper.go
@@ -736,23 +736,23 @@ func WriteBlobs(t testing.TB, testRepoPath string, n int) []string {
 // FeatureSet is a representation of a set of features that are enabled
 // This is useful in situations where a test needs to test any combination of features toggled on and off
 type FeatureSet struct {
-	features     map[featureflag.FeatureFlag]bool
-	rubyFeatures map[featureflag.FeatureFlag]bool
+	featureIndices map[featureflag.FeatureFlag]int
+	bitmask        int
 }
 
-func (f FeatureSet) IsEnabled(flag featureflag.FeatureFlag) bool {
-	on, ok := f.features[flag]
-	if !ok {
-		return flag.OnByDefault
-	}
-
-	return on
+func (f FeatureSet) IsEnabled(feature featureflag.FeatureFlag) bool {
+	return f.featureIndices[feature]&f.bitmask > 0
 }
 
 func (f FeatureSet) String() string {
 	features := make([]string, 0, len(f.enabledFeatures()))
 	for _, feature := range f.enabledFeatures() {
-		features = append(features, feature.Name)
+		switch feature.(type) {
+		case featureflag.GoFeatureFlag:
+			features = append(features, feature.GetName())
+		case featureflag.RubyFeatureFlag:
+			features = append(features, feature.GetName()+"(ruby)")
+		}
 	}
 
 	return strings.Join(features, ",")
@@ -761,20 +761,27 @@ func (f FeatureSet) String() string {
 func (f FeatureSet) enabledFeatures() []featureflag.FeatureFlag {
 	var enabled []featureflag.FeatureFlag
 
-	for feature := range f.features {
-		enabled = append(enabled, feature)
+	for feature, featureVal := range f.featureIndices {
+		if featureVal&f.bitmask > 0 {
+			enabled = append(enabled, feature)
+		}
 	}
 
 	return enabled
 }
 
 func (f FeatureSet) WithParent(ctx context.Context) context.Context {
-	for _, enabledFeature := range f.enabledFeatures() {
-		if _, ok := f.rubyFeatures[enabledFeature]; ok {
-			ctx = featureflag.OutgoingCtxWithRubyFeatureFlags(ctx, enabledFeature)
+	for feature, featureVal := range f.featureIndices {
+		if f.bitmask&featureVal == 0 {
 			continue
 		}
-		ctx = featureflag.OutgoingCtxWithFeatureFlags(ctx, enabledFeature)
+
+		switch feature.(type) {
+		case featureflag.GoFeatureFlag:
+			ctx = featureflag.OutgoingCtxWithRubyFeatureFlags(ctx, feature)
+		case featureflag.RubyFeatureFlag:
+			ctx = featureflag.OutgoingCtxWithFeatureFlags(ctx, feature)
+		}
 	}
 
 	return ctx
@@ -785,27 +792,21 @@ type FeatureSets []FeatureSet
 
 // NewFeatureSets takes a slice of go feature flags, and an optional variadic set of ruby feature flags
 // and returns a FeatureSets slice
-func NewFeatureSets(goFeatures []featureflag.FeatureFlag, rubyFeatures ...featureflag.FeatureFlag) (FeatureSets, error) {
-	rubyFeatureMap := make(map[featureflag.FeatureFlag]bool)
-	for _, rubyFeature := range rubyFeatures {
-		rubyFeatureMap[rubyFeature] = true
-	}
-
-	// start with an empty feature set
-	f := []FeatureSet{{features: make(map[featureflag.FeatureFlag]bool), rubyFeatures: rubyFeatureMap}}
+func NewFeatureSets(features []featureflag.FeatureFlag) (FeatureSets, error) {
+	featureIndices := make(map[featureflag.FeatureFlag]int)
 
-	allFeatures := append(goFeatures, rubyFeatures...)
-
-	for i := range allFeatures {
-		featureMap := make(map[featureflag.FeatureFlag]bool)
-		for j := 0; j <= i; j++ {
-			featureMap[allFeatures[j]] = true
-		}
+	featureVal := 1
+	for _, feature := range features {
+		featureIndices[feature] = featureVal
+		featureVal <<= 1
+	}
 
-		f = append(f, FeatureSet{features: featureMap, rubyFeatures: rubyFeatureMap})
+	featureSets := make([]FeatureSet, featureVal)
+	for i := 0; i < featureVal; i++ {
+		featureSets[i] = FeatureSet{featureIndices: featureIndices, bitmask: i}
 	}
 
-	return f, nil
+	return featureSets, nil
 }
 
 // ModifyEnvironment will change an environment variable and return a func suitable