diff --git a/controller/initializer/initializer.go b/controller/initializer/initializer.go new file mode 100644 index 0000000..4b1ae67 --- /dev/null +++ b/controller/initializer/initializer.go @@ -0,0 +1,213 @@ +/** + * Copyright 2023 The KusionStack Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package initializer + +import ( + "fmt" + "strings" + "sync" + + "github.com/spf13/pflag" + "k8s.io/apimachinery/pkg/util/sets" + "sigs.k8s.io/controller-runtime/pkg/manager" +) + +// InitFunc is used to launch a particular controller. It may run additional "should I activate checks". +// Any error returned will cause the controller process to `Fatal` +// The bool indicates whether the controller was enabled. +type InitFunc func(manager.Manager) (enabled bool, err error) + +// InitOption configures how we set up initializer +type InitOption interface { + apply(*options) +} + +type options struct { + disableByDefault bool +} +type optionFunc func(*options) + +func (o optionFunc) apply(opt *options) { + o(opt) +} + +// WithDisableByDefault disable controller by default +func WithDisableByDefault() InitOption { + return optionFunc(func(o *options) { + o.disableByDefault = true + }) +} + +// Interface knows how to set up controllers with manager +type Interface interface { + // Add add new controller setup function to initializer. + Add(controllerName string, setup InitFunc, options ...InitOption) error + + // KnownControllers returns a slice of strings describing the ControllerInitialzier's known controllers. + KnownControllers() []string + + // Enabled returns true if the controller is enabled. + Enabled(name string) bool + + // SetupWithManager add all enabled controllers to manager + SetupWithManager(mgr manager.Manager) error + + // BindFlag adds a flag for setting global feature gates to the specified FlagSet. + BindFlag(fs *pflag.FlagSet) +} + +// New returns a new instance of initializer interface +func New() Interface { + return &controllerInitializer{ + initializers: make(map[string]InitFunc), + all: sets.NewString(), + enabled: sets.NewString(), + disableByDefault: sets.NewString(), + } +} + +var _ Interface = &controllerInitializer{} + +type controllerInitializer struct { + lock sync.RWMutex + + initializers map[string]InitFunc + all sets.String + enabled sets.String + disableByDefault sets.String +} + +func (m *controllerInitializer) Add(controllerName string, setup InitFunc, opts ...InitOption) error { + m.lock.Lock() + defer m.lock.Unlock() + + if m.all.Has(controllerName) { + return fmt.Errorf("controller %q has already been registered", controllerName) + } + + opt := &options{} + for _, o := range opts { + o.apply(opt) + } + + m.all.Insert(controllerName) + + if opt.disableByDefault { + m.disableByDefault.Insert(controllerName) + } else { + m.enabled.Insert(controllerName) + } + + m.initializers[controllerName] = setup + return nil +} + +func (m *controllerInitializer) BindFlag(fs *pflag.FlagSet) { + all := m.all.List() + disabled := m.disableByDefault.List() + fs.Var(m, "controllers", fmt.Sprintf(""+ + "A list of controllers to enable. '*' enables all on-by-default controllers, 'foo' enables the controller "+ + "named 'foo', '-foo' disables the controller named 'foo'.\nAll controllers: %s\nDisabled-by-default controllers: %s", + strings.Join(all, ", "), strings.Join(disabled, ", "))) +} + +// KnownControllers implements ControllerInitialzier. +func (m *controllerInitializer) KnownControllers() []string { + m.lock.RLock() + defer m.lock.RUnlock() + return m.all.List() +} + +func (m *controllerInitializer) SetupWithManager(mgr manager.Manager) error { + m.lock.RLock() + defer m.lock.RUnlock() + + for _, name := range m.enabled.List() { + _, err := m.initializers[name](mgr) + if err != nil { + return fmt.Errorf("failed to initialize controller %q: %v", name, err) + } + } + return nil +} + +func (m *controllerInitializer) Enabled(name string) bool { + m.lock.RLock() + defer m.lock.RUnlock() + + return m.enabled.Has(name) +} + +func (m *controllerInitializer) isControllerEnabled(name string, controllers []string) bool { + hasStar := false + for _, ctrl := range controllers { + if ctrl == name { + return true + } + if ctrl == "-"+name { + return false + } + if ctrl == "*" { + hasStar = true + } + } + // if we get here, there was no explicit choice + if !hasStar { + // nothing on by default + return false + } + + return !m.disableByDefault.Has(name) +} + +// Set implements pflag.Value interface. +func (m *controllerInitializer) Set(value string) error { + m.lock.Lock() + defer m.lock.Unlock() + + controllers := strings.Split(strings.TrimSpace(value), ",") + all := m.all.List() + for _, name := range all { + if m.isControllerEnabled(name, controllers) { + m.enabled.Insert(name) + } else { + m.enabled.Delete(name) + } + } + return nil +} + +// Type implements pflag.Value interface. +func (m *controllerInitializer) Type() string { + return "stringSlice" +} + +// String implements pflag.Value interface. +func (m *controllerInitializer) String() string { + m.lock.RLock() + defer m.lock.RUnlock() + + pairs := []string{} + for _, name := range m.all.List() { + if m.enabled.Has(name) { + pairs = append(pairs, name) + } else { + pairs = append(pairs, "-"+name) + } + } + return strings.Join(pairs, ",") +} diff --git a/controller/initializer/initializer_test.go b/controller/initializer/initializer_test.go new file mode 100644 index 0000000..ea9c82c --- /dev/null +++ b/controller/initializer/initializer_test.go @@ -0,0 +1,86 @@ +/** + * Copyright 2023 KusionStack Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package initializer + +import ( + "testing" + + "github.com/spf13/pflag" + "github.com/stretchr/testify/assert" + "sigs.k8s.io/controller-runtime/pkg/manager" +) + +func Test_Initialzier(t *testing.T) { + initializer := New() + err := initializer.Add("test1", testInitFunc) + assert.NoError(t, err) + err = initializer.Add("test2", testInitFunc) + assert.NoError(t, err) + err = initializer.Add("test3", testInitFunc, WithDisableByDefault()) + assert.NoError(t, err) + + controllers := initializer.KnownControllers() + assert.EqualValues(t, []string{"test1", "test2", "test3"}, controllers) + assert.True(t, initializer.Enabled("test1")) + assert.True(t, initializer.Enabled("test2")) + assert.False(t, initializer.Enabled("test3")) + + // duplicate + err = initializer.Add("test1", testInitFunc) + assert.Error(t, err) + + // test bind flag + fs := pflag.NewFlagSet("test-*", pflag.PanicOnError) + initializer.BindFlag(fs) + fs.Set("controllers", "*") + err = fs.Parse(nil) + assert.NoError(t, err) + assert.True(t, initializer.Enabled("test1")) + assert.True(t, initializer.Enabled("test2")) + assert.False(t, initializer.Enabled("test3")) + + fs = pflag.NewFlagSet("test", pflag.PanicOnError) + initializer.BindFlag(fs) + fs.Set("controllers", "test1,test2") + err = fs.Parse(nil) + assert.NoError(t, err) + assert.True(t, initializer.Enabled("test1")) + assert.True(t, initializer.Enabled("test2")) + assert.False(t, initializer.Enabled("test3")) + + fs = pflag.NewFlagSet("test", pflag.PanicOnError) + initializer.BindFlag(fs) + fs.Set("controllers", "-test1,test3") + err = fs.Parse(nil) + assert.NoError(t, err) + assert.False(t, initializer.Enabled("test1")) + assert.False(t, initializer.Enabled("test2")) + assert.True(t, initializer.Enabled("test3")) + + fs = pflag.NewFlagSet("test", pflag.PanicOnError) + initializer.BindFlag(fs) + fs.Set("controllers", "-test1") + err = fs.Parse(nil) + assert.NoError(t, err) + assert.False(t, initializer.Enabled("test1")) + assert.False(t, initializer.Enabled("test2")) + assert.False(t, initializer.Enabled("test3")) +} + +func testInitFunc(manager.Manager) (enabled bool, err error) { + return true, nil +} diff --git a/controller/mixin/mixin.go b/controller/mixin/mixin.go index c24ce13..b38cb42 100644 --- a/controller/mixin/mixin.go +++ b/controller/mixin/mixin.go @@ -1,17 +1,18 @@ -// Copyright 2023 The KusionStack Authors -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - +/** + * Copyright 2023 KusionStack Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ package mixin import ( diff --git a/controller/mixin/mixin_test.go b/controller/mixin/mixin_test.go index 8eadb2b..d38a1a3 100644 --- a/controller/mixin/mixin_test.go +++ b/controller/mixin/mixin_test.go @@ -1,16 +1,18 @@ -// Copyright 2023 The KusionStack Authors -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. +/** + * Copyright 2023 KusionStack Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ package mixin diff --git a/controller/mixin/webhook.go b/controller/mixin/webhook.go index bb84424..b6f21cb 100644 --- a/controller/mixin/webhook.go +++ b/controller/mixin/webhook.go @@ -1,16 +1,18 @@ -// Copyright 2023 The KusionStack Authors -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. +/** + * Copyright 2023 KusionStack Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ package mixin diff --git a/controller/mixin/webhook_test.go b/controller/mixin/webhook_test.go index 7d3b5c3..4344677 100644 --- a/controller/mixin/webhook_test.go +++ b/controller/mixin/webhook_test.go @@ -1,17 +1,18 @@ -// Copyright 2023 The KusionStack Authors -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - +/** + * Copyright 2023 KusionStack Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ package mixin import ( diff --git a/go.mod b/go.mod index dac62c4..ef12f60 100644 --- a/go.mod +++ b/go.mod @@ -7,6 +7,7 @@ require ( github.com/onsi/ginkgo v1.16.5 github.com/onsi/gomega v1.27.6 github.com/prometheus/client_golang v1.16.0 + github.com/spf13/pflag v1.0.5 github.com/stretchr/testify v1.8.2 k8s.io/api v0.28.4 k8s.io/apimachinery v0.28.4 @@ -42,7 +43,6 @@ require ( github.com/prometheus/client_model v0.4.0 // indirect github.com/prometheus/common v0.44.0 // indirect github.com/prometheus/procfs v0.10.1 // indirect - github.com/spf13/pflag v1.0.5 // indirect go.uber.org/atomic v1.10.0 // indirect go.uber.org/multierr v1.11.0 // indirect go.uber.org/zap v1.19.0 // indirect