repository: Merge Workload Repository into master

cmd/tidb-server/BUILD.bazel

@@ -10,6 +10,7 @@ go_library(
         "//pkg/config",
         "//pkg/ddl",
         "//pkg/domain",
+        "//pkg/domain/repository",
         "//pkg/executor",
         "//pkg/executor/mppcoordmanager",
         "//pkg/extension",

cmd/tidb-server/main.go

@@ -34,6 +34,7 @@ import (
 	"github.com/pingcap/tidb/pkg/config"
 	"github.com/pingcap/tidb/pkg/ddl"
 	"github.com/pingcap/tidb/pkg/domain"
+	"github.com/pingcap/tidb/pkg/domain/repository"
 	"github.com/pingcap/tidb/pkg/executor"
 	"github.com/pingcap/tidb/pkg/executor/mppcoordmanager"
 	"github.com/pingcap/tidb/pkg/extension"
@@ -318,6 +319,7 @@ func main() {
 	executor.Start()
 	resourcemanager.InstanceResourceManager.Start()
 	storage, dom := createStoreAndDomain(keyspaceName)
+	repository.SetupRepository(dom)
 	svr := createServer(storage, dom)
 
 	exited := make(chan struct{})
@@ -918,6 +920,7 @@ func cleanup(svr *server.Server, storage kv.Storage, dom *domain.Domain) {
 	// See https://github.com/pingcap/tidb/issues/40038 for details.
 	svr.KillSysProcesses()
 	plugin.Shutdown(context.Background())
+	repository.StopRepository()
 	closeDomainAndStorage(storage, dom)
 	disk.CleanUp()
 	closeStmtSummary()

pkg/ddl/executor.go

@@ -4070,6 +4070,9 @@ var systemTables = map[string]struct{}{
 }
 
 func isUndroppableTable(schema, table string) bool {
+	if schema == "workload_schema" {
+		return true
+	}
 	if schema != mysql.SystemDB {
 		return false
 	}

pkg/domain/domain.go

@@ -2062,7 +2062,7 @@ func (do *Domain) LoadBindInfoLoop(ctxForHandle sessionctx.Context, ctxForEvolve
 		return err
 	}
 
-	owner := do.newOwnerManager(bindinfo.Prompt, bindinfo.OwnerKey)
+	owner := do.NewOwnerManager(bindinfo.Prompt, bindinfo.OwnerKey)
 	do.globalBindHandleWorkerLoop(owner)
 	return nil
 }
@@ -2385,7 +2385,7 @@ func (do *Domain) UpdateTableStatsLoop(ctx, initStatsCtx sessionctx.Context) err
 	}
 	variable.EnableStatsOwner = do.enableStatsOwner
 	variable.DisableStatsOwner = do.disableStatsOwner
-	do.statsOwner = do.newOwnerManager(handle.StatsPrompt, handle.StatsOwnerKey)
+	do.statsOwner = do.NewOwnerManager(handle.StatsPrompt, handle.StatsOwnerKey)
 	do.statsOwner.SetListener(owner.NewListenersWrapper(statsHandle, do.ddlNotifier))
 	do.wg.Run(func() {
 		do.indexUsageWorker()
@@ -2487,7 +2487,8 @@ func (do *Domain) StartLoadStatsSubWorkers(ctxList []sessionctx.Context) {
 	logutil.BgLogger().Info("start load stats sub workers", zap.Int("worker count", len(ctxList)))
 }
 
-func (do *Domain) newOwnerManager(prompt, ownerKey string) owner.Manager {
+// NewOwnerManager returns the owner manager for use outside of the domain.
+func (do *Domain) NewOwnerManager(prompt, ownerKey string) owner.Manager {
 	id := do.ddl.OwnerManager().ID()
 	var statsOwner owner.Manager
 	if do.etcdClient == nil {

pkg/domain/repository/BUILD.bazel

@@ -0,0 +1,39 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "repository",
+    srcs = [
+        "const.go",
+        "housekeeper.go",
+        "sampling.go",
+        "snapshot.go",
+        "table.go",
+        "utils.go",
+        "worker.go",
+    ],
+    importpath = "github.com/pingcap/tidb/pkg/domain/repository",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/domain",
+        "//pkg/domain/infosync",
+        "//pkg/infoschema",
+        "//pkg/kv",
+        "//pkg/meta/model",
+        "//pkg/owner",
+        "//pkg/parser/model",
+        "//pkg/sessionctx",
+        "//pkg/sessionctx/variable",
+        "//pkg/sessiontxn",
+        "//pkg/util",
+        "//pkg/util/chunk",
+        "//pkg/util/logutil",
+        "//pkg/util/slice",
+        "//pkg/util/sqlescape",
+        "//pkg/util/sqlexec",
+        "@com_github_ngaut_pools//:pools",
+        "@com_github_pingcap_errors//:errors",
+        "@io_etcd_go_etcd_client_v3//:client",
+        "@org_uber_go_atomic//:atomic",
+        "@org_uber_go_zap//:zap",
+    ],
+)

pkg/domain/repository/const.go

@@ -0,0 +1,38 @@
+// Copyright 2024 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 repository
+
+import (
+	"time"
+
+	"github.com/pingcap/tidb/pkg/parser/model"
+)
+
+const (
+	ownerKey  = "/tidb/repository/owner"
+	promptKey = "repository"
+	snapIDKey = "/tidb/repository/snap_id"
+
+	etcdOpTimeout = 5 * time.Second
+
+	// WorkloadSchema is the name of database for repository worker.
+	WorkloadSchema     = "WORKLOAD_SCHEMA"
+	histSnapshotsTable = "HIST_SNAPSHOTS"
+)
+
+var (
+	workloadSchemaCIStr = model.NewCIStr(WorkloadSchema)
+	zeroTime            = time.Time{}
+)

pkg/domain/repository/housekeeper.go

@@ -0,0 +1,149 @@
+// Copyright 2024 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 repository
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/pingcap/tidb/pkg/infoschema"
+	"github.com/pingcap/tidb/pkg/parser/model"
+	"github.com/pingcap/tidb/pkg/sessionctx"
+	"github.com/pingcap/tidb/pkg/sessiontxn"
+	"github.com/pingcap/tidb/pkg/util/logutil"
+	"github.com/pingcap/tidb/pkg/util/sqlescape"
+	"go.uber.org/atomic"
+	"go.uber.org/zap"
+)
+
+var (
+	retentionDays = atomic.NewInt32(-1)
+)
+
+func calcNextTick(now time.Time) time.Duration {
+	// only activated at 2am
+	next := time.Date(now.Year(), now.Month(), now.Day(), 2, 0, 0, 0, time.Local)
+	if !next.After(now) {
+		next = next.AddDate(0, 0, 1)
+	}
+	return next.Sub(now)
+}
+
+func (w *worker) createAllPartitions(ctx context.Context, sess sessionctx.Context, is infoschema.InfoSchema) error {
+	sb := &strings.Builder{}
+	for _, tbl := range workloadTables {
+		tbSchema, err := is.TableByName(ctx, workloadSchemaCIStr, model.NewCIStr(tbl.destTable))
+		if err != nil {
+			logutil.BgLogger().Info("repository cannot get table", zap.String("tbl", tbl.destTable), zap.NamedError("err", err))
+			return err
+		}
+		tbInfo := tbSchema.Meta()
+
+		sb.Reset()
+		sqlescape.MustFormatSQL(sb, "ALTER TABLE %n.%n ADD PARTITION (", WorkloadSchema, tbl.destTable)
+		if !generatePartitionRanges(sb, tbInfo) {
+			fmt.Fprintf(sb, ")")
+			_, err = w.execRetry(ctx, sess, sb.String())
+			if err != nil {
+				logutil.BgLogger().Info("repository cannot add partitions", zap.String("parts", sb.String()), zap.NamedError("err", err))
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func (w *worker) dropOldPartitions(ctx context.Context, sess sessionctx.Context, is infoschema.InfoSchema, now time.Time) error {
+	retention := int(retentionDays.Load())
+	if retention == -1 {
+		panic("Variable " + repositoryRetentionDays + " was not set before repository was started.")
+	} else if retention == 0 {
+		// disabled housekeeping
+		return nil
+	}
+
+	sb := &strings.Builder{}
+	for _, tbl := range workloadTables {
+		tbSchema, err := is.TableByName(ctx, workloadSchemaCIStr, model.NewCIStr(tbl.destTable))
+		if err != nil {
+			logutil.BgLogger().Info("repository cannot get table", zap.String("tbl", tbl.destTable), zap.NamedError("err", err))
+			continue
+		}
+		tbInfo := tbSchema.Meta()
+		for _, pt := range tbInfo.GetPartitionInfo().Definitions {
+			ot, err := time.Parse("p20060102", pt.Name.L)
+			if err != nil {
+				logutil.BgLogger().Info("repository cannot parse partition name", zap.String("part", pt.Name.L), zap.NamedError("err", err))
+				break
+			}
+			if int(now.Sub(ot).Hours()/24) < retention {
+				continue
+			}
+			sb.Reset()
+			sqlescape.MustFormatSQL(sb, "ALTER TABLE %s.%s DROP PARTITION %s",
+				WorkloadSchema, tbl.destTable, pt.Name.L)
+			_, err = w.execRetry(ctx, sess, sb.String())
+			if err != nil {
+				logutil.BgLogger().Info("repository cannot drop partition", zap.String("part", pt.Name.L), zap.NamedError("err", err))
+				break
+			}
+		}
+	}
+	return nil
+}
+
+func (w *worker) startHouseKeeper(ctx context.Context) func() {
+	return func() {
+		now := time.Now()
+		timer := time.NewTimer(calcNextTick(now))
+		defer timer.Stop()
+
+		_sessctx := w.getSessionWithRetry()
+		defer w.sesspool.Put(_sessctx)
+		sess := _sessctx.(sessionctx.Context)
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case now := <-timer.C:
+				// Owner only
+				if !w.owner.IsOwner() {
+					continue
+				}
+
+				// get the latest infoschema
+				is := sessiontxn.GetTxnManager(sess).GetTxnInfoSchema()
+
+				// create new partitions
+				if err := w.createAllPartitions(ctx, sess, is); err != nil {
+					continue
+				}
+
+				// drop old partitions
+				if err := w.dropOldPartitions(ctx, sess, is, now); err != nil {
+					continue
+				}
+
+				// reschedule, drain channel first
+				if !timer.Stop() {
+					<-timer.C
+				}
+				timer.Reset(calcNextTick(now))
+			}
+		}
+	}
+}