From befb1a570244b142d727a4c9232bda905976334b Mon Sep 17 00:00:00 2001 From: RemHero <1104304963@qq.com> Date: Wed, 17 Apr 2024 16:39:04 +0800 Subject: [PATCH 1/2] add memTracker for InitCheckExistence in newLockCtx Signed-off-by: RemHero <1104304963@qq.com> --- pkg/executor/executor.go | 12 ++++++++++++ pkg/util/hack/hack.go | 5 +++++ pkg/util/memory/tracker.go | 2 ++ 3 files changed, 19 insertions(+) diff --git a/pkg/executor/executor.go b/pkg/executor/executor.go index e30200b1fdc53..7d0bb7fd1eeb8 100644 --- a/pkg/executor/executor.go +++ b/pkg/executor/executor.go @@ -76,11 +76,13 @@ import ( "github.com/pingcap/tidb/pkg/util/deadlockhistory" "github.com/pingcap/tidb/pkg/util/disk" "github.com/pingcap/tidb/pkg/util/execdetails" + "github.com/pingcap/tidb/pkg/util/hack" "github.com/pingcap/tidb/pkg/util/intest" "github.com/pingcap/tidb/pkg/util/logutil" "github.com/pingcap/tidb/pkg/util/logutil/consistency" "github.com/pingcap/tidb/pkg/util/memory" "github.com/pingcap/tidb/pkg/util/resourcegrouptag" + "github.com/pingcap/tidb/pkg/util/set" "github.com/pingcap/tidb/pkg/util/sqlexec" "github.com/pingcap/tidb/pkg/util/syncutil" "github.com/pingcap/tidb/pkg/util/topsql" @@ -1210,6 +1212,8 @@ func (e *SelectLockExec) Next(ctx context.Context, req *chunk.Chunk) error { func newLockCtx(sctx sessionctx.Context, lockWaitTime int64, numKeys int) (*tikvstore.LockCtx, error) { seVars := sctx.GetSessionVars() + memTracker := memory.NewTracker(memory.LabelForLock, -1) + memTracker.AttachTo(seVars.MemTracker) forUpdateTS, err := sessiontxn.GetTxnManager(sctx).GetStmtForUpdateTS() if err != nil { return nil, err @@ -1246,8 +1250,16 @@ func newLockCtx(sctx sessionctx.Context, lockWaitTime int64, numKeys int) (*tikv rec := deadlockhistory.ErrDeadlockToDeadlockRecord(deadlock) deadlockhistory.GlobalDeadlockHistory.Push(rec) } + lockCtx.OnMemChange = func(capacity int) uint64 { + bucketMemoryUsage := hack.EstimateBucketMemoryUsageWithKVSize(lockCtx.GetValuesKSize(), lockCtx.GetValuesVSize()) + mapSize := set.EstimateMapSize(capacity, bucketMemoryUsage) + return mapSize + } if lockCtx.ForUpdateTS > 0 && seVars.AssertionLevel != variable.AssertionLevelOff { lockCtx.InitCheckExistence(numKeys) + if lockCtx.OnMemChange != nil { + memTracker.Consume(lockCtx.OnMemChange(numKeys)) + } } return lockCtx, nil } diff --git a/pkg/util/hack/hack.go b/pkg/util/hack/hack.go index 7af5e850bcb14..523cc787e51f6 100644 --- a/pkg/util/hack/hack.go +++ b/pkg/util/hack/hack.go @@ -83,3 +83,8 @@ const ( func EstimateBucketMemoryUsage[K comparable, V any]() uint64 { return (8*(1+uint64(unsafe.Sizeof(*new(K))+unsafe.Sizeof(*new(V)))) + 16) / 2 * 3 } + +// EstimateBucketMemoryUsage returns the estimated memory usage of a bucket in a map with key size and value size. +func EstimateBucketMemoryUsageWithKVSize(keySize uint64, valueSize uint64) uint64 { + return (8*(1+keySize+valueSize) + 16) / 2 * 3 +} diff --git a/pkg/util/memory/tracker.go b/pkg/util/memory/tracker.go index ec5a84c7b8ec5..874e49e36c837 100644 --- a/pkg/util/memory/tracker.go +++ b/pkg/util/memory/tracker.go @@ -860,6 +860,8 @@ const ( LabelForChunkDataInDiskByChunks int = -30 // LabelForSortPartition represents the label of the sort partition LabelForSortPartition = -31 + // LabelForLock represents the label of the lock + LabelForLock = -32 ) // MetricsTypes is used to get label for metrics From f725e45229efd6a73a0dcdd59f16161cafc24171 Mon Sep 17 00:00:00 2001 From: RemHero <1104304963@qq.com> Date: Fri, 19 Apr 2024 16:08:00 +0800 Subject: [PATCH 2/2] Add tests to OnMemChange of lockCtx Signed-off-by: RemHero <1104304963@qq.com> --- pkg/executor/executor_pkg_test.go | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/pkg/executor/executor_pkg_test.go b/pkg/executor/executor_pkg_test.go index 7070d4b3200f2..f9a4841bba290 100644 --- a/pkg/executor/executor_pkg_test.go +++ b/pkg/executor/executor_pkg_test.go @@ -16,6 +16,9 @@ package executor import ( "fmt" + "github.com/pingcap/tidb/pkg/util/hack" + "github.com/pingcap/tidb/pkg/util/set" + tikvstore "github.com/tikv/client-go/v2/kv" "runtime" "strconv" "strings" @@ -490,3 +493,18 @@ func TestErrLevelsForResetStmtContext(t *testing.T) { } } } + +func TestTikvstoreLockCtxMemConsume(t *testing.T) { + lockCtx := tikvstore.NewLockCtx(0, 0, time.Time{}) + lockCtx.OnMemChange = func(capacity int) uint64 { + bucketMemoryUsage := hack.EstimateBucketMemoryUsageWithKVSize(lockCtx.GetValuesKSize(), lockCtx.GetValuesVSize()) + mapSize := set.EstimateMapSize(capacity, bucketMemoryUsage) + return mapSize + } + numKeys := 100 + lockCtx.InitCheckExistence(numKeys) + memUsage := lockCtx.OnMemChange(numKeys) + // the type of (key,value) in lockCtx.Value is (string,ReturnedValue) with size (16,48) + expectedMemUsage := set.EstimateMapSize(numKeys, hack.EstimateBucketMemoryUsageWithKVSize(16, 48)) + require.Equal(t, expectedMemUsage, memUsage) +}