*: add sql variable @@tidb_accelerate_user_creation_update (#58512)

ref #55563

Author: tiancaiamao

pkg/domain/domain.go

@@ -1945,6 +1945,8 @@ func (do *Domain) LoadPrivilegeLoop(sctx sessionctx.Context) error {
 
 func privReloadEvent(h *privileges.Handle, event *PrivilegeEvent) (err error) {
 	switch {
+	case !vardef.AccelerateUserCreationUpdate.Load():
+		err = h.UpdateAll()
 	case event.All:
 		err = h.UpdateAllActive()
 	default:
@@ -2940,10 +2942,7 @@ func (do *Domain) notifyUpdatePrivilege(event PrivilegeEvent) error {
 		return nil
 	}
 
-	if event.All {
-		return do.PrivilegeHandle().UpdateAll()
-	}
-	return do.PrivilegeHandle().Update(event.UserList)
+	return privReloadEvent(do.PrivilegeHandle(), &event)
 }
 
 // NotifyUpdateSysVarCache updates the sysvar cache key in etcd, which other TiDB

pkg/privilege/privileges/BUILD.bazel

@@ -60,6 +60,7 @@ go_test(
     shard_count = 50,
     deps = [
         "//pkg/config",
+        "//pkg/domain",
         "//pkg/errno",
         "//pkg/kv",
         "//pkg/parser/auth",

pkg/privilege/privileges/cache.go

@@ -835,6 +835,10 @@ func (p *MySQLPrivilege) LoadDBTable(ctx sqlexec.RestrictedSQLExecutor) error {
 	if err != nil {
 		return err
 	}
+	p.db.Ascend(func(itm itemDB) bool {
+		slices.SortFunc(itm.data, compareDBRecord)
+		return true
+	})
 	return nil
 }
 
@@ -2013,8 +2017,8 @@ type Handle struct {
 	// Only load the active user's data to save memory
 	// username => struct{}
 	activeUsers sync.Map
-
-	globalVars variable.GlobalVarAccessor
+	fullData    atomic.Bool
+	globalVars  variable.GlobalVarAccessor
 }
 
 // NewHandle returns a Handle.
@@ -2033,6 +2037,10 @@ func (h *Handle) ensureActiveUser(ctx context.Context, user string) error {
 		visited := p.(*bool)
 		*visited = true
 	}
+	if h.fullData.Load() {
+		// All users data are in-memory, nothing to do
+		return nil
+	}
 
 	_, exist := h.activeUsers.Load(user)
 	if exist {
@@ -2067,7 +2075,7 @@ func (h *Handle) Get() *MySQLPrivilege {
 	return h.priv.Load()
 }
 
-// UpdateAll loads all the active users' privilege info from kv storage.
+// UpdateAll loads all the users' privilege info from kv storage.
 func (h *Handle) UpdateAll() error {
 	logutil.BgLogger().Warn("update all called")
 	priv := newMySQLPrivilege()
@@ -2076,11 +2084,13 @@ func (h *Handle) UpdateAll() error {
 		return errors.Trace(err)
 	}
 	h.priv.Store(priv)
+	h.fullData.Store(true)
 	return nil
 }
 
 // UpdateAllActive loads all the active users' privilege info from kv storage.
 func (h *Handle) UpdateAllActive() error {
+	h.fullData.Store(false)
 	userList := make([]string, 0, 20)
 	h.activeUsers.Range(func(key, _ any) bool {
 		userList = append(userList, key.(string))
@@ -2103,6 +2113,7 @@ func (h *Handle) UpdateAllActive() error {
 
 // Update loads the privilege info from kv storage for the list of users.
 func (h *Handle) Update(userList []string) error {
+	h.fullData.Store(false)
 	if len(userList) > 100 {
 		logutil.BgLogger().Warn("update user list is long", zap.Int("len", len(userList)))
 	}

pkg/privilege/privileges/privileges_test.go

@@ -30,6 +30,7 @@ import (
 	"time"
 
 	"github.com/pingcap/tidb/pkg/config"
+	"github.com/pingcap/tidb/pkg/domain"
 	"github.com/pingcap/tidb/pkg/errno"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/parser/auth"
@@ -2134,6 +2135,7 @@ func TestEnsureActiveUserCoverage(t *testing.T) {
 		{"set password for test = 'test2'", false},
 		{"show create user test", false},
 		{"create user test1", false},
+		{"grant select on test.* to test1", false},
 		{"show grants", true},
 		{"show grants for 'test'@'%'", true},
 	}
@@ -2151,3 +2153,43 @@ func TestEnsureActiveUserCoverage(t *testing.T) {
 		require.Equal(t, c.visited, visited, comment)
 	}
 }
+
+func TestSQLVariableAccelerateUserCreationUpdate(t *testing.T) {
+	store := createStoreAndPrepareDB(t)
+	tk := testkit.NewTestKit(t, store)
+	dom := domain.GetDomain(tk.Session())
+	// 1. check the default variable value
+	tk.MustQuery("select @@global.tidb_accelerate_user_creation_update").Check(testkit.Rows("0"))
+	// trigger priv reload
+	tk.MustExec("create user aaa")
+	handle := dom.PrivilegeHandle()
+	handle.CheckFullData(t, true)
+	priv := handle.Get()
+	require.False(t, priv.RequestVerification(nil, "bbb", "%", "test", "", "", mysql.SelectPriv))
+
+	// 2. change the variable and check
+	tk.MustExec("set @@global.tidb_accelerate_user_creation_update = on")
+	tk.MustQuery("select @@global.tidb_accelerate_user_creation_update").Check(testkit.Rows("1"))
+	require.True(t, vardef.AccelerateUserCreationUpdate.Load())
+	tk.MustExec("create user bbb")
+	handle.CheckFullData(t, false)
+	// trigger priv reload, but data for bbb is not really loaded
+	tk.MustExec("grant select on test.* to bbb")
+	priv = handle.Get()
+	// data for bbb is not loaded, because that user is not active
+	// So this is **counterintuitive**, but it's still the expected behavior.
+	require.False(t, priv.RequestVerification(nil, "bbb", "%", "test", "", "", mysql.SelectPriv))
+	tk1 := testkit.NewTestKit(t, store)
+	// if user bbb login, everything works as expected
+	require.NoError(t, tk1.Session().Auth(&auth.UserIdentity{Username: "bbb", Hostname: "localhost"}, nil, nil, nil))
+	priv = handle.Get()
+	require.True(t, priv.RequestVerification(nil, "bbb", "%", "test", "", "", mysql.SelectPriv))
+
+	// 3. change the variable and check again
+	tk.MustExec("set @@global.tidb_accelerate_user_creation_update = off")
+	tk.MustQuery("select @@global.tidb_accelerate_user_creation_update").Check(testkit.Rows("0"))
+	tk.MustExec("drop user aaa")
+	handle.CheckFullData(t, true)
+	priv = handle.Get()
+	require.True(t, priv.RequestVerification(nil, "bbb", "%", "test", "", "", mysql.SelectPriv))
+}

pkg/privilege/privileges/tidb_auth_token_test.go

@@ -486,4 +486,8 @@ func (p *MySQLPrivilege) RoleGraph() map[string]roleGraphEdgesTable {
 	return p.roleGraph
 }
 
+func (h *Handle) CheckFullData(t *testing.T, value bool) {
+	require.True(t, h.fullData.Load() == value)
+}
+
 var NewMySQLPrivilege = newMySQLPrivilege

pkg/sessionctx/vardef/tidb_vars.go

@@ -983,6 +983,9 @@ const (
 	// TiDBEnableSharedLockPromotion indicates whether the `select for share` statement would be executed
 	// as `select for update` statements which do acquire pessimistic locks.
 	TiDBEnableSharedLockPromotion = "tidb_enable_shared_lock_promotion"
+
+	// TiDBAccelerateUserCreationUpdate decides whether tidb will load & update the whole user's data in-memory.
+	TiDBAccelerateUserCreationUpdate = "tidb_accelerate_user_creation_update"
 )
 
 // TiDB vars that have only global scope
@@ -1568,6 +1571,7 @@ const (
 	DefTiDBEnableSharedLockPromotion                  = false
 	DefTiDBTSOClientRPCMode                           = TSOClientRPCModeDefault
 	DefTiDBCircuitBreakerPDMetaErrorRatePct           = 0
+	DefTiDBAccelerateUserCreationUpdate               = false
 )
 
 // Process global variables.
@@ -1690,8 +1694,9 @@ var (
 	IgnoreInlistPlanDigest          = atomic.NewBool(DefTiDBIgnoreInlistPlanDigest)
 	TxnEntrySizeLimit               = atomic.NewUint64(DefTiDBTxnEntrySizeLimit)
 
-	SchemaCacheSize           = atomic.NewUint64(DefTiDBSchemaCacheSize)
-	SchemaCacheSizeOriginText = atomic.NewString(strconv.Itoa(DefTiDBSchemaCacheSize))
+	SchemaCacheSize              = atomic.NewUint64(DefTiDBSchemaCacheSize)
+	SchemaCacheSizeOriginText    = atomic.NewString(strconv.Itoa(DefTiDBSchemaCacheSize))
+	AccelerateUserCreationUpdate = atomic.NewBool(DefTiDBAccelerateUserCreationUpdate)
 )
 
 func serverMemoryLimitDefaultValue() string {

pkg/sessionctx/variable/sysvar.go

@@ -3435,6 +3435,12 @@ var defaultSysVars = []*SysVar{
 			return nil
 		},
 	},
+	{Scope: vardef.ScopeGlobal, Name: vardef.TiDBAccelerateUserCreationUpdate, Value: BoolToOnOff(vardef.DefTiDBAccelerateUserCreationUpdate), Type: vardef.TypeBool,
+		SetGlobal: func(_ context.Context, s *SessionVars, val string) error {
+			vardef.AccelerateUserCreationUpdate.Store(TiDBOptOn(val))
+			return nil
+		},
+	},
 }
 
 // GlobalSystemVariableInitialValue gets the default value for a system variable including ones that are dynamically set (e.g. based on the store)