domain: add idle GC for cross-keyspace runtimes (#69013)

ref #68883

Author: D3Hunter

pkg/domain/crossks/BUILD.bazel

@@ -51,10 +51,11 @@ go_test(
     ],
     embed = [":crossks"],
     flaky = True,
-    shard_count = 4,
+    shard_count = 7,
     deps = [
         "//pkg/config",
         "//pkg/config/kerneltype",
+        "//pkg/ddl/schemaver",
         "//pkg/dxf/framework/storage",
         "//pkg/executor/importer",
         "//pkg/infoschema",

pkg/domain/crossks/cross_ks.go

@@ -48,7 +48,11 @@ import (
 	"go.uber.org/zap"
 )
 
-const crossKSSessPoolSize = 5
+const (
+	crossKSSessPoolSize         = 5
+	crossKSRuntimeIdleTimeout   = 30 * time.Minute
+	crossKSRuntimeSweepInterval = time.Minute
+)
 
 type runtimeEntry struct {
 	sessMgr       *SessionManager
@@ -326,7 +330,7 @@ func (*Manager) createSessionManager(
 
 // release releases the runtime handle for the specified keyspace and holderID.
 // the resources will be cleaned up if there is no active holder after enough
-// time. we will impl this part later.
+// time.
 func (m *Manager) release(targetKS string, holderID string) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
@@ -345,6 +349,57 @@ func (m *Manager) release(targetKS string, holderID string) {
 		zap.Int("activeHolderCount", len(entry.activeHolders)))
 }
 
+// RunSystemKSGCLoop periodically evicts idle cross keyspace runtimes.
+// as the name noted, this loop only runs in SYSTEM keyspace, user keyspace
+// only access the SYSTEM ks, and will be kept alive.
+// Note: there are raw caller to GetOrCreate which conflicts with the GC loop.
+// we will make the API more clear in the future after we can refactor those
+// callers to use Acquire instead of GetOrCreate directly.
+func (m *Manager) RunSystemKSGCLoop(ctx context.Context) {
+	interval := crossKSRuntimeSweepInterval
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			m.sweepIdleRuntimes(crossKSRuntimeIdleTimeout)
+		}
+	}
+}
+
+func (m *Manager) sweepIdleRuntimes(idleTimeout time.Duration) {
+	type evictedRuntime struct {
+		targetKS string
+		entry    *runtimeEntry
+	}
+
+	evicted := make([]evictedRuntime, 0, 1)
+	now := time.Now()
+	m.mu.Lock()
+	for targetKS, entry := range m.runtimes {
+		if len(entry.activeHolders) > 0 || entry.lastReleaseAt.IsZero() {
+			continue
+		}
+		if now.Sub(entry.lastReleaseAt) < idleTimeout {
+			continue
+		}
+		delete(m.runtimes, targetKS)
+		evicted = append(evicted, evictedRuntime{targetKS: targetKS, entry: entry})
+	}
+	m.mu.Unlock()
+
+	for _, item := range evicted {
+		logutil.BgLogger().Info("evict idle cross keyspace runtime",
+			zap.String("targetKS", item.targetKS),
+			zap.Duration("idleTimeout", idleTimeout),
+			zap.Time("lastReleaseAt", item.entry.lastReleaseAt))
+		item.entry.sessMgr.close()
+	}
+}
+
 // Close closes all session managers and their associated resources.
 func (m *Manager) Close() {
 	m.mu.Lock()

pkg/domain/crossks/cross_ks_internal_test.go

@@ -15,52 +15,76 @@
 package crossks
 
 import (
+	"context"
 	"fmt"
 	"testing"
+	"time"
 
 	"github.com/ngaut/pools"
 	"github.com/pingcap/tidb/pkg/config/kerneltype"
+	"github.com/pingcap/tidb/pkg/ddl/schemaver"
 	"github.com/pingcap/tidb/pkg/infoschema/validatorapi"
 	"github.com/pingcap/tidb/pkg/keyspace"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/testkit/testfailpoint"
 	"github.com/pingcap/tidb/pkg/util"
 	"github.com/stretchr/testify/require"
+	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
 type runtimeHandleTestStore struct {
 	kv.Storage
-	ks string
+	ks         string
+	closeCount int
 }
 
 func (s *runtimeHandleTestStore) GetKeyspace() string {
 	return s.ks
 }
 
+func (s *runtimeHandleTestStore) Close() error {
+	s.closeCount++
+	return nil
+}
+
 type runtimeHandleTestSessPool struct {
 	util.DestroyableSessionPool
 	closeCount int
+	onClose    func()
 }
 
 func (p *runtimeHandleTestSessPool) Close() {
 	p.closeCount++
+	if p.onClose != nil {
+		p.onClose()
+	}
 }
 
 func newRuntimeHandleTestManager(targetKS string) (*Manager, *runtimeEntry, *runtimeHandleTestStore, *runtimeHandleTestSessPool) {
 	mgr := NewManager(&runtimeHandleTestStore{ks: keyspace.System})
 	targetStore := &runtimeHandleTestStore{ks: targetKS}
 	sessPool := &runtimeHandleTestSessPool{}
 	entry := &runtimeEntry{
-		sessMgr: &SessionManager{
-			store:    targetStore,
-			sessPool: sessPool,
-		},
+		sessMgr:       newRuntimeHandleTestSessionManager(targetStore, sessPool),
 		activeHolders: make(map[string]struct{}),
 	}
 	mgr.runtimes[targetKS] = entry
 	return mgr, entry, targetStore, sessPool
 }
 
+func newRuntimeHandleTestSessionManager(targetStore *runtimeHandleTestStore, sessPool *runtimeHandleTestSessPool) *SessionManager {
+	ctx, cancel := context.WithCancel(context.Background())
+	return &SessionManager{
+		ctx:             ctx,
+		cancel:          cancel,
+		exitCh:          make(chan struct{}),
+		store:           targetStore,
+		etcdCli:         clientv3.NewCtxClient(context.Background()),
+		schemaVerSyncer: schemaver.NewMemSyncer(),
+		sessPool:        sessPool,
+	}
+}
+
 func unusedRuntimeHandleFactoryGetter(t *testing.T) func(string, validatorapi.Validator) pools.Factory {
 	return func(string, validatorapi.Validator) pools.Factory {
 		t.Fatal("test should use the pre-seeded runtime entry")
@@ -238,3 +262,126 @@ func TestAcquireRuntimeHandle(t *testing.T) {
 		require.Zero(t, sessPool.closeCount)
 	})
 }
+
+func TestEvictRuntime(t *testing.T) {
+	if kerneltype.IsClassic() {
+		t.Skip("cross keyspace runtime acquire is supported only in nextgen kernel")
+	}
+
+	t.Run("skips active holders", func(t *testing.T) {
+		targetKS := "ks-evict-active"
+		mgr, entry, _, sessPool := newRuntimeHandleTestManager(targetKS)
+		factoryGetter := unusedRuntimeHandleFactoryGetter(t)
+
+		first, err := mgr.Acquire(targetKS, "holder-1", factoryGetter)
+		require.NoError(t, err)
+		second, err := mgr.Acquire(targetKS, "holder-2", factoryGetter)
+		require.NoError(t, err)
+
+		first.Release()
+		entry.lastReleaseAt = time.Now().Add(-crossKSRuntimeIdleTimeout - time.Second)
+		mgr.sweepIdleRuntimes(crossKSRuntimeIdleTimeout)
+
+		_, ok := mgr.Get(targetKS)
+		require.True(t, ok)
+		require.Contains(t, entry.activeHolders, "holder-2")
+		require.Zero(t, sessPool.closeCount)
+		second.Release()
+	})
+
+	t.Run("closes idle entry outside manager lock", func(t *testing.T) {
+		targetKS := "ks-evict-idle"
+		mgr, entry, targetStore, sessPool := newRuntimeHandleTestManager(targetKS)
+		factoryGetter := unusedRuntimeHandleFactoryGetter(t)
+		sessPool.onClose = func() {
+			require.True(t, mgr.mu.TryLock())
+			mgr.mu.Unlock()
+		}
+
+		handle, err := mgr.Acquire(targetKS, "holder-1", factoryGetter)
+		require.NoError(t, err)
+		handle.Release()
+		entry.lastReleaseAt = time.Now().Add(-crossKSRuntimeIdleTimeout - time.Second)
+
+		mgr.sweepIdleRuntimes(crossKSRuntimeIdleTimeout)
+
+		_, ok := mgr.Get(targetKS)
+		require.False(t, ok)
+		require.Equal(t, 1, sessPool.closeCount)
+		require.Equal(t, 1, targetStore.closeCount)
+	})
+
+	t.Run("reacquire creates new runtime", func(t *testing.T) {
+		targetKS := "ks-evict-reacquire"
+		mgr, entry, oldStore, oldSessPool := newRuntimeHandleTestManager(targetKS)
+		factoryGetter := unusedRuntimeHandleFactoryGetter(t)
+		createCount := 0
+		testfailpoint.EnableCall(t, "github.com/pingcap/tidb/pkg/domain/crossks/mockCreateSessionManager",
+			func(createSessionManager *func(string, func(string, validatorapi.Validator) pools.Factory) (*SessionManager, error)) {
+				*createSessionManager = func(ks string, _ func(string, validatorapi.Validator) pools.Factory) (*SessionManager, error) {
+					createCount++
+					newStore := &runtimeHandleTestStore{ks: ks}
+					newSessPool := &runtimeHandleTestSessPool{}
+					return newRuntimeHandleTestSessionManager(newStore, newSessPool), nil
+				}
+			},
+		)
+
+		handle, err := mgr.Acquire(targetKS, "holder-1", factoryGetter)
+		require.NoError(t, err)
+		handle.Release()
+		entry.lastReleaseAt = time.Now().Add(-crossKSRuntimeIdleTimeout - time.Second)
+
+		mgr.sweepIdleRuntimes(crossKSRuntimeIdleTimeout)
+		require.Equal(t, 1, oldSessPool.closeCount)
+		require.Equal(t, 1, oldStore.closeCount)
+
+		reacquired, err := mgr.Acquire(targetKS, "holder-2", factoryGetter)
+		require.NoError(t, err)
+		require.NotSame(t, oldStore, reacquired.Store())
+		require.Equal(t, 1, createCount)
+		reacquired.Release()
+	})
+}
+
+func TestRuntimeHandleManagerCloseClosesAllEntriesRegardlessOfIdleTimeout(t *testing.T) {
+	firstKS := "ks-close-runtime-1"
+	secondKS := "ks-close-runtime-2"
+	mgr, firstEntry, firstStore, firstSessPool := newRuntimeHandleTestManager(firstKS)
+	secondStore := &runtimeHandleTestStore{ks: secondKS}
+	secondSessPool := &runtimeHandleTestSessPool{}
+	mgr.runtimes[secondKS] = &runtimeEntry{
+		sessMgr:       newRuntimeHandleTestSessionManager(secondStore, secondSessPool),
+		activeHolders: make(map[string]struct{}),
+	}
+	firstEntry.lastReleaseAt = time.Now()
+
+	mgr.Close()
+
+	require.Empty(t, mgr.GetAllKeyspace())
+	require.Equal(t, 1, firstSessPool.closeCount)
+	require.Equal(t, 1, firstStore.closeCount)
+	require.Equal(t, 1, secondSessPool.closeCount)
+	require.Equal(t, 1, secondStore.closeCount)
+}
+
+func TestGCLoopExitsWhenContextCancelled(t *testing.T) {
+	mgr := NewManager(&runtimeHandleTestStore{ks: keyspace.System})
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	go func() {
+		mgr.RunSystemKSGCLoop(ctx)
+		close(done)
+	}()
+
+	cancel()
+
+	require.Eventually(t, func() bool {
+		select {
+		case <-done:
+			return true
+		default:
+			return false
+		}
+	}, 5*time.Second, 10*time.Millisecond)
+}

pkg/domain/domain.go

@@ -852,18 +852,20 @@ func (do *Domain) Start(startMode ddl.StartMode) error {
 	return nil
 }
 
-// TODO: we should sync the system keyspace info schema, not just load it,
-// currently, we assume there is no upgrade, so we only load the info schema of
-// system keyspace once, and it will not change during the lifetime of the domain,
-// it's not right, but it's enough to push subtasks which depends on it forward,
-// we will fix it in the future.
 func (do *Domain) loadSysKSInfoSchema() error {
 	logutil.BgLogger().Info("loading system keyspace info schema")
+	// it will trigger the creation of system keyspace session manager,
+	// which will load the info schema cache.
 	_, err := do.GetKSStore(keyspace.System)
 	return err
 }
 
 // GetKSStore returns the kv.Storage for the given keyspace.
+// we should forbid direct access cross KS component through Domain. we should
+// use AcquireKSRuntime to manage their lifecycle.
+// but Session dependents on Domain, to create a session pool we need to access the
+// GetKSStore/GetKSInfoCache inside Session where we don't know the runtime holder.
+// and trying to refactor that part can cause import cycle easily.
 func (do *Domain) GetKSStore(targetKS string) (store kv.Storage, err error) {
 	mgr, err := do.crossKSSessMgr.GetOrCreate(targetKS, do.crossKSSessFactoryGetter)
 	if err != nil {
@@ -873,6 +875,7 @@ func (do *Domain) GetKSStore(targetKS string) (store kv.Storage, err error) {
 }
 
 // GetKSInfoCache returns the system keyspace info cache.
+// see comments of GetKSStore too.
 func (do *Domain) GetKSInfoCache(targetKS string) (*infoschema.InfoCache, error) {
 	mgr, err := do.crossKSSessMgr.GetOrCreate(targetKS, do.crossKSSessFactoryGetter)
 	if err != nil {