Change featureflags package to make it test friendly (#5160)

* change feature registry as a prototype to make test friendly

Signed-off-by: Kay Yan <kay.yan@daocloud.io>

* Change colTelemetry to work with a custom Registry, improve tests

Signed-off-by: Bogdan Drutu <bogdandrutu@gmail.com>

* Fix tests, windows, deprecate old API

Co-authored-by: Alex Boten <aboten@lightstep.com>

CHANGELOG.md

@@ -23,6 +23,7 @@
   - `pmetric.MetricValueTypeInt` is deprecated in favor of `NumberDataPointValueTypeInt`
   - `pmetric.MetricValueTypeDouble` is deprecated in favor of `NumberDataPointValueTypeDouble`
 - Deprecate `plog.LogRecord.SetName()` function (#5230)
+- Deprecate global `featuregate` funcs in favor of `GetRegistry` and a public `Registry` type (#5160)
 
 ### 💡 Enhancements 💡
 

service/collector.go

@@ -103,6 +103,10 @@ func New(set CollectorSettings) (*Collector, error) {
 		return nil, errors.New("invalid nil config provider")
 	}
 
+	if set.telemetry == nil {
+		set.telemetry = collectorTelemetry
+	}
+
 	return &Collector{
 		logger: zap.NewNop(), // Set a Nop logger as a place holder until a logger is created based on configuration
 
@@ -222,7 +226,7 @@ func (col *Collector) setupConfigurationComponents(ctx context.Context) error {
 	// TODO: This should be part of the service initialization, which should be responsible to create TelemetrySettings.
 	// For the moment happens here, since it needs service.Config and Logger.
 	// It is called once because that is how it is implemented using sync.Once.
-	if err = collectorTelemetry.init(col); err != nil {
+	if err = col.set.telemetry.init(col); err != nil {
 		return err
 	}
 
@@ -272,7 +276,7 @@ func (col *Collector) shutdown(ctx context.Context) error {
 		errs = multierr.Append(errs, fmt.Errorf("failed to shutdown service: %w", err))
 	}
 
-	if err := collectorTelemetry.shutdown(); err != nil {
+	if err := col.set.telemetry.shutdown(); err != nil {
 		errs = multierr.Append(errs, fmt.Errorf("failed to shutdown collector telemetry: %w", err))
 	}
 

service/collector_test.go

@@ -54,11 +54,6 @@ func TestStateString(t *testing.T) {
 // TestCollector_StartAsGoRoutine must be the first unit test on the file,
 // to test for initialization without setting CLI flags.
 func TestCollector_StartAsGoRoutine(t *testing.T) {
-	// use a mock AppTelemetry struct to return an error on shutdown
-	preservedAppTelemetry := collectorTelemetry
-	collectorTelemetry = &colTelemetry{}
-	defer func() { collectorTelemetry = preservedAppTelemetry }()
-
 	factories, err := testcomponents.NewDefaultFactories()
 	require.NoError(t, err)
 
@@ -73,6 +68,7 @@ func TestCollector_StartAsGoRoutine(t *testing.T) {
 		BuildInfo:      component.NewDefaultBuildInfo(),
 		Factories:      factories,
 		ConfigProvider: cfgProvider,
+		telemetry:      newColTelemetry(featuregate.NewRegistry()),
 	}
 	col, err := New(set)
 	require.NoError(t, err)
@@ -93,7 +89,7 @@ func TestCollector_StartAsGoRoutine(t *testing.T) {
 	assert.Equal(t, Closed, col.GetState())
 }
 
-func testCollectorStartHelper(t *testing.T) {
+func testCollectorStartHelper(t *testing.T, telemetry collectorTelemetryExporter) {
 	factories, err := testcomponents.NewDefaultFactories()
 	require.NoError(t, err)
 	var once sync.Once
@@ -123,6 +119,7 @@ func testCollectorStartHelper(t *testing.T) {
 		Factories:      factories,
 		ConfigProvider: cfgProvider,
 		LoggingOptions: []zap.Option{zap.Hooks(hook)},
+		telemetry:      telemetry,
 	})
 	require.NoError(t, err)
 
@@ -156,29 +153,16 @@ func testCollectorStartHelper(t *testing.T) {
 	assert.Equal(t, Closed, col.GetState())
 }
 
-// as telemetry instance is initialized only once, we need to reset it before each test so the metrics endpoint can
-// have correct handler spawned
-func resetCollectorTelemetry() {
-	collectorTelemetry = &colTelemetry{}
-}
-
 func TestCollector_Start(t *testing.T) {
-	resetCollectorTelemetry()
-	testCollectorStartHelper(t)
+	testCollectorStartHelper(t, newColTelemetry(featuregate.NewRegistry()))
 }
 
 func TestCollector_StartWithOtelInternalMetrics(t *testing.T) {
-	resetCollectorTelemetry()
-	originalFlag := featuregate.IsEnabled(useOtelForInternalMetricsfeatureGateID)
-	defer func() {
-		featuregate.Apply(map[string]bool{
-			useOtelForInternalMetricsfeatureGateID: originalFlag,
-		})
-	}()
-	featuregate.Apply(map[string]bool{
+	colTel := newColTelemetry(featuregate.NewRegistry())
+	colTel.registry.Apply(map[string]bool{
 		useOtelForInternalMetricsfeatureGateID: true,
 	})
-	testCollectorStartHelper(t)
+	testCollectorStartHelper(t, colTel)
 }
 
 // TestCollector_ShutdownNoop verifies that shutdown can be called even if a collector
@@ -206,11 +190,6 @@ func TestCollector_ShutdownNoop(t *testing.T) {
 }
 
 func TestCollector_ShutdownBeforeRun(t *testing.T) {
-	// use a mock AppTelemetry struct to return an error on shutdown
-	preservedAppTelemetry := collectorTelemetry
-	collectorTelemetry = &colTelemetry{}
-	defer func() { collectorTelemetry = preservedAppTelemetry }()
-
 	factories, err := testcomponents.NewDefaultFactories()
 	require.NoError(t, err)
 
@@ -224,6 +203,7 @@ func TestCollector_ShutdownBeforeRun(t *testing.T) {
 		BuildInfo:      component.NewDefaultBuildInfo(),
 		Factories:      factories,
 		ConfigProvider: cfgProvider,
+		telemetry:      newColTelemetry(featuregate.NewRegistry()),
 	}
 	col, err := New(set)
 	require.NoError(t, err)
@@ -244,10 +224,6 @@ func TestCollector_ShutdownBeforeRun(t *testing.T) {
 
 // TestCollector_ClosedStateOnStartUpError tests the collector changes it's state to Closed when a startup error occurs
 func TestCollector_ClosedStateOnStartUpError(t *testing.T) {
-	preservedAppTelemetry := collectorTelemetry
-	collectorTelemetry = &colTelemetry{}
-	defer func() { collectorTelemetry = preservedAppTelemetry }()
-
 	factories, err := testcomponents.NewDefaultFactories()
 	require.NoError(t, err)
 
@@ -262,6 +238,7 @@ func TestCollector_ClosedStateOnStartUpError(t *testing.T) {
 		BuildInfo:      component.NewDefaultBuildInfo(),
 		Factories:      factories,
 		ConfigProvider: cfgProvider,
+		telemetry:      newColTelemetry(featuregate.NewRegistry()),
 	}
 	col, err := New(set)
 	require.NoError(t, err)
@@ -284,11 +261,6 @@ func (tel *mockColTelemetry) shutdown() error {
 }
 
 func TestCollector_ReportError(t *testing.T) {
-	// use a mock AppTelemetry struct to return an error on shutdown
-	preservedAppTelemetry := collectorTelemetry
-	collectorTelemetry = &mockColTelemetry{}
-	defer func() { collectorTelemetry = preservedAppTelemetry }()
-
 	factories, err := testcomponents.NewDefaultFactories()
 	require.NoError(t, err)
 
@@ -302,6 +274,7 @@ func TestCollector_ReportError(t *testing.T) {
 		BuildInfo:      component.NewDefaultBuildInfo(),
 		Factories:      factories,
 		ConfigProvider: cfgProvider,
+		telemetry:      &mockColTelemetry{},
 	})
 	require.NoError(t, err)
 
@@ -322,11 +295,6 @@ func TestCollector_ReportError(t *testing.T) {
 
 // TestCollector_ContextCancel tests that the collector gracefully exits on context cancel
 func TestCollector_ContextCancel(t *testing.T) {
-	// use a mock AppTelemetry struct to return an error on shutdown
-	preservedAppTelemetry := collectorTelemetry
-	collectorTelemetry = &colTelemetry{}
-	defer func() { collectorTelemetry = preservedAppTelemetry }()
-
 	factories, err := testcomponents.NewDefaultFactories()
 	require.NoError(t, err)
 
@@ -341,6 +309,7 @@ func TestCollector_ContextCancel(t *testing.T) {
 		BuildInfo:      component.NewDefaultBuildInfo(),
 		Factories:      factories,
 		ConfigProvider: cfgProvider,
+		telemetry:      newColTelemetry(featuregate.NewRegistry()),
 	}
 	col, err := New(set)
 	require.NoError(t, err)

service/collector_windows.go

@@ -93,7 +93,7 @@ func (s *windowsService) start(elog *eventlog.Log, colErrorChannel chan error) e
 	if err := flags().Parse(os.Args[1:]); err != nil {
 		return err
 	}
-	featuregate.Apply(gatesList)
+	featuregate.GetRegistry().Apply(gatesList)
 	var err error
 	s.col, err = newWithWindowsEventLogCore(s.settings, elog)
 	if err != nil {

service/command.go

@@ -29,7 +29,7 @@ func NewCommand(set CollectorSettings) *cobra.Command {
 		Version:      set.BuildInfo.Version,
 		SilenceUsage: true,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			featuregate.Apply(gatesList)
+			featuregate.GetRegistry().Apply(gatesList)
 			if set.ConfigProvider == nil {
 				var err error
 				cfgSet := newDefaultConfigProviderSettings(getConfigFlag())

service/featuregate/gates.go

@@ -27,40 +27,40 @@ type Gate struct {
 	Enabled     bool
 }
 
-var reg = &registry{gates: make(map[string]Gate)}
+var reg = NewRegistry()
 
-// IsEnabled returns true if a registered feature gate is enabled and false otherwise.
-func IsEnabled(id string) bool {
-	return reg.isEnabled(id)
+// GetRegistry returns the global Registry.
+func GetRegistry() *Registry {
+	return reg
 }
 
-// List returns a slice of copies of all registered Gates.
-func List() []Gate {
-	return reg.list()
-}
+// Deprecated: [v0.50.0] Use GetRegistry().Apply.
+var Apply = GetRegistry().Apply
 
-// Register a Gate. May only be called in an init() function.
-// Will panic() if a Gate with the same ID is already registered.
-func Register(g Gate) {
-	if err := reg.add(g); err != nil {
-		panic(err)
-	}
-}
+// Deprecated: [v0.50.0] Use GetRegistry().IsEnabled.
+var IsEnabled = GetRegistry().IsEnabled
 
-// Apply a configuration in the form of a map of Gate identifiers to boolean values.
-// Sets only those values provided in the map, other gate values are not changed.
-func Apply(cfg map[string]bool) {
-	reg.apply(cfg)
+// Deprecated: [v0.50.0] Use GetRegistry().List.
+var List = GetRegistry().List
+
+// Deprecated: [v0.50.0] Use GetRegistry().MustRegister.
+var Register = GetRegistry().MustRegister
+
+// NewRegistry returns a new empty Registry.
+func NewRegistry() *Registry {
+	return &Registry{gates: make(map[string]Gate)}
 }
 
-type registry struct {
-	sync.RWMutex
+type Registry struct {
+	mu    sync.RWMutex
 	gates map[string]Gate
 }
 
-func (r *registry) apply(cfg map[string]bool) {
-	r.Lock()
-	defer r.Unlock()
+// Apply a configuration in the form of a map of Gate identifiers to boolean values.
+// Sets only those values provided in the map, other gate values are not changed.
+func (r *Registry) Apply(cfg map[string]bool) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
 	for id, val := range cfg {
 		if g, ok := r.gates[id]; ok {
 			g.Enabled = val
@@ -69,27 +69,36 @@ func (r *registry) apply(cfg map[string]bool) {
 	}
 }
 
-func (r *registry) add(g Gate) error {
-	r.Lock()
-	defer r.Unlock()
+// IsEnabled returns true if a registered feature gate is enabled and false otherwise.
+func (r *Registry) IsEnabled(id string) bool {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	g, ok := r.gates[id]
+	return ok && g.Enabled
+}
+
+// MustRegister like Register but panics if a Gate with the same ID is already registered.
+func (r *Registry) MustRegister(g Gate) {
+	if err := r.Register(g); err != nil {
+		panic(err)
+	}
+}
+
+// Register registers a Gate. May only be called in an init() function.
+func (r *Registry) Register(g Gate) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
 	if _, ok := r.gates[g.ID]; ok {
 		return fmt.Errorf("attempted to add pre-existing gate %q", g.ID)
 	}
-
 	r.gates[g.ID] = g
 	return nil
 }
 
-func (r *registry) isEnabled(id string) bool {
-	r.RLock()
-	defer r.RUnlock()
-	g, ok := r.gates[id]
-	return ok && g.Enabled
-}
-
-func (r *registry) list() []Gate {
-	r.RLock()
-	defer r.RUnlock()
+// List returns a slice of copies of all registered Gates.
+func (r *Registry) List() []Gate {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
 	ret := make([]Gate, len(r.gates))
 	i := 0
 	for _, gate := range r.gates {

service/featuregate/gates_test.go

@@ -21,46 +21,26 @@ import (
 )
 
 func TestRegistry(t *testing.T) {
-	r := registry{gates: map[string]Gate{}}
+	r := Registry{gates: map[string]Gate{}}
 
 	gate := Gate{
 		ID:          "foo",
 		Description: "Test Gate",
 		Enabled:     true,
 	}
 
-	assert.Empty(t, r.list())
-	assert.False(t, r.isEnabled(gate.ID))
+	assert.Empty(t, r.List())
+	assert.False(t, r.IsEnabled(gate.ID))
 
-	assert.NoError(t, r.add(gate))
-	assert.Len(t, r.list(), 1)
-	assert.True(t, r.isEnabled(gate.ID))
+	assert.NoError(t, r.Register(gate))
+	assert.Len(t, r.List(), 1)
+	assert.True(t, r.IsEnabled(gate.ID))
 
-	r.apply(map[string]bool{gate.ID: false})
-	assert.False(t, r.isEnabled(gate.ID))
+	r.Apply(map[string]bool{gate.ID: false})
+	assert.False(t, r.IsEnabled(gate.ID))
 
-	assert.Error(t, r.add(gate))
-}
-
-func TestGlobalRegistry(t *testing.T) {
-	gate := Gate{
-		ID:          "feature_gate_test.foo",
-		Description: "Test Gate",
-		Enabled:     true,
-	}
-
-	assert.NotContains(t, List(), gate)
-	assert.False(t, IsEnabled(gate.ID))
-
-	assert.NotPanics(t, func() { Register(gate) })
-	assert.Contains(t, List(), gate)
-	assert.True(t, IsEnabled(gate.ID))
-
-	Apply(map[string]bool{gate.ID: false})
-	assert.False(t, IsEnabled(gate.ID))
-
-	assert.Panics(t, func() { Register(gate) })
-	reg.Lock()
-	delete(reg.gates, gate.ID)
-	reg.Unlock()
+	assert.Error(t, r.Register(gate))
+	assert.Panics(t, func() {
+		r.MustRegister(gate)
+	})
 }

service/settings.go

@@ -63,4 +63,7 @@ type CollectorSettings struct {
 
 	// LoggingOptions provides a way to change behavior of zap logging.
 	LoggingOptions []zap.Option
+
+	// For testing purpose only.
+	telemetry collectorTelemetryExporter
 }