diff --git a/libbeat/feature/feature.go b/libbeat/feature/feature.go index 5dc13c0d5cf6..36dc4ef84ae7 100644 --- a/libbeat/feature/feature.go +++ b/libbeat/feature/feature.go @@ -101,11 +101,43 @@ func New(namespace, name string, factory interface{}, stability Stability) *Feat // RegisterBundle registers a bundle of features. func RegisterBundle(bundle *Bundle) error { for _, f := range bundle.Features() { - Registry.Register(f) + err := Registry.Register(f) + if err != nil { + return err + } } return nil } +// MustRegisterBundle register a new bundle and panic on error. +func MustRegisterBundle(bundle *Bundle) { + err := RegisterBundle(bundle) + if err != nil { + panic(err) + } +} + +// OverwriteBundle register a bundle of feature and replace any existing feature with a new +// implementation. +func OverwriteBundle(bundle *Bundle) error { + for _, f := range bundle.Features() { + err := Registry.Register(f) + if err != nil { + return err + } + } + return nil +} + +// MustOverwriteBundle register a bundle of feature, replace any existing feature with a new +// implementation and panic on error. +func MustOverwriteBundle(bundle *Bundle) { + err := OverwriteBundle(bundle) + if err != nil { + panic(err) + } +} + // Register register a new feature on the global registry. func Register(feature Featurable) error { return Registry.Register(feature) diff --git a/libbeat/feature/registry.go b/libbeat/feature/registry.go index 905f0f55273a..4d5751ece382 100644 --- a/libbeat/feature/registry.go +++ b/libbeat/feature/registry.go @@ -159,6 +159,19 @@ func (r *registry) LookupAll(namespace string) ([]Featurable, error) { return list, nil } +// Overwrite allow to replace an existing feature with a new implementation. +func (r *registry) Overwrite(feature Featurable) error { + _, err := r.Lookup(feature.Namespace(), feature.Name()) + if err == nil { + err := r.Unregister(feature.Namespace(), feature.Name()) + if err != nil { + return err + } + } + + return r.Register(feature) +} + // Size returns the number of registered features in the registry. func (r *registry) Size() int { r.RLock() diff --git a/libbeat/feature/registry_test.go b/libbeat/feature/registry_test.go index 57afb5facedc..622990bccc0c 100644 --- a/libbeat/feature/registry_test.go +++ b/libbeat/feature/registry_test.go @@ -176,3 +176,33 @@ func TestUnregister(t *testing.T) { assert.Equal(t, 0, r.Size()) }) } + +func TestOverwrite(t *testing.T) { + t.Run("when the feature doesn't exist", func(t *testing.T) { + f := func() {} + r := newRegistry() + assert.Equal(t, 0, r.Size()) + r.Overwrite(New("processor", "foo", f, Stable)) + assert.Equal(t, 1, r.Size()) + }) + + t.Run("overwrite when the feature exists", func(t *testing.T) { + f := func() {} + r := newRegistry() + r.Register(New("processor", "foo", f, Stable)) + assert.Equal(t, 1, r.Size()) + + check := 42 + r.Overwrite(New("processor", "foo", check, Stable)) + assert.Equal(t, 1, r.Size()) + + feature, err := r.Lookup("processor", "foo") + if !assert.NoError(t, err) { + return + } + + v, ok := feature.Factory().(int) + assert.True(t, ok) + assert.Equal(t, 42, v) + }) +}