extension: init tidb new extension framework (#38497)

close #38496

build/nogo_config.json

@@ -213,7 +213,8 @@
       "server/conn_stmt.go": "only for server/conn_stmt.go",
       "server/conn_test.go": "only for server/conn_test.go",
       "planner/core/plan.go": "only for planner/core/plan.go",
-      "errno/": "only for errno/"
+      "errno/": "only for errno/",
+      "extension/": "extension code"
     }
   },
   "gofmt": {
@@ -331,7 +332,8 @@
       "kv/": "kv code",
       "util/memory": "util/memory",
       "ddl/": "ddl",
-      "planner/": "planner"
+      "planner/": "planner",
+      "extension/": "extension code"
     }
   },
   "pkgfact": {
@@ -398,7 +400,8 @@
       "planner/core/util.go": "planner/core/util.go",
       "util/": "util code",
       "parser/": "parser code",
-      "meta/": "parser code"
+      "meta/": "parser code",
+      "extension/": "extension code"
     }
   },
   "shift": {
@@ -756,7 +759,8 @@
       "planner/core/plan.go": "planner/core/plan.go",
       "server/conn.go": "server/conn.go",
       "server/conn_stmt.go": "server/conn_stmt.go",
-      "server/conn_test.go": "server/conn_test.go"
+      "server/conn_test.go": "server/conn_test.go",
+      "extension/": "extension code"
     }
   },
   "SA2000": {

extension/BUILD.bazel

@@ -0,0 +1,29 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "extension",
+    srcs = [
+        "extensions.go",
+        "manifest.go",
+        "registry.go",
+        "util.go",
+    ],
+    importpath = "github.com/pingcap/tidb/extension",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//sessionctx/variable",
+        "@com_github_pingcap_errors//:errors",
+    ],
+)
+
+go_test(
+    name = "extension_test",
+    srcs = ["registry_test.go"],
+    deps = [
+        ":extension",
+        "//privilege/privileges",
+        "//sessionctx/variable",
+        "@com_github_pingcap_errors//:errors",
+        "@com_github_stretchr_testify//require",
+    ],
+)

extension/extensions.go

@@ -0,0 +1,30 @@
+// Copyright 2022 PingCAP, Inc.
+//
+// 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 extension
+
+// Extensions contains all extensions that have already setup
+type Extensions struct {
+	manifests []*Manifest
+}
+
+// Manifests returns a extension manifests
+func (es *Extensions) Manifests() []*Manifest {
+	if es == nil {
+		return nil
+	}
+	manifests := make([]*Manifest, len(es.manifests))
+	copy(manifests, es.manifests)
+	return manifests
+}

extension/manifest.go

@@ -0,0 +1,136 @@
+// Copyright 2022 PingCAP, Inc.
+//
+// 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 extension
+
+import (
+	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/sessionctx/variable"
+)
+
+// Option represents an option to initialize an extension
+type Option func(m *Manifest)
+
+// WithCustomSysVariables specifies custom variables of an extension
+func WithCustomSysVariables(vars []*variable.SysVar) Option {
+	return func(m *Manifest) {
+		m.sysVariables = vars
+	}
+}
+
+// WithCustomDynPrivs specifies dynamic privileges of an extension
+func WithCustomDynPrivs(privs []string) Option {
+	return func(m *Manifest) {
+		m.dynPrivs = privs
+	}
+}
+
+// WithClose specifies the close function of an extension.
+// It will be invoked when `extension.Reset` is called
+func WithClose(fn func()) Option {
+	return func(m *Manifest) {
+		m.close = fn
+	}
+}
+
+// Manifest is an extension's manifest
+type Manifest struct {
+	name         string
+	sysVariables []*variable.SysVar
+	dynPrivs     []string
+	close        func()
+}
+
+// Name returns the extension's name
+func (m *Manifest) Name() string {
+	return m.name
+}
+
+func newManifestWithSetup(name string, factory func() ([]Option, error)) (_ *Manifest, _ func(), err error) {
+	clearBuilder := &clearFuncBuilder{}
+	defer func() {
+		if err != nil {
+			clearBuilder.Build()()
+		}
+	}()
+
+	// new manifest with factory
+	m := &Manifest{name: name}
+	err = clearBuilder.DoWithCollectClear(func() (func(), error) {
+		options, err := factory()
+		if err != nil {
+			return nil, err
+		}
+
+		for _, opt := range options {
+			opt(m)
+		}
+
+		return m.close, nil
+	})
+
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// setup dynamic privileges
+	for i := range m.dynPrivs {
+		priv := m.dynPrivs[i]
+		err = clearBuilder.DoWithCollectClear(func() (func(), error) {
+			if err = RegisterDynamicPrivilege(priv); err != nil {
+				return nil, err
+			}
+			return func() {
+				RemoveDynamicPrivilege(priv)
+			}, nil
+		})
+		if err != nil {
+			return nil, nil, err
+		}
+	}
+
+	// setup sys vars
+	for i := range m.sysVariables {
+		sysVar := m.sysVariables[i]
+		err = clearBuilder.DoWithCollectClear(func() (func(), error) {
+			if sysVar == nil {
+				return nil, errors.New("system var should not be nil")
+			}
+
+			if sysVar.Name == "" {
+				return nil, errors.New("system var name should not be empty")
+			}
+
+			if variable.GetSysVar(sysVar.Name) != nil {
+				return nil, errors.Errorf("system var '%s' has already registered", sysVar.Name)
+			}
+
+			variable.RegisterSysVar(sysVar)
+			return func() {
+				variable.UnregisterSysVar(sysVar.Name)
+			}, nil
+		})
+
+		if err != nil {
+			return nil, nil, err
+		}
+	}
+	return m, clearBuilder.Build(), nil
+}
+
+// RegisterDynamicPrivilege is used to resolve dependency cycle
+var RegisterDynamicPrivilege func(string) error
+
+// RemoveDynamicPrivilege is used to resolve dependency cycle
+var RemoveDynamicPrivilege func(string) bool

