workload repo: Merge Workload Repository into master

cmd/tidb-server/BUILD.bazel

@@ -56,6 +56,7 @@ go_library(
         "//pkg/util/tiflashcompute",
         "//pkg/util/topsql",
         "//pkg/util/versioninfo",
+        "//pkg/util/workloadrepo",
         "@com_github_opentracing_opentracing_go//:opentracing-go",
         "@com_github_pingcap_errors//:errors",
         "@com_github_pingcap_failpoint//:failpoint",

cmd/tidb-server/main.go

@@ -80,6 +80,7 @@ import (
 	"github.com/pingcap/tidb/pkg/util/tiflashcompute"
 	"github.com/pingcap/tidb/pkg/util/topsql"
 	"github.com/pingcap/tidb/pkg/util/versioninfo"
+	repository "github.com/pingcap/tidb/pkg/util/workloadrepo"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/push"
 	"github.com/tikv/client-go/v2/tikv"
@@ -318,6 +319,7 @@ func main() {
 	executor.Start()
 	resourcemanager.InstanceResourceManager.Start()
 	storage, dom := createStoreDDLOwnerMgrAndDomain(keyspaceName)
+	repository.SetupRepository(dom)
 	svr := createServer(storage, dom)
 
 	exited := make(chan struct{})
@@ -921,6 +923,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()
 	closeDDLOwnerMgrDomainAndStorage(storage, dom)
 	disk.CleanUp()
 	closeStmtSummary()

pkg/ddl/executor.go

@@ -4066,6 +4066,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

@@ -2009,7 +2009,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)
 	err = owner.CampaignOwner()
 	if err != nil {
 		logutil.BgLogger().Warn("campaign owner failed", zap.Error(err))
@@ -2337,7 +2337,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))
 	if config.GetGlobalConfig().Instance.TiDBEnableStatsOwner.Load() {
 		err := do.statsOwner.CampaignOwner()
@@ -2446,7 +2446,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/util/workloadrepo/BUILD.bazel

@@ -0,0 +1,56 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "workloadrepo",
+    srcs = [
+        "const.go",
+        "housekeeper.go",
+        "sampling.go",
+        "snapshot.go",
+        "table.go",
+        "utils.go",
+        "worker.go",
+    ],
+    importpath = "github.com/pingcap/tidb/pkg/util/workloadrepo",
+    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_zap//:zap",
+    ],
+)
+
+go_test(
+    name = "workloadrepo_test",
+    timeout = "short",
+    srcs = ["worker_test.go"],
+    embed = [":workloadrepo"],
+    flaky = True,
+    deps = [
+        "//pkg/domain",
+        "//pkg/kv",
+        "//pkg/owner",
+        "//pkg/parser/model",
+        "//pkg/testkit",
+        "@com_github_stretchr_testify//require",
+        "@io_etcd_go_etcd_client_v3//:client",
+        "@io_etcd_go_etcd_server_v3//embed",
+    ],
+)

pkg/util/workloadrepo/const.go

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

pkg/util/workloadrepo/housekeeper.go

@@ -0,0 +1,145 @@
+// 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 workloadrepo
+
+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/zap"
+)
+
+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 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("workload 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 = execRetry(ctx, sess, sb.String())
+			if err != nil {
+				logutil.BgLogger().Info("workload 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 {
+	w.Lock()
+	retention := int(w.retentionDays)
+	w.Unlock()
+
+	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("workload 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("workload 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 = execRetry(ctx, sess, sb.String())
+			if err != nil {
+				logutil.BgLogger().Info("workload 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 := 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))
+			}
+		}
+	}
+}