extension/registry.go

@@ -0,0 +1,156 @@
+// Copyright 2022 PingCAP, Inc.
+//
+// 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 extension
+
+import (
+	"sort"
+	"sync"
+
+	"github.com/pingcap/errors"
+)
+
+type registry struct {
+	sync.RWMutex
+	factories      map[string]func() ([]Option, error)
+	extensionNames []string
+	setup          bool
+	extensions     *Extensions
+	close          func()
+}
+
+// Extensions returns the extensions after setup
+func (r *registry) Extensions() (*Extensions, error) {
+	r.RLock()
+	defer r.RUnlock()
+	if !r.setup {
+		return nil, errors.New("The extensions has not been setup")
+	}
+	return r.extensions, nil
+}
+
+// RegisterFactory registers a new extension with a factory
+func (r *registry) RegisterFactory(name string, factory func() ([]Option, error)) error {
+	r.Lock()
+	defer r.Unlock()
+
+	if r.setup {
+		return errors.New("Cannot register new extension because registry has already been setup")
+	}
+
+	if name == "" {
+		return errors.New("extension name should not be empty")
+	}
+
+	if _, ok := r.factories[name]; ok {
+		return errors.Errorf("extension with name '%s' already registered", name)
+	}
+
+	if r.factories == nil {
+		r.factories = make(map[string]func() ([]Option, error))
+	}
+
+	r.factories[name] = factory
+	r.extensionNames = append(r.extensionNames, name)
+	sort.Strings(r.extensionNames)
+	return nil
+}
+
+// Setup setups all extensions
+func (r *registry) Setup() (err error) {
+	r.Lock()
+	defer r.Unlock()
+	if r.setup {
+		return nil
+	}
+
+	if len(r.factories) == 0 {
+		r.extensions = nil
+		r.setup = true
+		return nil
+	}
+
+	clearBuilder := &clearFuncBuilder{}
+	defer func() {
+		if err != nil {
+			clearBuilder.Build()()
+		}
+	}()
+
+	manifests := make([]*Manifest, 0, len(r.factories))
+	for i := range r.extensionNames {
+		name := r.extensionNames[i]
+		err = clearBuilder.DoWithCollectClear(func() (func(), error) {
+			factory := r.factories[name]
+			m, clear, err := newManifestWithSetup(name, factory)
+			if err != nil {
+				return nil, err
+			}
+			manifests = append(manifests, m)
+			return clear, nil
+		})
+
+		if err != nil {
+			return err
+		}
+	}
+	r.extensions = &Extensions{manifests: manifests}
+	r.setup = true
+	r.close = clearBuilder.Build()
+	return nil
+}
+
+// Reset resets the registry. It is only used by test
+func (r *registry) Reset() {
+	r.Lock()
+	defer r.Unlock()
+
+	if r.close != nil {
+		r.close()
+		r.close = nil
+	}
+	r.factories = nil
+	r.extensionNames = nil
+	r.extensions = nil
+	r.setup = false
+}
+
+var globalRegistry registry
+
+// RegisterFactory registers a new extension with a factory
+func RegisterFactory(name string, factory func() ([]Option, error)) error {
+	return globalRegistry.RegisterFactory(name, factory)
+}
+
+// Register registers a new extension with options
+func Register(name string, options ...Option) error {
+	return RegisterFactory(name, func() ([]Option, error) {
+		return options, nil
+	})
+}
+
+// Setup setups extensions
+func Setup() error {
+	return globalRegistry.Setup()
+}
+
+// GetExtensions returns all extensions after setup
+func GetExtensions() (*Extensions, error) {
+	return globalRegistry.Extensions()
+}
+
+// Reset resets the registry. It is only used by test
+func Reset() {
+	globalRegistry.Reset()
+}